Update app.py

app.py

@@ -22,7 +22,7 @@ device="cuda"
 pipe = pipe.to(device)
 
 @spaces.GPU
-def run(prompt, negative_prompt=None, guidance_scale=7.0, pag_scale=3.0, pag_layers=["mid"], randomize_seed=True, seed=42, progress=gr.Progress(track_tqdm=True)):
+def run(prompt, negative_prompt=None, guidance_scale=7.0, pag_scale=3.0, pag_layers=["mid"], randomize_seed=True, seed=42, lora=None, progress=gr.Progress(track_tqdm=True)):
     prompt = prompt.strip()
     negative_prompt = negative_prompt.strip() if negative_prompt and negative_prompt.strip() else None
     print(f"Initial seed for prompt `{prompt}`", seed)
@@ -31,6 +31,10 @@ def run(prompt, negative_prompt=None, guidance_scale=7.0, pag_scale=3.0, pag_lay
     
     if not prompt and not negative_prompt:
         guidance_scale = 0.0
+    pipe.unload_lora_weights()
+    if lora is not None:
+        pipe.load_lora_weights(lora, adapter_name="custom")
+        pipe.fuse_lora(lora_scale=0.9)
     print(f"Seed before sending to generator for prompt: `{prompt}`", seed)
     generator = torch.Generator(device="cuda").manual_seed(seed)
     image_pag = pipe(prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, pag_scale=pag_scale, pag_applied_layers=pag_layers, generator=generator, num_inference_steps=25).images[0]    
@@ -56,6 +60,7 @@ with gr.Blocks(css=css, theme=theme) as demo:
         prompt = gr.Textbox(show_label=False, scale=4, placeholder="Your prompt", info="Leave blank to test unconditional generation")
         button = gr.Button("Generate", min_width=120)
       output = ImageSlider(label="Left: PAG, Right: No PAG", interactive=False)
+      lora = gr.Textbox(label="Custom LoRA path")
       with gr.Accordion("Advanced Settings", open=False):
         guidance_scale = gr.Number(label="CFG Guidance Scale", info="The guidance scale for CFG, ignored if no prompt is entered (unconditional generation)", value=7.0)
         negative_prompt = gr.Textbox(label="Negative prompt", info="Is only applied for the CFG part, leave blank for unconditional generation")
@@ -70,7 +75,7 @@ with gr.Blocks(css=css, theme=theme) as demo:
             prompt.submit
         ],
         fn=run,
-        inputs=[prompt, negative_prompt, guidance_scale, pag_scale, pag_layers, randomize_seed, seed],
+        inputs=[prompt, negative_prompt, guidance_scale, pag_scale, pag_layers, randomize_seed, seed, lora],
         outputs=[output, seed],
     )
 if __name__ == "__main__":