revert app.py changes

app.py

@@ -6,55 +6,53 @@ import contextlib
 import io
 import traceback
 import gradio as gr
-#import spaces
 from transformers import AutoModel, AutoTokenizer, AutoConfig, pipeline
 from huggingface_hub import snapshot_download
 import sys, pathlib
 import os
 os.environ["OPENAI_API_KEY"] = "test"
 
-print("=== VINE SPACE DEBUG START ===")
+# 1) Download the repo to a local cache dir
+repo_dir = snapshot_download(repo_id="KevinX-Penn28/testing", revision="main")
 
-os.environ.setdefault("OMP_NUM_THREADS", "1")
-os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+# 2) Register the snapshot as an importable package
+VINE_PACKAGE = "vine_remote_repo"
 
-print("Step 1: Loading Vine repo lazily...")
-VineConfig = VineModel = VinePipeline = None
+# Drop stale modules in case the script reloads
+for module_name in list(sys.modules):
+    if module_name == VINE_PACKAGE or module_name.startswith(f"{VINE_PACKAGE}."):
+        del sys.modules[module_name]
 
-def load_vine():
-    global VineConfig, VineModel, VinePipeline
-    if VineModel is not None:
-        return VineConfig, VineModel, VinePipeline
-    
-    print("Downloading repo...")
-    repo_dir = snapshot_download(repo_id="KevinX-Penn28/testing", revision="main")
-
-    print("Importing modules...")
-    package_spec = importlib.util.spec_from_file_location(
-        VINE_PACKAGE,
-        Path(repo_dir) / "__init__.py",
-        submodule_search_locations=[str(repo_dir)],
-    )
-    package_module = importlib.util.module_from_spec(package_spec)
-    sys.modules[VINE_PACKAGE] = package_module
-    package_spec.loader.exec_module(package_module)
+package_spec = importlib.util.spec_from_file_location(
+    VINE_PACKAGE,
+    Path(repo_dir) / "__init__.py",
+    submodule_search_locations=[str(repo_dir)],
+)
+if not package_spec or not package_spec.loader:
+    raise ImportError(f"Cannot create package spec for {VINE_PACKAGE} at {repo_dir}")
 
-    vine_config_module = importlib.import_module(f"{VINE_PACKAGE}.vine_config")
-    vine_model_module = importlib.import_module(f"{VINE_PACKAGE}.vine_model")
-    vine_pipeline_module = importlib.import_module(f"{VINE_PACKAGE}.vine_pipeline")
+package_module = importlib.util.module_from_spec(package_spec)
+sys.modules[VINE_PACKAGE] = package_module
+package_spec.loader.exec_module(package_module)
 
-    VineConfig = vine_config_module.VineConfig
-    VineModel = vine_model_module.VineModel
-    VinePipeline = vine_pipeline_module.VinePipeline
+# 3) Import and use via the registered package
+vine_config_module = importlib.import_module(f"{VINE_PACKAGE}.vine_config")
+vine_model_module = importlib.import_module(f"{VINE_PACKAGE}.vine_model")
+vine_pipeline_module = importlib.import_module(f"{VINE_PACKAGE}.vine_pipeline")
 
-    print("Vine loaded")
-    return VineConfig, VineModel, VinePipeline
+VineConfig = vine_config_module.VineConfig              # your config class
+VineModel = vine_model_module.VineModel                 # your model class
+VinePipeline = vine_pipeline_module.VinePipeline        
 
+current_dir = Path(__file__).parent
+sam_config_path = "//" + str(current_dir / "sam2_hiera_t.yaml")
+sam_checkpoint_path = "//" + str(current_dir / "sam2_hiera_tiny.pt")
+gd_config_path = "//" + str(current_dir / "GroundingDINO_SwinT_OGC.py")
+gd_checkpoint_path = "//" + str(current_dir / "groundingdino_swint_ogc.pth")
+visualization_dir = "//" + str(current_dir / "outputs")
 
-#@spaces.GPU(duration=200)
 def process_video(video_file, categorical_keywords, unary_keywords, binary_keywords, object_pairs, output_fps):
     log_buffer = io.StringIO()
-    VineConfig, VineModel, VinePipeline = load_vine()
     try:
         with contextlib.redirect_stdout(log_buffer), contextlib.redirect_stderr(log_buffer):
             categorical_keywords = [kw.strip() for kw in categorical_keywords.split(",")] if categorical_keywords else []
@@ -82,6 +80,7 @@ def process_video(video_file, categorical_keywords, unary_keywords, binary_keywo
                 visualization_dir=visualization_dir,
                 visualize=True,
                 debug_visualizations=False,
+                device = "cuda"
                 )
             model = VineModel(config)
             
@@ -93,6 +92,7 @@ def process_video(video_file, categorical_keywords, unary_keywords, binary_keywo
                 gd_config_path=gd_config_path,
                 gd_checkpoint_path=gd_checkpoint_path,
                 trust_remote_code=True,
+                device = "cuda"
                 )
             
             results = vine_pipe(
@@ -104,7 +104,8 @@ def process_video(video_file, categorical_keywords, unary_keywords, binary_keywo
                 segmentation_method="grounding_dino_sam2",        
                 return_top_k=5,
                 include_visualizations=True,
-                debug_visualizations=False
+                debug_visualizations=False,
+                device = 2
             )
             
             if isinstance(results, Mapping):