Update app.py

app.py

@@ -3,9 +3,10 @@ import numpy as np
 import random
 import os
 import torch
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, ControlNetModel
 from peft import PeftModel, LoraConfig
 
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_id_default = "stable-diffusion-v1-5/stable-diffusion-v1-5"
 
@@ -29,7 +30,14 @@ def infer(
     guidance_scale=7.0,
     lora_scale=1.0,
     num_inference_steps=20,
-    progress=gr.Progress(track_tqdm=True),
+    controlnet_checkbox=False,
+    controlnet_strength=0.0,
+    controlnet_mode="edge_detection",
+    controlnet_image=None,
+    ip_adapter_checkbox=False,
+    ip_adapter_scale=0.0,
+    ip_adapter_image=None,
+    progress=gr.Progress(track_tqdm=True),    
 ):  
     generator = torch.Generator(device).manual_seed(seed)
 
@@ -40,9 +48,52 @@ def infer(
     if model_id is None:
         raise ValueError("Please specify the base model name or path")
 
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, 
-                                                   torch_dtype=torch_dtype, 
-                                                   safety_checker=None).to(device)
+    if controlnet_checkbox:
+        if controlnet_mode == "depth_map":
+            controlnet = ControlNetModel.from_pretrained(
+                "lllyasviel/sd-controlnet-depth",
+                cache_dir="./models_cache",
+                torch_dtype=torch_dtype
+            )
+        elif controlnet_mode == "pose_estimation":
+            controlnet = ControlNetModel.from_pretrained(
+                "lllyasviel/sd-controlnet-openpose",
+                cache_dir="./models_cache",
+                torch_dtype=torch_dtype
+            )
+        elif controlnet_mode == "normal_map":
+            controlnet = ControlNetModel.from_pretrained(
+                "lllyasviel/sd-controlnet-normal",
+                cache_dir="./models_cache",
+                torch_dtype=torch_dtype
+            )
+        elif controlnet_mode == "scribbles":
+            controlnet = ControlNetModel.from_pretrained(
+                "lllyasviel/sd-controlnet-scribble",
+                cache_dir="./models_cache",
+                torch_dtype=torch_dtype
+            )
+        else:
+            controlnet_mode == "edge_detection":
+            controlnet = ControlNetModel.from_pretrained(
+                "lllyasviel/sd-controlnet-canny",
+                cache_dir="./models_cache",
+                torch_dtype=torch_dtype
+            )
+        pipe = StableDiffusionPipeline.from_pretrained(model_id, 
+                                                       controlnet=controlnet,
+                                                       torch_dtype=torch_dtype, 
+                                                       safety_checker=None).to(device)
+    else:
+        pipe = StableDiffusionPipeline.from_pretrained(model_id, 
+                                                       torch_dtype=torch_dtype, 
+                                                       safety_checker=None).to(device)
+
+
+    if ip_adapter_checkbox:
+        pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter-plus_sd15.bin")
+        pipe.set_ip_adapter_scale(ip_adapter_scale)
+    
     pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir)
     pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, text_encoder_sub_dir)
 
@@ -54,16 +105,31 @@ def infer(
         pipe.text_encoder.half()
 
     pipe.to(device)
-    
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+
+    if controlnet_checkbox:
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+            image=controlnet_image, 
+            controlnet_conditioning_scale=controlnet_strength,
+            ip_adapter_image=ip_adapter_image if ip_adapter_checkbox else None
+        ).images[0]
+    else:
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+            ip_adapter_image=ip_adapter_image if ip_adapter_checkbox else None
+        ).images[0]
     
     return image
 
@@ -138,25 +204,24 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                 label="ControlNet",
             )
             with gr.Column(visible=False) as controlnet_params:
-                control_strength = gr.Slider(
+                controlnet_strength = gr.Slider(
                     label="ControlNet conditioning scale",
                     minimum=0.0,
                     maximum=1.0,
                     step=0.01,
                     value=1.0,  
                 )
-                control_mode = gr.Dropdown(
+                controlnet_mode = gr.Dropdown(
                     label="ControlNet mode",
                     choices=["edge_detection", 
+                             "depth_map",
                              "pose_estimation", 
-                             "straight_line_detection",
-                             "hed_boundary",
-                             "scribbles",
-                             "human pose"],
+                             "normal_map",
+                             "scribbles"],
                     value="edge_detection",
                     max_choices=1
                 )
-                condition_image = gr.Image(
+                controlnet_image = gr.Image(
                     label="ControlNet condition image",
                     type="pil",
                     format="png"
@@ -168,27 +233,26 @@ with gr.Blocks(css=css, fill_height=True) as demo:
             )
 
         with gr.Row():
-            controlnet_checkbox = gr.Checkbox(
+            ip_adapter_checkbox = gr.Checkbox(
                 label="IPAdapter",
             )
-            with gr.Column(visible=False) as controlnet_params:
-                control_strength = gr.Slider(
-                    label="ControlNet conditioning scale",
+            with gr.Column(visible=False) as ip_adapter_params:
+                ip_adapter_scale = gr.Slider(
+                    label="IPAdapter scale",
                     minimum=0.0,
                     maximum=1.0,
                     step=0.01,
                     value=1.0,  
                 )
-                control_mode = gr.Dropdown(
-                    label="ControlNet mode",
-                    choices=["edge_detection", "other"],
-                    value="edge_detection",
-                    max_choices=1
+                ip_adapter_image = gr.Image(
+                    label="IPAdapter condition image",
+                    type="pil",
+                    format="png"
                 )
-            controlnet_checkbox.change(
+            ip_adapter_checkbox.change(
                 fn=lambda x: gr.Row.update(visible=x),
-                inputs=controlnet_checkbox,
-                outputs=controlnet_params
+                inputs=ip_adapter_checkbox,
+                outputs=ip_adapter_params
             )
             
         with gr.Accordion("Optional Settings", open=False):
@@ -225,8 +289,14 @@ with gr.Blocks(css=css, fill_height=True) as demo:
             seed,
             guidance_scale,      
             lora_scale,
-            num_inference_steps
-            
+            num_inference_steps,
+            controlnet_checkbox,
+            controlnet_strength,
+            controlnet_mode,
+            controlnet_image,
+            ip_adapter_checkbox,
+            ip_adapter_scale,
+            ip_adapter_image  
         ],
         outputs=[result],
     )