Add Jetson fp16 model serving assets

cosyvoice/cli/cosyvoice.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from contextlib import contextmanager
 import os
 import time
 from typing import Generator
@@ -24,6 +25,19 @@ from cosyvoice.utils.file_utils import logging
 from cosyvoice.utils.class_utils import get_model_type
 
 
+@contextmanager
+def _fp16_default_dtype(enabled):
+    if not enabled:
+        yield
+        return
+    previous_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(torch.float16)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(previous_dtype)
+
+
 class CosyVoice:
 
     def __init__(self, model_dir, load_jit=False, load_trt=False, fp16=False, trt_concurrent=1):
@@ -34,8 +48,9 @@ class CosyVoice:
         hyper_yaml_path = '{}/cosyvoice.yaml'.format(model_dir)
         if not os.path.exists(hyper_yaml_path):
             raise ValueError('{} not found!'.format(hyper_yaml_path))
-        with open(hyper_yaml_path, 'r') as f:
-            configs = load_hyperpyyaml(f)
+        with _fp16_default_dtype(fp16 and torch.cuda.is_available()):
+            with open(hyper_yaml_path, 'r') as f:
+                configs = load_hyperpyyaml(f)
         assert get_model_type(configs) == CosyVoiceModel, 'do not use {} for CosyVoice initialization!'.format(model_dir)
         self.frontend = CosyVoiceFrontEnd(configs['get_tokenizer'],
                                           configs['feat_extractor'],
@@ -146,8 +161,9 @@ class CosyVoice2(CosyVoice):
         hyper_yaml_path = '{}/cosyvoice2.yaml'.format(model_dir)
         if not os.path.exists(hyper_yaml_path):
             raise ValueError('{} not found!'.format(hyper_yaml_path))
-        with open(hyper_yaml_path, 'r') as f:
-            configs = load_hyperpyyaml(f, overrides={'qwen_pretrain_path': os.path.join(model_dir, 'CosyVoice-BlankEN')})
+        with _fp16_default_dtype(fp16 and torch.cuda.is_available()):
+            with open(hyper_yaml_path, 'r') as f:
+                configs = load_hyperpyyaml(f, overrides={'qwen_pretrain_path': os.path.join(model_dir, 'CosyVoice-BlankEN')})
         assert get_model_type(configs) == CosyVoice2Model, 'do not use {} for CosyVoice2 initialization!'.format(model_dir)
         self.frontend = CosyVoiceFrontEnd(configs['get_tokenizer'],
                                           configs['feat_extractor'],
@@ -196,8 +212,9 @@ class CosyVoice3(CosyVoice2):
         hyper_yaml_path = '{}/cosyvoice3.yaml'.format(model_dir)
         if not os.path.exists(hyper_yaml_path):
             raise ValueError('{} not found!'.format(hyper_yaml_path))
-        with open(hyper_yaml_path, 'r') as f:
-            configs = load_hyperpyyaml(f, overrides={'qwen_pretrain_path': os.path.join(model_dir, 'CosyVoice-BlankEN')})
+        with _fp16_default_dtype(fp16 and torch.cuda.is_available()):
+            with open(hyper_yaml_path, 'r') as f:
+                configs = load_hyperpyyaml(f, overrides={'qwen_pretrain_path': os.path.join(model_dir, 'CosyVoice-BlankEN')})
         assert get_model_type(configs) == CosyVoice3Model, 'do not use {} for CosyVoice3 initialization!'.format(model_dir)
         self.frontend = CosyVoiceFrontEnd(configs['get_tokenizer'],
                                           configs['feat_extractor'],

cosyvoice/cli/model.py

@@ -63,14 +63,54 @@ class CosyVoiceModel:
         self.silent_tokens = []
 
     def load(self, llm_model, flow_model, hift_model):
-        self.llm.load_state_dict(torch.load(llm_model, map_location=self.device, weights_only=True), strict=True)
-        self.llm.to(self.device).eval()
-        self.flow.load_state_dict(torch.load(flow_model, map_location=self.device, weights_only=True), strict=True)
-        self.flow.to(self.device).eval()
+        llm_model = self._select_checkpoint_path(llm_model)
+        flow_model = self._select_checkpoint_path(flow_model)
+        hift_model = self._select_checkpoint_path(hift_model)
+
+        self._load_module_checkpoint(self.llm, llm_model)
+        self._load_module_checkpoint(self.flow, flow_model)
+
         # in case hift_model is a hifigan model
-        hift_state_dict = {k.replace('generator.', ''): v for k, v in torch.load(hift_model, map_location=self.device, weights_only=True).items()}
-        self.hift.load_state_dict(hift_state_dict, strict=True)
-        self.hift.to(self.device).eval()
+        hift_state_dict = torch.load(hift_model, map_location='cpu', weights_only=True, mmap=True)
+        hift_state_dict = {
+            k.replace('generator.', ''): v
+            for k, v in hift_state_dict.items()
+        }
+        self._assign_state_dict(self.hift, hift_state_dict)
+
+    def _select_checkpoint_path(self, checkpoint_path: str) -> str:
+        if not self.fp16:
+            return checkpoint_path
+        base, ext = os.path.splitext(checkpoint_path)
+        fp16_path = f'{base}.fp16{ext}'
+        return fp16_path if os.path.exists(fp16_path) else checkpoint_path
+
+    def _assign_state_dict(self, module: torch.nn.Module, state_dict: dict):
+        module.load_state_dict(state_dict, strict=True, assign=True)
+        del state_dict
+        self._move_module_tensors(module)
+        module.eval()
+
+    def _load_module_checkpoint(self, module: torch.nn.Module, checkpoint_path: str):
+        state_dict = torch.load(checkpoint_path, map_location='cpu', weights_only=True, mmap=True)
+        self._assign_state_dict(module, state_dict)
+
+    def _move_module_tensors(self, module: torch.nn.Module):
+        target_dtype = torch.float16 if self.fp16 else torch.float32
+        for child in module.children():
+            self._move_module_tensors(child)
+        for name, param in list(module._parameters.items()):
+            if param is None:
+                continue
+            moved = param.to(device=self.device, dtype=target_dtype if param.is_floating_point() else None)
+            module._parameters[name] = torch.nn.Parameter(moved, requires_grad=param.requires_grad)
+        for name, buf in list(module._buffers.items()):
+            if buf is None:
+                continue
+            if buf.is_floating_point():
+                module._buffers[name] = buf.to(device=self.device, dtype=target_dtype)
+            else:
+                module._buffers[name] = buf.to(device=self.device)
 
     def load_jit(self, llm_text_encoder_model, llm_llm_model, flow_encoder_model):
         llm_text_encoder = torch.jit.load(llm_text_encoder_model, map_location=self.device)

flow.fp16.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50df4fe9f2e0a9a54dd2efd48f98ddcd4c4a3672f5e56212c672ac590b280b41
+size 664624971

hift.fp16.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b199a89141c2561e36e3311e96da32a0cdcdfd7bbdc80094a333f06b602c884
+size 41665547

inference.py

@@ -167,7 +167,7 @@ def remove_tail_click(audio, sr, search_s=0.20, burst_thresh=0.05,
 
 
 def synthesize(text, out_path):
-    model = CosyVoice3(MODEL_DIR)
+    model = CosyVoice3(MODEL_DIR, fp16=True)
     for output in model.inference_sft(
         INSTRUCT + text,
         spk_id=SPK_ID,

llm.fp16.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b4849441034fcfc4ab39821a7ae96fa59b53b0ade8eff60768c7f0b455041cb
+size 1284662368

scripts/convert_checkpoints_to_fp16.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+import argparse
+import os
+from pathlib import Path
+
+import torch
+
+
+def convert_state_dict_to_fp16(state_dict):
+    converted = {}
+    for key, value in state_dict.items():
+        if isinstance(value, torch.Tensor) and value.is_floating_point():
+            converted[key] = value.half()
+        else:
+            converted[key] = value
+    return converted
+
+
+def convert_checkpoint(src_path: Path, dst_path: Path):
+    state_dict = torch.load(src_path, map_location="cpu", weights_only=True)
+    converted = convert_state_dict_to_fp16(state_dict)
+    torch.save(converted, dst_path)
+    src_size = os.path.getsize(src_path)
+    dst_size = os.path.getsize(dst_path)
+    print(
+        f"{src_path.name} -> {dst_path.name} | "
+        f"{src_size / 1024**2:.2f} MiB -> {dst_size / 1024**2:.2f} MiB"
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Create fp16 checkpoint copies alongside the original fp32 files.")
+    parser.add_argument(
+        "--files",
+        nargs="*",
+        default=["llm.pt", "flow.pt", "hift.pt"],
+        help="Checkpoint files to convert.",
+    )
+    args = parser.parse_args()
+
+    for rel_path in args.files:
+        src_path = Path(rel_path)
+        if not src_path.exists():
+            raise FileNotFoundError(f"Missing checkpoint: {src_path}")
+        if src_path.suffix != ".pt":
+            raise ValueError(f"Expected a .pt checkpoint, got: {src_path}")
+        dst_path = src_path.with_suffix(f".fp16{src_path.suffix}")
+        convert_checkpoint(src_path, dst_path)
+
+
+if __name__ == "__main__":
+    main()

web-app/app.py

@@ -76,7 +76,7 @@ def load_model():
     global model
     if model is None:
         print(f"Loading model from {MODEL_DIR}...")
-        model = CosyVoice3(MODEL_DIR)
+        model = CosyVoice3(MODEL_DIR, fp16=True)
     return model
 
 def remove_tail_click(audio, sr, search_s=0.20, burst_thresh=0.05,