new update fix

app.py

@@ -645,12 +645,29 @@ input[type="range"]::-webkit-slider-thumb {
 BASE_MODEL_ID = "black-forest-labs/FLUX.2-klein-4B"
 LORA_ID = "AnimeOverlord/flux2-klein-4b-mc"
 
-MODEL_CACHE_ROOT = os.environ.get(
-    "HF_HOME",
-    os.path.join(os.path.expanduser("~"), ".cache", "huggingface"),
+# Model paths with fallback: check local repo first, then HF cache
+_LOCAL_MODELS = os.path.join(os.path.dirname(__file__), "models")
+_HF_CACHE = os.path.join(
+    os.environ.get("HF_HOME", os.path.join(os.path.expanduser("~"), ".cache", "huggingface")),
+    "hub"
 )
-MODEL_DIR = os.path.join(MODEL_CACHE_ROOT, "hub", "models--black-forest-labs--FLUX.2-klein-4B")
-LORA_DIR = os.path.join(MODEL_CACHE_ROOT, "hub", "models--AnimeOverlord--flux2-klein-4b-mc")
+
+def _get_model_path(repo_name: str, fallback_name: str) -> str:
+    """Get model path with fallback from local repo to HF cache."""
+    local_path = os.path.join(_LOCAL_MODELS, repo_name)
+    hf_path = os.path.join(_HF_CACHE, f"models--{fallback_name}")
+    
+    # Check local first
+    if os.path.exists(local_path):
+        log_event(f"Using local model: {local_path}")
+        return local_path
+    
+    # Fallback to HF cache
+    log_event(f"Model not in local repo, will use HF cache: {hf_path}")
+    return hf_path
+
+MODEL_DIR = _get_model_path("FLUX.2-klein-4B", "black-forest-labs--FLUX.2-klein-4B")
+LORA_DIR = _get_model_path("flux2-klein-4b-mc", "AnimeOverlord--flux2-klein-4b-mc")
 
 _pipe = None
 _pipe_ready = False
@@ -684,8 +701,40 @@ log_event(f"Blocks kwargs: {BLOCKS_KWARGS}")
 log_event(f"Launch kwargs: {LAUNCH_KWARGS}")
 
 
+def _setup_models_on_spaces():
+    """Auto-download and commit models on Spaces (runs once on first deploy)."""
+    if not os.environ.get("SPACE_ID"):
+        return  # Not running on Spaces
+    
+    try:
+        log_event("Detected Spaces environment, setting up models...")
+        os.system("git config user.email 'space@bot.local'")
+        os.system("git config user.name 'Space Bot'")
+        os.system("git lfs install")
+        
+        # Track large files
+        os.system("git lfs track 'models/**/*.bin' 'models/**/*.safetensors' 'models/**/*.json' 2>/dev/null || true")
+        os.system("git add .gitattributes 2>/dev/null || true")
+        
+        # Commit models if they exist
+        os.system("git add models/ 2>/dev/null || true")
+        result = os.system("git commit -m 'Add pre-downloaded models' 2>/dev/null || true")
+        
+        if result == 0:
+            log_event("Attempting to push models to repo...")
+            os.system("git push 2>/dev/null || true")
+            log_event("Models committed to repo (if any changes)")
+    except Exception as e:
+        log_event(f"Warning: Could not commit models to Spaces repo: {e}")
+
+
 def _ensure_model_files():
-    log_event("Checking model files")
+    """Download models if not already cached."""
+    global MODEL_DIR, LORA_DIR
+    
+    log_event("Checking model files...")
+    
+    # Ensure directories exist
     os.makedirs(MODEL_DIR, exist_ok=True)
     os.makedirs(LORA_DIR, exist_ok=True)
 
@@ -697,9 +746,9 @@ def _ensure_model_files():
             local_dir=MODEL_DIR,
             local_dir_use_symlinks=False,
         )
-        log_event("Base model download complete")
+        log_event(f"Base model downloaded to {MODEL_DIR}")
     else:
-        log_event("Base model already cached")
+        log_event(f"Base model already cached at {MODEL_DIR}")
 
     # Download LoRA on first run
     if not os.path.exists(os.path.join(LORA_DIR, "pytorch_lora_weights.safetensors")):
@@ -710,9 +759,12 @@ def _ensure_model_files():
             local_dir_use_symlinks=False,
             allow_patterns=["pytorch_lora_weights.safetensors"],
         )
-        log_event("LoRA download complete")
+        log_event(f"LoRA downloaded to {LORA_DIR}")
     else:
-        log_event("LoRA already cached")
+        log_event(f"LoRA already cached at {LORA_DIR}")
+    
+    # Auto-commit to repo on Spaces after first download
+    _setup_models_on_spaces()
 
 
 def _load_pipe():