deepgeek/MuseTalk
2
0
Fork 0
mirror of https://github.com/TMElyralab/MuseTalk.git synced 2026-02-05 09:59:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix some cuda related issue when run on M-Series Mac

Browse Source
This commit is contained in:
=
2024-04-05 22:03:28 +08:00
parent bc1379abad
commit 9a5212c8dd
2 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
musetalk/models/unet.py
View File

@@ -36,11 +36,11 @@ class UNet():
unet_config = json.load(f)
self.model = UNet2DConditionModel(**unet_config)
self.pe = PositionalEncoding(d_model=384)
self.weights = torch.load(model_path)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.weights = torch.load(model_path) if torch.cuda.is_available() else torch.load(model_path, map_location=self.device)
self.model.load_state_dict(self.weights)
if use_float16:
self.model = self.model.half()
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
if __name__ == "__main__":