flux.1-schnell_jupyter.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/camenduru/flux-jupyter/blob/main/flux.1-schnell_jupyter.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VjYy0F2gZIPR"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "!git clone -b totoro3 https://github.com/camenduru/ComfyUI /content/TotoroUI\n",
        "%cd /content/TotoroUI\n",
        "\n",
        "!pip install -q torchsde einops diffusers accelerate xformers==0.0.27\n",
        "!apt -y install -qq aria2\n",
        "\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors -d /content/TotoroUI/models/unet -o flux1-schnell.safetensors\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/ae.sft -d /content/TotoroUI/models/vae -o ae.sft\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/clip_l.safetensors -d /content/TotoroUI/models/clip -o clip_l.safetensors\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/t5xxl_fp8_e4m3fn.safetensors -d /content/TotoroUI/models/clip -o t5xxl_fp8_e4m3fn.safetensors\n",
        "\n",
        "import random\n",
        "import torch\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import nodes\n",
        "from nodes import NODE_CLASS_MAPPINGS\n",
        "from totoro_extras import nodes_custom_sampler\n",
        "from totoro import model_management\n",
        "\n",
        "DualCLIPLoader = NODE_CLASS_MAPPINGS[\"DualCLIPLoader\"]()\n",
        "UNETLoader = NODE_CLASS_MAPPINGS[\"UNETLoader\"]()\n",
        "RandomNoise = nodes_custom_sampler.NODE_CLASS_MAPPINGS[\"RandomNoise\"]()\n",
        "BasicGuider = nodes_custom_sampler.NODE_CLASS_MAPPINGS[\"BasicGuider\"]()\n",
        "KSamplerSelect = nodes_custom_sampler.NODE_CLASS_MAPPINGS[\"KSamplerSelect\"]()\n",
        "BasicScheduler = nodes_custom_sampler.NODE_CLASS_MAPPINGS[\"BasicScheduler\"]()\n",
        "SamplerCustomAdvanced = nodes_custom_sampler.NODE_CLASS_MAPPINGS[\"SamplerCustomAdvanced\"]()\n",
        "VAELoader = NODE_CLASS_MAPPINGS[\"VAELoader\"]()\n",
        "VAEDecode = NODE_CLASS_MAPPINGS[\"VAEDecode\"]()\n",
        "EmptyLatentImage = NODE_CLASS_MAPPINGS[\"EmptyLatentImage\"]()\n",
        "\n",
        "with torch.inference_mode():\n",
        "    clip = DualCLIPLoader.load_clip(\"t5xxl_fp8_e4m3fn.safetensors\", \"clip_l.safetensors\", \"flux\")[0]\n",
        "    unet = UNETLoader.load_unet(\"flux1-schnell.safetensors\", \"fp8_e4m3fn\")[0]\n",
        "    vae = VAELoader.load_vae(\"ae.sft\")[0]\n",
        "\n",
        "def closestNumber(n, m):\n",
        "    q = int(n / m)\n",
        "    n1 = m * q\n",
        "    if (n * m) > 0:\n",
        "        n2 = m * (q + 1)\n",
        "    else:\n",
        "        n2 = m * (q - 1)\n",
        "    if abs(n - n1) < abs(n - n2):\n",
        "        return n1\n",
        "    return n2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "with torch.inference_mode():\n",
        "    positive_prompt = \"black forest toast spelling out the words 'FLUX DEV', tasty, food photography, dynamic shot\"\n",
        "    width = 1024\n",
        "    height = 1024\n",
        "    seed = 0\n",
        "    steps = 4\n",
        "    sampler_name = \"euler\"\n",
        "    scheduler = \"simple\"\n",
        "\n",
        "    if seed == 0:\n",
        "        seed = random.randint(0, 18446744073709551615)\n",
        "    print(seed)\n",
        "\n",
        "    cond, pooled = clip.encode_from_tokens(clip.tokenize(positive_prompt), return_pooled=True)\n",
        "    cond = [[cond, {\"pooled_output\": pooled}]]\n",
        "    noise = RandomNoise.get_noise(seed)[0] \n",
        "    guider = BasicGuider.get_guider(unet, cond)[0]\n",
        "    sampler = KSamplerSelect.get_sampler(sampler_name)[0]\n",
        "    sigmas = BasicScheduler.get_sigmas(unet, scheduler, steps, 1.0)[0]\n",
        "    latent_image = EmptyLatentImage.generate(closestNumber(width, 16), closestNumber(height, 16))[0]\n",
        "    sample, sample_denoised = SamplerCustomAdvanced.sample(noise, guider, sampler, sigmas, latent_image)\n",
        "    model_management.soft_empty_cache()\n",
        "    decoded = VAEDecode.decode(vae, sample)[0].detach()\n",
        "    Image.fromarray(np.array(decoded*255, dtype=np.uint8)[0]).save(\"/content/flux.png\")\n",
        "\n",
        "Image.fromarray(np.array(decoded*255, dtype=np.uint8)[0])"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}