Integrate FastAI colorization with Firebase auth and Gradio UI - Replace main.py with FastAI implementation - Add Gradio interface for Space UI - Add Firebase authentication to /colorize endpoint - Add curl examples documentation - Update test.py with User-Agent headers

CURL_EXAMPLES.md

@@ -0,0 +1,83 @@
+# cURL Examples for Colorization API
+
+## Base URL
+```
+https://logicgoinfotechspaces-text-guided-image-colorization.hf.space
+```
+
+## 1. Health Check (No Auth Required)
+```bash
+curl -X GET "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space/health" \
+  -H "User-Agent: Mozilla/5.0"
+```
+
+## 2. Colorize Image (With Firebase Auth)
+
+### Using Firebase ID Token
+```bash
+curl -X POST "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space/colorize" \
+  -H "Authorization: Bearer YOUR_FIREBASE_ID_TOKEN" \
+  -H "User-Agent: Mozilla/5.0" \
+  -F "file=@/path/to/your/image.jpg"
+```
+
+### Using App Check Token
+```bash
+curl -X POST "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space/colorize" \
+  -H "X-Firebase-AppCheck: YOUR_APP_CHECK_TOKEN" \
+  -H "User-Agent: Mozilla/5.0" \
+  -F "file=@/path/to/your/image.jpg"
+```
+
+### Without Auth (if DISABLE_AUTH=true)
+```bash
+curl -X POST "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space/colorize" \
+  -H "User-Agent: Mozilla/5.0" \
+  -F "file=@/path/to/your/image.jpg" \
+  -o colorized_result.png
+```
+
+## 3. Windows PowerShell Example
+```powershell
+$imagePath = "C:\projects\colorize_text\pexels-andrey-grushnikov-223358-707676.jpg"
+$headers = @{
+    "User-Agent" = "Mozilla/5.0"
+}
+$form = @{
+    file = Get-Item $imagePath
+}
+Invoke-RestMethod -Uri "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space/colorize" `
+    -Method Post `
+    -Headers $headers `
+    -Form $form `
+    -OutFile "colorized_result.png"
+```
+
+## 4. Python requests Example
+```python
+import requests
+
+url = "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space/colorize"
+headers = {
+    "User-Agent": "Mozilla/5.0",
+    # "Authorization": "Bearer YOUR_FIREBASE_ID_TOKEN",  # Uncomment if auth enabled
+}
+files = {
+    "file": ("image.jpg", open("path/to/image.jpg", "rb"), "image/jpeg")
+}
+
+response = requests.post(url, headers=headers, files=files)
+if response.status_code == 200:
+    with open("colorized_result.png", "wb") as f:
+        f.write(response.content)
+    print("Colorized image saved!")
+else:
+    print(f"Error: {response.status_code} - {response.text}")
+```
+
+## Response
+- **Success (200)**: Returns PNG image file
+- **Error (400)**: Bad request (invalid file type)
+- **Error (401)**: Unauthorized (missing/invalid auth token)
+- **Error (503)**: Service unavailable (model not loaded)
+

Dockerfile

@@ -63,4 +63,4 @@ HEALTHCHECK --interval=30s --timeout=10s --start-period=60s --retries=3 \
 ENTRYPOINT ["/entrypoint.sh"]
 
 # Run the application (port will be set via environment variable)
-CMD ["sh", "-c", "uvicorn app.main:app --host 0.0.0.0 --port ${PORT:-7860}"]
+CMD ["sh", "-c", "uvicorn app.main:app --host 0.0.0.0 --port ${PORT:-7860}"]

app/main.py

@@ -1,43 +1,39 @@
 """
-FastAPI application for image colorization using ColorizeNet model
-with Firebase App Check integration
+FastAPI application for FastAI GAN Image Colorization
+with Firebase Authentication and Gradio UI
 """
 import os
-# Set environment variables BEFORE any imports to ensure they're used
-# Set OMP_NUM_THREADS before any torch imports to avoid libgomp warnings
+# Set environment variables BEFORE any imports
 os.environ["OMP_NUM_THREADS"] = "1"
-# Set HF cache directories to writable /tmp location BEFORE any HF imports
 os.environ["HF_HOME"] = "/tmp/hf_cache"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
 os.environ["HF_HUB_CACHE"] = "/tmp/hf_cache"
 os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/hf_cache"
 os.environ["XDG_CACHE_HOME"] = "/tmp/hf_cache"
-# Set matplotlib config directory to writable location
 os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib_config"
 
+import io
 import uuid
 import logging
 from pathlib import Path
 from typing import Optional
-from fastapi import FastAPI, File, UploadFile, HTTPException, Depends, Request
+
+from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, Request
 from fastapi.responses import FileResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 import firebase_admin
 from firebase_admin import credentials, app_check, auth as firebase_auth
-import numpy as np
-import torch
 from PIL import Image
-import io
+import torch
+import uvicorn
+import gradio as gr
 
-from app.colorize_model import ColorizeModel
-from app.config import settings
+# FastAI imports
+from fastai.vision.all import *
+from huggingface_hub import from_pretrained_fastai
 
-# Create writable directories (env vars already set above)
-Path("/tmp/hf_cache").mkdir(parents=True, exist_ok=True)
-Path("/tmp/matplotlib_config").mkdir(parents=True, exist_ok=True)
-Path("/tmp/colorize_uploads").mkdir(parents=True, exist_ok=True)
-Path("/tmp/colorize_results").mkdir(parents=True, exist_ok=True)
+from app.config import settings
 
 # Configure logging
 logging.basicConfig(
@@ -46,17 +42,23 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Create writable directories
+Path("/tmp/hf_cache").mkdir(parents=True, exist_ok=True)
+Path("/tmp/matplotlib_config").mkdir(parents=True, exist_ok=True)
+Path("/tmp/colorize_uploads").mkdir(parents=True, exist_ok=True)
+Path("/tmp/colorize_results").mkdir(parents=True, exist_ok=True)
+
 # Initialize FastAPI app
 app = FastAPI(
-    title="Colorize API",
-    description="Image colorization API using ColorizeNet model",
+    title="FastAI Image Colorizer API",
+    description="Image colorization using FastAI GAN model with Firebase authentication",
     version="1.0.0"
 )
 
 # CORS middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Configure appropriately for production
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -71,7 +73,10 @@ if os.path.exists(firebase_cred_path):
         logger.info("Firebase Admin SDK initialized")
     except Exception as e:
         logger.warning("Failed to initialize Firebase: %s", str(e))
-        firebase_admin.initialize_app()
+        try:
+            firebase_admin.initialize_app()
+        except:
+            pass
 else:
     logger.warning("Firebase credentials file not found. App Check will be disabled.")
     try:
@@ -79,74 +84,40 @@ else:
     except:
         pass
 
-# Always use /data in Spaces (writable). Allow override via env.
-DATA_ROOT = Path(os.getenv("DATA_DIR", "/data"))
+# Storage directories
+UPLOAD_DIR = Path("/tmp/colorize_uploads")
+RESULT_DIR = Path("/tmp/colorize_results")
 
-
-def _ensure_dir(preferred: Path, fallback: Path) -> Path:
-    try:
-        preferred.mkdir(parents=True, exist_ok=True)
-        return preferred
-    except Exception as exc:
-        logger.warning("Failed to create directory %s: %s. Falling back to %s", preferred, exc, fallback)
-        try:
-            fallback.mkdir(parents=True, exist_ok=True)
-            return fallback
-        except Exception as fallback_exc:
-            logger.error("Could not create fallback directory %s: %s", fallback, fallback_exc)
-            raise
-
-
-UPLOAD_DIR = _ensure_dir(
-    Path(os.getenv("UPLOAD_DIR", str(DATA_ROOT / "uploads"))),
-    Path("/tmp/colorize_uploads")
-)
-RESULT_DIR = _ensure_dir(
-    Path(os.getenv("RESULT_DIR", str(DATA_ROOT / "results"))),
-    Path("/tmp/colorize_results")
-)
-logger.info("Storage directories -> uploads: %s, results: %s", str(UPLOAD_DIR), str(RESULT_DIR))
-
-# Mount static files for serving results from resolved directories
+# Mount static files
 app.mount("/results", StaticFiles(directory=str(RESULT_DIR)), name="results")
 app.mount("/uploads", StaticFiles(directory=str(UPLOAD_DIR)), name="uploads")
 
-# Initialize ColorizeNet model
-colorize_model = None
+# Initialize FastAI model
+learn = None
 model_load_error: Optional[str] = None
 
-@app.get("/")
-async def root():
-    return {
-        "app": "Colorize API",
-        "version": "1.0.0",
-        "health": "/health",
-        "upload": "/upload",
-        "colorize": "/colorize"
-    }
-
 @app.on_event("startup")
 async def startup_event():
-    """Initialize the colorization model on startup"""
-    global colorize_model, model_load_error
+    """Load FastAI model on startup"""
+    global learn, model_load_error
     try:
-        logger.info("Loading ColorizeNet model with MODEL_ID: %s", settings.MODEL_ID)
-        logger.info("HF_TOKEN present: %s", "Yes" if os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN") else "No")
-        colorize_model = ColorizeModel(settings.MODEL_ID)
-        logger.info("ColorizeNet model loaded successfully")
+        model_id = os.getenv("MODEL_ID", "Hammad712/GAN-Colorization-Model")
+        logger.info("🔄 Loading FastAI GAN Colorization Model: %s", model_id)
+        learn = from_pretrained_fastai(model_id)
+        logger.info("✅ Model loaded successfully!")
         model_load_error = None
     except Exception as e:
         error_msg = str(e)
-        logger.error("Failed to load ColorizeNet model: %s", error_msg)
+        logger.error("❌ Failed to load model: %s", error_msg)
         model_load_error = error_msg
-        # Don't raise - allow health check to work even if model fails
+        # Don't raise - allow health check to work
 
 @app.on_event("shutdown")
 async def shutdown_event():
     """Cleanup on shutdown"""
-    global colorize_model
-    if colorize_model:
-        del colorize_model
+    global learn
+    if learn:
+        del learn
     logger.info("Application shutdown")
 
 def _extract_bearer_token(authorization_header: str | None) -> str | None:
@@ -157,9 +128,9 @@ def _extract_bearer_token(authorization_header: str | None) -> str | None:
         return parts[1].strip()
     return None
 
-
 async def verify_request(request: Request):
     """
+    Verify Firebase authentication
     Accept either:
       - Firebase Auth id_token via Authorization: Bearer <id_token>
       - Firebase App Check token via X-Firebase-AppCheck (when ENABLE_APP_CHECK=true)
@@ -173,12 +144,11 @@ async def verify_request(request: Request):
     if bearer:
         try:
             decoded = firebase_auth.verify_id_token(bearer)
-            request.state.user = decoded  # make claims available if needed
+            request.state.user = decoded
             logger.info("Firebase Auth id_token verified for uid: %s", decoded.get("uid"))
             return True
         except Exception as e:
             logger.warning("Auth token verification failed: %s", str(e))
-            # fall through to App Check if enabled
 
     # If App Check is enabled, require valid App Check token
     if settings.ENABLE_APP_CHECK:
@@ -196,159 +166,136 @@ async def verify_request(request: Request):
     # Neither token required nor provided → allow (App Check disabled)
     return True
 
+@app.get("/api")
+async def api_info():
+    """API info endpoint"""
+    return {
+        "app": "FastAI Image Colorizer API",
+        "version": "1.0.0",
+        "health": "/health",
+        "colorize": "/colorize",
+        "gradio": "/"
+    }
+
 @app.get("/health")
 async def health_check():
     """Health check endpoint"""
     response = {
         "status": "healthy",
-        "model_loaded": colorize_model is not None,
-        "model_id": settings.MODEL_ID
+        "model_loaded": learn is not None,
+        "model_id": os.getenv("MODEL_ID", "Hammad712/GAN-Colorization-Model")
     }
     if model_load_error:
         response["model_error"] = model_load_error
     return response
 
-@app.post("/upload")
-async def upload_image(
-    file: UploadFile = File(...),
-    verified: bool = Depends(verify_request)
-):
-    """
-    Upload an image and return the uploaded image URL
-    """
-    if not file.content_type or not file.content_type.startswith("image/"):
-        raise HTTPException(status_code=400, detail="File must be an image")
+def colorize_pil(image: Image.Image) -> Image.Image:
+    """Run model prediction and return colorized image"""
+    if learn is None:
+        raise RuntimeError("Model not loaded")
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    pred = learn.predict(image)
+    # Handle different return types from FastAI
+    if isinstance(pred, (list, tuple)):
+        colorized = pred[0] if len(pred) > 0 else image
+    else:
+        colorized = pred
     
-    # Generate unique filename
-    file_id = str(uuid.uuid4())
-    file_extension = Path(file.filename).suffix or ".jpg"
-    filename = f"{file_id}{file_extension}"
-    filepath = UPLOAD_DIR / filename
+    # Ensure we have a PIL Image
+    if not isinstance(colorized, Image.Image):
+        if isinstance(colorized, torch.Tensor):
+            # Convert tensor to PIL
+            if colorized.dim() == 4:
+                colorized = colorized[0]
+            if colorized.dim() == 3:
+                colorized = colorized.permute(1, 2, 0).cpu()
+                if colorized.dtype in (torch.float32, torch.float16):
+                    colorized = torch.clamp(colorized, 0, 1)
+                    colorized = (colorized * 255).byte()
+                colorized = Image.fromarray(colorized.numpy(), 'RGB')
+            else:
+                raise ValueError(f"Unexpected tensor shape: {colorized.shape}")
+        else:
+            raise ValueError(f"Unexpected prediction type: {type(colorized)}")
     
-    # Save uploaded file
-    try:
-        contents = await file.read()
-        with open(filepath, "wb") as f:
-            f.write(contents)
-        logger.info("Image uploaded: %s", filename)
-        
-        # Return the URL to access the uploaded image
-        base_url = os.getenv("BASE_URL", os.getenv("SPACE_HOST", "http://localhost:7860"))
-        image_url = f"{base_url}/uploads/{filename}"
-        
-        return {
-            "success": True,
-            "image_id": file_id,
-            "image_url": image_url,
-            "filename": filename
-        }
-    except Exception as e:
-        logger.error("Error uploading image: %s", str(e))
-        raise HTTPException(status_code=500, detail=f"Error uploading image: {str(e)}")
+    if colorized.mode != "RGB":
+        colorized = colorized.convert("RGB")
+    
+    return colorized
 
 @app.post("/colorize")
-async def colorize_image(
+async def colorize_api(
     file: UploadFile = File(...),
     verified: bool = Depends(verify_request)
 ):
     """
-    Colorize an uploaded grayscale image using ColorizeNet
-    Returns the colorized image URL
+    Upload a black & white image -> returns colorized image.
+    Requires Firebase authentication unless DISABLE_AUTH=true
     """
-    if colorize_model is None:
+    if learn is None:
         raise HTTPException(status_code=503, detail="Colorization model not loaded")
     
     if not file.content_type or not file.content_type.startswith("image/"):
         raise HTTPException(status_code=400, detail="File must be an image")
     
     try:
-        # Read image
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents))
-        
-        # Convert to RGB if needed
-        if image.mode != "RGB":
-            image = image.convert("RGB")
+        img_bytes = await file.read()
+        image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
         
-        # Colorize the image
         logger.info("Colorizing image...")
-        colorized_image, caption = colorize_model.colorize(image)
+        colorized = colorize_pil(image)
         
-        # Save colorized image
-        file_id = str(uuid.uuid4())
-        result_filename = f"{file_id}.jpg"
-        result_filepath = RESULT_DIR / result_filename
+        output_filename = f"{uuid.uuid4()}.png"
+        output_path = RESULT_DIR / output_filename
+        colorized.save(output_path, "PNG")
         
-        colorized_image.save(result_filepath, "JPEG", quality=95)
-        logger.info("Colorized image saved: %s", result_filename)
+        logger.info("Colorized image saved: %s", output_filename)
         
-        # Return URLs
-        base_url = os.getenv("BASE_URL", os.getenv("SPACE_HOST", "http://localhost:7860"))
-        download_url = f"{base_url}/results/{result_filename}"
-        api_download_url = f"{base_url}/download/{file_id}"
-        
-        return {
-            "success": True,
-            "result_id": file_id,
-            "download_url": download_url,
-            "api_download_url": api_download_url,
-            "filename": result_filename,
-            "caption": caption
-        }
+        # Return the image file
+        return FileResponse(
+            output_path,
+            media_type="image/png",
+            filename=f"colorized_{output_filename}"
+        )
     except Exception as e:
         logger.error("Error colorizing image: %s", str(e))
         raise HTTPException(status_code=500, detail=f"Error colorizing image: {str(e)}")
 
-@app.get("/download/{file_id}")
-async def download_result(
-    file_id: str,
-    verified: bool = Depends(verify_request)
-):
-    """
-    Download the colorized image by file ID
-    """
-    result_filepath = RESULT_DIR / f"{file_id}.jpg"
-    
-    if not result_filepath.exists():
-        raise HTTPException(status_code=404, detail="Result not found")
-    
-    return FileResponse(
-        result_filepath,
-        media_type="image/jpeg",
-        filename=f"colorized_{file_id}.jpg"
-    )
+# ==========================================================
+# Gradio Interface (for Space UI)
+# ==========================================================
+def gradio_colorize(image):
+    """Gradio colorization function"""
+    if image is None:
+        return None
+    try:
+        if learn is None:
+            return None
+        return colorize_pil(image)
+    except Exception as e:
+        logger.error("Gradio colorization error: %s", str(e))
+        return None
 
-@app.get("/results/{filename}")
-async def get_result_file(filename: str):
-    """
-    Serve result files directly (public endpoint for browser access)
-    """
-    result_filepath = RESULT_DIR / filename
-    
-    if not result_filepath.exists():
-        raise HTTPException(status_code=404, detail="File not found")
-    
-    return FileResponse(
-        result_filepath,
-        media_type="image/jpeg"
-    )
+title = "🎨 FastAI GAN Image Colorizer"
+description = "Upload a black & white photo to generate a colorized version using the FastAI GAN model."
 
-@app.get("/uploads/{filename}")
-async def get_upload_file(filename: str):
-    """
-    Serve uploaded files directly
-    """
-    upload_filepath = UPLOAD_DIR / filename
-    
-    if not upload_filepath.exists():
-        raise HTTPException(status_code=404, detail="File not found")
-    
-    return FileResponse(
-        upload_filepath,
-        media_type="image/jpeg"
-    )
+iface = gr.Interface(
+    fn=gradio_colorize,
+    inputs=gr.Image(type="pil", label="Upload B&W Image"),
+    outputs=gr.Image(type="pil", label="Colorized Image"),
+    title=title,
+    description=description,
+)
 
+# Mount Gradio app at root (this will be the Space UI)
+# Note: This will override the root endpoint, so use /api for API info
+app = gr.mount_gradio_app(app, iface, path="/")
+
+# ==========================================================
+# Run Server
+# ==========================================================
 if __name__ == "__main__":
-    import uvicorn
     port = int(os.getenv("PORT", "7860"))
     uvicorn.run(app, host="0.0.0.0", port=port)
+

app/main_fastai.py

@@ -0,0 +1,301 @@
+"""
+FastAPI application for FastAI GAN Image Colorization
+with Firebase Authentication and Gradio UI
+"""
+import os
+# Set environment variables BEFORE any imports
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["HF_HOME"] = "/tmp/hf_cache"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
+os.environ["HF_HUB_CACHE"] = "/tmp/hf_cache"
+os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/hf_cache"
+os.environ["XDG_CACHE_HOME"] = "/tmp/hf_cache"
+os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib_config"
+
+import io
+import uuid
+import logging
+from pathlib import Path
+from typing import Optional
+
+from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, Request
+from fastapi.responses import FileResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+import firebase_admin
+from firebase_admin import credentials, app_check, auth as firebase_auth
+from PIL import Image
+import torch
+import uvicorn
+import gradio as gr
+
+# FastAI imports
+from fastai.vision.all import *
+from huggingface_hub import from_pretrained_fastai
+
+from app.config import settings
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Create writable directories
+Path("/tmp/hf_cache").mkdir(parents=True, exist_ok=True)
+Path("/tmp/matplotlib_config").mkdir(parents=True, exist_ok=True)
+Path("/tmp/colorize_uploads").mkdir(parents=True, exist_ok=True)
+Path("/tmp/colorize_results").mkdir(parents=True, exist_ok=True)
+
+# Initialize FastAPI app
+app = FastAPI(
+    title="FastAI Image Colorizer API",
+    description="Image colorization using FastAI GAN model with Firebase authentication",
+    version="1.0.0"
+)
+
+# CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Initialize Firebase Admin SDK
+firebase_cred_path = os.getenv("FIREBASE_CREDENTIALS_PATH", "/tmp/firebase-adminsdk.json")
+if os.path.exists(firebase_cred_path):
+    try:
+        cred = credentials.Certificate(firebase_cred_path)
+        firebase_admin.initialize_app(cred)
+        logger.info("Firebase Admin SDK initialized")
+    except Exception as e:
+        logger.warning("Failed to initialize Firebase: %s", str(e))
+        try:
+            firebase_admin.initialize_app()
+        except:
+            pass
+else:
+    logger.warning("Firebase credentials file not found. App Check will be disabled.")
+    try:
+        firebase_admin.initialize_app()
+    except:
+        pass
+
+# Storage directories
+UPLOAD_DIR = Path("/tmp/colorize_uploads")
+RESULT_DIR = Path("/tmp/colorize_results")
+
+# Mount static files
+app.mount("/results", StaticFiles(directory=str(RESULT_DIR)), name="results")
+app.mount("/uploads", StaticFiles(directory=str(UPLOAD_DIR)), name="uploads")
+
+# Initialize FastAI model
+learn = None
+model_load_error: Optional[str] = None
+
+@app.on_event("startup")
+async def startup_event():
+    """Load FastAI model on startup"""
+    global learn, model_load_error
+    try:
+        model_id = os.getenv("MODEL_ID", "Hammad712/GAN-Colorization-Model")
+        logger.info("🔄 Loading FastAI GAN Colorization Model: %s", model_id)
+        learn = from_pretrained_fastai(model_id)
+        logger.info("✅ Model loaded successfully!")
+        model_load_error = None
+    except Exception as e:
+        error_msg = str(e)
+        logger.error("❌ Failed to load model: %s", error_msg)
+        model_load_error = error_msg
+        # Don't raise - allow health check to work
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    """Cleanup on shutdown"""
+    global learn
+    if learn:
+        del learn
+    logger.info("Application shutdown")
+
+def _extract_bearer_token(authorization_header: str | None) -> str | None:
+    if not authorization_header:
+        return None
+    parts = authorization_header.split(" ", 1)
+    if len(parts) == 2 and parts[0].lower() == "bearer":
+        return parts[1].strip()
+    return None
+
+async def verify_request(request: Request):
+    """
+    Verify Firebase authentication
+    Accept either:
+      - Firebase Auth id_token via Authorization: Bearer <id_token>
+      - Firebase App Check token via X-Firebase-AppCheck (when ENABLE_APP_CHECK=true)
+    """
+    # If Firebase is not initialized or auth is explicitly disabled, allow
+    if not firebase_admin._apps or os.getenv("DISABLE_AUTH", "false").lower() == "true":
+        return True
+
+    # Try Firebase Auth id_token first if present
+    bearer = _extract_bearer_token(request.headers.get("Authorization"))
+    if bearer:
+        try:
+            decoded = firebase_auth.verify_id_token(bearer)
+            request.state.user = decoded
+            logger.info("Firebase Auth id_token verified for uid: %s", decoded.get("uid"))
+            return True
+        except Exception as e:
+            logger.warning("Auth token verification failed: %s", str(e))
+
+    # If App Check is enabled, require valid App Check token
+    if settings.ENABLE_APP_CHECK:
+        app_check_token = request.headers.get("X-Firebase-AppCheck")
+        if not app_check_token:
+            raise HTTPException(status_code=401, detail="Missing App Check token")
+        try:
+            app_check_claims = app_check.verify_token(app_check_token)
+            logger.info("App Check token verified for: %s", app_check_claims.get("app_id"))
+            return True
+        except Exception as e:
+            logger.warning("App Check token verification failed: %s", str(e))
+            raise HTTPException(status_code=401, detail="Invalid App Check token")
+
+    # Neither token required nor provided → allow (App Check disabled)
+    return True
+
+@app.get("/api")
+async def api_info():
+    """API info endpoint"""
+    return {
+        "app": "FastAI Image Colorizer API",
+        "version": "1.0.0",
+        "health": "/health",
+        "colorize": "/colorize",
+        "gradio": "/"
+    }
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    response = {
+        "status": "healthy",
+        "model_loaded": learn is not None,
+        "model_id": os.getenv("MODEL_ID", "Hammad712/GAN-Colorization-Model")
+    }
+    if model_load_error:
+        response["model_error"] = model_load_error
+    return response
+
+def colorize_pil(image: Image.Image) -> Image.Image:
+    """Run model prediction and return colorized image"""
+    if learn is None:
+        raise RuntimeError("Model not loaded")
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    pred = learn.predict(image)
+    # Handle different return types from FastAI
+    if isinstance(pred, (list, tuple)):
+        colorized = pred[0] if len(pred) > 0 else image
+    else:
+        colorized = pred
+    
+    # Ensure we have a PIL Image
+    if not isinstance(colorized, Image.Image):
+        if isinstance(colorized, torch.Tensor):
+            # Convert tensor to PIL
+            if colorized.dim() == 4:
+                colorized = colorized[0]
+            if colorized.dim() == 3:
+                colorized = colorized.permute(1, 2, 0).cpu()
+                if colorized.dtype in (torch.float32, torch.float16):
+                    colorized = torch.clamp(colorized, 0, 1)
+                    colorized = (colorized * 255).byte()
+                colorized = Image.fromarray(colorized.numpy(), 'RGB')
+            else:
+                raise ValueError(f"Unexpected tensor shape: {colorized.shape}")
+        else:
+            raise ValueError(f"Unexpected prediction type: {type(colorized)}")
+    
+    if colorized.mode != "RGB":
+        colorized = colorized.convert("RGB")
+    
+    return colorized
+
+@app.post("/colorize")
+async def colorize_api(
+    file: UploadFile = File(...),
+    verified: bool = Depends(verify_request)
+):
+    """
+    Upload a black & white image -> returns colorized image.
+    Requires Firebase authentication unless DISABLE_AUTH=true
+    """
+    if learn is None:
+        raise HTTPException(status_code=503, detail="Colorization model not loaded")
+    
+    if not file.content_type or not file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="File must be an image")
+    
+    try:
+        img_bytes = await file.read()
+        image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+        
+        logger.info("Colorizing image...")
+        colorized = colorize_pil(image)
+        
+        output_filename = f"{uuid.uuid4()}.png"
+        output_path = RESULT_DIR / output_filename
+        colorized.save(output_path, "PNG")
+        
+        logger.info("Colorized image saved: %s", output_filename)
+        
+        # Return the image file
+        return FileResponse(
+            output_path,
+            media_type="image/png",
+            filename=f"colorized_{output_filename}"
+        )
+    except Exception as e:
+        logger.error("Error colorizing image: %s", str(e))
+        raise HTTPException(status_code=500, detail=f"Error colorizing image: {str(e)}")
+
+# ==========================================================
+# Gradio Interface (for Space UI)
+# ==========================================================
+def gradio_colorize(image):
+    """Gradio colorization function"""
+    if image is None:
+        return None
+    try:
+        if learn is None:
+            return None
+        return colorize_pil(image)
+    except Exception as e:
+        logger.error("Gradio colorization error: %s", str(e))
+        return None
+
+title = "🎨 FastAI GAN Image Colorizer"
+description = "Upload a black & white photo to generate a colorized version using the FastAI GAN model."
+
+iface = gr.Interface(
+    fn=gradio_colorize,
+    inputs=gr.Image(type="pil", label="Upload B&W Image"),
+    outputs=gr.Image(type="pil", label="Colorized Image"),
+    title=title,
+    description=description,
+)
+
+# Mount Gradio app at root (this will be the Space UI)
+# Note: This will override the root endpoint, so use /api for API info
+app = gr.mount_gradio_app(app, iface, path="/")
+
+# ==========================================================
+# Run Server
+# ==========================================================
+if __name__ == "__main__":
+    port = int(os.getenv("PORT", "7860"))
+    uvicorn.run(app, host="0.0.0.0", port=port)
+

requirements.txt

@@ -16,4 +16,5 @@ huggingface-hub>=0.16.0
 safetensors>=0.3.0
 fastai>=2.7.13
 toml>=0.10.2
+gradio>=4.0.0
 

test.py

@@ -33,9 +33,10 @@ def wait_for_model(base_url: str, timeout_seconds: int = 300) -> None:
 	health_url = f"{base_url}/health"
 	logging.info("Waiting for model to load at %s", health_url)
 	last_status = None
+	headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
 	while time.time() < deadline:
 		try:
-			resp = requests.get(health_url, timeout=15)
+			resp = requests.get(health_url, headers=headers, timeout=15)
 			if resp.ok:
 				data = resp.json()
 				last_status = data
@@ -53,7 +54,7 @@ def wait_for_model(base_url: str, timeout_seconds: int = 300) -> None:
 
 def upload_image(base_url: str, image_path: str, auth_bearer: Optional[str], app_check: Optional[str]) -> dict:
 	url = f"{base_url}/upload"
-	headers = {}
+	headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
 	if auth_bearer:
 		headers["Authorization"] = f"Bearer {auth_bearer}"
 	if app_check:
@@ -70,7 +71,7 @@ def upload_image(base_url: str, image_path: str, auth_bearer: Optional[str], app
 
 def colorize_image(base_url: str, image_path: str, auth_bearer: Optional[str], app_check: Optional[str]) -> dict:
 	url = f"{base_url}/colorize"
-	headers = {}
+	headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
 	if auth_bearer:
 		headers["Authorization"] = f"Bearer {auth_bearer}"
 	if app_check:
@@ -87,7 +88,7 @@ def colorize_image(base_url: str, image_path: str, auth_bearer: Optional[str], a
 
 def download_result(base_url: str, result_id: str, output_path: str, auth_bearer: Optional[str], app_check: Optional[str]) -> None:
 	url = f"{base_url}/download/{result_id}"
-	headers = {}
+	headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
 	if auth_bearer:
 		headers["Authorization"] = f"Bearer {auth_bearer}"
 	if app_check: