Update app.py

app.py

@@ -43,19 +43,34 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
+
 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:
         return False
     if alignment in ("Top", "Bottom") and source_height >= target_height:
         return False
     return True
 
-@spaces.GPU
+@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
     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)
+    
     if resize_option == "Full":
         resize_size = max(source.width, source.height)
     elif resize_option == "1/2":
@@ -74,7 +89,7 @@ def infer(image, width, height, overlap_width, num_inference_steps, resize_optio
 
     if not can_expand(source.width, source.height, target_size[0], target_size[1], alignment):
         alignment = "Middle"
-
+    # Calculate margins based on alignment
     if alignment == "Middle":
         margin_x = (target_size[0] - source.width) // 2
         margin_y = (target_size[1] - source.height) // 2
@@ -97,6 +112,7 @@ def infer(image, width, height, overlap_width, num_inference_steps, resize_optio
     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),
@@ -151,9 +167,11 @@ def infer(image, width, height, overlap_width, num_inference_steps, resize_optio
     yield background, cnet_image
 
 def clear_result():
+    """Clears the result ImageSlider."""
     return gr.update(value=None)
 
 def preload_presets(target_ratio, ui_width, ui_height):
+    """Updates the width and height sliders based on the selected aspect ratio."""
     if target_ratio == "9:16":
         changed_width = 720
         changed_height = 1280
@@ -174,6 +192,8 @@ def select_the_right_preset(user_width, user_height):
         return "9:16"
     elif user_width == 1280 and user_height == 720:
         return "16:9"
+    elif user_width == 1024 and user_height == 1024:
+        return "1:1"
     else:
         return "Custom"