Update app.py (#5)

- Update app.py (516e56ce390baf01fe4d33d4b4999c78fb6a1bd1)

app.py

@@ -17,44 +17,51 @@ 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",
+        filename="config_promax.json",
+    )
+    config = ControlNetModel_Union.load_config(config_file)
+    controlnet_model = ControlNetModel_Union.from_config(config)
+    model_file = hf_hub_download(
+        "xinsir/controlnet-union-sdxl-1.0",
+        filename="diffusion_pytorch_model_promax.safetensors",
+    )
+    state_dict = load_state_dict(model_file)
+    model, _,_, _,_ = ControlNetModel_Union._load_pretrained_model(
+        controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
+    )
+    model.to(device="cuda", dtype=torch.float16)
+    vae = AutoencoderKL.from_pretrained(
+        "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
+    ).to("cuda")
+    pipe = StableDiffusionXLFillPipeline.from_pretrained(
+        MODELS[model_name],
+        torch_dtype=torch.float16,
+        vae=vae,
+        controlnet=model,
+        variant="fp16",
+    )
+    pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
+    pipe.to("cuda")
+    return pipe
+
+# Initialize with the default model
+default_model_name = "RealVisXL V5.0 Lightning"
+pipe = init_pipeline(default_model_name)
 
-config_file = hf_hub_download(
-    "xinsir/controlnet-union-sdxl-1.0",
-    filename="config_promax.json",
-)
-config = ControlNetModel_Union.load_config(config_file)
-controlnet_model = ControlNetModel_Union.from_config(config)
-model_file = hf_hub_download(
-    "xinsir/controlnet-union-sdxl-1.0",
-    filename="diffusion_pytorch_model_promax.safetensors",
-)
-state_dict = load_state_dict(model_file)
-model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
-    controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
-)
-model.to(device="cuda", dtype=torch.float16)
-vae = AutoencoderKL.from_pretrained(
-    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
-).to("cuda")
-pipe = StableDiffusionXLFillPipeline.from_pretrained(
-    "SG161222/RealVisXL_V5.0_Lightning",
-    torch_dtype=torch.float16,
-    vae=vae,
-    controlnet=model,
-    variant="fp16",
-)
-pipe = StableDiffusionXLFillPipeline.from_pretrained(
-    "GraydientPlatformAPI/lustify-lightning",
-    torch_dtype=torch.float16,
-    vae=vae,
-    controlnet=model,
-)
-pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
-pipe.to("cuda")
+def update_pipeline(model_selection):
+    global pipe
+    if pipe.config.model_name != MODELS[model_selection]:
+        pipe = init_pipeline(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)
     if image is None:
         yield None, None
         return
@@ -191,13 +198,7 @@ def preview_image_and_mask(image, width, height, overlap_percentage, resize_opti
 @spaces.GPU(duration=12)
 def inpaint(prompt, image, inpaint_model, paste_back):
     global pipe
-    if pipe.config.model_name != MODELS[model_name]:
-        pipe = StableDiffusionXLFillPipeline.from_pretrained(
-            MODELS[model_name],
-            torch_dtype=torch.float16,
-            vae=vae,
-            controlnet=model,
-        ).to("cuda")
+    update_pipeline(inpaint_model)
     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}"
@@ -235,6 +236,8 @@ 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
     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"
@@ -259,7 +262,7 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
     image = image.convert("RGBA")
     cnet_image.paste(image, (0, 0), mask)
     yield background, cnet_image
-
+    
 def use_output_as_input(output_image):
     return gr.update(value=output_image[1])
 
@@ -360,6 +363,8 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                         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]
                 )
@@ -371,10 +376,24 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                     fn=lambda: gr.update(visible=False),
                     inputs=None,
                     outputs=use_as_input_button,
+                ).then(
+                    fn=lambda: gr.update(value="Loading Model...", visible=True),  # Show loading message
+                    inputs=None,
+                    outputs=[loading_message, use_as_input_button]
                 ).then(
                     fn=fill_image,
                     inputs=[prompt, input_image, model_selection, paste_back],
                     outputs=[result],
+                ).then(
+                    fn=lambda: gr.update(value="Model Loaded", visible=True),  # Show loaded message
+                    inputs=None,
+                    outputs=[loading_message],
+                    queue=False
+                ).then(
+                    fn=lambda: gr.update(value="", visible=False),  # Hide loading message
+                    inputs=None,
+                    outputs=[loading_message],
+                    queue=False
                 ).then(
                     fn=lambda: gr.update(visible=True),
                     inputs=None,
@@ -388,10 +407,24 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                     fn=lambda: gr.update(visible=False),
                     inputs=None,
                     outputs=use_as_input_button,
+                ).then(
+                    fn=lambda: gr.update(value="Loading Model...", visible=True),  # Show loading message
+                    inputs=None,
+                    outputs=[loading_message, use_as_input_button]
                 ).then(
                     fn=fill_image,
                     inputs=[prompt, input_image, model_selection, paste_back],
                     outputs=[result],
+                ).then(
+                    fn=lambda: gr.update(value="Model Loaded", visible=True),  # Show loaded message
+                    inputs=None,
+                    outputs=[loading_message],
+                    queue=False
+                ).then(
+                    fn=lambda: gr.update(value="", visible=False),  # Hide loading message
+                    inputs=None,
+                    outputs=[loading_message],
+                    queue=False
                 ).then(
                     fn=lambda: gr.update(visible=True),
                     inputs=None,
@@ -487,6 +520,8 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                         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
+
 
     target_ratio.change(
         fn=preload_presets,
@@ -525,16 +560,30 @@ with gr.Blocks(css=css, fill_height=True) as demo:
         fn=clear_result,
         inputs=None,
         outputs=result_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]
     ).then(
         fn=infer,
         inputs=[input_image_outpaint, width_slider, height_slider, overlap_percentage, num_inference_steps,
                 resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,
                 overlap_left, overlap_right, overlap_top, overlap_bottom],
         outputs=[result_outpaint],
+    ).then(
+        fn=lambda: gr.update(value="Model Loaded", visible=True),  # Show loaded message
+        inputs=None,
+        outputs=[loading_message_outpaint],
+        queue=False
     ).then(
         fn=lambda x, history: update_history(x[1], history),
         inputs=[result_outpaint, history_gallery],
         outputs=history_gallery,
+    ).then(
+        fn=lambda: gr.update(value="", visible=False),  # Hide loading message
+        inputs=None,
+        outputs=[loading_message_outpaint],
+        queue=False
     ).then(
         fn=lambda: gr.update(visible=True),
         inputs=None,
@@ -544,16 +593,30 @@ with gr.Blocks(css=css, fill_height=True) as demo:
         fn=clear_result,
         inputs=None,
         outputs=result_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]
     ).then(
         fn=infer,
         inputs=[input_image_outpaint, width_slider, height_slider, overlap_percentage, num_inference_steps,
                 resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,
                 overlap_left, overlap_right, overlap_top, overlap_bottom],
         outputs=[result_outpaint],
+    ).then(
+        fn=lambda: gr.update(value="Model Loaded", visible=True),  # Show loaded message
+        inputs=None,
+        outputs=[loading_message_outpaint],
+        queue=False
     ).then(
         fn=lambda x, history: update_history(x[1], history),
         inputs=[result_outpaint, history_gallery],
         outputs=history_gallery,
+    ).then(
+        fn=lambda: gr.update(value="", visible=False),  # Hide loading message
+        inputs=None,
+        outputs=[loading_message_outpaint],
+        queue=False
     ).then(
         fn=lambda: gr.update(visible=True),
         inputs=None,
@@ -566,5 +629,4 @@ with gr.Blocks(css=css, fill_height=True) as demo:
         outputs=[preview_image],
         queue=False
     )
-
 demo.launch(show_error=True)