Make it run locally with new sdxl-base model

.env

@@ -0,0 +1,7 @@
+# Hugging Face Configuration
+# Point Hugging Face cache to external SSD using official environment variables
+# This will create /Volumes/extssd/huggingface/hub and /Volumes/extssd/huggingface/datasets
+HF_HOME=/Volumes/extssd/huggingface
+
+# Model configuration (keep existing values)
+MODEL_ID=stabilityai/stable-diffusion-xl-base-1.0

Makefile

@@ -38,7 +38,17 @@ clean: ## Clean up cache files
 	find . -type f -name "*.pyc" -delete
 	find . -type f -name "*.pyo" -delete
 
-dev: ## Start development server
+dev: ## Start development server (with external SSD cache)
+	./run.sh
+
+dev-internal: ## Start development server (with internal cache)
 	uv run python app.py
 
+test-cache: ## Test external SSD cache configuration
+	@if [ -d "/Volumes/extssd" ]; then \
+		export HF_HOME="/Volumes/extssd/huggingface" && uv run python test_external_cache.py; \
+	else \
+		echo "❌ External SSD not found at /Volumes/extssd"; \
+	fi
+
 all: install setup quality test ## Run complete setup and checks

app.py

@@ -8,6 +8,10 @@ import sys
 import traceback
 
 import gradio as gr
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
 
 # Add src directory to path for imports
 src_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "src")

download_models.sh

@@ -5,6 +5,17 @@
 
 echo "🌸 Downloading Flowerfy models using Hugging Face CLI..."
 
+# Configure Hugging Face to use external SSD cache
+if [ -d "/Volumes/extssd" ]; then
+    export HF_HOME="/Volumes/extssd/huggingface"
+    echo "✅ Using external SSD cache at: $HF_HOME"
+    # Ensure directory exists
+    mkdir -p "$HF_HOME/hub"
+else
+    echo "⚠️  External SSD not found at /Volumes/extssd"
+    echo "   Models will be downloaded to default cache: ~/.cache/huggingface/hub"
+fi
+
 # Check if huggingface-hub is installed
 if ! command -v hf &> /dev/null; then
     echo "📦 Installing huggingface-hub CLI..."

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "diffusers>=0.35.1",
     "gradio>=5.44.0",
     "pillow>=11.3.0",
+    "python-dotenv>=1.0.0",
     "scikit-learn>=1.7.1",
     "torch>=2.8.0",
     "torchvision>=0.23.0",

run.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+# Flowerfy Startup Script with External SSD Model Support
+# This script sets up Hugging Face cache directories and launches the app
+
+echo "🌸 Starting Flowerfy with external SSD model cache..."
+
+# Check if external SSD is mounted
+if [ ! -d "/Volumes/extssd" ]; then
+    echo "❌ Error: External SSD not found at /Volumes/extssd"
+    echo "Please ensure your external SSD is mounted and try again."
+    exit 1
+fi
+
+# Create cache directories if they don't exist
+mkdir -p /Volumes/extssd/huggingface/hub
+
+# Export Hugging Face environment variables (using official HF_HOME method)
+export HF_HOME="/Volumes/extssd/huggingface"
+
+echo "✅ Cache directories configured:"
+echo "   HF_HOME: $HF_HOME"
+echo "   Models will be cached at: $HF_HOME/hub"
+echo "   Datasets will be cached at: $HF_HOME/datasets"
+
+# Launch the application with hot reload
+echo "🚀 Launching Flowerfy with hot reload..."
+uv run gradio app.py

src/core/config.py

@@ -14,9 +14,18 @@ class AppConfig:
         
     def _setup_device(self):
         """Setup device configuration for PyTorch."""
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
-        self.clf_device = 0 if torch.cuda.is_available() else -1
+        if torch.cuda.is_available():
+            self.device = "cuda"
+            self.dtype = torch.float16
+            self.clf_device = 0
+        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+            self.device = "mps"
+            self.dtype = torch.float16
+            self.clf_device = 0
+        else:
+            self.device = "cpu"
+            self.dtype = torch.float32
+            self.clf_device = -1
         
     @property
     def is_cuda_available(self):

src/core/constants.py

@@ -2,6 +2,10 @@
 
 import os
 
+# External storage configuration
+# If using external SSD, models will be cached at /Volumes/extssd/huggingface/hub
+# This is configured via environment variables (see .env file and run.sh script)
+
 # Model configuration  
 DEFAULT_MODEL_ID = os.getenv("MODEL_ID", "stabilityai/stable-diffusion-xl-base-1.0")
 FALLBACK_MODEL_ID = "black-forest-labs/FLUX.1-schnell"

src/services/models/image_generation.py

@@ -3,6 +3,7 @@
 from typing import Optional
 
 import torch
+import numpy as np
 from diffusers import AutoPipelineForText2Image, FluxPipeline
 from PIL import Image
 
@@ -110,7 +111,21 @@ class ImageGenerationService:
                 max_sequence_length=512,  # FLUX parameter for text encoding
             )
 
-        return result.images[0]
+        # Validate and clean the image before returning
+        image = result.images[0]
+        
+        # Convert to numpy array to check for invalid values
+        img_array = np.array(image)
+        
+        # Check for NaN or inf values and replace them
+        if np.any(np.isnan(img_array)) or np.any(np.isinf(img_array)):
+            print("⚠️  Warning: Image contains invalid values (NaN/inf), cleaning...")
+            img_array = np.nan_to_num(img_array, nan=0.0, posinf=255.0, neginf=0.0)
+            # Ensure values are in valid range [0, 255]
+            img_array = np.clip(img_array, 0, 255).astype(np.uint8)
+            image = Image.fromarray(img_array)
+        
+        return image
         
     def get_model_info(self) -> str:
         """Get information about the currently loaded model."""

test_external_cache.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+"""Test script to verify external SSD cache configuration works correctly."""
+
+import os
+import sys
+from pathlib import Path
+
+def test_cache_configuration():
+    """Test that the external cache configuration is working."""
+    
+    print("🧪 Testing External SSD Cache Configuration")
+    print("=" * 50)
+    
+    # Check if external SSD is mounted
+    external_path = Path("/Volumes/extssd")
+    if not external_path.exists():
+        print("❌ External SSD not found at /Volumes/extssd")
+        return False
+    
+    print("✅ External SSD is mounted")
+    
+    # Check if HF_HOME is set correctly
+    hf_home = os.environ.get("HF_HOME")
+    expected_hf_home = "/Volumes/extssd/huggingface"
+    
+    if hf_home != expected_hf_home:
+        print(f"⚠️  HF_HOME not set correctly. Expected: {expected_hf_home}, Got: {hf_home}")
+        print("   Set HF_HOME with: export HF_HOME=/Volumes/extssd/huggingface")
+        return False
+    
+    print(f"✅ HF_HOME correctly set to: {hf_home}")
+    
+    # Check if cache directories exist
+    hub_cache = Path(hf_home) / "hub"
+    if not hub_cache.exists():
+        print(f"❌ Hub cache directory not found at: {hub_cache}")
+        return False
+    
+    print(f"✅ Hub cache directory exists at: {hub_cache}")
+    
+    # Check if models are present
+    model_count = len(list(hub_cache.glob("models--*")))
+    print(f"✅ Found {model_count} models in cache")
+    
+    # Test importing Hugging Face libraries and check their cache detection
+    try:
+        from huggingface_hub import HfFolder
+        from transformers import AutoTokenizer
+        from diffusers import DiffusionPipeline
+        
+        print("✅ Hugging Face libraries imported successfully")
+        
+        # Test a small model to verify cache is working
+        print("🔄 Testing cache with a small model (this may take a moment)...")
+        
+        # This should use the external cache
+        tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+        
+        print("✅ Successfully loaded model from cache")
+        
+        # Check if the model files are in the expected location
+        clip_path = hub_cache / "models--openai--clip-vit-base-patch32"
+        if clip_path.exists():
+            print(f"✅ Model files found in external cache at: {clip_path}")
+        else:
+            print(f"⚠️  Model files not found at expected location: {clip_path}")
+            return False
+        
+        return True
+        
+    except Exception as e:
+        print(f"❌ Error loading model: {e}")
+        return False
+
+def main():
+    """Main test function."""
+    # Load .env file if available
+    env_file = Path(__file__).parent / ".env"
+    if env_file.exists():
+        print("📁 Loading .env file...")
+        from dotenv import load_dotenv
+        load_dotenv()
+    else:
+        print("⚠️  No .env file found, using system environment variables")
+    
+    success = test_cache_configuration()
+    
+    print("\n" + "=" * 50)
+    if success:
+        print("🎉 All tests passed! External SSD cache is working correctly.")
+        print("You can now run the application with: ./run.sh")
+    else:
+        print("❌ Some tests failed. Please check the configuration.")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()

uv.lock

@@ -266,6 +266,7 @@ dependencies = [
     { name = "diffusers" },
     { name = "gradio" },
     { name = "pillow" },
+    { name = "python-dotenv" },
     { name = "scikit-learn" },
     { name = "torch" },
     { name = "torchvision" },
@@ -285,6 +286,7 @@ requires-dist = [
     { name = "diffusers", specifier = ">=0.35.1" },
     { name = "gradio", specifier = ">=5.44.0" },
     { name = "pillow", specifier = ">=11.3.0" },
+    { name = "python-dotenv", specifier = ">=1.0.0" },
     { name = "scikit-learn", specifier = ">=1.7.1" },
     { name = "torch", specifier = ">=2.8.0" },
     { name = "torchvision", specifier = ">=0.23.0" },
@@ -1032,6 +1034,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" },
 ]
 
+[[package]]
+name = "python-dotenv"
+version = "1.1.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/f6/b0/4bc07ccd3572a2f9df7e6782f52b0c6c90dcbb803ac4a167702d7d0dfe1e/python_dotenv-1.1.1.tar.gz", hash = "sha256:a8a6399716257f45be6a007360200409fce5cda2661e3dec71d23dc15f6189ab", size = 41978, upload-time = "2025-06-24T04:21:07.341Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/5f/ed/539768cf28c661b5b068d66d96a2f155c4971a5d55684a514c1a0e0dec2f/python_dotenv-1.1.1-py3-none-any.whl", hash = "sha256:31f23644fe2602f88ff55e1f5c79ba497e01224ee7737937930c448e4d0e24dc", size = 20556, upload-time = "2025-06-24T04:21:06.073Z" },
+]
+
 [[package]]
 name = "python-multipart"
 version = "0.0.20"