Update app.py

app.py

@@ -3,7 +3,7 @@ import spaces
 import torch
 from diffusers import AutoencoderKL, TCDScheduler
 from diffusers.models.model_loading_utils import load_state_dict
-from gradio_imageslider import ImageSlider
+# Removed: from gradio_imageslider import ImageSlider
 from huggingface_hub import hf_hub_download
 
 from controlnet_union import ControlNetModel_Union
@@ -12,6 +12,7 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
 
+# --- Model Loading (unchanged) ---
 config_file = hf_hub_download(
     "xinsir/controlnet-union-sdxl-1.0",
     filename="config_promax.json",
@@ -45,6 +46,8 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
+# --- Helper Functions (unchanged, 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:
@@ -60,7 +63,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)
 
@@ -132,7 +135,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":
@@ -153,31 +156,40 @@ 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):
     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"
+        alignment = "Middle" # Default to middle if expansion not possible with current alignment
 
     cnet_image = background.copy()
-    cnet_image.paste(0, (0, 0), mask)
+    # Prepare the controlnet input image (original image with blacked-out mask area)
+    # Note: The pipeline expects the original image content where the mask is 0 (black)
+    # and the area to be filled where the mask is 255 (white).
+    # However, the current pipeline_fill_sd_xl seems to use the mask differently internally.
+    # Let's prepare the input image as per the original logic, which pastes black over the masked area.
+    black_fill = Image.new('RGB', cnet_image.size, (0, 0, 0))
+    # Invert the mask: white (255) becomes the area to keep, black (0) the area to fill
+    inverted_mask = Image.eval(mask, lambda x: 255 - x)
+    cnet_image.paste(black_fill, (0, 0), inverted_mask) # Paste black where the inverted mask is white (original mask was 0)
+
 
     final_prompt = f"{prompt_input} , high quality, 4k"
 
@@ -188,23 +200,41 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
         negative_pooled_prompt_embeds,
     ) = pipe.encode_prompt(final_prompt, "cuda", True)
 
-    for image in pipe(
+    # Generate the image content for the masked area
+    # The pipeline yields the generated content for the masked area
+    # We only need the final result from the generator
+    generated_content = None
+    for res 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,
+        image=cnet_image, # Pass the image with blacked-out area
+        mask_image=mask, # Pass the mask (white = area to fill)
         num_inference_steps=num_inference_steps
     ):
-        yield cnet_image, image
+        generated_content = res # Keep updating until the last step
 
-    image = image.convert("RGBA")
-    cnet_image.paste(image, (0, 0), mask)
+    # The pipeline directly returns the final composite image in recent versions
+    # If it returns only the filled part, we need to composite it
+    # Let's assume the pipeline returns the final composited image based on its name "FillPipeline"
+    final_image = generated_content
+
+    # --- OLD compositing logic (keep commented in case pipeline behavior differs) ---
+    # # Convert generated content to RGBA to handle potential transparency
+    # generated_content = generated_content.convert("RGBA")
+    # # Create the final composite image by pasting the generated content onto the background
+    # final_image = background.copy().convert("RGBA")
+    # # Paste the generated content using the original mask (white area = where to paste)
+    # final_image.paste(generated_content, (0, 0), mask)
+    # final_image = final_image.convert("RGB") # Convert back to RGB if needed
+
+    # Yield only the final composited image
+    yield final_image
 
-    yield background, cnet_image
 
 def clear_result():
-    """Clears the result ImageSlider."""
+    """Clears the result Image."""
     return gr.update(value=None)
 
 def preload_presets(target_ratio, ui_width, ui_height):
@@ -222,6 +252,7 @@ def preload_presets(target_ratio, ui_width, ui_height):
         changed_height = 1024
         return changed_width, changed_height, gr.update()
     elif target_ratio == "Custom":
+        # When switching to 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):
@@ -241,9 +272,16 @@ def update_history(new_image, history):
     """Updates the history gallery with the new image."""
     if history is None:
         history = []
-    history.insert(0, new_image)
+    # Ensure new_image is a PIL Image before inserting
+    if isinstance(new_image, Image.Image):
+        history.insert(0, new_image)
+    # Handle cases where the input might be None or not an image (e.g., during clearing)
+    elif new_image is not None:
+        print(f"Warning: Attempted to add non-image type to history: {type(new_image)}")
     return history
 
+
+# --- Gradio UI ---
 css = """
 .gradio-container {
     width: 1200px !important;
@@ -279,7 +317,7 @@ with gr.Blocks(css=css) as demo:
                         value="9:16",
                         scale=2
                     )
-                    
+
                     alignment_dropdown = gr.Dropdown(
                         choices=["Middle", "Left", "Right", "Top", "Bottom"],
                         value="Middle",
@@ -294,16 +332,16 @@ with gr.Blocks(css=css) as demo:
                                 minimum=720,
                                 maximum=1536,
                                 step=8,
-                                value=720,  # Set a default value
+                                value=720,  # Default for 9:16
                             )
                             height_slider = gr.Slider(
                                 label="Target Height",
                                 minimum=720,
                                 maximum=1536,
                                 step=8,
-                                value=1280,  # Set a default value
+                                value=1280, # Default for 9:16
                             )
-                        
+
                         num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=12, step=1, value=8)
                         with gr.Group():
                             overlap_percentage = gr.Slider(
@@ -316,7 +354,7 @@ with gr.Blocks(css=css) as demo:
                             with gr.Row():
                                 overlap_top = gr.Checkbox(label="Overlap Top", value=True)
                                 overlap_right = gr.Checkbox(label="Overlap Right", value=True)
-                            with gr.Row():
+                            with gr.Row(): # Changed nesting for better layout
                                 overlap_left = gr.Checkbox(label="Overlap Left", value=True)
                                 overlap_bottom = gr.Checkbox(label="Overlap Bottom", value=True)
                         with gr.Row():
@@ -333,11 +371,11 @@ with gr.Blocks(css=css) as demo:
                                 value=50,
                                 visible=False
                             )
-                        
-                        with gr.Column():
+
+                        with gr.Column(): # Keep preview button separate
                             preview_button = gr.Button("Preview alignment and mask")
-                            
-                            
+
+
                 gr.Examples(
                     examples=[
                         ["./examples/example_1.webp", 1280, 720, "Middle"],
@@ -346,29 +384,33 @@ with gr.Blocks(css=css) as demo:
                         ["./examples/example_3.jpg", 1024, 1024, "Bottom"],
                     ],
                     inputs=[input_image, width_slider, height_slider, alignment_dropdown],
+                     # Ensure examples don't try to set output components directly
+                     # outputs=[result], # Remove output mapping from examples
+                     # fn=infer, # Don't run infer on example click, just load inputs
                 )
 
-                
 
             with gr.Column():
-                result = ImageSlider(label="Generated Image", interactive=False, type="pil", slider_color="pink")
+                # *** MODIFICATION: Changed ImageSlider to Image ***
+                result = gr.Image(label="Generated Image", interactive=False, type="pil")
                 use_as_input_button = gr.Button("Use as Input Image", visible=False)
 
-                history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
-                preview_image = gr.Image(label="Preview")
+                history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False, type="pil")
+                preview_image = gr.Image(label="Preview", type="pil") # Ensure preview is also PIL
 
-        
+    # --- Event Handlers ---
 
     def use_output_as_input(output_image):
         """Sets the generated output as the new input image."""
-        return gr.update(value=output_image[1])
+        # *** MODIFICATION: Access the image directly, not output_image[1] ***
+        return gr.update(value=output_image)
 
     use_as_input_button.click(
         fn=use_output_as_input,
-        inputs=[result],
+        inputs=[result], # Input is the single result image
         outputs=[input_image]
     )
-    
+
     target_ratio.change(
         fn=preload_presets,
         inputs=[target_ratio, width_slider, height_slider],
@@ -376,73 +418,90 @@ with gr.Blocks(css=css) as demo:
         queue=False
     )
 
+    # Link sliders change to update the ratio selection to "Custom"
     width_slider.change(
         fn=select_the_right_preset,
         inputs=[width_slider, height_slider],
         outputs=[target_ratio],
         queue=False
+    ).then(
+        fn=lambda: gr.update(open=True), # Also open accordion on slider change
+        inputs=None,
+        outputs=settings_panel,
+        queue=False
     )
 
+
     height_slider.change(
         fn=select_the_right_preset,
         inputs=[width_slider, height_slider],
         outputs=[target_ratio],
         queue=False
+    ).then(
+        fn=lambda: gr.update(open=True), # Also open accordion on slider change
+        inputs=None,
+        outputs=settings_panel,
+        queue=False
     )
 
+
     resize_option.change(
         fn=toggle_custom_resize_slider,
         inputs=[resize_option],
         outputs=[custom_resize_percentage],
         queue=False
     )
-    
-    run_button.click(  # Clear the result
-        fn=clear_result,
+
+    # Combine run logic for Button and Textbox submission
+    run_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
+    ]
+
+    def run_generation(img, w, h, ov_perc, steps, res_opt, cust_res_perc, prompt, align, ov_l, ov_r, ov_t, ov_b, history):
+        # The infer function is a generator, we need to iterate to get the final value
+        final_image = None
+        for res_img in infer(img, w, h, ov_perc, steps, res_opt, cust_res_perc, prompt, align, ov_l, ov_r, ov_t, ov_b):
+            final_image = res_img
+        
+        # Update history with the final image
+        updated_history = update_history(final_image, history)
+        
+        # Return the final image for the result component and the updated history
+        return final_image, updated_history, gr.update(visible=True) # Also make button visible
+
+
+    run_button.click(
+        fn=clear_result, # First clear the previous result
         inputs=None,
         outputs=result,
-    ).then(  # Generate the new 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(  # Update the history gallery
-        fn=lambda x, history: update_history(x[1], history),
-        inputs=[result, history_gallery],
-        outputs=history_gallery,
-    ).then(  # Show the "Use as Input Image" button
-        fn=lambda: gr.update(visible=True),
-        inputs=None,
-        outputs=use_as_input_button,
+        queue=False # Clearing should be fast
+    ).then(
+        fn=run_generation, # Then run the generation and history update
+        inputs=run_inputs + [history_gallery], # Pass current history
+        outputs=[result, history_gallery, use_as_input_button], # Update result, history, and button visibility
     )
 
-    prompt_input.submit(  # Clear the result
-        fn=clear_result,
+    prompt_input.submit(
+        fn=clear_result, # First clear the previous result
         inputs=None,
         outputs=result,
-    ).then(  # Generate the new 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(  # Update the history gallery
-        fn=lambda x, history: update_history(x[1], history),
-        inputs=[result, history_gallery],
-        outputs=history_gallery,
-    ).then(  # Show the "Use as Input Image" button
-        fn=lambda: gr.update(visible=True),
-        inputs=None,
-        outputs=use_as_input_button,
+        queue=False # Clearing should be fast
+    ).then(
+        fn=run_generation, # Then run the generation and history update
+        inputs=run_inputs + [history_gallery], # Pass current history
+        outputs=[result, history_gallery, use_as_input_button], # Update result, history, and button visibility
     )
 
+
     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_image,
-        queue=False
+        queue=False # Preview should be fast
     )
 
+# Launch the demo
 demo.queue(max_size=20).launch(share=False, ssr_mode=False, show_error=True)