Update app.py

app.py

@@ -12,7 +12,6 @@ from pymongo import MongoClient
 import gridfs
 from bson.objectid import ObjectId
 from PIL import Image
-import torch
 
 # Hugging Face Inference client
 from huggingface_hub import InferenceClient
@@ -20,19 +19,19 @@ from huggingface_hub import InferenceClient
 # ---------------------------------------------------------------------
 # Environment & MongoDB setup
 # ---------------------------------------------------------------------
-HF_TOKEN = os.getenv("HF_TOKEN")  # Must set in .env
+HF_TOKEN = os.getenv("HF_TOKEN")  # Make sure to set this in your environment
 if not HF_TOKEN:
     raise RuntimeError("HF_TOKEN not set in environment variables")
 
-MONGODB_URI = "mongodb+srv://harilogicgo_db_user:pdnh6UCMsWvuTCoi@kiddoimages.k2a4nuv.mongodb.net/?appName=KiddoImages"
-DB_NAME = "kiddosnap_db"
+MONGODB_URI = "mongodb+srv://harilogicgo_db_user:EpWocJKRpXau8eBW@images.nh9bx9l.mongodb.net/?appName=images"
+DB_NAME = "polaroid_db"
 
 mongo = MongoClient(MONGODB_URI)
 db = mongo[DB_NAME]
 fs = gridfs.GridFS(db)
 logs_collection = db["logs"]
 
-# HF Inference client
+# Hugging Face Inference client
 hf_client = InferenceClient(token=HF_TOKEN)
 
 # ---------------------------------------------------------------------
@@ -60,6 +59,7 @@ class HealthResponse(BaseModel):
 # ---------------------------------------------------------------------
 @app.get("/health", response_model=HealthResponse)
 def health():
+    """Health check endpoint"""
     try:
         mongo.admin.command("ping")
         return HealthResponse(status="ok", db=db.name, model="Qwen/Qwen-Image-Edit")
@@ -73,7 +73,9 @@ async def generate(
     image1: UploadFile = File(...),
     image2: Optional[UploadFile] = File(None),
 ):
-    """Upload 1 or 2 images + prompt and get edited image via HF Inference"""
+    """
+    Upload 1 or 2 images + prompt and get edited image via HF Inference
+    """
     # -----------------------------
     # 1. Read first image
     # -----------------------------
@@ -108,23 +110,24 @@ async def generate(
         raise HTTPException(status_code=500, detail=f"Failed saving input images: {e}")
 
     # -----------------------------
-    # 4. Combine images if two
+    # 4. Prepare image for HF Inference
     # -----------------------------
     if pil_img2:
+        # Combine two images side by side
         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
+        combined_img = pil_img1  # single image
 
     # -----------------------------
     # 5. Run HF Inference
     # -----------------------------
     try:
         pil_output = hf_client.image_to_image(
-            image=combined_img,  # ✅ single image
+            image=combined_img,  # Always pass single combined image
             prompt=prompt,
             model="Qwen/Qwen-Image-Edit"
         )
@@ -170,7 +173,6 @@ 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"""