Added possibility to specify safetensors

app.py

@@ -60,7 +60,7 @@ def resize_image(input_path, output_path, target_height):
 
     return output_path
 
-def infer(use_custom_model, model_name, custom_lora_weight, image_in, prompt, negative_prompt, preprocessor, controlnet_conditioning_scale, guidance_scale, inf_steps, seed, progress=gr.Progress(track_tqdm=True)):
+def infer(use_custom_model, model_name, weight_name, custom_lora_weight, image_in, prompt, negative_prompt, preprocessor, controlnet_conditioning_scale, guidance_scale, inf_steps, seed, progress=gr.Progress(track_tqdm=True)):
     prompt = prompt
     negative_prompt = negative_prompt
     generator = torch.Generator(device=device).manual_seed(seed)
@@ -88,7 +88,7 @@ def infer(use_custom_model, model_name, custom_lora_weight, image_in, prompt, ne
         custom_model = model_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)
     
         lora_scale=custom_lora_weight
 
@@ -183,13 +183,14 @@ with gr.Blocks(css=css) as demo:
         use_custom_model = gr.Checkbox(label="Use a public custom model ?(optional)", value=False, info="To use a private model, you'll prefer to duplicate the space with your own access token.")
         with gr.Row():
             model_name = gr.Textbox(label="Custom Model to use", placeholder="username/my_custom_public_model")
+            weight_name = gr.Textbox(label="Specific safetensor", placeholder="specific_weight.safetensors")
             custom_lora_weight = gr.Slider(label="Custom model weights", minimum=0.1, maximum=0.9, step=0.1, value=0.9)
         submit_btn = gr.Button("Submit")
         result = gr.Image(label="Result")
 
     submit_btn.click(
         fn = infer,
-        inputs = [use_custom_model, model_name, custom_lora_weight, image_in, prompt, negative_prompt, preprocessor, controlnet_conditioning_scale, guidance_scale, inf_steps, seed],
+        inputs = [use_custom_model, model_name, weight_name, custom_lora_weight, image_in, prompt, negative_prompt, preprocessor, controlnet_conditioning_scale, guidance_scale, inf_steps, seed],
         outputs = [result]
     )