update

configs/i23d_args.json

@@ -30,7 +30,7 @@
   "log_interval": 50,
   "eval_interval": 5000,
   "save_interval": 10000,
-  "resume_checkpoint": "/nas/shared/V2V/yslan/logs/nips24/LSGM/t23d/FM/9cls/i23d/dit-L2-pixart-lognorm-rmsnorm-layernorm_before_pooled/gpu7-batch40-lr1e-4-bf16-qknorm-ctd3/model_joint_denoise_rec_model2990000.pt",
+  "resume_checkpoint": "checkpoints/objaverse/objaverse-dit/i23d/model_joint_denoise_rec_model2990000.safetensors",
   "resume_cldm_checkpoint": "",
   "resume_checkpoint_EG3D": "",
   "use_fp16": false,

nsr/train_util_diffusion.py

@@ -32,6 +32,8 @@ from guided_diffusion.train_util import (TrainLoop, calc_average_loss,
                                          parse_resume_step_from_filename)
 
 import dnnlib
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
 
 from nsr.camera_utils import FOV_to_intrinsics, LookAtPoseSampler
 
@@ -758,25 +760,40 @@ class TrainLoopDiffusionWithRec(TrainLoop):
                                   model=None,
                                   model_name='ddpm',
                                   resume_checkpoint=None):
-        if resume_checkpoint is None:
-            resume_checkpoint, self.resume_step = find_resume_checkpoint(
-                self.resume_checkpoint, model_name) or self.resume_checkpoint
+        # load safetensors from hf
+
+        hf_loading  = '.safetensors' in self.resume_checkpoint
+        if not hf_loading:
+            if resume_checkpoint is None:
+                resume_checkpoint, self.resume_step = find_resume_checkpoint(
+                    self.resume_checkpoint, model_name) or self.resume_checkpoint
 
         if model is None:
             model = self.model
 
-        if resume_checkpoint and Path(resume_checkpoint).exists():
+        if hf_loading or (resume_checkpoint and Path(resume_checkpoint).exists()):
             if dist_util.get_rank() == 0:
                 # ! rank 0 return will cause all other ranks to hang
-                logger.log(
-                    f"loading model from checkpoint: {resume_checkpoint}...")
                 map_location = {
                     'cuda:%d' % 0: 'cuda:%d' % dist_util.get_rank()
                 }  # configure map_location properly
-
                 logger.log(f'mark {model_name} loading ')
-                resume_state_dict = dist_util.load_state_dict(
-                    resume_checkpoint, map_location=map_location)
+
+                if hf_loading:
+                    logger.log(
+                        f"loading model from huggingface: yslan/LN3Diff/{self.resume_checkpoint}...")
+                else:
+                    logger.log(
+                        f"loading model from checkpoint: {resume_checkpoint}...")
+                
+                if hf_loading:
+                    model_path = hf_hub_download(repo_id="yslan/LN3Diff", 
+                            filename=self.resume_checkpoint)
+                    resume_state_dict = load_file(model_path)
+                else:
+                    resume_state_dict = dist_util.load_state_dict(
+                        resume_checkpoint, map_location=map_location)
+
                 logger.log(f'mark {model_name} loading finished')
 
                 model_state_dict = model.state_dict()