Update app.py

app.py

@@ -14,25 +14,21 @@ import gridfs
 from bson.objectid import ObjectId
 from PIL import Image
 from fastapi.concurrency import run_in_threadpool
+import logging
 import shutil
-import firebase_admin
-from firebase_admin import credentials, auth
 from PIL import Image
 from huggingface_hub import InferenceClient
+from dotenv import load_dotenv
+import google.generativeai as genai
 
 # ---------------------------------------------------------------------
-# Load Firebase Config from env (stringified JSON)
+# Load Firebase Config from env (stringified JSON) or from a credentials file
 # ---------------------------------------------------------------------
-firebase_config_json = os.getenv("firebase_config")
-if not firebase_config_json:
-    raise RuntimeError("❌ Missing Firebase config in environment variable 'firebase_config'")
+# Load .env file if present (allows using a .env in the project root)
+load_dotenv()
+
+# Firebase authentication removed — no initialization required
 
-try:
-    firebase_creds_dict = json.loads(firebase_config_json)
-    cred = credentials.Certificate(firebase_creds_dict)
-    firebase_admin.initialize_app(cred)
-except Exception as e:
-    raise RuntimeError(f"Failed to initialize Firebase Admin SDK: {e}")
 
 # ---------------------------------------------------------------------
 # Hugging Face setup
@@ -43,6 +39,14 @@ if not HF_TOKEN:
 
 hf_client = InferenceClient(token=HF_TOKEN)
 
+# Model configuration: defaults can be overridden via environment variables
+QWEN_MODEL = os.getenv("QWEN_MODEL", "Qwen/Qwen-Image-Edit")
+GEMINI_MODEL = os.getenv("GEMINI_MODEL", "gemini-2.5-flash-image")
+MODEL_MAP = {
+    "qwen": QWEN_MODEL,
+    "gemini": GEMINI_MODEL
+}
+
 # ---------------------------------------------------------------------
 # MongoDB setup
 # ---------------------------------------------------------------------
@@ -57,7 +61,7 @@ logs_collection = db["logs"]
 # ---------------------------------------------------------------------
 # FastAPI app setup
 # ---------------------------------------------------------------------
-app = FastAPI(title="Qwen Image Edit API with Firebase Auth")
+app = FastAPI(title="Qwen Image Edit API")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -66,22 +70,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ---------------------------------------------------------------------
-# Auth dependency
-# ---------------------------------------------------------------------
-async def verify_firebase_token(request: Request):
-    """Middleware-like dependency to verify Firebase JWT from Authorization header."""
-    auth_header = request.headers.get("Authorization")
-    if not auth_header or not auth_header.startswith("Bearer "):
-        raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
 
-    id_token = auth_header.split("Bearer ")[1]
-    try:
-        decoded_token = auth.verify_id_token(id_token)
-        request.state.user = decoded_token
-        return decoded_token
-    except Exception as e:
-        raise HTTPException(status_code=401, detail=f"Invalid or expired Firebase token: {e}")
 
 # ---------------------------------------------------------------------
 # Models
@@ -90,6 +79,7 @@ class HealthResponse(BaseModel):
     status: str
     db: str
     model: str
+    available_models: dict
 
 # --------------------- UTILS ---------------------
 def resize_image_if_needed(img: Image.Image, max_size=(1024, 1024)) -> Image.Image:
@@ -171,16 +161,19 @@ async def root():
 def health():
     """Public health check"""
     mongo.admin.command("ping")
-    return HealthResponse(status="ok", db=db.name, model="Qwen/Qwen-Image-Edit")
+    # Return default model (qwen) and available models mapping
+    return HealthResponse(status="ok", db=db.name, model=MODEL_MAP.get("qwen"), available_models=MODEL_MAP)
 
 @app.post("/generate")
 async def generate(
     prompt: str = Form(...),
+    model: str = Form("qwen"),
     image1: UploadFile = File(...),
     image2: Optional[UploadFile] = File(None),
     user_id: Optional[str] = Form(None),
     category_id: Optional[str] = Form(None),
-    user=Depends(verify_firebase_token)
+    # Firebase auth disabled — `user` can be provided by client if needed, otherwise it will be None
+    user: Optional[dict] = None
 ):
     start_time = time.time()
 
@@ -235,28 +228,21 @@ async def generate(
         try:
             admin_client = MongoClient(os.getenv("ADMIN_MONGODB_URI"))
             admin_db = admin_client["adminPanel"]
-    
+
             categories_col = admin_db.categories
             media_clicks_col = admin_db.media_clicks
-    
-            # Validate user_oid & category_oid
+
             user_oid = ObjectId(user_id)
             category_oid = ObjectId(category_id)
-    
-            # Check category exists
+
             category_doc = categories_col.find_one({"_id": category_oid})
             if not category_doc:
                 raise HTTPException(400, f"Invalid category_id: {category_id}")
-    
+
             now = datetime.utcnow()
-    
-            # Normalize dates (UTC midnight)
             today_date = datetime(now.year, now.month, now.day)
-            yesterday_date = today_date - timedelta(days=1)
-    
-            # --------------------------------------------------
-            # AI EDIT USAGE TRACKING (GLOBAL PER USER)
-            # --------------------------------------------------
+
+            # Update daily AI edit usage
             media_clicks_col.update_one(
                 {"userId": user_oid},
                 {
@@ -274,49 +260,27 @@ async def generate(
                 },
                 upsert=True
             )
-    
-            # --------------------------------------------------
-            # DAILY COUNT LOGIC
-            # --------------------------------------------------
-            now = datetime.utcnow()
-            today_date = datetime(now.year, now.month, now.day)
-            
-            doc = media_clicks_col.find_one(
-                {"userId": user_oid},
-                {"ai_edit_daily_count": 1}
-            )
-            
+
+            # Update category click counts
+            doc = media_clicks_col.find_one({"userId": user_oid}, {"ai_edit_daily_count": 1})
             daily_entries = doc.get("ai_edit_daily_count", []) if doc else []
-            
-            # Build UNIQUE date -> count map
+
             daily_map = {}
             for entry in daily_entries:
                 d = entry["date"]
                 d = datetime(d.year, d.month, d.day) if isinstance(d, datetime) else d
-                daily_map[d] = entry["count"]  # overwrite = no duplicates
-            
-            # Find last known date
+                daily_map[d] = entry["count"]
+
             last_date = max(daily_map.keys()) if daily_map else today_date
-            
-            # Fill ALL missing days with 0
             next_day = last_date + timedelta(days=1)
             while next_day < today_date:
                 daily_map.setdefault(next_day, 0)
                 next_day += timedelta(days=1)
-            
-            # Mark today as used (binary)
             daily_map[today_date] = 1
-            
-            # Rebuild list (OLD → NEW)
-            final_daily_entries = [
-                {"date": d, "count": daily_map[d]}
-                for d in sorted(daily_map.keys())
-            ]
-            
-            # Keep last 32 days only
+
+            final_daily_entries = [{"date": d, "count": daily_map[d]} for d in sorted(daily_map.keys())]
             final_daily_entries = final_daily_entries[-32:]
-            
-            # ATOMIC REPLACE (NO PUSH)
+
             media_clicks_col.update_one(
                 {"userId": user_oid},
                 {
@@ -328,10 +292,6 @@ async def generate(
                 }
             )
 
-
-            # --------------------------------------------------
-            # CATEGORY CLICK LOGIC
-            # --------------------------------------------------
             update_res = media_clicks_col.update_one(
                 {"userId": user_oid, "categories.categoryId": category_oid},
                 {
@@ -339,13 +299,10 @@ async def generate(
                         "updatedAt": now,
                         "categories.$.lastClickedAt": now
                     },
-                    "$inc": {
-                        "categories.$.click_count": 1
-                    }
+                    "$inc": {"categories.$.click_count": 1}
                 }
             )
-    
-            # If category does not exist → push new
+
             if update_res.matched_count == 0:
                 media_clicks_col.update_one(
                     {"userId": user_oid},
@@ -361,18 +318,55 @@ async def generate(
                     },
                     upsert=True
                 )
-    
+
         except Exception as e:
             print("CATEGORY_LOG_ERROR:", e)
+
     # -------------------------
-    # 4. HF INFERENCE
+    # 4. INFERENCE (Gemini direct + HF fallback)
     # -------------------------
+    used_model = None
+    pil_output = None
+    GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+
+    # determine user email safely (Firebase auth is disabled)
+    user_email = user.get("email") if user else None
+
     try:
-        pil_output = hf_client.image_to_image(
-            image=combined_img,
-            prompt=prompt,
-            model="Qwen/Qwen-Image-Edit"
-        )
+        chosen_key = model.lower()
+        chosen_model = MODEL_MAP.get(chosen_key, model)
+        fallback_model = MODEL_MAP.get("qwen") if chosen_key == "gemini" else MODEL_MAP.get("gemini")
+        used_model = chosen_model
+
+        try:
+            pil_output = hf_client.image_to_image(
+                image=combined_img,
+                prompt=prompt,
+                model=chosen_model
+            )
+        except Exception as e_primary:
+            logging.warning(f"Primary model {chosen_model} failed: {e_primary}. Trying fallback {fallback_model}")
+            try:
+                pil_output = hf_client.image_to_image(
+                    image=combined_img,
+                    prompt=prompt,
+                    model=fallback_model
+                )
+                used_model = fallback_model
+            except Exception as e_fallback:
+                response_time_ms = round((time.time() - start_time) * 1000)
+                logs_collection.insert_one({
+                    "timestamp": datetime.utcnow(),
+                    "status": "failure",
+                    "input1_id": str(input1_id),
+                    "input2_id": str(input2_id) if input2_id else None,
+                    "prompt": prompt,
+                    "user_email": user_email,
+                    "error": f"Primary error: {e_primary}; Fallback error: {e_fallback}",
+                    "response_time_ms": response_time_ms
+                })
+                raise HTTPException(500, detail=f"Inference failed: primary error: {e_primary}; fallback error: {e_fallback}")
+
     except Exception as e:
         response_time_ms = round((time.time() - start_time) * 1000)
         logs_collection.insert_one({
@@ -381,11 +375,11 @@ async def generate(
             "input1_id": str(input1_id),
             "input2_id": str(input2_id) if input2_id else None,
             "prompt": prompt,
-            "user_email": user.get("email"),
+            "user_email": user_email,
             "error": str(e),
             "response_time_ms": response_time_ms
         })
-        raise HTTPException(500, f"Inference failed: {e}")
+        raise HTTPException(500, detail=f"Inference failed: {e}")
 
     # -------------------------
     # 5. SAVE OUTPUT IMAGE
@@ -403,17 +397,13 @@ async def generate(
             "prompt": prompt,
             "input1_id": str(input1_id),
             "input2_id": str(input2_id) if input2_id else None,
-            "user_email": user.get("email"),
+            "user_email": user_email,
+            "used_model": used_model,
         }
     )
-    # -------------------------
-    # 5b. SAVE COMPRESSED IMAGE
-    # -------------------------
-    compressed_bytes = compress_pil_image_to_2mb(
-        pil_output,
-        max_dim=1280
-    )
-    
+
+    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",
@@ -422,11 +412,11 @@ async def generate(
             "role": "output_compressed",
             "original_output_id": str(out_id),
             "prompt": prompt,
-            "user_email": user.get("email")
+            "used_model": used_model,
+            "user_email": user_email
         }
     )
 
-
     response_time_ms = round((time.time() - start_time) * 1000)
 
     # -------------------------
@@ -439,18 +429,18 @@ async def generate(
         "input2_id": str(input2_id) if input2_id else None,
         "output_id": str(out_id),
         "prompt": prompt,
-        "user_email": user.get("email"),
+        "user_email": user_email,
+        "used_model": used_model,
         "response_time_ms": response_time_ms
     })
 
     return JSONResponse({
         "output_image_id": str(out_id),
-        "user": user.get("email"),
+        "user": user_email,
         "response_time_ms": response_time_ms,
-        "Compressed_Image_URL": (
-            f"https://logicgoinfotechspaces-polaroidimage.hf.space/image/{compressed_id}"
-        )
+        "Compressed_Image_URL": f"https://logicgoinfotechspaces-polaroidimage.hf.space/image/{compressed_id}"
     })
+
 ####################---------------------------------------------------------------------------------------------------###
 ###-----OLD CODE--------------####
 # @app.post("/generate")