Преглед изворни кода

Optimize firefly auto download

Lengyue пре 1 година
родитељ
комит
2a7c886940
1 измењених фајлова са 3 додато и 2 уклоњено
  1. 3 2
      fish_speech/models/vqgan/modules/firefly.py

+ 3 - 2
fish_speech/models/vqgan/modules/firefly.py

@@ -502,9 +502,10 @@ class FireflyBase(nn.Module):
         if ckpt_path is not None:
             self.load_state_dict(torch.load(ckpt_path, map_location="cpu"))
         elif pretrained:
-            state_dict = torch.load(
-                "checkpoints/firefly-gan-base-generator.ckpt",
+            state_dict = torch.hub.load_state_dict_from_url(
+                "https://github.com/fishaudio/vocoder/releases/download/1.0.0/firefly-gan-base-generator.ckpt",
                 map_location="cpu",
+                model_dir="checkpoints",
             )
 
         if "state_dict" in state_dict: