Update app.py

app.py

@@ -11,65 +11,67 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
 
-# Load configuration and models
+# Load VAE and ControlNet (shared components)
+vae = AutoencoderKL.from_pretrained(
+    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
+).to("cuda")
+
 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",
 )
-
 sstate_dict = load_state_dict(model_file)
-model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
+controlnet, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
     controlnet_model, sstate_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
 )
-model.to(device="cuda", dtype=torch.float16)
+controlnet.to(device="cuda", dtype=torch.float16)
 
-vae = AutoencoderKL.from_pretrained(
-    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
-).to("cuda")
+# Define available models
+models = {
+    "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
+    "RealVisXL V4.0 Lightning": "SG161222/RealVisXL_V4.0_Lightning",
+}
 
-# Initially load the default pipeline
+# Load default pipeline
+default_model = "RealVisXL V5.0 Lightning"
 pipe = StableDiffusionXLFillPipeline.from_pretrained(
-    "SG161222/RealVisXL_V5.0_Lightning",
+    models[default_model],
     torch_dtype=torch.float16,
     vae=vae,
-    controlnet=model,
+    controlnet=controlnet,
     variant="fp16",
 ).to("cuda")
-
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
-def load_model(selected_model):
-    global pipe
-    model_path = f"SG161222/{selected_model}"
-    pipe = StableDiffusionXLFillPipeline.from_pretrained(
-        model_path,
+# Function to load pipeline based on selected model
+def load_pipeline(model_name):
+    repo_id = models[model_name]
+    new_pipe = StableDiffusionXLFillPipeline.from_pretrained(
+        repo_id,
         torch_dtype=torch.float16,
         vae=vae,
-        controlnet=model,
+        controlnet=controlnet,
         variant="fp16",
     ).to("cuda")
-    pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
-    return f"Loaded model: {selected_model}"
+    new_pipe.scheduler = TCDScheduler.from_config(new_pipe.scheduler.config)
+    return new_pipe
 
+# Prepare image and mask function (unchanged)
 def prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
     target_size = (width, height)
 
-    # Calculate the scaling factor to fit the image within the target size
     scale_factor = min(target_size[0] / image.width, target_size[1] / image.height)
     new_width = int(image.width * scale_factor)
     new_height = int(image.height * scale_factor)
     
-    # Resize the source image to fit within target size
     source = image.resize((new_width, new_height), Image.LANCZOS)
 
-    # Apply resize option using percentages
     if resize_option == "Full":
         resize_percentage = 100
     elif resize_option == "50%":
@@ -81,27 +83,21 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
     else:  # Custom
         resize_percentage = custom_resize_percentage
 
-    # Calculate new dimensions based on percentage
     resize_factor = resize_percentage / 100
     new_width = int(source.width * resize_factor)
     new_height = int(source.height * resize_factor)
 
-    # Ensure minimum size of 64 pixels
     new_width = max(new_width, 64)
     new_height = max(new_height, 64)
 
-    # Resize the image
     source = source.resize((new_width, new_height), Image.LANCZOS)
 
-    # Calculate the overlap in pixels based on the percentage
     overlap_x = int(new_width * (overlap_percentage / 100))
     overlap_y = int(new_height * (overlap_percentage / 100))
 
-    # Ensure minimum overlap of 1 pixel
     overlap_x = max(overlap_x, 1)
     overlap_y = max(overlap_y, 1)
 
-    # Calculate margins based on alignment
     if alignment == "Middle":
         margin_x = (target_size[0] - new_width) // 2
         margin_y = (target_size[1] - new_height) // 2
@@ -118,19 +114,15 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
         margin_x = (target_size[0] - new_width) // 2
         margin_y = target_size[1] - new_height
 
-    # Adjust margins to eliminate gaps
     margin_x = max(0, min(margin_x, target_size[0] - new_width))
     margin_y = max(0, min(margin_y, target_size[1] - new_height))
 
-    # Create a new background image and paste the resized source image
     background = Image.new('RGB', target_size, (255, 255, 255))
     background.paste(source, (margin_x, margin_y))
 
-    # Create the mask
     mask = Image.new('L', target_size, 255)
     mask_draw = ImageDraw.Draw(mask)
 
-    # Calculate overlap areas
     white_gaps_patch = 2
 
     left_overlap = margin_x + overlap_x if overlap_left else margin_x + white_gaps_patch
@@ -147,7 +139,6 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
     elif alignment == "Bottom":
         bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height
 
-    # Draw the mask
     mask_draw.rectangle([
         (left_overlap, top_overlap),
         (right_overlap, bottom_overlap)
@@ -155,8 +146,9 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
 
     return background, mask
 
+# Updated inference function to use selected pipeline
 @spaces.GPU(duration=24)
-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):
+def infer(pipeline, 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)
     
     cnet_image = background.copy()
@@ -169,10 +161,9 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
         negative_prompt_embeds,
         pooled_prompt_embeds,
         negative_pooled_prompt_embeds,
-    ) = pipe.encode_prompt(final_prompt, "cuda", True)
+    ) = pipeline.encode_prompt(final_prompt, "cuda", True)
 
-    # Generate the image
-    for image in pipe(
+    for image in pipeline(
         prompt_embeds=prompt_embeds,
         negative_prompt_embeds=negative_prompt_embeds,
         pooled_prompt_embeds=pooled_prompt_embeds,
@@ -180,32 +171,25 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
         image=cnet_image,
         num_inference_steps=num_inference_steps
     ):
-        pass  # Wait for the generation to complete
-    generated_image = image  # Get the last image
+        pass
+    generated_image = image
 
     generated_image = generated_image.convert("RGBA")
     cnet_image.paste(generated_image, (0, 0), mask)
 
     return cnet_image
 
+# Utility functions (unchanged)
 def clear_result():
-    """Clears the result Image."""
     return gr.update(value=None)
 
 def preload_presets(target_ratio, ui_width, ui_height):
-    """Updates the width and height sliders based on the selected aspect ratio."""
     if target_ratio == "9:16":
-        changed_width = 720
-        changed_height = 1280
-        return changed_width, changed_height, gr.update()
+        return 720, 1280, gr.update()
     elif target_ratio == "16:9":
-        changed_width = 1280
-        changed_height = 720
-        return changed_width, changed_height, gr.update()
+        return 1280, 720, gr.update()
     elif target_ratio == "1:1":
-        changed_width = 1024
-        changed_height = 1024
-        return changed_width, changed_height, gr.update()
+        return 1024, 1024, gr.update()
     elif target_ratio == "Custom":
         return ui_width, ui_height, gr.update(open=True)
 
@@ -223,13 +207,12 @@ def toggle_custom_resize_slider(resize_option):
     return gr.update(visible=(resize_option == "Custom"))
 
 def update_history(new_image, history):
-    """Updates the history gallery with the new image."""
     if history is None:
         history = []
     history.insert(0, new_image)
     return history
 
-# CSS and Title
+# CSS and title (unchanged)
 css = """
 h1 {
   text-align: center;
@@ -240,6 +223,7 @@ h1 {
 title = """<h1 align="center">Diffusers Image Outpaint Lightning</h1>
 """
 
+# Gradio interface with model selection
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     with gr.Column():
         gr.HTML(title)
@@ -250,17 +234,20 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     type="pil",
                     label="Input Image"
                 )
-                model_selection = gr.Dropdown(
-                    choices=["RealVisXL_V5.0_Lightning", "RealVisXL_V4.0_Lightning"],
-                    value="RealVisXL_V5.0_Lightning",
-                    label="Select Model"
-                )
+
                 with gr.Row():
                     with gr.Column(scale=2):
                         prompt_input = gr.Textbox(label="Prompt (Optional)")
                     with gr.Column(scale=1):
                         run_button = gr.Button("Generate")
 
+                with gr.Row():
+                    model_selector = gr.Dropdown(
+                        label="Select Model",
+                        choices=list(models.keys()),
+                        value="RealVisXL V5.0 Lightning",
+                    )
+
                 with gr.Row():
                     target_ratio = gr.Radio(
                         label="Expected Ratio",
@@ -291,6 +278,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                                 step=8,
                                 value=1280,
                             )
+                        
                         num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=12, step=1, value=8)
                         with gr.Group():
                             overlap_percentage = gr.Slider(
@@ -320,7 +308,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                                 value=50,
                                 visible=False
                             )
-                status_text = gr.Textbox(label="Status", interactive=False)
+                            
                 gr.Examples(
                     examples=[
                         ["./examples/example_1.webp", 1280, 720, "Middle"],
@@ -339,62 +327,74 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 )
                 history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
 
-    # Event handlers
-    model_selection.change(fn=load_model, inputs=model_selection, outputs=status_text)
-    target_ratio.change(
-        fn=preload_presets,
-        inputs=[target_ratio, width_slider, height_slider],
-        outputs=[width_slider, height_slider, settings_panel],
-        queue=False
-    )
-    width_slider.change(
-        fn=select_the_right_preset,
-        inputs=[width_slider, height_slider],
-        outputs=[target_ratio],
-        queue=False
-    )
-    height_slider.change(
-        fn=select_the_right_preset,
-        inputs=[width_slider, height_slider],
-        outputs=[target_ratio],
-        queue=False
-    )
-    resize_option.change(
-        fn=toggle_custom_resize_slider,
-        inputs=[resize_option],
-        outputs=[custom_resize_percentage],
-        queue=False
-    )
-    run_button.click(
-        fn=clear_result,
-        inputs=None,
-        outputs=result,
-    ).then(
-        fn=infer,
-        inputs=[input_image, 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,
-    ).then(
-        fn=lambda x, history: update_history(x, history),
-        inputs=[result, history_gallery],
-        outputs=history_gallery,
-    )
-    prompt_input.submit(
-        fn=clear_result,
-        inputs=None,
-        outputs=result,
-    ).then(
-        fn=infer,
-        inputs=[input_image, 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,
-    ).then(
-        fn=lambda x, history: update_history(x, history),
-        inputs=[result, history_gallery],
-        outputs=history_gallery,
-    )
-    demo.load(fn=load_model, inputs=model_selection, outputs=status_text)
+        # State to hold the current pipeline
+        pipeline_state = gr.State(value=pipe)
+
+        # Update pipeline when model is selected
+        model_selector.change(
+            fn=load_pipeline,
+            inputs=model_selector,
+            outputs=pipeline_state,
+        )
+
+        target_ratio.change(
+            fn=preload_presets,
+            inputs=[target_ratio, width_slider, height_slider],
+            outputs=[width_slider, height_slider, settings_panel],
+            queue=False
+        )
+
+        width_slider.change(
+            fn=select_the_right_preset,
+            inputs=[width_slider, height_slider],
+            outputs=[target_ratio],
+            queue=False
+        )
+
+        height_slider.change(
+            fn=select_the_right_preset,
+            inputs=[width_slider, height_slider],
+            outputs=[target_ratio],
+            queue=False
+        )
+
+        resize_option.change(
+            fn=toggle_custom_resize_slider,
+            inputs=[resize_option],
+            outputs=[custom_resize_percentage],
+            queue=False
+        )
+        
+        run_button.click(
+            fn=clear_result,
+            inputs=None,
+            outputs=result,
+        ).then(
+            fn=infer,
+            inputs=[pipeline_state, input_image, 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,
+        ).then(
+            fn=lambda x, history: update_history(x, history),
+            inputs=[result, history_gallery],
+            outputs=history_gallery,
+        )
+
+        prompt_input.submit(
+            fn=clear_result,
+            inputs=None,
+            outputs=result,
+        ).then(
+            fn=infer,
+            inputs=[pipeline_state, input_image, 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,
+        ).then(
+            fn=lambda x, history: update_history(x, history),
+            inputs=[result, history_gallery],
+            outputs=history_gallery,
+        )
 
 demo.queue(max_size=20).launch(share=False, ssr_mode=False, show_error=True)