add huggingface_load

app.py

@@ -164,7 +164,7 @@ def save_video(frames, path="output.mp4", fps=10):
 def run(cfg: DictConfig):
 
     algo = run_local(cfg)
-    algo.to("cuda:0")
+    algo.to(device)
     
     actions = torch.zeros((1, 25))
     poses = torch.zeros((1, 5))
@@ -247,32 +247,31 @@ def run(cfg: DictConfig):
         gr.Markdown(
             """
             # WORLDMEM: Long-term Consistent World Generation with Memory
+            """
+            )
 
-            <div style="text-align: center;">
-            <!-- Public Website -->
-            <a style="display:inline-block" href="https://nirvanalan.github.io/projects/GA/">
-                <img src="https://img.shields.io/badge/public_website-8A2BE2">
-            </a>
+            # <div style="text-align: center;">
+            # <!-- Public Website -->
+            # <a style="display:inline-block" href="https://nirvanalan.github.io/projects/GA/">
+            #     <img src="https://img.shields.io/badge/public_website-8A2BE2">
+            # </a>
 
-            <!-- GitHub Stars -->
-            <a style="display:inline-block; margin-left: .5em" href="https://github.com/NIRVANALAN/GaussianAnything">
-                <img src="https://img.shields.io/github/stars/NIRVANALAN/GaussianAnything?style=social">
-            </a>
+            # <!-- GitHub Stars -->
+            # <a style="display:inline-block; margin-left: .5em" href="https://github.com/NIRVANALAN/GaussianAnything">
+            #     <img src="https://img.shields.io/github/stars/NIRVANALAN/GaussianAnything?style=social">
+            # </a>
 
-            <!-- Project Page -->
-            <a style="display:inline-block; margin-left: .5em" href="https://nirvanalan.github.io/projects/GA/">
-                <img src="https://img.shields.io/badge/project_page-blue">
-            </a>
+            # <!-- Project Page -->
+            # <a style="display:inline-block; margin-left: .5em" href="https://nirvanalan.github.io/projects/GA/">
+            #     <img src="https://img.shields.io/badge/project_page-blue">
+            # </a>
 
-            <!-- arXiv Paper -->
-            <a style="display:inline-block; margin-left: .5em" href="https://arxiv.org/abs/XXXX.XXXXX">
-                <img src="https://img.shields.io/badge/arXiv-paper-red">
-            </a>
-            </div>
+            # <!-- arXiv Paper -->
+            # <a style="display:inline-block; margin-left: .5em" href="https://arxiv.org/abs/XXXX.XXXXX">
+            #     <img src="https://img.shields.io/badge/arXiv-paper-red">
+            # </a>
+            # </div>
 
-            """
-            )
-        
         with gr.Row(variant="panel"):
             video_display = gr.Video(autoplay=True, loop=True)
             image_display = gr.Image(value=DEFAULT_IMAGE, interactive=False, label="Last Frame")
@@ -289,7 +288,7 @@ def run(cfg: DictConfig):
         sampling_timesteps_state = gr.State(algo.sampling_timesteps)
         
         example_actions = ["DDDDDDDDEEEEEEEEEESSSAAAAAAAAWWW", "DDDDDDDDDDDDQQQQQQQQQQQQQQQDDDDDDDDDDDD", 
-        "DDDDWWWDDDDDDDDDDDDDDDDDDDDSSSAAAAAAAAAAAAAAAAAAAAAAAA", "SSUNNWWEEEEEEEEEAAASSUNNWWEEEEEEEEEAAAAAAAAAAAAAAAAAAAAAA"]
+        "DDDDWWWDDDDDDDDDDDDDDDDDDDDSSSAAAAAAAAAAAAAAAAAAAAAAAA", "SSUNNWWEEEEEEEEEAAA1NNNNNNNNNSSUNNWW"]
     
         def set_action(action):
             return action

configurations/huggingface.yaml

@@ -30,9 +30,9 @@ experiment:
   load_t_to_r: false
   zero_init_gate: false
   only_tune_refer: false
-  diffusion_path: checkpoints/diffusion_only.ckpt
-  vae_path: checkpoints/vae_only.ckpt
-  pose_predictor_path: checkpoints/pose_prediction_model_only.ckpt
+  diffusion_path: yslan/worldmem_checkpoints/diffusion_only.ckpt
+  vae_path: yslan/worldmem_checkpoints/vae_only.ckpt
+  pose_predictor_path: yslan/worldmem_checkpoints/pose_prediction_model_only.ckpt
   customized_load: true
 
 algorithm:

experiments/exp_base.py

@@ -26,7 +26,7 @@ from utils.print_utils import cyan
 from utils.distributed_utils import is_rank_zero
 from safetensors.torch import load_model
 from pathlib import Path
-
+from huggingface_hub import hf_hub_download
 
 torch.set_float32_matmul_precision("high")
 
@@ -38,7 +38,16 @@ def load_custom_checkpoint(algo, optimizer, checkpoint_path):
     if not isinstance(checkpoint_path, Path):
         checkpoint_path = Path(checkpoint_path)
 
-    if checkpoint_path.suffix == ".pt":
+    if  "yslan" in str(checkpoint_path):
+        hf_ckpt = str(checkpoint_path).split('/')
+        repo_id = '/'.join(hf_ckpt[:2])
+        file_name = '/'.join(hf_ckpt[2:])
+        model_path = hf_hub_download(repo_id=repo_id, 
+                            filename=file_name)
+        ckpt = torch.load(model_path, map_location=torch.device('cpu'))
+        algo.load_state_dict(ckpt['state_dict'], strict=False)
+
+    elif checkpoint_path.suffix == ".pt":
         ckpt = torch.load(checkpoint_path, weights_only=True)
         algo.load_state_dict(ckpt, strict=False)
     elif checkpoint_path.suffix == ".ckpt":