Flux.1-dev-Union_CN-jupyter.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VjYy0F2gZIPR",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@markdown <center><h1>Install</h1></center>\n",
        "\n",
        "%cd /content\n",
        "!git clone -b totoro5 https://github.com/camenduru/ComfyUI /content/TotoroUI\n",
        "!git clone -b totoro https://github.com/LucipherDev/ComfyUI-GGUF /content/TotoroUI/custom_nodes/TotoroUI-GGUF\n",
        "%cd /content/TotoroUI\n",
        "\n",
        "!pip install -q torchsde einops diffusers accelerate xformers==0.0.28.post2\n",
        "!pip install -q git+https://github.com/LucipherDev/controlnet_aux.git\n",
        "!pip install -q -r /content/TotoroUI/custom_nodes/TotoroUI-GGUF/requirements.txt\n",
        "!apt -y install -qq aria2\n",
        "\n",
        "import nodes\n",
        "\n",
        "if not nodes.load_custom_node(\"custom_nodes/TotoroUI-GGUF\"):\n",
        "  raise Exception(\"Failed to load GGUF custom node\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown <center><h1>Load Models</h1></center>\n",
        "\n",
        "import torch\n",
        "from nodes import NODE_CLASS_MAPPINGS\n",
        "\n",
        "DualCLIPLoaderGGUF = NODE_CLASS_MAPPINGS[\"DualCLIPLoaderGGUF\"]()\n",
        "UnetLoaderGGUF = NODE_CLASS_MAPPINGS[\"UnetLoaderGGUF\"]()\n",
        "VAELoader = NODE_CLASS_MAPPINGS[\"VAELoader\"]()\n",
        "ControlNetLoader = NODE_CLASS_MAPPINGS[\"ControlNetLoader\"]()\n",
        "\n",
        "system_ram = \"low ram (Q2_K)\" # @param [\"low ram (Q2_K)\", \"high ram (Q8_0)\"]\n",
        "\n",
        "if system_ram == \"low ram (Q2_K)\":\n",
        "  print(f\"Downloading Flux1-dev-Q2_K...\")\n",
        "  !aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/city96/FLUX.1-dev-gguf/resolve/main/flux1-dev-Q2_K.gguf -d /content/TotoroUI/models/unet -o flux1-dev.gguf\n",
        "elif system_ram == \"high ram (Q8_0)\":\n",
        "  print(f\"Downloading Flux1-dev-Q8_0...\")\n",
        "  !aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/city96/FLUX.1-dev-gguf/resolve/main/flux1-dev-Q8_0.gguf -d /content/TotoroUI/models/unet -o flux1-dev.gguf\n",
        "\n",
        "print(\"Downloading VAE...\")\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -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",
        "\n",
        "print(\"Downloading Clips...\")\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -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 --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/city96/t5-v1_1-xxl-encoder-gguf/resolve/main/t5-v1_1-xxl-encoder-Q3_K_L.gguf -d /content/TotoroUI/models/clip -o t5-v1_1-xxl-encoder-Q3_K_L.gguf\n",
        "\n",
        "print(\"Downloading Controlnet Preprocessors...\")\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/Kijai/DepthAnythingV2-safetensors/resolve/main/depth_anything_v2_vitl_fp16.safetensors -d /content/TotoroUI/models/controlnet_preprocessors -o depth_anything_v2_vitl_fp16.safetensors\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetHED.pth -d /content/TotoroUI/models/controlnet_preprocessors -o ControlNetHED.pth\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/lllyasviel/Annotators/resolve/main/body_pose_model.pth -d /content/TotoroUI/models/controlnet_preprocessors -o body_pose_model.pth\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/lllyasviel/Annotators/resolve/main/hand_pose_model.pth -d /content/TotoroUI/models/controlnet_preprocessors -o hand_pose_model.pth\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/lllyasviel/Annotators/resolve/main/facenet.pth -d /content/TotoroUI/models/controlnet_preprocessors -o facenet.pth\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/lllyasviel/Annotators/resolve/main/sk_model.pth -d /content/TotoroUI/models/controlnet_preprocessors -o sk_model.pth\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/lllyasviel/Annotators/resolve/main/sk_model2.pth -d /content/TotoroUI/models/controlnet_preprocessors -o sk_model2.pth\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/fal/teed/resolve/main/5_model.pth -d /content/TotoroUI/models/controlnet_preprocessors -o 5_model.pth\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/lllyasviel/Annotators/resolve/main/table5_pidinet.pth -d /content/TotoroUI/models/controlnet_preprocessors -o table5_pidinet.pth\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/lllyasviel/Annotators/resolve/main/scannet.pt -d /content/TotoroUI/models/controlnet_preprocessors -o scannet.pt\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/lllyasviel/Annotators/resolve/main/mlsd_large_512_fp32.pth -d /content/TotoroUI/models/controlnet_preprocessors -omlsd_large_512_fp32.pth\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/lllyasviel/Annotators/resolve/main/netG.pth -d /content/TotoroUI/models/controlnet_preprocessors -o netG.pth\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/ybelkada/segment-anything/resolve/main/checkpoints/sam_vit_b_01ec64.pth -d /content/TotoroUI/models/controlnet_preprocessors -o sam_vit_b.pth\n",
        "\n",
        "print(\"Downloading Controlnet Union...\")\n",
        "!aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M https://huggingface.co/Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro/resolve/main/diffusion_pytorch_model.safetensors -d /content/TotoroUI/models/controlnet -o FLUX.1-dev-Controlnet-Union-Pro.safetensors\n",
        "\n",
        "with torch.inference_mode():\n",
        "    print(\"Loading Clips...\")\n",
        "    clip = DualCLIPLoaderGGUF.load_clip(\"t5-v1_1-xxl-encoder-Q3_K_L.gguf\", \"clip_l.safetensors\", \"flux\")[0]\n",
        "    print(\"Loading VAE...\")\n",
        "    vae = VAELoader.load_vae(\"ae.sft\")[0]\n",
        "    print(f\"Loading Flux.1-dev...\")\n",
        "    unet = UnetLoaderGGUF.load_unet(f\"flux1-dev.gguf\")[0]\n",
        "    print(\"Loading Controlnet Union...\")\n",
        "    controlnet = ControlNetLoader.load_controlnet(\"FLUX.1-dev-Controlnet-Union-Pro.safetensors\")[0]\n",
        "\n",
        "    unet_f, clip_f = unet, clip\n",
        "\n",
        "print(\"All Models Loaded!\")\n",
        "\n",
        "import warnings\n",
        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
        "warnings.simplefilter(action='ignore', category=UserWarning)\n",
        "\n",
        "import re\n",
        "import os\n",
        "import gc\n",
        "import random\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "from google.colab import files\n",
        "\n",
        "import nodes\n",
        "from totoro_extras import nodes_custom_sampler\n",
        "from totoro_extras import nodes_post_processing\n",
        "from totoro_extras import nodes_flux\n",
        "from totoro_extras import nodes_controlnet\n",
        "from totoro import model_management\n",
        "\n",
        "from controlnet_aux import CannyDetector, DepthAnythingDetector, HEDdetector, OpenposeDetector, PidiNetDetector, TEEDdetector, LineartDetector, LineartAnimeDetector, MLSDdetector, NormalBaeDetector, SamDetector\n",
        "\n",
        "CLIPTextEncodeFlux = nodes_flux.NODE_CLASS_MAPPINGS[\"CLIPTextEncodeFlux\"]()\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",
        "LoraLoader = NODE_CLASS_MAPPINGS[\"LoraLoader\"]()\n",
        "VAEDecode = NODE_CLASS_MAPPINGS[\"VAEDecode\"]()\n",
        "VAEEncode = NODE_CLASS_MAPPINGS[\"VAEEncode\"]()\n",
        "LoadImage = NODE_CLASS_MAPPINGS[\"LoadImage\"]()\n",
        "EmptyLatentImage = NODE_CLASS_MAPPINGS[\"EmptyLatentImage\"]()\n",
        "ImageScaleToTotalPixels = nodes_post_processing.NODE_CLASS_MAPPINGS[\"ImageScaleToTotalPixels\"]()\n",
        "SetUnionControlNetType = nodes_controlnet.NODE_CLASS_MAPPINGS[\"SetUnionControlNetType\"]()\n",
        "ControlNetApplyAdvanced = NODE_CLASS_MAPPINGS[\"ControlNetApplyAdvanced\"]()"
      ],
      "metadata": {
        "id": "k3aTOrdb8HxC",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @markdown <center><h1>Functions</h1></center>\n",
        "\n",
        "loras = {\n",
        "    \"xlabs_flux_anime\":\n",
        "     {\n",
        "         \"url\": \"https://huggingface.co/XLabs-AI/flux-lora-collection/resolve/main/anime_lora_comfy_converted.safetensors\",\n",
        "         \"filename\": \"xlabs_anime_lora.safetensors\",\n",
        "         \"triggers\": \"anime\"\n",
        "         },\n",
        "    \"xlabs_flux_art\":\n",
        "     {\n",
        "         \"url\": \"https://huggingface.co/XLabs-AI/flux-lora-collection/resolve/main/art_lora_comfy_converted.safetensors\",\n",
        "         \"filename\": \"xlabs_art_lora.safetensors\",\n",
        "         \"triggers\": \"art\"\n",
        "         },\n",
        "    \"xlabs_flux_disney\":\n",
        "     {\n",
        "         \"url\": \"https://huggingface.co/XLabs-AI/flux-lora-collection/resolve/main/disney_lora_comfy_converted.safetensors\",\n",
        "         \"filename\": \"xlabs_disney_lora.safetensors\",\n",
        "         \"triggers\": \"disney style\"\n",
        "         },\n",
        "    \"xlabs_flux_mjv6\":\n",
        "     {\n",
        "         \"url\": \"https://huggingface.co/XLabs-AI/flux-lora-collection/resolve/main/mjv6_lora_comfy_converted.safetensors\",\n",
        "         \"filename\": \"xlabs_mjv6_lora.safetensors\"\n",
        "         },\n",
        "    \"xlabs_flux_realism\":\n",
        "     {\n",
        "         \"url\": \"https://huggingface.co/XLabs-AI/flux-lora-collection/resolve/main/realism_lora_comfy_converted.safetensors\",\n",
        "         \"filename\": \"xlabs_realism_lora.safetensors\"\n",
        "         },\n",
        "    \"xlabs_flux_scenery\":\n",
        "     {\n",
        "         \"url\": \"https://huggingface.co/XLabs-AI/flux-lora-collection/resolve/main/scenery_lora_comfy_converted.safetensors\",\n",
        "         \"filename\": \"xlabs_scenery_lora.safetensors\",\n",
        "         \"triggers\": \"scenery style\"\n",
        "         },\n",
        "    \"xlabs_flux_furry\":\n",
        "     {\n",
        "         \"url\": \"https://huggingface.co/XLabs-AI/flux-lora-collection/resolve/main/furry_lora.safetensors\",\n",
        "        \"filename\": \"xlabs_flux_furry_lora.safetensors\"\n",
        "        }\n",
        "}\n",
        "\n",
        "\n",
        "def load_loras(prompt):\n",
        "  # @markdown <ul><li><h2>Load Loras</h2>- Add &lt;lora_name:model_strength&gt; to Prompt</li></ul>\n",
        "\n",
        "  global unet, clip, unet_f, clip_f\n",
        "\n",
        "  unet_f, clip_f = unet, clip\n",
        "\n",
        "  matches = re.findall(r\"<\\s*([^:]+?)\\s*:\\s*([0-9.]+)\\s*>\", prompt)\n",
        "\n",
        "  loras_list = [(name.strip(), float(value)) for name, value in matches]\n",
        "\n",
        "  if len(loras_list):\n",
        "    print(\"Loading Loras...\")\n",
        "\n",
        "  for lora_tuple in loras_list:\n",
        "    lora = loras.get(lora_tuple[0], None)\n",
        "\n",
        "    if lora:\n",
        "      !aria2c --quiet --console-log-level=error --auto-file-renaming=false --allow-overwrite=false -c -x 16 -s 16 -k 1M {lora[\"url\"]} -d /content/TotoroUI/models/loras -o {lora[\"filename\"]}\n",
        "\n",
        "      with torch.inference_mode():\n",
        "        unet_f, clip_f = LoraLoader.load_lora(unet_f, clip_f, lora[\"filename\"], lora_tuple[1], lora_tuple[1])\n",
        "\n",
        "      print(f\"Loaded Lora: {lora_tuple[0]}\")\n",
        "    else:\n",
        "      print(f\"Lora not listed: {lora_tuple[0]}\")\n",
        "\n",
        "def clean_prompt(prompt):\n",
        "  cleaned_prompt = re.sub(r\"<.*?>\", \"\", prompt)\n",
        "\n",
        "  return cleaned_prompt\n",
        "\n",
        "def cuda_gc():\n",
        "  try:\n",
        "    model_management.soft_empty_cache()\n",
        "    gc.collect()\n",
        "    torch.cuda.empty_cache()\n",
        "    torch.cuda.ipc_collect()\n",
        "  except:\n",
        "    pass\n",
        "\n",
        "def img_tensor_to_np(img_tensor):\n",
        "  img_tensor = img_tensor.clone() * 255.0\n",
        "  return img_tensor.squeeze().numpy().astype(np.uint8)\n",
        "\n",
        "def img_np_to_tensor(img_np_list):\n",
        "  return torch.from_numpy(img_np_list.astype(np.float32) / 255.0).unsqueeze(0)\n",
        "\n",
        "def controlnet_preprocess(width, height):\n",
        "  # @markdown <ul><li><h2>Controlnet Preprocessor</h2></li></ul>\n",
        "\n",
        "  process_image = True # @param {\"type\":\"boolean\"}\n",
        "  preprocessor_type = \"openpose\" #  @param [\"none\", \"openpose\", \"depth\", \"hed\", \"pidi\", \"scribble\", \"ted\", \"canny\", \"lineart\", \"anime_lineart\", \"mlsd\", \"normal\", \"segment\"]\n",
        "  input_image = \"/content/test.png\" # @param {\"type\":\"string\"}\n",
        "  resolution = 512 # @param {\"type\":\"slider\",\"min\":512,\"max\":2048,\"step\":1}\n",
        "  resize_mode = \"Just Resize\" # @param [\"Just Resize\",\"Resize and Fill\",\"Crop and Resize\"]\n",
        "\n",
        "  image_input = LoadImage.load_image(input_image)[0]\n",
        "\n",
        "  if process_image:\n",
        "    image_np = img_tensor_to_np(image_input)\n",
        "    img = Image.fromarray(image_np)\n",
        "\n",
        "    if resize_mode == \"Just Resize\":\n",
        "        img = img.resize((width, height), Image.Resampling.LANCZOS)\n",
        "\n",
        "    elif resize_mode == \"Resize and Fill\":\n",
        "        img.thumbnail((width, height), Image.ANTIALIAS)\n",
        "        result = Image.new(\"RGB\", (width, height), (0, 0, 0))\n",
        "        offset = ((width - img.width) // 2, (height - img.height) // 2)\n",
        "        result.paste(img, offset)\n",
        "        img = result\n",
        "\n",
        "    elif resize_mode == \"Crop and Resize\":\n",
        "        src_width, src_height = img.size\n",
        "        src_aspect = src_width / src_height\n",
        "        target_aspect = width / height\n",
        "\n",
        "        if src_aspect > target_aspect:\n",
        "            new_width = int(src_height * target_aspect)\n",
        "            left = (src_width - new_width) // 2\n",
        "            img = img.crop((left, 0, left + new_width, src_height))\n",
        "        else:\n",
        "            new_height = int(src_width / target_aspect)\n",
        "            top = (src_height - new_height) // 2\n",
        "            img = img.crop((0, top, src_width, top + new_height))\n",
        "\n",
        "        img = img.resize((width, height), Image.Resampling.LANCZOS)\n",
        "\n",
        "    image_np = np.array(img).astype(np.uint8)\n",
        "\n",
        "    if preprocessor_type == \"none\":\n",
        "      processed_image = Image.fromarray(image_np)\n",
        "\n",
        "    if preprocessor_type == \"canny\":\n",
        "      canny = CannyDetector()\n",
        "      processed_image = canny(image_np, detect_resolution=resolution, image_resolution=resolution)\n",
        "\n",
        "    elif preprocessor_type == \"depth\":\n",
        "      depth = DepthAnythingDetector.from_pretrained(\"/content/TotoroUI/models/controlnet_preprocessors\", filename=\"depth_anything_v2_vitl_fp16.safetensors\")\n",
        "      processed_image = depth(image_np, detect_resolution=resolution, image_resolution=resolution)\n",
        "\n",
        "    elif preprocessor_type == \"hed\":\n",
        "      hed = HEDdetector.from_pretrained(\"/content/TotoroUI/models/controlnet_preprocessors\", filename=\"ControlNetHED.pth\")\n",
        "      processed_image = hed(image_np, detect_resolution=resolution, image_resolution=resolution)\n",
        "\n",
        "    elif preprocessor_type == \"openpose\":\n",
        "      openpose = OpenposeDetector.from_pretrained(\"/content/TotoroUI/models/controlnet_preprocessors\", filename=\"body_pose_model.pth\", hand_filename=\"hand_pose_model.pth\", face_filename=\"facenet.pth\")\n",
        "      processed_image = openpose(image_np, detect_resolution=resolution, image_resolution=resolution, include_hand=True, include_face=True)\n",
        "\n",
        "    elif preprocessor_type == \"pidi\":\n",
        "      pidi = PidiNetDetector.from_pretrained(\"/content/TotoroUI/models/controlnet_preprocessors\", filename=\"table5_pidinet.pth\")\n",
        "      processed_image = pidi(image_np, detect_resolution=resolution, image_resolution=resolution, safe=True)\n",
        "\n",
        "    elif preprocessor_type == \"scribble\":\n",
        "      scribble = HEDdetector.from_pretrained(\"/content/TotoroUI/models/controlnet_preprocessors\", filename=\"ControlNetHED.pth\")\n",
        "      processed_image = scribble(image_np, detect_resolution=resolution, image_resolution=resolution, scribble=True)\n",
        "\n",
        "    elif preprocessor_type == \"ted\":\n",
        "      ted = TEEDdetector.from_pretrained(\"/content/TotoroUI/models/controlnet_preprocessors\", filename=\"5_model.pth\")\n",
        "      processed_image = ted(image_np, detect_resolution=resolution)\n",
        "\n",
        "    elif preprocessor_type == \"lineart\":\n",
        "      lineart = LineartDetector.from_pretrained(\"/content/TotoroUI/models/controlnet_preprocessors\", filename=\"sk_model.pth\", coarse_filename=\"sk_model2.pth\")\n",
        "      processed_image = lineart(image_np, detect_resolution=resolution, image_resolution=resolution, coarse=True)\n",
        "\n",
        "    elif preprocessor_type == \"anime_lineart\":\n",
        "      anime_lineart = LineartAnimeDetector.from_pretrained(\"/content/TotoroUI/models/controlnet_preprocessors\", filename=\"netG.pth\")\n",
        "      processed_image = anime_lineart(image_np, detect_resolution=resolution, image_resolution=resolution)\n",
        "\n",
        "    elif preprocessor_type == \"mlsd\":\n",
        "      mlsd = MLSDdetector.from_pretrained(\"/content/TotoroUI/models/controlnet_preprocessors\", filename=\"mlsd_large_512_fp32.pth\")\n",
        "      processed_image = mlsd(image_np, detect_resolution=resolution, image_resolution=resolution)\n",
        "\n",
        "    elif preprocessor_type == \"normal\":\n",
        "      normal = NormalBaeDetector.from_pretrained(\"/content/TotoroUI/models/controlnet_preprocessors\", filename=\"scannet.pt\")\n",
        "      processed_image = normal(image_np, detect_resolution=resolution, image_resolution=resolution)\n",
        "\n",
        "    elif preprocessor_type == \"segment\":\n",
        "      segment = SamDetector.from_pretrained(\"/content/TotoroUI/models/controlnet_preprocessors\", model_type=\"vit_b\", filename=\"sam_vit_b.pth\")\n",
        "      processed_image = segment(image_np, detect_resolution=resolution, image_resolution=resolution)\n",
        "\n",
        "    processed_image_np = np.array(processed_image)\n",
        "    processed_image_tensor = img_np_to_tensor(processed_image_np)\n",
        "\n",
        "    display(Image.fromarray(processed_image_np))\n",
        "\n",
        "    return processed_image_tensor\n",
        "\n",
        "  else:\n",
        "    return image_input\n",
        "\n",
        "def apply_controlnet(cond, width, height):\n",
        "  # @markdown <ul><li><h2>Union Controlnet</h2></li></ul>\n",
        "\n",
        "  controlnet_type = \"openpose\" # @param [\"openpose\", \"depth\", \"hed/pidi/scribble/ted\", \"canny/lineart/anime_lineart/mlsd\", \"normal\", \"segment\", \"tile\", \"repaint\" ]\n",
        "  strength = 0.5 # @param {\"type\":\"slider\",\"min\":0.0,\"max\":10.0,\"step\":0.01}\n",
        "  start_percent = 0 # @param {\"type\":\"slider\",\"min\":0.0,\"max\":1.0,\"step\":0.001}\n",
        "  end_percent = 0.225 # @param {\"type\":\"slider\",\"min\":0.0,\"max\":1.0,\"step\":0.001}\n",
        "\n",
        "  cond_neg = CLIPTextEncodeFlux.encode(clip_f, \"\", \"\", 0)[0]\n",
        "\n",
        "  print(\"Loading Controlnet...\")\n",
        "\n",
        "  preprocessed_image = controlnet_preprocess(width, height)\n",
        "\n",
        "  controlnet_f = SetUnionControlNetType.set_controlnet_type(controlnet, controlnet_type)[0]\n",
        "\n",
        "  cond = ControlNetApplyAdvanced.apply_controlnet(cond, cond_neg, controlnet_f, preprocessed_image, strength, start_percent, end_percent, vae)[0]\n",
        "\n",
        "  del controlnet_f\n",
        "\n",
        "  return cond\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\n",
        "\n",
        "def save_image(decoded, path, name, download=False):\n",
        "  full_path = os.path.abspath(os.path.join(path, name))\n",
        "  Image.fromarray(np.array(decoded*255, dtype=np.uint8)[0]).save( full_path)\n",
        "\n",
        "  img = Image.open(full_path)\n",
        "  display(img)\n",
        "\n",
        "  if download:\n",
        "    files.download(full_path)\n",
        "\n",
        "@torch.inference_mode()\n",
        "def generate(prompt, width, height, fixed_seed, guidance, steps, sampler_name, scheduler, batch_size, auto_download, mode=\"t2i\", input_img=None, denoise=1.0):\n",
        "  global unet, clip, unet_f, clip_f\n",
        "\n",
        "  print(\"Prompt Received\")\n",
        "\n",
        "  load_loras(prompt)\n",
        "  prompt = clean_prompt(prompt)\n",
        "\n",
        "  if mode == \"t2i\":\n",
        "    latent_image = EmptyLatentImage.generate(closestNumber(width, 16), closestNumber(height, 16))[0]\n",
        "\n",
        "  elif mode == \"i2i\":\n",
        "    image = LoadImage.load_image(input_img)[0]\n",
        "    latent_image = ImageScaleToTotalPixels.upscale(image, \"lanczos\", 1.0)[0]\n",
        "    latent_image = VAEEncode.encode(vae, latent_image)[0]\n",
        "\n",
        "  cond = CLIPTextEncodeFlux.encode(clip_f, prompt, prompt, guidance)[0]\n",
        "\n",
        "  cond = apply_controlnet(cond, width, height)\n",
        "\n",
        "  guider = BasicGuider.get_guider(unet_f, cond)[0]\n",
        "  sampler = KSamplerSelect.get_sampler(sampler_name)[0]\n",
        "  sigmas = BasicScheduler.get_sigmas(unet_f, scheduler, steps, denoise)[0]\n",
        "\n",
        "  for i in range(0, batch_size):\n",
        "    if fixed_seed == 0:\n",
        "      seed = random.randint(0, 18446744073709551615)\n",
        "    else:\n",
        "      seed = fixed_seed\n",
        "\n",
        "    print(\"Seed:\", seed)\n",
        "\n",
        "    noise = RandomNoise.get_noise(seed)[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",
        "\n",
        "    save_image(decoded, \"/content\", f\"flux_{mode}_{seed}_{i}.png\", auto_download)\n",
        "\n",
        "  cuda_gc()\n",
        "\n",
        "print(f\"{'Lora Name':<40} {'Trigger Words':<40}\")\n",
        "print(\"-\" * 80)\n",
        "\n",
        "for key_name, details in loras.items():\n",
        "    trigger_words = details.get(\"triggers\", \"N/A\")\n",
        "    print(f\"{key_name:<40} {trigger_words:<40}\")"
      ],
      "metadata": {
        "id": "mLGPKWvopwnC",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Ur9TmMNwC2kR"
      },
      "outputs": [],
      "source": [
        "#@markdown <center><h1>Txt2Img</h1></center>\n",
        "\n",
        "positive_prompt = \"\" # @param {\"type\":\"string\"}\n",
        "width = 512 # @param {\"type\":\"slider\",\"min\":256,\"max\":2048,\"step\":1}\n",
        "height = 512 # @param {\"type\":\"slider\",\"min\":256,\"max\":2048,\"step\":1}\n",
        "fixed_seed = 0 # @param {\"type\":\"slider\",\"min\":0,\"max\":18446744073709552000,\"step\":1}\n",
        "guidance = 4.5 # @param {\"type\":\"slider\",\"min\":0,\"max\":20,\"step\":0.5}\n",
        "steps = 25 # @param {\"type\":\"slider\",\"min\":4,\"max\":50,\"step\":1}\n",
        "sampler_name = \"euler\" # @param [\"euler\",\"heun\",\"heunpp2\",\"heunpp2\",\"dpm_2\",\"lms\",\"dpmpp_2m\",\"ipndm\",\"deis\",\"ddim\",\"uni_pc\",\"uni_pc_bh2\"]\n",
        "scheduler = \"simple\" # @param [\"normal\",\"sgm_uniform\",\"simple\",\"ddim_uniform\"]\n",
        "batch_size = 1 # @param {\"type\":\"slider\",\"min\":1,\"max\":20,\"step\":1}\n",
        "auto_download = False # @param {\"type\":\"boolean\"}\n",
        "\n",
        "generate(positive_prompt, width, height, fixed_seed, guidance, steps, sampler_name, scheduler, batch_size, auto_download)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Dpd2sfrePYoA"
      },
      "outputs": [],
      "source": [
        "#@markdown <center><h1>Img2Img</h1></center>\n",
        "\n",
        "positive_prompt = \"anime style\" # @param {\"type\":\"string\"}\n",
        "fixed_seed = 0 # @param {\"type\":\"slider\",\"min\":0,\"max\":18446744073709552000,\"step\":1}\n",
        "guidance = 4.5 # @param {\"type\":\"slider\",\"min\":0,\"max\":20,\"step\":0.5}\n",
        "steps = 25 # @param {\"type\":\"slider\",\"min\":4,\"max\":50,\"step\":1}\n",
        "sampler_name = \"euler\" # @param [\"euler\",\"heun\",\"heunpp2\",\"heunpp2\",\"dpm_2\",\"lms\",\"dpmpp_2m\",\"ipndm\",\"deis\",\"ddim\",\"uni_pc\",\"uni_pc_bh2\"]\n",
        "scheduler = \"simple\" # @param [\"normal\",\"sgm_uniform\",\"simple\",\"ddim_uniform\"]\n",
        "input_img = \"/content/test.png\" # @param {\"type\":\"string\"}\n",
        "denoise = 0.85 # @param {\"type\":\"slider\",\"min\":0,\"max\":1,\"step\":0.01}\n",
        "batch_size = 1 # @param {\"type\":\"slider\",\"min\":1,\"max\":20,\"step\":1}\n",
        "auto_download = False # @param {\"type\":\"boolean\"}\n",
        "\n",
        "\n",
        "generate(positive_prompt, 0, 0, fixed_seed, guidance, steps, sampler_name, scheduler, batch_size, auto_download, \"i2i\", input_img, denoise)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}