new

pipeline/kiss3d_wrapper.py

@@ -116,7 +116,7 @@ def init_wrapper_from_config(config_path):
     # unet_ckpt_path = hf_hub_download(repo_id="LTT/Kiss3DGen", filename="flexgen_19w.ckpt", repo_type="model", token=access_token)
     unet_ckpt_path = hf_hub_download(repo_id="LTT/Kiss3DGen", filename="flexgen.ckpt", repo_type="model", token=access_token)
     if unet_ckpt_path is not None:
-        state_dict = torch.load(unet_ckpt_path, map_location='cpu')['state_dict']
+        state_dict = torch.load(unet_ckpt_path, map_location='cpu')
         # state_dict = {k[10:]: v for k, v in state_dict.items() if k.startswith('unet.unet.')}
         multiview_pipeline.unet.load_state_dict(state_dict, strict=True)