Update app.py

app.py

@@ -46,9 +46,9 @@ os.environ['CUDA_LAUNCH_BLOCKING'] = '0'
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# ============================================================================
+# ============================================================================ #
 # GRADIO MONKEY PATCH (BUG FIX for gradio>=4.44.0)
-# ============================================================================
+# ============================================================================ #
 try:
     import gradio_client.utils as gc_utils
     original_get_type = gc_utils.get_type
@@ -65,58 +65,82 @@ try:
 except (ImportError, AttributeError) as e:
     logger.warning(f"⚠️ Could not apply Gradio monkey patch: {e}")
 
-# ============================================================================
+# ============================================================================ #
 
-# --------- Robust SAM2 loader for Hugging Face / local YAML configs ----------
-def load_sam2_predictor(device="cuda"):
+# --------- Robust SAM2 loader for Hugging Face / local YAML configs ---------- #
+def load_sam2_predictor(device="cuda", progress=gr.Progress()):
     """
     Loads the SAM2 model and returns a SAM2ImagePredictor instance.
     - Tries to load 'sam2_hiera_large' first.
     - Falls back to 'sam2_hiera_tiny' if large cannot be loaded.
-    - Assumes YAML configs are in ./configs/.
+    - Assumes YAML configs are in ./Configs/ (capital C), as required by upstream.
     """
     import hydra
     from omegaconf import OmegaConf
     import torch
-    import os
     import logging
 
     logger = logging.getLogger("SAM2Loader")
-    configs_dir = "./configs"
+    configs_dir = os.path.abspath("Configs")  # Capital C as in small app.py
+    logger.info(f"Looking for SAM2 configs in absolute path: {configs_dir}")
+
+    if not os.path.isdir(configs_dir):
+        logger.error(f"FATAL: Configs directory not found at '{configs_dir}'. Please ensure the 'Configs' folder is at the root of your repository.")
+        raise gr.Error(f"FATAL: SAM2 Configs directory not found. Check repository structure.")
+
     tried = []
 
     def try_load(config_name, checkpoint_name):
         try:
-            hydra.core.global_hydra.GlobalHydra.instance().clear()
-            hydra.initialize(config_path=configs_dir)
+            checkpoint_path = os.path.join("./checkpoints", checkpoint_name)
+            logger.info(f"Attempting to use checkpoint: {checkpoint_path}")
+
+            if not os.path.exists(checkpoint_path):
+                logger.info(f"Downloading {checkpoint_name} from Hugging Face Hub...")
+                progress(0.1, desc=f"Downloading {checkpoint_name}...")
+                from huggingface_hub import hf_hub_download
+                checkpoint_path = hf_hub_download(
+                    repo_id=f"facebook/{config_name.replace('.yaml', '')}",  # e.g. facebook/sam2_hiera_large
+                    filename=checkpoint_name,
+                    cache_dir="./checkpoints",  # Download to a local checkpoints dir
+                    local_dir_use_symlinks=False
+                )
+                logger.info(f"✅ Download complete: {checkpoint_path}")
+
+            if hydra.core.global_hydra.GlobalHydra.instance().is_initialized():
+                hydra.core.global_hydra.GlobalHydra.instance().clear()
+
+            hydra.initialize(config_path=os.path.relpath(configs_dir), job_name=f"sam2_load_{int(time.time())}")
             cfg = hydra.compose(config_name=config_name)
 
+            logger.info(f"Trying to load {config_name} on {device} with checkpoint {checkpoint_path}")
+            progress(0.3, desc=f"Loading {config_name}...")
+
             from sam2.build_sam import build_sam2
             from sam2.sam2_image_predictor import SAM2ImagePredictor
 
-            checkpoint_path = os.path.expanduser(f"~/.cache/sam2/{checkpoint_name}")
-            if not os.path.exists(checkpoint_path):
-                raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
-
-            logger.info(f"Trying {config_name} on {device} with checkpoint {checkpoint_path}")
-            sam2_model = build_sam2(cfg, checkpoint_path, device=device)
+            sam2_model = build_sam2(cfg.model, checkpoint_path)
+            sam2_model.to(device)
             predictor = SAM2ImagePredictor(sam2_model)
             logger.info(f"✅ Loaded {config_name} successfully on {device}")
             return predictor
         except Exception as e:
-            tried.append(f"{config_name}: {e}")
-            logger.warning(f"Failed to load {config_name}: {e}")
+            error_msg = f"Failed to load {config_name}: {e}\nTraceback: {traceback.format_exc()}"
+            tried.append(error_msg)
+            logger.warning(error_msg)
             return None
 
-    # Try large model first, then tiny
     predictor = try_load("sam2_hiera_large.yaml", "sam2_hiera_large.pt")
     if predictor is None:
+        logger.warning("Could not load large model, falling back to tiny model.")
         predictor = try_load("sam2_hiera_tiny.yaml", "sam2_hiera_tiny.pt")
         if predictor:
             logger.warning("⚠️ Using Tiny model as fallback (less accurate, but faster and lighter).")
+
     if predictor is None:
-        logger.error("❌ SAM2 loading failed for both large and tiny. Reasons: \n" + "\n".join(tried))
-        raise RuntimeError("SAM2 config/weights not found or could not be loaded. Make sure YAML and checkpoint files are present.")
+        error_message = "SAM2 loading failed for both large and tiny. Reasons: \n" + "\n".join(tried)
+        logger.error(f"❌ {error_message}")
+        raise gr.Error(error_message)
 
     return predictor
 # -------------------------------------------------------------------
@@ -128,7 +152,7 @@ models_loaded = False
 loading_lock = threading.Lock()
 
 # ------- Robust download_and_setup_models() using above loader --------
-def download_and_setup_models():
+def download_and_setup_models(progress=gr.Progress()):
     """
     Download and setup models (SAM2 and MatAnyone), robust to Hugging Face Spaces and local dev.
     Uses local YAML config, falls back to Tiny if Large can't be loaded.
@@ -143,7 +167,7 @@ def download_and_setup_models():
 
             # --- Load SAM2 ---
             device = "cuda" if torch.cuda.is_available() else "cpu"
-            sam2_predictor_local = load_sam2_predictor(device)
+            sam2_predictor_local = load_sam2_predictor(device, progress)
             sam2_predictor = sam2_predictor_local
 
             # --- Load MatAnyone (your original robust loader logic) ---
@@ -152,7 +176,7 @@ def download_and_setup_models():
             try:
                 from huggingface_hub import hf_hub_download
                 from matanyone import InferenceCore
-                matanyone_model_local = InferenceCore("PeiqingYang/MatAnyone")
+                matanyone_model_local = InferenceCore("PeiqingYang/MatAnyone-v1.0", device=device)
                 matanyone_loaded = True
                 logger.info("✅ MatAnyone loaded via HuggingFace Hub")
             except Exception as e:
@@ -164,6 +188,7 @@ def download_and_setup_models():
             matanyone_model = matanyone_model_local
 
             models_loaded = True
+            logger.info("--- ✅ All models loaded successfully ---")
             return "✅ SAM2 + MatAnyone loaded successfully!"
         except Exception as e:
             logger.error(f"❌ Enhanced loading failed: {str(e)}")
@@ -171,6 +196,15 @@ def download_and_setup_models():
             return f"❌ Enhanced loading failed: {str(e)}"
 # ------------------------------------------------------------------------------
 
+# [Now, everything below is unchanged from your previous full version.
+# If the code gets cut off due to token limits, just say "continue" and I will resume from the exact point!]
+
+# ... [Here comes the rest of your file: process_video_hq, create_interface, etc.]
+
+# =======================================================================
+# [START REST OF YOUR MAIN APP, UNCHANGED]
+# =======================================================================
+
 def process_video_hq(video_path, background_choice, custom_background_path, progress=gr.Progress()):
     """TWO-STAGE High-quality video processing: Original → Green Screen → Final Background"""
     if not models_loaded: