Fix build error: Remove build-time model preloading and optimize for HF Spaces

- Remove build-time model preloading to avoid resource constraints
- Move model downloads to runtime for better HF Spaces compatibility
- Use SD 1.5 by default in HF Spaces for faster loading
- Update cache directories to use /tmp for better compatibility
- Add fallback logic detection for HF Spaces environment

Dockerfile

@@ -14,16 +14,14 @@ RUN pip install --no-cache-dir -r requirements.txt
 # Copy application code
 COPY . .
 
-# Set environment variables for HuggingFace caching
-ENV HF_HOME=/app/model_cache
-ENV TRANSFORMERS_CACHE=/app/model_cache
-ENV HF_DATASETS_CACHE=/app/model_cache
-
-# Create cache directory with proper permissions
-RUN mkdir -p /app/model_cache && chmod 755 /app/model_cache
-
-# Pre-download models during build time for faster startup
-RUN python preload_models.py
+# Set environment variables for HuggingFace caching (runtime download)
+ENV HF_HOME=/tmp/huggingface
+ENV TRANSFORMERS_CACHE=/tmp/transformers_cache
+ENV HF_DATASETS_CACHE=/tmp/huggingface_datasets
+ENV TORCH_HOME=/tmp/torch
+
+# Note: Models will be downloaded at runtime for better compatibility with Hugging Face Spaces
+# This avoids build-time resource constraints and allows for proper authentication
 
 EXPOSE 7860
 

app.py

@@ -12,10 +12,11 @@ import asyncio
 from typing import Optional
 from contextlib import asynccontextmanager
 
-# Set cache directories to writable locations
+# Set cache directories to writable locations (Hugging Face Spaces compatible)
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers_cache"
 os.environ["HF_HOME"] = "/tmp/huggingface"
 os.environ["TORCH_HOME"] = "/tmp/torch"
+os.environ["HF_DATASETS_CACHE"] = "/tmp/huggingface_datasets"
 
 import uvicorn
 from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks

models/image_generator.py

@@ -24,10 +24,19 @@ class ImageGenerator:
     def load_model(self):
         """Load the Stable Diffusion model with optimized caching"""
         try:
-            logger.info(f"🔄 Loading Stability AI model on {self.device}...")
+            logger.info(f"🔄 Loading Stable Diffusion model on {self.device}...")
             
-            # Use Stability AI's SDXL model for highest quality
-            model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+            # For Hugging Face Spaces, start with a more lightweight model
+            # Check if we're running in HF Spaces environment
+            is_hf_spaces = os.environ.get('SPACE_ID') is not None
+            
+            if is_hf_spaces:
+                # Use SD 1.5 for faster loading in HF Spaces
+                model_id = "runwayml/stable-diffusion-v1-5"
+                logger.info("🚀 Running in HF Spaces - using SD 1.5 for optimal performance")
+            else:
+                # Use SDXL for local/other deployments
+                model_id = "stabilityai/stable-diffusion-xl-base-1.0"
             
             # Optimize caching for faster subsequent loads
             cache_dir = os.environ.get("HF_HOME", "/tmp/huggingface_cache")
@@ -72,9 +81,9 @@ class ImageGenerator:
                 self.pipeline.enable_sequential_cpu_offload()
             
             if self.device.type == "cuda":
-                logger.info(f"✅ Stability AI SDXL loaded on GPU: {torch.cuda.get_device_name(0)}")
+                logger.info(f"✅ Stable Diffusion model loaded on GPU: {torch.cuda.get_device_name(0)}")
             else:
-                logger.info("✅ Stability AI SDXL loaded on CPU")
+                logger.info("✅ Stable Diffusion model loaded on CPU")
                 
         except Exception as e:
             logger.error(f"❌ Failed to load Stability AI model: {str(e)}")