Deploy: HF cache dir, img2img fallback, auth bypass, root route

app/colorize_model.py

@@ -2,10 +2,11 @@
 ColorizeNet model wrapper for image colorization
 """
 import logging
+import os
 import torch
 import numpy as np
 from PIL import Image
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, StableDiffusionXLControlNetPipeline
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, StableDiffusionXLControlNetPipeline, StableDiffusionImg2ImgPipeline
 from diffusers.utils import load_image
 from transformers import pipeline
 from huggingface_hub import hf_hub_download
@@ -29,60 +30,86 @@ class ColorizeModel:
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         logger.info("Using device: %s", self.device)
         self.dtype = torch.float16 if self.device == "cuda" else torch.float32
+        self.hf_token = os.getenv("HF_TOKEN") or None
+
+        # Configure writable cache to avoid permission issues on Spaces
+        hf_cache_dir = os.getenv("HF_HOME", "./hf_cache")
+        os.environ.setdefault("HF_HOME", hf_cache_dir)
+        os.environ.setdefault("HUGGINGFACE_HUB_CACHE", hf_cache_dir)
+        os.environ.setdefault("TRANSFORMERS_CACHE", hf_cache_dir)
+        os.makedirs(hf_cache_dir, exist_ok=True)
+
+        # Avoid libgomp warning by setting a valid integer
+        os.environ.setdefault("OMP_NUM_THREADS", "1")
         
         try:
-            # Try loading as ControlNet with Stable Diffusion
-            logger.info("Attempting to load model as ControlNet: %s", self.model_id)
-            try:
-                # Load ControlNet model
-                self.controlnet = ControlNetModel.from_pretrained(
-                    self.model_id,
-                    torch_dtype=self.dtype
-                )
-                
-                # Try SDXL first, fallback to SD 1.5
+            # Decide whether to use ControlNet based on model_id
+            wants_controlnet = "control" in self.model_id.lower()
+
+            if wants_controlnet:
+                # Try loading as ControlNet with Stable Diffusion
+                logger.info("Attempting to load model as ControlNet: %s", self.model_id)
                 try:
-                    self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-                        "stabilityai/stable-diffusion-xl-base-1.0",
-                        controlnet=self.controlnet,
-                        torch_dtype=self.dtype,
-                        safety_checker=None,
-                        requires_safety_checker=False
-                    )
-                    logger.info("Loaded with SDXL base model")
-                except:
-                    self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
-                        "runwayml/stable-diffusion-v1-5",
-                        controlnet=self.controlnet,
+                    # Load ControlNet model
+                    self.controlnet = ControlNetModel.from_pretrained(
+                        self.model_id,
                         torch_dtype=self.dtype,
-                        safety_checker=None,
-                        requires_safety_checker=False
+                        token=self.hf_token,
+                        cache_dir=hf_cache_dir
                     )
-                    logger.info("Loaded with SD 1.5 base model")
-                
-                self.pipe.to(self.device)
-                
-                # Enable memory efficient attention if available
-                if hasattr(self.pipe, "enable_xformers_memory_efficient_attention"):
+                    
+                    # Try SDXL first, fallback to SD 1.5
                     try:
-                        self.pipe.enable_xformers_memory_efficient_attention()
-                        logger.info("XFormers memory efficient attention enabled")
-                    except Exception as e:
-                        logger.warning("Could not enable XFormers: %s", str(e))
-                
-                logger.info("ColorizeNet model loaded successfully as ControlNet")
-                self.model_type = "controlnet"
-                
-            except Exception as e:
-                logger.warning("Failed to load as ControlNet: %s", str(e))
-                # Fallback: try as image-to-image pipeline
+                        self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+                            "stabilityai/stable-diffusion-xl-base-1.0",
+                            controlnet=self.controlnet,
+                            torch_dtype=self.dtype,
+                            safety_checker=None,
+                            requires_safety_checker=False,
+                            token=self.hf_token,
+                            cache_dir=hf_cache_dir
+                        )
+                        logger.info("Loaded with SDXL base model")
+                    except Exception:
+                        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+                            "runwayml/stable-diffusion-v1-5",
+                            controlnet=self.controlnet,
+                            torch_dtype=self.dtype,
+                            safety_checker=None,
+                            requires_safety_checker=False,
+                            token=self.hf_token,
+                            cache_dir=hf_cache_dir
+                        )
+                        logger.info("Loaded with SD 1.5 base model")
+                    
+                    self.pipe.to(self.device)
+                    
+                    # Enable memory efficient attention if available
+                    if hasattr(self.pipe, "enable_xformers_memory_efficient_attention"):
+                        try:
+                            self.pipe.enable_xformers_memory_efficient_attention()
+                            logger.info("XFormers memory efficient attention enabled")
+                        except Exception as e:
+                            logger.warning("Could not enable XFormers: %s", str(e))
+                    
+                    logger.info("ColorizeNet model loaded successfully as ControlNet")
+                    self.model_type = "controlnet"
+                except Exception as e:
+                    logger.warning("Failed to load as ControlNet: %s", str(e))
+                    wants_controlnet = False  # fall through to pipeline
+
+            if not wants_controlnet:
+                # Load as image-to-image pipeline
                 logger.info("Trying to load as image-to-image pipeline...")
-                self.pipe = pipeline(
-                    "image-to-image",
-                    model=self.model_id,
-                    device=0 if self.device == "cuda" else -1,
-                    torch_dtype=self.dtype
-                )
+                self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+                    self.model_id,
+                    torch_dtype=self.dtype,
+                    safety_checker=None,
+                    requires_safety_checker=False,
+                    use_safetensors=True,
+                    cache_dir=hf_cache_dir,
+                    token=self.hf_token
+                ).to(self.device)
                 logger.info("ColorizeNet model loaded using image-to-image pipeline")
                 self.model_type = "pipeline"
                 

app/main.py

@@ -74,6 +74,16 @@ app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")
 # Initialize ColorizeNet model
 colorize_model = 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"""
@@ -109,6 +119,10 @@ async def verify_request(request: Request):
       - 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:

postman_collection.json

@@ -6,6 +6,65 @@
     "schema": "https://schema.getpostman.com/json/collection/v2.1.0/collection.json"
   },
   "item": [
+    {
+      "name": "Authentication",
+      "item": [
+        {
+          "name": "Login (Firebase Auth - email/password)",
+          "event": [
+            {
+              "listen": "test",
+              "script": {
+                "type": "text/javascript",
+                "exec": [
+                  "try {",
+                  "  const res = pm.response.json();",
+                  "  if (res.idToken) pm.collectionVariables.set('id_token', res.idToken);",
+                  "  if (res.refreshToken) pm.collectionVariables.set('refresh_token', res.refreshToken);",
+                  "  if (res.localId) pm.collectionVariables.set('local_id', res.localId);",
+                  "} catch (e) {",
+                  "  console.log('Failed to parse login response', e);",
+                  "}"
+                ]
+              }
+            }
+          ],
+          "request": {
+            "method": "POST",
+            "header": [
+              {
+                "key": "Content-Type",
+                "value": "application/json"
+              }
+            ],
+            "body": {
+              "mode": "raw",
+              "raw": "{\n  \"email\": \"{{email}}\",\n  \"password\": \"{{password}}\",\n  \"returnSecureToken\": true\n}"
+            },
+            "url": {
+              "raw": "https://identitytoolkit.googleapis.com/v1/accounts:signInWithPassword?key={{firebase_api_key}}",
+              "protocol": "https",
+              "host": [
+                "identitytoolkit",
+                "googleapis",
+                "com"
+              ],
+              "path": [
+                "v1",
+                "accounts:signInWithPassword"
+              ],
+              "query": [
+                {
+                  "key": "key",
+                  "value": "{{firebase_api_key}}"
+                }
+              ]
+            },
+            "description": "Obtain Firebase Auth id_token using email/password. Stores id_token in collection variable {{id_token}}."
+          }
+        }
+      ]
+    },
     {
       "name": "Health",
       "request": {
@@ -27,6 +86,11 @@
       "request": {
         "method": "POST",
         "header": [
+          {
+            "key": "Authorization",
+            "value": "Bearer {{id_token}}",
+            "type": "text"
+          },
           {
             "key": "X-Firebase-AppCheck",
             "value": "{{app_check_token}}",
@@ -60,6 +124,11 @@
       "request": {
         "method": "POST",
         "header": [
+          {
+            "key": "Authorization",
+            "value": "Bearer {{id_token}}",
+            "type": "text"
+          },
           {
             "key": "X-Firebase-AppCheck",
             "value": "{{app_check_token}}",
@@ -93,6 +162,11 @@
       "request": {
         "method": "GET",
         "header": [
+          {
+            "key": "Authorization",
+            "value": "Bearer {{id_token}}",
+            "type": "text"
+          },
           {
             "key": "X-Firebase-AppCheck",
             "value": "{{app_check_token}}",
@@ -155,6 +229,18 @@
       "key": "base_url",
       "value": "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space"
     },
+    {
+      "key": "firebase_api_key",
+      "value": "AIzaSyBIB6rcfyyqy5niERTXWvVD714Ter4Vx68"
+    },
+    {
+      "key": "email",
+      "value": "itisha.logico@gmail.com"
+    },
+    {
+      "key": "password",
+      "value": "123456"
+    },
     {
       "key": "app_check_token",
       "value": ""