hw6

app.py

@@ -1,110 +1,13 @@
 import gradio as gr
 import numpy as np
-import random
-
-import spaces #[uncomment to use ZeroGPU]
-from diffusers import (
-    DiffusionPipeline,
-    StableDiffusionPipeline,
-    StableDiffusionControlNetPipeline,
-    StableDiffusionControlNetImg2ImgPipeline,
-    DPMSolverMultistepScheduler,
-    PNDMScheduler,
-    ControlNetModel
-)
-import torch
-
-from peft import PeftModel, LoraConfig
-import os
-
-def get_lora_sd_pipeline(
-    ckpt_dir='./content/lora',
-    base_model_name_or_path=None,
-    dtype=torch.float16,
-    device="cuda",
-    adapter_name="default"
-):
-    unet_sub_dir = os.path.join(ckpt_dir, "unet")
-    text_encoder_sub_dir = os.path.join(ckpt_dir, "text_encoder")
-    if os.path.exists(text_encoder_sub_dir) and base_model_name_or_path is None:
-        config = LoraConfig.from_pretrained(text_encoder_sub_dir)
-        base_model_name_or_path = config.base_model_name_or_path
-
-    if base_model_name_or_path is None:
-        raise ValueError("Please specify the base model name or path")
-
-
-    pipe = StableDiffusionPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype).to(device)
-    pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)
-
-    if os.path.exists(text_encoder_sub_dir):
-        pipe.text_encoder = PeftModel.from_pretrained(
-            pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name
-        )
-
-    if dtype in (torch.float16, torch.bfloat16):
-        pipe.unet.half()
-        pipe.text_encoder.half()
-
-    return pipe
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
+from methods import infer
 
 model_id_default = "sd-legacy/stable-diffusion-v1-5"
 model_dropdown = ['stabilityai/sdxl-turbo', 'CompVis/stable-diffusion-v1-4', 'sd-legacy/stable-diffusion-v1-5']
 
-model_lora_default = "lora"
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-@spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    randomize_seed,
-    width=512,
-    height=512,
-    model_repo_id=model_id_default,
-    seed=22,
-    guidance_scale=7,
-    num_inference_steps=50,
-    model_lora_id=model_lora_default,
-    progress=gr.Progress(track_tqdm=True),
-    ):
-        
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-    
-    # добавляем обновление pipe по условию
-    if model_repo_id != model_id_default:
-        pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype).to(device)
-        pipe.safety_checker = None
-    else:
-        # добавляем lora
-        pipe = get_lora_sd_pipeline(ckpt_dir='./' + model_lora_id, base_model_name_or_path=model_id_default, 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,
-    }
-    
-    return pipe(**params).images[0], seed    
-
 
 examples = [
     "Cartoon sticker of sad Elon Musk",
@@ -115,91 +18,150 @@ examples = [
     "A parody cartoon sticker of Elon Musk arm-wrestling a robotic version of himself. The robot Musk has glowing red eyes and mechanical arms, while the real Musk smirks confidently. Sparks fly from the table as the intense match unfolds, and the background features a neon sign that reads 'Tesla vs. AI: Ultimate Showdown'."
 ]
 
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"): 
-        gr.Markdown("# Generate LoRa stickers")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            
-            model_repo_id = gr.Dropdown(
-                label="Model Id",
-                choices=model_dropdown,
-                info="Choose model",
-                visible=True,
-                allow_custom_value=True,
-                value=model_id_default,
-            )  
-
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=True,
-                value="monochrome, lowres, bad anatomy, worst quality, low quality, medical mask"
-            )
 
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=1000,
-                step=1,
-                value=22,
-            )
+def on_checkbox_change(use_advanced):
+    visible = use_advanced
+    return (gr.update(visible=visible, interactive=visible),
+            gr.update(visible=visible, interactive=visible),
+            gr.update(visible=visible, interactive=visible))
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+
+        with gr.Column(): 
+
+            gr.Markdown("## ControlNet")
+            use_advanced_controlnet = gr.Checkbox(label="ControlNet Settings")
+            control_strength = gr.Slider(
+                    label="control_strength",
+                    minimum=0,
+                    maximum=1,
+                    step=0.01,
+                    value=0.8,
+                    visible=False)
+            mode = gr.Dropdown(["edge_detection", "pose_estimation"], label="Выбор режима", visible=False)
+            image_upload_cn = gr.Image(label="Загрузите изображение", visible=False)
+            use_advanced_controlnet.change(on_checkbox_change, use_advanced_controlnet, [control_strength, mode, image_upload_cn])
+
+
+            gr.Markdown("## IPAdapter")
+            use_advanced_ip = gr.Checkbox(label="ControlNet Settings")
+            ip_adapter_scale = gr.Slider(
+                    label="ip_adapter_scale",
+                    minimum=0,
+                    maximum=1,
+                    step=0.01,
+                    value=0.8,
+                    visible=False)
+            image_upload_ip = gr.Image(label="Загрузите изображение", visible=False)
+            use_advanced_ip.change(on_checkbox_change, use_advanced_ip, [ip_adapter_scale, image_upload_ip])
+
+
+        with gr.Column(): 
+            gr.Markdown("## Generate")
+
+            with gr.Row():
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="Enter your prompt",
+                    container=False,
+                    )
+                run_button = gr.Button("Run", scale=0, variant="primary")
+
+            result = gr.Image(label="Result", show_label=False)
+
+            with gr.Accordion("Advanced Settings", open=False):
+                
+                model_repo_id = gr.Dropdown(
+                    label="Model Id",
+                    choices=model_dropdown,
+                    info="Choose model",
+                    visible=True,
+                    allow_custom_value=True,
+                    value=model_id_default,
+                )  
+
+                negative_prompt = gr.Text(
+                    label="Negative prompt",
+                    max_lines=1,
+                    placeholder="Enter a negative prompt",
+                    visible=True,
+                    value="monochrome, lowres, bad anatomy, worst quality, low quality, medical mask"
+                )
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=22,
+                )
 
-            with gr.Row():   
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,  # Replace with defaults that work for your model
+                guidance_scale = gr.Slider(
+                    label="guidance_scale",
+                    minimum=0,
+                    maximum=100,
+                    step=1,
+                    value=7,
                 )
 
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,  # Replace with defaults that work for your model
+                num_inference_steps = gr.Slider(
+                    label="num_inference_steps",
+                    minimum=0,
+                    maximum=100,
+                    step=1,
+                    value=50,
                 )
 
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            randomize_seed,
-            width,
-            height,
-            model_repo_id,
-            seed
-        ],
-        outputs=[result, seed],
-    )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+                with gr.Row():   
+                    width = gr.Slider(
+                        label="Width",
+                        minimum=256,
+                        maximum=MAX_IMAGE_SIZE,
+                        step=32,
+                        value=512,  # Replace with defaults that work for your model
+                    )
+
+                    height = gr.Slider(
+                        label="Height",
+                        minimum=256,
+                        maximum=MAX_IMAGE_SIZE,
+                        step=32,
+                        value=512,  # Replace with defaults that work for your model
+                    )
+
+            with gr.Accordion("Prompt examples", open=False):
+                gr.Examples(examples=examples, inputs=[prompt])
+
+            gr.on(
+                triggers=[run_button.click, prompt.submit],
+                fn=infer,
+                inputs=[
+                    prompt,
+                    negative_prompt,
+                    randomize_seed,
+                    width,
+                    height,
+                    model_repo_id,
+                    seed, 
+                    guidance_scale,
+                    num_inference_steps,
+
+                    use_advanced_controlnet,
+                    control_strength,
+                    image_upload_cn,
+
+                    use_advanced_ip,
+                    ip_adapter_scale,
+                    image_upload_ip
+                ],
+                outputs=[result, seed],
+            )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=False, debug=True)

methods.py

@@ -0,0 +1,212 @@
+import PIL
+from PIL import Image
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import cv2 as cv
+import random
+import os
+import spaces
+import gradio as gr
+
+from diffusers import DiffusionPipeline
+from peft import PeftModel, LoraConfig
+
+from diffusers import (
+    StableDiffusionPipeline,
+    StableDiffusionControlNetPipeline,
+    StableDiffusionControlNetImg2ImgPipeline,
+    DPMSolverMultistepScheduler,
+    PNDMScheduler,
+    ControlNetModel
+)
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg, retrieve_timesteps
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.utils import load_image, make_image_grid
+
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+model_id_default = "sd-legacy/stable-diffusion-v1-5"
+model_dropdown = ['stabilityai/sdxl-turbo', 'CompVis/stable-diffusion-v1-4', 'sd-legacy/stable-diffusion-v1-5']
+model_lora_default = "lora"
+
+
+def get_lora_sd_pipeline(
+    ckpt_dir='./' + model_lora_default,
+    base_model_name_or_path=None,
+    dtype=torch.float16,
+    device=DEVICE,
+    adapter_name="default",
+    controlnet=None,
+    ip_adapter=None
+):
+    unet_sub_dir = os.path.join(ckpt_dir, "unet")
+    text_encoder_sub_dir = os.path.join(ckpt_dir, "text_encoder")
+    if os.path.exists(text_encoder_sub_dir) and base_model_name_or_path is None:
+        config = LoraConfig.from_pretrained(text_encoder_sub_dir)
+        base_model_name_or_path = config.base_model_name_or_path
+
+    if base_model_name_or_path is None:
+        raise ValueError("Please specify the base model name or path")
+
+
+    if controlnet and ip_adapter:
+      print('Pipe with ControlNet and IpAdapter')
+
+      controlnet = ControlNetModel.from_pretrained(
+          "lllyasviel/sd-controlnet-canny",
+          cache_dir="./models_cache",
+          torch_dtype=torch.float16
+      )
+      pipe = StableDiffusionControlNetPipeline.from_pretrained(
+         base_model_name_or_path, 
+         torch_dtype=dtype, 
+         controlnet=controlnet).to(device)
+      
+      pipe.load_ip_adapter(
+          "h94/IP-Adapter",
+          subfolder="models",
+          weight_name="ip-adapter-plus_sd15.bin",
+      )
+
+
+    elif controlnet:
+      print('Pipe with ControlNet')
+      controlnet = ControlNetModel.from_pretrained(
+          "lllyasviel/sd-controlnet-canny",
+          cache_dir="./models_cache",
+          torch_dtype=torch.float16)
+      pipe = StableDiffusionControlNetPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype, controlnet=controlnet)
+
+
+    elif ip_adapter:
+      print('Pipe with IpAdapter')
+      pipe = StableDiffusionPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype)
+      pipe.load_ip_adapter(
+          "h94/IP-Adapter",
+          subfolder="models",
+          weight_name="ip-adapter-plus_sd15.bin")
+
+
+    else:
+      print('Pipe with only SD')
+      pipe = StableDiffusionPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype)
+
+
+    pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)
+    if os.path.exists(text_encoder_sub_dir):
+        pipe.text_encoder = PeftModel.from_pretrained(
+            pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name
+        )
+
+    if dtype in (torch.float16, torch.bfloat16):
+        pipe.unet.half()
+        pipe.text_encoder.half()
+
+    pipe.safety_checker = None
+    pipe.to(device)
+    return pipe
+
+
+@spaces.GPU
+def infer(
+    prompt,
+    negative_prompt,
+    randomize_seed,
+    width=512,
+    height=512,
+    model_repo_id=model_id_default, # в get_lora_sd_pipeline - base_model_name_or_path
+    seed=22,
+    guidance_scale=7,
+    num_inference_steps=50,
+
+    use_advanced_controlnet=None,
+    control_strength=None,
+    image_upload_cn=None,
+
+    use_advanced_ip=None,
+    ip_adapter_scale=None,
+    image_upload_ip=None,
+
+    model_lora_id=model_lora_default,
+    progress=gr.Progress(track_tqdm=True),
+    dtype=torch.float16,
+    device=DEVICE,
+    ):
+        
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+
+
+    #1. SD 1.5 + Lora
+    if use_advanced_controlnet == None and use_advanced_ip == None:
+        pipe = get_lora_sd_pipeline(base_model_name_or_path=model_repo_id, 
+                                    dtype=dtype).to(device)                   
+        
+        image = pipe(prompt,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                negative_prompt=negative_prompt,
+                width=width,
+                heigth=height,
+                generator=generator).images[0]
+
+    elif use_advanced_controlnet != None and use_advanced_ip == None:
+        #2. SD 1.5 + Lora + Controlnet
+        edges = cv.Canny(np.array(image_upload_cn))
+        edges = np.repeat(edges[:, :, None], 3, axis=2)
+        edges = Image.fromarray(edges)
+
+        pipe = get_lora_sd_pipeline(base_model_name_or_path=model_repo_id,
+                                    controlnet=True,
+                                    dtype=dtype).to(device)
+        
+        image = pipe(prompt,
+                edges,
+                num_inference_steps = num_inference_steps,
+                controlnet_conditioning_scale=control_strength,
+                negative_prompt=negative_prompt,
+                generator=generator).images[0]
+ 
+
+    elif use_advanced_ip != None and use_advanced_controlnet == None:
+        #3. SD 1.5 + Lora + IpAdapter
+        pipe = get_lora_sd_pipeline(base_model_name_or_path=model_repo_id,
+                                    ip_adapter=True,
+                                    dtype=dtype).to(device)
+        pipe.set_ip_adapter_scale(ip_adapter_scale)
+
+        image = pipe(
+            prompt,
+            edges,
+            ip_adapter_image=image_upload_ip,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            generator=generator).images[0]
+
+    elif use_advanced_ip != None and use_advanced_controlnet != None:
+        #4. SD 1.5 + Lora + IpAdapter + ControlNet
+        pipe = get_lora_sd_pipeline(base_model_name_or_path=model_repo_id,
+                                ip_adapter=True,
+                                controlnet=True,
+                                dtype=dtype).to(device)
+        
+        pipe.set_ip_adapter_scale(ip_adapter_scale)
+        image = pipe(prompt,
+            edges,
+            ip_adapter_image=image_upload_ip,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            controlnet_conditioning_scale=control_strength,
+            height=height,
+            width=width,
+            generator=generator,
+        ).images[0]
+
+    return image, seed