Update app.py

app.py

@@ -2,9 +2,10 @@ import gradio as gr
 import numpy as np
 import random
 from peft import PeftModel, LoraConfig
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, StableDiffusionControlNetPipeline
 from diffusers import ControlNetModel
 import torch
+from PIL import Image
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 if torch.cuda.is_available():
@@ -44,45 +45,54 @@ def infer(
     height,
     guidance_scale,
     num_inference_steps,
-    lscale,
-    controlnet_enabled,
-    control_strength,
-    control_mode,
-    control_image,
-    ip_adapter_enabled,
-    ip_adapter_scale,
-    ip_adapter_image,
+    lscale=0.0,
+    controlnet_enabled=False,
+    control_strength=0.0,
+    control_mode=None,
+    control_image=None,
+    ip_adapter_enabled=False,
+    ip_adapter_scale=0.0,
+    ip_adapter_image=None,
     progress=gr.Progress(track_tqdm=True),
 ):
+    control_strength=float(control_strength)
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
-    
+    if ip_adapter_enabled:
+        print("ip_adapter_image")
+        ip_adapter_image = ip_adapter_image.convert('RGB').resize((510, 510))
+        print("ip_adapter_image",ip_adapter_image.size)
     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")
-        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.safety_checker = None
-        pipe.unet = PeftModel.from_pretrained(pipe.unet, "um235/cartoon_cat_stickers")
+    if controlnet_enabled and control_image:
+        controlnet_model = ControlNetModel.from_pretrained(CONTROLNET_MODES.get(control_mode))
+        if model_id == "SD1.5 + lora Unet TextEncoder" or model_id == "SD1.5 + lora Unet":
+            pipe=StableDiffusionControlNetPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5",controlnet=controlnet_model)
+        else:
+            pipe=StableDiffusionControlNetPipeline.from_pretrained(model_id, controlnet=controlnet_model)
     else:
-        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
+        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.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.unet = PeftModel.from_pretrained(pipe.unet, "um235/cartoon_cat_stickers")
+        else:
+            pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
+    if ip_adapter_enabled:
+            print("ip_adapter_enabled",ip_adapter_enabled)
+            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.safety_checker = None
+           
     pipe = pipe.to(device)
 
     image = pipe(
         prompt=prompt,
+        image=control_image,
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
@@ -90,8 +100,8 @@ def infer(
         height=height,
         generator=generator,
         cross_attention_kwargs={"scale": lscale},
-        control_image=control_image,
-        controlnet_conditioning_scale=control_strength
+        controlnet_conditioning_scale=control_strength,
+        ip_adapter_image=ip_adapter_image,
     ).images[0]
 
     return image, seed
@@ -187,6 +197,7 @@ 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")
@@ -235,7 +246,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=9.0,  
+                    value=7.0,  
                 )
 
                 num_inference_steps = gr.Slider(