Update shadow_generator.py

shadow_generator.py

@@ -213,13 +213,13 @@
 # def initialize_once():
 #     _log("Initializing assets and SSN model...")
 #     load_ssn_once()
-    
-import io
+ import io
 import sys
 import shutil
 import tempfile
 from pathlib import Path
 from typing import Optional, Dict, Any, Tuple
+import subprocess
 
 import gdown
 from PIL import Image, ImageFilter
@@ -233,8 +233,8 @@ TMP_DIR = Path(tempfile.gettempdir()) / "pixhtlab"
 REPO_DIR = TMP_DIR / "PixHtLab-Src"
 WEIGHT_FILE = TMP_DIR / "weights" / "human_baseline_all_21-July-04-52-AM.pt"
 
-# Replace with your Drive folder & file link
-REPO_DRIVE_FOLDER = "https://drive.google.com/drive/folders/1YzxVaxoOXwBrdB9XHyoOgWz8z4BpLQFl?usp=sharing"
+# Hugging Face repo and Google Drive weight links
+REPO_HF_URL = "https://huggingface.co/karthikeya1212/test"
 DRIVE_WEIGHT_FILE = "https://drive.google.com/uc?id=1XwR2krb473vc46E_XB9ET611XPD1i8SH"
 
 _MIN_EXPECTED_WEIGHT_BYTES = 50 * 1024 * 1024
@@ -242,6 +242,7 @@ _MIN_EXPECTED_WEIGHT_BYTES = 50 * 1024 * 1024
 # -----------------------------
 # Logging
 # -----------------------------
+
 def _log(*args):
     print("[shadow]", *args, flush=True)
 
@@ -251,31 +252,24 @@ def _print_image(img: Image.Image, message: str):
 # -----------------------------
 # Repo & weight handling
 # -----------------------------
+
 def _download_repo():
-    """Download the repo folder from Google Drive into /tmp."""
     if REPO_DIR.exists():
         _log("Repo already exists:", REPO_DIR)
         return
 
-    _log("Downloading repository from Google Drive folder...")
-    temp_dir = TMP_DIR / "repo_dl"
-    if temp_dir.exists():
-        shutil.rmtree(temp_dir)
-    temp_dir.mkdir(parents=True, exist_ok=True)
-
+    _log("Cloning repository from Hugging Face...")
+    TMP_DIR.mkdir(parents=True, exist_ok=True)
     try:
-        gdown.download_folder(REPO_DRIVE_FOLDER, output=str(temp_dir), quiet=False, use_cookies=False)
-        candidates = list(temp_dir.glob("*"))
-        if not candidates:
-            raise RuntimeError("No files downloaded from the repo folder.")
-        shutil.move(str(candidates[0]), str(REPO_DIR))
-        _log("Repo downloaded successfully to:", REPO_DIR)
+        subprocess.run(["git", "lfs", "install"], check=False)
+        subprocess.run(["git", "clone", REPO_HF_URL, str(REPO_DIR)], check=True)
+        _log("Repo cloned successfully to:", REPO_DIR)
     except Exception as e:
-        _log("❌ Failed to download repo:", repr(e))
+        _log("❌ Failed to clone repo:", repr(e))
         raise RuntimeError("Cannot proceed without repository.")
 
+
 def _download_weights():
-    """Download .pt weight file into /tmp/weights/."""
     if WEIGHT_FILE.exists() and WEIGHT_FILE.stat().st_size > _MIN_EXPECTED_WEIGHT_BYTES:
         _log("Weights already exist:", WEIGHT_FILE)
         return
@@ -288,6 +282,7 @@ def _download_weights():
         raise RuntimeError(f"Downloaded weight file is invalid: {WEIGHT_FILE}")
     _log("✅ Weight file downloaded:", WEIGHT_FILE)
 
+
 def _validate_assets():
     if not REPO_DIR.exists():
         raise RuntimeError("Repo folder missing.")
@@ -295,9 +290,21 @@ def _validate_assets():
         raise RuntimeError("Weight file missing or too small.")
     _log("Assets validated successfully.")
 
+# -----------------------------
+# Auto-search for SSN_Model.py
+# -----------------------------
+
+def _find_ssn_model_file():
+    candidates = list(REPO_DIR.rglob("SSN_Model.py"))
+    if not candidates:
+        raise RuntimeError("SSN_Model.py not found in repo.")
+    _log("Found SSN_Model.py:", candidates[0])
+    return candidates[0].parent
+
 # -----------------------------
 # SSN model wrapper
 # -----------------------------
+
 class SSNWrapper:
     def __init__(self):
         self.model = None
@@ -307,11 +314,11 @@ class SSNWrapper:
         _download_weights()
         _validate_assets()
 
-        # Add repo code to sys.path
-        sys.path.insert(0, str(REPO_DIR.resolve()))
+        ssn_model_dir = _find_ssn_model_file()
+        sys.path.insert(0, str(ssn_model_dir.resolve()))
 
         try:
-            from Demo.SSN.models.SSN_Model import SSN_Model  # loaded from repo
+            from SSN_Model import SSN_Model  # dynamic import from found directory
             self.model = SSN_Model()
             state = torch.load(str(WEIGHT_FILE), map_location=self.device)
 
@@ -378,6 +385,7 @@ class SSNWrapper:
 # -----------------------------
 # Singleton
 # -----------------------------
+
 _ssn_wrapper: Optional[SSNWrapper] = None
 
 def load_ssn_once() -> SSNWrapper:
@@ -389,6 +397,7 @@ def load_ssn_once() -> SSNWrapper:
 # -----------------------------
 # Shadow compositing
 # -----------------------------
+
 def _hex_to_rgb(color: str) -> Tuple[int, int, int]:
     c = color.strip().lstrip("#")
     if len(c) == 3:
@@ -424,6 +433,7 @@ def _composite_shadow_and_image(original_img: Image.Image, matte_img: Image.Imag
 # -----------------------------
 # Public API
 # -----------------------------
+
 def _apply_params_defaults(params: Optional[Dict[str, Any]]) -> Dict[str, Any]:
     defaults = dict(softness=28.0, opacity=0.7, color="#000000")
     merged = {**defaults, **(params or {})}