Update app.py

app.py

@@ -40,7 +40,6 @@ sys.modules['sdib.utils.utils'] = mock
 
 # Configuration
 PRUNING_RATIOS = [10, 15, 20]
-LORA_CHECKPOINT_STEP = os.getenv("LORA_CHECKPOINT_STEP")
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
@@ -62,8 +61,8 @@ for ratio in PRUNING_RATIOS:
     try:
         print(f"Loading {ratio}% pruned model...")
         model_file = hf_hub_download(
-            repo_id="LWZ19/flux_dev_prune",
-            filename=f"pruned_model_{ratio}.pkl"
+            repo_id="LWZ19/flux_prune",
+            filename=f"dev/pruned_model_{ratio}.pkl"
         )
         
         with open(model_file, "rb") as f:
@@ -80,10 +79,10 @@ for ratio in PRUNING_RATIOS:
 print("📥 Preloading LoRA checkpoint for 20% pruning ratio...")
 try:
     lora_repo_path = snapshot_download(
-        repo_id="LWZ19/flux_dev_20_ckp_2",
-        allow_patterns=[f"lora/checkpoint-{LORA_CHECKPOINT_STEP}/*"]
+        repo_id="LWZ19/flux_retrain_weights",
+        allow_patterns=[f"dev/lora/prune_20/*"]
     )
-    lora_weights = load_file(os.path.join(lora_repo_path, "lora", f"checkpoint-{LORA_CHECKPOINT_STEP}", LORA_WEIGHT_NAME_SAFE))
+    lora_weights = load_file(os.path.join(lora_repo_path, "dev", "lora", "prune_20", LORA_WEIGHT_NAME_SAFE))
     print("✅ LoRA checkpoint loaded!")
 except Exception as e:
     print(f"❌ Failed to load LoRA checkpoint: {e}")