update app

app.py

@@ -3,6 +3,7 @@ import logging
 import hashlib
 import sys
 import traceback
+import copy
 
 import cv2
 import numpy as np
@@ -10,36 +11,50 @@ import torch
 import torch.nn.functional as F
 import gradio as gr
 from PIL import Image, ImageFilter, ImageChops, ImageDraw
-from huggingface_hub import hf_hub_download
-import spaces  # <--- NEW IMPORT
+from huggingface_hub import hf_hub_download  # <--- NEW IMPORT
+import spaces
 
 # --- IMPORT YOUR CUSTOM MODULES ---
-# Ensure the 'sam2' folder and 'plm_adapter_...' file are uploaded to your Space
 from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 from sam2.modeling.sam.mask_decoder import MaskDecoder
 from plm_adapter_lora_with_image_input_only_text_positions import PLMLanguageAdapter
 
 # ----------------- Configuration -----------------
-# UPDATE THESE TO MATCH YOUR HF REPO IF YOU STORE WEIGHTS THERE
 HF_REPO_ID = "aadarsh99/ConvSeg-Stage1" 
 SAM2_CONFIG = "sam2_hiera_l.yaml"
 
-# Checkpoint filenames (assumed to be in the root or downloaded)
+# Filenames
 BASE_CKPT_NAME = "sam2_hiera_large.pt"
-FINAL_CKPT_NAME = "fine_tuned_sam2_batched_100000.torch" # Update with your filename
-PLM_CKPT_NAME = "fine_tuned_sam2_batched_plm_100000.torch"      # Update with your filename
-LORA_CKPT_NAME = "lora_plm_adapter_100000"                   # Set filename if you use LoRA, else None
+FINAL_CKPT_NAME = "fine_tuned_sam2_batched_100000.torch"
+PLM_CKPT_NAME = "fine_tuned_sam2_batched_plm_100000.torch"
+LORA_CKPT_NAME = None
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 SQUARE_DIM = 1024
-
 logging.basicConfig(level=logging.INFO)
 
 # ----------------- Globals (Lazy Loading) -----------------
 MODEL_SAM = None
 PLM = None
 
+# ----------------- Helper: Download Logic -----------------
+def download_if_needed(filename):
+    """
+    Checks if file exists locally. If not, downloads from HF Repo.
+    Returns the valid path to the file.
+    """
+    if os.path.exists(filename):
+        logging.info(f"Found local file: {filename}")
+        return filename
+    
+    logging.info(f"{filename} not found locally. Downloading from {HF_REPO_ID}...")
+    try:
+        path = hf_hub_download(repo_id=HF_REPO_ID, filename=filename)
+        logging.info(f"Downloaded to: {path}")
+        return path
+    except Exception as e:
+        raise FileNotFoundError(f"Could not find {filename} locally or in HF repo {HF_REPO_ID}. Error: {e}")
+
 # ----------------- Overlay Style Helpers -----------------
 EDGE_COLORS_HEX = ["#3A86FF", "#FF006E", "#43AA8B", "#F3722C", "#8338EC", "#90BE6D"]
 
@@ -214,21 +229,14 @@ def load_models_lazy():
     print("Lazy loading models inside GPU context...")
     
     # 1. Base SAM2 Model
-    if not os.path.exists(BASE_CKPT_NAME):
-        raise FileNotFoundError(f"{BASE_CKPT_NAME} not found")
-        
-    # On ZeroGPU, we can load to 'cuda' directly, or 'cpu' then move.
-    # To be safe against the deepcopy error, we load to cpu then move.
-    # If the deepcopy error persists, we might need to load directly to 'cuda'.
-    # Let's try CPU load -> move to cuda.
+    base_path = download_if_needed(BASE_CKPT_NAME)
     
-    model = build_sam2(SAM2_CONFIG, BASE_CKPT_NAME, device="cpu")
+    # Init on CPU to avoid "deepcopy" errors, then move later
+    model = build_sam2(SAM2_CONFIG, base_path, device="cpu")
     
     # 2. Fine-tuned Weights
-    if not os.path.exists(FINAL_CKPT_NAME):
-        raise FileNotFoundError(f"{FINAL_CKPT_NAME} not found")
-
-    sd = torch.load(FINAL_CKPT_NAME, map_location="cpu")
+    final_path = download_if_needed(FINAL_CKPT_NAME)
+    sd = torch.load(final_path, map_location="cpu")
     model.load_state_dict(sd.get("model", sd), strict=True)
     
     # Move SAM to CUDA now
@@ -250,14 +258,13 @@ def load_models_lazy():
         device="cpu", # Init on CPU
     )
     
-    if not os.path.exists(PLM_CKPT_NAME):
-        raise FileNotFoundError(f"{PLM_CKPT_NAME} not found")
-
-    plm_sd = torch.load(PLM_CKPT_NAME, map_location="cpu")
+    plm_path = download_if_needed(PLM_CKPT_NAME)
+    plm_sd = torch.load(plm_path, map_location="cpu")
     plm.load_state_dict(plm_sd["plm"], strict=True)
 
-    if LORA_CKPT_NAME and os.path.exists(LORA_CKPT_NAME):
-        plm.load_lora(LORA_CKPT_NAME)
+    if LORA_CKPT_NAME:
+        lora_path = download_if_needed(LORA_CKPT_NAME)
+        plm.load_lora(lora_path)
     
     # Move PLM to CUDA
     plm.to("cuda")
@@ -268,7 +275,7 @@ def load_models_lazy():
     return MODEL_SAM, PLM
 
 
-@spaces.GPU(duration=120) # Increased duration for first-time load
+@spaces.GPU(duration=180) # Increased duration for download + load
 def run_prediction(image_pil, text_prompt):
     if image_pil is None or not text_prompt:
         return None, None, None
@@ -280,8 +287,6 @@ def run_prediction(image_pil, text_prompt):
         model_sam, plm = load_models_lazy()
         
         # 2. Instantiate Predictor
-        # We assume models are already on CUDA from load_models_lazy
-        # Just to be sure, we can call .to("cuda") again (cheap if already there)
         model_sam.to("cuda")
         plm.to("cuda")