add hw6

app.py

@@ -1,12 +1,15 @@
 import os
 import random
 
+import cv2
 import gradio as gr
 import numpy as np
 import torch
 # import spaces #[uncomment to use ZeroGPU]
-from diffusers import StableDiffusionPipeline
+from diffusers import (ControlNetModel, StableDiffusionControlNetPipeline,
+                       StableDiffusionPipeline)
 from peft import LoraConfig, PeftModel
+from PIL import Image
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 # model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
@@ -14,11 +17,12 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # model_dropdown = ["stabilityai/sdxl-turbo", "CompVis/stable-diffusion-v1-4"]
 
 models = [
-    "gstranger/kawaiicat-lora-1.4",
+    # "gstranger/kawaiicat-lora-1.4",
     "CompVis/stable-diffusion-v1-4",
     "stabilityai/sdxl-turbo",
     "sd-legacy/stable-diffusion-v1-5",
 ]
+controlnet_modes = ["canny", "Line Art"]
 
 model_dropdown = [
     "stabilityai/sdxl-turbo",
@@ -26,6 +30,14 @@ model_dropdown = [
     "sd-legacy/stable-diffusion-v1-5",
 ]
 
+def process_control_image(image, mode="canny"):
+  if mode == "canny":
+    image = np.array(image)
+    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+    canny = cv2.Canny(blurred, 50, 150)
+    return Image.fromarray(canny)
+  return image
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
@@ -81,47 +93,92 @@ def infer(
     model_id,
     prompt,
     negative_prompt,
-    randomize_seed,
-    width,
-    height,
-    # model_repo_id=model_repo_id,
+    randomize_seed=False,
+    width=512,
+    height=512,
+    
+    lora_scale=0.8,
+    lora_enable=True,
+    
+    controlnet_enable=False,
+    control_mode="Line Art",
+    control_strength=0.8,
+    control_image=None,
+
+    ip_adapter_enable=False,
+    ip_adapter_scale=0.8,
+    ip_image=None,
+
+
+    torch_dtype=torch_dtype,
     seed=42,
     guidance_scale=7,
     num_inference_steps=50,
     progress=gr.Progress(track_tqdm=True),
-    lora_scale=1,
 ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+    else:
+        seed = 488
 
     generator = torch.Generator().manual_seed(seed)
 
-    if model_id == "gstranger/kawaiicat-lora-1.4":
-        # добавляем lora
-        pipe = get_lora_sd_pipeline(
-            os.path.join(CKPT_DIR, ""), adapter_name="sd-14-lora", dtype=torch_dtype
-        ).to(device)
-        pipe.safety_checker = None
-        print(f"LoRA adapter loaded: {pipe.unet.active_adapters}")
+    params = {'prompt': prompt,
+            'negative_prompt': negative_prompt,
+            'guidance_scale': guidance_scale,
+            'num_inference_steps': num_inference_steps,
+            'width': width,
+            'height': height,
+            'generator': generator,
+            }
+
+    if controlnet_enable:
+        if control_mode == "canny":
+            controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_canny", 
+                                                         torch_dtype=torch_dtype, cache_dir="./models_cache")
+        elif control_mode == "Line Art":
+            controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_lineart", 
+                                                         torch_dtype=torch_dtype, cache_dir="./models_cache")
+            
+        pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id, 
+                                                                 controlnet=controlnet,
+                                                                 torch_dtype=torch_dtype, 
+                                                                 safety_checker=None) #.to(device)
+        
+        params['image'] = process_control_image(control_image, control_mode)
+        params['controlnet_conditioning_scale'] = float(control_strength)
+
 
     else:
-        pipe = StableDiffusionPipeline.from_pretrained(
-            model_id,
-            torch_dtype=torch_dtype,
-            requires_safety_checker=False,
-            safety_checker=None,
+        pipe = StableDiffusionPipeline.from_pretrained(model_id, 
+                                                       torch_dtype=torch_dtype, 
+                                                       safety_checker=None) #.to(device)
+
+    if lora_enable:
+        unet_sub_dir = os.path.join(CKPT_DIR, "unet")
+        text_encoder_sub_dir = os.path.join(CKPT_DIR, "text_encoder")
+        adapter_name="sd-14-lora"
+
+        pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)
+        pipe.text_encoder = PeftModel.from_pretrained(
+            pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name
         )
-        pipe = 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,
-        cross_attention_kwargs={"scale": lora_scale},
+        params['cross_attention_kwargs']={"scale": lora_scale}
+
+    if torch_dtype in (torch.float16, torch.bfloat16):
+        pipe.unet.half()
+        pipe.text_encoder.half()
+
+    if ip_adapter_enable:
+        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'] = process_control_image(ip_image, "")
+
+
+    # pipe.to(device)
+
+
+    image = pipe(**params
     ).images[0]
 
     return image, seed
@@ -155,9 +212,23 @@ with gr.Blocks(css=css) as demo:
             minimum=0,
             maximum=1,
             step=0.01,
-            value=1,
+            value=0.8,
         )
 
+        lora_enable = gr.Checkbox(label="Use LORA", value=True)
+
+        with gr.Columns():
+            controlnet_enable = gr.Checkbox(label="Enable ControlNet")
+            with gr.Accordion("ControlNet Settings", visible=False) as controlnet_accordion:
+                control_mode = gr.Dropdown(controlnet_modes, label="Control Mode", value="canny")
+                control_strength = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="Control Strength")
+                control_image = gr.Image(label="Control Image", type="pil")
+
+            ip_adapter_enable = gr.Checkbox(label="Enable IP-Adapter")
+            with gr.Accordion("IP-Adapter Settings", visible=False) as ipadapter_accordion:
+                ip_adapter_scale = gr.Slider(0, 1, value=0.5, label="IP-Adapter Scale")
+                ip_image = gr.Image(label="Reference Image", type="pil")
+
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -165,8 +236,16 @@ with gr.Blocks(css=css) as demo:
                 max_lines=1,
                 placeholder="Enter your prompt",
                 container=False,
+            )        
+            
+            negative_prompt = gr.Textbox(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter your negative prompt",
+                value="bad anatomy, crop image, bad face of the cat"
             )
 
+
             run_button = gr.Button("Run", scale=0, variant="primary")
 
         result = gr.Image(label="Result", show_label=False)
@@ -187,7 +266,7 @@ with gr.Blocks(css=css) as demo:
                 value=42,
             )
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
 
             with gr.Row():
                 width = gr.Slider(
@@ -224,6 +303,18 @@ with gr.Blocks(css=css) as demo:
                 )
 
         gr.Examples(examples=examples, inputs=[prompt])
+
+        controlnet_enable.change(
+            lambda x: gr.update(visible=x), 
+            controlnet_enable, 
+            controlnet_accordion
+        )
+        ip_adapter_enable.change(
+            lambda x: gr.update(visible=x), 
+            ip_adapter_enable, 
+            ipadapter_accordion
+        )
+
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
@@ -237,10 +328,14 @@ with gr.Blocks(css=css) as demo:
             seed,
             guidance_scale,
             num_inference_steps,
-            lora_scale,
+            lora_enable, lora_scale,
+            controlnet_enable, control_mode, control_strength, control_image,
+            ip_adapter_enable, ip_adapter_scale, ip_image
+
         ],
         outputs=[result, seed],
     )
 
+
 if __name__ == "__main__":
     demo.launch()