Function main

streamer.py:255–372

()

def main():
    parser = argparse.ArgumentParser(description="TTS inference with Asteroid model")
    parser.add_argument(
        "--jsonl",
        default="examples/examples.jsonl",
        help="Path to JSONL file (default: examples/examples.jsonl)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed for reproducibility (default: None)",
    )
    parser.add_argument(
        "--output_dir",
        default="outputs/streamer",
        help="Output directory for generated audio files (default: outputs)",
    )
    parser.add_argument(
        "--use_normalize",
        action="store_true",
        default=True,
        help="Whether to use text normalization (default: True)",
    )
    parser.add_argument(
        "--dtype",
        choices=["bf16", "fp16", "fp32"],
        default="bf16",
        help="Model data type (default: bf16)",
    )
    parser.add_argument(
        "--attn_implementation",
        choices=["flash_attention_2", "sdpa", "eager"],
        default="flash_attention_2",
        help="Attention implementation (default: flash_attention_2)",
    )
    parser.add_argument(
        "--use_tqdm",
        action="store_true",
        default=False,
        help="Whether to show progress bar using tqdm (default: False)",
    )

    args = parser.parse_args()

    # Convert dtype string to torch dtype
    dtype_mapping = {
        "bf16": torch.bfloat16,
        "fp16": torch.float16,
        "fp32": torch.float32,
    }
    torch_dtype = dtype_mapping[args.dtype]

    # Create output directory if it doesn't exist
    os.makedirs(args.output_dir, exist_ok=True)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")