Update app.py

app.py

@@ -43,8 +43,6 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
-# The key differences are below - don't modify the text encoder directly
-# We'll fix it in the pipeline code instead
 
 def can_expand(source_width, source_height, target_width, target_height, alignment):
     """Checks if the image can be expanded based on the alignment."""
@@ -54,108 +52,146 @@ def can_expand(source_width, source_height, target_width, target_height, alignme
         return False
     return True
 
-@spaces.GPU(duration=24)
-def infer(image, width, height, overlap_width, num_inference_steps, resize_option, custom_resize_size, prompt_input=None, alignment="Middle"):
-    source = image
+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)
-    overlap = overlap_width
-
-    # Upscale if source is smaller than target in both dimensions
-    if source.width < target_size[0] and source.height < target_size[1]:
-        scale_factor = min(target_size[0] / source.width, target_size[1] / source.height)
-        new_width = int(source.width * scale_factor)
-        new_height = int(source.height * scale_factor)
-        source = source.resize((new_width, new_height), Image.LANCZOS)
-
-    if source.width > target_size[0] or source.height > target_size[1]:
-        scale_factor = min(target_size[0] / source.width, target_size[1] / source.height)
-        new_width = int(source.width * scale_factor)
-        new_height = int(source.height * scale_factor)
-        source = source.resize((new_width, new_height), Image.LANCZOS)
+
+    # Calculate the scaling factor to fit the image within the target size
+    scale_factor = min(target_size[0] / image.width, target_size[1] / image.height)
+    new_width = int(image.width * scale_factor)
+    new_height = int(image.height * scale_factor)
     
+    # Resize the source image to fit within target size
+    source = image.resize((new_width, new_height), Image.LANCZOS)
+
+    # Apply resize option using percentages
     if resize_option == "Full":
-        resize_size = max(source.width, source.height)
-    elif resize_option == "1/2":
-        resize_size = max(source.width, source.height) // 2
-    elif resize_option == "1/3":
-        resize_size = max(source.width, source.height) // 3
-    elif resize_option == "1/4":
-        resize_size = max(source.width, source.height) // 4
+        resize_percentage = 100
+    elif resize_option == "50%":
+        resize_percentage = 50
+    elif resize_option == "33%":
+        resize_percentage = 33
+    elif resize_option == "25%":
+        resize_percentage = 25
     else:  # Custom
-        resize_size = custom_resize_size
+        resize_percentage = custom_resize_percentage
 
-    aspect_ratio = source.height / source.width
-    new_width = resize_size
-    new_height = int(resize_size * aspect_ratio)
+    # Calculate new dimensions based on percentage
+    resize_factor = resize_percentage / 100
+    new_width = int(source.width * resize_factor)
+    new_height = int(source.height * resize_factor)
+
+    # Ensure minimum size of 64 pixels
+    new_width = max(new_width, 64)
+    new_height = max(new_height, 64)
+
+    # Resize the image
     source = source.resize((new_width, new_height), Image.LANCZOS)
 
-    if not can_expand(source.width, source.height, target_size[0], target_size[1], alignment):
-        alignment = "Middle"
+    # Calculate the overlap in pixels based on the percentage
+    overlap_x = int(new_width * (overlap_percentage / 100))
+    overlap_y = int(new_height * (overlap_percentage / 100))
+
+    # Ensure minimum overlap of 1 pixel
+    overlap_x = max(overlap_x, 1)
+    overlap_y = max(overlap_y, 1)
 
     # Calculate margins based on alignment
     if alignment == "Middle":
-        margin_x = (target_size[0] - source.width) // 2
-        margin_y = (target_size[1] - source.height) // 2
+        margin_x = (target_size[0] - new_width) // 2
+        margin_y = (target_size[1] - new_height) // 2
     elif alignment == "Left":
         margin_x = 0
-        margin_y = (target_size[1] - source.height) // 2
+        margin_y = (target_size[1] - new_height) // 2
     elif alignment == "Right":
-        margin_x = target_size[0] - source.width
-        margin_y = (target_size[1] - source.height) // 2
+        margin_x = target_size[0] - new_width
+        margin_y = (target_size[1] - new_height) // 2
     elif alignment == "Top":
-        margin_x = (target_size[0] - source.width) // 2
+        margin_x = (target_size[0] - new_width) // 2
         margin_y = 0
     elif alignment == "Bottom":
-        margin_x = (target_size[0] - source.width) // 2
-        margin_y = target_size[1] - source.height
+        margin_x = (target_size[0] - new_width) // 2
+        margin_y = target_size[1] - new_height
+
+    # Adjust margins to eliminate gaps
+    margin_x = max(0, min(margin_x, target_size[0] - new_width))
+    margin_y = max(0, min(margin_y, target_size[1] - new_height))
 
+    # Create a new background image and paste the resized source image
     background = Image.new('RGB', target_size, (255, 255, 255))
     background.paste(source, (margin_x, margin_y))
 
+    # Create the mask
     mask = Image.new('L', target_size, 255)
     mask_draw = ImageDraw.Draw(mask)
 
-    # Adjust mask generation based on alignment
-    if alignment == "Middle":
-        mask_draw.rectangle([
-            (margin_x + overlap, margin_y + overlap),
-            (margin_x + source.width - overlap, margin_y + source.height - overlap)
-        ], fill=0)
-    elif alignment == "Left":
-        mask_draw.rectangle([
-            (margin_x, margin_y),
-            (margin_x + source.width - overlap, margin_y + source.height)
-        ], fill=0)
+    # Calculate overlap areas
+    white_gaps_patch = 2
+
+    left_overlap = margin_x + overlap_x if overlap_left else margin_x + white_gaps_patch
+    right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width - white_gaps_patch
+    top_overlap = margin_y + overlap_y if overlap_top else margin_y + white_gaps_patch
+    bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height - white_gaps_patch
+    
+    if alignment == "Left":
+        left_overlap = margin_x + overlap_x if overlap_left else margin_x
     elif alignment == "Right":
-        mask_draw.rectangle([
-            (margin_x + overlap, margin_y),
-            (margin_x + source.width, margin_y + source.height)
-        ], fill=0)
+        right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width
     elif alignment == "Top":
-        mask_draw.rectangle([
-            (margin_x, margin_y),
-            (margin_x + source.width, margin_y + source.height - overlap)
-        ], fill=0)
+        top_overlap = margin_y + overlap_y if overlap_top else margin_y
     elif alignment == "Bottom":
-        mask_draw.rectangle([
-            (margin_x, margin_y + overlap),
-            (margin_x + source.width, margin_y + source.height)
-        ], fill=0)
+        bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height
+
+
+    # Draw the mask
+    mask_draw.rectangle([
+        (left_overlap, top_overlap),
+        (right_overlap, bottom_overlap)
+    ], fill=0)
+
+    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
+    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"
 
     cnet_image = background.copy()
     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,
+        negative_prompt_embeds,
+        pooled_prompt_embeds,
+        negative_pooled_prompt_embeds,
+    ) = pipe.encode_prompt(final_prompt, "cuda", True)
 
-    # Important fix: Use the original pipeline's method without modifications
-    # Let it handle the dtype conversions internally
-    encoded_prompts = pipe.encode_prompt(final_prompt, "cuda", True)
-    
     for image in pipe(
-        prompt_embeds=encoded_prompts[0],
-        negative_prompt_embeds=encoded_prompts[1],
-        pooled_prompt_embeds=encoded_prompts[2],
-        negative_pooled_prompt_embeds=encoded_prompts[3],
+        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
     ):
@@ -175,15 +211,15 @@ 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)
+        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)
+        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)
+        return changed_width, changed_height, gr.update()
     elif target_ratio == "Custom":
         return ui_width, ui_height, gr.update(open=True)
 
@@ -200,14 +236,34 @@ def select_the_right_preset(user_width, user_height):
 def toggle_custom_resize_slider(resize_option):
     return gr.update(visible=(resize_option == "Custom"))
 
+def update_history(new_image, history):
+    """Updates the history gallery with the new image."""
+    if history is None:
+        history = []
+    history.insert(0, new_image)
+    return history
+
 css = """
-footer {
-    visibility: hidden;
+.gradio-container {
+    width: 1200px !important;
 }
 """
 
+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-outpout?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>
+"""
+
 with gr.Blocks(css=css) as demo:
     with gr.Column():
+        gr.HTML(title)
+
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(
@@ -232,59 +288,72 @@ with gr.Blocks(css=css) as demo:
                     alignment_dropdown = gr.Dropdown(
                         choices=["Middle", "Left", "Right", "Top", "Bottom"],
                         value="Middle",
-                        label="고정 위치 선택"
+                        label="Alignment"
                     )
 
                 with gr.Accordion(label="Advanced settings", open=False) as settings_panel:
                     with gr.Column():
                         with gr.Row():
                             width_slider = gr.Slider(
-                                label="Width",
-                                minimum=256,
+                                label="Target Width",
+                                minimum=720,
                                 maximum=1536,
                                 step=8,
-                                value=720,
+                                value=720,  # Set a default value
                             )
                             height_slider = gr.Slider(
-                                label="Height",
-                                minimum=256,
+                                label="Target Height",
+                                minimum=720,
                                 maximum=1536,
                                 step=8,
-                                value=480,
+                                value=1280,  # Set a default value
                             )
-                        with gr.Row():
-                            num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=12, step=1, value=8)
-                            overlap_width = gr.Slider(
-                                label="Mask overlap width",
+                        
+                        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=42,
+                                value=10,
                                 step=1
                             )
+                            with gr.Row():
+                                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",
-                                choices=["Full", "1/2", "1/3", "1/4", "Custom"],
+                                choices=["Full", "50%", "33%", "25%", "Custom"],
                                 value="Full"
                             )
-
-                            custom_resize_size = gr.Slider(
-                                label="Custom resize size",
-                                minimum=256,
-                                maximum=1024,
-                                step=8,
-                                value=512,
+                            custom_resize_percentage = gr.Slider(
+                                label="Custom resize (%)",
+                                minimum=1,
+                                maximum=100,
+                                step=1,
+                                value=50,
                                 visible=False
                             )
+                        
+                        with gr.Column():
+                            preview_button = gr.Button("Preview alignment and mask")
+                            
+                            
                 gr.Examples(
                     examples=[
-                        ["./examples/example_2.jpg", 1440, 810, "Middle"],
+                        ["./examples/example_2.jpg", 1440, 810, "Left"],
+                        ["./examples/example_3.jpg", 1024, 1024, "Top"],
                         ["./examples/example_3.jpg", 1024, 1024, "Bottom"],
-                        ["./examples/example_4.png", 1024, 1024, "Top"],                         
                     ],
                     inputs=[input_image, width_slider, height_slider, alignment_dropdown],
                 )
 
+                
+
             with gr.Column():
                 result = ImageSlider(
                     interactive=False,
@@ -292,6 +361,11 @@ with gr.Blocks(css=css) as demo:
                 )
                 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")
+
+        
+
     def use_output_as_input(output_image):
         """Sets the generated output as the new input image."""
         return gr.update(value=output_image[1])
@@ -326,38 +400,56 @@ with gr.Blocks(css=css) as demo:
     resize_option.change(
         fn=toggle_custom_resize_slider,
         inputs=[resize_option],
-        outputs=[custom_resize_size],
+        outputs=[custom_resize_percentage],
         queue=False
     )
     
-    run_button.click(
+    run_button.click(  # Clear the result
         fn=clear_result,
         inputs=None,
         outputs=result,
-    ).then(
+    ).then(  # Generate the new image
         fn=infer,
-        inputs=[input_image, width_slider, height_slider, overlap_width, num_inference_steps,
-                resize_option, custom_resize_size, prompt_input, alignment_dropdown],
-        outputs=result
-    ).then(
+        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,
     )
 
-    prompt_input.submit(
+    prompt_input.submit(  # Clear the result
         fn=clear_result,
         inputs=None,
         outputs=result,
-    ).then(
+    ).then(  # Generate the new image
         fn=infer,
-        inputs=[input_image, width_slider, height_slider, overlap_width, num_inference_steps,
-                resize_option, custom_resize_size, prompt_input, alignment_dropdown],
-        outputs=result
-    ).then(
+        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,
     )
 
+    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
+    )
+
 demo.queue(max_size=12).launch(share=False)