Update app.py

app.py

@@ -11,12 +11,6 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 
 from PIL import Image, ImageDraw
 import numpy as np
-import tempfile
-
-
-# ---------------------------
-# Load ControlNet-Union + VAE + SDXL Fill pipeline (same as your Space)
-# ---------------------------
 
 config_file = hf_hub_download(
     "xinsir/controlnet-union-sdxl-1.0",
@@ -45,6 +39,7 @@ result = ControlNetModel_Union._load_pretrained_model(
 model = result[0]
 model = model.to(device="cuda", dtype=torch.float16)
 
+
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
 ).to("cuda")
@@ -60,10 +55,6 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
 
-# ---------------------------
-# Helpers (unchanged behavior)
-# ---------------------------
-
 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:
@@ -72,17 +63,14 @@ def can_expand(source_width, source_height, target_width, target_height, alignme
         return False
     return True
 
-
-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 = (int(width), int(height))
+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)
 
@@ -95,20 +83,28 @@ def prepare_image_and_mask(image, width, height, overlap_percentage,
         resize_percentage = 33
     elif resize_option == "25%":
         resize_percentage = 25
-    elif resize_option == "Custom":
-        resize_percentage = max(1, min(400, int(custom_resize_percentage)))
-    else:
-        resize_percentage = 100
+    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)
 
-    # Apply the resize percentage to the already fitted source
-    resize_factor = resize_percentage / 100.0
-    new_width = max(64, int(source.width * resize_factor))
-    new_height = max(64, 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 = max(1, int(new_width * (float(overlap_percentage) / 100.0)))
-    overlap_y = max(1, int(new_height * (float(overlap_percentage) / 100.0)))
+    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":
@@ -126,9 +122,6 @@ def prepare_image_and_mask(image, width, height, overlap_percentage,
     elif alignment == "Bottom":
         margin_x = (target_size[0] - new_width) // 2
         margin_y = target_size[1] - new_height
-    else:
-        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))
@@ -149,8 +142,7 @@ def prepare_image_and_mask(image, width, height, overlap_percentage,
     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
-
-    # Tighten edges further depending on chosen alignment
+    
     if alignment == "Left":
         left_overlap = margin_x + overlap_x if overlap_left else margin_x
     elif alignment == "Right":
@@ -160,7 +152,8 @@ def prepare_image_and_mask(image, width, height, overlap_percentage,
     elif alignment == "Bottom":
         bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height
 
-    # Draw the mask (black = keep, white = generate)
+
+    # Draw the mask
     mask_draw.rectangle([
         (left_overlap, top_overlap),
         (right_overlap, bottom_overlap)
@@ -168,40 +161,28 @@ def prepare_image_and_mask(image, width, height, overlap_percentage,
 
     return background, mask
 
-
-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 overlay
+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')
-    red_overlay = Image.new('RGBA', background.size, (255, 0, 0, 64))
+    
+    # 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
 
-
-# ---------------------------
-# Main UI inference (returns ImageSlider tuple)
-# ---------------------------
-
 @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):
-
-    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
-    )
-
+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)
+    
     if not can_expand(background.width, background.height, width, height, alignment):
         alignment = "Middle"
 
@@ -210,7 +191,7 @@ def infer(image, width, height, overlap_percentage, num_inference_steps,
 
     final_prompt = f"{prompt_input} , high quality, 4k" if prompt_input else "high quality, 4k"
 
-    # Encode prompt + run pipeline yielding previews then final
+    # Use with torch.autocast to ensure consistent dtype
     with torch.autocast(device_type="cuda", dtype=torch.float16):
         (
             prompt_embeds,
@@ -227,80 +208,12 @@ def infer(image, width, height, overlap_percentage, num_inference_steps,
             image=cnet_image,
             num_inference_steps=num_inference_steps
         ):
-            # Streaming preview to slider (left = control, right = preview)
             yield cnet_image, image
 
-    # Final composite (place the original inside the masked area)
     image = image.convert("RGBA")
     cnet_image.paste(image, (0, 0), mask)
-    yield background, cnet_image
-
-
-# ---------------------------
-# img2img-style API: /process_images (single file path string)
-# ---------------------------
-
-@spaces.GPU(duration=24)
-def process_images(
-    image,                 # PIL image from handle_file
-    prompt="",             # str
-    strength=0.75,         # kept for client parity; unused
-    seed=0,                # int
-    inference_step=8,      # int
-    width=720,             # int
-    height=1280,           # int
-    overlap_percentage=10, # float
-    alignment="Middle",    # str
-):
-    if image is None:
-        return None
-
-    # Use same prep as UI
-    resize_option = "Full"
-    custom_resize_percentage = 50
-    overlap_left = overlap_right = overlap_top = overlap_bottom = True
-
-    background, mask = prepare_image_and_mask(
-        image, int(width), int(height), float(overlap_percentage),
-        resize_option, float(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} , high quality, 4k" if prompt else "high quality, 4k"
-
-    last_img = None
-    with torch.autocast(device_type="cuda", dtype=torch.float16):
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = pipe.encode_prompt(final_prompt, "cuda", True)
-
-        for gen_img 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=int(inference_step)
-        ):
-            last_img = gen_img
-
-    if last_img is None:
-        return None
-
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-        last_img.save(tmp.name)
-        return tmp.name
-
-
-# ---------------------------
-# Misc helpers & UI wiring
-# ---------------------------
+    yield background, cnet_image
 
 def clear_result():
     """Clears the result ImageSlider."""
@@ -320,33 +233,22 @@ def preload_presets(target_ratio, ui_width, ui_height):
         changed_width = 1024
         changed_height = 1024
         return changed_width, changed_height, gr.update()
-    else:
-        return ui_width, ui_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):
-    """Chooses the closest preset by ratio (for display)."""
-    ratio = user_width / max(1, user_height)
-    if abs(ratio - (9/16)) < 0.05:
+    if user_width == 720 and user_height == 1280:
         return "9:16"
-    if abs(ratio - (16/9)) < 0.05:
+    elif user_width == 1280 and user_height == 720:
         return "16:9"
-    if abs(ratio - 1.0) < 0.05:
+    elif user_width == 1024 and user_height == 1024:
         return "1:1"
-    return "Custom"
+    else:
+        return "Custom"
 
 def toggle_custom_resize_slider(resize_option):
-    """Controls visibility of the custom resize slider."""
     return gr.update(visible=(resize_option == "Custom"))
 
-def use_output_as_input(x):
-    """API bridge for ImageSlider -> Image. Returns right-hand image as next input."""
-    if not x:
-        return None
-    if isinstance(x, (list, tuple)) and len(x) >= 2:
-        # return the generated (right) image
-        return x[1]
-    return None
-
 def update_history(new_image, history):
     """Updates the history gallery with the new image."""
     if history is None:
@@ -354,15 +256,14 @@ def update_history(new_image, history):
     history.insert(0, new_image)
     return history
 
-
 css = """
 .gradio-container {
     width: 1200px !important;
 }
 """
 
+# Define the title HTML string
 title = """<h1 align="center">Re-Size Image Outpaint</h1>
-<p align="center">Extend images with ControlNet-Union SDXL fill — with an ImageSlider preview.</p>
 """
 
 with gr.Blocks(theme="soft", css=css) as demo:
@@ -379,16 +280,29 @@ with gr.Blocks(theme="soft", css=css) as demo:
                 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():
-                    with gr.Column(scale=2):
-                        target_ratio = gr.Radio(
-                            ["9:16", "16:9", "1:1", "Custom"], value="9:16", label="Expected Ratio"
-                        )
+                    target_ratio = gr.Radio(
+                        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="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",
-                                minimum=512,
+                                minimum=720,
                                 maximum=1536,
                                 step=8,
                                 value=720,
@@ -400,7 +314,7 @@ with gr.Blocks(theme="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(
@@ -416,66 +330,86 @@ with gr.Blocks(theme="soft", css=css) as demo:
                             with gr.Row():
                                 overlap_left = gr.Checkbox(label="Overlap Left", value=True)
                                 overlap_bottom = gr.Checkbox(label="Overlap Bottom", value=True)
-
-                    with gr.Column(scale=1):
-                        with gr.Group():
+                        with gr.Row():
                             resize_option = gr.Radio(
                                 label="Resize input image",
                                 choices=["Full", "50%", "33%", "25%", "Custom"],
                                 value="Full"
                             )
-                            # FIX: set visibility here, do NOT call .update() on a component
                             custom_resize_percentage = gr.Slider(
                                 label="Custom resize (%)",
                                 minimum=1,
                                 maximum=100,
                                 step=1,
                                 value=50,
-                                visible=False,
+                                visible=False
                             )
-
+                        
                         with gr.Column():
                             preview_button = gr.Button("Preview alignment and mask")
-
+                            
+                            
                 gr.Examples(
                     examples=[
                         ["./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, target_ratio],
-                    label="Quick examples",
+                    inputs=[input_image, width_slider, height_slider, alignment_dropdown],
                 )
 
+                
+
             with gr.Column():
-                preview_image = gr.Image(label="Preview", height=300)
                 result = ImageSlider(
-                    label="Generated Image",
-                    elem_id="gen_slider",
-                    show_label=True,
                     interactive=False,
+                    label="Generated Image",
                 )
-                run_button = gr.Button("Generate", variant="primary")
-                use_as_input_button = gr.Button("Use output as input", visible=False)
-                history_gallery = gr.Gallery(label="History", columns=4, height=220)
+                use_as_input_button = gr.Button("Use as Input Image", visible=False)
 
-    # Radio preset to width/height
+                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_image):
+        """Sets the generated output as the new input image."""
+        return gr.update(value=output_image[1])
+
+    use_as_input_button.click(
+        fn=use_output_as_input,
+        inputs=[result],
+        outputs=[input_image]
+    )
+    
     target_ratio.change(
         fn=preload_presets,
         inputs=[target_ratio, width_slider, height_slider],
-        outputs=[width_slider, height_slider, gr.State()],
+        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
     )
 
-    # Toggle custom resize slider visibility
     resize_option.change(
         fn=toggle_custom_resize_slider,
         inputs=[resize_option],
         outputs=[custom_resize_percentage],
         queue=False
     )
-
-    # Generate flow: clear slider -> stream infer -> update history -> show "use as input"
+    
     run_button.click(
         fn=clear_result,
         inputs=None,
@@ -483,11 +417,12 @@ with gr.Blocks(theme="soft", css=css) as demo:
     ).then(
         fn=infer,
         inputs=[input_image, width_slider, height_slider, overlap_percentage, num_inference_steps,
-                resize_option, custom_resize_percentage, prompt_input, target_ratio,
+                resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,
                 overlap_left, overlap_right, overlap_top, overlap_bottom],
         outputs=result,
     ).then(
-        # Safely update history only if the result is not None
+        # --- FIX APPLIED HERE ---
+        # Safely update history only if the result (x) is not None.
         fn=lambda x, history: update_history(x[1], history) if x else history,
         inputs=[result, history_gallery],
         outputs=history_gallery,
@@ -497,7 +432,6 @@ with gr.Blocks(theme="soft", css=css) as demo:
         outputs=use_as_input_button,
     )
 
-    # Enter in prompt also triggers generate flow
     prompt_input.submit(
         fn=clear_result,
         inputs=None,
@@ -505,10 +439,12 @@ with gr.Blocks(theme="soft", css=css) as demo:
     ).then(
         fn=infer,
         inputs=[input_image, width_slider, height_slider, overlap_percentage, num_inference_steps,
-                resize_option, custom_resize_percentage, prompt_input, target_ratio,
+                resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,
                 overlap_left, overlap_right, overlap_top, overlap_bottom],
         outputs=result,
     ).then(
+        # --- FIX APPLIED HERE ---
+        # Safely update history only if the result (x) is not None.
         fn=lambda x, history: update_history(x[1], history) if x else history,
         inputs=[result, history_gallery],
         outputs=history_gallery,
@@ -518,43 +454,12 @@ with gr.Blocks(theme="soft", css=css) as demo:
         outputs=use_as_input_button,
     )
 
-    # Preview button
     preview_button.click(
         fn=preview_image_and_mask,
-        inputs=[input_image, width_slider, height_slider, overlap_percentage,
-                resize_option, custom_resize_percentage, target_ratio,
+        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_image,
         queue=False
     )
 
-    # Use output as next input (ImageSlider -> Image)
-    use_as_input_button.click(
-        fn=use_output_as_input,
-        inputs=[result],
-        outputs=[input_image],
-        queue=False
-    )
-
-    # ===== Hidden API binding for img2img-compatible client =====
-    # Returns a single PATH string (so your client can copy/handle it exactly like img2img)
-    api_output_path = gr.Textbox(visible=False)
-    api_trigger = gr.Button(visible=False)
-    api_trigger.click(
-        fn=process_images,
-        inputs=[
-            input_image,         # image
-            prompt_input,        # prompt
-            gr.Number(value=0.75),   # strength (ignored)
-            gr.Number(value=0),      # seed
-            num_inference_steps,     # inference_step
-            width_slider,            # width
-            height_slider,           # height
-            overlap_percentage,      # overlap_percentage
-            target_ratio             # alignment (reusing same dropdown in this UI)
-        ],
-        outputs=[api_output_path],
-        api_name="/process_images"
-    )
-
-demo.queue(max_size=12).launch(share=False)
+demo.queue(max_size=12).launch(share=False)