Add lora

app.py

@@ -3,13 +3,16 @@ import numpy as np
 import random
 
 # import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, AutoPipelineForText2Image
+from peft import PeftModel
 import torch
 
 device = "cuda" if torch.cuda.is_available() \
         else "xpu" if torch.xpu.is_available() \
         else "cpu"
 current_model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+current_lora_repo = None
+current_lora_scale = 1.0
 
 if torch.cuda.is_available() or torch.xpu.is_available():
     torch_dtype = torch.float16
@@ -32,22 +35,39 @@ def clean_vram():
 def infer(
     prompt,
     model_repo,
+    lora_repo,
+    lora_scale,
     negative_prompt,
     seed,
     randomize_seed,
     width,
     height,
     guidance_scale,
+    pag_scale,
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
-    global current_model_repo_id, pipe
+    global current_model_repo_id, current_lora_repo, current_lora_scale, pipe
 
-    if model_repo != current_model_repo_id:
-        print(f"The model changed to {model_repo}, reloading pipeline...")
+    if lora_repo == "None":
+        lora_repo = None
+
+    if (model_repo != current_model_repo_id) or (lora_repo != current_lora_repo) or (current_lora_scale != lora_scale):
+        print(f"The model changed to {model_repo}, {lora_repo} lora, reloading pipeline...")
+        current_model_repo_id = model_repo
+        current_lora_repo = lora_repo
+        current_lora_scale = lora_scale
         del pipe
         clean_vram()
+
         pipe = DiffusionPipeline.from_pretrained(model_repo, torch_dtype=torch_dtype).to(device)
+        if lora_repo:
+            pipe.unet = PeftModel.from_pretrained(pipe.unet, lora_repo, subfolder="unet").to(device)
+            pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, lora_repo, subfolder="text_encoder").to(device)
+            pipe.unet.load_state_dict({k: lora_scale*v if 'lora' in k else v for k, v in pipe.unet.state_dict().items()})
+            pipe.text_encoder.load_state_dict({k: lora_scale*v if 'lora' in k else v for k, v in pipe.text_encoder.state_dict().items()})
+
+        pipe = AutoPipelineForText2Image.from_pipe(pipe, enable_pag=True)
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -58,6 +78,7 @@ def infer(
         prompt=prompt,
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
+        pag_scale=pag_scale,
         num_inference_steps=num_inference_steps,
         width=width,
         height=height,
@@ -88,7 +109,7 @@ with gr.Blocks(css=css) as demo:
 
         model_repo = gr.Dropdown(
             label="Model repository path",
-            choices=["stabilityai/sdxl-turbo", "CompVis/stable-diffusion-v1-4"],
+            choices=["stabilityai/sdxl-turbo", "CompVis/stable-diffusion-v1-4", "stable-diffusion-v1-5/stable-diffusion-v1-5"],
             allow_custom_value=True
         )
 
@@ -148,6 +169,14 @@ with gr.Blocks(css=css) as demo:
                     step=0.1,
                     value=0.0,  # Replace with defaults that work for your model
                 )
+                
+                pag_scale = gr.Slider(
+                    label="PAG scale",
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=0.0,  # Replace with defaults that work for your model
+                )
 
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
@@ -157,6 +186,21 @@ with gr.Blocks(css=css) as demo:
                     value=2,  # Replace with defaults that work for your model
                 )
 
+            with gr.Row():
+                lora_repo = gr.Dropdown(
+                    label="LoRA repository path",
+                    choices=["None", "AbstractQbit/biskvit_cat_lora"],
+                    allow_custom_value=True
+                )
+
+                lora_scale = gr.Slider(
+                    label="LoRA scale",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.1,
+                    value=1.0,  # Replace with defaults that work for your model
+                )
+
         gr.Examples(examples=examples, inputs=[prompt])
     gr.on(
         triggers=[run_button.click, prompt.submit],
@@ -164,12 +208,15 @@ with gr.Blocks(css=css) as demo:
         inputs=[
             prompt,
             model_repo,
+            lora_repo,
+            lora_scale,
             negative_prompt,
             seed,
             randomize_seed,
             width,
             height,
             guidance_scale,
+            pag_scale,
             num_inference_steps,
         ],
         outputs=[result, seed],