Update app.py

app.py

@@ -47,45 +47,66 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
-
-@spaces.GPU
-def infer(image, model_selection, width, height, overlap_width, num_inference_steps, prompt_input=None):
-    source = image
-    target_size = (width, height)
-    target_ratio = (width, height)  # Calculate aspect ratio from width and 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)
-
-    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))
-
+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)
+    
+    new_image = Image.new('RGB', target_size, (255, 255, 255))
+    
+    paste_x = (target_size[0] - resize_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([
-        (margin_x + overlap, margin_y + overlap),
-        (margin_x + source.width - overlap, margin_y + source.height - overlap)
-    ], fill=0)
+    mask_draw.rectangle([paste_x, paste_y, paste_x + resize_width, paste_y + new_height], fill=0)
+    
+    return new_image, mask
 
-    cnet_image = background.copy()
-    cnet_image.paste(0, (0, 0), mask)
+@spaces.GPU
+def infer(image, model_selection, width, height, overlap_width, num_inference_steps, prompt_input=None, expand_mode=False):
+    if expand_mode:
+        background, mask = resize_and_pad(image)
+        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)
+            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)
+
+        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))
+
+        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)
 
     final_prompt = "high quality"
-    if prompt_input.strip() != "":
+    if prompt_input and prompt_input.strip() != "":
         final_prompt += ", " + prompt_input
 
     (
@@ -112,27 +133,23 @@ def infer(image, model_selection, width, height, overlap_width, num_inference_st
 
 def preload_presets(target_ratio):
     if target_ratio == "9:16":
-        changed_width = 720
-        changed_height = 1280
-        return changed_width, changed_height, gr.update(open=False)
+        return 720, 1280, gr.update(visible=False), gr.update(open=False)
     elif target_ratio == "16:9":
-        changed_width = 1280
-        changed_height = 720
-        return changed_width, changed_height, gr.update(open=False)
+        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)
     elif target_ratio == "Custom":
-        return 720, 1280, gr.update(open=True)
+        return 720, 1280, gr.update(visible=False), 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;">
@@ -162,13 +179,15 @@ with gr.Blocks(css=css) as demo:
                 with gr.Row():
                     target_ratio = gr.Radio(
                         label = "Expected Ratio",
-                        choices = ["9:16", "16:9", "Custom"],
+                        choices = ["9:16", "16:9", "Expand", "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():
@@ -177,14 +196,14 @@ with gr.Blocks(css=css) as demo:
                                 minimum=720,
                                 maximum=1440,
                                 step=8,
-                                value=720,  # Set a default value
+                                value=720,
                             )
                             height_slider = gr.Slider(
                                 label="Height",
                                 minimum=720,
                                 maximum=1440,
                                 step=8,
-                                value=1280,  # Set a default value
+                                value=1280,
                             )
                         with gr.Row():
                             model_selection = gr.Dropdown(
@@ -192,7 +211,7 @@ with gr.Blocks(css=css) as demo:
                                 value="RealVisXL V5.0 Lightning",
                                 label="Model",
                             )
-                            num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=12, step=1, value=8 )
+                            num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=12, step=1, value=8)
 
                         overlap_width = gr.Slider(
                             label="Mask overlap width",
@@ -220,16 +239,17 @@ with gr.Blocks(css=css) as demo:
     target_ratio.change(
         fn = preload_presets,
         inputs = [target_ratio],
-        outputs = [width_slider, height_slider, settings_panel],
+        outputs = [width_slider, height_slider, expand_mode, 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],
+        inputs=[input_image, model_selection, width_slider, height_slider, overlap_width, num_inference_steps, prompt_input, expand_mode],
         outputs=result,
     )
 
@@ -239,7 +259,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],
+        inputs=[input_image, model_selection, width_slider, height_slider, overlap_width, num_inference_steps, prompt_input, expand_mode],
         outputs=result,
     )