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

<enhance>: support using float16 in inference to speed up

Browse Source
This commit is contained in:
czk32611
2024-04-27 14:26:50 +08:00
parent 2c52de01b4
commit 865a68c60e
6 changed files with 103 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
musetalk/models/unet.py
View File

@@ -37,11 +37,11 @@ class UNet():
self.model = UNet2DConditionModel(**unet_config)
self.pe = PositionalEncoding(d_model=384)
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)
weights = torch.load(model_path) if torch.cuda.is_available() else torch.load(model_path, map_location=self.device)
self.model.load_state_dict(weights)
if use_float16:
self.model = self.model.half()
self.model.to(self.device)
if __name__ == "__main__":
unet = UNet()
unet = UNet()