Update firefly.py (#295)

cminus01 · web-flow · commit 97e8e3c49104 · 2024-06-27T20:01:29.000-07:00
[fix bug]: Ensure state_dict is defined when ckpt_path is not None
diff --git a/fish_speech/models/vqgan/modules/firefly.py b/fish_speech/models/vqgan/modules/firefly.py
@@ -500,7 +500,7 @@ def __init__(self, ckpt_path: str = None, pretrained: bool = True):
         )
 
         if ckpt_path is not None:
-            self.load_state_dict(torch.load(ckpt_path, map_location="cpu"))
+            state_dict = torch.load(ckpt_path, map_location="cpu")
         elif pretrained:
             state_dict = torch.hub.load_state_dict_from_url(
                 "https://github.com/fishaudio/vocoder/releases/download/1.0.0/firefly-gan-base-generator.ckpt",

Original file line number	Diff line number	Diff line change
`@@ -500,7 +500,7 @@ def __init__(self, ckpt_path: str = None, pretrained: bool = True):`
`500`	`500`	`)`
`501`	`501`
`502`	`502`	`if ckpt_path is not None:`
`503`		`- self.load_state_dict(torch.load(ckpt_path, map_location="cpu"))`
	`503`	`+ state_dict = torch.load(ckpt_path, map_location="cpu")`
`504`	`504`	`elif pretrained:`
`505`	`505`	`state_dict = torch.hub.load_state_dict_from_url(`
`506`	`506`	`"https://github.com/fishaudio/vocoder/releases/download/1.0.0/firefly-gan-base-generator.ckpt",`