Convert to runtime weight downloading - Remove LFS tracked weights, add download functionality

.gitattributes

@@ -1,2 +1 @@
-weights/** filter=lfs diff=lfs merge=lfs -text
 examples/** filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -8,4 +8,20 @@ flagged
 Synthetic4Relight
 **/__pycache__/
 vis_*/
-src/
+src/
+
+# Downloaded model weights (large files)
+weights/*/checkpoints/*.ckpt
+weights/*/checkpoints/*.bin
+weights/*/checkpoints/*.safetensors
+
+# HuggingFace cache directory
+hf_cache/
+.cache/
+
+# Virtual environment (if not already excluded)
+test_env/
+
+# Temporary files
+*.tmp
+*.temp

app.py

@@ -14,11 +14,72 @@ import rembg
 import sys
 from loguru import logger
 
+# Add imports for downloading weights
+from huggingface_hub import hf_hub_download
+import shutil
+
 _SAMPLE_TAB_ID_ = 0
 _HIGHRES_TAB_ID_ = 1
 _FOREGROUND_TAB_ID_ = 2
 
 
+def download_model_weights():
+    """Download model weights from the original repository if they don't exist."""
+    
+    # Original repository ID
+    repo_id = "LittleFrog/IntrinsicAnything"
+    
+    # Define the paths and files to download
+    model_configs = [
+        {
+            "local_dir": "weights/albedo",
+            "files": [
+                "weights/albedo/checkpoints/last.ckpt",
+                "weights/albedo/configs/project.yaml"
+            ]
+        },
+        {
+            "local_dir": "weights/specular", 
+            "files": [
+                "weights/specular/checkpoints/last.ckpt",
+                "weights/specular/configs/project.yaml"
+            ]
+        }
+    ]
+    
+    for config in model_configs:
+        local_dir = config["local_dir"]
+        
+        # Create directories if they don't exist
+        os.makedirs(f"{local_dir}/checkpoints", exist_ok=True)
+        os.makedirs(f"{local_dir}/configs", exist_ok=True)
+        
+        for file_path in config["files"]:
+            local_file_path = file_path
+            
+            # Check if file exists and is not a LFS pointer (> 1KB)
+            if not os.path.exists(local_file_path) or os.path.getsize(local_file_path) < 1024:
+                logger.info(f"Downloading {file_path} from HuggingFace...")
+                
+                try:
+                    # Download the file
+                    downloaded_file = hf_hub_download(
+                        repo_id=repo_id,
+                        filename=file_path,
+                        repo_type="space",
+                        cache_dir="./hf_cache"
+                    )
+                    
+                    # Copy to the expected location
+                    shutil.copy2(downloaded_file, local_file_path)
+                    logger.info(f"Successfully downloaded {file_path}")
+                    
+                except Exception as e:
+                    logger.error(f"Failed to download {file_path}: {e}")
+                    raise e
+            else:
+                logger.info(f"{local_file_path} already exists and appears to be valid")
+
 def set_loggers(level):
     logger.remove()
     logger.add(sys.stderr, level=level)
@@ -92,6 +153,11 @@ description = \
 
 set_loggers("INFO")
 device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Download model weights if needed
+logger.info("Checking and downloading model weights if needed...")
+download_model_weights()
+
 logger.info(f"Loading Models...")
 model_dict = {
     "Albedo": InferenceModel(ckpt_path="weights/albedo", 

weights/albedo/checkpoints/last.ckpt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a5fa7a1caa7e1e3818119cd9a2e8715ee7b86a77fa66447cc4b0767d8ab550f8
-size 15458840153

weights/specular/checkpoints/last.ckpt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:361b7c3824f4603c4137657ff2fc8a127e0972954d425b440516522f034de7a0
-size 15458847705