Upload 3 files

app.py

@@ -24,6 +24,7 @@ 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'
+#controlnet_model_union_repo = 'InstantX/FLUX.1-dev-Controlnet-Union-alpha'
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
 controlnet_union = None
 controlnet = None
@@ -39,32 +40,32 @@ def change_base_model(repo_id: str, cn_on: bool, progress=gr.Progress(track_tqdm
     global last_model
     global last_cn_on
     dtype = torch.bfloat16
+    #dtype = torch.float8_e4m3fn
     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 (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 gr.update(visible=True)
         if cn_on:
-            progress(0, desc=f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
+            #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=dtype)
             controlnet = FluxMultiControlNetModel([controlnet_union])
             pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=dtype)
-            #pipe.enable_model_cpu_offload()
             last_model = repo_id
             last_cn_on = cn_on
-            progress(1, desc=f"Model loaded: {repo_id} / ControlNet Loaded: {controlnet_model_union_repo}")
+            #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}")
+            #progress(0, desc=f"Loading model: {repo_id}")
             print(f"Loading model: {repo_id}")
             clear_cache()
             pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=dtype)
-            #pipe.enable_model_cpu_offload()
             last_model = repo_id
             last_cn_on = cn_on
-            progress(1, desc=f"Model loaded: {repo_id}")
+            #progress(1, desc=f"Model loaded: {repo_id}")
             print(f"Model loaded: {repo_id}")
     except Exception as e:
-        print(e)
+        print(f"Model load Error: {e}")
+        raise gr.Error(f"Model load Error: {e}")
     return gr.update(visible=True)
 
 change_base_model.zerogpu = True
@@ -113,15 +114,16 @@ 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=gr.Progress(track_tqdm=True)):
-    pipe.to("cuda")
-    #controlnet.to("cuda")
-    #controlnet_union.to("cuda")
-    generator = torch.Generator(device="cuda").manual_seed(seed)
-    
-    with calculateDuration("Generating image"):
-        # Generate image
-        modes, images, scales = get_control_params()
-        try:
+    global pipe
+    global controlnet
+    global controlnet_union
+    try:
+        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 len(modes) == 0:
                 progress(0, desc="Start Inference.")
                 image = pipe(
@@ -135,6 +137,8 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
                 ).images[0]
             else:
                 progress(0, desc="Start Inference with ControlNet.")
+                if controlnet is not None: controlnet.to("cuda")
+                if controlnet_union is not None: controlnet_union.to("cuda")
                 image = pipe(
                     prompt=prompt_mash,
                     control_image=images,
@@ -147,14 +151,14 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
                     generator=generator,
                     joint_attention_kwargs={"scale": lora_scale},
                 ).images[0]
-        except Exception as e:
-            progress(1, desc=f"Inference Error: {e}")
-            print(e)
-            raise Exception(f"Inference Error: {e}")
+    except Exception as e:
+        print(e)
+        raise gr.Error(f"Inference Error: {e}")
     return image
 
 def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height,
               lora_scale, lora_json, cn_on, progress=gr.Progress(track_tqdm=True)):
+    global pipe
     if selected_index is None and not is_valid_lora(lora_json):
         gr.Info("LoRA isn't selected.")
     #    raise gr.Error("You must select a LoRA before proceeding.")
@@ -192,7 +196,7 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
     
-    progress(1, desc="Preparing Inference.")
+    progress(0, desc="Running Inference.")
 
     image = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, cn_on, progress)
     if is_valid_lora(lora_json):
@@ -200,6 +204,8 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
         pipe.unload_lora_weights()
     if selected_index is not None: pipe.unload_lora_weights()
     pipe.to("cpu")
+    if controlnet is not None: controlnet.to("cpu")
+    if controlnet_union is not None: controlnet_union.to("cpu")
     clear_cache()
     return image, seed  
 
@@ -390,6 +396,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
                             cn_mode = [None] * num_cns
                             cn_scale = [None] * num_cns
                             cn_image = [None] * num_cns
+                            cn_image_ref = [None] * num_cns
                             cn_res = [None] * num_cns
                             cn_num = [None] * num_cns
                             for i in range(num_cns):
@@ -399,42 +406,56 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
                                         cn_scale[i] = gr.Slider(label=f"ControlNet {int(i+1)} Weight", minimum=0.0, maximum=1.0, step=0.01, value=0.75)
                                         cn_res[i] = gr.Slider(label=f"ControlNet {int(i+1)} Preprocess resolution", minimum=128, maximum=512, value=384, step=1)
                                         cn_num[i] = gr.Number(i, visible=False)
-                                    cn_image[i] = gr.Image(type="pil", label="Control Image", height=256, show_share_button=False)
+                                    with gr.Row():
+                                        cn_image_ref[i] = gr.Image(label="Image Reference", type="pil", format="png", height=256, sources=["upload", "clipboard"], show_fullscreen_button=False, show_share_button=False)
+                                        cn_image[i] = gr.Image(label="Control Image", type="pil", format="png", height=256, show_share_button=False, show_fullscreen_button=False, interactive=False)
 
     gallery.select(
         update_selection,
         inputs=[width, height],
-        outputs=[prompt, selected_info, selected_index, width, height]
+        outputs=[prompt, selected_info, selected_index, width, height],
+        queue=False,
+        show_api=False,
     )
     custom_lora.input(
         add_custom_lora,
         inputs=[custom_lora],
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
+        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt],
+        queue=False,
+        show_api=False,
     )
     custom_lora_button.click(
         remove_custom_lora,
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
+        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora],
+        queue=False,
+        show_api=False,
     )
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=change_base_model,
         inputs=[model_name, cn_on],
-        outputs=[result]
+        outputs=[result],
+        queue=False,
+        show_api=False,
     ).success(
         fn=run_lora,
         inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height,
                  lora_scale, lora_repo_json, cn_on], 
-        outputs=[result, seed]
+        outputs=[result, seed],
+        queue=True,
+        show_api=True,
     )
 
-    deselect_lora_button.click(deselect_lora, None, [prompt, selected_info, selected_index, width, height])
+    deselect_lora_button.click(deselect_lora, None, [prompt, selected_info, selected_index, width, height], queue=False, show_api=False)
     gr.on(
         triggers=[model_name.change, cn_on.change],
         fn=change_base_model,
         inputs=[model_name, cn_on],
-        outputs=[result]
+        outputs=[result],
+        queue=True,
+        show_api=False,
     )
-    prompt_enhance.click(enhance_prompt, [prompt], [prompt])
+    prompt_enhance.click(enhance_prompt, [prompt], [prompt], queue=False, show_api=False)
 
     gr.on(
         triggers=[lora_search_civitai_submit.click, lora_search_civitai_query.submit],
@@ -449,7 +470,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
     lora_search_civitai_result.change(select_civitai_lora, [lora_search_civitai_result], [lora_download_url, lora_search_civitai_desc], scroll_to_output=True, queue=False, show_api=False)
 
     for i, l in enumerate(lora_repo):
-        deselect_lora_button.click(lambda: ("", 1.0), None, [lora_repo[i], lora_wt[i]])
+        deselect_lora_button.click(lambda: ("", 1.0), None, [lora_repo[i], lora_wt[i]], queue=False, show_api=False)
         gr.on(
             triggers=[lora_download[i].click],
             fn=download_my_lora,
@@ -479,8 +500,8 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
             outputs=[cn_on],
             queue=True,
             show_api=False,
-        )
-        cn_image[i].upload(set_control_union_image, [cn_num[i], cn_mode[i], cn_image[i], height, width, cn_res[i]], [cn_image[i]])
+        ).success(set_control_union_image, [cn_num[i], cn_mode[i], cn_image_ref[i], height, width, cn_res[i]], [cn_image[i]], queue=False, show_api=False)
+        cn_image_ref[i].upload(set_control_union_image, [cn_num[i], cn_mode[i], cn_image_ref[i], height, width, cn_res[i]], [cn_image[i]], queue=False, show_api=False)
 
     tagger_generate_from_image.click(
             lambda: ("", "", ""), None, [v2_series, v2_character, prompt], queue=False, show_api=False,
@@ -582,7 +603,8 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
     pg_create_caption_button.click(
         create_caption,
         inputs=[pg_input_image],
-        outputs=[pg_caption_output]
+        outputs=[pg_caption_output],
+        show_api=False,
     )
 
     pg_generate_button.click(
@@ -592,20 +614,23 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
                 pg_additional_details, pg_photography_styles, pg_device, pg_photographer,
                 pg_artist, pg_digital_artform,
                 pg_place, pg_lighting, pg_clothing, pg_composition, pg_pose, pg_background],
-        outputs=[pg_output, gr.Number(visible=False), pg_t5xxl_output, pg_clip_l_output, pg_clip_g_output]
+        outputs=[pg_output, gr.Number(visible=False), pg_t5xxl_output, pg_clip_l_output, pg_clip_g_output],
+        show_api=False,
     )
 
     pg_add_caption_button.click(
         prompt_generator.add_caption_to_prompt,
         inputs=[pg_output, pg_caption_output],
-        outputs=[pg_output]
+        outputs=[pg_output],
+        show_api=False,
     )
 
     pg_generate_text_button.click(
         huggingface_node.generate,
         inputs=[pg_model, pg_output, pg_happy_talk, pg_compress, pg_compression_level,
                 pg_poster, pg_custom_base_prompt],
-        outputs=pg_text_output
+        outputs=pg_text_output,
+        show_api=False,
     )
 
     def update_all_options(choice):
@@ -644,7 +669,8 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
             pg_roles, pg_hairstyles, pg_clothing,
             pg_place, pg_lighting, pg_composition, pg_pose, pg_background, pg_additional_details,
             pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform
-        ]
+        ],
+        show_api=False,
     )
 
 app.queue()

mod.py

@@ -169,6 +169,7 @@ def preprocess_image(image: Image.Image, control_mode: str, height: int, width:
     image_resolution = max(width, height)
     image_before = resize_image(expand2square(image.convert("RGB")), image_resolution, image_resolution, False)
     # generated control_
+    progress(0, desc="start to generate control image")
     print("start to generate control image")
     preprocessor = Preprocessor()
     if control_mode == "depth_midas":
@@ -219,7 +220,9 @@ def preprocess_image(image: Image.Image, control_mode: str, height: int, width:
         image_width, image_height = control_image.size
     
     image_after = resize_image(control_image, width, height, False)
-    print(f"generate control image success: {image_width}x{image_height} => {width}x{height}")
+    ref_width, ref_height = image.size
+    progress(1, desc=f"generate control image success: {ref_width}x{ref_height} => {image_width}x{image_height}")
+    print(f"generate control image success: {ref_width}x{ref_height} => {image_width}x{image_height}")
     return image_after
 
 
@@ -236,8 +239,9 @@ def set_control_union_mode(i: int, mode: str, scale: str):
     else: return gr.update(visible=True)
 
 
-def set_control_union_image(i: int, mode: str, image: Image.Image, height: int, width: int, preprocess_resolution: int):
+def set_control_union_image(i: int, mode: str, image: Image.Image | None, height: int, width: int, preprocess_resolution: int):
     global control_images
+    if image is None: return None
     control_images[i] = preprocess_image(image, mode, height, width, preprocess_resolution)
     return control_images[i]
 
@@ -267,8 +271,9 @@ def get_trigger_word(lorajson: list[dict]):
 
 # https://huggingface.co/docs/diffusers/v0.23.1/en/api/loaders#diffusers.loaders.LoraLoaderMixin.fuse_lora
 # https://github.com/huggingface/diffusers/issues/4919
-def fuse_loras(pipe, lorajson: list[dict]):
+def fuse_loras(pipe, lorajson: list[dict], progress=gr.Progress(track_tqdm=True)):
     if not lorajson or not isinstance(lorajson, list): return
+    progress(0, desc="Fusing LoRA.")
     a_list = []
     w_list = []
     for d in lorajson:

requirements.txt

@@ -4,7 +4,7 @@ torchvision
 huggingface_hub
 accelerate
 git+https://github.com/huggingface/diffusers
-git+https://github.com/huggingface/transformers
+transformers
 peft
 sentencepiece
 timm