init

.gitignore

@@ -0,0 +1,180 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+data
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+.idea/
+
+.DS_Store
+data_process/
+internvl_chat/work_dirs/
+internvl_chat/unittest/
+internvl_chat/data/
+Husky2/*
+data_process/
+*distillation*
+
+batchscript-*
+results/
+
+# *txt
+*csv
+*mp4
+temp/
+src/moviedubber/infer/basic_test.toml

README.md

@@ -1,12 +1,12 @@
 ---
-title: DeepDubber V1
-emoji: 🐠
-colorFrom: blue
-colorTo: pink
+title: Deepdubber V1
+emoji: 🌍
+colorFrom: red
+colorTo: green
 sdk: gradio
 sdk_version: 5.22.0
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,189 @@
+import os
+import tempfile
+
+import gradio as gr
+import librosa
+import soundfile
+import tomli
+import torch
+import torch.nn.functional as F
+import torchaudio
+from moviepy import VideoFileClip
+from pydub import AudioSegment
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+
+from src.moviedubber.infer.utils_infer import (
+    cfg_strength,
+    chunk_text,
+    nfe_step,
+    sway_sampling_coef,
+)
+from src.moviedubber.infer.video_preprocess import VideoFeatureExtractor
+from src.moviedubber.infer_with_mmlm_result import concat_movie_with_audio, get_spk_emb, load_models
+from src.moviedubber.model.utils import convert_char_to_pinyin
+
+
+def load_asr_model(model_id="openai/whisper-large-v3-turbo"):
+    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+    model = AutoModelForSpeechSeq2Seq.from_pretrained(
+        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
+    ).to(device)
+
+    processor = AutoProcessor.from_pretrained(model_id)
+
+    pipe = pipeline(
+        "automatic-speech-recognition",
+        model=model,
+        tokenizer=processor.tokenizer,
+        feature_extractor=processor.feature_extractor,
+        torch_dtype=torch_dtype,
+        device=device,
+    )
+    return pipe
+
+
+device = "cpu"
+config = tomli.load(open("src/moviedubber/infer/basic.toml", "rb"))
+
+
+ema_model, vocoder, ort_session = load_models(config, device=device)
+asr_pipe = load_asr_model()
+
+videofeature_extractor = VideoFeatureExtractor(device=device)
+
+
+def deepdubber(video_path: str, subtitle_text: str, audio_path: str = None) -> str:
+    print(f"Starting deepdubber with video_path: {video_path} and subtitle_text: {subtitle_text}")
+    gen_clip = videofeature_extractor.extract_features(video_path)
+    gen_text = subtitle_text
+
+    clip = VideoFileClip(video_path)
+    gen_audio_len = int(clip.duration * 24000 // 256)
+
+    gen_clip = gen_clip.unsqueeze(0).to(device=device, dtype=torch.float32).transpose(1, 2)
+    gen_clip = F.interpolate(gen_clip, size=(gen_audio_len,), mode="linear", align_corners=False).transpose(1, 2)
+
+    ref_audio_len = None
+    if audio_path is not None:
+        print("reference audio is not None, dubbing with reference audio")
+
+        if audio_path.endswith(".mp3"):
+            audio = AudioSegment.from_mp3(audio_path)
+
+            wav_file = audio_path.replace(".mp3", ".wav")
+            audio.export(wav_file, format="wav")
+        else:
+            wav_file = audio_path
+
+        ref_text = asr_pipe(librosa.load(wav_file, sr=16000)[0], generate_kwargs={"language": "english"})["text"]
+        ref_text = ref_text.replace("\n", " ").replace("\r", " ")
+        print(f"Reference text: {ref_text}")
+
+        spk_emb = get_spk_emb(wav_file, ort_session)
+        spk_emb = torch.tensor(spk_emb).to(device=device, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+
+        audio_data, sr = torchaudio.load(wav_file)
+        resampler = torchaudio.transforms.Resample(sr, 24000)
+        if sr != 24000:
+            audio_data = resampler(audio_data)
+
+        if audio_data.shape[0] > 1:
+            audio_data = torch.mean(audio_data, dim=0, keepdim=True)
+
+        audio_data = audio_data.to(device)
+
+        ref_audio_len = int(audio_data.shape[-1] // 256)
+        ref_clip = torch.zeros((1, ref_audio_len, 768)).to(device=device)
+
+        gen_clip = torch.cat((gen_clip, ref_clip), dim=1)
+
+        gen_audio_len = ref_audio_len + gen_audio_len
+
+        gen_text = ref_text + " " + gen_text
+
+    else:
+        spk_emb = torch.zeros((1, 1, 192)).to(device=device)
+        audio_data = torch.zeros((1, gen_audio_len, 100)).to(device=device)
+
+    gen_text_batches = chunk_text(gen_text, max_chars=1024)
+    final_text_list = convert_char_to_pinyin(gen_text_batches)
+
+    with torch.inference_mode():
+        generated, _ = ema_model.sample(
+            cond=audio_data,
+            text=final_text_list,
+            clip=gen_clip,
+            spk_emb=spk_emb,
+            duration=gen_audio_len,
+            steps=nfe_step,
+            cfg_strength=cfg_strength,
+            sway_sampling_coef=sway_sampling_coef,
+            no_ref_audio=False,
+        )
+
+        generated = generated.to(torch.float32)
+
+        if ref_audio_len is not None:
+            generated = generated[:, ref_audio_len:, :]
+
+        generated_mel_spec = generated.permute(0, 2, 1)
+        generated_wave = vocoder(generated_mel_spec)
+
+        generated_wave = generated_wave.squeeze().cpu().numpy()
+
+    # using a temporary wav file to save the generated audio
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_wav_file:
+        temp_wav_path = temp_wav_file.name
+        soundfile.write(temp_wav_path, generated_wave, samplerate=24000)
+
+    concated_video = concat_movie_with_audio(temp_wav_path, video_path, ".")
+
+    # Ensure the temporary file is deleted after use
+    os.remove(temp_wav_path)
+
+    print(f"Deepdubber completed successfully, output path: {concated_video}")
+    return concated_video
+
+
+def process_video_dubbing(video_path: str, subtitle_text: str, audio_path: str = None) -> str:
+    try:
+        print(f"Processing video: {video_path}")
+        if not os.path.exists(video_path):
+            raise ValueError("Video file does not exist")
+
+        if not subtitle_text.strip():
+            raise ValueError("Subtitle text cannot be empty")
+
+        output_path = deepdubber(video_path, subtitle_text, audio_path)
+
+        return output_path
+
+    except Exception as e:
+        print(f"Error in process_video_dubbing: {e}")
+
+        return None
+
+
+def create_ui():
+    with gr.Blocks(title="DeepDubber-V1") as app:
+        gr.Markdown("# DeepDubber-V1\nUpload your video file and enter the text you want to dub")
+
+        with gr.Row():
+            video_input = gr.Video(label="Upload video")
+            audio_input = gr.Audio(label="Upload audio", type="filepath")
+            subtitle_input = gr.Textbox(label="Enter the text", placeholder="Enter the text to be dubbed...", lines=5)
+
+        process_btn = gr.Button("Start Dubbing")
+
+        output_video = gr.Video(label="Dubbed Video")
+
+        process_btn.click(
+            fn=process_video_dubbing, inputs=[video_input, subtitle_input, audio_input], outputs=output_video
+        )
+
+    return app
+
+
+if __name__ == "__main__":
+    app = create_ui()
+    app.launch()

datasets/CoTMovieDubbing/README.md

@@ -0,0 +1,20 @@
+## MovieData Preprocessing
+
+TODO
+
+## Data Tree
+
+```
+moviecopy
+|-- movie1_name
+    |-- xxxx.mp4
+    |-- xxxx.mp3
+    |-- xxxx.txt
+    |-- ...
+|-- movie2_name
+    |-- xxxx.mp4
+    |-- xxxx.mp3
+    |-- xxxx.txt
+    |-- ...
+|-- ...
+```

datasets/Grid/README.md

@@ -0,0 +1 @@
+Refer to: [Grid](https://paperswithcode.com/dataset/grid)

datasets/V2C/README.md

@@ -0,0 +1 @@
+Refer to: [V2C](https://github.com/chenqi008/V2C)

requirements.txt

@@ -0,0 +1,237 @@
+absl-py==2.1.0
+accelerate==0.34.2
+addict==2.4.0
+aiofiles==24.1.0
+aiohappyeyeballs==2.4.6
+aiohttp==3.11.13
+aiosignal==1.3.2
+altair==5.5.0
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.8.0
+async-timeout==5.0.1
+attrs==25.1.0
+audioread==3.0.1
+beautifulsoup4==4.13.3
+bitsandbytes==0.42.0
+blinker==1.9.0
+boto3==1.37.12
+botocore==1.37.12
+cached_path==1.7.1
+cachetools==5.5.2
+certifi==2025.1.31
+cffi==1.17.1
+charset-normalizer==3.4.1
+click==8.1.8
+colorama==0.4.6
+coloredlogs==15.0.1
+contourpy==1.3.0
+cycler==0.12.1
+datasets==3.3.2
+decorator==5.2.1
+decord==0.6.0
+deepspeed==0.15.4
+dill==0.3.8
+docstring_parser==0.16
+einops==0.8.1
+einops-exts==0.0.4
+einx==0.3.0
+eval_type_backport==0.2.2
+exceptiongroup==1.2.2
+fastapi==0.115.8
+ffmpy==0.5.0
+filelock==3.13.1
+flash-attn==2.6.3
+flatbuffers==25.2.10
+fonttools==4.56.0
+frozendict==2.4.6
+frozenlist==1.5.0
+fsspec==2024.6.1
+future==1.0.0
+gdown==5.2.0
+gitdb==4.0.12
+GitPython==3.1.44
+google-api-core==2.24.2
+google-auth==2.38.0
+google-cloud-core==2.4.3
+google-cloud-storage==2.19.0
+google-crc32c==1.6.0
+google-resumable-media==2.7.2
+googleapis-common-protos==1.69.1
+gradio==3.35.2
+gradio_client==0.2.9
+grpcio==1.70.0
+h11==0.14.0
+hjson==3.1.0
+httpcore==0.17.3
+httpx==0.24.0
+huggingface-hub==0.27.1
+humanfriendly==10.0
+idna==3.10
+imageio==2.37.0
+imageio-ffmpeg==0.6.0
+importlib_metadata==8.6.1
+importlib_resources==6.5.2
+jieba==0.42.1
+Jinja2==3.1.4
+jmespath==1.0.1
+joblib==1.4.2
+jsonschema==4.23.0
+jsonschema-specifications==2024.10.1
+kiwisolver==1.4.7
+latex2mathml==3.77.0
+lazy_loader==0.4
+librosa==0.11.0
+liger_kernel==0.4.2
+linkify-it-py==2.0.3
+llvmlite==0.43.0
+loguru==0.7.3
+Markdown==3.7
+markdown-it-py==2.2.0
+markdown2==2.5.3
+MarkupSafe==2.1.5
+matplotlib==3.9.4
+mdit-py-plugins==0.3.3
+mdurl==0.1.2
+mmcls==0.25.0
+mmcv-full==1.6.2
+mmsegmentation==0.30.0
+model-index==0.1.11
+moviepy==2.1.2
+mpmath==1.3.0
+msgpack==1.1.0
+multidict==6.1.0
+multiprocess==0.70.16
+narwhals==1.28.0
+networkx==3.2.1
+ninja==1.11.1.3
+numba==0.60.0
+numpy==1.26.3
+nvidia-cublas-cu11==11.11.3.6
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu11==11.8.87
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu11==11.8.89
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu11==11.8.89
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu11==9.1.0.70
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu11==10.9.0.58
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu11==10.3.0.86
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu11==11.4.1.48
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu11==11.7.5.86
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-ml-py==12.570.86
+nvidia-nccl-cu11==2.20.5
+nvidia-nccl-cu12==2.20.5
+nvidia-nvjitlink-cu12==12.8.61
+nvidia-nvtx-cu11==11.8.86
+nvidia-nvtx-cu12==12.1.105
+omegaconf==2.3.0
+onnxruntime==1.18.0
+opencv-python==4.11.0.86
+opendatalab==0.0.10
+openmim==0.3.9
+openxlab==0.0.11
+ordered-set==4.1.0
+orjson==3.10.15
+packaging==24.2
+pandas==2.2.3
+peft==0.10.0
+pillow==10.4.0
+platformdirs==4.3.6
+pooch==1.8.2
+prettytable==3.14.0
+proglog==0.1.10
+propcache==0.3.0
+proto-plus==1.26.1
+protobuf==5.29.3
+psutil==7.0.0
+py-cpuinfo==9.0.0
+pyarrow==19.0.1
+pyasn1==0.6.1
+pyasn1_modules==0.4.1
+pycocoevalcap==1.2
+pycocotools==2.0.8
+pycparser==2.22
+pycryptodome==3.21.0
+pydantic==2.10.6
+pydantic_core==2.27.2
+pydeck==0.9.1
+pydub==0.25.1
+Pygments==2.19.1
+pyparsing==3.2.1
+pypinyin==0.53.0
+PySocks==1.7.1
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.20
+pytz==2025.1
+PyYAML==6.0.2
+referencing==0.36.2
+regex==2024.11.6
+requests==2.32.3
+rich==13.9.4
+rpds-py==0.23.1
+rsa==4.9
+s3transfer==0.11.4
+safetensors==0.5.3
+scikit-learn==1.6.1
+scipy==1.13.1
+semantic-version==2.10.0
+sentencepiece==0.1.99
+shortuuid==1.0.13
+shtab==1.7.1
+six==1.17.0
+smmap==5.0.2
+sniffio==1.3.1
+soundfile==0.13.1
+soupsieve==2.6
+soxr==0.5.0.post1
+starlette==0.45.3
+streamlit==1.42.2
+streamlit-image-select==0.6.0
+svgwrite==1.4.3
+sympy==1.13.1
+tabulate==0.9.0
+tenacity==9.0.0
+tensorboard==2.19.0
+tensorboard-data-server==0.7.2
+tensorboardX==2.6.2.2
+termcolor==2.5.0
+threadpoolctl==3.5.0
+timm==0.9.12
+tokenizers==0.19.1
+toml==0.10.2
+tomli==2.2.1
+torch==2.4.1
+torchaudio==2.4.1
+torchdiffeq==0.2.5
+torchvision==0.19.1
+tornado==6.4.2
+tqdm==4.67.1
+transformers==4.42.1
+triton==3.0.0
+trl==0.10.1
+typeguard==4.4.2
+typing_extensions==4.12.2
+tyro==0.9.16
+tzdata==2025.1
+uc-micro-py==1.0.3
+urllib3==1.26.20
+uvicorn==0.34.0
+watchdog==6.0.0
+wavedrom==2.0.3.post3
+wcwidth==0.2.13
+websockets==15.0
+Werkzeug==3.1.3
+x-transformers==2.1.37
+xxhash==3.5.0
+yacs==0.1.8
+yapf==0.40.1
+yarl==1.18.3
+zipp==3.21.0

ruff.toml

@@ -0,0 +1,11 @@
+line-length = 120
+target-version = "py310"
+
+[lint]
+# Only ignore variables with names starting with "_".
+dummy-variable-rgx = "^_.*$"
+ignore = ["E402"]
+
+[lint.isort]
+force-single-line = false
+lines-after-imports = 2

src/internvl/eval.py

@@ -0,0 +1,337 @@
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+import torchvision.transforms as T
+from decord import VideoReader, cpu
+from PIL import Image
+from torchvision.transforms.functional import InterpolationMode
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+
+sys.path.insert(0, os.path.join(str(Path(__file__).resolve().parents[2]), "src/third_party/InternVL/internvl_chat"))
+from internvl.model.internvl_chat.modeling_internvl_chat import InternVLChatModel  # type: ignore
+
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+
+def build_transform(input_size):
+    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
+    transform = T.Compose(
+        [
+            T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
+            T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
+            T.ToTensor(),
+            T.Normalize(mean=MEAN, std=STD),
+        ]
+    )
+    return transform
+
+
+def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+    best_ratio_diff = float("inf")
+    best_ratio = (1, 1)
+    area = width * height
+    for ratio in target_ratios:
+        target_aspect_ratio = ratio[0] / ratio[1]
+        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+        if ratio_diff < best_ratio_diff:
+            best_ratio_diff = ratio_diff
+            best_ratio = ratio
+        elif ratio_diff == best_ratio_diff:
+            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                best_ratio = ratio
+    return best_ratio
+
+
+def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
+    orig_width, orig_height = image.size
+    aspect_ratio = orig_width / orig_height
+
+    # calculate the existing image aspect ratio
+    target_ratios = set(
+        (i, j)
+        for n in range(min_num, max_num + 1)
+        for i in range(1, n + 1)
+        for j in range(1, n + 1)
+        if i * j <= max_num and i * j >= min_num
+    )
+    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+    # find the closest aspect ratio to the target
+    target_aspect_ratio = find_closest_aspect_ratio(aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+    # calculate the target width and height
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    # resize the image
+    resized_img = image.resize((target_width, target_height))
+    processed_images = []
+    for i in range(blocks):
+        box = (
+            (i % (target_width // image_size)) * image_size,
+            (i // (target_width // image_size)) * image_size,
+            ((i % (target_width // image_size)) + 1) * image_size,
+            ((i // (target_width // image_size)) + 1) * image_size,
+        )
+        # split the image
+        split_img = resized_img.crop(box)
+        processed_images.append(split_img)
+    assert len(processed_images) == blocks
+    if use_thumbnail and len(processed_images) != 1:
+        thumbnail_img = image.resize((image_size, image_size))
+        processed_images.append(thumbnail_img)
+    return processed_images
+
+
+def load_image(image_file, input_size=448, max_num=12):
+    image = Image.open(image_file).convert("RGB")
+    transform = build_transform(input_size=input_size)
+    images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
+    pixel_values = [transform(image) for image in images]
+    pixel_values = torch.stack(pixel_values)
+    return pixel_values
+
+
+def get_index(bound, fps, max_frame, first_idx=0, num_segments=32):
+    if bound:
+        start, end = bound[0], bound[1]
+    else:
+        start, end = -100000, 100000
+    start_idx = max(first_idx, round(start * fps))
+    end_idx = min(round(end * fps), max_frame)
+    seg_size = float(end_idx - start_idx) / num_segments
+    frame_indices = np.array(
+        [int(start_idx + (seg_size / 2) + np.round(seg_size * idx)) for idx in range(num_segments)]
+    )
+    return frame_indices
+
+
+def load_video(
+    video_path,
+    bound=None,
+    input_size=448,
+    max_num=1,
+    num_segments=32,
+    cache_dir=".cache/expcache",
+):
+    vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
+    max_frame = len(vr) - 1
+    fps = float(vr.get_avg_fps())
+
+    video_cache_dir = video_path.split("/")[-2] + "_" + os.path.basename(video_path).split(".")[0]
+    video_cache_dir = os.path.join(cache_dir, video_cache_dir)
+    cache_filename = os.path.join(
+        video_cache_dir,
+        f"_bound-{bound}_input_size-{input_size}_max_num-{max_num}_num_segments-{num_segments}.pt",
+    )
+    if os.path.exists(cache_filename) and os.path.isfile(cache_filename):
+        cache = torch.load(cache_filename, weights_only=True)
+        pixel_values = cache["pixel_values"]
+        num_patches_list = cache["num_patches_list"]
+
+    else:
+        pixel_values_list, num_patches_list = [], []
+        transform = build_transform(input_size=input_size)
+        frame_indices = get_index(bound, fps, max_frame, first_idx=0, num_segments=num_segments)
+
+        frame_indices = np.append(0, frame_indices)  # Add 0 at the beginning of the list
+        frame_indices = np.append(frame_indices, max_frame)  # Add max_frame at the end of the list
+
+        os.makedirs(video_cache_dir, exist_ok=True)
+
+        idx = 0
+        for frame_index in frame_indices:
+            img = Image.fromarray(vr[frame_index].asnumpy()).convert("RGB")
+
+            img.save(os.path.join(video_cache_dir, f"frame_{frame_index}_tile_{idx}.png"))
+
+            img = dynamic_preprocess(img, image_size=input_size, use_thumbnail=True, max_num=max_num)
+
+            pixel_values = [transform(tile) for tile in img]
+            pixel_values = torch.stack(pixel_values)
+            num_patches_list.append(pixel_values.shape[0])
+            pixel_values_list.append(pixel_values)
+
+            idx += 1
+        pixel_values = torch.cat(pixel_values_list)
+
+        os.makedirs(cache_dir, exist_ok=True)
+        torch.save({"pixel_values": pixel_values, "num_patches_list": num_patches_list}, cache_filename)
+
+    return pixel_values, num_patches_list
+
+
+def analyze_predictions(file_path):
+    # Read the CSV file
+    df = pd.read_csv(file_path)
+
+    # Calculate overall accuracy
+    total_samples = len(df)
+    correct_predictions = df["is_correct"].value_counts().get(True, 0)
+    overall_accuracy = correct_predictions / total_samples
+
+    # Initialize metrics for each class
+    classes = ["A", "B", "C"]
+    class_metrics = {}
+
+    for cls in classes:
+        # Filter for samples where target is this class
+        true_class = df[df["target"] == cls]
+        # Filter for samples where prediction is this class
+        # pred_class = df[df["predict"] == cls]
+
+        # Calculate TP, FP, FN
+        TP = len(df[(df["target"] == cls) & (df["predict"] == cls)])
+        FP = len(df[(df["target"] != cls) & (df["predict"] == cls)])
+        FN = len(df[(df["target"] == cls) & (df["predict"] != cls)])
+
+        # Calculate precision, recall, F1
+        precision = TP / (TP + FP) if (TP + FP) > 0 else 0
+        recall = TP / (TP + FN) if (TP + FN) > 0 else 0
+        f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
+
+        # Store metrics
+        class_metrics[cls] = {
+            "total_samples": len(true_class),
+            "precision": precision,
+            "recall": recall,
+            "f1": f1,
+            "true_positives": TP,
+            "false_positives": FP,
+            "false_negatives": FN,
+        }
+
+    print(f"Overall Accuracy: {overall_accuracy:.4f} ({correct_predictions}/{total_samples})")
+    print()
+    print("Indicators for each category:")
+
+    for cls in classes:
+        metrics = class_metrics[cls]
+        print(f"  Class {cls}:")
+        print(f"    Total Samples: {metrics['total_samples']}")
+        print(f"    Precision: {metrics['precision']:.4f}")
+        print(f"    Recall: {metrics['recall']:.4f}")
+        print(f"    F1 Score: {metrics['f1']:.4f}")
+        print(f"    True Positives: {metrics['true_positives']}")
+        print(f"    False Positives: {metrics['false_positives']}")
+        print(f"    False Negatives: {metrics['false_negatives']}")
+
+    return overall_accuracy, class_metrics
+
+
+def s_thread(video_dir, model_path, device, chunk, idx, queue):
+    model = InternVLChatModel.from_pretrained(
+        model_path,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+        use_flash_attn=True,
+    )
+    model = model.eval().to(device)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=False)
+
+    generation_config = dict(max_new_tokens=1024, do_sample=False)
+
+    res = []
+    for line in tqdm(chunk, position=idx, desc=f"Device {device}"):
+        data = json.loads(line)
+
+        video_path = os.path.join(video_dir, data["video"])
+        ques = data["conversations"][0]["value"]
+
+        target_ans = data["conversations"][1]["value"].split("<CONCLUSION>")[1].split("</CONCLUSION>")[0].strip()
+
+        pixel_values, num_patches_list = load_video(video_path, num_segments=8, max_num=1)
+        pixel_values = pixel_values.to(torch.bfloat16).to(device)
+        video_prefix = "".join([f"Frame{i + 1}: <image>\n" for i in range(len(num_patches_list))])
+        question = video_prefix + f"{ques}"
+        response = model.chat(
+            tokenizer,
+            pixel_values,
+            question,
+            generation_config,
+            num_patches_list=num_patches_list,
+            history=None,
+            return_history=False,
+        )
+
+        try:
+            ans = response.split("<CONCLUSION>")[1].split("</CONCLUSION>")[0].strip()
+        except Exception as e:
+            print(f"Error: {e}, response: {response}")
+            ans = response.strip()[0]
+
+        is_correct = False
+        if ans == target_ans:
+            is_correct = True
+
+        res.append(f"{video_path},{is_correct},{target_ans},{ans}")
+
+    queue.put(res)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    import torch.multiprocessing as mp
+
+    parser = argparse.ArgumentParser(description="eval script for mmlm")
+    parser.add_argument("--model_path", type=str, help="Path to the model checkpoint.")
+    parser.add_argument("--test_file", type=str, help="Path to the test file.")
+    parser.add_argument("--video_dir", type=str, help="Path to the test video directory.")
+    parser.add_argument("--gpuids", type=str, help="GPU ids to use.")
+
+    # python eval.py --model_path /path/to/model --test_file /path/to/test_file --video_dir /path/to/video_dir --gpuids 0,1,2,3
+
+    args = parser.parse_args()
+
+    model_path = args.model_path
+    test_file = args.test_file
+    video_dir = args.video_dir
+
+    gpu_ids = args.gpuids.split(",") if args.gpuids else ["0"]
+
+    cot_test = Path(test_file).read_text().splitlines()
+
+    chunks = np.array_split(cot_test, len(gpu_ids))
+
+    mp.set_start_method("spawn", force=True)
+
+    queue = mp.Queue()
+
+    processes = []
+    for idx, chunk in enumerate(chunks):
+        device = gpu_ids[idx % len(gpu_ids)]
+        device = f"cuda:{device}"
+
+        p = mp.Process(target=s_thread, args=(video_dir, model_path, device, chunk, idx, queue))
+        processes.append(p)
+        p.start()
+
+    for process in processes:
+        process.join()
+
+    result = []
+    for _ in range(len(chunks)):
+        res = queue.get()
+        result.extend(res)
+
+    res_saved = f"{'__'.join(model_path.split('/'))}_res.csv"
+    with open(res_saved, "w") as f:
+        f.write("video_id,is_correct,target,predict\n")
+        for res in result:
+            f.write(f"{res}\n")
+
+    accuracy, metrics = analyze_predictions(res_saved)
+
+    print("All processes finished.\n\n")

src/moviedubber/configs/basemodel.yaml

@@ -0,0 +1,9 @@
+model:
+  arch:
+    dim: 1024
+    depth: 22
+    heads: 16
+    ff_mult: 2
+    text_dim: 512
+    conv_layers: 4
+

src/moviedubber/eval.py

@@ -0,0 +1,245 @@
+import argparse
+import os
+import string
+from concurrent.futures import ProcessPoolExecutor
+from pathlib import Path
+
+import librosa
+import numpy as np
+import torch
+from evaluate import load
+from pymcd.mcd import Calculate_MCD
+from tqdm import tqdm
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, Wav2Vec2FeatureExtractor, WavLMForXVector, pipeline
+
+
+def convert_numbers_to_words(text):
+    """Convert single digits in text to words with spaces"""
+    number_word_map = {
+        "0": "zero",
+        "1": "one",
+        "2": "two",
+        "3": "three",
+        "4": "four",
+        "5": "five",
+        "6": "six",
+        "7": "seven",
+        "8": "eight",
+        "9": "nine",
+    }
+
+    words = text.split()
+    converted_words = []
+
+    for word in words:
+        # Check if the word contains both letters and numbers (like 'j4')
+        if any(c.isdigit() for c in word) and any(c.isalpha() for c in word):
+            # Split the word into parts and convert digits
+            new_word = ""
+            for c in word:
+                if c.isdigit():
+                    new_word += " " + number_word_map[c]
+                else:
+                    new_word += c
+            converted_words.append(new_word)
+        # Check if the word is a single digit
+        elif word.isdigit() and len(word) == 1:
+            converted_words.append(number_word_map[word])
+        else:
+            converted_words.append(word)
+
+    return " ".join(converted_words)
+
+
+def clean_text(text):
+    text = convert_numbers_to_words(text)
+    text = text.translate(str.maketrans("", "", string.punctuation))
+    text = text.lower()
+    return text
+
+
+def wer_pipe(gen_dir: str, target_dir: str, model_id="openai/whisper-large-v3-turbo"):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+    print(f"Using Model: {model_id} for WER Evaluation")
+    model = AutoModelForSpeechSeq2Seq.from_pretrained(
+        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
+    ).to(device)
+
+    processor = AutoProcessor.from_pretrained(model_id)
+
+    pipe = pipeline(
+        "automatic-speech-recognition",
+        model=model,
+        tokenizer=processor.tokenizer,
+        feature_extractor=processor.feature_extractor,
+        torch_dtype=torch_dtype,
+        device=device,
+    )
+
+    gen_list = list(Path(gen_dir).glob("*.wav"))
+    for line in tqdm(gen_list, desc="Processing audio files"):
+        wav = line
+        if not wav.exists():
+            continue
+
+        text = pipe(librosa.load(wav, sr=16000)[0], generate_kwargs={"language": "english"})["text"]
+        with open(wav.with_suffix(".asrtxt"), "w") as fw:
+            fw.write(text)
+
+    wer_metric = load("wer")
+
+    val_list = list(Path(target_dir).glob("*.txt"))
+
+    wer = []
+    for txt in tqdm(val_list, desc="Calculating WER"):
+        try:
+            # Since the original text is automatically transcribed and has not been manually verified, all texts will be cleaned here.
+
+            target_text = " ".join(set(txt.read_text().splitlines()))
+            target_text = clean_text(target_text)
+
+            gen_text = " ".join(Path(os.path.join(gen_dir, txt.with_suffix(".asrtxt").name)).read_text().splitlines())
+            gen_text = clean_text(gen_text)
+
+            if target_text == "" or gen_text == "":
+                continue
+
+            wer_ = wer_metric.compute(references=[target_text], predictions=[gen_text])
+
+        except Exception as e:
+            print("Error in wer calculation: ", e)
+            continue
+
+        wer.append(wer_)
+
+    return np.mean(wer)
+
+
+def spk_sim_pipe(gen_dir, target_dir):
+    feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("microsoft/wavlm-base-sv")
+    model = WavLMForXVector.from_pretrained("microsoft/wavlm-base-sv").cuda()
+
+    cosine_sim = torch.nn.CosineSimilarity(dim=-1)
+
+    val_list = list(Path(target_dir).glob("*.wav"))
+
+    scos = []
+
+    for target_wav in tqdm(val_list, desc="Calculating speaker similarity"):
+        target = librosa.load(target_wav, sr=16000)[0]
+        gen = librosa.load(os.path.join(gen_dir, target_wav.name), sr=16000)[0]
+
+        try:
+            input1 = feature_extractor(gen, return_tensors="pt", sampling_rate=16000).to("cuda")
+            embeddings1 = model(**input1).embeddings
+
+            input2 = feature_extractor(target, return_tensors="pt", sampling_rate=16000).to("cuda")
+            embeddings2 = model(**input2).embeddings
+
+            similarity = cosine_sim(embeddings1[0], embeddings2[0])
+
+        except Exception as e:
+            print(f"Error in {target_wav}, {e}")
+            continue
+
+        scos.append(similarity.detach().cpu().numpy())
+
+    return np.mean(scos)
+
+
+def calculate_mcd_for_wav(target_wav, gen_dir, mcd_toolbox_dtw, mcd_toolbox_dtw_sl):
+    _mcd_dtw = mcd_toolbox_dtw.calculate_mcd(target_wav, os.path.join(gen_dir, target_wav.name))
+    _mcd_dtw_sl = mcd_toolbox_dtw_sl.calculate_mcd(target_wav, os.path.join(gen_dir, target_wav.name))
+    return _mcd_dtw, _mcd_dtw_sl
+
+
+def mcd_pipe(gen_dir, target_dir, num_processes=16):
+    mcd_toolbox_dtw = Calculate_MCD(MCD_mode="dtw")
+    mcd_toolbox_dtw_sl = Calculate_MCD(MCD_mode="dtw_sl")
+
+    val_list = list(Path(target_dir).glob("*.wav"))
+
+    mcd_dtw = []
+    mcd_dtw_sl = []
+
+    with ProcessPoolExecutor(max_workers=num_processes) as executor:
+        futures = [
+            executor.submit(calculate_mcd_for_wav, target_wav, gen_dir, mcd_toolbox_dtw, mcd_toolbox_dtw_sl)
+            for target_wav in val_list
+        ]
+        for future in tqdm(futures, desc="Calculating MCD"):
+            _mcd_dtw, _mcd_dtw_sl = future.result()
+            mcd_dtw.append(_mcd_dtw)
+            mcd_dtw_sl.append(_mcd_dtw_sl)
+
+    return np.mean(mcd_dtw), np.mean(mcd_dtw_sl)
+
+
+def run_all_metrics(gen_dir, target_dir, whisper_model="openai/whisper-large-v3-turbo"):
+    """Run all evaluation metrics and return results"""
+    results = {}
+
+    print("Running WER evaluation...")
+    results["wer"] = wer_pipe(gen_dir, target_dir, model_id=whisper_model)
+
+    print("Running speaker similarity evaluation...")
+    results["speaker_similarity"] = spk_sim_pipe(gen_dir, target_dir)
+
+    print("Running MCD evaluation...")
+    mcd_dtw, mcd_dtw_sl = mcd_pipe(gen_dir, target_dir)
+    results["mcd_dtw"] = mcd_dtw
+    results["mcd_dtw_sl"] = mcd_dtw_sl
+
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Audio evaluation metrics")
+    parser.add_argument("--gen_dir", type=str, required=True, help="Directory containing generated audio files")
+    parser.add_argument("--target_dir", type=str, required=True, help="Directory containing target audio files")
+    parser.add_argument(
+        "--metric",
+        type=str,
+        default="all",
+        choices=["wer", "spk_sim", "mcd", "all"],
+        help="Evaluation metric to use",
+    )
+    parser.add_argument(
+        "--whisper_model",
+        type=str,
+        default="openai/whisper-large-v3-turbo",
+        help="Whisper model to use for WER evaluation",
+    )
+    # python eval.py --gen_dir path/to/generated --target_dir path/to/target
+    # keep the name of gen_wav and target_wav the same
+    args = parser.parse_args()
+
+    gen_dir = args.gen_dir
+    target_dir = args.target_dir
+
+    if not os.path.exists(gen_dir):
+        raise ValueError(f"Generated audio directory does not exist: {gen_dir}")
+    if not os.path.exists(target_dir):
+        raise ValueError(f"Target audio directory does not exist: {target_dir}")
+
+    if args.metric == "all":
+        results = run_all_metrics(gen_dir, target_dir, args.whisper_model)
+        print("\nEvaluation Results:")
+        print(f"WER: {results['wer']:.4f}")
+        print(f"Speaker Similarity: {results['speaker_similarity']:.4f}")
+        print(f"MCD (DTW): {results['mcd_dtw']:.4f}")
+        print(f"MCD (DTW-SL): {results['mcd_dtw_sl']:.4f}")
+
+    elif args.metric == "wer":
+        wer = wer_pipe(gen_dir, target_dir, model_id=args.whisper_model)
+        print(f"WER: {wer:.4f}")
+
+    elif args.metric == "spk_sim":
+        spk_sim = spk_sim_pipe(gen_dir, target_dir)
+        print(f"Speaker Similarity: {spk_sim:.4f}")
+
+    elif args.metric == "mcd":
+        mcd_dtw, mcd_dtw_sl = mcd_pipe(gen_dir, target_dir)
+        print(f"MCD (DTW): {mcd_dtw:.4f}")
+        print(f"MCD (DTW-SL): {mcd_dtw_sl:.4f}")

src/moviedubber/infer/basic.toml

@@ -0,0 +1,4 @@
+ckpt_file = "/path/to/ckpt_file.pth"
+vocab_file = "/path/to/vocab_file.txt"
+vocoder_local_path = "/path/to/bigvgan"
+campplus_path = "/path/to/campplus.onnx"

src/moviedubber/infer/utils_infer.py

@@ -0,0 +1,399 @@
+# A unified script for inference process
+# Make adjustments inside functions, and consider both gradio and cli scripts if need to change func output format
+
+import re
+from importlib.resources import files
+
+import matplotlib
+
+
+matplotlib.use("Agg")
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchaudio
+import tqdm
+
+from src.moviedubber.model import CFM
+from src.moviedubber.model.utils import convert_char_to_pinyin, get_tokenizer
+
+
+device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+
+# -----------------------------------------
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+win_length = 1024
+n_fft = 1024
+mel_spec_type = "bigvgan"
+target_rms = 0.1
+cross_fade_duration = 0.15
+ode_method = "euler"
+nfe_step = 32  # 16, 32
+# cfg_strength = 2.0
+cfg_strength = 1
+sway_sampling_coef = -1.0
+speed = 1.0
+fix_duration = None
+
+# -----------------------------------------
+
+
+# chunk text into smaller pieces
+
+
+def chunk_text(text, max_chars=135):
+    """
+    Splits the input text into chunks, each with a maximum number of characters.
+
+    Args:
+        text (str): The text to be split.
+        max_chars (int): The maximum number of characters per chunk.
+
+    Returns:
+        List[str]: A list of text chunks.
+    """
+    chunks = []
+    current_chunk = ""
+    # Split the text into sentences based on punctuation followed by whitespace
+    sentences = re.split(r"(?<=[;:,.!?])\s+|(?<=[；：，。！？])", text)
+
+    for sentence in sentences:
+        if len(current_chunk.encode("utf-8")) + len(sentence.encode("utf-8")) <= max_chars:
+            current_chunk += sentence + " " if sentence and len(sentence[-1].encode("utf-8")) == 1 else sentence
+        else:
+            if current_chunk:
+                chunks.append(current_chunk.strip())
+            current_chunk = sentence + " " if sentence and len(sentence[-1].encode("utf-8")) == 1 else sentence
+
+    if current_chunk:
+        chunks.append(current_chunk.strip())
+
+    return chunks
+
+
+# load vocoder
+def load_vocoder(local_path, device=device):
+    from src.third_party.BigVGAN import bigvgan
+
+    vocoder = bigvgan.BigVGAN.from_pretrained(local_path, use_cuda_kernel=False)
+
+    vocoder.remove_weight_norm()
+    vocoder = vocoder.eval().to(device)
+    return vocoder
+
+
+# load model checkpoint for inference
+
+
+def load_checkpoint(model, ckpt_path, device: str, dtype=None, use_ema=True):
+    if dtype is None:
+        dtype = (
+            torch.float16
+            if "cuda" in device
+            and torch.cuda.get_device_properties(device).major >= 6
+            and not torch.cuda.get_device_name().endswith("[ZLUDA]")
+            else torch.float32
+        )
+    model = model.to(dtype)
+
+    ckpt_type = ckpt_path.split(".")[-1]
+    if ckpt_type == "safetensors":
+        from safetensors.torch import load_file
+
+        checkpoint = load_file(ckpt_path, device=device)
+    else:
+        checkpoint = torch.load(ckpt_path, map_location=device, weights_only=True)
+
+    if use_ema:
+        if ckpt_type == "safetensors":
+            checkpoint = {"ema_model_state_dict": checkpoint}
+        checkpoint["model_state_dict"] = {
+            k.replace("ema_model.", ""): v
+            for k, v in checkpoint["ema_model_state_dict"].items()
+            if k not in ["initted", "step"]
+        }
+
+        # patch for backward compatibility, 305e3ea
+        for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
+            if key in checkpoint["model_state_dict"]:
+                del checkpoint["model_state_dict"][key]
+
+        state_dict_result = model.load_state_dict(checkpoint["model_state_dict"], strict=False)
+        if state_dict_result.unexpected_keys:
+            print("\nUnexpected keys in state_dict:", state_dict_result.unexpected_keys)
+        if state_dict_result.missing_keys:
+            print("\nMissing keys in state_dict:", state_dict_result.missing_keys)
+    else:
+        if ckpt_type == "safetensors":
+            checkpoint = {"model_state_dict": checkpoint}
+        model.load_state_dict(checkpoint["model_state_dict"], strict=True)
+
+    del checkpoint
+    torch.cuda.empty_cache()
+
+    return model.to(device)
+
+
+# load model for inference
+
+
+def load_model(
+    model_cls,
+    model_cfg,
+    ckpt_path,
+    controlnet=None,
+    mel_spec_type=mel_spec_type,
+    vocab_file="",
+    ode_method=ode_method,
+    use_ema=True,
+    device=device,
+):
+    tokenizer = "custom"
+
+    print("\nvocab : ", vocab_file)
+    print("token : ", tokenizer)
+    print("model : ", ckpt_path, "\n")
+
+    vocab_char_map, vocab_size = get_tokenizer(vocab_file, tokenizer)
+
+    if controlnet is not None:
+        controlnet = controlnet(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels)
+
+    model = CFM(
+        transformer=model_cls(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels),
+        mel_spec_kwargs=dict(
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            n_mel_channels=n_mel_channels,
+            target_sample_rate=target_sample_rate,
+            mel_spec_type=mel_spec_type,
+        ),
+        odeint_kwargs=dict(
+            method=ode_method,
+        ),
+        vocab_char_map=vocab_char_map,
+        controlnet=controlnet,
+    ).to(device)
+
+    dtype = torch.float32 if mel_spec_type == "bigvgan" else None
+    model = load_checkpoint(model, ckpt_path, device, dtype=dtype, use_ema=use_ema)
+
+    return model
+
+
+def infer_process(
+    ref_audio,
+    ref_text,
+    ref_clip,
+    ref_lip,
+    gen_text,
+    gen_clip,
+    gen_lip,
+    model_obj,
+    vocoder,
+    gen_caption=None,
+    mel_spec_type=mel_spec_type,
+    progress=tqdm,
+    target_rms=target_rms,
+    cross_fade_duration=cross_fade_duration,
+    nfe_step=nfe_step,
+    cfg_strength=cfg_strength,
+    sway_sampling_coef=sway_sampling_coef,
+    speed=speed,
+    fix_duration=fix_duration,
+    device=device,
+):
+    # Split the input text into batches
+    audio, sr = torchaudio.load(ref_audio)
+    max_chars = int(len(ref_text.encode("utf-8")) / (audio.shape[-1] / sr) * (25 - audio.shape[-1] / sr))
+    gen_text_batches = chunk_text(gen_text, max_chars=max_chars)
+
+    return infer_batch_process(
+        (audio, sr),
+        ref_text,
+        ref_clip,
+        ref_lip,
+        gen_text_batches,
+        gen_clip,
+        gen_lip,
+        model_obj,
+        vocoder,
+        gen_caption=gen_caption,
+        mel_spec_type=mel_spec_type,
+        progress=progress,
+        target_rms=target_rms,
+        cross_fade_duration=cross_fade_duration,
+        nfe_step=nfe_step,
+        cfg_strength=cfg_strength,
+        sway_sampling_coef=sway_sampling_coef,
+        speed=speed,
+        fix_duration=fix_duration,
+        device=device,
+    )
+
+
+# infer batches
+
+
+def infer_batch_process(
+    ref_audio,
+    ref_text,
+    ref_clip,
+    ref_lip,
+    gen_text_batches,
+    gen_clip,
+    gen_lip,
+    model_obj,
+    vocoder,
+    gen_caption=None,
+    mel_spec_type="vocos",
+    target_rms=0.1,
+    cross_fade_duration=0.15,
+    nfe_step=32,
+    cfg_strength=2.0,
+    sway_sampling_coef=-1,
+    speed=1,
+    fix_duration=None,
+    device=None,
+):
+    audio, sr = ref_audio
+    if audio.shape[0] > 1:
+        audio = torch.mean(audio, dim=0, keepdim=True)
+
+    rms = torch.sqrt(torch.mean(torch.square(audio)))
+    if rms < target_rms:
+        audio = audio * target_rms / rms
+    if sr != target_sample_rate:
+        resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
+        audio = resampler(audio)
+    audio = audio.to(device)
+
+    generated_waves = []
+    spectrograms = []
+
+    if len(ref_text[-1].encode("utf-8")) == 1:
+        ref_text = ref_text + " "
+
+    for i, gen_text in enumerate(gen_text_batches):
+        # Prepare the text
+        text_list = [ref_text + gen_text]
+        final_text_list = convert_char_to_pinyin(text_list)
+
+        ref_audio_len = audio.shape[-1] // hop_length
+        if ref_clip is not None:
+            ref_clip = F.interpolate(
+                ref_clip.unsqueeze(0).transpose(1, 2), size=ref_audio_len, mode="linear", align_corners=False
+            )
+        else:
+            ref_clip = torch.zeros(1, 768, ref_audio_len).to(device)
+
+        if fix_duration is not None:
+            duration = int(fix_duration * target_sample_rate / hop_length)
+            gen_audio_len = duration - ref_audio_len
+
+            gen_clip = F.interpolate(
+                gen_clip.unsqueeze(0).transpose(1, 2), size=gen_audio_len, mode="linear", align_corners=False
+            )
+
+        else:
+            # Calculate duration
+            ref_text_len = len(ref_text.encode("utf-8"))
+            gen_text_len = len(gen_text.encode("utf-8"))
+
+            gen_audio_len = int(ref_audio_len / ref_text_len * gen_text_len)
+            gen_clip = F.interpolate(
+                gen_clip.unsqueeze(0).transpose(1, 2), size=gen_audio_len, mode="linear", align_corners=False
+            )
+
+            duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
+
+        if ref_lip is None:
+            ref_lip = torch.zeros(ref_audio_len // 4, 512)
+
+        clip = torch.cat([ref_clip, gen_clip], dim=-1).permute(0, 2, 1).to(device)
+
+        if gen_lip is not None:
+            lip = torch.cat([ref_lip.unsqueeze(0).transpose(1, 2), gen_lip.unsqueeze(0).transpose(1, 2)], dim=-1).to(
+                device
+            )
+            lip = F.pad(lip, (0, duration - lip.size(-1)), value=0).permute(0, 2, 1)
+        else:
+            lip = None
+
+        # inference
+        with torch.inference_mode():
+            generated, _ = model_obj.sample(
+                cond=audio,
+                text=final_text_list,
+                clip=clip,
+                lip=lip,
+                caption_emb=gen_caption,
+                duration=duration,
+                steps=nfe_step,
+                cfg_strength=cfg_strength,
+                sway_sampling_coef=sway_sampling_coef,
+                no_ref_audio=False,
+            )
+
+            generated = generated.to(torch.float32)
+            generated = generated[:, ref_audio_len:, :]
+            generated_mel_spec = generated.permute(0, 2, 1)
+            if mel_spec_type == "vocos":
+                generated_wave = vocoder.decode(generated_mel_spec)
+            elif mel_spec_type == "bigvgan":
+                generated_wave = vocoder(generated_mel_spec)
+            if rms < target_rms:
+                generated_wave = generated_wave * rms / target_rms
+
+            # wav -> numpy
+            generated_wave = generated_wave.squeeze().cpu().numpy()
+
+            generated_waves.append(generated_wave)
+            spectrograms.append(generated_mel_spec[0].cpu().numpy())
+
+    # Combine all generated waves with cross-fading
+    if cross_fade_duration <= 0:
+        # Simply concatenate
+        final_wave = np.concatenate(generated_waves)
+    else:
+        final_wave = generated_waves[0]
+        for i in range(1, len(generated_waves)):
+            prev_wave = final_wave
+            next_wave = generated_waves[i]
+
+            # Calculate cross-fade samples, ensuring it does not exceed wave lengths
+            cross_fade_samples = int(cross_fade_duration * target_sample_rate)
+            cross_fade_samples = min(cross_fade_samples, len(prev_wave), len(next_wave))
+
+            if cross_fade_samples <= 0:
+                # No overlap possible, concatenate
+                final_wave = np.concatenate([prev_wave, next_wave])
+                continue
+
+            # Overlapping parts
+            prev_overlap = prev_wave[-cross_fade_samples:]
+            next_overlap = next_wave[:cross_fade_samples]
+
+            # Fade out and fade in
+            fade_out = np.linspace(1, 0, cross_fade_samples)
+            fade_in = np.linspace(0, 1, cross_fade_samples)
+
+            # Cross-faded overlap
+            cross_faded_overlap = prev_overlap * fade_out + next_overlap * fade_in
+
+            # Combine
+            new_wave = np.concatenate(
+                [prev_wave[:-cross_fade_samples], cross_faded_overlap, next_wave[cross_fade_samples:]]
+            )
+
+            final_wave = new_wave
+
+    # Create a combined spectrogram
+    combined_spectrogram = np.concatenate(spectrograms, axis=1)
+
+    return final_wave, target_sample_rate, combined_spectrogram

src/moviedubber/infer/video_preprocess.py

@@ -0,0 +1,315 @@
+import argparse
+import glob
+import logging
+import os
+import os.path as osp
+from pathlib import Path
+from typing import Optional, Union
+
+import cv2
+import imageio
+import numpy as np
+import torch
+import torch.multiprocessing as mp
+from decord import AudioReader, VideoReader, cpu
+from PIL import Image
+from tqdm import tqdm
+from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
+
+
+logging.basicConfig(level=logging.ERROR, format="%(asctime)s - %(levelname)s - %(message)s")
+
+
+NUM_FRAMES = None  # NUM_FRAMES = 160
+MAX_FRAMES = None  # MAX_FRAMES = 256
+NUM_FRAMES_PER_SECOND = 10
+
+
+def get_full_indices(reader: Union[VideoReader, AudioReader]) -> np.ndarray:
+    if isinstance(reader, VideoReader):
+        return np.linspace(0, len(reader) - 1, len(reader), dtype=int)
+    elif isinstance(reader, AudioReader):
+        return np.linspace(0, reader.shape[-1] - 1, reader.shape[-1], dtype=int)
+
+
+def create_output_directories(output_dir):
+    try:
+        os.makedirs(osp.join(output_dir, "audio"), exist_ok=True)
+        os.makedirs(osp.join(output_dir, "video"), exist_ok=True)
+    except OSError as e:
+        print(f"Error creating directories: {e}")
+        raise
+
+
+def frame_sample(duration, mode="uniform", num_frames=None, fps=None):
+    if mode == "uniform":
+        assert num_frames is not None, "Number of frames must be provided for uniform sampling."
+        # NOTE: v1 version
+        # Calculate the size of each segment from which a frame will be extracted
+        seg_size = float(duration - 1) / num_frames
+
+        frame_ids = []
+        for i in range(num_frames):
+            # Calculate the start and end indices of each segment
+            start = seg_size * i
+            end = seg_size * (i + 1)
+            # Append the middle index of the segment to the list
+            frame_ids.append((start + end) / 2)
+
+        return np.round(np.array(frame_ids) + 1e-6).astype(int)
+        # NOTE: v0 version
+        # return np.linspace(0, duration-1, num_frames, dtype=int)
+    elif mode == "fps":
+        assert fps is not None, "FPS must be provided for FPS sampling."
+        segment_len = min(fps // NUM_FRAMES_PER_SECOND, duration)
+        return np.arange(segment_len // 2, duration, segment_len, dtype=int)
+    else:
+        raise ImportError(f"Unsupported frame sampling mode: {mode}")
+
+
+def expand2square(pil_img, background_color):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+def process_video(video_path, processor, s=None, e=None, aspect_ratio="pad", num_frames=NUM_FRAMES):
+    if isinstance(video_path, str):
+        if s is not None and e is not None:
+            s = s if s >= 0.0 else 0.0
+            e = e if e >= 0.0 else 0.0
+            if s > e:
+                s, e = e, s
+            elif s == e:
+                e = s + 1
+
+        # 1. Loading Video
+        if os.path.isdir(video_path):
+            frame_files = sorted(os.listdir(video_path))
+
+            fps = 3
+            num_frames_of_video = len(frame_files)
+        elif video_path.endswith(".gif"):
+            gif_reader = imageio.get_reader(video_path)
+
+            fps = 25
+            num_frames_of_video = len(gif_reader)
+        else:
+            try:
+                vreader = VideoReader(video_path, ctx=cpu(0), num_threads=1)
+            except:  # noqa: E722
+                return None
+
+            fps = vreader.get_avg_fps()
+            num_frames_of_video = len(vreader)
+
+        # 2. Determine frame range & Calculate frame indices
+        f_start = 0 if s is None else max(int(s * fps) - 1, 0)
+        f_end = num_frames_of_video - 1 if e is None else min(int(e * fps) - 1, num_frames_of_video - 1)
+        frame_indices = list(range(f_start, f_end + 1))
+
+        duration = len(frame_indices)
+        # 3. Sampling frame indices
+        if num_frames is None:
+            sampled_frame_indices = [frame_indices[i] for i in frame_sample(duration, mode="fps", fps=fps)]
+        else:
+            sampled_frame_indices = [
+                frame_indices[i] for i in frame_sample(duration, mode="uniform", num_frames=num_frames)
+            ]
+
+        # 4. Acquire frame data
+        if os.path.isdir(video_path):
+            video_data = [Image.open(os.path.join(video_path, frame_files[f_idx])) for f_idx in sampled_frame_indices]
+        elif video_path.endswith(".gif"):
+            video_data = [
+                Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_RGBA2RGB))
+                for idx, frame in enumerate(gif_reader)
+                if idx in sampled_frame_indices
+            ]
+        else:
+            video_data = [Image.fromarray(frame) for frame in vreader.get_batch(sampled_frame_indices).asnumpy()]
+
+    elif isinstance(video_path, np.ndarray):
+        video_data = [Image.fromarray(f) for f in video_path]
+    elif isinstance(video_path, list) and isinstance(video_path[0], np.ndarray):
+        video_data = [Image.fromarray(f) for f in video_path]
+    elif isinstance(video_path, list) and isinstance(video_path[0], str):
+        video_data = [Image.open(f) for f in video_path]
+    elif isinstance(video_path, list) and isinstance(video_path[0], Image.Image):
+        video_data = video_path
+    else:
+        raise ValueError(f"Unsupported video path type: {type(video_path)}")
+
+    while num_frames is not None and len(video_data) < num_frames:
+        video_data.append(Image.fromarray(np.zeros((*video_data[-1].size, 3), dtype=np.uint8)))
+
+    # MAX_FRAMES filter
+    if MAX_FRAMES:
+        video_data = video_data[:MAX_FRAMES]
+
+    if aspect_ratio == "pad":
+        images = [expand2square(f, tuple(int(x * 255) for x in processor.image_mean)) for f in video_data]
+    else:
+        images = list(video_data)
+    video = processor.preprocess(images, return_tensors="pt")["pixel_values"]
+    return video
+
+
+class VideoFeatureExtractor:
+    def __init__(
+        self,
+        pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = "openai/clip-vit-large-patch14",
+        device: str = "cuda",
+    ):
+        self.device = device
+
+        self.processor = CLIPImageProcessor.from_pretrained(pretrained_model_name_or_path)
+        self.model = CLIPVisionModelWithProjection.from_pretrained(pretrained_model_name_or_path).to(self.device).half()
+
+    def extract_features(self, video_path):
+        images = process_video(video_path, self.processor)
+        if images is None:
+            return None
+        clip_feature = self.model(images.to(self.device).half()).image_embeds
+
+        return clip_feature
+
+
+def video_processor(item, feature_extractor, output_dir=None):
+    video_path = Path(item)
+    if not os.path.exists(video_path):
+        return
+
+    clip_feature = feature_extractor.extract_features(str(video_path))
+    if clip_feature is None:
+        return
+
+    if output_dir is not None and not os.path.exists(output_dir):
+        os.makedirs(output_dir, exist_ok=True)
+        output_path = osp.join(output_dir, f"{video_path.stem}.pt")
+    else:
+        output_path = video_path.with_suffix(".clip")
+
+    torch.save(clip_feature, output_path)
+
+
+def s_thread(items, id, device, output_dir):
+    feature_extractor = VideoFeatureExtractor(device=device)
+    for i, data in tqdm(enumerate(items), total=len(items), position=id):
+        video_processor(data, feature_extractor, output_dir)
+
+
+def load_tensor(file_path, map_location="cpu", weights_only=True):
+    try:
+        return torch.load(file_path, map_location=map_location, weights_only=weights_only)
+    except FileNotFoundError:
+        logging.error(f"File not found: {file_path}")
+    except torch.serialization.pickle.UnpicklingError:
+        logging.error(f"Failed to unpickle file: {file_path}")
+    except Exception as e:
+        logging.error(f"An error occurred while loading {file_path}: {e}")
+    return None
+
+
+def post_check(directory):
+    if not osp.isdir(directory):
+        logging.error(f"Invalid directory: {directory}")
+        return
+
+    video_dir = osp.join(directory, "video")
+    pt_files = glob.glob(f"{video_dir}/*.pt")
+
+    for file_path in tqdm(pt_files):
+        embeds = load_tensor(file_path)
+        if embeds is None:
+            continue
+
+        audio_file_path = file_path.replace("video", "audio")
+        audio_text_embeds = load_tensor(audio_file_path)
+        if audio_text_embeds is None:
+            logging.error(f"Failed to load audio file: {audio_file_path}")
+            continue
+
+        text = audio_text_embeds.get("text")
+        mel = audio_text_embeds.get("mel")
+        if text is None or mel is None:
+            logging.error(f"Missing 'text' or 'mel' in {audio_file_path}")
+
+
+def args_parse():
+    args = argparse.ArgumentParser()
+    args.add_argument("--data_type", "-d", type=str, default="video", help="'audio' or 'video'")
+    args.add_argument("--check", action="store_true", help="post check, if any pt file was damaged")
+
+    args.add_argument(
+        "--num_threads",
+        "-n",
+        type=int,
+        default=1,
+        required=False,
+        help="num_threads",
+    )
+    args.add_argument(
+        "--input",
+        "-i",
+        type=str,
+        required=True,
+        help="input file path",
+    )
+    args.add_argument(
+        "--output_dir",
+        "-o",
+        type=str,
+        default=None,
+        help="output folder path",
+    )
+    args.add_argument("--multi_gpu", "-m", nargs="+", type=str, default=None, required=False, help="GPU ids")
+
+    args = args.parse_args()
+    return args
+
+
+if __name__ == "__main__":
+    args_main = args_parse()
+
+    if args_main.check:
+        post_check(args_main.output_dir)
+        exit(0)
+
+    gpu_ids = ["cuda:0"]
+    if args_main.multi_gpu is not None:
+        gpu_ids = [f"cuda:{gpu}" for gpu in args_main.multi_gpu]
+
+    output_dir = args_main.output_dir
+    if output_dir is not None:
+        create_output_directories(output_dir)
+
+    rows = None
+    rows = [it.strip() for it in Path(args_main.input).read_text().split("\n") if it.strip() != ""]
+
+    chunks = np.array_split(rows, args_main.num_threads)
+    chunks = [chunk.tolist() for chunk in chunks]
+
+    processes = []
+    mp.set_start_method("spawn", force=True)
+    for idx, chunk in enumerate(chunks):
+        device = gpu_ids[idx % len(gpu_ids)]
+        p = mp.Process(target=s_thread, args=(chunk, idx, device, output_dir))
+        processes.append(p)
+        p.start()
+
+    for process in processes:
+        process.join()
+
+    # DEBUG
+    # s_thread(args_main, input_dir, output_dir, chunks[0], 0, "cuda:0")
+
+    print("process done!")

src/moviedubber/infer_with_mmlm_result.py

@@ -0,0 +1,339 @@
+import argparse
+import os
+import os.path as osp
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import onnxruntime
+import soundfile
+import tomli
+import torch
+import torch.nn.functional as F
+import torchaudio
+import torchaudio.compliance.kaldi as kaldi
+from moviepy import AudioFileClip, VideoFileClip
+from omegaconf import OmegaConf
+from pydub import AudioSegment
+from tqdm import tqdm
+
+
+src_path = Path(osp.dirname(__file__)).parent.parent
+sys.path.insert(0, str(src_path))
+sys.path.append(str(src_path / "src/third_party/BigVGAN"))
+
+from src.moviedubber.infer.utils_infer import (
+    cfg_strength,
+    chunk_text,
+    load_model,
+    load_vocoder,
+    mel_spec_type,
+    nfe_step,
+    sway_sampling_coef,
+)
+from src.moviedubber.infer.video_preprocess import VideoFeatureExtractor
+from src.moviedubber.model import ControlNetDiT, DiT
+from src.moviedubber.model.utils import convert_char_to_pinyin
+
+
+def concat_movie_with_audio(wav, video_path, out_dir):
+    if not os.path.exists(wav):
+        raise FileNotFoundError(f"Audio file {wav} does not exist")
+
+    if not os.path.exists(video_path):
+        raise FileNotFoundError(f"Video file {video_path} does not exist")
+
+    try:
+        with (
+            AudioFileClip(str(wav)) as audio_clip,
+            VideoFileClip(str(video_path)) as video_clip,
+        ):
+            duration = min(video_clip.duration, audio_clip.duration)
+
+            video_subclip = video_clip.subclipped(0, duration)
+            audio_subclip = audio_clip.subclipped(0, duration)
+
+            final_video = video_subclip.with_audio(audio_subclip)
+
+            output_path = wav.replace(".wav", ".mp4")
+
+            final_video.write_videofile(
+                str(output_path),
+                codec="libx264",
+                audio_codec="mp3",
+                fps=25,
+                logger=None,
+                threads=1,
+                temp_audiofile_path=out_dir,
+            )
+
+    except Exception as e:
+        print(f"Error processing {wav} {video_path}: {str(e)}")
+
+    return output_path
+
+
+def get_spk_emb(audio_path, ort_session):
+    audio, sample_rate = torchaudio.load(str(audio_path))
+    if sample_rate != 16000:
+        audio = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(audio)
+    feat = kaldi.fbank(audio, num_mel_bins=80, dither=0, sample_frequency=16000)
+    feat = feat - feat.mean(dim=0, keepdim=True)
+    embedding = (
+        ort_session.run(None, {ort_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()})[0]
+        .flatten()
+        .tolist()
+    )
+    return embedding
+
+
+def load_models(config, device):
+    model_cfg = config.get("model_cfg", "src/moviedubber/configs/basemodel.yaml")
+    ckpt_file = config.get("ckpt_file", None)
+    campplus_path = config.get("campplus_path", None)
+    vocab_file = config.get("vocab_file", None)
+
+    vocoder_local_path = config.get("vocoder_local_path", None)
+
+    if ckpt_file is None or vocab_file is None or vocoder_local_path is None or campplus_path is None:
+        raise ValueError("ckpt_file, vocab_file and vocoder_local_path must be specified")
+
+    vocoder_name = config.get("vocoder_name", mel_spec_type)
+
+    vocoder = load_vocoder(local_path=vocoder_local_path, device=device)
+
+    model_cls = DiT
+    model_cfg = OmegaConf.load(model_cfg).model.arch
+    controlnet = ControlNetDiT
+
+    ema_model = load_model(
+        model_cls,
+        model_cfg,
+        ckpt_file,
+        mel_spec_type=vocoder_name,
+        vocab_file=vocab_file,
+        controlnet=controlnet,
+        device=device,
+    )
+
+    option = onnxruntime.SessionOptions()
+    option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+    option.intra_op_num_threads = 1
+    providers = ["CPUExecutionProvider"]
+    ort_session = onnxruntime.InferenceSession(
+        campplus_path,
+        sess_options=option,
+        providers=providers,
+    )
+    return ema_model, vocoder, ort_session
+
+
+def main(config, device, chunk, gen_dir, target_dir, out_dir, idx):
+    ema_model, vocoder, ort_session = load_models(config, device=device)
+
+    videofeature_extractor = VideoFeatureExtractor(device=device)
+
+    for it in tqdm(chunk, total=len(chunk), position=idx, desc=f"Processing {idx}"):
+        wav, video, text, ref_wav = it
+
+        with open(f"{target_dir}/{wav.split('/')[-1].split('.')[0]}.txt", "a") as f:
+            f.write(text + "\n")
+
+        if wav.endswith(".mp3"):
+            audio = AudioSegment.from_mp3(wav)
+
+            wav_file = wav.replace(".mp3", ".wav")
+            audio.export(wav_file, format="wav")
+
+        wav = Path(wav).with_suffix(".wav")
+        if wav.exists() is False:
+            continue
+
+        os.system(f"cp {wav} {target_dir}/")
+
+        gen_audio, sr = torchaudio.load(str(wav))
+        resampler = torchaudio.transforms.Resample(sr, 24000)
+        if sr != 24000:
+            gen_audio = resampler(gen_audio)
+
+        if gen_audio.shape[0] > 1:
+            gen_audio = torch.mean(gen_audio, dim=0, keepdim=True)
+
+        gen_video = video
+        gen_clip_path = gen_video.replace(".mp4", ".clip")
+
+        if not os.path.exists(gen_clip_path):
+            gen_clip = videofeature_extractor.extract_features(gen_video)
+
+            torch.save(gen_clip.detach().cpu(), gen_clip_path)
+
+        else:
+            gen_clip = torch.load(gen_clip_path, weights_only=True).to(device=device, dtype=torch.float32)
+
+        if ref_wav == "None":
+            use_ref_audio = False
+            gen_text_ = text
+
+            gen_clip_ = gen_clip
+
+            ref_audio_ = gen_audio
+
+            spk_emb = torch.zeros(1, 1, 192).to(device=device, dtype=torch.float32)
+
+        else:
+            use_ref_audio = True
+            ref_audio = Path(ref_wav)
+
+            spk_emb = get_spk_emb(ref_audio, ort_session)
+            spk_emb = torch.tensor(spk_emb).to(device=device, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+
+            ref_text = ref_audio.with_suffix(".txt").read_text().strip()
+            gen_text_ = ref_text + " " + text
+
+            if ref_audio.exists() is False:
+                raise Exception(f"ref_audio {ref_audio} not found")
+
+            if ref_audio.suffix == ".mp3":
+                audio = AudioSegment.from_mp3(ref_audio)
+
+                wav_file = ref_audio.with_suffix(".wav")
+                audio.export(wav_file, format="wav")
+
+            ref_audio_, _ = torchaudio.load(str(ref_audio.with_suffix(".wav")))
+            resampler = torchaudio.transforms.Resample(sr, 24000)
+            if sr != 24000:
+                ref_audio_ = resampler(ref_audio_)
+
+            if ref_audio_.shape[0] > 1:
+                ref_audio_ = torch.mean(ref_audio_, dim=0, keepdim=True)
+
+            ref_video = ref_audio.with_suffix(".mp4")
+            ref_clip_path = ref_video.with_suffix(".clip")
+
+            if not ref_clip_path.exists():
+                ref_clip = videofeature_extractor.extract_features(str(ref_video))
+
+                torch.save(ref_clip.detach().cpu(), ref_clip_path)
+
+            else:
+                ref_clip = torch.load(ref_clip_path, weights_only=True).to(device=device, dtype=torch.float32)
+
+            gen_clip_ = torch.cat([ref_clip, gen_clip], dim=0)
+
+        gen_audio_len = gen_audio.shape[1] // 256
+
+        if use_ref_audio:
+            ref_audio_len = ref_audio_.shape[1] // 256
+            duration = ref_audio_len + gen_audio_len
+        else:
+            duration = gen_audio_len
+
+        gen_clip_ = gen_clip_.unsqueeze(0).to(device=device, dtype=torch.float32).transpose(1, 2)
+        gen_clip_ = F.interpolate(gen_clip_, size=duration, mode="linear", align_corners=False).transpose(1, 2)
+
+        gen_text_batches = chunk_text(gen_text_, max_chars=1024)
+        final_text_list = convert_char_to_pinyin(gen_text_batches)
+
+        with torch.inference_mode():
+            generated, _ = ema_model.sample(
+                cond=ref_audio_.to(device),
+                text=final_text_list,
+                clip=gen_clip_,
+                spk_emb=spk_emb,
+                duration=duration,
+                steps=nfe_step,
+                cfg_strength=cfg_strength,
+                sway_sampling_coef=sway_sampling_coef,
+                no_ref_audio=not use_ref_audio,
+            )
+
+            generated = generated.to(torch.float32)
+
+            if use_ref_audio:
+                generated = generated[:, ref_audio_len:, :]
+
+            generated_mel_spec = generated.permute(0, 2, 1)
+            generated_wave = vocoder(generated_mel_spec)
+
+            generated_wave = generated_wave.squeeze().cpu().numpy()
+
+            out_path = osp.join(gen_dir, f"{wav.stem}.wav")
+            soundfile.write(out_path, generated_wave, samplerate=24000)
+            _ = concat_movie_with_audio(out_path, gen_video, out_dir)
+
+
+if __name__ == "__main__":
+    import torch.multiprocessing as mp
+
+    parser = argparse.ArgumentParser(
+        prog="python3 infer-cli.py",
+        description="Commandline interface for moviedubber infer with Advanced Batch Processing.",
+        epilog="Specify options above to override one or more settings from config.",
+    )
+    parser.add_argument(
+        "-c",
+        "--config",
+        type=str,
+        default="src/moviedubber/infer/basic.toml",
+        help="The configuration file, default see infer/basic.toml",
+    )
+    parser.add_argument("-i", "--input_list", type=str, required=True, help="The val list file")
+    parser.add_argument("-s", "--ref_spk_list", type=str, required=True, help="The spk list file")
+    parser.add_argument("-o", "--out_dir", type=str, default="data/dubberout", help="The output directory")
+    parser.add_argument("--gpuids", type=str, help="GPU ids to use, split by comma")
+    parser.add_argument("--nums_workers", type=int, default=1, help="Number of workers for per gpu")
+
+    args = parser.parse_args()
+
+    out_dir = args.out_dir
+    input_list = args.input_list
+    gpu_ids = args.gpuids.split(",") if args.gpuids else ["0"]
+    num_pre = args.nums_workers
+    spk_ref_path = args.ref_spk_list
+
+    config = tomli.load(open(args.config, "rb"))
+
+    gen_lst = Path(input_list).read_text().splitlines()[1:]
+
+    gen_pre_conf = []
+
+    spk_lines = Path(spk_ref_path).read_text().splitlines()
+
+    for idx, line in enumerate(gen_lst):
+        if line.strip():
+            mp4_path, is_correc, _, _ = line.split(",")
+
+            wav_path = mp4_path.replace(".mp4", ".mp3")
+            text = Path(wav_path.replace(".mp3", ".txt")).read_text().strip()
+
+            if is_correc == "True":
+                ref_wav = spk_lines[idx].split(",")[1].strip()
+            else:
+                ref_wav = random.choice(spk_lines).split(",")[-1].strip()  # Use random speaker for incorrect samples
+
+            gen_pre_conf.append([wav_path, mp4_path, text, ref_wav])
+
+    chunks = np.array_split(gen_pre_conf, len(gpu_ids) * num_pre)
+
+    gen_dir = os.path.join(out_dir, "generated")
+    target_dir = os.path.join(out_dir, "target")
+
+    if os.path.exists(gen_dir) is False or os.path.exists(target_dir) is False:
+        os.makedirs(gen_dir)
+        os.makedirs(target_dir)
+
+    mp.set_start_method("spawn", force=True)
+    processes = []
+    for idx, chunk in enumerate(chunks):
+        device = gpu_ids[idx % len(gpu_ids)]
+
+        device = f"cuda:{device}"
+        p = mp.Process(target=main, args=(config, device, chunk, gen_dir, target_dir, out_dir, idx))
+        processes.append(p)
+        p.start()
+
+    for process in processes:
+        process.join()
+
+    print("All processes finished.")

src/moviedubber/model/__init__.py

@@ -0,0 +1,5 @@
+from .cfm import CFM
+from .dit import ControlNetDiT, DiT
+
+
+__all__ = ["CFM", "UNetT", "DiT", "ControlNetDiT", "MMDiT", "Trainer"]

src/moviedubber/model/cfm.py

@@ -0,0 +1,209 @@
+# modified from https://github.com/SWivid/F5-TTS/blob/main/src/f5_tts/model/cfm.py
+
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.utils.rnn import pad_sequence
+from torchdiffeq import odeint
+
+from .modules import MelSpec
+from .utils import (
+    default,
+    exists,
+    lens_to_mask,
+    list_str_to_idx,
+    list_str_to_tensor,
+)
+
+
+class CFM(nn.Module):
+    def __init__(
+        self,
+        transformer: nn.Module,
+        sigma=0.0,
+        odeint_kwargs: dict = dict(
+            method="euler"  # 'midpoint'
+        ),
+        num_channels=None,
+        mel_spec_module: nn.Module | None = None,
+        mel_spec_kwargs: dict = dict(),
+        frac_lengths_mask: tuple[float, float] = (0.7, 1.0),
+        vocab_char_map: dict[str:int] | None = None,
+        controlnet: nn.Module | None = None,
+    ):
+        super().__init__()
+
+        self.frac_lengths_mask = frac_lengths_mask
+
+        self.mel_spec = default(mel_spec_module, MelSpec(**mel_spec_kwargs))
+        num_channels = default(num_channels, self.mel_spec.n_mel_channels)
+        self.num_channels = num_channels
+
+        self.transformer = transformer
+        dim = transformer.dim
+        self.dim = dim
+
+        self.sigma = sigma
+
+        self.odeint_kwargs = odeint_kwargs
+
+        self.vocab_char_map = vocab_char_map
+
+        self.controlnet = controlnet
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @torch.no_grad()
+    def sample(
+        self,
+        cond: float["b n d"] | float["b nw"],  # noqa: F722
+        text: int["b nt"] | list[str],  # noqa: F722
+        clip: float["b n d"],  # noqa: F722
+        duration: int | int["b"],  # noqa: F821
+        *,
+        caption_emb: float["b n d"] | None = None,  # noqa: F722
+        spk_emb: float["b n d"] | None = None,  # noqa: F722
+        lens: int["b"] | None = None,  # noqa: F821
+        steps=32,
+        cfg_strength=1.0,
+        sway_sampling_coef=None,
+        seed: int | None = None,
+        max_duration=4096,
+        vocoder: Callable[[float["b d n"]], float["b nw"]] | None = None,  # noqa: F722
+        no_ref_audio=False,
+        duplicate_test=False,
+        t_inter=0.1,
+        edit_mask=None,
+    ):
+        self.eval()
+
+        if cond.ndim == 2:
+            cond = self.mel_spec(cond)
+            cond = cond.permute(0, 2, 1)
+            assert cond.shape[-1] == self.num_channels
+
+        cond = cond.to(next(self.parameters()).dtype)
+
+        batch, cond_seq_len, device = *cond.shape[:2], cond.device
+        if not exists(lens):
+            lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long)
+
+        if isinstance(text, list):
+            if exists(self.vocab_char_map):
+                text = list_str_to_idx(text, self.vocab_char_map).to(device)
+            else:
+                text = list_str_to_tensor(text).to(device)
+            assert text.shape[0] == batch
+
+        if exists(text):
+            text_lens = (text != -1).sum(dim=-1)
+            lens = torch.maximum(text_lens, lens)
+
+        cond_mask = lens_to_mask(lens)
+        if edit_mask is not None:
+            cond_mask = cond_mask & edit_mask
+
+        if isinstance(duration, int):
+            duration = torch.full((batch,), duration, device=device, dtype=torch.long)
+
+        # duration = torch.maximum(lens + 1, duration)
+
+        duration = duration.clamp(max=max_duration)
+        max_duration = duration.amax()
+
+        if duplicate_test:
+            test_cond = F.pad(cond, (0, 0, cond_seq_len, max_duration - 2 * cond_seq_len), value=0.0)
+
+        cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0)
+        cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value=False)
+        cond_mask = cond_mask.unsqueeze(-1)
+        step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))
+
+        if batch > 1:
+            mask = lens_to_mask(duration)
+        else:
+            mask = None
+
+        if no_ref_audio:
+            cond = torch.zeros_like(cond)
+
+        def fn(t, x):
+            step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))
+
+            controlnet_embeds = self.controlnet(
+                x=x,
+                text=text,
+                clip=clip,
+                spk_emb=spk_emb,
+                caption=caption_emb,
+                time=t,
+            )
+
+            cond_pred = self.transformer(
+                x=x,
+                cond=step_cond,
+                text=text,
+                time=t,
+                mask=mask,
+                drop_audio_cond=[False],
+                drop_text=[False],
+                controlnet_embeds=controlnet_embeds,
+            )
+
+            null_pred = self.transformer(
+                x=x,
+                cond=step_cond,
+                text=text,
+                time=t,
+                mask=mask,
+                drop_audio_cond=[True],
+                drop_text=[True],
+                controlnet_embeds=None,
+            )
+
+            return null_pred + (cond_pred - null_pred) * 2
+
+        y0 = []
+        for dur in duration:
+            if exists(seed):
+                torch.manual_seed(seed)
+            y0.append(torch.randn(dur, self.num_channels, device=self.device, dtype=step_cond.dtype))
+        y0 = pad_sequence(y0, padding_value=0, batch_first=True)
+
+        t_start = 0
+
+        if duplicate_test:
+            t_start = t_inter
+            y0 = (1 - t_start) * y0 + t_start * test_cond
+            steps = int(steps * (1 - t_start))
+
+        t = torch.linspace(t_start, 1, steps + 1, device=self.device, dtype=step_cond.dtype)
+        if sway_sampling_coef is not None:
+            t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
+
+        trajectory = odeint(fn, y0, t, **self.odeint_kwargs)
+
+        sampled = trajectory[-1]
+        out = sampled
+        out = torch.where(cond_mask, cond, out)
+
+        if exists(vocoder):
+            out = out.permute(0, 2, 1)
+            out = vocoder(out)
+
+        return out, trajectory

src/moviedubber/model/dit.py

@@ -0,0 +1,297 @@
+# modified from https://github.com/SWivid/F5-TTS/blob/main/src/f5_tts/model/backbones/dit.py
+
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from x_transformers.x_transformers import RotaryEmbedding
+
+from .modules import (
+    AdaLayerNormZero_Final,
+    ConvNeXtV2Block,
+    ConvPositionEmbedding,
+    DiTBlock,
+    TimestepEmbedding,
+    get_pos_embed_indices,
+    precompute_freqs_cis,
+)
+
+
+# Text embedding
+
+
+class TextEmbedding(nn.Module):
+    def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2):
+        super().__init__()
+        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token
+
+        if conv_layers > 0:
+            self.extra_modeling = True
+            self.precompute_max_pos = 4096  # ~44s of 24khz audio
+            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)
+            self.text_blocks = nn.Sequential(
+                *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]
+            )
+        else:
+            self.extra_modeling = False
+
+    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F722
+        text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()
+        text = text[:, :seq_len]  # curtail if character tokens are more than the mel spec tokens
+        batch, text_len = text.shape[0], text.shape[1]
+        text = F.pad(text, (0, seq_len - text_len), value=0)
+
+        for idx, _drop in enumerate(drop_text):  # cfg for text
+            if _drop:
+                text[idx] = torch.zeros_like(text[idx])
+
+        text = self.text_embed(text)  # b n -> b n d
+
+        # possible extra modeling
+        if self.extra_modeling:
+            # sinus pos emb
+            batch_start = torch.zeros((batch,), dtype=torch.long)
+            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)
+            text_pos_embed = self.freqs_cis[pos_idx]
+            text = text + text_pos_embed
+
+            # convnextv2 blocks
+            text = self.text_blocks(text)
+
+        return text
+
+
+# noised input audio and context mixing embedding
+
+
+class InputEmbedding(nn.Module):
+    def __init__(self, mel_dim, text_dim, out_dim):
+        super().__init__()
+        self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)
+        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
+
+    def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], drop_audio_cond=False):  # noqa: F722
+        for idx, _drop in enumerate(drop_audio_cond):  # cfg for cond audio
+            if _drop:
+                cond[idx] = torch.zeros_like(cond[idx])
+
+        x = self.proj(torch.cat((x, cond, text_embed), dim=-1))
+        x = self.conv_pos_embed(x) + x
+        return x
+
+
+class InputEmbeddingO(nn.Module):
+    def __init__(self, mel_dim, text_dim, out_dim):
+        super().__init__()
+        self.proj = nn.Linear(mel_dim + 512 + text_dim + 192 + 32, out_dim)
+        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
+
+    def forward(
+        self,
+        x: float["b n d"],  # noqa: F722
+        text_emb: float["b n d"],  # noqa: F722
+        video_emb: float["b n d"],  # noqa: F722
+        spk_emb: float["b n d"],  # noqa: F722
+        caption_emb: float["b n d"],  # noqa: F722
+    ):
+        x = self.proj(torch.cat((x, text_emb, video_emb, spk_emb, caption_emb), dim=-1))
+        x = self.conv_pos_embed(x) + x
+        return x
+
+
+# Transformer backbone using DiT blocks
+
+
+class DiT(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth=8,
+        heads=8,
+        dim_head=64,
+        dropout=0.1,
+        ff_mult=4,
+        mel_dim=100,
+        text_num_embeds=256,
+        text_dim=None,
+        conv_layers=0,
+        long_skip_connection=False,
+    ):
+        super().__init__()
+
+        self.time_embed = TimestepEmbedding(dim)
+        if text_dim is None:
+            text_dim = mel_dim
+        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers)
+        self.input_embed = InputEmbedding(mel_dim, text_dim, dim)
+
+        self.rotary_embed = RotaryEmbedding(dim_head)
+
+        self.dim = dim
+        self.depth = depth
+
+        self.transformer_blocks = nn.ModuleList(
+            [DiTBlock(dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, dropout=dropout) for _ in range(depth)]
+        )
+        self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None
+
+        self.norm_out = AdaLayerNormZero_Final(dim)  # final modulation
+        self.proj_out = nn.Linear(dim, mel_dim)
+
+    def forward(
+        self,
+        x: float["b n d"],  # nosied input audio  # noqa: F722
+        cond: float["b n d"],  # masked cond audio  # noqa: F722
+        text: int["b nt"],  # text  # noqa: F722
+        time: float["b"] | float[""],  # time step  # noqa: F821 F722
+        drop_audio_cond,  # cfg for cond audio
+        drop_text,  # cfg for text
+        mask: bool["b n"] | None = None,  # noqa: F722
+        controlnet_embeds: float["b n d"] | None = None,  # noqa: F722
+    ):
+        batch, seq_len = x.shape[0], x.shape[1]
+        if time.ndim == 0:
+            time = time.repeat(batch)
+
+        t = self.time_embed(time)
+        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)
+        x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)
+
+        rope = self.rotary_embed.forward_from_seq_len(seq_len)
+
+        if self.long_skip_connection is not None:
+            residual = x
+
+        for i, block in enumerate(self.transformer_blocks):
+            if controlnet_embeds is not None and i < 12:
+                x += controlnet_embeds[i]
+
+            x = block(x, t, mask=mask, rope=rope)
+
+        if self.long_skip_connection is not None:
+            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))
+
+        x = self.norm_out(x, t)
+        output = self.proj_out(x)
+
+        return output
+
+
+class ControlNetDiT(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth=8,
+        heads=8,
+        dim_head=64,
+        dropout=0.1,
+        ff_mult=4,
+        mel_dim=100,
+        text_num_embeds=256,
+        text_dim=None,
+        conv_layers=0,
+        long_skip_connection=False,
+        checkpoint_activations=False,
+        duration_predictor=None,
+    ):
+        super().__init__()
+
+        if text_dim is None:
+            text_dim = mel_dim
+
+        self.time_embed = TimestepEmbedding(dim)
+
+        self.rotary_embed = RotaryEmbedding(dim_head)
+
+        self.dim = dim
+        self.depth = depth // 2 + 1
+
+        self.transformer_blocks1 = nn.ModuleList(
+            [
+                DiTBlock(dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, dropout=dropout)
+                for _ in range(self.depth)
+            ]
+        )
+
+        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers)
+
+        self.input_embed = InputEmbeddingO(mel_dim, text_dim, dim)
+
+        self.spk_embed_affine_layer = torch.nn.Linear(192, 192)
+        self.clip_embed_affine_layer = torch.nn.Linear(768, 512)
+        self.caption_embed_affine_layer = torch.nn.Linear(512, 32)
+
+        self.zero_linear = nn.ModuleList([nn.Linear(dim, dim, bias=False) for _ in range(12)])
+        for zero_linear in self.zero_linear:
+            nn.init.zeros_(zero_linear.weight)
+
+        self.duration_predictor = duration_predictor
+
+    def forward(
+        self,
+        x: float["b n d"],  # nosied input audio  # noqa: F722
+        text: int["b nt"],  # text  # noqa: F722
+        clip: float["b n d"],  # video clip # noqa: F722
+        spk_emb: float["b d"],  # speaker embedding  # noqa: F722
+        time: float["b"] | float[""],  # time step  # noqa: F821 F722
+        caption: float["b nt"] | None = None,  # caption  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        lens: int["b"] | None = None,  # noqa: F722, F821
+        return_dur: bool = False,  # return duration prediction
+    ):
+        batch, seq_len = x.shape[0], x.shape[1]
+
+        if time.ndim == 0:
+            time = time.repeat(batch)
+
+        t = self.time_embed(time)
+
+        clip_emb = F.normalize(clip, dim=-1)
+        clip_emb = self.clip_embed_affine_layer(clip)
+
+        spk_emb = F.normalize(spk_emb, dim=-1)
+        spk_emb = self.spk_embed_affine_layer(spk_emb)
+        spk_emb = torch.repeat_interleave(spk_emb, seq_len, dim=1)
+
+        if caption is None:
+            caption = torch.zeros(1, seq_len, 512).to(device=x.device)
+
+        caption_emb = F.normalize(caption, dim=-1)
+        caption_emb = self.caption_embed_affine_layer(caption_emb)
+
+        text_embed = self.text_embed(text, seq_len, drop_text=[False])
+
+        x = self.input_embed(x, text_embed, clip_emb, spk_emb, caption_emb)
+
+        rope = self.rotary_embed.forward_from_seq_len(seq_len)
+
+        info = []
+        for i, block in enumerate(self.transformer_blocks1):
+            x = block(x, t, mask=mask, rope=rope)  # 'b n 1024'
+
+            info.append(x)
+
+        out_info = []
+        for i, linear in enumerate(self.zero_linear):
+            h = linear(info[i])
+            out_info.append(h)
+
+        if return_dur and self.duration_predictor is not None:
+            dur_loss = self.duration_predictor(x=x, text=clip_emb, lens=lens)
+
+            return out_info, dur_loss
+
+        else:
+            return out_info

src/moviedubber/model/modules.py

@@ -0,0 +1,467 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from librosa.filters import mel as librosa_mel_fn
+from torch import nn
+from x_transformers.x_transformers import apply_rotary_pos_emb
+
+
+# raw wav to mel spec
+
+
+mel_basis_cache = {}
+hann_window_cache = {}
+
+
+def get_bigvgan_mel_spectrogram(
+    waveform,
+    n_fft=1024,
+    n_mel_channels=100,
+    target_sample_rate=24000,
+    hop_length=256,
+    win_length=1024,
+    fmin=0,
+    fmax=None,
+    center=False,
+):  # Copy from https://github.com/NVIDIA/BigVGAN/tree/main
+    device = waveform.device
+    key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}"
+
+    if key not in mel_basis_cache:
+        mel = librosa_mel_fn(sr=target_sample_rate, n_fft=n_fft, n_mels=n_mel_channels, fmin=fmin, fmax=fmax)
+        mel_basis_cache[key] = torch.from_numpy(mel).float().to(device)  # TODO: why they need .float()?
+        hann_window_cache[key] = torch.hann_window(win_length).to(device)
+
+    mel_basis = mel_basis_cache[key]
+    hann_window = hann_window_cache[key]
+
+    padding = (n_fft - hop_length) // 2
+    waveform = torch.nn.functional.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
+
+    spec = torch.stft(
+        waveform,
+        n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=hann_window,
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=True,
+    )
+    spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
+
+    mel_spec = torch.matmul(mel_basis, spec)
+    mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5))
+
+    return mel_spec
+
+
+class MelSpec(nn.Module):
+    def __init__(
+        self,
+        n_fft=1024,
+        hop_length=256,
+        win_length=1024,
+        n_mel_channels=100,
+        target_sample_rate=24_000,
+        mel_spec_type="bigvgan",
+    ):
+        super().__init__()
+
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_mel_channels = n_mel_channels
+        self.target_sample_rate = target_sample_rate
+
+        self.extractor = get_bigvgan_mel_spectrogram
+
+        self.register_buffer("dummy", torch.tensor(0), persistent=False)
+
+    def forward(self, wav):
+        if self.dummy.device != wav.device:
+            self.to(wav.device)
+
+        mel = self.extractor(
+            waveform=wav,
+            n_fft=self.n_fft,
+            n_mel_channels=self.n_mel_channels,
+            target_sample_rate=self.target_sample_rate,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+        )
+
+        return mel
+
+
+# sinusoidal position embedding
+
+
+class SinusPositionEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x, scale=1000):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+# convolutional position embedding
+
+
+class ConvPositionEmbedding(nn.Module):
+    def __init__(self, dim, kernel_size=31, groups=16):
+        super().__init__()
+        assert kernel_size % 2 != 0
+        self.conv1d = nn.Sequential(
+            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
+            nn.Mish(),
+            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
+            nn.Mish(),
+        )
+
+    def forward(self, x: float["b n d"], mask: bool["b n"] | None = None):  # noqa: F722
+        if mask is not None:
+            mask = mask[..., None]
+            x = x.masked_fill(~mask, 0.0)
+
+        x = x.permute(0, 2, 1)
+        x = self.conv1d(x)
+        out = x.permute(0, 2, 1)
+
+        if mask is not None:
+            out = out.masked_fill(~mask, 0.0)
+
+        return out
+
+
+# rotary positional embedding related
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):
+    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+    # has some connection to NTK literature
+    # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
+    # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
+    theta *= theta_rescale_factor ** (dim / (dim - 2))
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t, freqs).float()  # type: ignore
+    freqs_cos = torch.cos(freqs)  # real part
+    freqs_sin = torch.sin(freqs)  # imaginary part
+    return torch.cat([freqs_cos, freqs_sin], dim=-1)
+
+
+def get_pos_embed_indices(start, length, max_pos, scale=1.0):
+    # length = length if isinstance(length, int) else length.max()
+    scale = scale * torch.ones_like(start, dtype=torch.float32)  # in case scale is a scalar
+    pos = (
+        start.unsqueeze(1)
+        + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()
+    )
+    # avoid extra long error.
+    pos = torch.where(pos < max_pos, pos, max_pos - 1)
+    return pos
+
+
+# Global Response Normalization layer (Instance Normalization ?)
+
+
+class GRN(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=1, keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py
+# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108
+
+
+class ConvNeXtV2Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        dilation: int = 1,
+    ):
+        super().__init__()
+        padding = (dilation * (7 - 1)) // 2
+        self.dwconv = nn.Conv1d(
+            dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation
+        )  # depthwise conv
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.grn = GRN(intermediate_dim)
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = x.transpose(1, 2)  # b n d -> b d n
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # b d n -> b n d
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        return residual + x
+
+
+# AdaLayerNormZero
+# return with modulated x for attn input, and params for later mlp modulation
+
+
+class AdaLayerNormZero(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 6)
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb=None):
+        emb = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
+
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+# AdaLayerNormZero for final layer
+# return only with modulated x for attn input, cuz no more mlp modulation
+
+
+class AdaLayerNormZero_Final(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 2)
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb):
+        emb = self.linear(self.silu(emb))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+# FeedForward
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        activation = nn.GELU(approximate=approximate)
+        project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
+        self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
+
+    def forward(self, x):
+        return self.ff(x)
+
+
+# Attention with possible joint part
+# modified from diffusers/src/diffusers/models/attention_processor.py
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        processor: AttnProcessor,
+        dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        context_dim: Optional[int] = None,  # if not None -> joint attention
+        context_pre_only=None,
+    ):
+        super().__init__()
+
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+        self.processor = processor
+
+        self.dim = dim
+        self.heads = heads
+        self.inner_dim = dim_head * heads
+        self.dropout = dropout
+
+        self.context_dim = context_dim
+        self.context_pre_only = context_pre_only
+
+        self.to_q = nn.Linear(dim, self.inner_dim)
+        self.to_k = nn.Linear(dim, self.inner_dim)
+        self.to_v = nn.Linear(dim, self.inner_dim)
+
+        if self.context_dim is not None:
+            self.to_k_c = nn.Linear(context_dim, self.inner_dim)
+            self.to_v_c = nn.Linear(context_dim, self.inner_dim)
+            if self.context_pre_only is not None:
+                self.to_q_c = nn.Linear(context_dim, self.inner_dim)
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(nn.Linear(self.inner_dim, dim))
+        self.to_out.append(nn.Dropout(dropout))
+
+        if self.context_pre_only is not None and not self.context_pre_only:
+            self.to_out_c = nn.Linear(self.inner_dim, dim)
+
+    def forward(
+        self,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        c: float["b n d"] = None,  # context c  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding for x
+        c_rope=None,  # rotary position embedding for c
+    ) -> torch.Tensor:
+        if c is not None:
+            return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope)
+        else:
+            return self.processor(self, x, mask=mask, rope=rope)
+
+
+# Attention processor
+
+
+class AttnProcessor:
+    def __init__(self):
+        pass
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding
+    ) -> torch.FloatTensor:
+        batch_size = x.shape[0]
+
+        # `sample` projections.
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # apply rotary position embedding
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+
+            query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+            key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        # attention
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # mask. e.g. inference got a batch with different target durations, mask out the padding
+        if mask is not None:
+            attn_mask = mask
+            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # 'b n -> b 1 1 n'
+            attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
+        else:
+            attn_mask = None
+
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        x = x.to(query.dtype)
+
+        # linear proj
+        x = attn.to_out[0](x)
+        # dropout
+        x = attn.to_out[1](x)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.0)
+
+        return x
+
+
+# DiT Block
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1):
+        super().__init__()
+
+        self.attn_norm = AdaLayerNormZero(dim)
+        self.attn = Attention(
+            processor=AttnProcessor(),
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            dropout=dropout,
+        )
+
+        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+
+    def forward(self, x, t, mask=None, rope=None):  # x: noised input, t: time embedding
+        # pre-norm & modulation for attention input
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
+
+        # attention
+        attn_output = self.attn(x=norm, mask=mask, rope=rope)
+
+        # process attention output for input x
+        x = x + gate_msa.unsqueeze(1) * attn_output
+
+        norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff(norm)
+        x = x + gate_mlp.unsqueeze(1) * ff_output
+
+        return x
+
+
+# time step conditioning embedding
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, dim, freq_embed_dim=256):
+        super().__init__()
+        self.time_embed = SinusPositionEmbedding(freq_embed_dim)
+        self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+
+    def forward(self, timestep: float["b"]):  # noqa: F821
+        time_hidden = self.time_embed(timestep)
+        time_hidden = time_hidden.to(timestep.dtype)
+        time = self.time_mlp(time_hidden)  # b d
+        return time

src/moviedubber/model/utils.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+
+import jieba
+import torch
+from pypinyin import Style, lazy_pinyin
+from torch.nn.utils.rnn import pad_sequence
+
+
+def exists(v):
+    return v is not None
+
+
+def default(v, d):
+    return v if exists(v) else d
+
+
+# tensor helpers
+
+
+def lens_to_mask(t: int["b"], length: int | None = None) -> bool["b n"]:  # noqa: F722 F821
+    if not exists(length):
+        length = t.amax()
+
+    seq = torch.arange(length, device=t.device)
+    return seq[None, :] < t[:, None]
+
+
+# simple utf-8 tokenizer, since paper went character based
+def list_str_to_tensor(text: list[str], padding_value=-1) -> int["b nt"]:  # noqa: F722
+    list_tensors = [torch.tensor([*bytes(t, "UTF-8")]) for t in text]  # ByT5 style
+    text = pad_sequence(list_tensors, padding_value=padding_value, batch_first=True)
+    return text
+
+
+# char tokenizer, based on custom dataset's extracted .txt file
+def list_str_to_idx(
+    text: list[str] | list[list[str]],
+    vocab_char_map: dict[str, int],  # {char: idx}
+    padding_value=-1,
+) -> int["b nt"]:  # noqa: F722
+    list_idx_tensors = [torch.tensor([vocab_char_map.get(c, 0) for c in t]) for t in text]  # pinyin or char style
+    text = pad_sequence(list_idx_tensors, padding_value=padding_value, batch_first=True)
+    return text
+
+
+# Get tokenizer
+
+
+def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):
+    """
+    tokenizer   - "pinyin" do g2p for only chinese characters, need .txt vocab_file
+                - "char" for char-wise tokenizer, need .txt vocab_file
+                - "byte" for utf-8 tokenizer
+                - "custom" if you're directly passing in a path to the vocab.txt you want to use
+    vocab_size  - if use "pinyin", all available pinyin types, common alphabets (also those with accent) and symbols
+                - if use "char", derived from unfiltered character & symbol counts of custom dataset
+                - if use "byte", set to 256 (unicode byte range)
+    """
+    if tokenizer in ["pinyin", "char"]:
+        # tokenizer_path = os.path.join(files("f5_tts").joinpath("../../data"), f"{dataset_name}_{tokenizer}/vocab.txt")
+        tokenizer_path = "/ailab-train/speech/zhengjunjie/huggingface/models/F5-TTS/F5TTS_Base/vocab.txt"
+        print(f"Loading {tokenizer} tokenizer from {tokenizer_path}")
+        with open(tokenizer_path, "r", encoding="utf-8") as f:
+            vocab_char_map = {}
+            for i, char in enumerate(f):
+                vocab_char_map[char[:-1]] = i
+        vocab_size = len(vocab_char_map)
+        assert vocab_char_map[" "] == 0, "make sure space is of idx 0 in vocab.txt, cuz 0 is used for unknown char"
+
+    elif tokenizer == "byte":
+        vocab_char_map = None
+        vocab_size = 256
+
+    elif tokenizer == "custom":
+        with open(dataset_name, "r", encoding="utf-8") as f:
+            vocab_char_map = {}
+            for i, char in enumerate(f):
+                vocab_char_map[char[:-1]] = i
+        vocab_size = len(vocab_char_map)
+
+    return vocab_char_map, vocab_size
+
+
+# convert char to pinyin
+
+jieba.initialize()
+print("Word segmentation module jieba initialized.\n")
+
+
+def convert_char_to_pinyin(text_list, polyphone=True):
+    final_text_list = []
+    custom_trans = str.maketrans(
+        {";": ",", "“": '"', "”": '"', "‘": "'", "’": "'"}
+    )  # add custom trans here, to address oov
+
+    def is_chinese(c):
+        return (
+            "\u3100" <= c <= "\u9fff"  # common chinese characters
+        )
+
+    for text in text_list:
+        char_list = []
+        text = text.translate(custom_trans)
+        for seg in jieba.cut(text):
+            seg_byte_len = len(bytes(seg, "UTF-8"))
+            if seg_byte_len == len(seg):  # if pure alphabets and symbols
+                if char_list and seg_byte_len > 1 and char_list[-1] not in " :'\"":
+                    char_list.append(" ")
+                char_list.extend(seg)
+            elif polyphone and seg_byte_len == 3 * len(seg):  # if pure east asian characters
+                seg_ = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True)
+                for i, c in enumerate(seg):
+                    if is_chinese(c):
+                        char_list.append(" ")
+                    char_list.append(seg_[i])
+            else:  # if mixed characters, alphabets and symbols
+                for c in seg:
+                    if ord(c) < 256:
+                        char_list.extend(c)
+                    elif is_chinese(c):
+                        char_list.append(" ")
+                        char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True))
+                    else:
+                        char_list.append(c)
+        final_text_list.append(char_list)
+
+    return final_text_list

src/third_party/BigVGAN/.gitignore

@@ -0,0 +1,146 @@
+# BigVGAN
+alias_free_activation/cuda/build/
+exp/
+tmp/
+
+# Symlinks for bundled LibriTTS filelists
+filelists/LibriTTS/train-clean-100
+filelists/LibriTTS/train-clean-360
+filelists/LibriTTS/train-other-500
+filelists/LibriTTS/dev-clean
+filelists/LibriTTS/dev-other
+filelists/LibriTTS/test-clean
+filelists/LibriTTS/test-other
+
+# VSCode configs
+.vscode/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+.idea/

src/third_party/BigVGAN/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 NVIDIA CORPORATION.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

src/third_party/BigVGAN/README.md

@@ -0,0 +1,266 @@
+## BigVGAN: A Universal Neural Vocoder with Large-Scale Training
+
+#### Sang-gil Lee, Wei Ping, Boris Ginsburg, Bryan Catanzaro, Sungroh Yoon
+
+[[Paper]](https://arxiv.org/abs/2206.04658) - [[Code]](https://github.com/NVIDIA/BigVGAN) - [[Showcase]](https://bigvgan-demo.github.io/) - [[Project Page]](https://research.nvidia.com/labs/adlr/projects/bigvgan/) - [[Weights]](https://huggingface.co/collections/nvidia/bigvgan-66959df3d97fd7d98d97dc9a) - [[Demo]](https://huggingface.co/spaces/nvidia/BigVGAN)
+
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bigvgan-a-universal-neural-vocoder-with-large/speech-synthesis-on-libritts)](https://paperswithcode.com/sota/speech-synthesis-on-libritts?p=bigvgan-a-universal-neural-vocoder-with-large)
+
+<center><img src="https://user-images.githubusercontent.com/15963413/218609148-881e39df-33af-4af9-ab95-1427c4ebf062.png" width="800"></center>
+
+## News
+- **Sep 2024 (v2.4):**
+  - We have updated the pretrained checkpoints trained for 5M steps. This is final release of the BigVGAN-v2 checkpoints.
+
+- **Jul 2024 (v2.3):**
+  - General refactor and code improvements for improved readability.
+  - Fully fused CUDA kernel of anti-alised activation (upsampling + activation + downsampling) with inference speed benchmark.
+
+- **Jul 2024 (v2.2):** The repository now includes an interactive local demo using gradio.
+
+- **Jul 2024 (v2.1):** BigVGAN is now integrated with 🤗 Hugging Face Hub with easy access to inference using pretrained checkpoints. We also provide an interactive demo on Hugging Face Spaces.
+
+- **Jul 2024 (v2):** We release BigVGAN-v2 along with pretrained checkpoints. Below are the highlights:
+  - Custom CUDA kernel for inference: we provide a fused upsampling + activation kernel written in CUDA for accelerated inference speed. Our test shows 1.5 - 3x faster speed on a single A100 GPU.
+  - Improved discriminator and loss: BigVGAN-v2 is trained using a [multi-scale sub-band CQT discriminator](https://arxiv.org/abs/2311.14957) and a [multi-scale mel spectrogram loss](https://arxiv.org/abs/2306.06546).
+  - Larger training data: BigVGAN-v2 is trained using datasets containing diverse audio types, including speech in multiple languages, environmental sounds, and instruments.
+  - We provide pretrained checkpoints of BigVGAN-v2 using diverse audio configurations, supporting up to 44 kHz sampling rate and 512x upsampling ratio.
+
+## Installation
+
+The codebase has been tested on Python `3.10` and PyTorch `2.3.1` conda packages with either `pytorch-cuda=12.1` or `pytorch-cuda=11.8`. Below is an example command to create the conda environment:
+
+```shell
+conda create -n bigvgan python=3.10 pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
+conda activate bigvgan
+```
+
+Clone the repository and install dependencies:
+
+```shell
+git clone https://github.com/NVIDIA/BigVGAN
+cd BigVGAN
+pip install -r requirements.txt
+```
+
+## Inference Quickstart using 🤗 Hugging Face Hub
+
+Below example describes how you can use BigVGAN: load the pretrained BigVGAN generator from Hugging Face Hub, compute mel spectrogram from input waveform, and generate synthesized waveform using the mel spectrogram as the model's input.
+
+```python
+device = 'cuda'
+
+import torch
+import bigvgan
+import librosa
+from meldataset import get_mel_spectrogram
+
+# instantiate the model. You can optionally set use_cuda_kernel=True for faster inference.
+model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_24khz_100band_256x', use_cuda_kernel=False)
+
+# remove weight norm in the model and set to eval mode
+model.remove_weight_norm()
+model = model.eval().to(device)
+
+# load wav file and compute mel spectrogram
+wav_path = '/path/to/your/audio.wav'
+wav, sr = librosa.load(wav_path, sr=model.h.sampling_rate, mono=True) # wav is np.ndarray with shape [T_time] and values in [-1, 1]
+wav = torch.FloatTensor(wav).unsqueeze(0) # wav is FloatTensor with shape [B(1), T_time]
+
+# compute mel spectrogram from the ground truth audio
+mel = get_mel_spectrogram(wav, model.h).to(device) # mel is FloatTensor with shape [B(1), C_mel, T_frame]
+
+# generate waveform from mel
+with torch.inference_mode():
+    wav_gen = model(mel) # wav_gen is FloatTensor with shape [B(1), 1, T_time] and values in [-1, 1]
+wav_gen_float = wav_gen.squeeze(0).cpu() # wav_gen is FloatTensor with shape [1, T_time]
+
+# you can convert the generated waveform to 16 bit linear PCM
+wav_gen_int16 = (wav_gen_float * 32767.0).numpy().astype('int16') # wav_gen is now np.ndarray with shape [1, T_time] and int16 dtype
+```
+
+## Local gradio demo <a href='https://github.com/gradio-app/gradio'><img src='https://img.shields.io/github/stars/gradio-app/gradio'></a>
+
+You can run a local gradio demo using below command:
+
+```python
+pip install -r demo/requirements.txt
+python demo/app.py
+```
+
+## Training
+
+Create symbolic link to the root of the dataset. The codebase uses filelist with the relative path from the dataset. Below are the example commands for LibriTTS dataset:
+
+```shell
+cd filelists/LibriTTS && \
+ln -s /path/to/your/LibriTTS/train-clean-100 train-clean-100 && \
+ln -s /path/to/your/LibriTTS/train-clean-360 train-clean-360 && \
+ln -s /path/to/your/LibriTTS/train-other-500 train-other-500 && \
+ln -s /path/to/your/LibriTTS/dev-clean dev-clean && \
+ln -s /path/to/your/LibriTTS/dev-other dev-other && \
+ln -s /path/to/your/LibriTTS/test-clean test-clean && \
+ln -s /path/to/your/LibriTTS/test-other test-other && \
+cd ../..
+```
+
+Train BigVGAN model. Below is an example command for training BigVGAN-v2 using LibriTTS dataset at 24kHz with a full 100-band mel spectrogram as input:
+
+```shell
+python train.py \
+--config configs/bigvgan_v2_24khz_100band_256x.json \
+--input_wavs_dir filelists/LibriTTS \
+--input_training_file filelists/LibriTTS/train-full.txt \
+--input_validation_file filelists/LibriTTS/val-full.txt \
+--list_input_unseen_wavs_dir filelists/LibriTTS filelists/LibriTTS \
+--list_input_unseen_validation_file filelists/LibriTTS/dev-clean.txt filelists/LibriTTS/dev-other.txt \
+--checkpoint_path exp/bigvgan_v2_24khz_100band_256x
+```
+
+## Synthesis
+
+Synthesize from BigVGAN model. Below is an example command for generating audio from the model.
+It computes mel spectrograms using wav files from `--input_wavs_dir` and saves the generated audio to `--output_dir`.
+
+```shell
+python inference.py \
+--checkpoint_file /path/to/your/bigvgan_v2_24khz_100band_256x/bigvgan_generator.pt \
+--input_wavs_dir /path/to/your/input_wav \
+--output_dir /path/to/your/output_wav
+```
+
+`inference_e2e.py` supports synthesis directly from the mel spectrogram saved in `.npy` format, with shapes `[1, channel, frame]` or `[channel, frame]`.
+It loads mel spectrograms from `--input_mels_dir` and saves the generated audio to `--output_dir`.
+
+Make sure that the STFT hyperparameters for mel spectrogram are the same as the model, which are defined in `config.json` of the corresponding model.
+
+```shell
+python inference_e2e.py \
+--checkpoint_file /path/to/your/bigvgan_v2_24khz_100band_256x/bigvgan_generator.pt \
+--input_mels_dir /path/to/your/input_mel \
+--output_dir /path/to/your/output_wav
+```
+
+## Using Custom CUDA Kernel for Synthesis
+
+You can apply the fast CUDA inference kernel by using a parameter `use_cuda_kernel` when instantiating BigVGAN:
+
+```python
+generator = BigVGAN(h, use_cuda_kernel=True)
+```
+
+You can also pass `--use_cuda_kernel` to `inference.py` and `inference_e2e.py` to enable this feature.
+
+When applied for the first time, it builds the kernel using `nvcc` and `ninja`. If the build succeeds, the kernel is saved to `alias_free_activation/cuda/build` and the model automatically loads the kernel. The codebase has been tested using CUDA `12.1`.
+
+Please make sure that both are installed in your system and `nvcc` installed in your system matches the version your PyTorch build is using.
+
+We recommend running `test_cuda_vs_torch_model.py` first to build and check the correctness of the CUDA kernel. See below example command and its output, where it returns `[Success] test CUDA fused vs. plain torch BigVGAN inference`:
+
+```python
+python tests/test_cuda_vs_torch_model.py \
+--checkpoint_file /path/to/your/bigvgan_generator.pt
+```
+
+```shell
+loading plain Pytorch BigVGAN
+...
+loading CUDA kernel BigVGAN with auto-build
+Detected CUDA files, patching ldflags
+Emitting ninja build file /path/to/your/BigVGAN/alias_free_activation/cuda/build/build.ninja..
+Building extension module anti_alias_activation_cuda...
+...
+Loading extension module anti_alias_activation_cuda...
+...
+Loading '/path/to/your/bigvgan_generator.pt'
+...
+[Success] test CUDA fused vs. plain torch BigVGAN inference
+ > mean_difference=0.0007238413265440613
+...
+```
+
+If you see `[Fail] test CUDA fused vs. plain torch BigVGAN inference`, it means that the CUDA kernel inference is incorrect. Please check if `nvcc` installed in your system is compatible with your PyTorch version.
+
+## Pretrained Models
+
+We provide the [pretrained models on Hugging Face Collections](https://huggingface.co/collections/nvidia/bigvgan-66959df3d97fd7d98d97dc9a).
+One can download the checkpoints of the generator weight (named `bigvgan_generator.pt`) and its discriminator/optimizer states (named `bigvgan_discriminator_optimizer.pt`) within the listed model repositories.
+
+| Model Name                                                                                               | Sampling Rate | Mel band | fmax  | Upsampling Ratio | Params | Dataset                    | Steps | Fine-Tuned |
+|:--------------------------------------------------------------------------------------------------------:|:-------------:|:--------:|:-----:|:----------------:|:------:|:--------------------------:|:-----:|:----------:|
+| [bigvgan_v2_44khz_128band_512x](https://huggingface.co/nvidia/bigvgan_v2_44khz_128band_512x)             | 44 kHz        | 128      | 22050 | 512              | 122M   | Large-scale Compilation    | 5M    | No         |
+| [bigvgan_v2_44khz_128band_256x](https://huggingface.co/nvidia/bigvgan_v2_44khz_128band_256x)             | 44 kHz        | 128      | 22050 | 256              | 112M   | Large-scale Compilation    | 5M    | No         |
+| [bigvgan_v2_24khz_100band_256x](https://huggingface.co/nvidia/bigvgan_v2_24khz_100band_256x)             | 24 kHz        | 100      | 12000 | 256              | 112M   | Large-scale Compilation    | 5M    | No         |
+| [bigvgan_v2_22khz_80band_256x](https://huggingface.co/nvidia/bigvgan_v2_22khz_80band_256x)               | 22 kHz        | 80       | 11025 | 256              | 112M   | Large-scale Compilation    | 5M    | No         |
+| [bigvgan_v2_22khz_80band_fmax8k_256x](https://huggingface.co/nvidia/bigvgan_v2_22khz_80band_fmax8k_256x) | 22 kHz        | 80       | 8000  | 256              | 112M   | Large-scale Compilation    | 5M    | No         |
+| [bigvgan_24khz_100band](https://huggingface.co/nvidia/bigvgan_24khz_100band)                             | 24 kHz        | 100      | 12000 | 256              | 112M   | LibriTTS                   | 5M    | No         |
+| [bigvgan_base_24khz_100band](https://huggingface.co/nvidia/bigvgan_base_24khz_100band)                   | 24 kHz        | 100      | 12000 | 256              | 14M    | LibriTTS                   | 5M    | No         |
+| [bigvgan_22khz_80band](https://huggingface.co/nvidia/bigvgan_22khz_80band)                               | 22 kHz        | 80       | 8000  | 256              | 112M   | LibriTTS + VCTK + LJSpeech | 5M    | No         |
+| [bigvgan_base_22khz_80band](https://huggingface.co/nvidia/bigvgan_base_22khz_80band)                     | 22 kHz        | 80       | 8000  | 256              | 14M    | LibriTTS + VCTK + LJSpeech | 5M    | No         |
+
+The paper results are based on the original 24kHz BigVGAN models (`bigvgan_24khz_100band` and `bigvgan_base_24khz_100band`) trained on LibriTTS dataset.
+We also provide 22kHz BigVGAN models with band-limited setup (i.e., fmax=8000) for TTS applications.
+Note that the checkpoints use `snakebeta` activation with log scale parameterization, which have the best overall quality.
+
+You can fine-tune the models by:
+
+1. downloading the checkpoints (both the generator weight and its discriminator/optimizer states)
+2. resuming training using your audio dataset by specifying `--checkpoint_path` that includes the checkpoints when launching `train.py`
+
+## Training Details of BigVGAN-v2
+
+Comapred to the original BigVGAN, the pretrained checkpoints of BigVGAN-v2 used `batch_size=32` with a longer `segment_size=65536` and are trained using 8 A100 GPUs.
+
+Note that the BigVGAN-v2 `json` config files in `./configs` use `batch_size=4` as default to fit in a single A100 GPU for training. You can fine-tune the models adjusting `batch_size` depending on your GPUs.
+
+When training BigVGAN-v2 from scratch with small batch size, it can potentially encounter the early divergence problem mentioned in the paper. In such case, we recommend lowering the `clip_grad_norm` value (e.g. `100`) for the early training iterations (e.g. 20000 steps) and increase the value to the default `500`.
+
+## Evaluation Results of BigVGAN-v2
+
+Below are the objective results of the 24kHz model (`bigvgan_v2_24khz_100band_256x`) obtained from the LibriTTS `dev` sets. BigVGAN-v2 shows noticeable improvements of the metrics. The model also exhibits reduced perceptual artifacts, especially for non-speech audio.
+
+| Model      | Dataset                 | Steps | PESQ(↑)   | M-STFT(↓)  | MCD(↓)     | Periodicity(↓) | V/UV F1(↑) |
+|:----------:|:-----------------------:|:-----:|:---------:|:----------:|:----------:|:--------------:|:----------:|
+| BigVGAN    | LibriTTS                | 1M    | 4.027     | 0.7997     | 0.3745     | 0.1018         | 0.9598     |
+| BigVGAN    | LibriTTS                | 5M    | 4.256     | 0.7409     | 0.2988     | 0.0809         | 0.9698     |
+| BigVGAN-v2 | Large-scale Compilation | 3M    | 4.359     | 0.7134     | 0.3060     | 0.0621         | 0.9777     |
+| BigVGAN-v2 | Large-scale Compilation | 5M    | **4.362** | **0.7026** | **0.2903** | **0.0593**     | **0.9793** |
+
+## Speed Benchmark
+
+Below are the speed and VRAM usage benchmark results of BigVGAN from `tests/test_cuda_vs_torch_model.py`, using `bigvgan_v2_24khz_100band_256x` as a reference model.
+
+| GPU                        | num_mel_frame | use_cuda_kernel | Speed (kHz) | Real-time Factor | VRAM (GB) |
+|:--------------------------:|:-------------:|:---------------:|:-----------:|:----------------:|:---------:|
+| NVIDIA A100                | 256           | False           | 1672.1      | 69.7x            | 1.3       |
+|                            |               | True            | 3916.5      | 163.2x           | 1.3       |
+|                            | 2048          | False           | 1899.6      | 79.2x            | 1.7       |
+|                            |               | True            | 5330.1      | 222.1x           | 1.7       |
+|                            | 16384         | False           | 1973.8      | 82.2x            | 5.0       |
+|                            |               | True            | 5761.7      | 240.1x           | 4.4       |
+| NVIDIA GeForce RTX 3080    | 256           | False           | 841.1       | 35.0x            | 1.3       |
+|                            |               | True            | 1598.1      | 66.6x            | 1.3       |
+|                            | 2048          | False           | 929.9       | 38.7x            | 1.7       |
+|                            |               | True            | 1971.3      | 82.1x            | 1.6       |
+|                            | 16384         | False           | 943.4       | 39.3x            | 5.0       |
+|                            |               | True            | 2026.5      | 84.4x            | 3.9       |
+| NVIDIA GeForce RTX 2080 Ti | 256           | False           | 515.6       | 21.5x            | 1.3       |
+|                            |               | True            | 811.3       | 33.8x            | 1.3       |
+|                            | 2048          | False           | 576.5       | 24.0x            | 1.7       |
+|                            |               | True            | 1023.0      | 42.6x            | 1.5       |
+|                            | 16384         | False           | 589.4       | 24.6x            | 5.0       |
+|                            |               | True            | 1068.1      | 44.5x            | 3.2       |
+
+## Acknowledgements
+
+We thank Vijay Anand Korthikanti and Kevin J. Shih for their generous support in implementing the CUDA kernel for inference.
+
+## References
+
+- [HiFi-GAN](https://github.com/jik876/hifi-gan) (for generator and multi-period discriminator)
+- [Snake](https://github.com/EdwardDixon/snake) (for periodic activation)
+- [Alias-free-torch](https://github.com/junjun3518/alias-free-torch) (for anti-aliasing)
+- [Julius](https://github.com/adefossez/julius) (for low-pass filter)
+- [UnivNet](https://github.com/mindslab-ai/univnet) (for multi-resolution discriminator)
+- [descript-audio-codec](https://github.com/descriptinc/descript-audio-codec) and [vocos](https://github.com/gemelo-ai/vocos) (for multi-band multi-scale STFT discriminator and multi-scale mel spectrogram loss)
+- [Amphion](https://github.com/open-mmlab/Amphion) (for multi-scale sub-band CQT discriminator)

src/third_party/BigVGAN/activations.py

@@ -0,0 +1,126 @@
+# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license.
+#   LICENSE is in incl_licenses directory.
+
+import torch
+from torch import nn, sin, pow
+from torch.nn import Parameter
+
+
+class Snake(nn.Module):
+    """
+    Implementation of a sine-based periodic activation function
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter
+    References:
+        - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snake(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha: trainable parameter
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            alpha will be trained along with the rest of your model.
+        """
+        super(Snake, self).__init__()
+        self.in_features = in_features
+
+        # Initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # Log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+        else:  # Linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        Snake ∶= x + 1/a * sin^2 (xa)
+        """
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # Line up with x to [B, C, T]
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+        x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x
+
+
+class SnakeBeta(nn.Module):
+    """
+    A modified Snake function which uses separate parameters for the magnitude of the periodic components
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter that controls frequency
+        - beta - trainable parameter that controls magnitude
+    References:
+        - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snakebeta(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha - trainable parameter that controls frequency
+            - beta - trainable parameter that controls magnitude
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            beta is initialized to 1 by default, higher values = higher-magnitude.
+            alpha will be trained along with the rest of your model.
+        """
+        super(SnakeBeta, self).__init__()
+        self.in_features = in_features
+
+        # Initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # Log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+            self.beta = Parameter(torch.zeros(in_features) * alpha)
+        else:  # Linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+            self.beta = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        SnakeBeta ∶= x + 1/b * sin^2 (xa)
+        """
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # Line up with x to [B, C, T]
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x

src/third_party/BigVGAN/alias_free_activation/cuda/activation1d.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+#   Licensed under the MIT license.
+
+import torch
+import torch.nn as nn
+from alias_free_activation.torch.resample import UpSample1d, DownSample1d
+
+# load fused CUDA kernel: this enables importing anti_alias_activation_cuda
+from alias_free_activation.cuda import load
+
+anti_alias_activation_cuda = load.load()
+
+
+class FusedAntiAliasActivation(torch.autograd.Function):
+    """
+    Assumes filter size 12, replication padding on upsampling/downsampling, and logscale alpha/beta parameters as inputs.
+    The hyperparameters are hard-coded in the kernel to maximize speed.
+    NOTE: The fused kenrel is incorrect for Activation1d with different hyperparameters.
+    """
+
+    @staticmethod
+    def forward(ctx, inputs, up_ftr, down_ftr, alpha, beta):
+        activation_results = anti_alias_activation_cuda.forward(
+            inputs, up_ftr, down_ftr, alpha, beta
+        )
+
+        return activation_results
+
+    @staticmethod
+    def backward(ctx, output_grads):
+        raise NotImplementedError
+        return output_grads, None, None
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self,
+        activation,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+        fused: bool = True,
+    ):
+        super().__init__()
+        self.up_ratio = up_ratio
+        self.down_ratio = down_ratio
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+        self.fused = fused  # Whether to use fused CUDA kernel or not
+
+    def forward(self, x):
+        if not self.fused:
+            x = self.upsample(x)
+            x = self.act(x)
+            x = self.downsample(x)
+            return x
+        else:
+            if self.act.__class__.__name__ == "Snake":
+                beta = self.act.alpha.data  # Snake uses same params for alpha and beta
+            else:
+                beta = (
+                    self.act.beta.data
+                )  # Snakebeta uses different params for alpha and beta
+            alpha = self.act.alpha.data
+            if (
+                not self.act.alpha_logscale
+            ):  # Exp baked into cuda kernel, cancel it out with a log
+                alpha = torch.log(alpha)
+                beta = torch.log(beta)
+
+            x = FusedAntiAliasActivation.apply(
+                x, self.upsample.filter, self.downsample.lowpass.filter, alpha, beta
+            )
+            return x

src/third_party/BigVGAN/alias_free_activation/cuda/anti_alias_activation.cpp

@@ -0,0 +1,23 @@
+/* coding=utf-8
+ * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * 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.
+ */
+
+ #include <torch/extension.h>
+
+extern "C" torch::Tensor fwd_cuda(torch::Tensor const &input, torch::Tensor const &up_filter, torch::Tensor const &down_filter, torch::Tensor const &alpha, torch::Tensor const &beta);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &fwd_cuda, "Anti-Alias Activation forward (CUDA)");
+}

src/third_party/BigVGAN/alias_free_activation/cuda/anti_alias_activation_cuda.cu

@@ -0,0 +1,246 @@
+/* coding=utf-8
+ * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * 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.
+ */
+
+#include <ATen/ATen.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+#include "type_shim.h"
+#include <assert.h>
+#include <cfloat>
+#include <limits>
+#include <stdint.h>
+#include <c10/macros/Macros.h>
+
+namespace
+{
+    // Hard-coded hyperparameters
+    // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and
+    constexpr int ELEMENTS_PER_LDG_STG = 1; //(WARP_ITERATIONS < 4) ? 1 : 4;
+    constexpr int BUFFER_SIZE = 32;
+    constexpr int FILTER_SIZE = 12;
+    constexpr int HALF_FILTER_SIZE = 6;
+    constexpr int UPSAMPLE_REPLICATION_PAD = 5; // 5 on each side, matching torch impl
+    constexpr int DOWNSAMPLE_REPLICATION_PAD_LEFT = 5; // matching torch impl
+    constexpr int DOWNSAMPLE_REPLICATION_PAD_RIGHT = 6; // matching torch impl
+
+    template <typename input_t, typename output_t, typename acc_t>
+    __global__ void anti_alias_activation_forward(
+        output_t *dst,
+        const input_t *src,
+        const input_t *up_ftr,
+        const input_t *down_ftr,
+        const input_t *alpha,
+        const input_t *beta,
+        int batch_size,
+        int channels,
+        int seq_len)
+    {
+        // Up and downsample filters
+        input_t up_filter[FILTER_SIZE];
+        input_t down_filter[FILTER_SIZE];
+
+        // Load data from global memory including extra indices reserved for replication paddings
+        input_t elements[2 * FILTER_SIZE + 2 * BUFFER_SIZE + 2 * UPSAMPLE_REPLICATION_PAD] = {0};
+        input_t intermediates[2 * FILTER_SIZE + 2 * BUFFER_SIZE + DOWNSAMPLE_REPLICATION_PAD_LEFT + DOWNSAMPLE_REPLICATION_PAD_RIGHT] = {0};
+
+        // Output stores downsampled output before writing to dst
+        output_t output[BUFFER_SIZE];
+
+        // blockDim/threadIdx = (128, 1, 1)
+        // gridDim/blockIdx = (seq_blocks, channels, batches)
+        int block_offset = (blockIdx.x * 128 * BUFFER_SIZE + seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
+        int local_offset = threadIdx.x * BUFFER_SIZE;
+        int seq_offset = blockIdx.x * 128 * BUFFER_SIZE + local_offset;
+
+        // intermediate have double the seq_len
+        int intermediate_local_offset = threadIdx.x * BUFFER_SIZE * 2;
+        int intermediate_seq_offset = blockIdx.x * 128 * BUFFER_SIZE * 2 + intermediate_local_offset;
+
+        // Get values needed for replication padding before moving pointer
+        const input_t *right_most_pntr = src + (seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
+        input_t seq_left_most_value = right_most_pntr[0];
+        input_t seq_right_most_value = right_most_pntr[seq_len - 1];
+
+        // Move src and dst pointers
+        src += block_offset + local_offset;
+        dst += block_offset + local_offset;
+
+        // Alpha and beta values for snake activatons. Applies exp by default
+        alpha = alpha + blockIdx.y;
+        input_t alpha_val = expf(alpha[0]);
+        beta = beta + blockIdx.y;
+        input_t beta_val = expf(beta[0]);
+
+        #pragma unroll
+        for (int it = 0; it < FILTER_SIZE; it += 1)
+        {
+            up_filter[it] = up_ftr[it];
+            down_filter[it] = down_ftr[it];
+        }
+
+        // Apply replication padding for upsampling, matching torch impl
+        #pragma unroll
+        for (int it = -HALF_FILTER_SIZE; it < BUFFER_SIZE + HALF_FILTER_SIZE; it += 1)
+        {
+            int element_index = seq_offset + it; // index for element
+            if ((element_index < 0) && (element_index >= -UPSAMPLE_REPLICATION_PAD))
+            {
+                elements[2 * (HALF_FILTER_SIZE + it)] = 2 * seq_left_most_value;
+            }
+            if ((element_index >= seq_len) && (element_index < seq_len + UPSAMPLE_REPLICATION_PAD))
+            {
+                elements[2 * (HALF_FILTER_SIZE + it)] = 2 * seq_right_most_value;
+            }
+            if ((element_index >= 0) && (element_index < seq_len))
+            {
+                elements[2 * (HALF_FILTER_SIZE + it)] = 2 * src[it];
+            }
+        }
+
+        // Apply upsampling strided convolution and write to intermediates. It reserves DOWNSAMPLE_REPLICATION_PAD_LEFT for replication padding of the downsampilng conv later
+        #pragma unroll
+        for (int it = 0; it < (2 * BUFFER_SIZE + 2 * FILTER_SIZE); it += 1)
+        {
+            input_t acc = 0.0;
+            int element_index = intermediate_seq_offset + it; // index for intermediate
+            #pragma unroll
+            for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx += 1)
+            {
+                if ((element_index + f_idx) >= 0)
+                {
+                    acc += up_filter[f_idx] * elements[it + f_idx];
+                }
+            }
+            intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] = acc;
+        }
+
+        // Apply activation function. It reserves DOWNSAMPLE_REPLICATION_PAD_LEFT and DOWNSAMPLE_REPLICATION_PAD_RIGHT for replication padding of the downsampilng conv later
+        double no_div_by_zero = 0.000000001;
+        #pragma unroll
+        for (int it = 0; it < 2 * BUFFER_SIZE + 2 * FILTER_SIZE; it += 1)
+        {
+            intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] += (1.0 / (beta_val + no_div_by_zero)) * sinf(intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] * alpha_val) * sinf(intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] * alpha_val);
+        }
+
+        // Apply replication padding before downsampling conv from intermediates
+        #pragma unroll
+        for (int it = 0; it < DOWNSAMPLE_REPLICATION_PAD_LEFT; it += 1)
+        {
+            intermediates[it] = intermediates[DOWNSAMPLE_REPLICATION_PAD_LEFT];
+        }
+        #pragma unroll
+        for (int it = DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE; it < DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE + DOWNSAMPLE_REPLICATION_PAD_RIGHT; it += 1)
+        {
+            intermediates[it] = intermediates[DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE - 1];
+        }
+
+        // Apply downsample strided convolution (assuming stride=2) from intermediates
+        #pragma unroll
+        for (int it = 0; it < BUFFER_SIZE; it += 1)
+        {
+            input_t acc = 0.0;
+            #pragma unroll
+            for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx += 1)
+            {
+                // Add constant DOWNSAMPLE_REPLICATION_PAD_RIGHT to match torch implementation
+                acc += down_filter[f_idx] * intermediates[it * 2 + f_idx + DOWNSAMPLE_REPLICATION_PAD_RIGHT];
+            }
+            output[it] = acc;
+        }
+
+        // Write output to dst
+        #pragma unroll
+        for (int it = 0;  it < BUFFER_SIZE;  it += ELEMENTS_PER_LDG_STG)
+        {
+            int element_index = seq_offset + it;
+            if (element_index < seq_len)
+            {
+                dst[it] = output[it];
+            }
+        }
+
+    }
+
+    template <typename input_t, typename output_t, typename acc_t>
+    void dispatch_anti_alias_activation_forward(
+        output_t *dst,
+        const input_t *src,
+        const input_t *up_ftr,
+        const input_t *down_ftr,
+        const input_t *alpha,
+        const input_t *beta,
+        int batch_size,
+        int channels,
+        int seq_len)
+    {
+        if (seq_len == 0)
+        {
+            return;
+        }
+        else
+        {
+            // Use 128 threads per block to maximimize gpu utilization
+            constexpr int threads_per_block = 128;
+            constexpr int seq_len_per_block = 4096;
+            int blocks_per_seq_len = (seq_len + seq_len_per_block - 1) / seq_len_per_block;
+            dim3 blocks(blocks_per_seq_len, channels, batch_size);
+            dim3 threads(threads_per_block, 1, 1);
+
+            anti_alias_activation_forward<input_t, output_t, acc_t>
+                <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, up_ftr, down_ftr, alpha, beta, batch_size, channels, seq_len);
+        }
+    }
+}
+
+extern "C" torch::Tensor fwd_cuda(torch::Tensor const &input, torch::Tensor const &up_filter, torch::Tensor const &down_filter, torch::Tensor const &alpha, torch::Tensor const &beta)
+{
+    // Input is a 3d tensor with dimensions [batches, channels, seq_len]
+    const int batches = input.size(0);
+    const int channels = input.size(1);
+    const int seq_len = input.size(2);
+
+    // Output
+    auto act_options = input.options().requires_grad(false);
+
+    torch::Tensor anti_alias_activation_results =
+        torch::empty({batches, channels, seq_len}, act_options);
+
+    void *input_ptr = static_cast<void *>(input.data_ptr());
+    void *up_filter_ptr = static_cast<void *>(up_filter.data_ptr());
+    void *down_filter_ptr = static_cast<void *>(down_filter.data_ptr());
+    void *alpha_ptr = static_cast<void *>(alpha.data_ptr());
+    void *beta_ptr = static_cast<void *>(beta.data_ptr());
+    void *anti_alias_activation_results_ptr = static_cast<void *>(anti_alias_activation_results.data_ptr());
+
+    DISPATCH_FLOAT_HALF_AND_BFLOAT(
+        input.scalar_type(),
+        "dispatch anti alias activation_forward",
+        dispatch_anti_alias_activation_forward<scalar_t, scalar_t, float>(
+            reinterpret_cast<scalar_t *>(anti_alias_activation_results_ptr),
+            reinterpret_cast<const scalar_t *>(input_ptr),
+            reinterpret_cast<const scalar_t *>(up_filter_ptr),
+            reinterpret_cast<const scalar_t *>(down_filter_ptr),
+            reinterpret_cast<const scalar_t *>(alpha_ptr),
+            reinterpret_cast<const scalar_t *>(beta_ptr),
+            batches,
+            channels,
+            seq_len););
+    return anti_alias_activation_results;
+}

src/third_party/BigVGAN/alias_free_activation/cuda/compat.h

@@ -0,0 +1,29 @@
+/* coding=utf-8
+ * Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * 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.
+ */
+
+/*This code is copied fron NVIDIA apex:
+ *     https://github.com/NVIDIA/apex
+ *     with minor changes. */
+
+#ifndef TORCH_CHECK
+#define TORCH_CHECK AT_CHECK
+#endif
+
+#ifdef VERSION_GE_1_3
+#define DATA_PTR data_ptr
+#else
+#define DATA_PTR data
+#endif

src/third_party/BigVGAN/alias_free_activation/cuda/load.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+#   Licensed under the MIT license.
+
+import os
+import pathlib
+import subprocess
+
+from torch.utils import cpp_extension
+
+"""
+Setting this param to a list has a problem of generating different compilation commands (with diferent order of architectures) and leading to recompilation of fused kernels. 
+Set it to empty stringo avoid recompilation and assign arch flags explicity in extra_cuda_cflags below
+"""
+os.environ["TORCH_CUDA_ARCH_LIST"] = ""
+
+
+def load():
+    # Check if cuda 11 is installed for compute capability 8.0
+    cc_flag = []
+    _, bare_metal_major, _ = _get_cuda_bare_metal_version(cpp_extension.CUDA_HOME)
+    if int(bare_metal_major) >= 11:
+        cc_flag.append("-gencode")
+        cc_flag.append("arch=compute_80,code=sm_80")
+
+    # Build path
+    srcpath = pathlib.Path(__file__).parent.absolute()
+    buildpath = srcpath / "build"
+    _create_build_dir(buildpath)
+
+    # Helper function to build the kernels.
+    def _cpp_extention_load_helper(name, sources, extra_cuda_flags):
+        return cpp_extension.load(
+            name=name,
+            sources=sources,
+            build_directory=buildpath,
+            extra_cflags=[
+                "-O3",
+            ],
+            extra_cuda_cflags=[
+                "-O3",
+                "-gencode",
+                "arch=compute_70,code=sm_70",
+                "--use_fast_math",
+            ]
+            + extra_cuda_flags
+            + cc_flag,
+            verbose=True,
+        )
+
+    extra_cuda_flags = [
+        "-U__CUDA_NO_HALF_OPERATORS__",
+        "-U__CUDA_NO_HALF_CONVERSIONS__",
+        "--expt-relaxed-constexpr",
+        "--expt-extended-lambda",
+    ]
+
+    sources = [
+        srcpath / "anti_alias_activation.cpp",
+        srcpath / "anti_alias_activation_cuda.cu",
+    ]
+    anti_alias_activation_cuda = _cpp_extention_load_helper(
+        "anti_alias_activation_cuda", sources, extra_cuda_flags
+    )
+
+    return anti_alias_activation_cuda
+
+
+def _get_cuda_bare_metal_version(cuda_dir):
+    raw_output = subprocess.check_output(
+        [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
+    )
+    output = raw_output.split()
+    release_idx = output.index("release") + 1
+    release = output[release_idx].split(".")
+    bare_metal_major = release[0]
+    bare_metal_minor = release[1][0]
+
+    return raw_output, bare_metal_major, bare_metal_minor
+
+
+def _create_build_dir(buildpath):
+    try:
+        os.mkdir(buildpath)
+    except OSError:
+        if not os.path.isdir(buildpath):
+            print(f"Creation of the build directory {buildpath} failed")

src/third_party/BigVGAN/alias_free_activation/cuda/type_shim.h

@@ -0,0 +1,92 @@
+/* coding=utf-8
+ * Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * 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.
+ */
+
+#include <ATen/ATen.h>
+#include "compat.h"
+
+#define DISPATCH_FLOAT_HALF_AND_BFLOAT(TYPE, NAME, ...)                 \
+	switch (TYPE)                                                       \
+	{                                                                   \
+	case at::ScalarType::Float:                                         \
+	{                                                                   \
+		using scalar_t = float;                                         \
+		__VA_ARGS__;                                                    \
+		break;                                                          \
+	}                                                                   \
+	case at::ScalarType::Half:                                          \
+	{                                                                   \
+		using scalar_t = at::Half;                                      \
+		__VA_ARGS__;                                                    \
+		break;                                                          \
+	}                                                                   \
+	case at::ScalarType::BFloat16:                                      \
+	{                                                                   \
+		using scalar_t = at::BFloat16;                                  \
+		__VA_ARGS__;                                                    \
+		break;                                                          \
+	}                                                                   \
+	default:                                                            \
+		AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \
+	}
+
+#define DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(TYPEIN, TYPEOUT, NAME, ...) \
+	switch (TYPEIN)                                                            \
+	{                                                                          \
+	case at::ScalarType::Float:                                                \
+	{                                                                          \
+		using scalar_t_in = float;                                             \
+		switch (TYPEOUT)                                                       \
+		{                                                                      \
+		case at::ScalarType::Float:                                            \
+		{                                                                      \
+			using scalar_t_out = float;                                        \
+			__VA_ARGS__;                                                       \
+			break;                                                             \
+		}                                                                      \
+		case at::ScalarType::Half:                                             \
+		{                                                                      \
+			using scalar_t_out = at::Half;                                     \
+			__VA_ARGS__;                                                       \
+			break;                                                             \
+		}                                                                      \
+		case at::ScalarType::BFloat16:                                         \
+		{                                                                      \
+			using scalar_t_out = at::BFloat16;                                 \
+			__VA_ARGS__;                                                       \
+			break;                                                             \
+		}                                                                      \
+		default:                                                               \
+			AT_ERROR(#NAME, " not implemented for '", toString(TYPEOUT), "'"); \
+		}                                                                      \
+		break;                                                                 \
+	}                                                                          \
+	case at::ScalarType::Half:                                                 \
+	{                                                                          \
+		using scalar_t_in = at::Half;                                          \
+		using scalar_t_out = at::Half;                                         \
+		__VA_ARGS__;                                                           \
+		break;                                                                 \
+	}                                                                          \
+	case at::ScalarType::BFloat16:                                             \
+	{                                                                          \
+		using scalar_t_in = at::BFloat16;                                      \
+		using scalar_t_out = at::BFloat16;                                     \
+		__VA_ARGS__;                                                           \
+		break;                                                                 \
+	}                                                                          \
+	default:                                                                   \
+		AT_ERROR(#NAME, " not implemented for '", toString(TYPEIN), "'");      \
+	}

src/third_party/BigVGAN/alias_free_activation/torch/__init__.py

@@ -0,0 +1,6 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+from .filter import *
+from .resample import *
+from .act import *

src/third_party/BigVGAN/alias_free_activation/torch/act.py

@@ -0,0 +1,30 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch.nn as nn
+from alias_free_activation.torch.resample import UpSample1d, DownSample1d
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self,
+        activation,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+    ):
+        super().__init__()
+        self.up_ratio = up_ratio
+        self.down_ratio = down_ratio
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    # x: [B,C,T]
+    def forward(self, x):
+        x = self.upsample(x)
+        x = self.act(x)
+        x = self.downsample(x)
+
+        return x

src/third_party/BigVGAN/alias_free_activation/torch/filter.py

@@ -0,0 +1,101 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+if "sinc" in dir(torch):
+    sinc = torch.sinc
+else:
+    # This code is adopted from adefossez's julius.core.sinc under the MIT License
+    # https://adefossez.github.io/julius/julius/core.html
+    #   LICENSE is in incl_licenses directory.
+    def sinc(x: torch.Tensor):
+        """
+        Implementation of sinc, i.e. sin(pi * x) / (pi * x)
+        __Warning__: Different to julius.sinc, the input is multiplied by `pi`!
+        """
+        return torch.where(
+            x == 0,
+            torch.tensor(1.0, device=x.device, dtype=x.dtype),
+            torch.sin(math.pi * x) / math.pi / x,
+        )
+
+
+# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
+# https://adefossez.github.io/julius/julius/lowpass.html
+#   LICENSE is in incl_licenses directory.
+def kaiser_sinc_filter1d(
+    cutoff, half_width, kernel_size
+):  # return filter [1,1,kernel_size]
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+
+    # For kaiser window
+    delta_f = 4 * half_width
+    A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if A > 50.0:
+        beta = 0.1102 * (A - 8.7)
+    elif A >= 21.0:
+        beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+    # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
+    if even:
+        time = torch.arange(-half_size, half_size) + 0.5
+    else:
+        time = torch.arange(kernel_size) - half_size
+    if cutoff == 0:
+        filter_ = torch.zeros_like(time)
+    else:
+        filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+        """
+        Normalize filter to have sum = 1, otherwise we will have a small leakage of the constant component in the input signal.
+        """
+        filter_ /= filter_.sum()
+        filter = filter_.view(1, 1, kernel_size)
+
+    return filter
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(
+        self,
+        cutoff=0.5,
+        half_width=0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ):
+        """
+        kernel_size should be even number for stylegan3 setup, in this implementation, odd number is also possible.
+        """
+        super().__init__()
+        if cutoff < -0.0:
+            raise ValueError("Minimum cutoff must be larger than zero.")
+        if cutoff > 0.5:
+            raise ValueError("A cutoff above 0.5 does not make sense.")
+        self.kernel_size = kernel_size
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.padding = padding
+        self.padding_mode = padding_mode
+        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+        self.register_buffer("filter", filter)
+
+    # Input [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        if self.padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+
+        return out

src/third_party/BigVGAN/alias_free_activation/torch/resample.py

@@ -0,0 +1,58 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch.nn as nn
+from torch.nn import functional as F
+from alias_free_activation.torch.filter import LowPassFilter1d
+from alias_free_activation.torch.filter import kaiser_sinc_filter1d
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.stride = ratio
+        self.pad = self.kernel_size // ratio - 1
+        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+        self.pad_right = (
+            self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        )
+        filter = kaiser_sinc_filter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
+        )
+        self.register_buffer("filter", filter)
+
+    # x: [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        x = self.ratio * F.conv_transpose1d(
+            x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C
+        )
+        x = x[..., self.pad_left : -self.pad_right]
+
+        return x
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio,
+            half_width=0.6 / ratio,
+            stride=ratio,
+            kernel_size=self.kernel_size,
+        )
+
+    def forward(self, x):
+        xx = self.lowpass(x)
+
+        return xx

src/third_party/BigVGAN/bigvgan.py

@@ -0,0 +1,493 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+#   Licensed under the MIT license.
+
+# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
+#   LICENSE is in incl_licenses directory.
+
+import os
+import json
+from pathlib import Path
+from typing import Optional, Union, Dict
+
+import torch
+import torch.nn as nn
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+import activations
+from utils import init_weights, get_padding
+from alias_free_activation.torch.act import Activation1d as TorchActivation1d
+from env import AttrDict
+
+from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
+
+
+def load_hparams_from_json(path) -> AttrDict:
+    with open(path) as f:
+        data = f.read()
+    return AttrDict(json.loads(data))
+
+
+class AMPBlock1(torch.nn.Module):
+    """
+    AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer.
+    AMPBlock1 has additional self.convs2 that contains additional Conv1d layers with a fixed dilation=1 followed by each layer in self.convs1
+
+    Args:
+        h (AttrDict): Hyperparameters.
+        channels (int): Number of convolution channels.
+        kernel_size (int): Size of the convolution kernel. Default is 3.
+        dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5).
+        activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None.
+    """
+
+    def __init__(
+        self,
+        h: AttrDict,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple = (1, 3, 5),
+        activation: str = None,
+    ):
+        super().__init__()
+        
+        self.h = h
+
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                )
+                for _ in range(len(dilation))
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+        self.num_layers = len(self.convs1) + len(
+            self.convs2
+        )  # Total number of conv layers
+
+        # Select which Activation1d, lazy-load cuda version to ensure backward compatibility
+        if self.h.get("use_cuda_kernel", False):
+            from alias_free_activation.cuda.activation1d import (
+                Activation1d as CudaActivation1d,
+            )
+
+            Activation1d = CudaActivation1d
+        else:
+            Activation1d = TorchActivation1d
+
+        # Activation functions
+        if activation == "snake":
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=activations.Snake(
+                            channels, alpha_logscale=h.snake_logscale
+                        )
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        elif activation == "snakebeta":
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=activations.SnakeBeta(
+                            channels, alpha_logscale=h.snake_logscale
+                        )
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        else:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+    def forward(self, x):
+        acts1, acts2 = self.activations[::2], self.activations[1::2]
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+            xt = a1(x)
+            xt = c1(xt)
+            xt = a2(xt)
+            xt = c2(xt)
+            x = xt + x
+
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class AMPBlock2(torch.nn.Module):
+    """
+    AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer.
+    Unlike AMPBlock1, AMPBlock2 does not contain extra Conv1d layers with fixed dilation=1
+
+    Args:
+        h (AttrDict): Hyperparameters.
+        channels (int): Number of convolution channels.
+        kernel_size (int): Size of the convolution kernel. Default is 3.
+        dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5).
+        activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None.
+    """
+
+    def __init__(
+        self,
+        h: AttrDict,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple = (1, 3, 5),
+        activation: str = None,
+    ):
+        super().__init__()
+        
+        self.h = h
+
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
+        self.convs.apply(init_weights)
+
+        self.num_layers = len(self.convs)  # Total number of conv layers
+
+        # Select which Activation1d, lazy-load cuda version to ensure backward compatibility
+        if self.h.get("use_cuda_kernel", False):
+            from alias_free_activation.cuda.activation1d import (
+                Activation1d as CudaActivation1d,
+            )
+
+            Activation1d = CudaActivation1d
+        else:
+            Activation1d = TorchActivation1d
+
+        # Activation functions
+        if activation == "snake":
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=activations.Snake(
+                            channels, alpha_logscale=h.snake_logscale
+                        )
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        elif activation == "snakebeta":
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=activations.SnakeBeta(
+                            channels, alpha_logscale=h.snake_logscale
+                        )
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        else:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+    def forward(self, x):
+        for c, a in zip(self.convs, self.activations):
+            xt = a(x)
+            xt = c(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class BigVGAN(
+    torch.nn.Module,
+    PyTorchModelHubMixin,
+    library_name="bigvgan",
+    repo_url="https://github.com/NVIDIA/BigVGAN",
+    docs_url="https://github.com/NVIDIA/BigVGAN/blob/main/README.md",
+    pipeline_tag="audio-to-audio",
+    license="mit",
+    tags=["neural-vocoder", "audio-generation", "arxiv:2206.04658"],
+):
+    """
+    BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks).
+    New in BigVGAN-v2: it can optionally use optimized CUDA kernels for AMP (anti-aliased multi-periodicity) blocks.
+
+    Args:
+        h (AttrDict): Hyperparameters.
+        use_cuda_kernel (bool): If set to True, loads optimized CUDA kernels for AMP. This should be used for inference only, as training is not supported with CUDA kernels.
+
+    Note:
+        - The `use_cuda_kernel` parameter should be used for inference only, as training with CUDA kernels is not supported.
+        - Ensure that the activation function is correctly specified in the hyperparameters (h.activation).
+    """
+
+    def __init__(self, h: AttrDict, use_cuda_kernel: bool = False):
+        super().__init__()
+        self.h = h
+        self.h["use_cuda_kernel"] = use_cuda_kernel
+
+        # Select which Activation1d, lazy-load cuda version to ensure backward compatibility
+        if self.h.get("use_cuda_kernel", False):
+            from alias_free_activation.cuda.activation1d import (
+                Activation1d as CudaActivation1d,
+            )
+
+            Activation1d = CudaActivation1d
+        else:
+            Activation1d = TorchActivation1d
+
+        self.num_kernels = len(h.resblock_kernel_sizes)
+        self.num_upsamples = len(h.upsample_rates)
+
+        # Pre-conv
+        self.conv_pre = weight_norm(
+            Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)
+        )
+
+        # Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
+        if h.resblock == "1":
+            resblock_class = AMPBlock1
+        elif h.resblock == "2":
+            resblock_class = AMPBlock2
+        else:
+            raise ValueError(
+                f"Incorrect resblock class specified in hyperparameters. Got {h.resblock}"
+            )
+
+        # Transposed conv-based upsamplers. does not apply anti-aliasing
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        weight_norm(
+                            ConvTranspose1d(
+                                h.upsample_initial_channel // (2**i),
+                                h.upsample_initial_channel // (2 ** (i + 1)),
+                                k,
+                                u,
+                                padding=(k - u) // 2,
+                            )
+                        )
+                    ]
+                )
+            )
+
+        # Residual blocks using anti-aliased multi-periodicity composition modules (AMP)
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = h.upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)
+            ):
+                self.resblocks.append(
+                    resblock_class(h, ch, k, d, activation=h.activation)
+                )
+
+        # Post-conv
+        activation_post = (
+            activations.Snake(ch, alpha_logscale=h.snake_logscale)
+            if h.activation == "snake"
+            else (
+                activations.SnakeBeta(ch, alpha_logscale=h.snake_logscale)
+                if h.activation == "snakebeta"
+                else None
+            )
+        )
+        if activation_post is None:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+        self.activation_post = Activation1d(activation=activation_post)
+
+        # Whether to use bias for the final conv_post. Default to True for backward compatibility
+        self.use_bias_at_final = h.get("use_bias_at_final", True)
+        self.conv_post = weight_norm(
+            Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)
+        )
+
+        # Weight initialization
+        for i in range(len(self.ups)):
+            self.ups[i].apply(init_weights)
+        self.conv_post.apply(init_weights)
+
+        # Final tanh activation. Defaults to True for backward compatibility
+        self.use_tanh_at_final = h.get("use_tanh_at_final", True)
+
+    def forward(self, x):
+        # Pre-conv
+        x = self.conv_pre(x)
+
+        for i in range(self.num_upsamples):
+            # Upsampling
+            for i_up in range(len(self.ups[i])):
+                x = self.ups[i][i_up](x)
+            # AMP blocks
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+
+        # Post-conv
+        x = self.activation_post(x)
+        x = self.conv_post(x)
+        # Final tanh activation
+        if self.use_tanh_at_final:
+            x = torch.tanh(x)
+        else:
+            x = torch.clamp(x, min=-1.0, max=1.0)  # Bound the output to [-1, 1]
+
+        return x
+
+    def remove_weight_norm(self):
+        try:
+            print("Removing weight norm...")
+            for l in self.ups:
+                for l_i in l:
+                    remove_weight_norm(l_i)
+            for l in self.resblocks:
+                l.remove_weight_norm()
+            remove_weight_norm(self.conv_pre)
+            remove_weight_norm(self.conv_post)
+        except ValueError:
+            print("[INFO] Model already removed weight norm. Skipping!")
+            pass
+
+    # Additional methods for huggingface_hub support
+    def _save_pretrained(self, save_directory: Path) -> None:
+        """Save weights and config.json from a Pytorch model to a local directory."""
+
+        model_path = save_directory / "bigvgan_generator.pt"
+        torch.save({"generator": self.state_dict()}, model_path)
+
+        config_path = save_directory / "config.json"
+        with open(config_path, "w") as config_file:
+            json.dump(self.h, config_file, indent=4)
+
+    @classmethod
+    def _from_pretrained(
+        cls,
+        *,
+        model_id: str,
+        revision: str,
+        cache_dir: str,
+        force_download: bool,
+        proxies: Optional[Dict],
+        resume_download: bool,
+        local_files_only: bool,
+        token: Union[str, bool, None],
+        map_location: str = "cpu",  # Additional argument
+        strict: bool = False,  # Additional argument
+        use_cuda_kernel: bool = False,
+        **model_kwargs,
+    ):
+        """Load Pytorch pretrained weights and return the loaded model."""
+
+        # Download and load hyperparameters (h) used by BigVGAN
+        if os.path.isdir(model_id):
+            print("Loading config.json from local directory")
+            config_file = os.path.join(model_id, "config.json")
+        else:
+            config_file = hf_hub_download(
+                repo_id=model_id,
+                filename="config.json",
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                resume_download=resume_download,
+                token=token,
+                local_files_only=local_files_only,
+            )
+        h = load_hparams_from_json(config_file)
+
+        # instantiate BigVGAN using h
+        if use_cuda_kernel:
+            print(
+                f"[WARNING] You have specified use_cuda_kernel=True during BigVGAN.from_pretrained(). Only inference is supported (training is not implemented)!"
+            )
+            print(
+                f"[WARNING] You need nvcc and ninja installed in your system that matches your PyTorch build is using to build the kernel. If not, the model will fail to initialize or generate incorrect waveform!"
+            )
+            print(
+                f"[WARNING] For detail, see the official GitHub repository: https://github.com/NVIDIA/BigVGAN?tab=readme-ov-file#using-custom-cuda-kernel-for-synthesis"
+            )
+        model = cls(h, use_cuda_kernel=use_cuda_kernel)
+
+        # Download and load pretrained generator weight
+        if os.path.isdir(model_id):
+            print("Loading weights from local directory")
+            model_file = os.path.join(model_id, "bigvgan_generator.pt")
+        else:
+            print(f"Loading weights from {model_id}")
+            model_file = hf_hub_download(
+                repo_id=model_id,
+                filename="bigvgan_generator.pt",
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                resume_download=resume_download,
+                token=token,
+                local_files_only=local_files_only,
+            )
+
+        checkpoint_dict = torch.load(model_file, map_location=map_location)
+
+        try:
+            model.load_state_dict(checkpoint_dict["generator"])
+        except RuntimeError:
+            print(
+                f"[INFO] the pretrained checkpoint does not contain weight norm. Loading the checkpoint after removing weight norm!"
+            )
+            model.remove_weight_norm()
+            model.load_state_dict(checkpoint_dict["generator"])
+
+        return model

src/third_party/BigVGAN/configs/bigvgan_22khz_80band.json

@@ -0,0 +1,45 @@
+{
+    "resblock": "1",
+    "num_gpus": 0,
+    "batch_size": 32,
+    "learning_rate": 0.0001,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.9999996,
+    "seed": 1234,
+
+    "upsample_rates": [4,4,2,2,2,2],
+    "upsample_kernel_sizes": [8,8,4,4,4,4],
+    "upsample_initial_channel": 1536,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+
+    "activation": "snakebeta",
+    "snake_logscale": true,
+
+    "resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
+    "mpd_reshapes": [2, 3, 5, 7, 11],
+    "use_spectral_norm": false,
+    "discriminator_channel_mult": 1,
+
+    "segment_size": 8192,
+    "num_mels": 80,
+    "num_freq": 1025,
+    "n_fft": 1024,
+    "hop_size": 256,
+    "win_size": 1024,
+
+    "sampling_rate": 22050,
+
+    "fmin": 0,
+    "fmax": 8000,
+    "fmax_for_loss": null,
+
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}

src/third_party/BigVGAN/configs/bigvgan_24khz_100band.json

@@ -0,0 +1,45 @@
+{
+    "resblock": "1",
+    "num_gpus": 0,
+    "batch_size": 32,
+    "learning_rate": 0.0001,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.9999996,
+    "seed": 1234,
+
+    "upsample_rates": [4,4,2,2,2,2],
+    "upsample_kernel_sizes": [8,8,4,4,4,4],
+    "upsample_initial_channel": 1536,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+
+    "activation": "snakebeta",
+    "snake_logscale": true,
+
+    "resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
+    "mpd_reshapes": [2, 3, 5, 7, 11],
+    "use_spectral_norm": false,
+    "discriminator_channel_mult": 1,
+
+    "segment_size": 8192,
+    "num_mels": 100,
+    "num_freq": 1025,
+    "n_fft": 1024,
+    "hop_size": 256,
+    "win_size": 1024,
+
+    "sampling_rate": 24000,
+
+    "fmin": 0,
+    "fmax": 12000,
+    "fmax_for_loss": null,
+
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}

src/third_party/BigVGAN/configs/bigvgan_base_22khz_80band.json

@@ -0,0 +1,45 @@
+{
+    "resblock": "1",
+    "num_gpus": 0,
+    "batch_size": 32,
+    "learning_rate": 0.0001,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.9999996,
+    "seed": 1234,
+
+    "upsample_rates": [8,8,2,2],
+    "upsample_kernel_sizes": [16,16,4,4],
+    "upsample_initial_channel": 512,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+
+    "activation": "snakebeta",
+    "snake_logscale": true,
+
+    "resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
+    "mpd_reshapes": [2, 3, 5, 7, 11],
+    "use_spectral_norm": false,
+    "discriminator_channel_mult": 1,
+
+    "segment_size": 8192,
+    "num_mels": 80,
+    "num_freq": 1025,
+    "n_fft": 1024,
+    "hop_size": 256,
+    "win_size": 1024,
+
+    "sampling_rate": 22050,
+
+    "fmin": 0,
+    "fmax": 8000,
+    "fmax_for_loss": null,
+
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}

src/third_party/BigVGAN/configs/bigvgan_base_24khz_100band.json

@@ -0,0 +1,45 @@
+{
+    "resblock": "1",
+    "num_gpus": 0,
+    "batch_size": 32,
+    "learning_rate": 0.0001,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.9999996,
+    "seed": 1234,
+
+    "upsample_rates": [8,8,2,2],
+    "upsample_kernel_sizes": [16,16,4,4],
+    "upsample_initial_channel": 512,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+
+    "activation": "snakebeta",
+    "snake_logscale": true,
+
+    "resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
+    "mpd_reshapes": [2, 3, 5, 7, 11],
+    "use_spectral_norm": false,
+    "discriminator_channel_mult": 1,
+
+    "segment_size": 8192,
+    "num_mels": 100,
+    "num_freq": 1025,
+    "n_fft": 1024,
+    "hop_size": 256,
+    "win_size": 1024,
+
+    "sampling_rate": 24000,
+
+    "fmin": 0,
+    "fmax": 12000,
+    "fmax_for_loss": null,
+
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}

src/third_party/BigVGAN/configs/bigvgan_v2_22khz_80band_256x.json

@@ -0,0 +1,61 @@
+{
+    "resblock": "1",
+    "num_gpus": 0,
+    "batch_size": 4,
+    "learning_rate": 0.0001,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.9999996,
+    "seed": 1234,
+
+    "upsample_rates": [4,4,2,2,2,2],
+    "upsample_kernel_sizes": [8,8,4,4,4,4],
+    "upsample_initial_channel": 1536,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+
+    "use_tanh_at_final": false,
+    "use_bias_at_final": false,
+
+    "activation": "snakebeta",
+    "snake_logscale": true,
+
+    "use_cqtd_instead_of_mrd": true,
+    "cqtd_filters": 128,
+    "cqtd_max_filters": 1024,
+    "cqtd_filters_scale": 1,
+    "cqtd_dilations": [1, 2, 4],
+    "cqtd_hop_lengths": [512, 256, 256],
+    "cqtd_n_octaves": [9, 9, 9],
+    "cqtd_bins_per_octaves": [24, 36, 48],
+
+    "mpd_reshapes": [2, 3, 5, 7, 11],
+    "use_spectral_norm": false,
+    "discriminator_channel_mult": 1,
+    
+    "use_multiscale_melloss": true,
+    "lambda_melloss": 15,
+
+    "clip_grad_norm": 500,
+
+    "segment_size": 65536,
+    "num_mels": 80,
+    "num_freq": 1025,
+    "n_fft": 1024,
+    "hop_size": 256,
+    "win_size": 1024,
+
+    "sampling_rate": 22050,
+
+    "fmin": 0,
+    "fmax": null,
+    "fmax_for_loss": null,
+    
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}

src/third_party/BigVGAN/configs/bigvgan_v2_22khz_80band_fmax8k_256x.json

@@ -0,0 +1,61 @@
+{
+    "resblock": "1",
+    "num_gpus": 0,
+    "batch_size": 4,
+    "learning_rate": 0.0001,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.9999996,
+    "seed": 1234,
+
+    "upsample_rates": [4,4,2,2,2,2],
+    "upsample_kernel_sizes": [8,8,4,4,4,4],
+    "upsample_initial_channel": 1536,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+
+    "use_tanh_at_final": false,
+    "use_bias_at_final": false,
+
+    "activation": "snakebeta",
+    "snake_logscale": true,
+
+    "use_cqtd_instead_of_mrd": true,
+    "cqtd_filters": 128,
+    "cqtd_max_filters": 1024,
+    "cqtd_filters_scale": 1,
+    "cqtd_dilations": [1, 2, 4],
+    "cqtd_hop_lengths": [512, 256, 256],
+    "cqtd_n_octaves": [9, 9, 9],
+    "cqtd_bins_per_octaves": [24, 36, 48],
+    
+    "mpd_reshapes": [2, 3, 5, 7, 11],
+    "use_spectral_norm": false,
+    "discriminator_channel_mult": 1,
+
+    "use_multiscale_melloss": true,
+    "lambda_melloss": 15,
+
+    "clip_grad_norm": 500,
+
+    "segment_size": 65536,
+    "num_mels": 80,
+    "num_freq": 1025,
+    "n_fft": 1024,
+    "hop_size": 256,
+    "win_size": 1024,
+
+    "sampling_rate": 22050,
+
+    "fmin": 0,
+    "fmax": 8000,
+    "fmax_for_loss": null,
+
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}

src/third_party/BigVGAN/configs/bigvgan_v2_24khz_100band_256x.json

@@ -0,0 +1,61 @@
+{
+    "resblock": "1",
+    "num_gpus": 0,
+    "batch_size": 4,
+    "learning_rate": 0.0001,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.9999996,
+    "seed": 1234,
+
+    "upsample_rates": [4,4,2,2,2,2],
+    "upsample_kernel_sizes": [8,8,4,4,4,4],
+    "upsample_initial_channel": 1536,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+
+    "use_tanh_at_final": false,
+    "use_bias_at_final": false,
+
+    "activation": "snakebeta",
+    "snake_logscale": true,
+
+    "use_cqtd_instead_of_mrd": true,
+    "cqtd_filters": 128,
+    "cqtd_max_filters": 1024,
+    "cqtd_filters_scale": 1,
+    "cqtd_dilations": [1, 2, 4],
+    "cqtd_hop_lengths": [512, 256, 256],
+    "cqtd_n_octaves": [9, 9, 9],
+    "cqtd_bins_per_octaves": [24, 36, 48],
+
+    "mpd_reshapes": [2, 3, 5, 7, 11],
+    "use_spectral_norm": false,
+    "discriminator_channel_mult": 1,
+
+    "use_multiscale_melloss": true,
+    "lambda_melloss": 15,
+
+    "clip_grad_norm": 500,
+
+    "segment_size": 65536,
+    "num_mels": 100,
+    "num_freq": 1025,
+    "n_fft": 1024,
+    "hop_size": 256,
+    "win_size": 1024,
+
+    "sampling_rate": 24000,
+
+    "fmin": 0,
+    "fmax": null,
+    "fmax_for_loss": null,
+    
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}

src/third_party/BigVGAN/configs/bigvgan_v2_44khz_128band_256x.json

@@ -0,0 +1,61 @@
+{
+    "resblock": "1",
+    "num_gpus": 0,
+    "batch_size": 4,
+    "learning_rate": 0.0001,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.9999996,
+    "seed": 1234,
+
+    "upsample_rates": [4,4,2,2,2,2],
+    "upsample_kernel_sizes": [8,8,4,4,4,4],
+    "upsample_initial_channel": 1536,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+
+    "use_tanh_at_final": false,
+    "use_bias_at_final": false,
+
+    "activation": "snakebeta",
+    "snake_logscale": true,
+
+    "use_cqtd_instead_of_mrd": true,
+    "cqtd_filters": 128,
+    "cqtd_max_filters": 1024,
+    "cqtd_filters_scale": 1,
+    "cqtd_dilations": [1, 2, 4],
+    "cqtd_hop_lengths": [512, 256, 256],
+    "cqtd_n_octaves": [9, 9, 9],
+    "cqtd_bins_per_octaves": [24, 36, 48],
+
+    "mpd_reshapes": [2, 3, 5, 7, 11],
+    "use_spectral_norm": false,
+    "discriminator_channel_mult": 1,
+
+    "use_multiscale_melloss": true,
+    "lambda_melloss": 15,
+
+    "clip_grad_norm": 500,
+
+    "segment_size": 65536,
+    "num_mels": 128,
+    "num_freq": 1025,
+    "n_fft": 1024,
+    "hop_size": 256,
+    "win_size": 1024,
+
+    "sampling_rate": 44100,
+
+    "fmin": 0,
+    "fmax": null,
+    "fmax_for_loss": null,
+    
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}

src/third_party/BigVGAN/configs/bigvgan_v2_44khz_128band_512x.json

@@ -0,0 +1,61 @@
+{
+    "resblock": "1",
+    "num_gpus": 0,
+    "batch_size": 4,
+    "learning_rate": 0.0001,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.9999996,
+    "seed": 1234,
+
+    "upsample_rates": [8,4,2,2,2,2],
+    "upsample_kernel_sizes": [16,8,4,4,4,4],
+    "upsample_initial_channel": 1536,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+
+    "use_tanh_at_final": false,
+    "use_bias_at_final": false,
+
+    "activation": "snakebeta",
+    "snake_logscale": true,
+
+    "use_cqtd_instead_of_mrd": true,
+    "cqtd_filters": 128,
+    "cqtd_max_filters": 1024,
+    "cqtd_filters_scale": 1,
+    "cqtd_dilations": [1, 2, 4],
+    "cqtd_hop_lengths": [512, 256, 256],
+    "cqtd_n_octaves": [9, 9, 9],
+    "cqtd_bins_per_octaves": [24, 36, 48],
+
+    "mpd_reshapes": [2, 3, 5, 7, 11],
+    "use_spectral_norm": false,
+    "discriminator_channel_mult": 1,
+
+    "use_multiscale_melloss": true,
+    "lambda_melloss": 15,
+
+    "clip_grad_norm": 500,
+
+    "segment_size": 65536,
+    "num_mels": 128,
+    "num_freq": 2049,
+    "n_fft": 2048,
+    "hop_size": 512,
+    "win_size": 2048,
+
+    "sampling_rate": 44100,
+
+    "fmin": 0,
+    "fmax": null,
+    "fmax_for_loss": null,
+
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}