Upload backend/model_downloader.py with huggingface_hub

backend/model_downloader.py

@@ -0,0 +1,88 @@
+"""
+Automatic model downloader utility.
+Downloads models from Hugging Face if they don't exist locally.
+"""
+import os
+import sys
+from pathlib import Path
+from huggingface_hub import snapshot_download, hf_hub_download
+
+# Hugging Face repository
+HF_REPO = "Rasta02/dataku"
+
+def ensure_models_exist(model_paths_dict):
+    """
+    Check if models exist, download from HF if missing.
+    
+    Args:
+        model_paths_dict: Dictionary of model paths to check
+                         e.g., {'lama': '/path/to/lama', 'sttn': '/path/to/sttn.pth'}
+    """
+    print("[Model Checker] Checking for required models...")
+    
+    for model_name, model_path in model_paths_dict.items():
+        # Check if model exists
+        if os.path.exists(model_path):
+            if os.path.isdir(model_path):
+                # Check if directory has files
+                if len(os.listdir(model_path)) > 0:
+                    print(f"[✓] {model_name} model found at {model_path}")
+                    continue
+            else:
+                # File exists
+                print(f"[✓] {model_name} model found at {model_path}")
+                continue
+        
+        # Model doesn't exist, download from HF
+        print(f"[!] {model_name} model not found. Downloading from Hugging Face...")
+        try:
+            # Determine the relative path from backend directory
+            backend_dir = os.path.dirname(os.path.abspath(__file__))
+            rel_path = os.path.relpath(model_path, backend_dir)
+            
+            if os.path.isabs(model_path) or model_path.endswith('.pth') or model_path.endswith('.pt'):
+                # Download specific file
+                filename = os.path.basename(model_path)
+                repo_path = rel_path.replace(filename, '')
+                
+                hf_hub_download(
+                    repo_id=HF_REPO,
+                    filename=os.path.join(repo_path, filename).replace('\\', '/'),
+                    local_dir=backend_dir,
+                    local_dir_use_symlinks=False
+                )
+                print(f"[✓] Downloaded {model_name} model successfully!")
+            else:
+                # Download entire directory
+                os.makedirs(model_path, exist_ok=True)
+                snapshot_download(
+                    repo_id=HF_REPO,
+                    allow_patterns=f"{rel_path}/*",
+                    local_dir=backend_dir,
+                    local_dir_use_symlinks=False
+                )
+                print(f"[✓] Downloaded {model_name} model directory successfully!")
+        except Exception as e:
+            print(f"[✗] Failed to download {model_name} model: {e}")
+            print(f"[!] Please manually download models from https://huggingface.co/{HF_REPO}")
+            sys.exit(1)
+    
+    print("[Model Checker] All models are ready!")
+
+
+def download_backend_models():
+    """Download all backend models if they don't exist"""
+    import config
+    
+    model_paths = {
+        'LAMA': config.LAMA_MODEL_PATH,
+        'STTN': config.STTN_MODEL_PATH,
+        'ProPainter': config.VIDEO_INPAINT_MODEL_PATH,
+        'Detection': config.DET_MODEL_PATH,
+    }
+    
+    ensure_models_exist(model_paths)
+
+
+if __name__ == '__main__':
+    download_backend_models()