Update app.py

app.py

@@ -11,56 +11,52 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
 
-# --- Configuration and Model Loading ---
-
-# Load ControlNet Union
+# Load configuration and models
 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_model, sstate_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
 )
 model.to(device="cuda", dtype=torch.float16)
 
-# Load VAE
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
 ).to("cuda")
 
-# --- Load Multiple Pipelines ---
-pipelines = {}
-
-# Load RealVisXL V5.0 Lightning
-pipe_v5 = StableDiffusionXLFillPipeline.from_pretrained(
+# Initially load the default pipeline
+pipe = StableDiffusionXLFillPipeline.from_pretrained(
     "SG161222/RealVisXL_V5.0_Lightning",
     torch_dtype=torch.float16,
     vae=vae,
-    controlnet=model, # Use the same controlnet
-    variant="fp16",
-).to("cuda")
-pipe_v5.scheduler = TCDScheduler.from_config(pipe_v5.scheduler.config)
-pipelines["RealVisXL V5.0 Lightning"] = pipe_v5
-
-# Load RealVisXL V4.0 Lightning
-pipe_v4 = StableDiffusionXLFillPipeline.from_pretrained(
-    "SG161222/RealVisXL_V4.0_Lightning",
-    torch_dtype=torch.float16,
-    vae=vae, # Use the same VAE
-    controlnet=model, # Use the same controlnet
+    controlnet=model,
     variant="fp16",
 ).to("cuda")
-pipe_v4.scheduler = TCDScheduler.from_config(pipe_v4.scheduler.config)
-pipelines["RealVisXL V4.0 Lightning"] = pipe_v4
 
-# --- Helper Functions ---
+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,
+        torch_dtype=torch.float16,
+        vae=vae,
+        controlnet=model,
+        variant="fp16",
+    ).to("cuda")
+    pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
+    return f"Loaded model: {selected_model}"
 
 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)
@@ -69,7 +65,7 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
     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)
 
@@ -121,10 +117,6 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
     elif alignment == "Bottom":
         margin_x = (target_size[0] - new_width) // 2
         margin_y = target_size[1] - new_height
-    else: # Default to Middle if alignment is somehow invalid
-        margin_x = (target_size[0] - new_width) // 2
-        margin_y = (target_size[1] - new_height) // 2
-
 
     # Adjust margins to eliminate gaps
     margin_x = max(0, min(margin_x, target_size[0] - new_width))
@@ -135,126 +127,66 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
     background.paste(source, (margin_x, margin_y))
 
     # Create the mask
-    mask = Image.new('L', target_size, 255) # White background (area to be filled)
+    mask = Image.new('L', target_size, 255)
     mask_draw = ImageDraw.Draw(mask)
 
-    # Calculate overlap areas (where the mask should be black = keep original)
-    white_gaps_patch = 2 # Small value to ensure no tiny gaps at edges if overlap is off
+    # Calculate overlap areas
+    white_gaps_patch = 2
 
-    # Determine the coordinates for the black rectangle (the non-masked area)
-    # Start with the full area covered by the pasted image
-    left_black = margin_x
-    top_black = margin_y
-    right_black = margin_x + new_width
-    bottom_black = margin_y + new_height
-
-    # Adjust the black area based on overlap checkboxes
-    if overlap_left:
-        left_black += overlap_x
-    else:
-        # If not overlapping left, ensure the black mask starts exactly at the image edge or slightly inside
-        left_black += white_gaps_patch if alignment != "Left" else 0
-
-    if overlap_right:
-        right_black -= overlap_x
-    else:
-         # If not overlapping right, ensure the black mask ends exactly at the image edge or slightly inside
-        right_black -= white_gaps_patch if alignment != "Right" else 0
-
-    if overlap_top:
-        top_black += overlap_y
-    else:
-        # If not overlapping top, ensure the black mask starts exactly at the image edge or slightly inside
-        top_black += white_gaps_patch if alignment != "Top" else 0
+    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":
+        right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width
+    elif alignment == "Top":
+        top_overlap = margin_y + overlap_y if overlap_top else margin_y
+    elif alignment == "Bottom":
+        bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height
 
-    if overlap_bottom:
-        bottom_black -= overlap_y
-    else:
-        # If not overlapping bottom, ensure the black mask ends exactly at the image edge or slightly inside
-        bottom_black -= white_gaps_patch if alignment != "Bottom" else 0
-
-    # Ensure coordinates are valid (left < right, top < bottom)
-    left_black = min(left_black, target_size[0])
-    top_black = min(top_black, target_size[1])
-    right_black = max(left_black, right_black) # Ensure right >= left
-    bottom_black = max(top_black, bottom_black) # Ensure bottom >= top
-    right_black = min(right_black, target_size[0])
-    bottom_black = min(bottom_black, target_size[1])
-
-
-    # Draw the black rectangle onto the white mask
-    # The area *inside* this rectangle will be kept (mask value 0)
-    # The area *outside* this rectangle will be filled (mask value 255)
-    if right_black > left_black and bottom_black > top_black:
-         mask_draw.rectangle(
-             [(left_black, top_black), (right_black, bottom_black)],
-             fill=0 # Black means keep this area
-         )
+    # Draw the mask
+    mask_draw.rectangle([
+        (left_overlap, top_overlap),
+        (right_overlap, bottom_overlap)
+    ], fill=0)
 
     return background, mask
 
-
 @spaces.GPU(duration=24)
-def infer(selected_model_name, image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
-    if image is None:
-        raise gr.Error("Please upload an input image.")
-    try:
-        # Select the pipeline based on the dropdown choice
-        pipe = pipelines[selected_model_name]
-
-        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)
-
-        # Create the controlnet input image (original image pasted on white bg, with masked area blacked out)
-        cnet_image = background.copy()
-        # Create a black image of the same size as the mask
-        black_fill = Image.new('RGB', mask.size, (0, 0, 0))
-        # Paste the black fill using the mask (where mask is 255/white, paste black)
-        cnet_image.paste(black_fill, (0, 0), mask)
-
-
-        final_prompt = f"{prompt_input} , high quality, 4k" if prompt_input else "high quality, 4k"
-
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = pipe.encode_prompt(final_prompt, "cuda", True)
-
-        # Generate the image
-        generator = torch.Generator(device="cuda").manual_seed(np.random.randint(0, 2**32)) # Add random seed
-
-        # The pipeline expects the 'image' argument to be the background with the original content
-        # and the 'mask_image' argument to define the area to *inpaint* (white area in our mask)
-        result_image = pipe(
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-            image=background, # The background containing the original image
-            mask_image=mask, # The mask (white = fill, black = keep)
-            control_image=cnet_image, # ControlNet input image
-            num_inference_steps=num_inference_steps,
-            generator=generator, # Use generator for reproducibility if needed
-            output_type="pil" # Ensure PIL output
-        ).images[0]
-
-        # The pipeline directly returns the final composited image.
-        # No need for manual pasting like before.
-
-        return result_image
-    except Exception as e:
-        print(f"Error during inference: {e}")
-        import traceback
-        traceback.print_exc()
-        # Return the background image or raise a Gradio error for clarity
-        # raise gr.Error(f"Inference failed: {e}")
-        # Or return the prepared background/mask for debugging
-        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)
-        # Combine background and mask for visualization
-        debug_img = Image.blend(background.convert("RGBA"), mask.convert("RGBA"), 0.5)
-        return debug_img # Return a debug image or None
-
+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):
+    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"
+
+    (
+        prompt_embeds,
+        negative_prompt_embeds,
+        pooled_prompt_embeds,
+        negative_pooled_prompt_embeds,
+    ) = pipe.encode_prompt(final_prompt, "cuda", True)
+
+    # Generate the image
+    for image in pipe(
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+        image=cnet_image,
+        num_inference_steps=num_inference_steps
+    ):
+        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
 
 def clear_result():
     """Clears the result Image."""
@@ -265,21 +197,17 @@ def preload_presets(target_ratio, ui_width, ui_height):
     if target_ratio == "9:16":
         changed_width = 720
         changed_height = 1280
-        return changed_width, changed_height, gr.update(open=False) # Close accordion on preset
+        return changed_width, changed_height, gr.update()
     elif target_ratio == "16:9":
         changed_width = 1280
         changed_height = 720
-        return changed_width, changed_height, gr.update(open=False) # Close accordion on preset
+        return changed_width, changed_height, gr.update()
     elif target_ratio == "1:1":
         changed_width = 1024
         changed_height = 1024
-        return changed_width, changed_height, gr.update(open=False) # Close accordion on preset
+        return changed_width, changed_height, gr.update()
     elif target_ratio == "Custom":
-        # When switching to Custom, keep current slider values but open accordion
         return ui_width, ui_height, gr.update(open=True)
-    # Should not happen, but return current values if it does
-    return ui_width, ui_height, gr.update()
-
 
 def select_the_right_preset(user_width, user_height):
     if user_width == 720 and user_height == 1280:
@@ -296,71 +224,54 @@ def toggle_custom_resize_slider(resize_option):
 
 def update_history(new_image, history):
     """Updates the history gallery with the new image."""
-    if new_image is None: # Don't add None to history (e.g., on clear or error)
-        return history
     if history is None:
         history = []
-    # Prepend the new image (as PIL or path depending on Gallery config)
     history.insert(0, new_image)
-    # Limit history size if desired (e.g., keep last 12)
-    max_history = 12
-    if len(history) > max_history:
-        history = history[:max_history]
     return history
 
-# --- CSS and Title ---
+# CSS and Title
 css = """
 h1 {
-    text-align: center;
-    display: block;
-}
-.gradio-container {
-    max-width: 1280px !important;
-    margin: auto !important;
+  text-align: center;
+  display: block;
 }
 """
 
 title = """<h1 align="center">Diffusers Image Outpaint Lightning</h1>
-<p align="center">Expand images using ControlNet Union and Lightning models. Choose a base model below.</p>
 """
 
-# --- Gradio UI ---
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     with gr.Column():
         gr.HTML(title)
 
         with gr.Row():
-            with gr.Column(scale=2): # Input column
+            with gr.Column():
                 input_image = gr.Image(
                     type="pil",
                     label="Input Image"
                 )
-
-                # --- Model Selector ---
-                model_selector = gr.Dropdown(
-                    label="Select Model",
-                    choices=list(pipelines.keys()),
-                    value="RealVisXL V5.0 Lightning", # Default model
+                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 (Describe the desired output)", placeholder="e.g., beautiful landscape, photorealistic")
-                    with gr.Column(scale=1, min_width=120):
-                        run_button = gr.Button("Generate", variant="primary")
+                        prompt_input = gr.Textbox(label="Prompt (Optional)")
+                    with gr.Column(scale=1):
+                        run_button = gr.Button("Generate")
 
                 with gr.Row():
                     target_ratio = gr.Radio(
-                        label="Target Ratio",
+                        label="Expected Ratio",
                         choices=["9:16", "16:9", "1:1", "Custom"],
-                        value="9:16", # Default ratio
+                        value="9:16",
                         scale=2
                     )
-
                     alignment_dropdown = gr.Dropdown(
                         choices=["Middle", "Left", "Right", "Top", "Bottom"],
                         value="Middle",
-                        label="Align Input Image"
+                        label="Alignment"
                     )
 
                 with gr.Accordion(label="Advanced settings", open=False) as settings_panel:
@@ -368,43 +279,38 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                         with gr.Row():
                             width_slider = gr.Slider(
                                 label="Target Width",
-                                minimum=512, # Lowered minimum slightly
+                                minimum=720,
                                 maximum=1536,
-                                step=64, # Steps of 64 common for SDXL
-                                value=720, # Default width
+                                step=8,
+                                value=720,
                             )
                             height_slider = gr.Slider(
                                 label="Target Height",
-                                minimum=512, # Lowered minimum slightly
+                                minimum=720,
                                 maximum=1536,
-                                step=64, # Steps of 64
-                                value=1280, # Default height
+                                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(
                                 label="Mask overlap (%)",
-                                info="Percentage of the input image edge to keep (reduces seams)",
                                 minimum=1,
                                 maximum=50,
-                                value=10, # Default overlap
+                                value=10,
                                 step=1
                             )
-                            gr.Markdown("Select edges to apply overlap:")
                             with gr.Row():
-                                overlap_top = gr.Checkbox(label="Top", value=True)
-                                overlap_right = gr.Checkbox(label="Right", value=True)
-                                overlap_left = gr.Checkbox(label="Left", value=True)
-                                overlap_bottom = gr.Checkbox(label="Bottom", value=True)
-
+                                overlap_top = gr.Checkbox(label="Overlap Top", value=True)
+                                overlap_right = gr.Checkbox(label="Overlap Right", value=True)
+                            with gr.Row():
+                                overlap_left = gr.Checkbox(label="Overlap Left", value=True)
+                                overlap_bottom = gr.Checkbox(label="Overlap Bottom", value=True)
                         with gr.Row():
                             resize_option = gr.Radio(
-                                label="Resize input image before placing",
-                                info="Scale the input image relative to its fitted size",
+                                label="Resize input image",
                                 choices=["Full", "50%", "33%", "25%", "Custom"],
-                                value="Full" # Default resize option
+                                value="Full"
                             )
                             custom_resize_percentage = gr.Slider(
                                 label="Custom resize (%)",
@@ -412,48 +318,35 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                                 maximum=100,
                                 step=1,
                                 value=50,
-                                visible=False # Initially hidden
+                                visible=False
                             )
-
+                status_text = gr.Textbox(label="Status", interactive=False)
                 gr.Examples(
                     examples=[
-                        ["./examples/example_1.webp", "RealVisXL V5.0 Lightning", 1280, 720, "Middle"],
-                        ["./examples/example_2.jpg", "RealVisXL V4.0 Lightning", 1440, 810, "Left"],
-                        ["./examples/example_3.jpg", "RealVisXL V5.0 Lightning", 1024, 1024, "Top"],
-                        ["./examples/example_3.jpg", "RealVisXL V5.0 Lightning", 1024, 1024, "Bottom"],
+                        ["./examples/example_1.webp", 1280, 720, "Middle"],
+                        ["./examples/example_2.jpg", 1440, 810, "Left"],
+                        ["./examples/example_3.jpg", 1024, 1024, "Top"],
+                        ["./examples/example_3.jpg", 1024, 1024, "Bottom"],
                     ],
-                    inputs=[input_image, model_selector, width_slider, height_slider, alignment_dropdown],
-                    label="Examples (Prompt is optional)"
+                    inputs=[input_image, width_slider, height_slider, alignment_dropdown],
                 )
 
-            with gr.Column(scale=3): # Output column
+            with gr.Column():
                 result = gr.Image(
                     interactive=False,
                     label="Generated Image",
                     format="png",
                 )
-                history_gallery = gr.Gallery(
-                    label="History",
-                    columns=4, # Adjust columns as needed
-                    object_fit="contain",
-                    interactive=False,
-                    show_label=True,
-                    allow_preview=True,
-                    preview=True
-                 )
-
-
-    # --- Event Listeners ---
+                history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
 
-    # Update sliders and accordion based on ratio selection
+    # 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
     )
-
-    # Update ratio selection based on slider changes
     width_slider.change(
         fn=select_the_right_preset,
         inputs=[width_slider, height_slider],
@@ -466,55 +359,42 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         outputs=[target_ratio],
         queue=False
     )
-
-    # Show/hide custom resize slider
     resize_option.change(
         fn=toggle_custom_resize_slider,
         inputs=[resize_option],
         outputs=[custom_resize_percentage],
         queue=False
     )
-
-    # Define inputs for the main inference function
-    infer_inputs = [
-        model_selector, 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
-    ]
-
-    # --- Run Button Click ---
     run_button.click(
         fn=clear_result,
         inputs=None,
-        outputs=[result], # Clear only the main result image
-        queue=False # Clearing should be fast
+        outputs=result,
     ).then(
         fn=infer,
-        inputs=infer_inputs,
-        outputs=[result], # Output to the main result image
+        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=update_history, # Use the specific update function
-        inputs=[result, history_gallery], # Pass the result and current history
-        outputs=[history_gallery], # Update the history gallery
+        fn=lambda x, history: update_history(x, history),
+        inputs=[result, history_gallery],
+        outputs=history_gallery,
     )
-
-    # --- Prompt Submit (Enter Key) ---
     prompt_input.submit(
-         fn=clear_result,
+        fn=clear_result,
         inputs=None,
-        outputs=[result],
-        queue=False
+        outputs=result,
     ).then(
         fn=infer,
-        inputs=infer_inputs,
-        outputs=[result],
+        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=update_history,
+        fn=lambda x, history: update_history(x, history),
         inputs=[result, history_gallery],
-        outputs=[history_gallery],
+        outputs=history_gallery,
     )
+    demo.load(fn=load_model, inputs=model_selection, outputs=status_text)
 
-# --- Launch App ---
-# Make sure you have example images at the specified paths or remove/update the gr.Examples section
-# Create an 'examples' directory and place images like 'example_1.webp', 'example_2.jpg', 'example_3.jpg' inside it.
 demo.queue(max_size=20).launch(share=False, ssr_mode=False, show_error=True)