Upload 2 files

app.py

@@ -5,8 +5,9 @@ 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 diffusers.pipelines.flux.pipeline_flux_controlnet import FluxControlNetPipeline
+#from diffusers.models.controlnet_flux import FluxControlNetModel, FluxMultiControlNetModel
+from diffusers import FluxControlNetPipeline, FluxControlNetModel, FluxMultiControlNetModel
 from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
 import copy
 import random
@@ -14,7 +15,7 @@ import time
 
 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, num_cns, set_control_union_image,
+                 get_trigger_word, enhance_prompt, deselect_lora, 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,
@@ -26,6 +27,8 @@ from tagger.fl2flux import predict_tags_fl2_flux
 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_union = None
+controlnet = None
 last_model = models[0]
 last_cn_on = False
 
@@ -33,6 +36,8 @@ last_cn_on = False
 # 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_union
+    global controlnet
     global last_model
     global last_cn_on
     try:
@@ -115,31 +120,35 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
     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(
-                prompt=prompt_mash,
-                num_inference_steps=steps,
-                guidance_scale=cfg_scale,
-                width=width,
-                height=height,
-                generator=generator,
-                joint_attention_kwargs={"scale": lora_scale},
-            ).images[0]
-        else:
-            progress(0, desc="Start Inference with ControlNet.")
-            image = pipe(
-                prompt=prompt_mash,
-                control_image=images,
-                control_mode=modes,
-                num_inference_steps=steps,
-                guidance_scale=cfg_scale,
-                width=width,
-                height=height,
-                controlnet_conditioning_scale=scales,
-                generator=generator,
-                joint_attention_kwargs={"scale": lora_scale},
-            ).images[0]
+        try:
+            if not cn_on or len(modes) == 0:
+                progress(0, desc="Start Inference.")
+                image = pipe(
+                    prompt=prompt_mash,
+                    num_inference_steps=steps,
+                    guidance_scale=cfg_scale,
+                    width=width,
+                    height=height,
+                    generator=generator,
+                    joint_attention_kwargs={"scale": lora_scale},
+                ).images[0]
+            else:
+                progress(0, desc="Start Inference with ControlNet.")
+                image = pipe(
+                    prompt=prompt_mash,
+                    control_image=images,
+                    control_mode=modes,
+                    num_inference_steps=steps,
+                    guidance_scale=cfg_scale,
+                    width=width,
+                    height=height,
+                    controlnet_conditioning_scale=scales,
+                    generator=generator,
+                    joint_attention_kwargs={"scale": lora_scale},
+                ).images[0]
+        except Exception as e:
+            print(e)
+            raise Exception(f"Inference Error: {e}")
     return image
 
 def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height,
@@ -320,6 +329,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
                     gr.Markdown("[Check the list of FLUX LoRas](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.1-dev)", elem_id="lora_list")
                 custom_lora_info = gr.HTML(visible=False)
                 custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
+                deselect_lora_button = gr.Button("Deselect LoRA", variant="secondary")
             with gr.Column(scale=4):
                 result = gr.Image(label="Generated Image", format="png", show_share_button=False)
 
@@ -422,6 +432,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
         outputs=[result, seed]
     )
 
+    deselect_lora_button.click(deselect_lora, None, [prompt, selected_info, selected_index, width, height])
     gr.on(
         triggers=[model_name.change, cn_on.change],
         fn=change_base_model,
@@ -443,6 +454,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]])
         gr.on(
             triggers=[lora_download[i].click],
             fn=download_my_lora,
@@ -463,7 +475,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
         ).success(get_repo_safetensors, [lora_repo[i]], [lora_weights[i]], queue=False, show_api=False
         ).success(apply_lora_prompt, [lora_info[i]], [lora_trigger[i]], queue=False, show_api=False
         ).success(compose_lora_json, [lora_repo_json, lora_num[i], lora_repo[i], lora_wt[i], lora_weights[i], lora_trigger[i]], [lora_repo_json], queue=False, show_api=False)
-
+        
     for i, m in enumerate(cn_mode):
         gr.on(
             triggers=[cn_mode[i].change, cn_scale[i].change],

mod.py

@@ -62,6 +62,21 @@ def clear_cache():
     gc.collect()
 
 
+def deselect_lora():
+    selected_index = gr.State(None)
+    new_placeholder = "Type a prompt"
+    updated_text = ""
+    width = 1024
+    height = 1024
+    return (
+        gr.update(placeholder=new_placeholder),
+        updated_text,
+        selected_index,
+        width,
+        height,
+    )
+
+
 def get_repo_safetensors(repo_id: str):
     from huggingface_hub import HfApi
     api = HfApi()
@@ -311,4 +326,5 @@ def enhance_prompt(input_prompt):
 
 load_prompt_enhancer.zerogpu = True
 fuse_loras.zerogpu = True
-
+preprocess_image.zerogpu = True
+get_control_params.zerogpu = True