Update app.py

app.py

@@ -51,12 +51,20 @@ DO_SPACES_KEY = os.getenv("DO_SPACES_KEY")
 DO_SPACES_SECRET = os.getenv("DO_SPACES_SECRET")
 DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET")
 
-# NEW admin DB
+# NEW admin DB (with error handling for missing env vars)
 ADMIN_MONGO_URL = os.getenv("ADMIN_MONGO_URL")
-admin_client = AsyncIOMotorClient(ADMIN_MONGO_URL)
-admin_db = admin_client.adminPanel
-subcategories_col = admin_db.subcategories
-media_clicks_col = admin_db.media_clicks
+admin_client = None
+admin_db = None
+subcategories_col = None
+media_clicks_col = None
+if ADMIN_MONGO_URL:
+    try:
+        admin_client = AsyncIOMotorClient(ADMIN_MONGO_URL)
+        admin_db = admin_client.adminPanel
+        subcategories_col = admin_db.subcategories
+        media_clicks_col = admin_db.media_clicks
+    except Exception as e:
+        logger.warning(f"MongoDB admin connection failed (optional): {e}")
 
 # OLD logs DB
 MONGODB_URL = os.getenv("MONGODB_URL")
@@ -65,67 +73,88 @@ database = None
 
 # --------------------- Download Models ---------------------
 def download_models():
-    logger.info("Downloading models...")
-    inswapper_path = hf_hub_download(
-        repo_id=REPO_ID,
-        filename="models/inswapper_128.onnx",
-        repo_type="model",
-        local_dir=MODELS_DIR,
-        token=HF_TOKEN
-    )
-
-    buffalo_files = ["1k3d68.onnx", "2d106det.onnx", "genderage.onnx", "det_10g.onnx", "w600k_r50.onnx"]
-    for f in buffalo_files:
-        hf_hub_download(
+    try:
+        logger.info("Downloading models...")
+        inswapper_path = hf_hub_download(
             repo_id=REPO_ID,
-            filename=f"models/buffalo_l/" + f,
+            filename="models/inswapper_128.onnx",
             repo_type="model",
             local_dir=MODELS_DIR,
             token=HF_TOKEN
         )
 
-    logger.info("Models downloaded.")
-    return inswapper_path
-
+        buffalo_files = ["1k3d68.onnx", "2d106det.onnx", "genderage.onnx", "det_10g.onnx", "w600k_r50.onnx"]
+        for f in buffalo_files:
+            hf_hub_download(
+                repo_id=REPO_ID,
+                filename=f"models/buffalo_l/" + f,
+                repo_type="model",
+                local_dir=MODELS_DIR,
+                token=HF_TOKEN
+            )
 
-inswapper_path = download_models()
+        logger.info("Models downloaded successfully.")
+        return inswapper_path
+    except Exception as e:
+        logger.error(f"Model download failed: {e}")
+        raise
 
-# --------------------- Face Analysis + Swapper ---------------------
-providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
-face_analysis_app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=providers)
-face_analysis_app.prepare(ctx_id=0, det_size=(640, 640))
-swapper = insightface.model_zoo.get_model(inswapper_path, providers=providers)
+try:
+    inswapper_path = download_models()
+    
+    # --------------------- Face Analysis + Swapper ---------------------
+    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+    face_analysis_app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=providers)
+    face_analysis_app.prepare(ctx_id=0, det_size=(640, 640))
+    swapper = insightface.model_zoo.get_model(inswapper_path, providers=providers)
+    logger.info("Face analysis models loaded successfully")
+except Exception as e:
+    logger.error(f"Failed to initialize face analysis models: {e}")
+    # Set defaults to prevent crash
+    inswapper_path = None
+    face_analysis_app = None
+    swapper = None
 
 # --------------------- CodeFormer ---------------------
 CODEFORMER_PATH = "CodeFormer/inference_codeformer.py"
 
-# def ensure_codeformer():
-#     if not os.path.exists("CodeFormer"):
-#         subprocess.run("git clone https://github.com/sczhou/CodeFormer.git", shell=True, check=True)
-#         subprocess.run("pip install -r CodeFormer/requirements.txt", shell=True, check=True)
-#         subprocess.run("python CodeFormer/basicsr/setup.py develop", shell=True, check=True)
-#         subprocess.run("python CodeFormer/scripts/download_pretrained_models.py facelib", shell=True, check=True)
-#         subprocess.run("python CodeFormer/scripts/download_pretrained_models.py CodeFormer", shell=True, check=True)
 def ensure_codeformer():
-    if not os.path.exists("CodeFormer"):
-        subprocess.run("git clone https://github.com/sczhou/CodeFormer.git", shell=True, check=True)
-        subprocess.run("pip install -r CodeFormer/requirements.txt", shell=True, check=True)
-    
-    # Always ensure BasicSR is installed from local directory
-    # This is needed for Hugging Face Spaces where BasicSR can't be installed from GitHub
-    if os.path.exists("CodeFormer/basicsr/setup.py"):
-        subprocess.run("python CodeFormer/basicsr/setup.py develop", shell=True, check=True)
-    
-    # Install realesrgan after BasicSR is installed (realesrgan depends on BasicSR)
-    # This must be done after BasicSR installation to avoid PyPI install issues
     try:
-        import realesrgan
-    except ImportError:
-        subprocess.run("pip install --no-cache-dir realesrgan", shell=True, check=True)
-    
-    if not os.path.exists("CodeFormer"):
-        subprocess.run("python CodeFormer/scripts/download_pretrained_models.py facelib", shell=True, check=True)
-        subprocess.run("python CodeFormer/scripts/download_pretrained_models.py CodeFormer", shell=True, check=True)
+        if not os.path.exists("CodeFormer"):
+            logger.info("CodeFormer not found, cloning repository...")
+            subprocess.run("git clone https://github.com/sczhou/CodeFormer.git", shell=True, check=True)
+            subprocess.run("pip install -r CodeFormer/requirements.txt", shell=True, check=False)  # Non-critical deps
+        
+        # Always ensure BasicSR is installed from local directory
+        # This is needed for Hugging Face Spaces where BasicSR can't be installed from GitHub
+        if os.path.exists("CodeFormer/basicsr/setup.py"):
+            logger.info("Installing BasicSR from local directory...")
+            subprocess.run("python CodeFormer/basicsr/setup.py develop", shell=True, check=True)
+            logger.info("BasicSR installed successfully")
+        
+        # Install realesrgan after BasicSR is installed (realesrgan depends on BasicSR)
+        # This must be done after BasicSR installation to avoid PyPI install issues
+        try:
+            import realesrgan
+            logger.info("RealESRGAN already installed")
+        except ImportError:
+            logger.info("Installing RealESRGAN...")
+            subprocess.run("pip install --no-cache-dir realesrgan", shell=True, check=True)
+            logger.info("RealESRGAN installed successfully")
+        
+        # Download models if CodeFormer exists (fixed logic)
+        if os.path.exists("CodeFormer"):
+            try:
+                subprocess.run("python CodeFormer/scripts/download_pretrained_models.py facelib", shell=True, check=False, timeout=300)
+            except (subprocess.TimeoutExpired, subprocess.CalledProcessError):
+                logger.warning("Failed to download facelib models (optional)")
+            try:
+                subprocess.run("python CodeFormer/scripts/download_pretrained_models.py CodeFormer", shell=True, check=False, timeout=300)
+            except (subprocess.TimeoutExpired, subprocess.CalledProcessError):
+                logger.warning("Failed to download CodeFormer models (optional)")
+    except Exception as e:
+        logger.error(f"CodeFormer setup failed: {e}")
+        logger.warning("Continuing without CodeFormer features...")
 
 ensure_codeformer()
 
@@ -545,10 +574,18 @@ fastapi_app = FastAPI()
 @fastapi_app.on_event("startup")
 async def startup_db():
     global client, database
-    logger.info("Initializing MongoDB for API logs...")
-    client = AsyncIOMotorClient(MONGODB_URL)
-    database = client.FaceSwap
-    logger.info("MongoDB initialized for API logs")
+    if MONGODB_URL:
+        try:
+            logger.info("Initializing MongoDB for API logs...")
+            client = AsyncIOMotorClient(MONGODB_URL)
+            database = client.FaceSwap
+            logger.info("MongoDB initialized for API logs")
+        except Exception as e:
+            logger.warning(f"MongoDB connection failed (optional): {e}")
+            client = None
+            database = None
+    else:
+        logger.warning("MONGODB_URL not set, skipping MongoDB initialization")
 
 @fastapi_app.on_event("shutdown")
 async def shutdown_db():
@@ -660,6 +697,8 @@ def multi_face_swap(src_img, tgt_img):
 
     for src_face, _ in pairs:
         # 🔁 re-detect current target faces
+        if face_analysis_app is None:
+            raise ValueError("Face analysis models not initialized. Please ensure models are downloaded.")
         current_faces = face_analysis_app.get(result_img)
         current_faces = sorted(current_faces, key=face_sort_key)
 
@@ -670,6 +709,8 @@ def multi_face_swap(src_img, tgt_img):
 
         target_face = candidates[0]
 
+        if swapper is None:
+            raise ValueError("Face swap models not initialized. Please ensure models are downloaded.")
         result_img = swapper.get(
             result_img,
             target_face,
@@ -696,10 +737,14 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir=None):
 
             src_faces = face_analysis_app.get(src_bgr)
             tgt_faces = face_analysis_app.get(tgt_bgr)
+            if face_analysis_app is None:
+                return None, None, "❌ Face analysis models not initialized. Please ensure models are downloaded."
             if not src_faces or not tgt_faces:
                 return None, None, "❌ Face not detected in one of the images"
 
             swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
+            if swapper is None:
+                return None, None, "❌ Face swap models not initialized. Please ensure models are downloaded."
             swapped_bgr = swapper.get(tgt_bgr, tgt_faces[0], src_faces[0])
             if swapped_bgr is None:
                 return None, None, "❌ Face swap failed"
@@ -1590,9 +1635,13 @@ async def face_swap_api(
         with swap_lock:
             result_img = tgt_bgr.copy()
             for src_face, _ in pairs:
+                if face_analysis_app is None:
+                    raise HTTPException(status_code=500, detail="Face analysis models not initialized. Please ensure models are downloaded.")
                 current_faces = sorted(face_analysis_app.get(result_img), key=face_sort_key)
                 candidates = [f for f in current_faces if f.gender == src_face.gender] or current_faces
                 target_face = candidates[0]
+                if swapper is None:
+                    raise HTTPException(status_code=500, detail="Face swap models not initialized. Please ensure models are downloaded.")
                 result_img = swapper.get(result_img, target_face, src_face, paste_back=True)
 
         result_rgb = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)