Update app.py

app.py

@@ -44,6 +44,11 @@ MAX_SEED = 1000000
 MAX_PIXEL_BUDGET = 8192 * 8192  # Increased for tiling support
 
 
+def make_divisible_by_16(size):
+    """Adjust size to nearest multiple of 16, stretching if necessary"""
+    return ((size // 16) * 16) if (size % 16) < 8 else ((size // 16 + 1) * 16)
+
+
 def process_input(input_image, upscale_factor):
     """Process input image and handle size constraints"""
     w, h = input_image.size
@@ -101,6 +106,8 @@ def tiled_flux_img2img(pipe, prompt, image, strength, steps, guidance, generator
         for y in range(0, h, tile_size - overlap):
             tile_w = min(tile_size, w - x)
             tile_h = min(tile_size, h - y)
+            if tile_h < 16 or tile_w < 16:  # Skip tiny tiles
+                continue
             tile = image.crop((x, y, x + tile_w, y + tile_h))
 
             # Run Flux on tile
@@ -240,6 +247,13 @@ def enhance_image(
         w, h = input_image.size
         control_image = input_image.resize((w * upscale_factor, h * upscale_factor), resample=Image.LANCZOS)
 
+    # Resize control_image to divisible by 16 (stretching)
+    control_w, control_h = control_image.size
+    new_control_w = make_divisible_by_16(control_w)
+    new_control_h = make_divisible_by_16(control_h)
+    if (new_control_w, new_control_h) != (control_w, control_h):
+        control_image = control_image.resize((new_control_w, new_control_h), resample=Image.LANCZOS)
+
     # Tiled Flux Img2Img for refinement
     image = tiled_flux_img2img(
         pipe,
@@ -253,9 +267,14 @@ def enhance_image(
         overlap=32
     )
 
+    # Resize back to original target size if stretched
+    target_w, target_h = w_original * upscale_factor, h_original * upscale_factor
+    if image.size != (target_w, target_h):
+        image = image.resize((target_w, target_h), resample=Image.LANCZOS)
+
     if was_resized:
-        gr.Info(f"📏 Resizing output to target size: {w_original * upscale_factor}x{h_original * upscale_factor}")
-        image = image.resize((w_original * upscale_factor, h_original * upscale_factor), resample=Image.LANCZOS)
+        gr.Info(f"📏 Resizing output to target size: {target_w}x{target_h}")
+        image = image.resize((target_w, target_h), resample=Image.LANCZOS)
     
     # Resize input image to match output size for slider alignment
     resized_input = true_input_image.resize(image.size, resample=Image.LANCZOS)
@@ -469,4 +488,5 @@ with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX") as demo:
     """)
 
 if __name__ == "__main__":
-    demo.queue().launch(share=True, server_name="0.0.0.0", server_port=7860)
+    demo.queue().launch(share=True, server_name="0.0.0.0", server_port=7860)
+</xai:function_call