(case: dict[str, Any], args: argparse.Namespace)
| 968 | |
| 969 | |
| 970 | def pocket_tts_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]: |
| 971 | import yaml |
| 972 | import torch |
| 973 | |
| 974 | set_threads(args.threads) |
| 975 | ref_root = REPO_ROOT / "reference" / "pocket-tts" |
| 976 | if str(ref_root) not in sys.path: |
| 977 | sys.path.insert(0, str(ref_root)) |
| 978 | from pocket_tts import TTSModel |
| 979 | |
| 980 | model_dir = ensure_absolute_path(case["model"]) |
| 981 | language = case.get("session_options", {}).get("language", "english") |
| 982 | config_src = ref_root / "pocket_tts" / "config" / f"{language}.yaml" |
| 983 | config = yaml.safe_load(config_src.read_text(encoding="utf-8")) |
| 984 | bundle_dir = model_dir / "languages" / language if (model_dir / "languages" / language).exists() else model_dir |
| 985 | config["weights_path"] = str(bundle_dir / "model.safetensors") |
| 986 | config["weights_path_without_voice_cloning"] = str(bundle_dir / "model.safetensors") |
| 987 | config["flow_lm"]["lookup_table"]["tokenizer_path"] = str(bundle_dir / "tokenizer.model") |
| 988 | config_dir = REPO_ROOT / "build" / "perf" / "pocket_tts" |
| 989 | config_dir.mkdir(parents=True, exist_ok=True) |
| 990 | config_path = config_dir / f"{language}.yaml" |
| 991 | config_path.write_text(yaml.safe_dump(config, sort_keys=False), encoding="utf-8") |
| 992 | device = torch.device(f"cuda:{args.device}" if args.backend == "cuda" else "cpu") |
| 993 | model = TTSModel.load_model(config=str(config_path)).to(device).eval() |
| 994 | |
| 995 | def run(request: dict[str, Any], _args: argparse.Namespace) -> tuple[float, float]: |
| 996 | voice_state = model.get_state_for_audio_prompt(request["voice_ref"]) |
| 997 | started = time.perf_counter() |
| 998 | audio = model.generate_audio(voice_state, request["text"]) |
| 999 | if args.backend == "cuda": |
| 1000 | torch.cuda.synchronize(args.device) |
| 1001 | audio_np = audio.detach().cpu().numpy().astype(np.float32, copy=False) |
| 1002 | wall_ms = (time.perf_counter() - started) * 1000.0 |
| 1003 | return wall_ms, audio_duration_sec(audio_np, int(model.sample_rate)) |
| 1004 | |
| 1005 | return model, run |
| 1006 | |
| 1007 | |
| 1008 | def omnivoice_loader(case: dict[str, Any], args: argparse.Namespace) -> tuple[Any, RunnerFn]: |
nothing calls this directly
no test coverage detected