Update app.py

app.py

@@ -18,9 +18,9 @@ pipe = DiffusionPipeline.from_pretrained(
 
 device="cuda" if torch.cuda.is_available() else "cpu"
 
-def load_model(custom_model):
+def load_model(custom_model, weight_name):
     # This is where you load your trained weights
-    pipe.load_lora_weights(custom_model, use_auth_token=True)
+    pipe.load_lora_weights(custom_model, weight_name=weight_name, use_auth_token=True)
     pipe.to(device)
     return "Model loaded!"
 
@@ -88,6 +88,7 @@ with gr.Blocks(css=css) as demo:
         with gr.Row():
             with gr.Column():
                 custom_model = gr.Textbox(label="Your custom model ID", placeholder="your_username/your_trained_model_name", info="Make sure your model is set to PUBLIC ")
+                weight_name = gr.Textbox(label="Safetensors file", value="pytorch_lora_weights.safetensors")
                 model_status = gr.Textbox(label="model status", interactive=False)
             load_model_btn = gr.Button("Load my model")
         
@@ -125,7 +126,7 @@ with gr.Blocks(css=css) as demo:
 
     load_model_btn.click(
         fn = load_model,
-        inputs=[custom_model],
+        inputs=[custom_model, weight_name],
         outputs = [model_status]
     )
     submit_btn.click(