Update app.py

app.py

@@ -17,6 +17,7 @@ MODELS = {
     "Juggernaut-XL-V9-GE-RDPhoto2": "AiWise/Juggernaut-XL-V9-GE-RDPhoto2-Lightning_4S",
     "SatPony-Lightning": "John6666/satpony-lightning-v2-sdxl"
 }
+
 def init_pipeline(model_name):
     config_file = hf_hub_download(
         "xinsir/controlnet-union-sdxl-1.0",
@@ -50,22 +51,24 @@ def init_pipeline(model_name):
 # Initialize with the default model
 default_model_name = "RealVisXL V5.0 Lightning"
 pipe = init_pipeline(default_model_name)
+loaded_model_name = default_model_name  # Track the loaded model
 
 def update_pipeline(model_selection):
-    global pipe
-    if pipe.config.model_name != MODELS[model_selection]:
+    global pipe, loaded_model_name
+    if model_selection != loaded_model_name:
+        print(f"Loading new model: {model_selection}")
         pipe = init_pipeline(model_selection)
+        loaded_model_name = model_selection
     return pipe
-
+    
 @spaces.GPU(duration=12)
 def fill_image(prompt, image, model_selection, paste_back):
-    print(f"Received image: {image}")
     global pipe
     update_pipeline(model_selection)
+    print(f"Received image: {image}")
     if image is None:
         yield None, None
         return
-
     (
         prompt_embeds,
         negative_prompt_embeds,
@@ -78,7 +81,6 @@ def fill_image(prompt, image, model_selection, paste_back):
     binary_mask = alpha_channel.point(lambda p: p > 0 and 255)
     cnet_image = source.copy()
     cnet_image.paste(0, (0, 0), binary_mask)
-
     for image in pipe(
         prompt_embeds=prompt_embeds,
         negative_prompt_embeds=negative_prompt_embeds,
@@ -87,7 +89,6 @@ def fill_image(prompt, image, model_selection, paste_back):
         image=cnet_image,
     ):
         yield image, cnet_image
-
     print(f"{model_selection=}")
     print(f"{paste_back=}")
     if paste_back:
@@ -196,9 +197,9 @@ def preview_image_and_mask(image, width, height, overlap_percentage, resize_opti
     return preview
 
 @spaces.GPU(duration=12)
-def inpaint(prompt, image, inpaint_model, paste_back):
+def inpaint(prompt, image, model_selection, paste_back):
     global pipe
-    update_pipeline(inpaint_model)
+    update_pipeline(model_selection)
     mask = Image.fromarray(image["mask"]).convert("L")
     image = Image.fromarray(image["image"])
     inpaint_final_prompt = f"score_9, score_8_up, score_7_up, {prompt}"
@@ -206,7 +207,7 @@ def inpaint(prompt, image, inpaint_model, paste_back):
     if paste_back:
         result.paste(image, (0, 0), Image.fromarray(255 - np.array(mask)))
     return result
-
+    
 @spaces.GPU(duration=12)
 def outpaint(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
     background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
@@ -237,7 +238,7 @@ def outpaint(image, width, height, overlap_percentage, num_inference_steps, resi
 @spaces.GPU(duration=12)
 def infer(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
     global pipe
-    update_pipeline(model_selection)  # Added this line
+    update_pipeline(model_selection)  # Ensure model_selection is defined
     background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
     if not can_expand(background.width, background.height, width, height, alignment):
         alignment = "Middle"
@@ -348,9 +349,11 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                     with gr.Column():
                         model_selection = gr.Dropdown(
                             choices=list(MODELS.keys()),
-                            value="RealVisXL V5.0 Lightning",
+                            value=default_model_name,
                             label="Model",
                         )
+                        loading_message = gr.Label(label="Status", value="", visible=False)  # Added loading message label
+
                         with gr.Row():
                             run_button = gr.Button("Generate")
                             paste_back = gr.Checkbox(True, label="Paste back original")
@@ -362,8 +365,8 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                         interactive=False,
                         label="Generated Image",
                     )
+
                 use_as_input_button = gr.Button("Use as Input Image", visible=False)
-                loading_message = gr.Label(label="Status", value="", visible=False)  # Added loading message label
 
                 use_as_input_button.click(
                     fn=use_output_as_input, inputs=[result], outputs=[input_image]
@@ -379,7 +382,7 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                 ).then(
                     fn=lambda: gr.update(value="Loading Model...", visible=True),  # Show loading message
                     inputs=None,
-                    outputs=[loading_message, use_as_input_button]
+                    outputs=loading_message
                 ).then(
                     fn=fill_image,
                     inputs=[prompt, input_image, model_selection, paste_back],
@@ -387,12 +390,12 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                 ).then(
                     fn=lambda: gr.update(value="Model Loaded", visible=True),  # Show loaded message
                     inputs=None,
-                    outputs=[loading_message],
+                    outputs=loading_message,
                     queue=False
                 ).then(
                     fn=lambda: gr.update(value="", visible=False),  # Hide loading message
                     inputs=None,
-                    outputs=[loading_message],
+                    outputs=loading_message,
                     queue=False
                 ).then(
                     fn=lambda: gr.update(visible=True),
@@ -410,7 +413,7 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                 ).then(
                     fn=lambda: gr.update(value="Loading Model...", visible=True),  # Show loading message
                     inputs=None,
-                    outputs=[loading_message, use_as_input_button]
+                    outputs=loading_message
                 ).then(
                     fn=fill_image,
                     inputs=[prompt, input_image, model_selection, paste_back],
@@ -418,12 +421,12 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                 ).then(
                     fn=lambda: gr.update(value="Model Loaded", visible=True),  # Show loaded message
                     inputs=None,
-                    outputs=[loading_message],
+                    outputs=loading_message,
                     queue=False
                 ).then(
                     fn=lambda: gr.update(value="", visible=False),  # Hide loading message
                     inputs=None,
-                    outputs=[loading_message],
+                    outputs=loading_message,
                     queue=False
                 ).then(
                     fn=lambda: gr.update(visible=True),
@@ -477,7 +480,7 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                                     overlap_percentage = gr.Slider(
                                         label="Mask overlap (%)",
                                         minimum=1,
-                                        maximum=50,
+                                        maximum=100,
                                         value=10,
                                         step=1
                                     )
@@ -517,10 +520,12 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                             interactive=False,
                             label="Generated Image",
                         )
-                        use_as_input_button_outpaint = gr.Button("Use as Input Image", visible=False)
-                        history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
-                        preview_image = gr.Image(label="Preview")
                         loading_message_outpaint = gr.Label(label="Status", value="", visible=False)  # Added loading message label
+                        use_as_input_button_outpaint = gr.Button("Use as Input Image", visible=False)
+                        with gr.Accordion():
+                            history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
+                            preview_image = gr.Image(label="Preview")
+                        
 
 
     target_ratio.change(
@@ -556,6 +561,7 @@ with gr.Blocks(css=css, fill_height=True) as demo:
         inputs=[result_outpaint],
         outputs=[input_image_outpaint]
     )
+    
     runout_button.click(
         fn=clear_result,
         inputs=None,
@@ -563,7 +569,7 @@ with gr.Blocks(css=css, fill_height=True) as demo:
     ).then(
         fn=lambda: gr.update(value="Loading Model...", visible=True),  # Show loading message
         inputs=None,
-        outputs=[loading_message_outpaint, use_as_input_button_outpaint]
+        outputs=loading_message_outpaint
     ).then(
         fn=infer,
         inputs=[input_image_outpaint, width_slider, height_slider, overlap_percentage, num_inference_steps,
@@ -573,7 +579,7 @@ with gr.Blocks(css=css, fill_height=True) as demo:
     ).then(
         fn=lambda: gr.update(value="Model Loaded", visible=True),  # Show loaded message
         inputs=None,
-        outputs=[loading_message_outpaint],
+        outputs=loading_message_outpaint,
         queue=False
     ).then(
         fn=lambda x, history: update_history(x[1], history),
@@ -582,21 +588,25 @@ with gr.Blocks(css=css, fill_height=True) as demo:
     ).then(
         fn=lambda: gr.update(value="", visible=False),  # Hide loading message
         inputs=None,
-        outputs=[loading_message_outpaint],
+        outputs=loading_message_outpaint,
         queue=False
     ).then(
         fn=lambda: gr.update(visible=True),
         inputs=None,
-        outputs=[use_as_input_button_outpaint],
+        outputs=use_as_input_button_outpaint,
     )
     prompt_input.submit(
         fn=clear_result,
         inputs=None,
         outputs=result_outpaint,
+    ).then(
+        fn=lambda: gr.update(visible=False),
+        inputs=None,
+        outputs=use_as_input_button_outpaint,
     ).then(
         fn=lambda: gr.update(value="Loading Model...", visible=True),  # Show loading message
         inputs=None,
-        outputs=[loading_message_outpaint, use_as_input_button_outpaint]
+        outputs=loading_message_outpaint
     ).then(
         fn=infer,
         inputs=[input_image_outpaint, width_slider, height_slider, overlap_percentage, num_inference_steps,
@@ -606,7 +616,7 @@ with gr.Blocks(css=css, fill_height=True) as demo:
     ).then(
         fn=lambda: gr.update(value="Model Loaded", visible=True),  # Show loaded message
         inputs=None,
-        outputs=[loading_message_outpaint],
+        outputs=loading_message_outpaint,
         queue=False
     ).then(
         fn=lambda x, history: update_history(x[1], history),
@@ -615,12 +625,12 @@ with gr.Blocks(css=css, fill_height=True) as demo:
     ).then(
         fn=lambda: gr.update(value="", visible=False),  # Hide loading message
         inputs=None,
-        outputs=[loading_message_outpaint],
+        outputs=loading_message_outpaint,
         queue=False
     ).then(
         fn=lambda: gr.update(visible=True),
         inputs=None,
-        outputs=[use_as_input_button_outpaint],
+        outputs=use_as_input_button_outpaint,
     )
     preview_button.click(
         fn=preview_image_and_mask,