[ADD] Add IP adapter and ControlNet

app.py

@@ -25,7 +25,13 @@ else:
 # Cache to avoid re-initializing pipelines repeatedly
 model_cache = {}
 
-def load_pipeline(model_id: str):
+def load_pipeline(model_id,
+                  lora_scale,
+                  controlnet_checkbox,
+                  controlnet_mode,
+                  ip_adapter_checkbox,
+                  ip_adapter_scale
+                  ):
     """
     Loads or retrieves a cached DiffusionPipeline.
     
@@ -34,11 +40,52 @@ def load_pipeline(model_id: str):
     """
     if model_id in model_cache:
         return model_cache[model_id]
-    
+
+    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 = ControlNetModel.from_pretrained(
+                "lllyasviel/sd-controlnet-canny",
+                cache_dir="./models_cache",
+                torch_dtype=torch_dtype
+            )
+        pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id,
+                                                                 controlnet=controlnet,
+                                                                 torch_dtype=torch_dtype,
+                                                                 safety_checker=None).to(device)
+        # params['image'] = controlnet_image
+        # params['controlnet_conditioning_scale'] = float(controlnet_strength)
+    else:
+        pipe = StableDiffusionPipeline.from_pretrained(model_id,
+                                                       torch_dtype=torch_dtype,
+                                                       safety_checker=None).to(device)
+
     if model_id == "YaArtemNosenko/dino_stickers":
         # Use the specified base model for your LoRA adapter.
         base_model = "CompVis/stable-diffusion-v1-4"
-        pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch_dtype)
         # Load the LoRA weights
         pipe.unet = PeftModel.from_pretrained(
             pipe.unet, 
@@ -52,9 +99,21 @@ def load_pipeline(model_id: str):
             subfolder="text_encoder", 
             torch_dtype=torch_dtype
         )
+        pipe.unet.load_state_dict({k: lora_scale * v for k, v in pipe.unet.state_dict().items()})
+        pipe.text_encoder.load_state_dict({k: lora_scale * v for k, v in pipe.text_encoder.state_dict().items()})
     else:
-        pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
-    
+        pipe = DiffusionPipeline.from_pretrained(model_id, 
+                                                 torch_dtype=torch_dtype
+                                                 )
+
+    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)
+        # params['ip_adapter_image'] = ip_adapter_image
+
     pipe.to(device)
     model_cache[model_id] = pipe
     return pipe
@@ -72,17 +131,36 @@ def infer(
     height,
     guidance_scale,
     num_inference_steps,
-    lora_scale,  # New parameter for adjusting LoRA scale
+    lora_scale,                         # New parameter for adjusting LoRA scale
+    controlnet_checkbox=False,          # используем ли мы controlnet
+    controlnet_conditioning_scale=0.0,  # вес для controlnet
+    controlnet_mode="edge_detection",   # вариант controlnet
+    controlnet_image=None,              # картинка для controlnet
+    ip_adapter_checkbox=False,          # используется ли ip адаптера
+    ip_adapter_scale=0.0,               # вес для ip адаптера
+    ip_adapter_image=None,              # картинка для ip адаптера
     progress=gr.Progress(track_tqdm=True),
 ):
     # Load the pipeline for the chosen model
-    pipe = load_pipeline(model_id)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    params = {'prompt': prompt,
+              'negative_prompt': negative_prompt,
+              'guidance_scale': guidance_scale,
+              'num_inference_steps': num_inference_steps,
+              'width': width,
+              'height': height,
+              'generator': generator
+              }
+    pipe = load_pipeline(lora_scale,
+                         controlnet_checkbox,
+                         controlnet_mode,
+                         ip_adapter_checkbox,
+                         ip_adapter_scale
+                         )
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    generator = torch.Generator(device=device).manual_seed(seed)
-
     # If using the LoRA model, update the LoRA scale if supported.
     if model_id == "YaArtemNosenko/dino_stickers":
         # This assumes your pipeline's unet has a method to update the LoRA scale.
@@ -90,17 +168,15 @@ def infer(
             pipe.unet.set_lora_scale(lora_scale)
         else:
             print("Warning: LoRA scale adjustment method not found on UNet.")
+    # если используем controlnet
+    if controlnet_checkbox:
+        params['image'] = controlnet_image
+        params['controlnet_conditioning_scale'] = float(controlnet_conditioning_scale)
+    # если используем IP адаптер
+    if ip_adapter_checkbox:
+        params['ip_adapter_image'] = ip_adapter_image
 
-    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]
-
+    image = pipe(**params).images[0]
     return image, seed
 
 examples = [
@@ -201,6 +277,61 @@ with gr.Blocks(css=css) as demo:
                 value=1.0,
                 info="Adjust the influence of the LoRA weights",
             )
+        with gr.Row():
+            controlnet_checkbox = gr.Checkbox(
+                label="ControlNet",
+                value=False
+            )
+            with gr.Column(visible=False) as controlnet_params:
+                controlnet_conditioning_scale = gr.Slider(
+                    label="ControlNet conditioning scale",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=1.0,  
+                )
+                controlnet_mode = gr.Dropdown(
+                    label="ControlNet mode",
+                    choices=["edge_detection", 
+                             "depth_map",
+                             "pose_estimation", 
+                             "normal_map",
+                             "scribbles"],
+                    value="edge_detection",
+                    max_choices=1
+                )
+                controlnet_image = gr.Image(
+                    label="ControlNet condition image",
+                    type="pil",
+                    format="png"
+                )
+            controlnet_checkbox.change(
+                    fn=lambda x: gr.Row.update(visible=x),
+                    inputs=controlnet_checkbox,
+                    outputs=controlnet_params
+            )
+        with gr.Row():
+            ip_adapter_checkbox = gr.Checkbox(
+                label="IPAdapter",
+                value=False
+            )
+            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,  
+                )
+                ip_adapter_image = gr.Image(
+                    label="IPAdapter condition image",
+                    type="pil"
+                )
+            ip_adapter_checkbox.change(
+                fn=lambda x: gr.Row.update(visible=x),
+                inputs=ip_adapter_checkbox,
+                outputs=ip_adapter_params
+            )
 
         gr.Examples(examples=examples, inputs=[prompt])
     gr.on(