Spaces:

Wuvin
/

Unique3D

Running on Zero

Wuvin commited on May 31, 2024

Commit

8d1277c

1 Parent(s): 8cb0437

bug fix

Files changed (1) hide show

gradio_app/custom_models/utils.py CHANGED Viewed

@@ -59,7 +59,7 @@ def load_pipeline(config_path, ckpt_path, pipeline_filter=lambda x: True, weight
         shared_modules = trainer.init_shared_modules(shared_modules)
     if load_from_checkpoint is not None:
-        state_dict = torch.load(load_from_checkpoint)
         configurable_unet.unet.load_state_dict(state_dict, strict=False)
     # Move unet, vae and text_encoder to device and cast to weight_dtype
     configurable_unet.unet.to(device, dtype=weight_dtype)

         shared_modules = trainer.init_shared_modules(shared_modules)
     if load_from_checkpoint is not None:
+        state_dict = torch.load(load_from_checkpoint, map_location="cpu")
         configurable_unet.unet.load_state_dict(state_dict, strict=False)
     # Move unet, vae and text_encoder to device and cast to weight_dtype
     configurable_unet.unet.to(device, dtype=weight_dtype)