Update app.py

app.py

@@ -17,7 +17,7 @@ from fastapi.concurrency import run_in_threadpool
 import shutil
 import firebase_admin
 from firebase_admin import credentials, auth
-
+from PIL import Image
 from huggingface_hub import InferenceClient
 
 # ---------------------------------------------------------------------
@@ -116,6 +116,41 @@ def prepare_image(file_bytes: bytes) -> Image.Image:
     img = resize_image_if_needed(img, max_size=(1024, 1024))
     return img
     
+MAX_COMPRESSED_SIZE = 2 * 1024 * 1024  # 2 MB
+def compress_pil_image_to_2mb(
+    pil_img: Image.Image,
+    max_dim: int = 1280
+) -> bytes:
+    """
+    Resize + compress PIL image to <= 2MB.
+    Returns JPEG bytes.
+    """
+    img = pil_img.convert("RGB")
+
+    # Resize (maintain aspect ratio)
+    img.thumbnail((max_dim, max_dim), Image.LANCZOS)
+
+    quality = 85
+    buffer = io.BytesIO()
+
+    while quality >= 40:
+        buffer.seek(0)
+        buffer.truncate()
+
+        img.save(
+            buffer,
+            format="JPEG",
+            quality=quality,
+            optimize=True,
+            progressive=True
+        )
+
+        if buffer.tell() <= MAX_COMPRESSED_SIZE:
+            break
+
+        quality -= 5
+
+    return buffer.getvalue()
 # ---------------------------------------------------------------------
 # Endpoints
 # ---------------------------------------------------------------------
@@ -358,6 +393,26 @@ async def generate(
             "user_email": user.get("email"),
         }
     )
+    # -------------------------
+    # 5b. SAVE COMPRESSED IMAGE
+    # -------------------------
+    compressed_bytes = compress_pil_image_to_2mb(
+        pil_output,
+        max_dim=1280
+    )
+    
+    compressed_id = fs.put(
+        compressed_bytes,
+        filename=f"result_{input1_id}_compressed.jpg",
+        contentType="image/jpeg",
+        metadata={
+            "role": "output_compressed",
+            "original_output_id": str(out_id),
+            "prompt": prompt,
+            "user_email": user.get("email")
+        }
+    )
+
 
     response_time_ms = round((time.time() - start_time) * 1000)
 
@@ -378,7 +433,10 @@ async def generate(
     return JSONResponse({
         "output_image_id": str(out_id),
         "user": user.get("email"),
-        "response_time_ms": response_time_ms
+        "response_time_ms": response_time_ms,
+        "Compressed_Image_URL": (
+            f"https://logicgoinfotechspaces-polaroidimage.hf.space/image/{compressed_id}"
+        )
     })
 ####################---------------------------------------------------------------------------------------------------###
 ###-----OLD CODE--------------####