Update app.py

app.py

@@ -74,13 +74,18 @@ async def generate(
     image2: Optional[UploadFile] = File(None),
 ):
     """Upload 1 or 2 images + prompt and get edited image via HF Inference"""
-    # Read images
+    # -----------------------------
+    # 1. Read first image
+    # -----------------------------
     try:
         img1_bytes = await image1.read()
         pil_img1 = Image.open(io.BytesIO(img1_bytes)).convert("RGB")
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Failed to read first image: {e}")
 
+    # -----------------------------
+    # 2. Read second image if provided
+    # -----------------------------
     pil_img2 = None
     img2_bytes = None
     if image2:
@@ -91,7 +96,9 @@ async def generate(
         except Exception as e:
             raise HTTPException(status_code=400, detail=f"Failed to read second image: {e}")
 
-    # Save input images to GridFS
+    # -----------------------------
+    # 3. Save input images to GridFS
+    # -----------------------------
     try:
         input1_id = fs.put(img1_bytes, filename=image1.filename, contentType=image1.content_type, role="input")
         input2_id = None
@@ -100,28 +107,39 @@ async def generate(
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Failed saving input images: {e}")
 
-    # Run HF Inference
+    # -----------------------------
+    # 4. Combine images if two
+    # -----------------------------
+    if pil_img2:
+        total_width = pil_img1.width + pil_img2.width
+        max_height = max(pil_img1.height, pil_img2.height)
+        combined_img = Image.new("RGB", (total_width, max_height))
+        combined_img.paste(pil_img1, (0, 0))
+        combined_img.paste(pil_img2, (pil_img1.width, 0))
+    else:
+        combined_img = pil_img1
+
+    # -----------------------------
+    # 5. Run HF Inference
+    # -----------------------------
     try:
-        images_to_pass = [pil_img1]
-        if pil_img2:
-            images_to_pass.append(pil_img2)
-
         pil_output = hf_client.image_to_image(
-            images=images_to_pass,
+            image=combined_img,  # ✅ single image
             prompt=prompt,
             model="Qwen/Qwen-Image-Edit"
         )
 
-        output_image = pil_output[0] if isinstance(pil_output, list) else pil_output
-
+        # Convert output PIL to bytes
         out_buf = io.BytesIO()
-        output_image.save(out_buf, format="PNG")
+        pil_output.save(out_buf, format="PNG")
         out_bytes = out_buf.getvalue()
     except Exception as e:
         traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"Inference failed: {e}")
 
-    # Save output image
+    # -----------------------------
+    # 6. Save output image
+    # -----------------------------
     try:
         out_id = fs.put(
             out_bytes,
@@ -135,7 +153,9 @@ async def generate(
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Failed saving output image: {e}")
 
-    # Log request
+    # -----------------------------
+    # 7. Log the request
+    # -----------------------------
     try:
         logs_collection.insert_one({
             "timestamp": datetime.utcnow(),
@@ -150,6 +170,7 @@ async def generate(
     return JSONResponse({"output_image_id": str(out_id)})
 
 
+
 @app.get("/image/{image_id}")
 def get_image(image_id: str, download: Optional[bool] = False):
     """Retrieve image by GridFS ID"""