Update app.py

app.py

@@ -47,35 +47,44 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
-def resize_and_pad(image, target_size=(1024, 1024), resize_width=512):
-    aspect_ratio = image.height / image.width
-    new_height = int(resize_width * aspect_ratio)
-    
-    resized_image = image.resize((resize_width, new_height), Image.LANCZOS)
+def resize_and_pad(image, target_size):
+    aspect_ratio = image.width / image.height
+    target_aspect_ratio = target_size[0] / target_size[1]
+
+    if aspect_ratio > target_aspect_ratio:
+        # Image is wider than the target ratio
+        new_width = target_size[0]
+        new_height = int(new_width / aspect_ratio)
+    else:
+        # Image is taller than the target ratio
+        new_height = target_size[1]
+        new_width = int(new_height * aspect_ratio)
+
+    resized_image = image.resize((new_width, new_height), Image.LANCZOS)
     
     new_image = Image.new('RGB', target_size, (255, 255, 255))
     
-    paste_x = (target_size[0] - resize_width) // 2
+    paste_x = (target_size[0] - new_width) // 2
     paste_y = (target_size[1] - new_height) // 2
     
     new_image.paste(resized_image, (paste_x, paste_y))
     
     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)
+    mask_draw.rectangle([paste_x, paste_y, paste_x + new_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):
+    target_size = (width, height)
+    
     if expand_mode:
-        background, mask = resize_and_pad(image)
+        background, mask = resize_and_pad(image, target_size)
         cnet_image = background.copy()
         cnet_image.paste(0, (0, 0), mask)
     else:
         source = image
-        target_size = (width, height)
-        overlap = overlap_width
 
         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)
@@ -98,8 +107,8 @@ def infer(image, model_selection, width, height, overlap_width, num_inference_st
         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)
+            (margin_x + overlap_width, margin_y + overlap_width),
+            (margin_x + source.width - overlap_width, margin_y + source.height - overlap_width)
         ], fill=0)
 
         cnet_image = background.copy()
@@ -133,13 +142,11 @@ 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(visible=False), gr.update(open=False)
+        return 720, 1280, gr.update(open=False)
     elif target_ratio == "16:9":
-        return 1280, 720, gr.update(visible=False), gr.update(open=False)
-    elif target_ratio == "Expand":
-        return 1024, 1024, gr.update(visible=True), gr.update(open=False)
+        return 1280, 720, gr.update(open=False)
     elif target_ratio == "Custom":
-        return 720, 1280, gr.update(visible=False), gr.update(open=True)
+        return 720, 1280, gr.update(open=True)
 
 def clear_result():
     return gr.update(value=None)
@@ -179,15 +186,13 @@ with gr.Blocks(css=css) as demo:
                 with gr.Row():
                     target_ratio = gr.Radio(
                         label = "Expected Ratio",
-                        choices = ["9:16", "16:9", "Expand", "Custom"],
+                        choices = ["9:16", "16:9", "Custom"],
                         value = "9:16",
                         scale = 2
                     )
                     
                     run_button = gr.Button("Generate", scale=1)
 
-                expand_mode = gr.Checkbox(label="Use Expand Mode", visible=False)
-
                 with gr.Accordion(label="Advanced settings", open=False) as settings_panel:
                     with gr.Column(): 
                         with gr.Row():
@@ -220,6 +225,8 @@ with gr.Blocks(css=css) as demo:
                             value=42,
                             step=1
                         )
+                        
+                        expand_mode = gr.Checkbox(label="Use Expand Mode", value=False)
 
                 gr.Examples(
                     examples=[
@@ -239,7 +246,7 @@ with gr.Blocks(css=css) as demo:
     target_ratio.change(
         fn = preload_presets,
         inputs = [target_ratio],
-        outputs = [width_slider, height_slider, expand_mode, settings_panel],
+        outputs = [width_slider, height_slider, settings_panel],
         queue = False
     )