lora fixes

app.py

@@ -3,10 +3,7 @@ import numpy as np
 import random
 import spaces
 import torch
-from diffusers import DiffusionPipeline#, FlowMatchEulerDiscreteScheduler
-# from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
-# from huggingface_hub import hf_hub_download
-# import os
+from diffusers import DiffusionPipeline
 
 dtype = torch.bfloat16
 
@@ -27,6 +24,12 @@ print(f"Device is: {device}")
 # Initialize the pipeline globally
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)
 
+lora_weights = {
+    "jerry": {"path": "bryanbrunetti/jerryfluxlora", "weight_name": "jerry_000003000.safetensors"},
+    "sebastian": {"path": "bryanbrunetti/sebastianfluxlora", "weight_name": "sadie_flux.safetensors"},
+    "sadie": {"path": "bryanbrunetti/sadiefluxlora", "weight_name": "sebastian_flux_000001500.safetensors"},
+}
+
 
 @spaces.GPU(duration=120)
 def infer(prompt, lora_models, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0,
@@ -35,10 +38,10 @@ def infer(prompt, lora_models, seed=42, randomize_seed=False, width=1024, height
     
     # Load LoRAs if specified
     if lora_models:
-        print(f"Loading Loras: {lora_models}")
         try:
             for lora_model in lora_models:
-                pipe.load_lora_weights(lora_model)
+                print(f"loading LoRA: {lora_model}")
+                pipe.load_lora_weights(lora_weights[lora_model], weight_name=lora_weights[lora_model]["weight_name"])
         except Exception as e:
             return None, seed, f"Failed to load LoRA model: {str(e)}"
     
@@ -88,11 +91,8 @@ with gr.Blocks(css=css) as demo:
         #     label="LoRA Model ID (optional)",
         #     placeholder="Enter Hugging Face LoRA model ID",
         # )
-        lora_models = gr.Dropdown([
-            ("jerry", "bryanbrunetti/jerryfluxlora"),
-            ("sebastian", "bryanbrunetti/sebastianfluxlora"),
-            ("sadie", "bryanbrunetti/sadiefluxlora")
-        ], multiselect=True, info="load lora (optional) use the name in the prompt", label="Choose People")
+        lora_models = gr.Dropdown(list(lora_weights.keys()), multiselect=True,
+                                  info="Load LoRA (optional) use the name in the prompt", label="Choose LoRAs")
         
         result = gr.Image(label="Result", show_label=False)
         
@@ -139,7 +139,7 @@ with gr.Blocks(css=css) as demo:
                 )
         
         output_message = gr.Textbox(label="Output Message")
-        
+    
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
@@ -147,4 +147,4 @@ with gr.Blocks(css=css) as demo:
         outputs=[result, seed, output_message]
     )
 
-demo.launch()
+demo.launch()