fix lora and sd

app.py

@@ -2,99 +2,204 @@ import gradio as gr
 import numpy as np
 import random
 
-import spaces  # [uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
+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='./lora', 
-    base_model_name_or_path=None, 
-    dtype=torch.float16, 
+    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 = DiffusionPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype)
+
+
+    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)
-    pipe.unet.set_adapter(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)
-    
+        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"
+
+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="sd-legacy/stable-diffusion-v1-5",
-    seed=42,
+    model_repo_id=model_id_default,
+    seed=22,
     guidance_scale=7,
-    num_inference_steps=20,
-    model_lora_id="lora",
-    lora_scale=0.5,
-    use_controlnet=False,
-    controlnet_image=None,
-    control_strength=0.5,
-    control_mode="edge_detection",
-    use_ip_adapter=False,
-    ip_adapter_image=None,
-    ip_adapter_scale=0.5
+    num_inference_steps=50,
+    model_lora_id=model_lora_default,
+    progress=gr.Progress(track_tqdm=True),
     ):
-    
+        
     if randomize_seed:
-        seed = random.randint(0, 1000)
+        seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
-    pipe = get_lora_sd_pipeline(ckpt_dir=f'./{model_lora_id}', base_model_name_or_path=model_repo_id).to("cuda")
-    
-    if use_controlnet and controlnet_image is not None:
-        pipe.enable_controlnet(control_mode, controlnet_image, control_strength)
-    
-    if use_ip_adapter and ip_adapter_image is not None:
-        pipe.enable_ip_adapter(ip_adapter_image, ip_adapter_scale)
     
-    return pipe(prompt=prompt, negative_prompt=negative_prompt, width=width, height=height, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, generator=generator).images[0], 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}")
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Generate LoRa stickers with ControlNet & IP-Adapter")
-    
-    prompt = gr.Text(label="Prompt", placeholder="Enter your prompt")
-    negative_prompt = gr.Text(label="Negative Prompt", placeholder="Enter a negative prompt")
-    randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
-    width = gr.Slider(label="Width", minimum=256, maximum=1024, step=32, value=512)
-    height = gr.Slider(label="Height", minimum=256, maximum=1024, step=32, value=512)
-    
-    use_controlnet = gr.Checkbox(label="Use ControlNet", value=False)
-    controlnet_image = gr.Image(label="ControlNet Image")
-    control_strength = gr.Slider(label="ControlNet Strength", minimum=0, maximum=1, step=0.1, value=0.5)
-    control_mode = gr.Dropdown(label="ControlNet Mode", choices=["edge_detection", "pose_estimation"], value="edge_detection")
-    
-    use_ip_adapter = gr.Checkbox(label="Use IP-Adapter", value=False)
-    ip_adapter_image = gr.Image(label="IP-Adapter Image")
-    ip_adapter_scale = gr.Slider(label="IP-Adapter Scale", minimum=0, maximum=1, step=0.1, value=0.5)
-    
-    run_button = gr.Button("Run")
-    result = gr.Image(label="Result")
-    
-    run_button.click(infer, inputs=[prompt, negative_prompt, randomize_seed, width, height, use_controlnet, controlnet_image, control_strength, control_mode, use_ip_adapter, ip_adapter_image, ip_adapter_scale], outputs=[result])
+    # на вызов pipe с эмбеддингами
+    params = {
+        'prompt_embeds': prompt,
+        'negative_prompt_embeds': 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 = [
+    "A cartoon-style sticker of Elon Musk shaking hands with Donald Trump. Both figures have exaggerated facial expressions, with Musk grinning confidently and Trump giving a signature thumbs-up. The background features a patriotic red, white, and blue color scheme with fireworks exploding behind them.",
+    "A cyberpunk-themed cartoon sticker of Elon Musk standing atop a futuristic Tesla spaceship. He wears a sleek, neon-lit jacket with glowing circuits, while the city skyline behind him is filled with holographic billboards displaying SpaceX and Neuralink logos. His sunglasses reflect the distant stars, adding to the sci-fi aesthetic.",
+    "A medieval fantasy sticker of Elon Musk depicted as a wizard. He holds a glowing blue orb in one hand and a spellbook in the other, wearing a long, starry robe with intricate golden details. His expression is both wise and mischievous, as if he's about to reveal the secrets of the universe. The background features a mystical castle and a dragon flying in the sky.",
+    "A sticker of Elon Musk dressed as a cowboy in the Wild West. He wears a wide-brimmed hat, leather boots, and a long trench coat, standing in front of a saloon with a SpaceX rocket docked nearby instead of a horse. A wanted poster on the wall reads 'Wanted: Mars Pioneer', adding to the playful western theme.",
+    "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="bad face, bad quality, artifacts, low-res, black and white, blurry, low quality, distorted, low resolution, medical mask"
+            )
+
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=1000,
+                step=1,
+                value=42,
+            )
+
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+
+            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
+                )
+
+        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],
+    )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()