Update app.py

app.py

@@ -11,16 +11,19 @@ from pydantic import BaseModel
 from pymongo import MongoClient
 import gridfs
 from bson.objectid import ObjectId
-
-import gradio as gr
-import transformers_gradio
-import spaces
-import torch
 from PIL import Image
+import torch
+
+# Hugging Face Inference client
+from huggingface_hub import InferenceClient
 
 # ---------------------------------------------------------------------
-# MongoDB setup
+# Environment & MongoDB setup
 # ---------------------------------------------------------------------
+HF_TOKEN = os.getenv("HF_TOKEN")  # Must set in .env
+if not HF_TOKEN:
+    raise RuntimeError("HF_TOKEN not set in environment variables")
+
 MONGODB_URI = "mongodb+srv://harilogicgo_db_user:jFhyDM4oA4dklUsp@api-logs.i7rqf9p.mongodb.net/?appName=API-LOGS"
 DB_NAME = "polaroid_db"
 
@@ -29,8 +32,11 @@ db = mongo[DB_NAME]
 fs = gridfs.GridFS(db)
 logs_collection = db["logs"]
 
+# HF Inference client
+hf_client = InferenceClient(token=HF_TOKEN)
+
 # ---------------------------------------------------------------------
-# FastAPI app setup
+# FastAPI app
 # ---------------------------------------------------------------------
 app = FastAPI(title="Qwen Image Edit API")
 app.add_middleware(
@@ -47,7 +53,6 @@ app.add_middleware(
 BEARER_TOKEN = "logicgo@123"
 
 def verify_token(authorization: str = Header(None)):
-    """Bearer token verification."""
     if not authorization or not authorization.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
     token = authorization.split(" ")[1]
@@ -55,14 +60,6 @@ def verify_token(authorization: str = Header(None)):
         raise HTTPException(status_code=403, detail="Invalid bearer token")
     return True
 
-# ---------------------------------------------------------------------
-# Load Qwen Image Edit model through Gradio backend
-# ---------------------------------------------------------------------
-demo = gr.load(name="Qwen/Qwen-Image-Edit", src=transformers_gradio.registry)
-demo.fn = spaces.GPU()(demo.fn)  # GPU acceleration
-for fn in demo.fns.values():
-    fn.api_name = False  # disable Gradio API names
-
 # ---------------------------------------------------------------------
 # Models
 # ---------------------------------------------------------------------
@@ -76,7 +73,6 @@ 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")
@@ -90,24 +86,16 @@ async def generate(
     image2: Optional[UploadFile] = File(None),
     authorized: bool = Depends(verify_token)
 ):
-    """
-    Upload 1 or 2 images + prompt, generate edited image, store both input/output in GridFS.
-    """
-    # -----------------------------
-    # 1. Read first image
-    # -----------------------------
+    """Upload 1 or 2 images + prompt and get edited image via HF Inference"""
+    # Read images
     try:
         img1_bytes = await image1.read()
-        if not img1_bytes:
-            raise HTTPException(status_code=400, detail="First image is empty")
         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:
         try:
             img2_bytes = await image2.read()
@@ -116,9 +104,7 @@ async def generate(
         except Exception as e:
             raise HTTPException(status_code=400, detail=f"Failed to read second image: {e}")
 
-    # -----------------------------
-    # 3. Save input images to GridFS
-    # -----------------------------
+    # Save input images to GridFS
     try:
         input1_id = fs.put(img1_bytes, filename=image1.filename, contentType=image1.content_type, role="input")
         input2_id = None
@@ -127,18 +113,25 @@ async def generate(
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Failed saving input images: {e}")
 
-    # -----------------------------
-    # 4. Run Qwen Image Edit model
-    # -----------------------------
+    # Run HF Inference
     try:
         images_to_pass = [pil_img1]
         if pil_img2:
             images_to_pass.append(pil_img2)
 
-        # Run Gradio model fn directly
-        output_image = demo.fn(images_to_pass, prompt)
+        # The "Qwen/Qwen-Image-Edit" expects a list of PIL images and a prompt
+        pil_output = hf_client.image_to_image(
+            images=images_to_pass,
+            prompt=prompt,
+            model="Qwen/Qwen-Image-Edit"
+        )
+
+        # Handle list return
+        if isinstance(pil_output, list):
+            output_image = pil_output[0]
+        else:
+            output_image = pil_output
 
-        # Convert output PIL to bytes
         out_buf = io.BytesIO()
         output_image.save(out_buf, format="PNG")
         out_bytes = out_buf.getvalue()
@@ -146,17 +139,21 @@ async def generate(
         traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"Inference failed: {e}")
 
-    # -----------------------------
-    # 5. Save output image
-    # -----------------------------
+    # Save output image
     try:
-        out_id = fs.put(out_bytes, filename=f"result_{input1_id}.png", contentType="image/png", prompt=prompt, role="output", input1_id=str(input1_id), input2_id=str(input2_id) if input2_id else None)
+        out_id = fs.put(
+            out_bytes,
+            filename=f"result_{input1_id}.png",
+            contentType="image/png",
+            prompt=prompt,
+            role="output",
+            input1_id=str(input1_id),
+            input2_id=str(input2_id) if input2_id else None
+        )
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Failed saving output image: {e}")
 
-    # -----------------------------
-    # 6. Log the request
-    # -----------------------------
+    # Log request
     try:
         logs_collection.insert_one({
             "timestamp": datetime.utcnow(),
@@ -170,6 +167,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"""
@@ -187,3 +185,10 @@ def get_image(image_id: str, download: Optional[bool] = False):
         headers["Content-Disposition"] = f'attachment; filename="{grid_out.filename}"'
 
     return StreamingResponse(iterfile(), media_type=grid_out.content_type or "application/octet-stream", headers=headers)
+
+# ---------------------------------------------------------------------
+# Run locally
+# ---------------------------------------------------------------------
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)