new update fix

app.py

@@ -655,7 +655,11 @@ HF_HUB_CACHE = os.path.join(HF_HOME, "hub")
 # Store the downloaded model files here
 MODEL_ROOT = os.path.join(SPACE_DATA_DIR, "models")
 MODEL_DIR = os.path.join(MODEL_ROOT, "FLUX.2-klein-4B")
-LORA_DIR = os.path.join(MODEL_ROOT, "flux2-klein-4b-mc")
+import tempfile
+import shutil
+
+LORA_DIR = os.path.join(tempfile.gettempdir(), "flux2-klein-4b-mc-v2")
+LORA_CACHE_DIR = os.path.join(tempfile.gettempdir(), "hf-lora-cache")
 
 # Set these before any HF downloads happen
 os.environ["HF_HOME"] = HF_HOME
@@ -700,10 +704,8 @@ def _ensure_model_files():
     log_event("Checking model files...")
 
     os.makedirs(MODEL_DIR, exist_ok=True)
-    os.makedirs(LORA_DIR, exist_ok=True)
 
     base_marker = os.path.join(MODEL_DIR, "model_index.json")
-    lora_marker = os.path.join(LORA_DIR, "pytorch_lora_weights.safetensors")
 
     if not os.path.exists(base_marker):
         log_event(f"Downloading base model: {BASE_MODEL_ID}")
@@ -716,18 +718,21 @@ def _ensure_model_files():
     else:
         log_event(f"Base model already cached at {MODEL_DIR}")
 
-    if not os.path.exists(lora_marker):
-        log_event(f"Downloading LoRA: {LORA_ID}")
-        snapshot_download(
-            repo_id=LORA_ID,
-            local_dir=LORA_DIR,
-            local_dir_use_symlinks=False,
-            allow_patterns=["pytorch_lora_weights.safetensors"],
-        )
-        log_event(f"LoRA downloaded to {LORA_DIR}")
-    else:
-        log_event(f"LoRA already cached at {LORA_DIR}")
-
+    # Always refresh LoRA in temp storage
+    shutil.rmtree(LORA_DIR, ignore_errors=True)
+    shutil.rmtree(LORA_CACHE_DIR, ignore_errors=True)
+    os.makedirs(LORA_DIR, exist_ok=True)
+    os.makedirs(LORA_CACHE_DIR, exist_ok=True)
+
+    log_event(f"Downloading LoRA fresh to temp: {LORA_ID}")
+    snapshot_download(
+        repo_id=LORA_ID,
+        local_dir=LORA_DIR,
+        cache_dir=LORA_CACHE_DIR,
+        local_dir_use_symlinks=False,
+        allow_patterns=["pytorch_lora_weights.safetensors"],
+    )
+    log_event(f"LoRA downloaded to temp: {LORA_DIR}")
 
 def _load_pipe():
     global _pipe, _pipe_ready
@@ -757,6 +762,9 @@ def _load_pipe():
 
     pipe.set_progress_bar_config(disable=False)
 
+    shutil.rmtree(LORA_DIR, ignore_errors=True)
+    shutil.rmtree(LORA_CACHE_DIR, ignore_errors=True)
+
     _pipe = pipe
     _pipe_ready = True
     log_event("Pipeline ready")