SpongeBobFan2002/openaudio-s1-mini

1import os2from argparse import ArgumentParser3from pathlib import Path4 5import pyrootutils6import torch7from loguru import logger8 9pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)10 11from fish_speech.inference_engine import TTSInferenceEngine12from fish_speech.models.dac.inference import load_model as load_decoder_model13from fish_speech.models.text2semantic.inference import launch_thread_safe_queue14from fish_speech.utils.schema import ServeTTSRequest15from tools.webui import build_app16from tools.webui.inference import get_inference_wrapper17 18# Make einx happy19os.environ["EINX_FILTER_TRACEBACK"] = "false"20 21 22def parse_args():23    parser = ArgumentParser()24    parser.add_argument(25        "--llama-checkpoint-path",26        type=Path,27        default="checkpoints/openaudio-s1-mini",28    )29    parser.add_argument(30        "--decoder-checkpoint-path",31        type=Path,32        default="checkpoints/openaudio-s1-mini/codec.pth",33    )34    parser.add_argument("--decoder-config-name", type=str, default="modded_dac_vq")35    parser.add_argument("--device", type=str, default="cuda")36    parser.add_argument("--half", action="store_true")37    parser.add_argument("--compile", action="store_true")38    parser.add_argument("--max-gradio-length", type=int, default=0)39    parser.add_argument("--theme", type=str, default="light")40 41    return parser.parse_args()42 43 44if __name__ == "__main__":45    args = parse_args()46    args.precision = torch.half if args.half else torch.bfloat1647 48    # Check if MPS or CUDA is available49    if torch.backends.mps.is_available():50        args.device = "mps"51        logger.info("mps is available, running on mps.")52    elif not torch.cuda.is_available():53        logger.info("CUDA is not available, running on CPU.")54        args.device = "cpu"55 56    logger.info("Loading Llama model...")57    llama_queue = launch_thread_safe_queue(58        checkpoint_path=args.llama_checkpoint_path,59        device=args.device,60        precision=args.precision,61        compile=args.compile,62    )63 64    logger.info("Loading VQ-GAN model...")65    decoder_model = load_decoder_model(66        config_name=args.decoder_config_name,67        checkpoint_path=args.decoder_checkpoint_path,68        device=args.device,69    )70 71    logger.info("Decoder model loaded, warming up...")72 73    # Create the inference engine74    inference_engine = TTSInferenceEngine(75        llama_queue=llama_queue,76        decoder_model=decoder_model,77        compile=args.compile,78        precision=args.precision,79    )80 81    # Dry run to check if the model is loaded correctly and avoid the first-time latency82    list(83        inference_engine.inference(84            ServeTTSRequest(85                text="Hello world.",86                references=[],87                reference_id=None,88                max_new_tokens=1024,89                chunk_length=200,90                top_p=0.7,91                repetition_penalty=1.5,92                temperature=0.7,93                format="wav",94            )95        )96    )97 98    logger.info("Warming up done, launching the web UI...")99 100    # Get the inference function with the immutable arguments101    inference_fct = get_inference_wrapper(inference_engine)102 103    app = build_app(inference_fct, args.theme)104    app.launch(show_api=True)105