Update app.py

app.py

@@ -47,42 +47,23 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
-def resize_and_pad(image, target_size, resize_width=512):
-    # Calculate the new height to maintain aspect ratio
-    aspect_ratio = image.height / image.width
-    new_height = int(resize_width * aspect_ratio)
-    
-    # Resize the image
-    resized_image = image.resize((resize_width, new_height), Image.LANCZOS)
-    
-    # Create a new white image with the target size
-    new_image = Image.new('RGB', target_size, (255, 255, 255))
-    
-    # Calculate position to paste the resized image (center it)
-    paste_x = (target_size[0] - resize_width) // 2
-    paste_y = (target_size[1] - new_height) // 2
-    
-    # Paste the resized image onto the new image
-    new_image.paste(resized_image, (paste_x, paste_y))
-    
-    # Create a mask
-    mask = Image.new('L', target_size, 255)
-    mask_draw = ImageDraw.Draw(mask)
-    mask_draw.rectangle([paste_x, paste_y, paste_x + resize_width, paste_y + new_height], fill=0)
-    
-    return new_image, mask
 
 @spaces.GPU
-def infer(image, model_selection, width, height, overlap_width, num_inference_steps, prompt_input=None, expand_mode=False):
+def infer(image, model_selection, width, height, overlap_width, num_inference_steps, resize_size, prompt_input=None):
+    source = image
     target_size = (width, height)
-    
-    if expand_mode:
-        background, mask = resize_and_pad(image, target_size, resize_width=512)
-        cnet_image = background.copy()
-        cnet_image.paste(0, (0, 0), mask)
-    else:
-        source = image
-
+    target_ratio = (width, height)
+    overlap = overlap_width
+
+    # New resizing logic based on resize_size
+    if resize_size != -1:
+        # Calculate new height to maintain aspect ratio
+        aspect_ratio = source.height / source.width
+        new_height = int(resize_size * aspect_ratio)
+        source = source.resize((resize_size, new_height), Image.LANCZOS)
+
+    # Existing resizing logic (now only applied if resize_size is -1)
+    if resize_size == -1:
         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)
@@ -95,24 +76,24 @@ def infer(image, model_selection, width, height, overlap_width, num_inference_st
             new_height = int(source.height * scale_factor)
             source = source.resize((new_width, new_height), Image.LANCZOS)
 
-        margin_x = (target_size[0] - source.width) // 2
-        margin_y = (target_size[1] - source.height) // 2
+    margin_x = (target_size[0] - source.width) // 2
+    margin_y = (target_size[1] - source.height) // 2
 
-        background = Image.new('RGB', target_size, (255, 255, 255))
-        background.paste(source, (margin_x, margin_y))
+    background = Image.new('RGB', target_size, (255, 255, 255))
+    background.paste(source, (margin_x, margin_y))
 
-        mask = Image.new('L', target_size, 255)
-        mask_draw = ImageDraw.Draw(mask)
-        mask_draw.rectangle([
-            (margin_x + overlap_width, margin_y + overlap_width),
-            (margin_x + source.width - overlap_width, margin_y + source.height - overlap_width)
-        ], fill=0)
+    mask = Image.new('L', target_size, 255)
+    mask_draw = ImageDraw.Draw(mask)
+    mask_draw.rectangle([
+        (margin_x + overlap, margin_y + overlap),
+        (margin_x + source.width - overlap, margin_y + source.height - overlap)
+    ], fill=0)
 
-        cnet_image = background.copy()
-        cnet_image.paste(0, (0, 0), mask)
+    cnet_image = background.copy()
+    cnet_image.paste(0, (0, 0), mask)
 
     final_prompt = "high quality"
-    if prompt_input and prompt_input.strip() != "":
+    if prompt_input.strip() != "":
         final_prompt += ", " + prompt_input
 
     (
@@ -139,21 +120,27 @@ def infer(image, model_selection, width, height, overlap_width, num_inference_st
 
 def preload_presets(target_ratio):
     if target_ratio == "9:16":
-        return 720, 1280, gr.update(open=False)
+        changed_width = 720
+        changed_height = 1280
+        return changed_width, changed_height, gr.update(open=False)
     elif target_ratio == "16:9":
-        return 1280, 720, gr.update(open=False)
+        changed_width = 1280
+        changed_height = 720
+        return changed_width, changed_height, gr.update(open=False)
     elif target_ratio == "Custom":
         return 720, 1280, gr.update(open=True)
 
 def clear_result():
     return gr.update(value=None)
 
+
 css = """
 .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;">
@@ -223,7 +210,11 @@ with gr.Blocks(css=css) as demo:
                             step=1
                         )
                         
-                        expand_mode = gr.Checkbox(label="Use Expand Mode", value=False)
+                        resize_size = gr.Number(
+                            label="Resize Size (-1 for default behavior)",
+                            value=-1,
+                            precision=0
+                        )
 
                 gr.Examples(
                     examples=[
@@ -246,14 +237,13 @@ with gr.Blocks(css=css) as demo:
         outputs = [width_slider, height_slider, settings_panel],
         queue = False
     )
-    
     run_button.click(
         fn=clear_result,
         inputs=None,
         outputs=result,
     ).then(
         fn=infer,
-        inputs=[input_image, model_selection, width_slider, height_slider, overlap_width, num_inference_steps, prompt_input, expand_mode],
+        inputs=[input_image, model_selection, width_slider, height_slider, overlap_width, num_inference_steps, resize_size, prompt_input],
         outputs=result,
     )
 
@@ -263,7 +253,7 @@ with gr.Blocks(css=css) as demo:
         outputs=result,
     ).then(
         fn=infer,
-        inputs=[input_image, model_selection, width_slider, height_slider, overlap_width, num_inference_steps, prompt_input, expand_mode],
+        inputs=[input_image, model_selection, width_slider, height_slider, overlap_width, num_inference_steps, resize_size, prompt_input],
         outputs=result,
     )