Update app.py

app.py

@@ -91,6 +91,23 @@ class HealthResponse(BaseModel):
     db: str
     model: str
 
+# --------------------- UTILS ---------------------
+def resize_image_if_needed(img: Image.Image, max_size=(1024, 1024)) -> Image.Image:
+    """
+    Resize image to fit within max_size while keeping aspect ratio.
+    """
+    if img.width > max_size[0] or img.height > max_size[1]:
+        img.thumbnail(max_size, Image.ANTIALIAS)
+    return img
+
+def prepare_image(file_bytes: bytes) -> Image.Image:
+    """
+    Open image and resize if larger than 1024x1024
+    """
+    img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
+    img = resize_image_if_needed(img, max_size=(1024, 1024))
+    return img
+    
 # ---------------------------------------------------------------------
 # Endpoints
 # ---------------------------------------------------------------------
@@ -113,7 +130,7 @@ async def generate(
     # PREPARE IMAGE(S)
     try:
         img1_bytes = await image1.read()
-        pil_img1 = Image.open(io.BytesIO(img1_bytes)).convert("RGB")
+        pil_img1 = prepare_image(img1_bytes)
         input1_id = fs.put(
             img1_bytes,
             filename=image1.filename,
@@ -130,7 +147,7 @@ async def generate(
     if image2:
         try:
             img2_bytes = await image2.read()
-            pil_img2 = Image.open(io.BytesIO(img2_bytes)).convert("RGB")
+            pil_img2 = prepare_image(img2_bytes)
             input2_id = fs.put(
                 img2_bytes,
                 filename=image2.filename,