Upload web-ui.py

web-ui.py

@@ -60,7 +60,7 @@ def convert_model(checkpoint_path, output_path):
 def load_model(model_name):
     if model_name in model_cache:
         return model_cache[model_name]
-
+    print(f"loading model {model_name}")
     # Limit cache size
     if len(model_cache) >= max_cache_size:
         model_cache.pop(next(iter(model_cache)))
@@ -96,9 +96,10 @@ def load_model(model_name):
     return ip_model
 
 # Function to process image and generate output
-def generate_image(input_image, positive_prompt, negative_prompt, width, height, model_name, num_inference_steps, seed, randomize_seed, num_images, batch_size, enable_shortcut, s_scale):
+def generate_image(input_image, positive_prompt, negative_prompt, width, height, model_name, num_inference_steps, seed, randomize_seed, num_images, batch_size, enable_shortcut, s_scale, custom_model_path):
     saved_images = []
-
+    if custom_model_path:
+        model_name = custom_model_path
     # Load and prepare the model
     ip_model = load_model(model_name)
 
@@ -170,6 +171,7 @@ with gr.Blocks() as demo:
         negative_prompt = gr.Textbox(label="Negative Prompt")    
     with gr.Row():            
         model_selector = gr.Dropdown(label="Select Model", choices=static_model_names, value=static_model_names[0])
+        custom_model_path = gr.Textbox(label="Custom Model Path (Optional)")
 
     with gr.Column():
         output_gallery = gr.Gallery(label="Generated Images")
@@ -184,8 +186,8 @@ with gr.Blocks() as demo:
 
     generate_btn.click(
         generate_image,
-        inputs=[input_image, positive_prompt, negative_prompt, width, height, model_selector, num_inference_steps, seed, randomize_seed, num_images, batch_size, enable_shortcut, s_scale],
-        outputs=[output_gallery, output_text, display_seed],
+        inputs=[input_image, positive_prompt, negative_prompt, width, height, model_selector, num_inference_steps, seed, randomize_seed, num_images, batch_size, enable_shortcut, s_scale, custom_model_path],
+        outputs=[output_gallery, output_text, display_seed]
     )
 	
     convert_btn.click(