Update app.py

app.py

@@ -2,8 +2,6 @@ import gradio as gr
 import numpy as np
 import random
 from peft import PeftModel, LoraConfig
-
-# import spaces #[uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline
 import torch
 
@@ -16,6 +14,23 @@ else:
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+# ControlNet modes list with aliases
+CONTROLNET_MODES = {
+    "Canny Edge Detection": "lllyasviel/control_v11p_sd15_canny",
+    "Pixel to Pixel": "lllyasviel/control_v11e_sd15_ip2p",
+    "Inpainting": "lllyasviel/control_v11p_sd15_inpaint",
+    "Multi-Level Line Segments": "lllyasviel/control_v11p_sd15_mlsd",
+    "Depth Estimation": "lllyasviel/control_v11f1p_sd15_depth",
+    "Surface Normal Estimation": "lllyasviel/control_v11p_sd15_normalbae",
+    "Image Segmentation": "lllyasviel/control_v11p_sd15_seg",
+    "Line Art Generation": "lllyasviel/control_v11p_sd15_lineart",
+    "Anime Line Art": "lllyasviel/control_v11p_sd15_lineart_anime",
+    "Human Pose Estimation": "lllyasviel/control_v11p_sd15_openpose",
+    "Scribble-Based Generation": "lllyasviel/control_v11p_sd15_scribble",
+    "Soft Edge Generation": "lllyasviel/control_v11p_sd15_softedge",
+    "Image Shuffling": "lllyasviel/control_v11e_sd15_shuffle",
+    "Image Tiling": "lllyasviel/control_v11f1e_sd15_tile",
+}
 
 # @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
@@ -44,18 +59,25 @@ def infer(
     generator = torch.Generator().manual_seed(seed)
     
     pipe = None
-    if (model_id=="SD1.5 + lora Unet TextEncoder"): 
-        pipe=DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch_dtype)
-        pipe.unet = PeftModel.from_pretrained(pipe.unet,"um235/VanillaCat",subfolder="unet")
+    if model_id == "SD1.5 + lora Unet TextEncoder": 
+        pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch_dtype)
+        pipe.unet = PeftModel.from_pretrained(pipe.unet, "um235/VanillaCat", subfolder="unet")
         pipe.safety_checker = None
-        pipe.text_encoder= PeftModel.from_pretrained(pipe.text_encoder,"um235/VanillaCat",subfolder="text_encoder")
-    elif  (model_id=="SD1.5 + lora Unet"): 
-        pipe=DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch_dtype)
+        pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, "um235/VanillaCat", subfolder="text_encoder")
+    elif model_id == "SD1.5 + lora Unet": 
+        pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch_dtype)
         pipe.safety_checker = None
-        pipe.unet = PeftModel.from_pretrained(pipe.unet,"um235/cartoon_cat_stickers")
+        pipe.unet = PeftModel.from_pretrained(pipe.unet, "um235/cartoon_cat_stickers")
     else:
-        pipe=DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
+        pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
         pipe.safety_checker = None
+        
+    if controlnet_enabled:
+        controlnet_model = CONTROLNET_MODES.get(control_mode)
+        if controlnet_model:
+            controlnet_model = ControlNetModel.from_pretrained(controlnet_model)
+            pipe.controlnet = controlnet_model
+    
     pipe = pipe.to(device)
 
     image = pipe(
@@ -66,7 +88,9 @@ def infer(
         width=width,
         height=height,
         generator=generator,
-        cross_attention_kwargs={"scale": lscale}
+        cross_attention_kwargs={"scale": lscale},
+        control_image=control_image,
+        controlnet_conditioning_scale=control_strength
     ).images[0]
 
     return image, seed
@@ -86,40 +110,36 @@ css = """
 """
 
 def update_controlnet_visibility(controlnet_enabled):
-    # Возвращаем два значения для обновления видимости control_strength и control_mode
     return gr.update(visible=controlnet_enabled), gr.update(visible=controlnet_enabled), gr.update(visible=controlnet_enabled)
 
 def update_ip_adapter_visibility(ip_adapter_enabled):
-    # Возвращаем два значения для обновления видимости ip_adapter_scale и ip_adapter_image
     return gr.update(visible=ip_adapter_enabled), gr.update(visible=ip_adapter_enabled)
 
 
-
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(" # UM235 DIFFUSION Space")
         
         model_id_input = gr.Dropdown(
-                label="Choose Model",
-                choices=[
-                    "stable-diffusion-v1-5/stable-diffusion-v1-5",  
-                    "CompVis/stable-diffusion-v1-4", 
-                    "SD1.5 + lora Unet TextEncoder",  
-                    "SD1.5 + lora Unet"
-                ],
-                value="SD1.5 + lora Unet TextEncoder",  
-                show_label=True,
-                type="value",  
+            label="Choose Model",
+            choices=[
+                "stable-diffusion-v1-5/stable-diffusion-v1-5",  
+                "CompVis/stable-diffusion-v1-4", 
+                "SD1.5 + lora Unet TextEncoder",  
+                "SD1.5 + lora Unet"
+            ],
+            value="SD1.5 + lora Unet TextEncoder",  
+            show_label=True,
+            type="value",  
         )
         with gr.Row():
-                lscale = gr.Slider(
-                    label="Lora scale",
-                    minimum=0,
-                    maximum=2,
-                    step=0.05,
-                    value=1,  
-                )
-        
+            lscale = gr.Slider(
+                label="Lora scale",
+                minimum=0,
+                maximum=2,
+                step=0.05,
+                value=1,  
+            )
         
         with gr.Row():
             prompt = gr.Text(
@@ -134,20 +154,20 @@ with gr.Blocks(css=css) as demo:
             controlnet_enabled = gr.Checkbox(label="Enable ControlNet", value=False)
 
         with gr.Row():
-                control_strength = gr.Slider(
-                    label="ControlNet scale",
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.05,
-                    value=0.75,
-                    visible=False,
-                )
+            control_strength = gr.Slider(
+                label="ControlNet scale",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.05,
+                value=0.75,
+                visible=False,
+            )
 
-                control_mode = gr.Dropdown(
-                    label="ControlNet Mode",
-                    choices=["edge_detection", "pose_estimation", "depth_estimation"],
-                    value="edge_detection",
-                    visible=False,
+            control_mode = gr.Dropdown(
+                label="ControlNet Mode",
+                choices=list(CONTROLNET_MODES.keys()),
+                value="Canny Edge Detection",
+                visible=False,
             )
 
         control_image = gr.Image(label="ControlNet Image", type="pil", visible=False)
@@ -168,7 +188,6 @@ with gr.Blocks(css=css) as demo:
             ip_adapter_image = gr.Image(label="IP-Adapter Image", type="pil", visible=False)
 
         with gr.Row():
-            
             run_button = gr.Button("Run", scale=0, variant="primary")
 
         result = gr.Image(label="Result", show_label=False)
@@ -198,7 +217,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,  # Replace with defaults that work for your model
+                    value=512, 
                 )
 
                 height = gr.Slider(
@@ -206,7 +225,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,  # Replace with defaults that work for your model
+                    value=512,  
                 )
 
             with gr.Row():
@@ -215,7 +234,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=9.0,  # Replace with defaults that work for your model
+                    value=9.0,  
                 )
 
                 num_inference_steps = gr.Slider(
@@ -223,10 +242,11 @@ with gr.Blocks(css=css) as demo:
                     minimum=1,
                     maximum=50,
                     step=1,
-                    value=36,  # Replace with defaults that work for your model
+                    value=36,  
                 )
 
         gr.Examples(examples=examples, inputs=[prompt])
+    
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
@@ -251,13 +271,13 @@ with gr.Blocks(css=css) as demo:
         ],
         outputs=[result, seed],
     )
+    
     controlnet_enabled.change(
         fn=update_controlnet_visibility,
         inputs=[controlnet_enabled],
         outputs=[control_strength, control_mode, control_image],
     )
 
-    # Updates visibility when the checkbox for IP-Adapter is toggled
     ip_adapter_enabled.change(
         fn=update_ip_adapter_visibility,
         inputs=[ip_adapter_enabled],
@@ -266,4 +286,4 @@ with gr.Blocks(css=css) as demo:
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()