Upload 49 files

app.py

@@ -5,14 +5,16 @@ import torch
 from PIL import Image
 import spaces
 from diffusers import DiffusionPipeline
+from diffusers.pipelines.flux.pipeline_flux_controlnet import FluxControlNetPipeline
+from diffusers.models.controlnet_flux import FluxControlNetModel, FluxMultiControlNetModel
 from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
 import copy
 import random
 import time
 
-from mod import (models, clear_cache, get_repo_safetensors, change_base_model,
+from mod import (models, clear_cache, get_repo_safetensors, is_repo_name, is_repo_exists,
                  description_ui, num_loras, compose_lora_json, is_valid_lora, fuse_loras,
-                 get_trigger_word, enhance_prompt, pipe, controlnet, num_cns, set_control_union_image,
+                 get_trigger_word, enhance_prompt, num_cns, set_control_union_image,
                  get_control_union_mode, set_control_union_mode, get_control_params)
 from flux import (search_civitai_lora, select_civitai_lora, search_civitai_lora_json,
                   download_my_lora, get_all_lora_tupled_list, apply_lora_prompt,
@@ -21,6 +23,44 @@ from tagger.tagger import predict_tags_wd, compose_prompt_to_copy
 from tagger.fl2cog import predict_tags_fl2_cog
 from tagger.fl2flux import predict_tags_fl2_flux
 
+# Initialize the base model
+base_model = models[0]
+controlnet_model_union_repo = 'InstantX/FLUX.1-dev-Controlnet-Union'
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
+last_model = models[0]
+last_cn_on = False
+
+# https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union
+# https://huggingface.co/spaces/jiuface/FLUX.1-dev-Controlnet-Union
+def change_base_model(repo_id: str, cn_on: bool, progress=gr.Progress(track_tqdm=True)):
+    global pipe
+    global last_model
+    global last_cn_on
+    try:
+        if (repo_id == last_model and cn_on is last_cn_on) or not is_repo_name(repo_id) or not is_repo_exists(repo_id): return
+        if cn_on:
+            progress(0, desc=f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
+            print(f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
+            clear_cache()
+            controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union_repo, torch_dtype=torch.bfloat16)
+            controlnet = FluxMultiControlNetModel([controlnet_union])
+            pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=torch.bfloat16)
+            last_model = repo_id
+            progress(1, desc=f"Model loaded: {repo_id} / ControlNet Loaded: {controlnet_model_union_repo}")
+            print(f"Model loaded: {repo_id} / ControlNet Loaded: {controlnet_model_union_repo}")
+        else:
+            progress(0, desc=f"Loading model: {repo_id}")
+            print(f"Loading model: {repo_id}")
+            clear_cache()
+            pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
+            last_model = repo_id
+            progress(1, desc=f"Model loaded: {repo_id}")
+            print(f"Model loaded: {repo_id}")
+    except Exception as e:
+        print(e)
+    return gr.update(visible=True)
+
+change_base_model.zerogpu = True
 
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
@@ -65,14 +105,14 @@ def update_selection(evt: gr.SelectData, width, height):
     )
 
 @spaces.GPU(duration=70)
-def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, cn_on, progress):
+def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, cn_on, progress=gr.Progress(track_tqdm=True)):
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
     
     with calculateDuration("Generating image"):
         # Generate image
         modes, images, scales = get_control_params()
-        if not cn_on or controlnet is None or len(modes) == 0:
+        if not cn_on or len(modes) == 0:
             progress(0, desc="Start Inference.")
             image = pipe(
                 prompt=prompt_mash,
@@ -85,7 +125,6 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
             ).images[0]
         else:
             progress(0, desc="Start Inference with ControlNet.")
-            print(modes, scales) #
             image = pipe(
                 prompt=prompt_mash,
                 control_image=images,
@@ -337,7 +376,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
                                     for i in range(num_loras):
                                         lora_download[i] = gr.Button(f"Get and set LoRA to {int(i+1)}")
 
-                    with gr.Accordion("ControlNet", open=True):
+                    with gr.Accordion("ControlNet", open=False):
                         with gr.Column():
                             cn_on = gr.Checkbox(False, label="Use ControlNet")
                             cn_mode = [None] * num_cns

mod.py

@@ -1,9 +1,7 @@
 import gradio as gr
 import torch
 import spaces
-from diffusers import DiffusionPipeline
-from diffusers.pipelines.flux.pipeline_flux_controlnet import FluxControlNetPipeline
-from diffusers.models.controlnet_flux import FluxControlNetModel, FluxMultiControlNetModel
+
 from pathlib import Path
 import gc
 import subprocess
@@ -35,16 +33,9 @@ models = [
 
 num_loras = 3
 num_cns = 2
-# Initialize the base model
-base_model = models[0]
-controlnet_model_union_repo = 'InstantX/FLUX.1-dev-Controlnet-Union'
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
-controlnet = None
 control_images = [None] * num_cns
 control_modes = [-1] * num_cns
 control_scales = [0] * num_cns
-last_model = models[0]
-last_cn_on = False
 
 
 def is_repo_name(s):
@@ -84,39 +75,6 @@ def get_repo_safetensors(repo_id: str):
     else: return gr.update(value=files[0], choices=files)
 
 
-# https://huggingface.co/spaces/DamarJati/FLUX.1-DEV-Canny
-# https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union
-# https://huggingface.co/spaces/jiuface/FLUX.1-dev-Controlnet-Union
-def change_base_model(repo_id: str, cn_on: bool, progress=gr.Progress(track_tqdm=True)):
-    global pipe
-    global controlnet
-    global last_model
-    global last_cn_on
-    try:
-        if (repo_id == last_model and cn_on is last_cn_on) or not is_repo_name(repo_id) or not is_repo_exists(repo_id): return
-        if cn_on:
-            progress(0, desc=f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
-            print(f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
-            clear_cache()
-            controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union_repo, torch_dtype=torch.bfloat16)
-            controlnet = FluxMultiControlNetModel([controlnet_union])
-            pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=torch.bfloat16)
-            last_model = repo_id
-            progress(1, desc=f"Model loaded: {repo_id} / ControlNet Loaded: {controlnet_model_union_repo}")
-            print(f"Model loaded: {repo_id} / ControlNet Loaded: {controlnet_model_union_repo}")
-        else:
-            progress(0, desc=f"Loading model: {repo_id}")
-            print(f"Loading model: {repo_id}")
-            clear_cache()
-            pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
-            last_model = repo_id
-            progress(1, desc=f"Model loaded: {repo_id}")
-            print(f"Model loaded: {repo_id}")
-    except Exception as e:
-        print(e)
-    return gr.update(visible=True)
-
-
 def expand2square(pil_img: Image.Image, background_color: tuple=(0, 0, 0)):
     width, height = pil_img.size
     if width == height:
@@ -183,7 +141,8 @@ def get_control_params():
 
 
 from preprocessor import Preprocessor
-def preprocess_image(image: Image.Image, control_mode: str, height: int, width: int, preprocess_resolution: int):
+def preprocess_image(image: Image.Image, control_mode: str, height: int, width: int,
+                     preprocess_resolution: int, progress=gr.Progress(track_tqdm=True)):
     if control_mode == "None": return image
     image_resolution = max(width, height)
     image_before = resize_image(expand2square(image), image_resolution, image_resolution, False)
@@ -239,7 +198,6 @@ def preprocess_image(image: Image.Image, control_mode: str, height: int, width:
     
     image_after = resize_image(control_image, width, height, False)
     print(f"generate control image success: {image_width}x{image_height} => {width}x{height}")
-
     return image_after
 
 
@@ -347,5 +305,4 @@ def enhance_prompt(input_prompt):
 
 
 load_prompt_enhancer.zerogpu = True
-change_base_model.zerogpu = True
 fuse_loras.zerogpu = True