before multi-frame gen

app.py

@@ -27,13 +27,19 @@ import requests
 torch.set_float32_matmul_precision("high")
 
 def load_custom_checkpoint(algo, 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)
+    try:
+        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)
+        print("Load: ", model_path)
+    except:
+        ckpt = torch.load(checkpoint_path, map_location=torch.device('cpu'))
+        algo.load_state_dict(ckpt['state_dict'], strict=False)      
+        print("Load: ", checkpoint_path)  
 
 def download_assets_if_needed():
     ASSETS_URL_BASE = "https://huggingface.co/spaces/yslan/worldmem/resolve/main/assets/examples"

configurations/huggingface.yaml

@@ -53,6 +53,6 @@ noise_level: random_all
 causal: True
 x_shape: [3, 360, 640]
 context_frames: 1
-diffusion_path: yslan/worldmem_checkpoints/diffusion_only.ckpt
+diffusion_path: /mnt/xiaozeqi/worldmem_github/diffusion_only_2025-04-04_4-56-27_epoch1step560000.ckpt
 vae_path: yslan/worldmem_checkpoints/vae_only.ckpt
 pose_predictor_path: yslan/worldmem_checkpoints/pose_prediction_model_only.ckpt