Update app.py

app.py

@@ -1,11 +1,8 @@
-# -*- coding: utf-8 -*-
 import gradio as gr
 import spaces
 import torch
 from diffusers import AutoencoderKL, TCDScheduler
 from diffusers.models.model_loading_utils import load_state_dict
-# Remove ImageSlider import as it's no longer needed
-# from gradio_imageslider import ImageSlider
 from huggingface_hub import hf_hub_download
 
 from controlnet_union import ControlNetModel_Union
@@ -14,7 +11,6 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
 
-# --- Model Loading (Keep as is) ---
 config_file = hf_hub_download(
     "xinsir/controlnet-union-sdxl-1.0",
     filename="config_promax.json",
@@ -26,9 +22,10 @@ model_file = hf_hub_download(
     "xinsir/controlnet-union-sdxl-1.0",
     filename="diffusion_pytorch_model_promax.safetensors",
 )
-state_dict = load_state_dict(model_file)
+
+sstate_dict = load_state_dict(model_file)
 model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
-    controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
+    controlnet_model, sstate_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
 )
 model.to(device="cuda", dtype=torch.float16)
 
@@ -46,8 +43,6 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
-# --- Helper Functions (Keep as is, except infer) ---
-
 def can_expand(source_width, source_height, target_width, target_height, alignment):
     """Checks if the image can be expanded based on the alignment."""
     if alignment in ("Left", "Right") and source_width >= target_width:
@@ -63,7 +58,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)
 
@@ -135,7 +130,7 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
     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":
@@ -145,7 +140,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),
@@ -156,47 +150,33 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
 
 def preview_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, 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)
-
+    
     # Create a preview image showing the mask
     preview = background.copy().convert('RGBA')
-
+    
     # Create a semi-transparent red overlay
     red_overlay = Image.new('RGBA', background.size, (255, 0, 0, 64))  # Reduced alpha to 64 (25% opacity)
-
+    
     # Convert black pixels in the mask to semi-transparent red
     red_mask = Image.new('RGBA', background.size, (0, 0, 0, 0))
     red_mask.paste(red_overlay, (0, 0), mask)
-
+    
     # Overlay the red mask on the background
     preview = Image.alpha_composite(preview, red_mask)
-
+    
     return preview
 
 @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):
-    if image is None:
-        raise gr.Error("Please upload an input image.")
-
     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):
-        # Optionally provide feedback or default to middle
-        # gr.Warning(f"Cannot expand image with '{alignment}' alignment as source dimension is larger than target. Defaulting to 'Middle'.")
         alignment = "Middle"
-        # Recalculate background and mask if alignment changed due to this check
-        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()
-    # Apply mask to create the input for controlnet (black out non-masked area)
-    # cnet_image.paste(0, (0, 0), mask) # This line seems incorrect for inpainting/outpainting, usually the unmasked area is kept
-    # The pipeline expects the original image content where mask=0 and potentially noise/latents where mask=1
-    # Let's keep the original image content in the unmasked area and let the pipeline handle the masked area.
-    # The `StableDiffusionXLFillPipeline` likely uses the `image` input and `mask` differently than standard inpainting.
-    # Based on typical diffusers pipelines, `image` is often the *original* content placed on the canvas.
-    # Let's pass `background` as the image input for the pipeline.
+    cnet_image.paste(0, (0, 0), mask)
 
-    final_prompt = f"{prompt_input} , high quality, 4k" if prompt_input else "high quality, 4k"
+    final_prompt = f"{prompt_input} , high quality, 4k"
 
     (
         prompt_embeds,
@@ -205,42 +185,25 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
         negative_pooled_prompt_embeds,
     ) = pipe.encode_prompt(final_prompt, "cuda", True)
 
-    # The pipeline expects the `image` and `mask_image` arguments for inpainting/outpainting
-    # `image` should be the canvas with the original image placed.
-    # `mask_image` defines the area to be filled (white=fill, black=keep).
-    # Our mask is inverted (black=keep, white=fill). Invert it.
-    inverted_mask = Image.fromarray(255 - np.array(mask))
-
-    # Run the pipeline
-    # Note: The generator inside the pipeline call is not used here as we only need the final result.
-    # We iterate once to get the final image.
-    generated_image = None
-    for img_output in pipe(
+    # 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=background, # Pass the background with the source image placed
-        mask_image=inverted_mask, # Pass the inverted mask (white = area to fill)
-        control_image=background, # ControlNet Union might need the full image context
-        num_inference_steps=num_inference_steps,
-        output_type="pil" # Ensure PIL images are returned
+        image=cnet_image,
+        num_inference_steps=num_inference_steps
     ):
-        generated_image = img_output[0] # The pipeline returns a list containing the image
-
-    if generated_image is None:
-         raise gr.Error("Image generation failed.")
-
-    # The pipeline should return the complete image already composited.
-    # No need to manually paste.
-    final_image = generated_image.convert("RGB")
+        pass  # Wait for the generation to complete
+    generated_image = image  # Get the last image
 
-    # Yield only the final generated image
-    yield final_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 component."""
+    """Clears the result Image."""
     return gr.update(value=None)
 
 def preload_presets(target_ratio, ui_width, ui_height):
@@ -248,21 +211,19 @@ 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
+        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
+        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
+        return changed_width, changed_height, gr.update()
     elif target_ratio == "Custom":
-        # Keep current slider values but open the accordion
         return ui_width, ui_height, gr.update(open=True)
 
 def select_the_right_preset(user_width, user_height):
-    """Selects the preset radio button based on current width/height."""
     if user_width == 720 and user_height == 1280:
         return "9:16"
     elif user_width == 1280 and user_height == 720:
@@ -273,49 +234,24 @@ def select_the_right_preset(user_width, user_height):
         return "Custom"
 
 def toggle_custom_resize_slider(resize_option):
-    """Shows/hides the custom resize slider."""
     return gr.update(visible=(resize_option == "Custom"))
 
 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
-        return history
     if history is None:
         history = []
-    # Prepend the new image (as PIL) to the history list
     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
 
-# --- Gradio UI ---
-
 css = """
 .gradio-container {
-    max-width: 1200px !important; /* Limit overall width */
-    margin: auto; /* Center the container */
-}
-/* Ensure gallery items are reasonably sized */
-#history_gallery .thumbnail-item {
-    height: 100px !important; /* Adjust as needed */
-}
-#history_gallery .gallery {
-    grid-template-columns: repeat(auto-fill, minmax(100px, 1fr)) !important; /* Adjust column size */
+    width: 1200px !important;
 }
-
+h1 { text-align: center; }
+footer { visibility: hidden; }
 """
 
-title = """<h1 align="center">Diffusers Image Outpaint</h1>
-<div align="center">Drop an image you would like to extend, pick your expected ratio and hit Generate.</div>
-<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
-    <p style="display: flex;gap: 6px;">
-         <a href="https://huggingface.co/spaces/fffiloni/diffusers-image-outpaint?duplicate=true">
-            <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-md.svg" alt="Duplicate this Space">
-        </a> to skip the queue and enjoy faster inference on the GPU of your choice
-    </p>
-</div>
+title = """<h1 align="center">Diffusers Image Outpaint Lightning</h1>
 """
 
 with gr.Blocks(css=css) as demo:
@@ -323,7 +259,7 @@ with gr.Blocks(css=css) as demo:
         gr.HTML(title)
 
         with gr.Row():
-            with gr.Column(scale=1): # Input column
+            with gr.Column():
                 input_image = gr.Image(
                     type="pil",
                     label="Input Image"
@@ -331,65 +267,62 @@ with gr.Blocks(css=css) as demo:
 
                 with gr.Row():
                     with gr.Column(scale=2):
-                        prompt_input = gr.Textbox(label="Prompt (Optional)", placeholder="Describe the desired extended scene...")
-                    with gr.Column(scale=1, min_width=150):
-                        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",
                         scale=2
                     )
-
+                    
                     alignment_dropdown = gr.Dropdown(
                         choices=["Middle", "Left", "Right", "Top", "Bottom"],
                         value="Middle",
-                        label="Align Source Image"
+                        label="Alignment"
                     )
 
                 with gr.Accordion(label="Advanced settings", open=False) as settings_panel:
                     with gr.Column():
                         with gr.Row():
                             width_slider = gr.Slider(
-                                label="Target Width (px)",
-                                minimum=512, # Lowered min slightly
-                                maximum=2048, # Increased max slightly
-                                step=64, # SDXL optimal step
+                                label="Target Width",
+                                minimum=720,
+                                maximum=1536,
+                                step=8,
                                 value=720,
                             )
                             height_slider = gr.Slider(
-                                label="Target Height (px)",
-                                minimum=512, # Lowered min slightly
-                                maximum=2048, # Increased max slightly
-                                step=64, # SDXL optimal step
+                                label="Target Height",
+                                minimum=720,
+                                maximum=1536,
+                                step=8,
                                 value=1280,
                             )
-
-                        num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=20, step=1, value=8) # Increased max steps slightly
+                        
+                        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 (%)",
                                 minimum=1,
                                 maximum=50,
                                 value=10,
-                                step=1,
-                                info="How much the new area overlaps the original image."
+                                step=1
                             )
-                            gr.Markdown("Select sides to overlap (influences mask generation):")
                             with gr.Row():
-                                overlap_top = gr.Checkbox(label="Top", value=True)
-                                overlap_right = gr.Checkbox(label="Right", 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="Left", value=True)
-                                overlap_bottom = gr.Checkbox(label="Bottom", value=True)
+                                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",
+                                label="Resize input image",
                                 choices=["Full", "50%", "33%", "25%", "Custom"],
-                                value="Full",
-                                info="Scales the source image down before placing it on the target canvas."
+                                value="Full"
                             )
                             custom_resize_percentage = gr.Slider(
                                 label="Custom resize (%)",
@@ -399,67 +332,37 @@ with gr.Blocks(css=css) as demo:
                                 value=50,
                                 visible=False
                             )
-
+                        
                         with gr.Column():
-                            preview_button = gr.Button("Preview Alignment & Mask")
-
-
+                            preview_button = gr.Button("Preview alignment and mask")
+                            
                 gr.Examples(
                     examples=[
-                        ["./examples/example_1.webp", 1280, 720, "Middle", "A wide landscape view of the mountains"],
-                        ["./examples/example_2.jpg", 1440, 810, "Left", "Full body shot of the cat sitting on a rug"],
-                        ["./examples/example_3.jpg", 1024, 1024, "Top", "The cloudy sky above the building"],
-                        ["./examples/example_3.jpg", 1024, 1024, "Bottom", "The street below the building"],
+                        ["./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, width_slider, height_slider, alignment_dropdown, prompt_input],
-                    label="Examples (Click to load)"
-                )
-
-            with gr.Column(scale=1): # Output column
-                # Replace ImageSlider with gr.Image
-                result_image = gr.Image(
-                    label="Generated Image",
-                    interactive=False,
-                    show_download_button=True,
-                    type="pil" # Ensure output is PIL for history
+                    inputs=[input_image, width_slider, height_slider, alignment_dropdown],
                 )
-                with gr.Row():
-                    use_as_input_button = gr.Button("Use as Input", visible=False)
-                    clear_button = gr.Button("Clear Output") # Added clear button
-
-                preview_mask_image = gr.Image(label="Alignment & Mask Preview", interactive=False)
-
-                history_gallery = gr.Gallery(
-                    label="History",
-                    columns=4, # Adjust columns as needed
-                    object_fit="contain",
-                    interactive=False,
-                    show_label=True,
-                    elem_id="history_gallery",
-                    height=300 # Set a fixed height for the gallery area
-                    )
 
+            with gr.Column():
+                result = gr.Image(label="Generated Image", type="pil")
+                use_as_input_button = gr.Button("Use as Input Image", visible=False)
 
-    # --- Event Handlers ---
+                history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
+                preview_image = gr.Image(label="Preview")
 
-    def use_output_as_input(output_pil_image):
-        """Sets the generated output PIL image as the new input image."""
-        # output_image comes directly from result_image which is PIL type
-        return gr.update(value=output_pil_image)
+    def use_output_as_input(output_image):
+        """Sets the generated output as the new input image."""
+        return gr.update(value=output_image)
 
     use_as_input_button.click(
         fn=use_output_as_input,
-        inputs=[result_image], # Input is the single result image
+        inputs=[result],
         outputs=[input_image]
     )
-
-    clear_button.click(
-        fn=lambda: (gr.update(value=None), gr.update(visible=False), gr.update(value=None)), # Clear image, hide button, clear preview
-        inputs=None,
-        outputs=[result_image, use_as_input_button, preview_mask_image],
-        queue=False
-    )
-
+    
     target_ratio.change(
         fn=preload_presets,
         inputs=[target_ratio, width_slider, height_slider],
@@ -467,18 +370,17 @@ with gr.Blocks(css=css) as demo:
         queue=False
     )
 
-    # Link sliders back to ratio selector and potentially open accordion
     width_slider.change(
-        fn=lambda w, h: (select_the_right_preset(w, h), gr.update() if select_the_right_preset(w, h) == "Custom" else gr.update()),
+        fn=select_the_right_preset,
         inputs=[width_slider, height_slider],
-        outputs=[target_ratio, settings_panel],
+        outputs=[target_ratio],
         queue=False
     )
 
     height_slider.change(
-         fn=lambda w, h: (select_the_right_preset(w, h), gr.update() if select_the_right_preset(w, h) == "Custom" else gr.update()),
+        fn=select_the_right_preset,
         inputs=[width_slider, height_slider],
-        outputs=[target_ratio, settings_panel],
+        outputs=[target_ratio],
         queue=False
     )
 
@@ -488,59 +390,53 @@ with gr.Blocks(css=css) as demo:
         outputs=[custom_resize_percentage],
         queue=False
     )
-
-    # Define common inputs for generation
-    gen_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
-    ]
-
-    # Define common steps after generation
-    def handle_output(generated_image, current_history):
-        # generated_image is the single PIL image from infer
-        new_history = update_history(generated_image, current_history)
-        button_visibility = gr.update(visible=True) if generated_image else gr.update(visible=False)
-        return generated_image, new_history, button_visibility
-
+    
     run_button.click(
-        fn=lambda: (gr.update(value=None), gr.update(visible=False)), # Clear result and hide button first
+        fn=clear_result,
         inputs=None,
-        outputs=[result_image, use_as_input_button],
-        queue=False # Don't queue the clearing part
+        outputs=result,
     ).then(
-        fn=infer, # Run the generation
-        inputs=gen_inputs,
-        outputs=result_image, # Output is the single generated image
+        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,
     ).then(
-        fn=handle_output, # Process output: update history, show button
-        inputs=[result_image, history_gallery],
-        outputs=[result_image, history_gallery, use_as_input_button] # Update result again (no change), history, and button visibility
+        fn=lambda: gr.update(visible=True),
+        inputs=None,
+        outputs=use_as_input_button,
     )
 
     prompt_input.submit(
-         fn=lambda: (gr.update(value=None), gr.update(visible=False)), # Clear result and hide button first
+        fn=clear_result,
         inputs=None,
-        outputs=[result_image, use_as_input_button],
-        queue=False # Don't queue the clearing part
+        outputs=result,
     ).then(
-        fn=infer, # Run the generation
-        inputs=gen_inputs,
-        outputs=result_image, # Output is the single generated image
+        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,
     ).then(
-        fn=handle_output, # Process output: update history, show button
-        inputs=[result_image, history_gallery],
-        outputs=[result_image, history_gallery, use_as_input_button] # Update result again (no change), history, and button visibility
+        fn=lambda: gr.update(visible=True),
+        inputs=None,
+        outputs=use_as_input_button,
     )
 
-
     preview_button.click(
         fn=preview_image_and_mask,
         inputs=[input_image, width_slider, height_slider, overlap_percentage, resize_option, custom_resize_percentage, alignment_dropdown,
                 overlap_left, overlap_right, overlap_top, overlap_bottom],
-        outputs=preview_mask_image, # Output to the preview image component
-        queue=False # Preview should be fast
+        outputs=preview_image,
+        queue=False
     )
 
-# Launch the app
-demo.queue(max_size=12).launch(share=False, ssr_mode=False, show_error=True)
+demo.queue(max_size=20).launch(share=False, ssr_mode=False, show_error=True)