load from checkpoint_path

app.py

@@ -30,7 +30,6 @@ from hf_utils.css_and_html import (
     get_header_html,
 )
 from hf_utils.visual_util import predictions_to_glb
-from mapanything.models import MapAnything
 from mapanything.utils.geometry import depthmap_to_world_frame, points_to_normals
 from mapanything.utils.image import load_images, rgb
 
@@ -53,7 +52,7 @@ print("Initializing and loading MapAnything model...")
 def load_hf_token():
     """Load HuggingFace access token from local file"""
     token_file_paths = [
-        "~/hf_token.txt",
+        "/home/aknapitsch/hf_token.txt",
     ]
 
     for token_path in token_file_paths:
@@ -66,6 +65,8 @@ def load_hf_token():
             except Exception as e:
                 print(f"Error reading token from {token_path}: {e}")
                 continue
+        else:
+            print(token_path, "token_path doesnt exist")
 
     # Also try environment variable
     # see https://huggingface.co/docs/hub/spaces-overview#managing-secrets on options
@@ -103,13 +104,14 @@ def init_hydra_config(config_path, overrides=None):
 # MapAnything Configuration
 high_level_config = {
     "path": "configs/train.yaml",
-    "hf_model_name": "facebook/map-anything",
+    "hf_model_name": "facebook/map-anything-apache",
     "config_overrides": [
         "machine=aws",
         "model=mapanything",
         "model/task=images_only",
         "model.encoder.uses_torch_hub=false",
     ],
+    "checkpoint_path": "https://huggingface.co/facebook/MapAnything/resolve/main/mapa_curri_24v_13d_48ipg_64g.pth",
     "trained_with_amp": True,
     "trained_with_amp_dtype": "fp16",
     "data_norm_type": "dinov2",
@@ -130,6 +132,8 @@ def run_model(target_dir, model_placeholder, apply_mask=True, mask_edges=True):
     Run the MapAnything model on images in the 'target_dir/images' folder and return predictions.
     """
     global model
+    import torch  # Ensure torch is available in function scope
+
     print(f"Processing images from {target_dir}")
 
     # Device check
@@ -140,11 +144,76 @@ def run_model(target_dir, model_placeholder, apply_mask=True, mask_edges=True):
     if model is None:
         print("Initializing MapAnything model...")
 
-        print("Loading CC-BY-NC 4.0 licensed MapAnything model...")
-        model = MapAnything.from_pretrained(high_level_config["hf_model_name"]).to(
-            device
+        # Initialize Hydra config and create model from configuration
+        cfg = init_hydra_config(
+            high_level_config["path"], overrides=high_level_config["config_overrides"]
         )
 
+        print("Loading MapAnything model...")
+        # Create model from local configuration instead of using from_pretrained
+        from mapanything.models import init_model
+
+        model = init_model(
+            model_str=cfg.model.model_str,
+            model_config=cfg.model.model_config,
+            torch_hub_force_reload=high_level_config.get(
+                "torch_hub_force_reload", False
+            ),
+        )
+
+        # Load the pretrained weights from HuggingFace Hub
+        try:
+            from huggingface_hub import hf_hub_download, list_repo_files
+
+            # First, let's see what files are available in the repository
+            try:
+                repo_files = list_repo_files(
+                    repo_id=high_level_config["hf_model_name"], token=load_hf_token()
+                )
+                print(f"Available files in repository: {repo_files}")
+
+                checkpoint_filename = "model.safetensors"
+
+                # Download the model weights
+                checkpoint_path = hf_hub_download(
+                    repo_id=high_level_config["hf_model_name"],
+                    filename=checkpoint_filename,
+                    token=load_hf_token(),
+                )
+
+                # Load the weights
+                print("start loading checkpoint")
+                if checkpoint_filename.endswith(".safetensors"):
+                    from safetensors.torch import load_file
+
+                    checkpoint = load_file(checkpoint_path)
+                else:
+                    checkpoint = torch.load(
+                        checkpoint_path, map_location="cpu", weights_only=True
+                    )
+
+                print("start loading state_dict")
+                if "model" in checkpoint:
+                    model.load_state_dict(checkpoint["model"])
+                elif "state_dict" in checkpoint:
+                    model.load_state_dict(checkpoint["state_dict"])
+                else:
+                    model.load_state_dict(checkpoint)
+
+                print(
+                    f"Successfully loaded pretrained weights from HuggingFace Hub ({checkpoint_filename})"
+                )
+
+            except Exception as inner_e:
+                print(f"Error listing repository files or loading weights: {inner_e}")
+                raise inner_e
+
+        except Exception as e:
+            print(f"Warning: Could not load pretrained weights: {e}")
+            print("Proceeding with randomly initialized model...")
+
+        model = model.to(device)
+
     else:
         model = model.to(device)