load lora weights

app.py

@@ -10,7 +10,7 @@ import torch
 import torch.nn.functional as F
 import gradio as gr
 from PIL import Image, ImageFilter, ImageChops
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, snapshot_download # <--- Added snapshot_download
 import spaces
 
 # --- IMPORT YOUR CUSTOM MODULES ---
@@ -28,6 +28,7 @@ SAM2_CONFIG = "sam2_hiera_l.yaml"
 BASE_CKPT_NAME = "sam2_hiera_large.pt"
 FINE_TUNED_SAM = "fine_tuned_sam2_batched_90000.torch"
 FINE_TUNED_PLM = "fine_tuned_sam2_batched_plm_90000.torch"
+FINE_TUNED_LORA = "lora_plm_adapter_90000" # Folder name in the HF Repo
 
 SQUARE_DIM = 1024
 
@@ -85,7 +86,7 @@ def ensure_models_loaded():
     sd = torch.load(sam_ckpt_path, map_location="cpu")
     model.load_state_dict(sd.get("model", sd), strict=True)
     
-    # 2. Load PLM Adapter
+    # 2. Load PLM Adapter Base
     plm_path = download_if_needed(REPO_ID, FINE_TUNED_PLM)
     plm = PLMLanguageAdapter(
         model_name="Qwen/Qwen2.5-VL-3B-Instruct", 
@@ -96,6 +97,21 @@ def ensure_models_loaded():
     )
     plm_sd = torch.load(plm_path, map_location="cpu")
     plm.load_state_dict(plm_sd["plm"], strict=True)
+    
+    # 3. Load LoRA Weights
+    try:
+        logging.info(f"Downloading LoRA folder: {FINE_TUNED_LORA}...")
+        # snapshot_download returns the root cache folder; we use allow_patterns to get just the LoRA folder
+        cache_root = snapshot_download(repo_id=REPO_ID, allow_patterns=f"{FINE_TUNED_LORA}/*")
+        
+        # Construct the full path to the directory containing the LoRA files
+        lora_dir_path = os.path.join(cache_root, FINE_TUNED_LORA)
+        
+        logging.info(f"Loading LoRA from {lora_dir_path}...")
+        plm.load_lora(lora_dir_path)
+    except Exception as e:
+        raise RuntimeError(f"Failed to load LoRA weights: {e}")
+
     plm.eval()
     
     MODEL_CACHE["sam"] = model