Fix permission errors and environment variables for model loading - Use /tmp directory for writable model storage - Set OMP_NUM_THREADS and cache directories at startup - Add better error handling and permission checks

app/main_sdxl.py

@@ -3,6 +3,15 @@ FastAPI application for Text-Guided Image Colorization using SDXL + ControlNet
 Based on fffiloni/text-guided-image-colorization
 """
 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
@@ -174,17 +183,39 @@ async def startup_event():
     try:
         logger.info("🔄 Loading SDXL + ControlNet colorization models...")
         
-        # Ensure required directories exist
-        os.makedirs("sdxl_light_caption_output", exist_ok=True)
-        
-        # Download controlnet model snapshot
+        # Use writable directory for model downloads
+        controlnet_dir = "/tmp/sdxl_light_caption_output"
         try:
-            snapshot_download(
-                repo_id='nickpai/sdxl_light_caption_output',
-                local_dir='sdxl_light_caption_output'
-            )
+            os.makedirs(controlnet_dir, exist_ok=True)
+            # Test write permissions
+            test_file = os.path.join(controlnet_dir, ".test_write")
+            with open(test_file, "w") as f:
+                f.write("test")
+            os.remove(test_file)
+            logger.info(f"Using directory: {controlnet_dir}")
+        except PermissionError as e:
+            logger.error(f"Permission denied for directory {controlnet_dir}: {e}")
+            raise
         except Exception as e:
-            logger.warning(f"Could not download controlnet snapshot: {e}")
+            logger.error(f"Failed to create directory {controlnet_dir}: {e}")
+            raise
+        
+        # Download controlnet model snapshot
+        controlnet_path = os.path.join(controlnet_dir, "checkpoint-30000", "controlnet")
+        if os.path.exists(controlnet_path):
+            logger.info(f"ControlNet model already exists at {controlnet_path}")
+        else:
+            try:
+                logger.info("Downloading ControlNet model...")
+                snapshot_download(
+                    repo_id='nickpai/sdxl_light_caption_output',
+                    local_dir=controlnet_dir
+                )
+                logger.info("ControlNet model downloaded successfully")
+            except Exception as e:
+                logger.error(f"Could not download controlnet snapshot: {e}")
+                if not os.path.exists(controlnet_path):
+                    raise
         
         # Device and precision setup
         accelerator = Accelerator(mixed_precision="fp16")
@@ -196,7 +227,7 @@ async def startup_event():
         # Pretrained paths
         base_model_path = settings.BASE_MODEL_ID
         safetensors_ckpt = settings.LIGHTNING_WEIGHTS
-        controlnet_path = "sdxl_light_caption_output/checkpoint-30000/controlnet"
+        # controlnet_path already defined above
         
         # Load diffusion components
         logger.info("Loading VAE...")