flux_image_to_image.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyNJFuD+VuwNdDFjxkfCkG67",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/detektor777/colab_list_image/blob/main/flux_image_to_image.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MxH_tozq-CIy"
      },
      "outputs": [],
      "source": [
        "#@title ##**Install** { display-mode: \"form\" }\n",
        "%%capture\n",
        "%cd /content\n",
        "!git clone -b totoro3 https://github.com/camenduru/ComfyUI /content/TotoroUI\n",
        "%cd /content/TotoroUI\n",
        "\n",
        "!pip install torch\n",
        "!pip install torchsde\n",
        "\n",
        "import torch\n",
        "import torchvision\n",
        "import torchaudio\n",
        "!apt -y install -qq aria2\n",
        "!apt -y install -qq aria2\n",
        "\n",
        "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/flux1-dev-fp8.safetensors -d /content/TotoroUI/models/unet -o flux1-dev-fp8.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",
        "!wget https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/test.png -O /content/test.png\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_extras import nodes_post_processing\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",
        "VAEEncode = NODE_CLASS_MAPPINGS[\"VAEEncode\"]()\n",
        "EmptyLatentImage = NODE_CLASS_MAPPINGS[\"EmptyLatentImage\"]()\n",
        "ImageScaleToTotalPixels = nodes_post_processing.NODE_CLASS_MAPPINGS[\"ImageScaleToTotalPixels\"]()\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-dev-fp8.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",
      "source": [
        "#@title ##**Upload image** { display-mode: \"form\" }\n",
        "\n",
        "from google.colab import files\n",
        "from IPython.display import Image as IPythonImage\n",
        "import io\n",
        "\n",
        "uploaded = files.upload()\n",
        "\n",
        "for filename in uploaded.keys():\n",
        "    with open('/content/input.png', 'wb') as f:\n",
        "        f.write(uploaded[filename])\n"
      ],
      "metadata": {
        "id": "5yVdZ_Jt-eGr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Run** { display-mode: \"form\" }\n",
        "from PIL import Image\n",
        "import random\n",
        "import numpy as np\n",
        "import torch\n",
        "import gc\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "\n",
        "def resize_and_pad_to_nearest_multiple(image, multiple):\n",
        "    width, height = image.size\n",
        "    new_width = ((width - 1) // multiple + 1) * multiple\n",
        "    new_height = ((height - 1) // multiple + 1) * multiple\n",
        "    image = image.copy()\n",
        "    image.thumbnail((new_width, new_height), Image.LANCZOS)\n",
        "    new_image = Image.new('RGB', (new_width, new_height))\n",
        "    left = (new_width - image.width) // 2\n",
        "    top = (new_height - image.height) // 2\n",
        "    new_image.paste(image, (left, top))\n",
        "    return new_image, (left, top, image.width, image.height)\n",
        "\n",
        "with torch.inference_mode():\n",
        "    positive_prompt = \"photo\" #@param {type:\"string\"}\n",
        "    seed = 0 #@param {type:\"integer\"}\n",
        "    steps = 30 #@param {type:\"slider\", min:1, max:100, step:1}\n",
        "    sampler_name = \"euler\" #@param [\"euler\", \"ddim\", \"heun\", \"euler_cfg_pp\", \"euler_ancestral\", \"euler_ancestral_cfg_pp\", \"heunpp2\", \"dpm_2\", \"dpm_2_ancestral\", \"dpm_fast\", \"dpm_adaptive\", \"dpmpp_2s_ancestral\", \"dpmpp_sde\", \"dpmpp_sde_gpu\", \"dpmpp_2m\", \"dpmpp_2m_sde\", \"dpmpp_2m_sde_gpu\", \"dpmpp_3m_sde\", \"dpmpp_3m_sde_gpu\", \"ddpm\", \"Icm\", \"ipndm\", \"ipndm_v\", \"deis\", \"uni_pc\", \"uni_pc_bh2\"]\n",
        "    scheduler = \"simple\" #@param [\"simple\", \"karras\", \"normal\", \"exponential\", \"sgm_uniform\", \"ddim_uniform\", \"beta\"]\n",
        "    sigma_value = 0.35 #@param {type:\"slider\", min:0.0, max:1.0, step:0.01}\n",
        "\n",
        "    if seed == 0:\n",
        "        seed = random.randint(0, 18446744073709551615)\n",
        "    print(f\"seed: {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, sigma_value)[0]\n",
        "\n",
        "    image = Image.open(\"/content/input.png\")\n",
        "    image, padding_info = resize_and_pad_to_nearest_multiple(image, 64)\n",
        "    width, height = image.size\n",
        "\n",
        "    image.save(\"/content/resized_input.png\")\n",
        "    if torch.cuda.is_available():\n",
        "      torch.cuda.empty_cache()\n",
        "    gc.collect()\n",
        "\n",
        "    image = nodes.LoadImage().load_image(\"/content/resized_input.png\")[0]\n",
        "    latent_image = VAEEncode.encode(vae, image)[0]\n",
        "\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",
        "\n",
        "    result_image_path = \"/content/flux.png\"\n",
        "    decoded_image = Image.fromarray(np.array(decoded * 255, dtype=np.uint8)[0])\n",
        "\n",
        "    left, top, width, height = padding_info\n",
        "    cropped_image = decoded_image.crop((left, top, left + width, top + height))\n",
        "    cropped_image.save(result_image_path)\n",
        "\n",
        "    from IPython.display import Image as IPythonImage\n",
        "    display(IPythonImage(filename=result_image_path))\n",
        "    if torch.cuda.is_available():\n",
        "      torch.cuda.empty_cache()\n",
        "    gc.collect()\n"
      ],
      "metadata": {
        "id": "Zo80ApBm-3Cq"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}