Update app.py

app.py

@@ -18,57 +18,47 @@ MODELS = {
     "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)
-loaded_model_name = default_model_name  # Track the loaded model
-
-def update_pipeline(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
-    
+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")
+
 @spaces.GPU(duration=12)
 def fill_image(prompt, image, model_selection, paste_back):
-    global pipe
-    update_pipeline(model_selection)
     print(f"Received image: {image}")
     if image is None:
         yield None, None
         return
+
     (
         prompt_embeds,
         negative_prompt_embeds,
@@ -81,6 +71,7 @@ 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,
@@ -89,6 +80,7 @@ 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:
@@ -197,9 +189,15 @@ def preview_image_and_mask(image, width, height, overlap_percentage, resize_opti
     return preview
 
 @spaces.GPU(duration=12)
-def inpaint(prompt, image, model_selection, paste_back):
+def inpaint(prompt, image, inpaint_model, paste_back):
     global pipe
-    update_pipeline(model_selection)
+    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")
     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}"
@@ -207,7 +205,7 @@ def inpaint(prompt, image, model_selection, 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,8 +235,6 @@ 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)  # 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"
@@ -263,7 +259,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])
 
@@ -311,16 +307,13 @@ def update_history(new_image, history):
     return history
 
 css = """
-.nukgradio-container {
+.nulgradio-container {
     width: 86vw !important;
 }
-.contain {
+.nulcontain {
     overflow-y: scroll !important;
     padding: 10px 40px !important;
 }
-div#component-8 {
-    min-height: fit-content !important;
-}
 div#component-17 {
     height: auto !important;
 }
@@ -338,41 +331,35 @@ title = """<h1 align="center">Diffusers Image Outpaint</h1>
 """
 
 with gr.Blocks(css=css, fill_height=True) as demo:
-    gr.Markdown("# SDXL Lightning Diffusers Inpaint and Outpaint")
+    gr.Markdown("# Diffusers Inpaint and Outpaint")
     with gr.Tabs():
         with gr.TabItem("Inpaint"):
             with gr.Column():
                 with gr.Row():
-                    with gr.Column():
-                        input_image = gr.ImageMask(
-                            type="pil", label="Input Image", layers=True
-                        )
-
                     with gr.Column():
                         prompt = gr.Textbox(
                             label="Prompt",
                             info="Describe what to inpaint the mask with",
                             lines=3,
                         )
+                    with gr.Column():
                         model_selection = gr.Dropdown(
                             choices=list(MODELS.keys()),
-                            value=default_model_name,
+                            value="RealVisXL V5.0 Lightning",
                             label="Model",
                         )
-                        paste_back = gr.Checkbox(True, label="Paste back original")
-                        loading_message = gr.Label(label="Status", value="", visible=False)  # Added loading message label
-
-                        result = ImageSlider(
-                            interactive=False,
-                            label="Generated Image",
-                        )
-
-
-                run_button = gr.Button("Generate")
-                
-
+                        with gr.Row():
+                            run_button = gr.Button("Generate")
+                            paste_back = gr.Checkbox(True, label="Paste back original")
+                with gr.Row(equal_height=False):
+                    input_image = gr.ImageMask(
+                        type="pil", label="Input Image", layers=True
+                    )
+                    result = ImageSlider(
+                        interactive=False,
+                        label="Generated Image",
+                    )
                 use_as_input_button = gr.Button("Use as Input Image", visible=False)
-
                 use_as_input_button.click(
                     fn=use_output_as_input, inputs=[result], outputs=[input_image]
                 )
@@ -384,24 +371,10 @@ 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
                 ).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,
@@ -415,24 +388,10 @@ 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
                 ).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,
@@ -485,7 +444,7 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                                     overlap_percentage = gr.Slider(
                                         label="Mask overlap (%)",
                                         minimum=1,
-                                        maximum=100,
+                                        maximum=50,
                                         value=10,
                                         step=1
                                     )
@@ -525,13 +484,9 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                             interactive=False,
                             label="Generated Image",
                         )
-                        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")
-                        
-
+                        history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
+                        preview_image = gr.Image(label="Preview")
 
     target_ratio.change(
         fn=preload_presets,
@@ -566,76 +521,43 @@ 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,
         outputs=result_outpaint,
-    ).then(
-        fn=lambda: gr.update(value="Loading Model...", visible=True),  # Show loading message
-        inputs=None,
-        outputs=loading_message_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,
-        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
     ).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,
-        outputs=use_as_input_button_outpaint,
+        outputs=[use_as_input_button_outpaint],
     )
     preview_button.click(
         fn=preview_image_and_mask,
@@ -644,4 +566,5 @@ with gr.Blocks(css=css, fill_height=True) as demo:
         outputs=[preview_image],
         queue=False
     )
+
 demo.launch(show_error=True)