Update app.py

app.py

@@ -11,67 +11,62 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
 
-# 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)
-controlnet, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
+model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
     controlnet_model, sstate_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
 )
-controlnet.to(device="cuda", dtype=torch.float16)
-
-# Define available models
-models = {
-    "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
-    "RealVisXL V4.0 Lightning": "SG161222/RealVisXL_V4.0_Lightning",
-}
+model.to(device="cuda", dtype=torch.float16)
 
-# Load default pipeline
-default_model = "RealVisXL V5.0 Lightning"
-pipe = StableDiffusionXLFillPipeline.from_pretrained(
-    models[default_model],
-    torch_dtype=torch.float16,
-    vae=vae,
-    controlnet=controlnet,
-    variant="fp16",
+vae = AutoencoderKL.from_pretrained(
+    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
 ).to("cuda")
-pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
-# 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,
+# Initialize both pipelines and store them in a dictionary
+pipelines = {
+    "RealVisXL V5.0 Lightning": StableDiffusionXLFillPipeline.from_pretrained(
+        "SG161222/RealVisXL_V5.0_Lightning",
         torch_dtype=torch.float16,
         vae=vae,
-        controlnet=controlnet,
+        controlnet=model,
         variant="fp16",
-    ).to("cuda")
-    new_pipe.scheduler = TCDScheduler.from_config(new_pipe.scheduler.config)
-    return new_pipe
+    ).to("cuda"),
+    "RealVisXL V4.0 Lightning": StableDiffusionXLFillPipeline.from_pretrained(
+        "SG161222/RealVisXL_V4.0_Lightning",
+        torch_dtype=torch.float16,
+        vae=vae,
+        controlnet=model,
+        variant="fp16",
+    ).to("cuda"),
+}
+
+for pipe in pipelines.values():
+    pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
+
 
-# 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%":
@@ -83,21 +78,27 @@ 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
@@ -114,22 +115,26 @@ 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
     right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width - white_gaps_patch
     top_overlap = margin_y + overlap_y if overlap_top else margin_y + white_gaps_patch
     bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height - white_gaps_patch
-    
+
     if alignment == "Left":
         left_overlap = margin_x + overlap_x if overlap_left else margin_x
     elif alignment == "Right":
@@ -139,6 +144,7 @@ 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)
@@ -146,24 +152,27 @@ 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(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):
+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, selected_model):
     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()
     cnet_image.paste(0, (0, 0), mask)
 
     final_prompt = f"{prompt_input} , high quality, 4k"
 
+    # Access the selected pipeline from the dictionary
+    pipe = pipelines[selected_model]
+
     (
         prompt_embeds,
         negative_prompt_embeds,
         pooled_prompt_embeds,
         negative_pooled_prompt_embeds,
-    ) = pipeline.encode_prompt(final_prompt, "cuda", True)
+    ) = pipe.encode_prompt(final_prompt, "cuda", True)
 
-    for image in pipeline(
+    # Generate the image
+    for image in pipe(
         prompt_embeds=prompt_embeds,
         negative_prompt_embeds=negative_prompt_embeds,
         pooled_prompt_embeds=pooled_prompt_embeds,
@@ -171,28 +180,38 @@ def infer(pipeline, image, width, height, overlap_percentage, num_inference_step
         image=cnet_image,
         num_inference_steps=num_inference_steps
     ):
-        pass
-    generated_image = image
+        pass  # Wait for the generation to complete
+    generated_image = image  # Get the last 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":
-        return 720, 1280, gr.update()
+        changed_width = 720
+        changed_height = 1280
+        return changed_width, changed_height, gr.update()
     elif target_ratio == "16:9":
-        return 1280, 720, gr.update()
+        changed_width = 1280
+        changed_height = 720
+        return changed_width, changed_height, gr.update()
     elif target_ratio == "1:1":
-        return 1024, 1024, gr.update()
+        changed_width = 1024
+        changed_height = 1024
+        return changed_width, changed_height, gr.update()
     elif target_ratio == "Custom":
         return ui_width, ui_height, gr.update(open=True)
 
+
 def select_the_right_preset(user_width, user_height):
     if user_width == 720 and user_height == 1280:
         return "9:16"
@@ -203,16 +222,20 @@ def select_the_right_preset(user_width, user_height):
     else:
         return "Custom"
 
+
 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 (unchanged)
+
+# --- CSS and Title (unchanged) ---
 css = """
 h1 {
   text-align: center;
@@ -223,7 +246,6 @@ 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)
@@ -241,13 +263,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     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",
@@ -255,11 +270,18 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                         value="9:16",
                         scale=2
                     )
+
                     alignment_dropdown = gr.Dropdown(
                         choices=["Middle", "Left", "Right", "Top", "Bottom"],
                         value="Middle",
                         label="Alignment"
                     )
+                with gr.Row():
+                    model_selector = gr.Dropdown(
+                        label="Select Model",
+                        choices=list(pipelines.keys()),
+                        value="RealVisXL V5.0 Lightning",
+                    )
 
                 with gr.Accordion(label="Advanced settings", open=False) as settings_panel:
                     with gr.Column():
@@ -278,7 +300,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(
@@ -308,7 +330,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                                 value=50,
                                 visible=False
                             )
-                            
+
                 gr.Examples(
                     examples=[
                         ["./examples/example_1.webp", 1280, 720, "Middle"],
@@ -327,74 +349,64 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 )
                 history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
 
-        # 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,
-        )
+    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, model_selector],
+        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, model_selector],
+        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)