lora model

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 import numpy as np
 import random
-
 # import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionPipeline
+from peft import PeftModel, PeftConfig
 import torch
 from sympy.core.random import choice
 from rembg import remove
@@ -14,7 +14,6 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 # model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
 
-
 # @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
         model_id,
@@ -35,11 +34,17 @@ def infer(
         torch_dtype = torch.float16
     else:
         torch_dtype = torch.float32
-    #print(model_id)
-    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
+
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype, use_safetensors=True)
+    is_lora = False
     if model_id == "CompVis/stable-diffusion-v1-4" and lora == "pepe":
         lora_id = "seregasmirnov/pepe-lora"
-        pipe.load_lora_weights(lora_id)
+        pipe.unet = PeftModel.from_pretrained(
+            pipe.unet, 
+            lora_id, 
+            adapter_name="default"
+          )
+        is_lora = True
     pipe = pipe.to(device)
 
     if randomize_seed:
@@ -47,16 +52,28 @@ def infer(
 
     generator = torch.Generator().manual_seed(seed)
 
-    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]
-
+    if is_lora:
+        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}
+        ).images[0]
+    else:
+        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]
+        
     if del_back:
         image = remove(image)
         
@@ -65,6 +82,7 @@ def infer(
 
 
 examples = [
+    "sticker of a happy cat climbing a tree",
     "cute animal",
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
     "An astronaut riding a green horse",
@@ -97,7 +115,7 @@ with gr.Blocks(css=css) as demo:
                 step=0.1,
                 value=1.0,
                 visible=False,  
-                info="Adjust LoRA adapter strength"
+                info="setup lora strength"
             )
 
         def setup_lora(sel_model, sel_lora):
@@ -108,12 +126,6 @@ with gr.Blocks(css=css) as demo:
                     return [gr.Dropdown(choices=["None", "pepe"], info="Choose lora", visible=True), gr.Slider(visible=True)]
             else:
                 return [gr.Dropdown(choices=["None", "pepe"], info="Choose lora", visible=False), gr.Slider(visible=False)]
-
-        def setup_lora_scale(selected_option):
-            if selected_option == "None":
-                return gr.Slider(visible=False)
-            else:
-                return gr.Slider(visible=True)
                
         model_id.change(
             fn=setup_lora,
@@ -125,12 +137,6 @@ with gr.Blocks(css=css) as demo:
             inputs=[model_id, lora],
             outputs=[lora, lora_scale])
         
-
-        #lora.change(
-        #    fn=setup_lora_scale, 
-        #    inputs=lora, 
-        #    outputs=lora_scale)
-       
         with gr.Row():
             del_back = gr.Checkbox(label="Delete background", value=False)
 
@@ -140,7 +146,7 @@ with gr.Blocks(css=css) as demo:
                 max_lines=1,
                 placeholder="Enter your prompt",
                 container=False,
-                value="cute_animal",
+                value="sticker of a happy cat climbing a tree",
             )
 
             run_button = gr.Button("Run", scale=0, variant="primary")
@@ -153,7 +159,7 @@ with gr.Blocks(css=css) as demo:
                 max_lines=1,
                 placeholder="Enter a negative prompt",
                 visible=True,
-                value="cat, dog",
+                value="",
             )
             gr.Examples(examples=neg_examples, inputs=[negative_prompt])
 
@@ -173,7 +179,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                    value=512,  
                 )
 
                 height = gr.Slider(
@@ -181,7 +187,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                    value=512,  
                 )
 
                 guidance_scale = gr.Slider(
@@ -189,7 +195,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=7.0,#0.0,  # Replace with defaults that work for your model
+                    value=4.0,
                 )
 
                 num_inference_steps = gr.Slider(
@@ -197,7 +203,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=1,
                     maximum=50,
                     step=1,
-                    value=2,  # Replace with defaults that work for your model
+                    value=50,  
                 )
 
         gr.Examples(examples=examples, inputs=[prompt])