Update app/main.py

app/main.py

@@ -1,43 +1,69 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException, Header
 from fastapi.responses import FileResponse
 from huggingface_hub import hf_hub_download
-from firebase_admin import credentials, initialize_app, app_check
 import uuid
 import os
+import io
+import json
 from PIL import Image
 import torch
-import io
 from torchvision import transforms
 
+# -------------------------------------------------
+# 🚀 FastAPI App
+# -------------------------------------------------
 app = FastAPI(title="Text-Guided Image Colorization API")
 
 # -------------------------------------------------
-# 🔐 Firebase App Check Initialization
+# 🔐 Firebase Initialization (ENV-based)
 # -------------------------------------------------
-cred = credentials.Certificate("firebase-key.json")  # Your service account key
-initialize_app(cred)
+try:
+    import firebase_admin
+    from firebase_admin import credentials, app_check
+
+    firebase_json = os.getenv("FIREBASE_CREDENTIALS")
+
+    if firebase_json:
+        print("🔥 Loading Firebase credentials from ENV...")
+        firebase_dict = json.loads(firebase_json)
+        cred = credentials.Certificate(firebase_dict)
+        firebase_admin.initialize_app(cred)
+    else:
+        print("⚠️ No Firebase credentials found. Firebase disabled.")
 
+except Exception as e:
+    print("❌ Firebase initialization failed:", e)
+
+# -------------------------------------------------
+# 📁 Directories
+# -------------------------------------------------
 UPLOAD_DIR = "uploads"
 RESULTS_DIR = "results"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(RESULTS_DIR, exist_ok=True)
 
 # -------------------------------------------------
-# 🧠 Load ColorizeNet Model
+# 🧠 Load GAN Colorization Model
 # -------------------------------------------------
 MODEL_REPO = "Hammad712/GAN-Colorization-Model"
 MODEL_FILENAME = "generator.pt"
+
+print("⬇️ Downloading model...")
 model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)
+
+print("📦 Loading model weights...")
 state_dict = torch.load(model_path, map_location="cpu")
 
-# (Example model structure – replace with your actual ColorizeNet)
-# from your_model import ColorizeNet
+# NOTE:
+# Replace this with your actual model class.
+# This is just example placeholder.
+# from model import ColorizeNet
 # model = ColorizeNet()
 # model.load_state_dict(state_dict)
 # model.eval()
 
-# Dummy colorization function
 def colorize_image(img: Image.Image):
+    """ Dummy colorizer (convert grayscale → fake color) """
     transform = transforms.ToTensor()
     tensor = transform(img.convert("L")).unsqueeze(0)
     tensor = tensor.repeat(1, 3, 1, 1)
@@ -45,23 +71,22 @@ def colorize_image(img: Image.Image):
     return output_img
 
 # -------------------------------------------------
-# 🩺 1. Health Check
+# 🩺 Health Check
 # -------------------------------------------------
 @app.get("/health")
 def health_check():
     return {"status": "healthy", "model_loaded": True}
 
 # -------------------------------------------------
-# ✅ Firebase App Check Token Validation
+# 🔐 Firebase Token Validator
 # -------------------------------------------------
 def verify_app_check_token(token: str):
-    # In production, verify token with Firebase REST API or Admin SDK.
     if not token or len(token) < 20:
-        raise HTTPException(status_code=401, detail="Missing or invalid Firebase App Check token")
+        raise HTTPException(status_code=401, detail="Invalid Firebase App Check token")
     return True
 
 # -------------------------------------------------
-# 📤 2. Upload Image
+# 📤 Upload Image
 # -------------------------------------------------
 @app.post("/upload")
 async def upload_image(
@@ -73,22 +98,22 @@ async def upload_image(
     if not file.content_type.startswith("image/"):
         raise HTTPException(status_code=400, detail="Invalid file type")
 
-    image_id = str(uuid.uuid4())
-    filename = f"{image_id}.jpg"
-    path = os.path.join(UPLOAD_DIR, filename)
+    image_id = str(uuid.uuid4()) + ".jpg"
+    file_path = os.path.join(UPLOAD_DIR, image_id)
 
-    with open(path, "wb") as f:
+    with open(file_path, "wb") as f:
         f.write(await file.read())
 
+    base_url = "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space"
+
     return {
         "success": True,
-        "image_id": image_id,
-        "image_url": f"https://logicgoinfotechspaces-text-guided-image-colorization.hf.space/uploads/{filename}",
-        "filename": filename
+        "image_id": image_id.replace(".jpg", ""),
+        "file_url": f"{base_url}/uploads/{image_id}"
     }
 
 # -------------------------------------------------
-# 🎨 3. Colorize Image
+# 🎨 Colorize Image
 # -------------------------------------------------
 @app.post("/colorize")
 async def colorize(
@@ -103,43 +128,46 @@ async def colorize(
     img = Image.open(io.BytesIO(await file.read()))
     output_img = colorize_image(img)
 
-    result_id = str(uuid.uuid4())
-    filename = f"{result_id}.jpg"
-    path = os.path.join(RESULTS_DIR, filename)
-    output_img.save(path)
+    result_id = str(uuid.uuid4()) + ".jpg"
+    output_path = os.path.join(RESULTS_DIR, result_id)
+    output_img.save(output_path)
 
     base_url = "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space"
+
     return {
         "success": True,
-        "result_id": result_id,
-        "download_url": f"{base_url}/results/{filename}",
-        "api_download_url": f"{base_url}/download/{result_id}",
-        "filename": filename
+        "result_id": result_id.replace(".jpg", ""),
+        "download_url": f"{base_url}/results/{result_id}",
+        "api_download": f"{base_url}/download/{result_id.replace('.jpg','')}"
     }
 
 # -------------------------------------------------
-# ⬇️ 4. Download Processed Image
+# ⬇️ Download via API (Secure)
 # -------------------------------------------------
 @app.get("/download/{file_id}")
 def download_result(file_id: str, x_firebase_appcheck: str = Header(None)):
     verify_app_check_token(x_firebase_appcheck)
-    path = os.path.join(RESULTS_DIR, f"{file_id}.jpg")
+
+    filename = f"{file_id}.jpg"
+    path = os.path.join(RESULTS_DIR, filename)
+
     if not os.path.exists(path):
-        raise HTTPException(status_code=404, detail="File not found")
+        raise HTTPException(status_code=404, detail="Result not found")
+
     return FileResponse(path, media_type="image/jpeg")
 
 # -------------------------------------------------
-# 🌈 5. Get Result (Public URL)
+# 🌐 Public Result File
 # -------------------------------------------------
 @app.get("/results/{filename}")
 def get_result(filename: str):
     path = os.path.join(RESULTS_DIR, filename)
     if not os.path.exists(path):
-        raise HTTPException(status_code=404, detail="File not found")
+        raise HTTPException(status_code=404, detail="Result not found")
     return FileResponse(path, media_type="image/jpeg")
 
 # -------------------------------------------------
-# 🖼️ 6. Get Uploaded Image (Public URL)
+# 🌐 Public Uploaded File
 # -------------------------------------------------
 @app.get("/uploads/{filename}")
 def get_upload(filename: str):