Update app.py

app.py

@@ -57,10 +57,14 @@ def load_model(custom_model):
 
     print(f"Safetensors available: {sfts_available_files}")
 
-    return "Model Ready", gr.update(choices=sfts_available_files, value=sfts_available_files[0], visible=True), gr.update(value=instance_prompt, visible=True)
+    return custom_model, "Model Ready", gr.update(choices=sfts_available_files, value=sfts_available_files[0], visible=True), gr.update(value=instance_prompt, visible=True)
 
-def custom_model_changed(custom_model):
-    return "Model changed, you must reload before re-run"
+def custom_model_changed(custom_model, previous_model):
+    if custom_model != previous_model:
+        status_message = "Model changed, you must reload before re-run"
+    else:
+        status_message = "Model ready"
+    return status_message
 
 @spaces.GPU
 def infer (custom_model, weight_name, prompt, inf_steps, guidance_scale, seed, lora_weight, progress=gr.Progress(track_tqdm=True)):
@@ -159,41 +163,45 @@ with gr.Blocks(css=css) as demo:
 <h2 style="text-align: center;">SD-XL Custom Model Inference</h2>
 <p style="text-align: center;">Use this demo to check results from your previously trained LoRa model.</p>
         """)
-        with gr.Row():
-            with gr.Column():
-                if not is_shared_ui:
-                    your_username = api.whoami()["name"]
-                    my_models = api.list_models(author=your_username, filter=["diffusers", "stable-diffusion-xl", 'lora'])
-                    model_names = [item.modelId for item in my_models]
-
-                if not is_shared_ui:
-                    custom_model = gr.Dropdown(
-                        label = "Your custom model ID",
-                        choices = model_names,
-                        allow_custom_value = True
-                        #placeholder = "username/model_id"
+        with gr.Box():
+            with gr.Row():
+                with gr.Column():
+                    if not is_shared_ui:
+                        your_username = api.whoami()["name"]
+                        my_models = api.list_models(author=your_username, filter=["diffusers", "stable-diffusion-xl", 'lora'])
+                        model_names = [item.modelId for item in my_models]
+    
+                    if not is_shared_ui:
+                        custom_model = gr.Dropdown(
+                            label = "Your custom model ID",
+                            choices = model_names,
+                            allow_custom_value = True
+                            #placeholder = "username/model_id"
+                        )
+                    else:
+                        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.Dropdown(
+                        label="Safetensors file", 
+                        #value="pytorch_lora_weights.safetensors", 
+                        info="specify which one if model has several .safetensors files",
+                        allow_custom_value=True,
+                        visible = False
                     )
-                else:
-                    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"
+                with gr.Column():
+                    load_model_btn = gr.Button("Load my model")
+                    previous_model = gr.Textbox(
+                        visible=False
                     )
-                
-                weight_name = gr.Dropdown(
-                    label="Safetensors file", 
-                    #value="pytorch_lora_weights.safetensors", 
-                    info="specify which one if model has several .safetensors files",
-                    allow_custom_value=True,
-                    visible = False
-                )
-            with gr.Column():
-                load_model_btn = gr.Button("Load my model")
-                model_status = gr.Textbox(
-                    label = "model status",
-                    show_label = False
-                )
-                trigger_word = gr.Textbox(label="Trigger word", interactive=False, visible=False)
+                    model_status = gr.Textbox(
+                        label = "model status",
+                        show_label = False
+                    )
+                    trigger_word = gr.Textbox(label="Trigger word", interactive=False, visible=False)
         
         prompt_in = gr.Textbox(
             label="Your Prompt",
@@ -242,14 +250,14 @@ with gr.Blocks(css=css) as demo:
 
     custom_model.blur(
         fn=custom_model_changed,
-        inputs = [custom_model],
+        inputs = [custom_model, previous_model],
         outputs = [model_status],
         queue = False
     )
     load_model_btn.click(
         fn = load_model,
         inputs=[custom_model],
-        outputs = [model_status, weight_name, trigger_word],
+        outputs = [previous_model, model_status, weight_name, trigger_word],
         queue = False
     )
     submit_btn.click(