diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..b908d4cbe37ccea3388a2f308595ae63a255755d --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +*.pyc +.DS_Store diff --git a/README.md b/README.md index bd3322ac62524cd8d92f1dcfd0e1c8b95f390bce..c08009b667119bbd16673694f4590c4d1903fc03 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ emoji: 🐢 colorFrom: red colorTo: gray sdk: gradio -sdk_version: 6.20.0 +sdk_version: 5.28.0 python_version: '3.12' app_file: app.py pinned: false diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..a00ab2609a9b8c697cb7dc4edb917a345ffd5acc --- /dev/null +++ b/app.py @@ -0,0 +1,337 @@ +import sys +sys.stdout.reconfigure(line_buffering=True) + +import gc +import tempfile +import threading +import traceback +import types + +import torch +import yaml +import gradio as gr +from huggingface_hub import hf_hub_download +from diffusers import FlowMatchEulerDiscreteScheduler + +try: + import spaces + def gpu_decorator(func): return spaces.GPU(func, duration=120) +except ImportError: + def gpu_decorator(func): return func + +from pyharp import ModelCard, build_endpoint +from unison.models.mmaudio.features_utils import FeaturesUtils +from unison.pipelines.infer import ( + init_text_hidden_extractor, + sync_omni_dim_with_text_encoder, + _load_model, + sample_latents, + decode_and_save, + decode_and_save_full, + load_source_audio, + load_ref_audio, + make_edit_mask, + downsample_mask, + join_ref_target_text, + transcribe_ref_audio, + MAX_AUDIO_DURATION, +) + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +OMNI_MODEL_ID = "Qwen/Qwen2.5-Omni-7B" +UNISON_REPO = "jac22/UNISON" +MMAUDIO_REPO = "hkchengrex/MMAudio" +DEFAULT_VARIANT = "Balanced (44kHz)" + +VARIANTS = { + "Balanced (44kHz)": { + "model_config": "unison/config/D20S0_O_40ch.yaml", + "ckpt_file": "unison_D20S0_O_40ch/model.safetensors", + "vae_mode": "44k", + "vae_ckpt_file": "ext_weights/v1-44.pth", + "vocoder_ckpt_file": None, # 44k mode auto-pulls BigVGANv2 from HF Hub + "sample_rate": 44100, + }, + "High detail (16kHz)": { + "model_config": "unison/config/D24S0_O_20ch.yaml", + "ckpt_file": "unison_D24S0_O_20ch/model.safetensors", + "vae_mode": "16k", + "vae_ckpt_file": "ext_weights/v1-16.pth", + "vocoder_ckpt_file": "ext_weights/best_netG.pt", + "sample_rate": 16000, + }, +} + +# Only one variant's full stack (text encoder + DiT + VAE) is kept on the GPU at a +# time: the two variants need different Qwen2.5-Omni-7B layer-sampling depths, so +# the encoder can't be shared as-is (QwenOmniThinkerExtractor bakes dit_depth into +# which LLM layers it extracts at construction time). Switching variants evicts the +# previous one and rebuilds from the local HF cache (no re-download). +_active_name = None +_active_entry = None +_loading = True +_load_error = None + +# Only one request runs at a time. The scheduler object is shared and stateful, +# so two requests running together would corrupt each other's results. This lock +# also covers _active_name/_active_entry below, so no separate lock is needed. +_gpu_lock = threading.Lock() + + +def _build_variant(name): + """Download and construct one variant's full stack: text encoder, DiT backbone, + and VAE. Mirrors the model-loading steps in unison/pipelines/infer.py's main().""" + spec = VARIANTS[name] + print(f"Building variant: {name} ...") + + with open(spec["model_config"]) as f: + model_config = yaml.safe_load(f) + dit_depth = model_config.get("mm_double_blocks_depth", 0) + model_config.get("mm_single_blocks_depth", 0) + omni_last_layer_idx = model_config.get("omni_last_layer_idx", -1) + + extractor = init_text_hidden_extractor( + "omni", OMNI_MODEL_ID, None, dit_depth, DEVICE, omni_last_layer_idx=omni_last_layer_idx, + ) + sync_omni_dim_with_text_encoder(model_config, extractor) + + ckpt_path = hf_hub_download(repo_id=UNISON_REPO, filename=spec["ckpt_file"]) + model = _load_model(types.SimpleNamespace(model_ckpt=ckpt_path), DEVICE, model_config) + + vae_ckpt_path = hf_hub_download(repo_id=MMAUDIO_REPO, filename=spec["vae_ckpt_file"]) + vocoder_ckpt_path = ( + hf_hub_download(repo_id=MMAUDIO_REPO, filename=spec["vocoder_ckpt_file"]) + if spec["vocoder_ckpt_file"] else None + ) + audio_vae = FeaturesUtils( + tod_vae_ckpt=vae_ckpt_path, + bigvgan_vocoder_ckpt=vocoder_ckpt_path, + mode=spec["vae_mode"], + ) + audio_vae.to(DEVICE).eval() + + # probe latent length for a MAX_AUDIO_DURATION-long clip (needed for generation mode). + sample_rate = spec["sample_rate"] + dummy_len = int(MAX_AUDIO_DURATION * sample_rate) + with torch.no_grad(): + dummy_lat = audio_vae.wrapped_encode(torch.zeros(1, dummy_len, device=DEVICE)) + gen_target_frames = int(dummy_lat.shape[-1]) + + scheduler = FlowMatchEulerDiscreteScheduler() + print(f"Variant ready: {name}") + return (model, extractor, audio_vae, 0.5, sample_rate, gen_target_frames, scheduler) + + +def load_default_variant(): + """Background-thread target: build DEFAULT_VARIANT at startup and publish it + as the active variant, so the Space doesn't block its HTTP server on the load.""" + global _active_name, _active_entry, _loading, _load_error + try: + entry = _build_variant(DEFAULT_VARIANT) + # No lock needed: _loading stays True until right after this write, and + # process_fn won't touch _active_name/_active_entry while _loading is True. + _active_name, _active_entry = DEFAULT_VARIANT, entry + except Exception: + _load_error = traceback.format_exc() + print(f"Variant load error: {_load_error}") + finally: + _loading = False + + +threading.Thread(target=load_default_variant, daemon=True).start() + + +def get_variant(name: str): + """Return the active variant's (model, extractor, vae, ...) tuple, building and + switching to `name` first if a different variant is currently active. + + Caller must hold _gpu_lock — that's what makes the unlocked reads/writes of + _active_name/_active_entry here safe.""" + global _active_name, _active_entry + if name == _active_name: + return _active_entry + + print(f"Switching variant: {_active_name!r} -> {name!r}") + entry = _build_variant(name) + _active_name, _active_entry = name, entry + gc.collect() + torch.cuda.empty_cache() + return entry + + +model_card = ModelCard( + name="UNISON", + description=( + "Unified sound generation and editing: text-to-audio, text-to-speech, " + "audio-scene editing, and zero-shot voice cloning from a single model." + ), + author="Zhaoqing Li, Haoning Xu, Jingran Su, Yaofang Liu, Zhefan Rao, Huimeng Wang, " + "Jiajun Deng, Tianzi Wang, Zengrui Jin, Rui Liu, Haoxuan Che, Xunying Liu", + tags=["text-to-audio", "text-to-speech", "audio-editing", "zero-shot-tts"], +) + + +@gpu_decorator +@torch.inference_mode() +def process_fn( + input_audio_path: str, + mode: str, + sound_type: str, + voice: str, + prompt: str, + background: str, + ref_text: str, + model_variant: str, + steps: int, + guidance: float, + duration: float, +) -> str: + if _loading: + raise gr.Error("Model is still loading, please wait a moment and try again.") + if _active_entry is None: + raise gr.Error(f"Model failed to load: {_load_error}") + + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: + out_path = f.name + + with _gpu_lock: + model, omni_extractor, audio_vae, vae_scale_factor, vae_sample_rate, gen_target_frames, scheduler = get_variant(model_variant) + + if mode == "Generate": + # Text-to-audio/speech, no input audio involved. sound_type picks which of + # the model's trained prompt templates to build (see README task table) — + # voice/background only apply to the Speech templates, unused otherwise. + if sound_type == "Sound Effect": + tagged_prompt = f"[Audio] {prompt}" + elif sound_type == "Speech": + tagged_prompt = f'[Speech] A {voice.lower()} voice saying "{prompt}"' + else: # "Speech + Background" + tagged_prompt = f'[Speech] A {voice.lower()} voice saying "{prompt}" [Audio] {background}' + + latents = sample_latents( + model, scheduler, omni_extractor, [tagged_prompt], + num_inference_steps=steps, guidance_scale=guidance, device=DEVICE, + target_frames=gen_target_frames, + ) + latents = latents * (1.0 / vae_scale_factor) + decode_and_save(audio_vae, latents, [duration], [out_path], sample_rate=vae_sample_rate) + + elif mode == "Edit": + # Source audio + instruction. Mask is all-ones: the whole clip is editable, + # source_latents just gives the model something to condition on. + # sound_type again picks [Audio] vs [Speech] as the edit's sub-tag; "Speech + # + Background" isn't a real edit template, so it also falls back to [Audio]. + if not input_audio_path: + raise gr.Error("Edit mode requires an input audio track.") + edit_target = "Speech" if sound_type == "Speech" else "Audio" + tgt_wav_len = int(MAX_AUDIO_DURATION * vae_sample_rate) + src_wav = load_source_audio(input_audio_path, target_sr=vae_sample_rate, + target_length=tgt_wav_len, device=DEVICE) + src_latent = audio_vae.wrapped_encode(src_wav) * vae_scale_factor + mask_wav = make_edit_mask(src_wav.shape[-1], DEVICE) + mask_lat = downsample_mask(mask_wav, src_latent.shape[-1]) + + latents = sample_latents( + model, scheduler, omni_extractor, [f"[Edit] [{edit_target}] {prompt}"], + source_latents=src_latent, masks=mask_lat, + num_inference_steps=steps, guidance_scale=guidance, device=DEVICE, + ) + latents = latents * (1.0 / vae_scale_factor) + edit_duration = src_wav.shape[-1] / vae_sample_rate + decode_and_save(audio_vae, latents, [edit_duration], [out_path], sample_rate=vae_sample_rate) + + elif mode == "Clone Voice": + # Reference clip + text to speak in that voice. + if not input_audio_path: + raise gr.Error("Clone Voice mode requires a reference audio track.") + + # 3.0s matches REF_DURATION, the reference-clip length the model was trained with. + ref_wav = load_ref_audio(input_audio_path, target_sr=vae_sample_rate, + max_ref_duration=3.0, device=DEVICE) + ref_audio_duration = ref_wav.shape[-1] / vae_sample_rate + if ref_audio_duration + 1.0 > duration: # 1.0s floor so some cloned speech always fits + raise gr.Error(f"Reference clip ({ref_audio_duration:.1f}s) leaves under 1s for " + f"the cloned speech at Duration={duration:.1f}s. Increase Duration.") + + # Get the ref transcript (typed one wins over auto-transcription), then + # combine it with what the user wants said next into one prompt. + # Exclude load_ref_audio's trailing silence pad (tail_pad_s=0.1) so Whisper + # doesn't hallucinate tokens over silence. + speech_samples = max(ref_wav.shape[-1] - int(0.1 * vae_sample_rate), 1) + resolved_ref_text = ref_text.strip() or transcribe_ref_audio( + ref_wav[..., :speech_samples], sr=vae_sample_rate, + ) + combined_text = join_ref_target_text(resolved_ref_text, prompt) + full_prompt = f"[Speech with voice] {combined_text}" + + # Build one waveform [ref audio | silence] and encode it as a single clip — + # the model generates the target portion conditioned on the ref portion. + total_wav_len = int(duration * vae_sample_rate) + ref_wav_1d = ref_wav.squeeze(0) if ref_wav.dim() == 2 else ref_wav + ref_wav_len = min(ref_wav_1d.shape[-1], total_wav_len) + source_wav = torch.zeros(1, total_wav_len, device=DEVICE) + source_wav[:, :ref_wav_len] = ref_wav_1d[:ref_wav_len] + source_latent = audio_vae.wrapped_encode(source_wav) * vae_scale_factor + + mask_wav = torch.zeros(1, 1, total_wav_len, device=DEVICE) + mask_wav[:, :, :ref_wav_len] = 2.0 # 2 = reference region, 0 = target (see sample_latents docstring) + mask_latent = torch.nn.functional.interpolate( + mask_wav, size=source_latent.shape[-1] + ).to(torch.bfloat16) + + latents = sample_latents( + model, scheduler, omni_extractor, [full_prompt], + source_latents=source_latent, masks=mask_latent, + num_inference_steps=steps, guidance_scale=guidance, device=DEVICE, + ) + latents_full = latents * (1.0 / vae_scale_factor) + + # Decode the full ref+target latent, then crop out just the target — + # the decoder needs the ref portion as context to decode cleanly. + ref_samples = int(ref_audio_duration * vae_sample_rate) + decode_and_save_full( + audio_vae, latents_full, ref_samples, + [duration - ref_audio_duration], [out_path], sample_rate=vae_sample_rate, + ) + + else: + raise gr.Error(f"Unknown mode: {mode}") + + return out_path + + +with gr.Blocks() as demo: + input_components = [ + gr.Audio(type="filepath", label="Input Audio").harp_required(False) + .set_info("Source track for Edit mode, or reference voice for Clone Voice mode. Unused in Generate mode."), + gr.Dropdown(choices=["Generate", "Edit", "Clone Voice"], value="Generate", label="Mode"), + gr.Dropdown(choices=["Sound Effect", "Speech", "Speech + Background"], value="Sound Effect", + label="Sound Type", + info="Generate/Edit only: what kind of content Prompt describes."), + # Voice/Background are only meaningful for the Speech sound types. HARP has no + # way to hide a control based on another control's value, so they're always + # shown and just ignored (e.g. for Sound Effect) rather than hidden. + gr.Dropdown(choices=["Female", "Male"], value="Female", label="Voice", + info="Speech sound types only."), + gr.Textbox(label="Prompt", + info="Generate: describe the sound, or what's said. Edit: describe the change. Clone Voice: text to speak."), + gr.Textbox(label="Background (optional)", + info="\"Speech + Background\" sound type only: the background sound to mix in."), + gr.Textbox(label="Reference Transcript (optional)", + info="Only used in Clone Voice mode. Leave blank to auto-transcribe the input audio."), + gr.Dropdown(choices=list(VARIANTS), value="Balanced (44kHz)", label="Model"), + gr.Slider(minimum=10, maximum=100, step=5, value=50, label="Generation Steps"), + gr.Slider(minimum=1.0, maximum=10.0, step=0.5, value=4.5, label="Prompt Strength"), + gr.Slider(minimum=1.0, maximum=float(MAX_AUDIO_DURATION), step=0.5, value=10.0, label="Duration (s)"), + ] + output_components = [ + gr.Audio(type="filepath", label="Output Audio").set_info("Generated or edited audio."), + ] + + build_endpoint( + model_card=model_card, + input_components=input_components, + output_components=output_components, + process_fn=process_fn, + ) + +demo.queue().launch(pwa=True) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..c83d435829aab212c1f87067e1b8b043d180c893 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,19 @@ +git+https://github.com/TEAMuP-dev/pyharp.git@v0.3.0 +# model-specific deps below: +torch>=2.6.0 +torchaudio>=2.6.0 +numpy>=1.26.0 +einops>=0.8.0 +transformers>=4.53.0 +accelerate>=0.34.0 +diffusers>=0.30.0 +safetensors>=0.4.5 +huggingface-hub>=0.34.0 +soundfile>=0.12.0 +librosa>=0.9.0 +scipy>=1.11.0 +openai-whisper>=20240930 +pyyaml>=6.0 +omegaconf>=2.3.0 +loguru>=0.7.0 +tqdm>=4.67.0 diff --git a/unison/__init__.py b/unison/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7be4f2d5d0cf63c0c3c64669e58ca5d2f7b8d205 --- /dev/null +++ b/unison/__init__.py @@ -0,0 +1,44 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +import socket +from .commons import get_gpu_memory + + +if 'TOKENIZERS_PARALLELISM' not in os.environ: + os.environ['TOKENIZERS_PARALLELISM'] = 'false' + +def find_free_port(): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(('localhost', 0)) + port = sock.getsockname()[1] + sock.close() + return port + +def __initialize_default_distributed_environment(): + if 'RANK' not in os.environ: + os.environ['RANK'] = '0' + if 'WORLD_SIZE' not in os.environ: + os.environ['WORLD_SIZE'] = '1' + if 'LOCAL_RANK' not in os.environ: + os.environ['LOCAL_RANK'] = '0' + if 'MASTER_ADDR' not in os.environ: + os.environ['MASTER_ADDR'] = 'localhost' + if 'MASTER_PORT' not in os.environ: + os.environ['MASTER_PORT'] = str(find_free_port()) + +__initialize_default_distributed_environment() \ No newline at end of file diff --git a/unison/commons/__init__.py b/unison/commons/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8168cf433434caa4203ba32fbb797ed2b50a1db6 --- /dev/null +++ b/unison/commons/__init__.py @@ -0,0 +1,242 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +import torch +from itertools import repeat +from contextlib import contextmanager +from torch import nn +import collections.abc + +def _ntuple(n): + """Create a function that converts input to n-tuple.""" + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + x = tuple(x) + if len(x) == 1: + x = tuple(repeat(x[0], n)) + return x + return tuple(repeat(x, n)) + return parse + +# Convenience functions for common tuple sizes +to_1tuple = _ntuple(1) +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) +to_4tuple = _ntuple(4) + +PRECISION_TO_TYPE = { + 'fp32': torch.float32, + 'fp16': torch.float16, + 'bf16': torch.bfloat16, +} + +# Default generation pipeline configurations +PIPELINE_CONFIGS = { + '480p_t2v': { + 'guidance_scale': 6.0, + 'embedded_guidance_scale': None, + 'flow_shift': 5.0, + }, + '480p_i2v': { + 'guidance_scale': 6.0, + 'embedded_guidance_scale': None, + 'flow_shift': 5.0, + }, + '720p_t2v': { + 'guidance_scale': 6.0, + 'embedded_guidance_scale': None, + 'flow_shift': 9.0, + }, + '720p_i2v': { + 'guidance_scale': 6.0, + 'embedded_guidance_scale': None, + 'flow_shift': 7.0, + }, + '480p_t2v_distilled': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 5.0, + }, + '480p_i2v_distilled': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 5.0, + }, + '720p_t2v_distilled': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 9.0, + }, + '720p_i2v_distilled': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 7.0, + }, + '720p_t2v_distilled_sparse': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 7.0, + }, + '720p_i2v_distilled_sparse': { + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'flow_shift': 9.0, + }, +} + +# Default super-resolution pipeline configurations +SR_PIPELINE_CONFIGS = { + '720p_sr_distilled': { + 'flow_shift': 2.0, + 'base_resolution': '480p', + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'num_inference_steps': 6, + }, + '1080p_sr_distilled': { + 'flow_shift': 2.0, + 'base_resolution': '720p', + 'guidance_scale': 1.0, + 'embedded_guidance_scale': None, + 'num_inference_steps': 8, + }, +} + +TRANSFORMER_VERSION_TO_SR_VERSION = { + '480p_t2v': '720p_sr_distilled', + '720p_t2v': '1080p_sr_distilled', + '480p_i2v': '720p_sr_distilled', + '720p_i2v': '1080p_sr_distilled', + '480p_t2v_distilled': '720p_sr_distilled', + '720p_t2v_distilled': '1080p_sr_distilled', + '480p_i2v_distilled': '720p_sr_distilled', + '720p_i2v_distilled': '1080p_sr_distilled', + '480p_t2v_distilled_sparse': '720p_sr_distilled', + '720p_t2v_distilled_sparse': '1080p_sr_distilled', + '480p_i2v_distilled_sparse': '720p_sr_distilled', + '720p_i2v_distilled_sparse': '1080p_sr_distilled', +} + +def is_flash2_available(): + try: + from flash_attn import flash_attn_varlen_qkvpacked_func + return True + except Exception: + return False + +def is_flash3_available(): + try: + from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_v3 # noqa: F401 + return True + except Exception: + return False + +def is_flash_available(): + return is_flash2_available() or is_flash3_available() + +def is_sparse_attn_supported(): + return 'nvidia h' in torch.cuda.get_device_properties(0).name.lower() + +def is_sparse_attn_available(): + if not is_sparse_attn_supported(): + return False + try: + from flex_block_attn import flex_block_attn_func # noqa: F401 + return True + except Exception: + return False + +def maybe_fallback_attn_mode(attn_mode, infer_state=None, block_idx=None): + """ + Determine the final attention mode based on configuration and availability. + + Args: + attn_mode: Requested attention mode + infer_state: Inference configuration object (optional) + block_idx: Current block index (optional) + + Returns: + Final attention mode to use + """ + import warnings + + # Check for sageattn and flex-block-attn conflict + enable_sageattn = False + if infer_state is not None: + enable_sageattn = (infer_state.enable_sageattn and + block_idx in infer_state.sage_blocks_range) + + assert not (enable_sageattn and attn_mode == 'flex-block-attn'), \ + ("SageAttention cannot be used with flex-block-attn mode. " + "Please disable enable_sageattn or use a different attention mode.") + + # Use SageAttention if configured + if enable_sageattn: + attn_mode = 'sageattn' + return attn_mode + + # Handle flash attention modes + if attn_mode == 'flash': + if is_flash3_available(): + attn_mode = 'flash3' + elif is_flash2_available(): + attn_mode = 'flash2' + else: + warnings.warn("flash is not available. Falling back to torch attention.") + attn_mode = 'torch' + elif attn_mode == 'flash3': + if not is_flash3_available(): + warnings.warn("flash3 is not available. Falling back to torch attention.") + attn_mode = 'torch' + elif attn_mode == 'flash2': + if not is_flash2_available(): + warnings.warn("flash2 is not available. Falling back to torch attention.") + attn_mode = 'torch' + if attn_mode in ('flex-block-attn'): + from unison.commons import is_sparse_attn_available + if not is_sparse_attn_available(): + raise ValueError(f"{attn_mode} is not available for your GPU or flex-block-attn is not properly installed.") + return attn_mode + +@contextmanager +def auto_offload_model(models, device, enabled=True): + from diffusers.hooks.group_offloading import _is_group_offload_enabled + if enabled: + if isinstance(models, nn.Module): + models = [models] + for model in models: + if model is not None: + model.to(device) + yield + if enabled: + for model in models: + if model is not None: + model.to(torch.device('cpu')) + +def get_gpu_memory(device=None): + if not torch.cuda.is_available(): + return 0 + device = device if device is not None else torch.cuda.current_device() + props = torch.cuda.get_device_properties(device) + if hasattr(torch.cuda, 'get_per_process_memory_fraction'): + memory_fraction = torch.cuda.get_per_process_memory_fraction() + else: + memory_fraction = 1.0 + return props.total_memory * memory_fraction + +def get_rank(): + return int(os.environ.get('RANK', '0')) diff --git a/unison/commons/infer_state.py b/unison/commons/infer_state.py new file mode 100644 index 0000000000000000000000000000000000000000..a991994be05f46493ea8e9f2dd15c9edbf92c739 --- /dev/null +++ b/unison/commons/infer_state.py @@ -0,0 +1,48 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Optional +from dataclasses import dataclass + +@dataclass +class InferState: + enable_sageattn: bool = False # whether to use SageAttention + sage_blocks_range: Optional[range] = None # block range to use SageAttention + enable_torch_compile: bool = False # whether to use torch compile + +__infer_state = None + +def parse_range(value): + if '-' in value: + start, end = map(int, value.split('-')) + return list(range(start, end + 1)) + else: + return [int(x) for x in value.split(',')] + +def initialize_infer_state(args): + global __infer_state + sage_blocks_range = parse_range(args.sage_blocks_range) + # Map CLI argument use_sageattn to internal enable_sageattn field + use_sageattn = getattr(args, 'use_sageattn', False) + __infer_state = InferState( + enable_sageattn = use_sageattn, + sage_blocks_range = sage_blocks_range, + enable_torch_compile = args.enable_torch_compile, + ) + return __infer_state + +def get_infer_state(): + return __infer_state diff --git a/unison/commons/parallel_states.py b/unison/commons/parallel_states.py new file mode 100644 index 0000000000000000000000000000000000000000..830ce0198c2d40407ce8f2dd3e341fc75da22f86 --- /dev/null +++ b/unison/commons/parallel_states.py @@ -0,0 +1,85 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import os +from dataclasses import dataclass + +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh + +@dataclass +class ParallelDims: + sp: int = 1 + world_size: int = -1 + + def __post_init__(self): + if self.world_size == -1: + if dist.is_initialized(): + self.world_size = dist.get_world_size() + else: + self.world_size = int(os.getenv("WORLD_SIZE", "1")) + if dist.is_initialized(): + self.build_mesh("cuda") + else: + self.world_mesh = None + + def build_mesh(self, device_type): + assert self.world_size % self.sp == 0, "world_size must be divisible by sp" + mesh = init_device_mesh( + device_type, + [self.world_size // self.sp, self.sp], + mesh_dim_names=["dp", "sp"] + ) + self.world_mesh = mesh + return mesh + + @property + def sp_enabled(self): + return self.sp > 1 + + @property + def sp_group(self): + return self.world_mesh['sp'].get_group() + + @property + def sp_mesh(self): + return self.world_mesh['sp'] + + @property + def sp_rank(self): + if self.sp_enabled: + return self.world_mesh['sp'].get_local_rank() + else: + return dist.get_rank() + + @property + def dp_enabled(self): + return self.sp > 1 + +__parallel_dims = None + +def initialize_parallel_state( + sp: int = 1, +): + global __parallel_dims + __parallel_dims = ParallelDims(sp=sp) + return __parallel_dims + +def get_parallel_state(): + if __parallel_dims is None: + # create default parallel states (without enabling any parallelism) + initialize_parallel_state() + return __parallel_dims diff --git a/unison/config/D20S0_O_40ch.yaml b/unison/config/D20S0_O_40ch.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bfc5e8f06c9decc389c835c8b72003c68e9a92af --- /dev/null +++ b/unison/config/D20S0_O_40ch.yaml @@ -0,0 +1,27 @@ +in_channels: 40 # MMAudio 44 kHz VAE latent channels (C=40); model receives 2C+1=81 channels when concat_condition=True +out_channels: 40 +patch_size: [1, 2, 2] +hidden_size: 1024 +heads_num: 8 +mlp_act_type: gelu_tanh +mlp_width_ratio: 4.0 +mm_double_blocks_depth: 20 +mm_single_blocks_depth: 0 +qkv_bias: true +qk_norm: true +qk_norm_type: rms +attn_mode: flash +attn_param: null +rope_dim_list: [96, 16, 16] +rope_theta: 10000 +omni_dim: 3584 +is_reshape_temporal_channels: false +use_duration_embedding: false +is_audio_type: true +use_omni_embedding: true +use_omni_last_embedding: false +omni_last_layer_idx: -2 # -1 = last LLM layer, -2 = second-to-last +guidance_embed: false +temporal_rope_scaling_factor: 1.0 +repa_z_dim: null +repa_layer_num: null diff --git a/unison/config/D24S0_O_20ch.yaml b/unison/config/D24S0_O_20ch.yaml new file mode 100644 index 0000000000000000000000000000000000000000..46ee1e404c3ca46baaabe3e612ec48cc2d2471f8 --- /dev/null +++ b/unison/config/D24S0_O_20ch.yaml @@ -0,0 +1,27 @@ +in_channels: 20 # MMAudio 16 kHz VAE latent channels (C=20); model receives 2C+1=41 channels when concat_condition=True +out_channels: 20 +patch_size: [1, 2, 2] +hidden_size: 1024 +heads_num: 8 +mlp_act_type: gelu_tanh +mlp_width_ratio: 4.0 +mm_double_blocks_depth: 24 +mm_single_blocks_depth: 0 +qkv_bias: true +qk_norm: true +qk_norm_type: rms +attn_mode: flash +attn_param: null +rope_dim_list: [96, 16, 16] +rope_theta: 10000 +omni_dim: 3584 +is_reshape_temporal_channels: false +use_duration_embedding: false +is_audio_type: true +use_omni_embedding: true +use_omni_last_embedding: false +omni_last_layer_idx: -2 # -1 = last LLM layer, -2 = second-to-last +guidance_embed: false +temporal_rope_scaling_factor: 1.0 +repa_z_dim: null +repa_layer_num: null diff --git a/unison/models/__init__.py b/unison/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2fe721640f8f0188339c70082ecac3ff9b74dcfb --- /dev/null +++ b/unison/models/__init__.py @@ -0,0 +1,16 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + diff --git a/unison/models/mmaudio/__init__.py b/unison/models/mmaudio/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e5601dd86c2b9f45c050ca4e33e3661b244bfefb --- /dev/null +++ b/unison/models/mmaudio/__init__.py @@ -0,0 +1 @@ +# MMAudio package \ No newline at end of file diff --git a/unison/models/mmaudio/ext/__init__.py b/unison/models/mmaudio/ext/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/unison/models/mmaudio/ext/__init__.py @@ -0,0 +1 @@ + diff --git a/unison/models/mmaudio/ext/autoencoder/__init__.py b/unison/models/mmaudio/ext/autoencoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b69d4b793eb8dd9efe20367dc8806860a59b2c9f --- /dev/null +++ b/unison/models/mmaudio/ext/autoencoder/__init__.py @@ -0,0 +1 @@ +from .autoencoder import AutoEncoderModule \ No newline at end of file diff --git a/unison/models/mmaudio/ext/autoencoder/autoencoder.py b/unison/models/mmaudio/ext/autoencoder/autoencoder.py new file mode 100644 index 0000000000000000000000000000000000000000..4ae5e03a3ea7b1d28d10c5048cab7a022487a646 --- /dev/null +++ b/unison/models/mmaudio/ext/autoencoder/autoencoder.py @@ -0,0 +1,60 @@ +from typing import Literal, Optional + +import torch +import torch.nn as nn +import numpy as np + +from .vae import VAE, get_my_vae +from .distributions import DiagonalGaussianDistribution +from ..bigvgan import BigVGAN +from ..bigvgan_v2.bigvgan import BigVGAN as BigVGANv2 + + + + +class AutoEncoderModule(nn.Module): + + def __init__(self, + *, + vae_ckpt_path, + vocoder_ckpt_path: Optional[str] = None, + mode: Literal['16k', '44k'], + need_vae_encoder: bool = True): + super().__init__() + self.vae: VAE = get_my_vae(mode).eval() + vae_state_dict = torch.load(vae_ckpt_path, weights_only=True, map_location='cpu') + self.vae.load_state_dict(vae_state_dict) + self.vae.remove_weight_norm() + + if mode == '16k': + assert vocoder_ckpt_path is not None + self.vocoder = BigVGAN(vocoder_ckpt_path).eval() + elif mode == '44k': + # vocoder_ckpt_path can be: + # - None (default): pull 'nvidia/bigvgan_v2_44khz_128band_512x' from HF Hub + # - HF repo id (e.g., 'nvidia/bigvgan_v2_44khz_128band') + # - Local directory containing bigvgan_generator.pt + config.json + bigvgan_v2_id = vocoder_ckpt_path or 'nvidia/bigvgan_v2_44khz_128band_512x' + self.vocoder = BigVGANv2.from_pretrained(bigvgan_v2_id, + use_cuda_kernel=False) + self.vocoder.remove_weight_norm() + else: + raise ValueError(f'Unknown mode: {mode}') + + for param in self.parameters(): + param.requires_grad = False + + if not need_vae_encoder: + del self.vae.encoder + + @torch.inference_mode() + def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution: + return self.vae.encode(x) + + @torch.inference_mode() + def decode(self, z: torch.Tensor) -> torch.Tensor: + return self.vae.decode(z) + + @torch.inference_mode() + def vocode(self, spec: torch.Tensor) -> torch.Tensor: + return self.vocoder(spec) diff --git a/unison/models/mmaudio/ext/autoencoder/distributions.py b/unison/models/mmaudio/ext/autoencoder/distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..996f641c90ca9b793446d3a01cf5464516227537 --- /dev/null +++ b/unison/models/mmaudio/ext/autoencoder/distributions.py @@ -0,0 +1,45 @@ +from typing import Optional +import torch +import numpy as np + + +class DiagonalGaussianDistribution: + + def __init__(self, parameters, deterministic=False): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device) + + def sample(self, rng: Optional[torch.Generator] = None): + # x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device) + + r = torch.empty_like(self.mean).normal_(generator=rng) + x = self.mean + self.std * r + + return x + + def kl(self, other=None): + if self.deterministic: + return torch.Tensor([0.]) + else: + if other is None: + + return 0.5 * torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar + else: + return 0.5 * (torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var - 1.0 - self.logvar + other.logvar) + + def nll(self, sample, dims=[1, 2, 3]): + if self.deterministic: + return torch.Tensor([0.]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims) + + def mode(self): + return self.mean \ No newline at end of file diff --git a/unison/models/mmaudio/ext/autoencoder/edm2_utils.py b/unison/models/mmaudio/ext/autoencoder/edm2_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a18ffba5cc42214fddf1300034be2eff2760025c --- /dev/null +++ b/unison/models/mmaudio/ext/autoencoder/edm2_utils.py @@ -0,0 +1,168 @@ +# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# This work is licensed under a Creative Commons +# Attribution-NonCommercial-ShareAlike 4.0 International License. +# You should have received a copy of the license along with this +# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/ +"""Improved diffusion model architecture proposed in the paper +"Analyzing and Improving the Training Dynamics of Diffusion Models".""" + +import numpy as np +import torch + +#---------------------------------------------------------------------------- +# Variant of constant() that inherits dtype and device from the given +# reference tensor by default. + +_constant_cache = dict() + + +def constant(value, shape=None, dtype=None, device=None, memory_format=None): + value = np.asarray(value) + if shape is not None: + shape = tuple(shape) + if dtype is None: + dtype = torch.get_default_dtype() + if device is None: + device = torch.device('cpu') + if memory_format is None: + memory_format = torch.contiguous_format + + key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format) + tensor = _constant_cache.get(key, None) + if tensor is None: + tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device) + if shape is not None: + tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape)) + tensor = tensor.contiguous(memory_format=memory_format) + _constant_cache[key] = tensor + return tensor + + +def const_like(ref, value, shape=None, dtype=None, device=None, memory_format=None): + if dtype is None: + dtype = ref.dtype + if device is None: + device = ref.device + return constant(value, shape=shape, dtype=dtype, device=device, memory_format=memory_format) + + +#---------------------------------------------------------------------------- +# Normalize given tensor to unit magnitude with respect to the given +# dimensions. Default = all dimensions except the first. + + +def normalize(x, dim=None, eps=1e-4): + if dim is None: + dim = list(range(1, x.ndim)) + norm = torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32) + norm = torch.add(eps, norm, alpha=np.sqrt(norm.numel() / x.numel())) + return x / norm.to(x.dtype) + + +class Normalize(torch.nn.Module): + + def __init__(self, dim=None, eps=1e-4): + super().__init__() + self.dim = dim + self.eps = eps + + def forward(self, x): + return normalize(x, dim=self.dim, eps=self.eps) + + +#---------------------------------------------------------------------------- +# Upsample or downsample the given tensor with the given filter, +# or keep it as is. + + +def resample(x, f=[1, 1], mode='keep'): + if mode == 'keep': + return x + f = np.float32(f) + assert f.ndim == 1 and len(f) % 2 == 0 + pad = (len(f) - 1) // 2 + f = f / f.sum() + f = np.outer(f, f)[np.newaxis, np.newaxis, :, :] + f = const_like(x, f) + c = x.shape[1] + if mode == 'down': + return torch.nn.functional.conv2d(x, + f.tile([c, 1, 1, 1]), + groups=c, + stride=2, + padding=(pad, )) + assert mode == 'up' + return torch.nn.functional.conv_transpose2d(x, (f * 4).tile([c, 1, 1, 1]), + groups=c, + stride=2, + padding=(pad, )) + + +#---------------------------------------------------------------------------- +# Magnitude-preserving SiLU (Equation 81). + + +def mp_silu(x): + return torch.nn.functional.silu(x) / 0.596 + + +class MPSiLU(torch.nn.Module): + + def forward(self, x): + return mp_silu(x) + + +#---------------------------------------------------------------------------- +# Magnitude-preserving sum (Equation 88). + + +def mp_sum(a, b, t=0.5): + return a.lerp(b, t) / np.sqrt((1 - t)**2 + t**2) + + +#---------------------------------------------------------------------------- +# Magnitude-preserving concatenation (Equation 103). + + +def mp_cat(a, b, dim=1, t=0.5): + Na = a.shape[dim] + Nb = b.shape[dim] + C = np.sqrt((Na + Nb) / ((1 - t)**2 + t**2)) + wa = C / np.sqrt(Na) * (1 - t) + wb = C / np.sqrt(Nb) * t + return torch.cat([wa * a, wb * b], dim=dim) + + +#---------------------------------------------------------------------------- +# Magnitude-preserving convolution or fully-connected layer (Equation 47) +# with force weight normalization (Equation 66). + + +class MPConv1D(torch.nn.Module): + + def __init__(self, in_channels, out_channels, kernel_size): + super().__init__() + self.out_channels = out_channels + self.weight = torch.nn.Parameter(torch.randn(out_channels, in_channels, kernel_size)) + + self.weight_norm_removed = False + + def forward(self, x, gain=1): + assert self.weight_norm_removed, 'call remove_weight_norm() before inference' + + w = self.weight * gain + if w.ndim == 2: + return x @ w.t() + assert w.ndim == 3 + return torch.nn.functional.conv1d(x, w, padding=(w.shape[-1] // 2, )) + + def remove_weight_norm(self): + w = self.weight.to(torch.float32) + w = normalize(w) # traditional weight normalization + w = w / np.sqrt(w[0].numel()) + w = w.to(self.weight.dtype) + self.weight.data.copy_(w) + + self.weight_norm_removed = True + return self diff --git a/unison/models/mmaudio/ext/autoencoder/vae.py b/unison/models/mmaudio/ext/autoencoder/vae.py new file mode 100644 index 0000000000000000000000000000000000000000..b4fe72b51fad7392328f8403aef1c6979c0cfb75 --- /dev/null +++ b/unison/models/mmaudio/ext/autoencoder/vae.py @@ -0,0 +1,363 @@ +import logging +from typing import Optional + +import torch +import torch.nn as nn + +from .edm2_utils import MPConv1D +from .vae_modules import (AttnBlock1D, Downsample1D, ResnetBlock1D, + Upsample1D, nonlinearity) +from .distributions import DiagonalGaussianDistribution + +log = logging.getLogger() + +DATA_MEAN_80D = [ + -1.6058, -1.3676, -1.2520, -1.2453, -1.2078, -1.2224, -1.2419, -1.2439, -1.2922, -1.2927, + -1.3170, -1.3543, -1.3401, -1.3836, -1.3907, -1.3912, -1.4313, -1.4152, -1.4527, -1.4728, + -1.4568, -1.5101, -1.5051, -1.5172, -1.5623, -1.5373, -1.5746, -1.5687, -1.6032, -1.6131, + -1.6081, -1.6331, -1.6489, -1.6489, -1.6700, -1.6738, -1.6953, -1.6969, -1.7048, -1.7280, + -1.7361, -1.7495, -1.7658, -1.7814, -1.7889, -1.8064, -1.8221, -1.8377, -1.8417, -1.8643, + -1.8857, -1.8929, -1.9173, -1.9379, -1.9531, -1.9673, -1.9824, -2.0042, -2.0215, -2.0436, + -2.0766, -2.1064, -2.1418, -2.1855, -2.2319, -2.2767, -2.3161, -2.3572, -2.3954, -2.4282, + -2.4659, -2.5072, -2.5552, -2.6074, -2.6584, -2.7107, -2.7634, -2.8266, -2.8981, -2.9673 +] + +DATA_STD_80D = [ + 1.0291, 1.0411, 1.0043, 0.9820, 0.9677, 0.9543, 0.9450, 0.9392, 0.9343, 0.9297, 0.9276, 0.9263, + 0.9242, 0.9254, 0.9232, 0.9281, 0.9263, 0.9315, 0.9274, 0.9247, 0.9277, 0.9199, 0.9188, 0.9194, + 0.9160, 0.9161, 0.9146, 0.9161, 0.9100, 0.9095, 0.9145, 0.9076, 0.9066, 0.9095, 0.9032, 0.9043, + 0.9038, 0.9011, 0.9019, 0.9010, 0.8984, 0.8983, 0.8986, 0.8961, 0.8962, 0.8978, 0.8962, 0.8973, + 0.8993, 0.8976, 0.8995, 0.9016, 0.8982, 0.8972, 0.8974, 0.8949, 0.8940, 0.8947, 0.8936, 0.8939, + 0.8951, 0.8956, 0.9017, 0.9167, 0.9436, 0.9690, 1.0003, 1.0225, 1.0381, 1.0491, 1.0545, 1.0604, + 1.0761, 1.0929, 1.1089, 1.1196, 1.1176, 1.1156, 1.1117, 1.1070 +] + +DATA_MEAN_128D = [ + -3.3462, -2.6723, -2.4893, -2.3143, -2.2664, -2.3317, -2.1802, -2.4006, -2.2357, -2.4597, + -2.3717, -2.4690, -2.5142, -2.4919, -2.6610, -2.5047, -2.7483, -2.5926, -2.7462, -2.7033, + -2.7386, -2.8112, -2.7502, -2.9594, -2.7473, -3.0035, -2.8891, -2.9922, -2.9856, -3.0157, + -3.1191, -2.9893, -3.1718, -3.0745, -3.1879, -3.2310, -3.1424, -3.2296, -3.2791, -3.2782, + -3.2756, -3.3134, -3.3509, -3.3750, -3.3951, -3.3698, -3.4505, -3.4509, -3.5089, -3.4647, + -3.5536, -3.5788, -3.5867, -3.6036, -3.6400, -3.6747, -3.7072, -3.7279, -3.7283, -3.7795, + -3.8259, -3.8447, -3.8663, -3.9182, -3.9605, -3.9861, -4.0105, -4.0373, -4.0762, -4.1121, + -4.1488, -4.1874, -4.2461, -4.3170, -4.3639, -4.4452, -4.5282, -4.6297, -4.7019, -4.7960, + -4.8700, -4.9507, -5.0303, -5.0866, -5.1634, -5.2342, -5.3242, -5.4053, -5.4927, -5.5712, + -5.6464, -5.7052, -5.7619, -5.8410, -5.9188, -6.0103, -6.0955, -6.1673, -6.2362, -6.3120, + -6.3926, -6.4797, -6.5565, -6.6511, -6.8130, -6.9961, -7.1275, -7.2457, -7.3576, -7.4663, + -7.6136, -7.7469, -7.8815, -8.0132, -8.1515, -8.3071, -8.4722, -8.7418, -9.3975, -9.6628, + -9.7671, -9.8863, -9.9992, -10.0860, -10.1709, -10.5418, -11.2795, -11.3861 +] + +DATA_STD_128D = [ + 2.3804, 2.4368, 2.3772, 2.3145, 2.2803, 2.2510, 2.2316, 2.2083, 2.1996, 2.1835, 2.1769, 2.1659, + 2.1631, 2.1618, 2.1540, 2.1606, 2.1571, 2.1567, 2.1612, 2.1579, 2.1679, 2.1683, 2.1634, 2.1557, + 2.1668, 2.1518, 2.1415, 2.1449, 2.1406, 2.1350, 2.1313, 2.1415, 2.1281, 2.1352, 2.1219, 2.1182, + 2.1327, 2.1195, 2.1137, 2.1080, 2.1179, 2.1036, 2.1087, 2.1036, 2.1015, 2.1068, 2.0975, 2.0991, + 2.0902, 2.1015, 2.0857, 2.0920, 2.0893, 2.0897, 2.0910, 2.0881, 2.0925, 2.0873, 2.0960, 2.0900, + 2.0957, 2.0958, 2.0978, 2.0936, 2.0886, 2.0905, 2.0845, 2.0855, 2.0796, 2.0840, 2.0813, 2.0817, + 2.0838, 2.0840, 2.0917, 2.1061, 2.1431, 2.1976, 2.2482, 2.3055, 2.3700, 2.4088, 2.4372, 2.4609, + 2.4731, 2.4847, 2.5072, 2.5451, 2.5772, 2.6147, 2.6529, 2.6596, 2.6645, 2.6726, 2.6803, 2.6812, + 2.6899, 2.6916, 2.6931, 2.6998, 2.7062, 2.7262, 2.7222, 2.7158, 2.7041, 2.7485, 2.7491, 2.7451, + 2.7485, 2.7233, 2.7297, 2.7233, 2.7145, 2.6958, 2.6788, 2.6439, 2.6007, 2.4786, 2.2469, 2.1877, + 2.1392, 2.0717, 2.0107, 1.9676, 1.9140, 1.7102, 0.9101, 0.7164 +] + + +class VAE(nn.Module): + + def __init__( + self, + *, + data_dim: int, + embed_dim: int, + hidden_dim: int, + ): + super().__init__() + + if data_dim == 80: + self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_80D, dtype=torch.float32)) + self.data_std = nn.Buffer(torch.tensor(DATA_STD_80D, dtype=torch.float32)) + elif data_dim == 128: + self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32)) + self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32)) + + self.data_mean = self.data_mean.view(1, -1, 1) + self.data_std = self.data_std.view(1, -1, 1) + + self.encoder = Encoder1D( + dim=hidden_dim, + ch_mult=(1, 2, 4), + num_res_blocks=2, + attn_layers=[3], + down_layers=[0], + in_dim=data_dim, + embed_dim=embed_dim, + ) + self.decoder = Decoder1D( + dim=hidden_dim, + ch_mult=(1, 2, 4), + num_res_blocks=2, + attn_layers=[3], + down_layers=[0], + in_dim=data_dim, + out_dim=data_dim, + embed_dim=embed_dim, + ) + + self.embed_dim = embed_dim + # self.quant_conv = nn.Conv1d(2 * embed_dim, 2 * embed_dim, 1) + # self.post_quant_conv = nn.Conv1d(embed_dim, embed_dim, 1) + + self.initialize_weights() + + def initialize_weights(self): + pass + + def encode(self, x: torch.Tensor, normalize: bool = True) -> DiagonalGaussianDistribution: + if normalize: + x = self.normalize(x) + moments = self.encoder(x) + posterior = DiagonalGaussianDistribution(moments) + return posterior + + def decode(self, z: torch.Tensor, unnormalize: bool = True) -> torch.Tensor: + dec = self.decoder(z) + if unnormalize: + dec = self.unnormalize(dec) + return dec + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x - self.data_mean) / self.data_std + + def unnormalize(self, x: torch.Tensor) -> torch.Tensor: + return x * self.data_std + self.data_mean + + def forward( + self, + x: torch.Tensor, + sample_posterior: bool = True, + rng: Optional[torch.Generator] = None, + normalize: bool = True, + unnormalize: bool = True, + ) -> tuple[torch.Tensor, DiagonalGaussianDistribution]: + + posterior = self.encode(x, normalize=normalize) + if sample_posterior: + z = posterior.sample(rng) + else: + z = posterior.mode() + dec = self.decode(z, unnormalize=unnormalize) + return dec, posterior + + def load_weights(self, src_dict) -> None: + self.load_state_dict(src_dict, strict=True) + + @property + def device(self) -> torch.device: + return next(self.parameters()).device + + def get_last_layer(self): + return self.decoder.conv_out.weight + + def remove_weight_norm(self): + for name, m in self.named_modules(): + if isinstance(m, MPConv1D): + m.remove_weight_norm() + log.debug(f"Removed weight norm from {name}") + return self + + +class Encoder1D(nn.Module): + + def __init__(self, + *, + dim: int, + ch_mult: tuple[int] = (1, 2, 4, 8), + num_res_blocks: int, + attn_layers: list[int] = [], + down_layers: list[int] = [], + resamp_with_conv: bool = True, + in_dim: int, + embed_dim: int, + double_z: bool = True, + kernel_size: int = 3, + clip_act: float = 256.0): + super().__init__() + self.dim = dim + self.num_layers = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.in_channels = in_dim + self.clip_act = clip_act + self.down_layers = down_layers + self.attn_layers = attn_layers + self.conv_in = MPConv1D(in_dim, self.dim, kernel_size=kernel_size) + + in_ch_mult = (1, ) + tuple(ch_mult) + self.in_ch_mult = in_ch_mult + # downsampling + self.down = nn.ModuleList() + for i_level in range(self.num_layers): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = dim * in_ch_mult[i_level] + block_out = dim * ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append( + ResnetBlock1D(in_dim=block_in, + out_dim=block_out, + kernel_size=kernel_size, + use_norm=True)) + block_in = block_out + if i_level in attn_layers: + attn.append(AttnBlock1D(block_in)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level in down_layers: + down.downsample = Downsample1D(block_in, resamp_with_conv) + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock1D(in_dim=block_in, + out_dim=block_in, + kernel_size=kernel_size, + use_norm=True) + self.mid.attn_1 = AttnBlock1D(block_in) + self.mid.block_2 = ResnetBlock1D(in_dim=block_in, + out_dim=block_in, + kernel_size=kernel_size, + use_norm=True) + + # end + self.conv_out = MPConv1D(block_in, + 2 * embed_dim if double_z else embed_dim, + kernel_size=kernel_size) + + self.learnable_gain = nn.Parameter(torch.zeros([])) + + def forward(self, x): + + # downsampling + hs = [self.conv_in(x)] + for i_level in range(self.num_layers): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + h = h.clamp(-self.clip_act, self.clip_act) + hs.append(h) + if i_level in self.down_layers: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + h = h.clamp(-self.clip_act, self.clip_act) + + # end + h = nonlinearity(h) + h = self.conv_out(h, gain=(self.learnable_gain + 1)) + return h + + +class Decoder1D(nn.Module): + + def __init__(self, + *, + dim: int, + out_dim: int, + ch_mult: tuple[int] = (1, 2, 4, 8), + num_res_blocks: int, + attn_layers: list[int] = [], + down_layers: list[int] = [], + kernel_size: int = 3, + resamp_with_conv: bool = True, + in_dim: int, + embed_dim: int, + clip_act: float = 256.0): + super().__init__() + self.ch = dim + self.num_layers = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.in_channels = in_dim + self.clip_act = clip_act + self.down_layers = [i + 1 for i in down_layers] # each downlayer add one + + # compute in_ch_mult, block_in and curr_res at lowest res + block_in = dim * ch_mult[self.num_layers - 1] + + # z to block_in + self.conv_in = MPConv1D(embed_dim, block_in, kernel_size=kernel_size) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True) + self.mid.attn_1 = AttnBlock1D(block_in) + self.mid.block_2 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_layers)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = dim * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + block.append(ResnetBlock1D(in_dim=block_in, out_dim=block_out, use_norm=True)) + block_in = block_out + if i_level in attn_layers: + attn.append(AttnBlock1D(block_in)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level in self.down_layers: + up.upsample = Upsample1D(block_in, resamp_with_conv) + self.up.insert(0, up) # prepend to get consistent order + + # end + self.conv_out = MPConv1D(block_in, out_dim, kernel_size=kernel_size) + self.learnable_gain = nn.Parameter(torch.zeros([])) + + def forward(self, z): + # z to block_in + h = self.conv_in(z) + + # middle + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + h = h.clamp(-self.clip_act, self.clip_act) + + # upsampling + for i_level in reversed(range(self.num_layers)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + h = h.clamp(-self.clip_act, self.clip_act) + if i_level in self.down_layers: + h = self.up[i_level].upsample(h) + + h = nonlinearity(h) + h = self.conv_out(h, gain=(self.learnable_gain + 1)) + return h + + +def VAE_16k(**kwargs) -> VAE: + return VAE(data_dim=80, embed_dim=20, hidden_dim=384, **kwargs) + + +def VAE_44k(**kwargs) -> VAE: + return VAE(data_dim=128, embed_dim=40, hidden_dim=512, **kwargs) + + +def get_my_vae(name: str, **kwargs) -> VAE: + if name == '16k': + return VAE_16k(**kwargs) + if name == '44k': + return VAE_44k(**kwargs) + raise ValueError(f'Unknown model: {name}') + + diff --git a/unison/models/mmaudio/ext/autoencoder/vae_modules.py b/unison/models/mmaudio/ext/autoencoder/vae_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..f083a0efb572f3c0ee595aac28a78fcba08acb40 --- /dev/null +++ b/unison/models/mmaudio/ext/autoencoder/vae_modules.py @@ -0,0 +1,117 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from .edm2_utils import (MPConv1D, mp_silu, mp_sum, normalize) + + +def nonlinearity(x): + # swish + return mp_silu(x) + + +class ResnetBlock1D(nn.Module): + + def __init__(self, *, in_dim, out_dim=None, conv_shortcut=False, kernel_size=3, use_norm=True): + super().__init__() + self.in_dim = in_dim + out_dim = in_dim if out_dim is None else out_dim + self.out_dim = out_dim + self.use_conv_shortcut = conv_shortcut + self.use_norm = use_norm + + self.conv1 = MPConv1D(in_dim, out_dim, kernel_size=kernel_size) + self.conv2 = MPConv1D(out_dim, out_dim, kernel_size=kernel_size) + if self.in_dim != self.out_dim: + if self.use_conv_shortcut: + self.conv_shortcut = MPConv1D(in_dim, out_dim, kernel_size=kernel_size) + else: + self.nin_shortcut = MPConv1D(in_dim, out_dim, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + + # pixel norm + if self.use_norm: + x = normalize(x, dim=1) + + h = x + h = nonlinearity(h) + h = self.conv1(h) + + h = nonlinearity(h) + h = self.conv2(h) + + if self.in_dim != self.out_dim: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + + return mp_sum(x, h, t=0.3) + + +class AttnBlock1D(nn.Module): + + def __init__(self, in_channels, num_heads=1): + super().__init__() + self.in_channels = in_channels + + self.num_heads = num_heads + self.qkv = MPConv1D(in_channels, in_channels * 3, kernel_size=1) + self.proj_out = MPConv1D(in_channels, in_channels, kernel_size=1) + + def forward(self, x): + h = x + y = self.qkv(h) + y = y.reshape(y.shape[0], self.num_heads, -1, 3, y.shape[-1]) + q, k, v = normalize(y, dim=2).unbind(3) + + q = rearrange(q, 'b h c l -> b h l c') + k = rearrange(k, 'b h c l -> b h l c') + v = rearrange(v, 'b h c l -> b h l c') + + h = F.scaled_dot_product_attention(q, k, v) + h = rearrange(h, 'b h l c -> b (h c) l') + + h = self.proj_out(h) + + return mp_sum(x, h, t=0.3) + + +class Upsample1D(nn.Module): + + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + self.conv = MPConv1D(in_channels, in_channels, kernel_size=3) + + def forward(self, x): + x = F.interpolate(x, scale_factor=2.0, mode='nearest-exact') # support 3D tensor(B,C,T) + if self.with_conv: + x = self.conv(x) + return x + + +class Downsample1D(nn.Module): + + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + self.conv1 = MPConv1D(in_channels, in_channels, kernel_size=1) + self.conv2 = MPConv1D(in_channels, in_channels, kernel_size=1) + + def forward(self, x): + + if self.with_conv: + x = self.conv1(x) + + x = F.avg_pool1d(x, kernel_size=2, stride=2) + + if self.with_conv: + x = self.conv2(x) + + return x diff --git a/unison/models/mmaudio/ext/bigvgan/LICENSE b/unison/models/mmaudio/ext/bigvgan/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..e9663595cc28938f88d6299acd3ba791542e4c0c --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2022 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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/__init__.py b/unison/models/mmaudio/ext/bigvgan/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..00f13e9bf9ccb0b4ec37e1c70869f9a9a538871f --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/__init__.py @@ -0,0 +1 @@ +from .bigvgan import BigVGAN diff --git a/unison/models/mmaudio/ext/bigvgan/activations.py b/unison/models/mmaudio/ext/bigvgan/activations.py new file mode 100644 index 0000000000000000000000000000000000000000..61f2808a5466b3cf4d041059700993af5527dd29 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/activations.py @@ -0,0 +1,120 @@ +# 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 \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/alias_free_torch/__init__.py b/unison/models/mmaudio/ext/bigvgan/alias_free_torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a2318b63198250856809c0cb46210a4147b829bc --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/alias_free_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 * \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/alias_free_torch/act.py b/unison/models/mmaudio/ext/bigvgan/alias_free_torch/act.py new file mode 100644 index 0000000000000000000000000000000000000000..028debd697dd60458aae75010057df038bd3518a --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/alias_free_torch/act.py @@ -0,0 +1,28 @@ +# 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 .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 \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/alias_free_torch/filter.py b/unison/models/mmaudio/ext/bigvgan/alias_free_torch/filter.py new file mode 100644 index 0000000000000000000000000000000000000000..7ad6ea87c1f10ddd94c544037791d7a4634d5ae1 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/alias_free_torch/filter.py @@ -0,0 +1,95 @@ +# 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., 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.: + beta = 0.1102 * (A - 8.7) + elif A >= 21.: + beta = 0.5842 * (A - 21)**0.4 + 0.07886 * (A - 21.) + else: + beta = 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.: + 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 \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/alias_free_torch/resample.py b/unison/models/mmaudio/ext/bigvgan/alias_free_torch/resample.py new file mode 100644 index 0000000000000000000000000000000000000000..750e6c3402cc5ac939c4b9d075246562e0e1d1a7 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/alias_free_torch/resample.py @@ -0,0 +1,49 @@ +# 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 .filter import LowPassFilter1d +from .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 \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/bigvgan.py b/unison/models/mmaudio/ext/bigvgan/bigvgan.py new file mode 100644 index 0000000000000000000000000000000000000000..4955f6918a229a56a55bdbf6f1cc744d9a698e19 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/bigvgan.py @@ -0,0 +1,32 @@ +from pathlib import Path + +import torch +import torch.nn as nn +from omegaconf import OmegaConf + +from .models import BigVGANVocoder + +_bigvgan_vocoder_path = Path(__file__).parent / 'bigvgan_vocoder.yml' + + +class BigVGAN(nn.Module): + + def __init__(self, ckpt_path, config_path=_bigvgan_vocoder_path): + super().__init__() + vocoder_cfg = OmegaConf.load(config_path) + self.vocoder = BigVGANVocoder(vocoder_cfg).eval() + vocoder_ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=True)['generator'] + self.vocoder.load_state_dict(vocoder_ckpt) + + self.weight_norm_removed = False + self.remove_weight_norm() + + @torch.inference_mode() + def forward(self, x): + assert self.weight_norm_removed, 'call remove_weight_norm() before inference' + return self.vocoder(x) + + def remove_weight_norm(self): + self.vocoder.remove_weight_norm() + self.weight_norm_removed = True + return self diff --git a/unison/models/mmaudio/ext/bigvgan/bigvgan_vocoder.yml b/unison/models/mmaudio/ext/bigvgan/bigvgan_vocoder.yml new file mode 100644 index 0000000000000000000000000000000000000000..d4db31ec45336e757d94d5099ed16cb3c906c24a --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/bigvgan_vocoder.yml @@ -0,0 +1,63 @@ +resblock: '1' +num_gpus: 0 +batch_size: 64 +num_mels: 80 +learning_rate: 0.0001 +adam_b1: 0.8 +adam_b2: 0.99 +lr_decay: 0.999 +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 +num_workers: 4 +dist_config: + dist_backend: nccl + dist_url: tcp://localhost:54341 + world_size: 1 diff --git a/unison/models/mmaudio/ext/bigvgan/env.py b/unison/models/mmaudio/ext/bigvgan/env.py new file mode 100644 index 0000000000000000000000000000000000000000..b8be238d4db710c8c9a338d336baea0138f18d1f --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/env.py @@ -0,0 +1,18 @@ +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. +# LICENSE is in incl_licenses directory. + +import os +import shutil + + +class AttrDict(dict): + def __init__(self, *args, **kwargs): + super(AttrDict, self).__init__(*args, **kwargs) + self.__dict__ = self + + +def build_env(config, config_name, path): + t_path = os.path.join(path, config_name) + if config != t_path: + os.makedirs(path, exist_ok=True) + shutil.copyfile(config, os.path.join(path, config_name)) \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_1 b/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_1 new file mode 100644 index 0000000000000000000000000000000000000000..5afae394d6b37da0e12ba6b290d2512687f421ac --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_1 @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2020 Jungil Kong + +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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_2 b/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_2 new file mode 100644 index 0000000000000000000000000000000000000000..322b758863c4219be68291ae3826218baa93cb4c --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_2 @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2020 Edward Dixon + +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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_3 b/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_3 new file mode 100644 index 0000000000000000000000000000000000000000..56ee3c8c4cc2b4b32e0975d17258f9ba515fdbcc --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_3 @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_4 b/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_4 new file mode 100644 index 0000000000000000000000000000000000000000..48fd1a1ba8d81a94b6c7d1c2ff1a1f307cc5371d --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_4 @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2019, Seungwon Park 박승원 +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_5 b/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_5 new file mode 100644 index 0000000000000000000000000000000000000000..01ae5538e6b7c787bb4f5d6f2cd9903520d6e465 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/incl_licenses/LICENSE_5 @@ -0,0 +1,16 @@ +Copyright 2020 Alexandre Défossez + +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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan/models.py b/unison/models/mmaudio/ext/bigvgan/models.py new file mode 100644 index 0000000000000000000000000000000000000000..b39e4715586028938058365b70e0b1bd20d98d3c --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/models.py @@ -0,0 +1,255 @@ +# Copyright (c) 2022 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 torch +import torch.nn as nn +from torch.nn import Conv1d, ConvTranspose1d +from torch.nn.utils.parametrizations import weight_norm +from torch.nn.utils.parametrize import remove_parametrizations + +from . import activations +from .alias_free_torch import * +from .utils import get_padding, init_weights + +LRELU_SLOPE = 0.1 + + +class AMPBlock1(torch.nn.Module): + + def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5), activation=None): + super(AMPBlock1, self).__init__() + self.h = h + + self.convs1 = nn.ModuleList([ + weight_norm( + Conv1d(channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]))), + weight_norm( + Conv1d(channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]))), + weight_norm( + Conv1d(channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + padding=get_padding(kernel_size, dilation[2]))) + ]) + self.convs1.apply(init_weights) + + self.convs2 = nn.ModuleList([ + weight_norm( + Conv1d(channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1))), + weight_norm( + Conv1d(channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1))), + weight_norm( + Conv1d(channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1))) + ]) + self.convs2.apply(init_weights) + + self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers + + if activation == 'snake': # periodic nonlinearity with snake function and anti-aliasing + self.activations = nn.ModuleList([ + Activation1d( + activation=activations.Snake(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ]) + elif activation == 'snakebeta': # periodic nonlinearity with snakebeta function and anti-aliasing + 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_parametrizations(l, 'weight') + for l in self.convs2: + remove_parametrizations(l, 'weight') + + +class AMPBlock2(torch.nn.Module): + + def __init__(self, h, channels, kernel_size=3, dilation=(1, 3), activation=None): + super(AMPBlock2, self).__init__() + self.h = h + + self.convs = nn.ModuleList([ + weight_norm( + Conv1d(channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]))), + weight_norm( + Conv1d(channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]))) + ]) + self.convs.apply(init_weights) + + self.num_layers = len(self.convs) # total number of conv layers + + if activation == 'snake': # periodic nonlinearity with snake function and anti-aliasing + self.activations = nn.ModuleList([ + Activation1d( + activation=activations.Snake(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ]) + elif activation == 'snakebeta': # periodic nonlinearity with snakebeta function and anti-aliasing + 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_parametrizations(l, 'weight') + + +class BigVGANVocoder(torch.nn.Module): + # this is our main BigVGAN model. Applies anti-aliased periodic activation for resblocks. + def __init__(self, h): + super().__init__() + self.h = h + + 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 + resblock = AMPBlock1 if h.resblock == '1' else AMPBlock2 + + # 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(h, ch, k, d, activation=h.activation)) + + # post conv + if h.activation == "snake": # periodic nonlinearity with snake function and anti-aliasing + activation_post = activations.Snake(ch, alpha_logscale=h.snake_logscale) + self.activation_post = Activation1d(activation=activation_post) + elif h.activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing + activation_post = activations.SnakeBeta(ch, alpha_logscale=h.snake_logscale) + self.activation_post = Activation1d(activation=activation_post) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3)) + + # weight initialization + for i in range(len(self.ups)): + self.ups[i].apply(init_weights) + self.conv_post.apply(init_weights) + + 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) + x = torch.tanh(x) + + return x + + def remove_weight_norm(self): + print('Removing weight norm...') + for l in self.ups: + for l_i in l: + remove_parametrizations(l_i, 'weight') + for l in self.resblocks: + l.remove_weight_norm() + remove_parametrizations(self.conv_pre, 'weight') + remove_parametrizations(self.conv_post, 'weight') diff --git a/unison/models/mmaudio/ext/bigvgan/utils.py b/unison/models/mmaudio/ext/bigvgan/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..aff7e653533d3390756c53a0215801b06cc924b5 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan/utils.py @@ -0,0 +1,31 @@ +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. +# LICENSE is in incl_licenses directory. + +import os + +import torch +from torch.nn.utils.parametrizations import weight_norm + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +def apply_weight_norm(m): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + weight_norm(m) + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +def load_checkpoint(filepath, device): + assert os.path.isfile(filepath) + print("Loading '{}'".format(filepath)) + checkpoint_dict = torch.load(filepath, map_location=device) + print("Complete.") + return checkpoint_dict diff --git a/unison/models/mmaudio/ext/bigvgan_v2/LICENSE b/unison/models/mmaudio/ext/bigvgan_v2/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..4c78361c86d4f685117d60d6623e2197fcfed706 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/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. diff --git a/unison/models/mmaudio/ext/bigvgan_v2/__init__.py b/unison/models/mmaudio/ext/bigvgan_v2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/unison/models/mmaudio/ext/bigvgan_v2/activations.py b/unison/models/mmaudio/ext/bigvgan_v2/activations.py new file mode 100644 index 0000000000000000000000000000000000000000..4f08ddab5b55d6dcaf3e968af98889e0770c44f5 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/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 diff --git a/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/__init__.py b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/activation1d.py b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/activation1d.py new file mode 100644 index 0000000000000000000000000000000000000000..fbc0fd8f28a37ad949fbdb9832f51b5b933c6ff2 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/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 diff --git a/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/anti_alias_activation.cpp b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/anti_alias_activation.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c5651f77143bd678169eb11564a7cf7a7969a59e --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/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 + +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)"); +} \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/anti_alias_activation_cuda.cu b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/anti_alias_activation_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..8c442334869fe72d639ec203fa4fac07f96a0ee1 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/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 +#include +#include +#include +#include +#include +#include +#include "type_shim.h" +#include +#include +#include +#include +#include + +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 + __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 + 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 + <<>>(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(input.data_ptr()); + void *up_filter_ptr = static_cast(up_filter.data_ptr()); + void *down_filter_ptr = static_cast(down_filter.data_ptr()); + void *alpha_ptr = static_cast(alpha.data_ptr()); + void *beta_ptr = static_cast(beta.data_ptr()); + void *anti_alias_activation_results_ptr = static_cast(anti_alias_activation_results.data_ptr()); + + DISPATCH_FLOAT_HALF_AND_BFLOAT( + input.scalar_type(), + "dispatch anti alias activation_forward", + dispatch_anti_alias_activation_forward( + reinterpret_cast(anti_alias_activation_results_ptr), + reinterpret_cast(input_ptr), + reinterpret_cast(up_filter_ptr), + reinterpret_cast(down_filter_ptr), + reinterpret_cast(alpha_ptr), + reinterpret_cast(beta_ptr), + batches, + channels, + seq_len);); + return anti_alias_activation_results; +} \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/compat.h b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/compat.h new file mode 100644 index 0000000000000000000000000000000000000000..25818b2edf4cb0dc9130e62c7c4de8d16a01baa5 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/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 diff --git a/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/load.py b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/load.py new file mode 100644 index 0000000000000000000000000000000000000000..ca5d01de398249e75e9e2298958764acb436edba --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/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") diff --git a/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/type_shim.h b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/cuda/type_shim.h new file mode 100644 index 0000000000000000000000000000000000000000..5db7e8a397e982d4d30d16ab6060814b98b7ab83 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/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 +#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), "'"); \ + } diff --git a/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/torch/__init__.py b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8f756ed83f87f9839e457b240f60469bc187707d --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/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 * diff --git a/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/torch/act.py b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/torch/act.py new file mode 100644 index 0000000000000000000000000000000000000000..f25dda1b5626319819a7749e3bad0b7e7bcd34f5 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/torch/act.py @@ -0,0 +1,32 @@ +# 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 .resample import (DownSample1d, UpSample1d) + + +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 diff --git a/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/torch/filter.py b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/torch/filter.py new file mode 100644 index 0000000000000000000000000000000000000000..0fa35b0d5ddf8d6cb04cd9d47364ca033cebcd32 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/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 diff --git a/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/torch/resample.py b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/torch/resample.py new file mode 100644 index 0000000000000000000000000000000000000000..9ee1225f37fe9e876a982836d8e0a4aa87e93eee --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/alias_free_activation/torch/resample.py @@ -0,0 +1,53 @@ +# 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 .filter import (LowPassFilter1d, 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 diff --git a/unison/models/mmaudio/ext/bigvgan_v2/bigvgan.py b/unison/models/mmaudio/ext/bigvgan_v2/bigvgan.py new file mode 100644 index 0000000000000000000000000000000000000000..cbcb21fcae85fd85210db13fd4ad6c20bd5c47db --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/bigvgan.py @@ -0,0 +1,439 @@ +# 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 json +import os +from pathlib import Path +from typing import Dict, Optional, Union + +import torch +import torch.nn as nn +from huggingface_hub import PyTorchModelHubMixin, hf_hub_download +from torch.nn import Conv1d, ConvTranspose1d +from torch.nn.utils.parametrizations import weight_norm +from torch.nn.utils.parametrize import remove_parametrizations + +from . import activations +from .alias_free_activation.torch.act import \ + Activation1d as TorchActivation1d +from .env import AttrDict +from .utils import get_padding, init_weights + + +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_parametrizations(l, 'weight') + for l in self.convs2: + remove_parametrizations(l, 'weight') + + +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_parametrizations(l_i, 'weight') + for l in self.resblocks: + l.remove_weight_norm() + remove_parametrizations(self.conv_pre, 'weight') + remove_parametrizations(self.conv_post, 'weight') + 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, weights_only=True) + + 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 diff --git a/unison/models/mmaudio/ext/bigvgan_v2/env.py b/unison/models/mmaudio/ext/bigvgan_v2/env.py new file mode 100644 index 0000000000000000000000000000000000000000..b8be238d4db710c8c9a338d336baea0138f18d1f --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/env.py @@ -0,0 +1,18 @@ +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. +# LICENSE is in incl_licenses directory. + +import os +import shutil + + +class AttrDict(dict): + def __init__(self, *args, **kwargs): + super(AttrDict, self).__init__(*args, **kwargs) + self.__dict__ = self + + +def build_env(config, config_name, path): + t_path = os.path.join(path, config_name) + if config != t_path: + os.makedirs(path, exist_ok=True) + shutil.copyfile(config, os.path.join(path, config_name)) \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_1 b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_1 new file mode 100644 index 0000000000000000000000000000000000000000..5afae394d6b37da0e12ba6b290d2512687f421ac --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_1 @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2020 Jungil Kong + +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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_2 b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_2 new file mode 100644 index 0000000000000000000000000000000000000000..322b758863c4219be68291ae3826218baa93cb4c --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_2 @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2020 Edward Dixon + +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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_3 b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_3 new file mode 100644 index 0000000000000000000000000000000000000000..56ee3c8c4cc2b4b32e0975d17258f9ba515fdbcc --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_3 @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_4 b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_4 new file mode 100644 index 0000000000000000000000000000000000000000..48fd1a1ba8d81a94b6c7d1c2ff1a1f307cc5371d --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_4 @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2019, Seungwon Park 박승원 +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_5 b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_5 new file mode 100644 index 0000000000000000000000000000000000000000..01ae5538e6b7c787bb4f5d6f2cd9903520d6e465 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_5 @@ -0,0 +1,16 @@ +Copyright 2020 Alexandre Défossez + +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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_6 b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_6 new file mode 100644 index 0000000000000000000000000000000000000000..2569ec0b6c85f94f3cd071ba16e9028ccf156be2 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_6 @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023-present, Descript + +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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_7 b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_7 new file mode 100644 index 0000000000000000000000000000000000000000..c37bdaf99c6921f5849425d546069e972f52d7fa --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_7 @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 Charactr Inc. + +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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_8 b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_8 new file mode 100644 index 0000000000000000000000000000000000000000..ab3d7ffe795779f54e339078e4e752ad9019aae8 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/incl_licenses/LICENSE_8 @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 Amphion + +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. \ No newline at end of file diff --git a/unison/models/mmaudio/ext/bigvgan_v2/utils.py b/unison/models/mmaudio/ext/bigvgan_v2/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3b1d41670fa1ee257b2ed22c61086ba7a32c7cb0 --- /dev/null +++ b/unison/models/mmaudio/ext/bigvgan_v2/utils.py @@ -0,0 +1,31 @@ +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. +# LICENSE is in incl_licenses directory. + +import os + +import torch +from torch.nn.utils import weight_norm + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +def apply_weight_norm(m): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + weight_norm(m) + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +def load_checkpoint(filepath, device): + assert os.path.isfile(filepath) + print(f"Loading '{filepath}'") + checkpoint_dict = torch.load(filepath, map_location=device) + print("Complete.") + return checkpoint_dict diff --git a/unison/models/mmaudio/ext/mel_converter.py b/unison/models/mmaudio/ext/mel_converter.py new file mode 100644 index 0000000000000000000000000000000000000000..dc5d740e6b42b197f49a22e9b54570ebea7bd81d --- /dev/null +++ b/unison/models/mmaudio/ext/mel_converter.py @@ -0,0 +1,113 @@ +# Reference: # https://github.com/bytedance/Make-An-Audio-2 +from typing import Literal + +import torch +import torch.nn as nn +from librosa.filters import mel as librosa_mel_fn + + +def dynamic_range_compression_torch(x, C=1, clip_val=1e-5, *, norm_fn): + return norm_fn(torch.clamp(x, min=clip_val) * C) + + +def spectral_normalize_torch(magnitudes, norm_fn): + output = dynamic_range_compression_torch(magnitudes, norm_fn=norm_fn) + return output + + +class MelConverter(nn.Module): + + def __init__( + self, + *, + sampling_rate: float, + n_fft: int, + num_mels: int, + hop_size: int, + win_size: int, + fmin: float, + fmax: float, + norm_fn, + ): + super().__init__() + self.sampling_rate = sampling_rate + self.n_fft = n_fft + self.num_mels = num_mels + self.hop_size = hop_size + self.win_size = win_size + self.fmin = fmin + self.fmax = fmax + self.norm_fn = norm_fn + + mel = librosa_mel_fn(sr=self.sampling_rate, + n_fft=self.n_fft, + n_mels=self.num_mels, + fmin=self.fmin, + fmax=self.fmax) + mel_basis = torch.from_numpy(mel).float() + hann_window = torch.hann_window(self.win_size) + + self.register_buffer('mel_basis', mel_basis) + self.register_buffer('hann_window', hann_window) + + @property + def device(self): + return self.mel_basis.device + + def forward(self, waveform: torch.Tensor, center: bool = False) -> torch.Tensor: + waveform = waveform.clamp(min=-1., max=1.).to(self.device) + + # Dithering: only for 16kHz where fmax=Nyquist causes edge artifacts + if self.sampling_rate <= 16000: + dither_scale = 1e-4 + noise = torch.randn_like(waveform) * dither_scale + silence_mask = (torch.abs(waveform) < dither_scale) + waveform = torch.where(silence_mask, noise, waveform) + + waveform = torch.nn.functional.pad( + waveform.unsqueeze(1), + [int((self.n_fft - self.hop_size) / 2), + int((self.n_fft - self.hop_size) / 2)], + mode='reflect') + waveform = waveform.squeeze(1) + + spec = torch.stft(waveform, + self.n_fft, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=center, + pad_mode='reflect', + normalized=False, + onesided=True, + return_complex=True) + + spec = torch.view_as_real(spec) + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)).float() + spec = torch.matmul(self.mel_basis, spec) + spec = spectral_normalize_torch(spec, self.norm_fn) + + return spec + + +def get_mel_converter(mode: Literal['16k', '44k']) -> MelConverter: + if mode == '16k': + return MelConverter(sampling_rate=16_000, + n_fft=1024, + num_mels=80, + hop_size=256, + win_size=1024, + fmin=0, + fmax=8_000, + norm_fn=torch.log10) + elif mode == '44k': + return MelConverter(sampling_rate=44_100, + n_fft=2048, + num_mels=128, + hop_size=512, + win_size=2048, + fmin=0, + fmax=44100 / 2, + norm_fn=torch.log) + else: + raise ValueError(f'Unknown mode: {mode}') diff --git a/unison/models/mmaudio/ext/rotary_embeddings.py b/unison/models/mmaudio/ext/rotary_embeddings.py new file mode 100644 index 0000000000000000000000000000000000000000..1ea9d56278cb68b7577ed13148227c30ed98fd02 --- /dev/null +++ b/unison/models/mmaudio/ext/rotary_embeddings.py @@ -0,0 +1,35 @@ +from typing import Union + +import torch +from einops import rearrange +from torch import Tensor + +# Ref: https://github.com/black-forest-labs/flux/blob/main/src/flux/math.py +# Ref: https://github.com/lucidrains/rotary-embedding-torch + + +def compute_rope_rotations(length: int, + dim: int, + theta: int, + *, + freq_scaling: float = 1.0, + device: Union[torch.device, str] = 'cpu') -> Tensor: + assert dim % 2 == 0 + + with torch.amp.autocast(device_type='cuda', enabled=False): + pos = torch.arange(length, dtype=torch.float32, device=device) + freqs = 1.0 / (theta**(torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) + freqs *= freq_scaling + + rot = torch.einsum('..., f -> ... f', pos, freqs) + rot = torch.stack([torch.cos(rot), -torch.sin(rot), torch.sin(rot), torch.cos(rot)], dim=-1) + rot = rearrange(rot, 'n d (i j) -> 1 n d i j', i=2, j=2) + return rot + + +def apply_rope(x: Tensor, rot: Tensor) -> tuple[Tensor, Tensor]: + with torch.amp.autocast(device_type='cuda', enabled=False): + _x = x.float() + _x = _x.view(*_x.shape[:-1], -1, 1, 2) + x_out = rot[..., 0] * _x[..., 0] + rot[..., 1] * _x[..., 1] + return x_out.reshape(*x.shape).to(dtype=x.dtype) diff --git a/unison/models/mmaudio/ext/stft_converter.py b/unison/models/mmaudio/ext/stft_converter.py new file mode 100644 index 0000000000000000000000000000000000000000..62922067ef3b1d3b8727ec39e7d664ccb304d9fe --- /dev/null +++ b/unison/models/mmaudio/ext/stft_converter.py @@ -0,0 +1,183 @@ +# Reference: # https://github.com/bytedance/Make-An-Audio-2 + +import torch +import torch.nn as nn +import torchaudio +from einops import rearrange +from librosa.filters import mel as librosa_mel_fn + + +def dynamic_range_compression_torch(x, C=1, clip_val=1e-5, norm_fn=torch.log10): + return norm_fn(torch.clamp(x, min=clip_val) * C) + + +def spectral_normalize_torch(magnitudes, norm_fn): + output = dynamic_range_compression_torch(magnitudes, norm_fn=norm_fn) + return output + + +class STFTConverter(nn.Module): + + def __init__( + self, + *, + sampling_rate: float = 16_000, + n_fft: int = 1024, + num_mels: int = 128, + hop_size: int = 256, + win_size: int = 1024, + fmin: float = 0, + fmax: float = 8_000, + norm_fn=torch.log, + ): + super().__init__() + self.sampling_rate = sampling_rate + self.n_fft = n_fft + self.num_mels = num_mels + self.hop_size = hop_size + self.win_size = win_size + self.fmin = fmin + self.fmax = fmax + self.norm_fn = norm_fn + + mel = librosa_mel_fn(sr=self.sampling_rate, + n_fft=self.n_fft, + n_mels=self.num_mels, + fmin=self.fmin, + fmax=self.fmax) + mel_basis = torch.from_numpy(mel).float() + hann_window = torch.hann_window(self.win_size) + + self.register_buffer('mel_basis', mel_basis) + self.register_buffer('hann_window', hann_window) + + @property + def device(self): + return self.hann_window.device + + def forward(self, waveform: torch.Tensor) -> torch.Tensor: + # input: batch_size * length + bs = waveform.shape[0] + waveform = waveform.clamp(min=-1., max=1.) + + spec = torch.stft(waveform, + self.n_fft, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=True, + pad_mode='reflect', + normalized=False, + onesided=True, + return_complex=True) + + spec = torch.view_as_real(spec) + # print('After stft', spec.shape, spec.min(), spec.max(), spec.mean()) + + power = spec.pow(2).sum(-1) + angle = torch.atan2(spec[..., 1], spec[..., 0]) + + print('power', power.shape, power.min(), power.max(), power.mean()) + print('angle', angle.shape, angle.min(), angle.max(), angle.mean()) + + # print('mel', self.mel_basis.shape, self.mel_basis.min(), self.mel_basis.max(), + # self.mel_basis.mean()) + + # spec = rearrange(spec, 'b f t c -> (b c) f t') + + # spec = self.mel_transform(spec) + + # spec = torch.matmul(self.mel_basis, spec) + + # print('After mel', spec.shape, spec.min(), spec.max(), spec.mean()) + + # spec = spectral_normalize_torch(spec, self.norm_fn) + + # print('After norm', spec.shape, spec.min(), spec.max(), spec.mean()) + + # compute magnitude + # magnitude = torch.sqrt((spec**2).sum(-1)) + # normalize by magnitude + # scaled_magnitude = torch.log10(magnitude.clamp(min=1e-5)) * 10 + # spec = spec / magnitude.unsqueeze(-1) * scaled_magnitude.unsqueeze(-1) + + # power = torch.log10(power.clamp(min=1e-5)) * 10 + power = torch.log10(power.clamp(min=1e-5)) + + print('After scaling', power.shape, power.min(), power.max(), power.mean()) + + spec = torch.stack([power, angle], dim=-1) + + # spec = rearrange(spec, '(b c) f t -> b c f t', b=bs) + spec = rearrange(spec, 'b f t c -> b c f t', b=bs) + + # spec[:, :, 400:] = 0 + + return spec + + def invert(self, spec: torch.Tensor, length: int) -> torch.Tensor: + bs = spec.shape[0] + + # spec = rearrange(spec, 'b c f t -> (b c) f t') + # print(spec.shape, self.mel_basis.shape) + # spec = torch.linalg.lstsq(self.mel_basis.unsqueeze(0), spec).solution + # spec = torch.linalg.pinv(self.mel_basis.unsqueeze(0)) @ spec + + # spec = self.invmel_transform(spec) + + spec = rearrange(spec, 'b c f t -> b f t c', b=bs).contiguous() + + # spec[..., 0] = 10**(spec[..., 0] / 10) + + power = spec[..., 0] + power = 10**power + + # print('After unscaling', spec[..., 0].shape, spec[..., 0].min(), spec[..., 0].max(), + # spec[..., 0].mean()) + + unit_vector = torch.stack([ + torch.cos(spec[..., 1]), + torch.sin(spec[..., 1]), + ], dim=-1) + + spec = torch.sqrt(power) * unit_vector + + # spec = rearrange(spec, '(b c) f t -> b f t c', b=bs).contiguous() + spec = torch.view_as_complex(spec) + + waveform = torch.istft( + spec, + self.n_fft, + length=length, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=True, + normalized=False, + onesided=True, + return_complex=False, + ) + + return waveform + + +if __name__ == '__main__': + + converter = STFTConverter(sampling_rate=16000) + + signal = torchaudio.load('./output/ZZ6GRocWW38_000090.wav')[0] + # resample signal at 44100 Hz + # signal = torchaudio.transforms.Resample(16_000, 44_100)(signal) + + L = signal.shape[1] + print('Input signal', signal.shape) + spec = converter(signal) + + print('Final spec', spec.shape) + + signal_recon = converter.invert(spec, length=L) + print('Output signal', signal_recon.shape, signal_recon.min(), signal_recon.max(), + signal_recon.mean()) + + print('MSE', torch.nn.functional.mse_loss(signal, signal_recon)) + torchaudio.save('./output/ZZ6GRocWW38_000090_recon.wav', signal_recon, 16000) diff --git a/unison/models/mmaudio/ext/stft_converter_mel.py b/unison/models/mmaudio/ext/stft_converter_mel.py new file mode 100644 index 0000000000000000000000000000000000000000..f6b32d4cb9a23cd74f723e7d8307fd82fa1abba0 --- /dev/null +++ b/unison/models/mmaudio/ext/stft_converter_mel.py @@ -0,0 +1,234 @@ +# Reference: # https://github.com/bytedance/Make-An-Audio-2 + +import torch +import torch.nn as nn +import torchaudio +from einops import rearrange +from librosa.filters import mel as librosa_mel_fn + + +def dynamic_range_compression_torch(x, C=1, clip_val=1e-5, norm_fn=torch.log10): + return norm_fn(torch.clamp(x, min=clip_val) * C) + + +def spectral_normalize_torch(magnitudes, norm_fn): + output = dynamic_range_compression_torch(magnitudes, norm_fn=norm_fn) + return output + + +class STFTConverter(nn.Module): + + def __init__( + self, + *, + sampling_rate: float = 16_000, + n_fft: int = 1024, + num_mels: int = 128, + hop_size: int = 256, + win_size: int = 1024, + fmin: float = 0, + fmax: float = 8_000, + norm_fn=torch.log, + ): + super().__init__() + self.sampling_rate = sampling_rate + self.n_fft = n_fft + self.num_mels = num_mels + self.hop_size = hop_size + self.win_size = win_size + self.fmin = fmin + self.fmax = fmax + self.norm_fn = norm_fn + + mel = librosa_mel_fn(sr=self.sampling_rate, + n_fft=self.n_fft, + n_mels=self.num_mels, + fmin=self.fmin, + fmax=self.fmax) + mel_basis = torch.from_numpy(mel).float() + hann_window = torch.hann_window(self.win_size) + + self.register_buffer('mel_basis', mel_basis) + self.register_buffer('hann_window', hann_window) + + @property + def device(self): + return self.hann_window.device + + def forward(self, waveform: torch.Tensor) -> torch.Tensor: + # input: batch_size * length + bs = waveform.shape[0] + waveform = waveform.clamp(min=-1., max=1.) + + spec = torch.stft(waveform, + self.n_fft, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=True, + pad_mode='reflect', + normalized=False, + onesided=True, + return_complex=True) + + spec = torch.view_as_real(spec) + # print('After stft', spec.shape, spec.min(), spec.max(), spec.mean()) + + power = (spec.pow(2).sum(-1))**(0.5) + angle = torch.atan2(spec[..., 1], spec[..., 0]) + + print('power 1', power.shape, power.min(), power.max(), power.mean()) + print('angle 1', angle.shape, angle.min(), angle.max(), angle.mean(), angle[:, :2, :2]) + + # print('mel', self.mel_basis.shape, self.mel_basis.min(), self.mel_basis.max(), + # self.mel_basis.mean()) + + # spec = self.mel_transform(spec) + + # power = torch.matmul(self.mel_basis, power) + + spec = rearrange(spec, 'b f t c -> (b c) f t') + spec = self.mel_basis.unsqueeze(0) @ spec + spec = rearrange(spec, '(b c) f t -> b f t c', b=bs) + + power = (spec.pow(2).sum(-1))**(0.5) + angle = torch.atan2(spec[..., 1], spec[..., 0]) + + print('power', power.shape, power.min(), power.max(), power.mean()) + print('angle', angle.shape, angle.min(), angle.max(), angle.mean(), angle[:, :2, :2]) + + # print('After mel', spec.shape, spec.min(), spec.max(), spec.mean()) + + # spec = spectral_normalize_torch(spec, self.norm_fn) + + # print('After norm', spec.shape, spec.min(), spec.max(), spec.mean()) + + # compute magnitude + # magnitude = torch.sqrt((spec**2).sum(-1)) + # normalize by magnitude + # scaled_magnitude = torch.log10(magnitude.clamp(min=1e-5)) * 10 + # spec = spec / magnitude.unsqueeze(-1) * scaled_magnitude.unsqueeze(-1) + + # power = torch.log10(power.clamp(min=1e-5)) * 10 + power = torch.log10(power.clamp(min=1e-8)) + + print('After scaling', power.shape, power.min(), power.max(), power.mean()) + + # spec = torch.stack([power, angle], dim=-1) + + # spec = rearrange(spec, '(b c) f t -> b c f t', b=bs) + # spec = rearrange(spec, 'b f t c -> b c f t', b=bs) + + # spec[:, :, 400:] = 0 + + return power, angle + # return spec[..., 0], spec[..., 1] + + def invert(self, spec: torch.Tensor, length: int) -> torch.Tensor: + + power, angle = spec + + bs = power.shape[0] + + # spec = rearrange(spec, 'b c f t -> (b c) f t') + # print(spec.shape, self.mel_basis.shape) + # spec = torch.linalg.lstsq(self.mel_basis.unsqueeze(0), spec).solution + # spec = torch.linalg.pinv(self.mel_basis.unsqueeze(0)) @ spec + + # spec = self.invmel_transform(spec) + + # spec = rearrange(spec, 'b c f t -> b f t c', b=bs).contiguous() + + # spec[..., 0] = 10**(spec[..., 0] / 10) + + # power = spec[..., 0] + power = 10**power + + # print('After unscaling', spec[..., 0].shape, spec[..., 0].min(), spec[..., 0].max(), + # spec[..., 0].mean()) + + unit_vector = torch.stack([ + torch.cos(angle), + torch.sin(angle), + ], dim=-1) + + spec = power.unsqueeze(-1) * unit_vector + + # power = torch.linalg.lstsq(self.mel_basis.unsqueeze(0), power).solution + spec = rearrange(spec, 'b f t c -> (b c) f t') + spec = torch.linalg.pinv(self.mel_basis.unsqueeze(0)) @ spec + # spec = torch.linalg.lstsq(self.mel_basis.unsqueeze(0), spec).solution + spec = rearrange(spec, '(b c) f t -> b f t c', b=bs).contiguous() + + power = (spec.pow(2).sum(-1))**(0.5) + angle = torch.atan2(spec[..., 1], spec[..., 0]) + + print('power 2', power.shape, power.min(), power.max(), power.mean()) + print('angle 2', angle.shape, angle.min(), angle.max(), angle.mean(), angle[:, :2, :2]) + + # spec = rearrange(spec, '(b c) f t -> b f t c', b=bs).contiguous() + spec = torch.view_as_complex(spec) + + waveform = torch.istft( + spec, + self.n_fft, + length=length, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=True, + normalized=False, + onesided=True, + return_complex=False, + ) + + return waveform + + +if __name__ == '__main__': + + converter = STFTConverter(sampling_rate=16000) + + signal = torchaudio.load('./output/ZZ6GRocWW38_000090.wav')[0] + # resample signal at 44100 Hz + # signal = torchaudio.transforms.Resample(16_000, 44_100)(signal) + + L = signal.shape[1] + print('Input signal', signal.shape) + spec = converter(signal) + + power, angle = spec + + # print(power.shape, angle.shape) + # print(power, power.min(), power.max(), power.mean()) + # power = power.clamp(-1, 1) + # angle = angle.clamp(-1, 1) + + import matplotlib.pyplot as plt + + # Visualize power + plt.figure() + plt.imshow(power[0].detach().numpy(), aspect='auto', origin='lower') + plt.colorbar() + plt.title('Power') + plt.xlabel('Time') + plt.ylabel('Frequency') + plt.savefig('./output/power.png') + + # Visualize angle + plt.figure() + plt.imshow(angle[0].detach().numpy(), aspect='auto', origin='lower') + plt.colorbar() + plt.title('Angle') + plt.xlabel('Time') + plt.ylabel('Frequency') + plt.savefig('./output/angle.png') + + # print('Final spec', spec.shape) + + signal_recon = converter.invert(spec, length=L) + print('Output signal', signal_recon.shape, signal_recon.min(), signal_recon.max(), + signal_recon.mean()) + + print('MSE', torch.nn.functional.mse_loss(signal, signal_recon)) + torchaudio.save('./output/ZZ6GRocWW38_000090_recon.wav', signal_recon, 16000) diff --git a/unison/models/mmaudio/features_utils.py b/unison/models/mmaudio/features_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7fea9eec4907927037efc4a784d7a01c8197967a --- /dev/null +++ b/unison/models/mmaudio/features_utils.py @@ -0,0 +1,77 @@ +from typing import Literal, Optional + +import torch +import torch.nn as nn + +from .ext.autoencoder import AutoEncoderModule +from .ext.autoencoder.distributions import DiagonalGaussianDistribution +from .ext.mel_converter import get_mel_converter + + +class FeaturesUtils(nn.Module): + + def __init__( + self, + *, + tod_vae_ckpt: str, + bigvgan_vocoder_ckpt: Optional[str] = None, + mode=Literal['16k', '44k'], + need_vae_encoder: bool = True, + ): + super().__init__() + + self.mel_converter = get_mel_converter(mode) + self.tod = AutoEncoderModule(vae_ckpt_path=tod_vae_ckpt, + vocoder_ckpt_path=bigvgan_vocoder_ckpt, + mode=mode, + need_vae_encoder=need_vae_encoder) + + def compile(self): + + self.decode = torch.compile(self.decode) + self.vocode = torch.compile(self.vocode) + + def train(self, mode: bool) -> None: + return super().train(False) + + @torch.inference_mode() + def encode_audio(self, x) -> DiagonalGaussianDistribution: + assert self.tod is not None, 'VAE is not loaded' + # x: (B * L) + mel = self.mel_converter(x) + dist = self.tod.encode(mel) + + return dist + + @torch.inference_mode() + def vocode(self, mel: torch.Tensor) -> torch.Tensor: + assert self.tod is not None, 'VAE is not loaded' + return self.tod.vocode(mel) + + @torch.inference_mode() + def decode(self, z: torch.Tensor) -> torch.Tensor: + assert self.tod is not None, 'VAE is not loaded' + return self.tod.decode(z) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + @torch.no_grad() + def wrapped_decode(self, z): + with torch.amp.autocast('cuda', dtype=self.dtype): + mel_decoded = self.decode(z) + audio = self.vocode(mel_decoded) + + return audio + + @torch.no_grad() + def wrapped_encode(self, audio): + with torch.amp.autocast('cuda', dtype=self.dtype): + dist = self.encode_audio(audio) + + return dist.mean \ No newline at end of file diff --git a/unison/models/text_encoders/omni_encoder.py b/unison/models/text_encoders/omni_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..b86ab11ce64a0a6fd9d97e3d3a39be807bcfca54 --- /dev/null +++ b/unison/models/text_encoders/omni_encoder.py @@ -0,0 +1,232 @@ +""" +Qwen2.5-Omni text encoder with layer-wise hidden state extraction for deep LLM fusion. + +Extracts hidden states from uniformly-sampled layers of a frozen Qwen2.5-Omni-7B +text backbone. Each extracted state is injected into the corresponding MM-DiT double- +stream block via a learned linear projection (deep fusion). + +Key design: + - Only the Thinker's language model is retained; vision, audio tower, and Talker + sub-modules are dropped immediately after loading to save GPU memory. + - A domain-specific system prompt is prepended; its token span is sliced off from + every extracted hidden state so the DiT only receives user-prompt tokens. + - padding_side='right' ensures the system prompt is always at the front (index 0), + making the slice offset stable across batches. +""" + +import gc +from typing import Optional + +import numpy as np +import torch +import torch.nn as nn +from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor + + +OMNI_PRESET_MODEL_IDS = { + "omni-3b": "Qwen/Qwen2.5-Omni-3B", + "omni-7b": "Qwen/Qwen2.5-Omni-7B", +} + + +def resolve_omni_model_path(preset: str, model_path: Optional[str]) -> str: + """Resolve Qwen2.5-Omni checkpoint path. + + - preset == 'omni' -> use `model_path` directly (user-supplied path). + - preset == 'omni-3b/7b' -> use `model_path` if given, else fall back to HF hub id. + """ + if preset == "omni": + if not model_path: + raise ValueError( + "text_encoder_type='omni' requires --omni_model_path to be set." + ) + return model_path + if preset not in OMNI_PRESET_MODEL_IDS: + raise ValueError( + f"Unknown Omni preset {preset!r}. " + f"Choose one of {list(OMNI_PRESET_MODEL_IDS) + ['omni']}." + ) + return model_path or OMNI_PRESET_MODEL_IDS[preset] + + +class QwenOmniThinkerExtractor(nn.Module): + """Frozen Qwen2.5-Omni Thinker language model used as a text-feature extractor. + + Returns per-layer hidden states for deep fusion into the DiT backbone. + """ + + def __init__( + self, + model_path: str = "Qwen/Qwen2.5-Omni-7B", + dit_depth: int = 20, + select_mode: str = "interval", + device: str = "cuda", + dtype: torch.dtype = torch.bfloat16, + omni_last_layer_idx: Optional[int] = None, + ): + """ + Args: + model_path: Path to Qwen2.5-Omni model (local or HF hub). + dit_depth: Number of MM-DiT double-stream blocks; determines + how many LLM layers to sample. + select_mode: 'interval' = uniformly-spaced layers; + 'last' = last dit_depth layers. + device: Device to place the text backbone on. + dtype: Computation dtype (bfloat16 recommended). + omni_last_layer_idx: If set, also return a single "global" embedding from + this LLM layer index (0-indexed from end when negative). + None = only return per-block hidden states. + """ + super().__init__() + self.device = device + self.dtype = dtype + + print(f"Loading Qwen2.5-Omni Thinker from {model_path}...") + + # Load full model then discard everything except the language backbone. + full_model = Qwen2_5OmniForConditionalGeneration.from_pretrained( + model_path, + torch_dtype=dtype, + device_map="cpu", + attn_implementation="flash_attention_2", + ) + self.text_backbone = full_model.thinker.model.eval() + + for param in self.text_backbone.parameters(): + param.requires_grad = False + + print("Dropping unused sub-modules (talker, token2wav, visual, audio_tower)...") + del full_model.talker + del full_model.token2wav + del full_model.thinker.visual + del full_model.thinker.audio_tower + del full_model + + gc.collect() + + print(f"Moving text backbone to {device}...") + self.text_backbone = self.text_backbone.to(device) + torch.cuda.empty_cache() + + # Load processor; force right-side padding so the system prompt is always + # at index 0, making the slice offset constant. + self.processor = Qwen2_5OmniProcessor.from_pretrained(model_path) + self.processor.tokenizer.padding_side = "right" + + # Domain-specific system prompt for audio generation conditioning. + self.system_prompt_content = [{ + "type": "text", + "text": ( + "You are an expert Audio Director and Sound Engineer. " + "Your task is to mentally simulate and analyze the acoustic scene described in the text. " + "Pay attention to the following three layers: " + "1. Voice & Prosody: If there is speech, identify the speaker's gender, age, emotion, accent, and exact spoken content. " + "2. Environmental Sounds: Identify background ambience, distinct sound effects, and their textures. " + "3. Spatial Mix: Determine how the voice and background sounds overlap and interact in the scene. " + "Extract deep semantic representations to guide the generation of high-fidelity speech, audio, or a mixture of both." + ) + }] + + # Pre-compute system-prompt token length so we can slice it off during forward. + print("Pre-computing system prompt token length...") + sys_str = self.processor.apply_chat_template( + [{"role": "system", "content": self.system_prompt_content}], + add_generation_prompt=False, + tokenize=False, + ) + sys_tokens = self.processor( + sys_str, + return_tensors="pt", + padding=False, + truncation=False, + add_special_tokens=False, + ) + self.system_token_len = sys_tokens.input_ids.shape[1] + print(f"System prompt length: {self.system_token_len} tokens (will be sliced off in forward)") + + # Map LLM layer indices to DiT block indices. + self.total_llm_layers = len(self.text_backbone.layers) + self.layer_indices = self._get_layer_indices(self.total_llm_layers, dit_depth, select_mode) + print(f"Selected LLM layer indices: {self.layer_indices}") + + # omni_last_layer_idx: index into hidden_states tuple for an optional global + # embedding (e.g. -1 = last layer, -2 = second-to-last). + # When None, forward() returns (extracted_states, mask) only. + self.omni_last_layer_idx = omni_last_layer_idx + if omni_last_layer_idx is not None: + resolved = self.total_llm_layers + 1 + omni_last_layer_idx + print(f"Global embedding from LLM layer index {omni_last_layer_idx} (= layer {resolved} of {self.total_llm_layers})") + + def _get_layer_indices(self, total_llm: int, dit_depth: int, mode: str): + """Return list of LLM layer indices to sample for each DiT block.""" + available = list(range(1, total_llm + 1)) + if total_llm >= dit_depth: + if mode == "interval": + idxs = np.linspace(0, len(available) - 1, dit_depth, dtype=int) + return [available[i] for i in idxs] + else: + return available[-dit_depth:] + else: + # Repeat last layer to pad up to dit_depth. + return available + [available[-1]] * (dit_depth - total_llm) + + @torch.no_grad() + def forward(self, text_list: list): + """Extract per-block hidden states for a batch of text prompts. + + Args: + text_list: List of prompt strings. + + Returns: + extracted_states: List[Tensor], length == dit_depth. + Each tensor is [B, user_seq_len, hidden_dim]. + user_attention_mask: [B, user_seq_len] — system-prompt tokens removed. + omni_last_emb (only if omni_last_layer_idx is set): + [B, user_seq_len, hidden_dim] from a single chosen layer. + """ + batch_conversations = [ + [ + {"role": "system", "content": self.system_prompt_content}, + {"role": "user", "content": [{"type": "text", "text": txt}]}, + ] + for txt in text_list + ] + + text_prompts = self.processor.apply_chat_template( + batch_conversations, + add_generation_prompt=True, + tokenize=False, + ) + + # Tokenize with right-padding: layout is [System][User][Pad…] + inputs = self.processor( + text_prompts, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=256, + add_special_tokens=False, + ).to(self.device) + + outputs = self.text_backbone( + input_ids=inputs.input_ids, + attention_mask=inputs.attention_mask, + output_hidden_states=True, + return_dict=True, + use_cache=False, + ) + + all_states = outputs.hidden_states + + # Slice off system-prompt tokens from each selected layer. + extracted_states = [ + all_states[idx].to(self.dtype)[:, self.system_token_len:, :] + for idx in self.layer_indices + ] + user_attention_mask = inputs.attention_mask[:, self.system_token_len:] + + if self.omni_last_layer_idx is not None: + omni_last_emb = all_states[self.omni_last_layer_idx].to(self.dtype)[:, self.system_token_len:, :] + return extracted_states, user_attention_mask, omni_last_emb + + return extracted_states, user_attention_mask diff --git a/unison/models/transformers/backbone.py b/unison/models/transformers/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..b7b38efa956a81c6a163504ba095b7e83fc404ec --- /dev/null +++ b/unison/models/transformers/backbone.py @@ -0,0 +1,764 @@ +# UnisonBackbone — MM-DiT backbone for UNISON. +# +# Task identity and source audio are encoded entirely through the input tensor: +# model input shape: [B, 2*C+1, T] +# [:, :C, :] — noisy target latent +# [:, C:2C, :] — source/reference latent (zeros for pure generation) +# [:, 2C, :] — task mask (0=generate, 1=edit, 2=zero-shot-TTS reference region) +# +# The UniversalPatchEmbed layer (concat_condition=True) handles the 2C+1 → hidden_size +# projection; C is set by in_channels in the model config and matches the VAE latent dim. +# +# Deep LLM fusion: each MM-DiT block receives the hidden state from the corresponding +# layer of a frozen MLLM (Qwen2.5-Omni or Qwen3) via a learned linear projection, +# providing depth-matched semantic conditioning. + +from typing import Any, List, Tuple, Optional, Union, Dict + +import torch +import torch.nn as nn +from einops import rearrange +from loguru import logger + +from diffusers.models import ModelMixin +from diffusers.configuration_utils import ConfigMixin, register_to_config + +from .modules.activation_layers import get_activation_layer +from .modules.norm_layers import get_norm_layer +from .modules.embed_layers import ( + TimestepEmbedder, + PatchEmbed, + VisionProjection, + UniversalPatchEmbed, + DurationEmbedder +) +from .modules.attention import parallel_attention +from .modules.posemb_layers import apply_rotary_emb, get_nd_rotary_pos_embed, get_audio_rotary_pos_embed +from .modules.mlp_layers import MLP, MLPEmbedder, FinalLayer, LinearWarpforSingle +from .modules.modulate_layers import ModulateDiT, modulate, apply_gate +from .modules.token_refiner import SingleTokenRefiner + +from unison.utils.communications import all_gather +from unison.utils.infer_utils import torch_compile_wrapper + +from unison.commons.parallel_states import get_parallel_state + +class MMDoubleStreamBlock(nn.Module): + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float, + mlp_act_type: str = "gelu_tanh", + attn_mode: str = None, + qk_norm: bool = True, + qk_norm_type: str = "rms", + qkv_bias: bool = False, + has_text_ffn: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.heads_num = heads_num + self.attn_mode = attn_mode + self.has_text_ffn = has_text_ffn + + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.img_mod = ModulateDiT( + hidden_size, factor=6, act_layer=get_activation_layer("silu"), **factory_kwargs + ) + self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.img_attn_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + self.img_attn_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + self.img_attn_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + qk_norm_layer = get_norm_layer(qk_norm_type) + self.img_attn_q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.img_attn_k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.img_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.img_mlp = MLP(hidden_size, mlp_hidden_dim, act_layer=get_activation_layer(mlp_act_type), bias=True, **factory_kwargs) + + txt_mod_factor = 6 if has_text_ffn else 2 + self.txt_mod = ModulateDiT( + hidden_size, factor=txt_mod_factor, act_layer=get_activation_layer("silu"), **factory_kwargs + ) + self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.txt_attn_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + self.txt_attn_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + self.txt_attn_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.txt_attn_q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.txt_attn_k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + + if has_text_ffn: + self.txt_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, act_layer=get_activation_layer(mlp_act_type), bias=True, **factory_kwargs) + + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + @torch_compile_wrapper() + def forward( + self, + img: torch.Tensor, + txt: torch.Tensor, + vec: torch.Tensor, + freqs_cis: tuple = None, + text_mask=None, + attn_param=None, + is_flash=False, + block_idx=None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + ( + img_mod1_shift, + img_mod1_scale, + img_mod1_gate, + img_mod2_shift, + img_mod2_scale, + img_mod2_gate, + ) = self.img_mod(vec).chunk(6, dim=-1) + + if self.has_text_ffn: + ( + txt_mod1_shift, + txt_mod1_scale, + txt_mod1_gate, + txt_mod2_shift, + txt_mod2_scale, + txt_mod2_gate, + ) = self.txt_mod(vec).chunk(6, dim=-1) + else: + txt_mod1_shift, txt_mod1_scale = self.txt_mod(vec).chunk(2, dim=-1) + + img_modulated = self.img_norm1(img) + img_modulated = modulate(img_modulated, shift=img_mod1_shift, scale=img_mod1_scale) + + img_q = self.img_attn_q(img_modulated) + img_k = self.img_attn_k(img_modulated) + img_v = self.img_attn_v(img_modulated) + img_q = rearrange(img_q, "B L (H D) -> B L H D", H=self.heads_num) + img_k = rearrange(img_k, "B L (H D) -> B L H D", H=self.heads_num) + img_v = rearrange(img_v, "B L (H D) -> B L H D", H=self.heads_num) + img_q = self.img_attn_q_norm(img_q).to(img_v) + img_k = self.img_attn_k_norm(img_k).to(img_v) + + if freqs_cis is not None: + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) + assert ( + img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + txt_modulated = self.txt_norm1(txt) + txt_modulated = modulate(txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale) + txt_q = self.txt_attn_q(txt_modulated) + txt_k = self.txt_attn_k(txt_modulated) + txt_v = self.txt_attn_v(txt_modulated) + txt_q = rearrange(txt_q, "B L (H D) -> B L H D", H=self.heads_num) + txt_k = rearrange(txt_k, "B L (H D) -> B L H D", H=self.heads_num) + txt_v = rearrange(txt_v, "B L (H D) -> B L H D", H=self.heads_num) + txt_q = self.txt_attn_q_norm(txt_q).to(txt_v) + txt_k = self.txt_attn_k_norm(txt_k).to(txt_v) + + attn_mode = 'flash' if is_flash else self.attn_mode + attn = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + attn_mode=attn_mode, + attn_param=attn_param, + block_idx=block_idx, + ) + + img_attn, txt_attn = attn[:, :img_q.shape[1]].contiguous(), attn[:, img_q.shape[1]:].contiguous() + + img = img + apply_gate(self.img_attn_proj(img_attn), gate=img_mod1_gate) + img = img + apply_gate( + self.img_mlp( + modulate(self.img_norm2(img), shift=img_mod2_shift, scale=img_mod2_scale) + ), + gate=img_mod2_gate, + ) + + if self.has_text_ffn: + txt = txt + apply_gate(self.txt_attn_proj(txt_attn), gate=txt_mod1_gate) + txt = txt + apply_gate( + self.txt_mlp(modulate(self.txt_norm2(txt), shift=txt_mod2_shift, scale=txt_mod2_scale)), + gate=txt_mod2_gate, + ) + + return img, txt + + +class MMSingleStreamBlock(nn.Module): + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = "gelu_tanh", + attn_mode: str = None, + qk_norm: bool = True, + qk_norm_type: str = "rms", + qk_scale: float = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.attn_mode = attn_mode + + self.hidden_size = hidden_size + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + self.mlp_hidden_dim = mlp_hidden_dim + self.scale = qk_scale or head_dim ** -0.5 + + self.linear1_q = nn.Linear(hidden_size, hidden_size, **factory_kwargs) + self.linear1_k = nn.Linear(hidden_size, hidden_size, **factory_kwargs) + self.linear1_v = nn.Linear(hidden_size, hidden_size, **factory_kwargs) + self.linear1_mlp = nn.Linear(hidden_size, mlp_hidden_dim, **factory_kwargs) + self.linear2 = LinearWarpforSingle(hidden_size + mlp_hidden_dim, hidden_size, bias=True, **factory_kwargs) + self.mlp_act = get_activation_layer(mlp_act_type)() + + qk_norm_layer = get_norm_layer(qk_norm_type) + self.q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.modulation = ModulateDiT(hidden_size, factor=3, act_layer=get_activation_layer("silu"), **factory_kwargs) + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward( + self, + x: torch.Tensor, + vec: torch.Tensor, + txt_len: int, + freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None, + text_mask=None, + attn_param=None, + is_flash=False, + ) -> torch.Tensor: + mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1) + x_mod = modulate(self.pre_norm(x), shift=mod_shift, scale=mod_scale) + + q = self.linear1_q(x_mod) + k = self.linear1_k(x_mod) + v = self.linear1_v(x_mod) + + q = rearrange(q, "B L (H D) -> B L H D", H=self.heads_num) + k = rearrange(k, "B L (H D) -> B L H D", H=self.heads_num) + v = rearrange(v, "B L (H D) -> B L H D", H=self.heads_num) + + mlp = self.linear1_mlp(x_mod) + + q = self.q_norm(q).to(v) + k = self.k_norm(k).to(v) + + img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :] + img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :] + img_v, txt_v = v[:, :-txt_len, :, :], v[:, -txt_len:, :, :] + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) + assert ( + img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + if is_flash: + attn_mode = 'flash' + else: + attn_mode = self.attn_mode + attn = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + attn_mode=attn_mode, + attn_param=attn_param, + ) + output = self.linear2(attn, self.mlp_act(mlp)) + + return x + apply_gate(output, gate=mod_gate) + + +class UnisonBackbone(ModelMixin, ConfigMixin): + """Channel-cat backbone: no ref_latents, no ref-isolated AdaLN, continuous RoPE.""" + + @register_to_config + def __init__( + self, + in_channels: int = 48, + out_channels: int = None, + patch_size: list = [1, 2, 2], + hidden_size: int = 1024, + heads_num: int = 8, + concat_condition: bool = False, + mlp_act_type: str = "gelu_tanh", + mlp_width_ratio: float = 4.0, + mm_double_blocks_depth: int = 18, + mm_single_blocks_depth: int = 0, + qkv_bias: bool = True, + qk_norm: bool = True, + qk_norm_type: str = "rms", + attn_mode: str = "flash", + attn_param: dict = None, + rope_dim_list: list = [16, 56, 56], + rope_theta: int = 10000, + omni_dim: int = 1280, + is_reshape_temporal_channels: bool = False, + use_duration_embedding: bool = False, + is_audio_type: bool = False, + use_omni_embedding: bool = False, + use_omni_last_embedding: bool = False, + projector_type: str = "linear", + guidance_embed: bool = False, + temporal_rope_scaling_factor: float = 1.0, + repa_z_dim: int = None, + repa_layer_num: int = None, + audio_patch_type: str = "conv_mlp", + ): + super().__init__() + + self.hidden_size = hidden_size + self.heads_num = heads_num + self.is_audio_type = is_audio_type + self.concat_condition = concat_condition + self.out_channels = out_channels or in_channels + self.patch_size = patch_size + self.rope_dim_list = rope_dim_list + self.rope_theta = rope_theta + self.temporal_rope_scaling_factor = temporal_rope_scaling_factor + self.mm_double_blocks_depth = mm_double_blocks_depth + self.mm_single_blocks_depth = mm_single_blocks_depth + self.total_blocks_depth = mm_double_blocks_depth + mm_single_blocks_depth + self.guidance_embed = guidance_embed + self.use_duration_embedding = use_duration_embedding + self.use_omni_embedding = use_omni_embedding + self.projector_type = projector_type + self.repa_z_dim = repa_z_dim + self.repa_layer_num = repa_layer_num + self.use_omni_last_embedding = use_omni_last_embedding + self.gradient_checkpointing = False + + self.img_in = UniversalPatchEmbed( + patch_size=patch_size, + in_chans=in_channels, + embed_dim=hidden_size, + is_audio=is_audio_type, + is_reshape_temporal_channels=is_reshape_temporal_channels, + audio_kernel_size=7, + audio_padding=3, + concat_condition=concat_condition, + audio_patch_type=audio_patch_type, + ) + + self.time_in = TimestepEmbedder(hidden_size, get_activation_layer("silu")) + self.guidance_in = TimestepEmbedder(hidden_size, get_activation_layer("silu")) if guidance_embed else None + + if use_duration_embedding: + self.duration_embedder = DurationEmbedder(hidden_size, min_value=0, max_value=30) + else: + self.duration_embedder = None + + if use_omni_embedding: + if self.projector_type == "linear": + self.deep_fusion_projs = nn.ModuleList([ + nn.Linear(omni_dim, hidden_size) for _ in range(self.total_blocks_depth) + ]) + elif self.projector_type == "mlp": + self.deep_fusion_projs = nn.ModuleList([ + nn.Sequential( + nn.Linear(omni_dim, hidden_size), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size) + ) for _ in range(self.total_blocks_depth) + ]) + else: + raise ValueError(f"Invalid projector type: {self.projector_type}") + else: + self.deep_fusion_projs = None + + if self.use_omni_last_embedding: + if self.projector_type == "linear": + self.omni_last_proj = nn.Linear(omni_dim, hidden_size) + elif self.projector_type == "mlp": + self.omni_last_proj = nn.Sequential( + nn.Linear(omni_dim, hidden_size), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size) + ) + else: + raise ValueError(f"Invalid projector type: {self.projector_type}") + else: + self.omni_last_proj = None + + need_text_ffn = self.use_omni_last_embedding or (self.duration_embedder is not None) + + self.double_blocks = nn.ModuleList([ + MMDoubleStreamBlock( + hidden_size=hidden_size, + heads_num=heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + attn_mode="flash", + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + has_text_ffn=need_text_ffn, + ) for _ in range(mm_double_blocks_depth) + ]) + self.single_blocks = nn.ModuleList([ + MMSingleStreamBlock( + hidden_size=hidden_size, + heads_num=heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + attn_mode="flash", + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + ) for _ in range(mm_single_blocks_depth) + ]) + + if self.is_audio_type: + final_patch_size = (1, 1, 1) + else: + final_patch_size = patch_size + + self.final_layer = FinalLayer(hidden_size, final_patch_size, self.out_channels, get_activation_layer("silu")) + + if self.repa_z_dim is not None and self.repa_layer_num is not None: + self.repa_proj = nn.Linear(hidden_size, repa_z_dim) + else: + self.repa_proj = None + + def enable_gradient_checkpointing(self): + self.gradient_checkpointing = True + + def disable_gradient_checkpointing(self): + self.gradient_checkpointing = False + + def get_rotary_pos_embed(self, grid_sizes): + if self.is_audio_type: + return get_audio_rotary_pos_embed( + rope_dim_list=self.rope_dim_list, + length=grid_sizes[0], + theta=self.rope_theta, + freqs_scaling=self.temporal_rope_scaling_factor, + use_real=True + ) + else: + return get_nd_rotary_pos_embed( + self.rope_dim_list, + grid_sizes, + theta=self.rope_theta, + use_real=True, + theta_rescale_factor=1 + ) + + def count_parameters(self, verbose=True): + total_params = 0 + trainable_params = 0 + + module_params = { + "Patch Embed (img_in)": 0, + "Timestep Embed (time_in)": 0, + "Guidance Embed": 0, + "Duration Embed": 0, + "Omni Last Projection": 0, + "Deep Fusion Projections": 0, + "Double Blocks": 0, + "Single Blocks": 0, + "Final Layer": 0, + "RepA Projection": 0, + "Others": 0, + } + + for name, param in self.named_parameters(): + num_params = param.numel() + total_params += num_params + if param.requires_grad: + trainable_params += num_params + + if "img_in" in name: + module_params["Patch Embed (img_in)"] += num_params + elif "time_in" in name: + module_params["Timestep Embed (time_in)"] += num_params + elif "guidance_in" in name: + module_params["Guidance Embed"] += num_params + elif "duration_embedder" in name: + module_params["Duration Embed"] += num_params + elif "omni_last_proj" in name: + module_params["Omni Last Projection"] += num_params + elif "deep_fusion_projs" in name: + module_params["Deep Fusion Projections"] += num_params + elif "double_blocks" in name: + module_params["Double Blocks"] += num_params + elif "single_blocks" in name: + module_params["Single Blocks"] += num_params + elif "final_layer" in name: + module_params["Final Layer"] += num_params + elif "repa_proj" in name: + module_params["RepA Projection"] += num_params + else: + module_params["Others"] += num_params + + if verbose: + print(f"\n{'='*50}") + print(f" Model Parameter Statistics") + print(f"{'='*50}") + for category, count in module_params.items(): + if count > 0: + print(f" {category:<30}: {count / 1e6:>8.2f} M") + print(f" {'-'*48}") + print(f" {'Total Parameters':<30}: {total_params / 1e6:>8.2f} M ({total_params / 1e9:.3f} B)") + print(f" {'Trainable Parameters':<30}: {trainable_params / 1e6:>8.2f} M ({trainable_params / 1e9:.3f} B)") + print(f"{'='*50}\n") + + return total_params + + def prepare_packed_indices(self, mask_list): + total_mask = torch.cat(mask_list, dim=1).bool() + sorted_mask, sort_indices = torch.sort( + total_mask.int(), + dim=1, + descending=True, + stable=True + ) + gather_indices = sort_indices.unsqueeze(-1).expand(-1, -1, self.hidden_size) + return sorted_mask.bool(), gather_indices + + def apply_packing(self, embedding_list, gather_indices, sorted_mask): + raw_concat = torch.cat(embedding_list, dim=1) + packed_embedding = torch.gather(raw_concat, dim=1, index=gather_indices) + mask_broadcaster = sorted_mask.unsqueeze(-1).to(dtype=packed_embedding.dtype) + packed_embedding = packed_embedding * mask_broadcaster + return packed_embedding + + def forward( + self, + x: torch.Tensor, + t: torch.Tensor, + duration: torch.Tensor = None, + omni_emb_list: List[torch.Tensor] = None, + omni_last_emb: torch.Tensor = None, + omni_mask: torch.Tensor = None, + guidance: torch.Tensor = None, + ): + assert self.use_omni_embedding or self.use_omni_last_embedding, \ + "At least one of use_omni_embedding / use_omni_last_embedding must be True" + assert omni_mask is not None, "omni_mask is required" + if self.use_omni_embedding: + assert omni_emb_list is not None and len(omni_emb_list) >= self.total_blocks_depth, \ + f"omni_emb_list required. Expected len >= {self.total_blocks_depth}, got {len(omni_emb_list) if omni_emb_list else 0}" + if self.use_omni_last_embedding: + if omni_last_emb is None: + assert omni_emb_list is not None and len(omni_emb_list) > 0 + omni_last_emb = omni_emb_list[-1] + if omni_mask.dim() == 3: + omni_mask = omni_mask.squeeze(1) + + # --- 1. Patching & Embeddings --- + img = self.img_in(x) + vec = self.time_in(t) + repa_feat = None + if self.guidance_embed and guidance is not None: + vec = vec + self.guidance_in(guidance) + + # --- 2. RoPE (always continuous, no offset) --- + grid_sizes_target = None + if self.is_audio_type: + freqs_cos, freqs_sin = get_audio_rotary_pos_embed( + rope_dim_list=self.rope_dim_list, + length=img.shape[1], + theta=self.rope_theta, + freqs_scaling=self.temporal_rope_scaling_factor, + use_real=True, + ) + freqs_cis = (freqs_cos.to(img.device), freqs_sin.to(img.device)) + else: + _, _, ot, oh, ow = x.shape + pt, ph, pw = self.patch_size + grid_sizes_target = (ot // pt, oh // ph, ow // pw) + freqs_cos, freqs_sin = self.get_rotary_pos_embed(grid_sizes_target) + freqs_cis = (freqs_cos.to(img.device), freqs_sin.to(img.device)) + + # --- 3. Build base text tokens --- + txt_base = None + txt_mask = None + + if self.use_omni_last_embedding: + txt_base = self.omni_last_proj(omni_last_emb) + txt_mask = omni_mask + + if self.duration_embedder is not None and duration is not None: + dur_emb = self.duration_embedder(duration) + dur_mask = torch.ones(omni_mask.shape[0], 1, device=omni_mask.device, dtype=omni_mask.dtype) + if txt_base is not None: + txt_base = torch.cat([dur_emb, txt_base], dim=1) + txt_mask = torch.cat([dur_mask, txt_mask], dim=1) + else: + txt_base = dur_emb + txt_mask = dur_mask + + has_base = txt_base is not None + txt_base_len = txt_base.shape[1] if has_base else 0 + + if self.use_omni_embedding and has_base: + txt_mask_sorted, gather_indices = self.prepare_packed_indices([txt_mask, omni_mask]) + elif self.use_omni_embedding and not has_base: + txt_mask_sorted = omni_mask + gather_indices = None + else: + txt_mask_sorted = txt_mask + gather_indices = None + + # --- 4. Double Stream Loop --- + for i, block in enumerate(self.double_blocks): + if self.use_omni_embedding: + omni_curr = self.deep_fusion_projs[i](omni_emb_list[i]) + if has_base: + txt_packed = self.apply_packing([txt_base, omni_curr], gather_indices, txt_mask_sorted) + else: + txt_packed = omni_curr + else: + txt_packed = txt_base + + if self.training and self.gradient_checkpointing: + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs, freqs_cis=freqs_cis, text_mask=txt_mask_sorted, is_flash=True) + return custom_forward + + img, txt_out_packed = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + img, txt_packed, vec, + use_reentrant=False + ) + else: + img, txt_out_packed = block( + img=img, txt=txt_packed, vec=vec, + freqs_cis=freqs_cis, text_mask=txt_mask_sorted, + is_flash=True + ) + + if self.repa_proj is not None and i == self.repa_layer_num: + repa_feat = self.repa_proj(img) + + if self.use_omni_embedding and has_base: + B, Total_Len, D = txt_packed.shape + restored_txt = torch.zeros(B, Total_Len, D, device=img.device, dtype=img.dtype) + restored_txt.scatter_(dim=1, index=gather_indices, src=txt_out_packed) + txt_base = restored_txt[:, :txt_base_len, :] + elif not self.use_omni_embedding: + txt_base = txt_out_packed + + # --- 5. Single Stream Loop --- + img_len = img.shape[1] + block_idx_offset = self.mm_double_blocks_depth + + for i, block in enumerate(self.single_blocks): + if self.use_omni_embedding: + omni_curr = self.deep_fusion_projs[block_idx_offset + i](omni_emb_list[block_idx_offset + i]) + if has_base: + txt_packed = self.apply_packing([txt_base, omni_curr], gather_indices, txt_mask_sorted) + else: + txt_packed = omni_curr + else: + txt_packed = txt_base + + x_in = torch.cat([img, txt_packed], dim=1) + txt_len_packed = txt_packed.shape[1] + + if self.training and self.gradient_checkpointing: + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs, txt_len=txt_len_packed, freqs_cis=freqs_cis, text_mask=txt_mask_sorted, is_flash=True) + return custom_forward + + x_out = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x_in, vec, + use_reentrant=False + ) + else: + x_out = block( + x=x_in, vec=vec, txt_len=txt_len_packed, + freqs_cis=freqs_cis, text_mask=txt_mask_sorted, + is_flash=True + ) + + img = x_out[:, :img_len, :] + + if self.repa_proj is not None and i + block_idx_offset == self.repa_layer_num: + repa_feat = self.repa_proj(img) + + txt_out_packed = x_out[:, img_len:, :] + + if self.use_omni_embedding and has_base: + B, Total_Len, D = txt_packed.shape + restored_txt = torch.zeros(B, Total_Len, D, device=img.device, dtype=img.dtype) + restored_txt.scatter_(dim=1, index=gather_indices, src=txt_out_packed) + txt_base = restored_txt[:, :txt_base_len, :] + elif not self.use_omni_embedding: + txt_base = txt_out_packed + + # --- 6. Final Layer --- + img_out = self.final_layer(img, vec) + out = self.unpatchify(img_out, shape_info=None if self.is_audio_type else grid_sizes_target) + + return out, 0, repa_feat + + def unpatchify(self, x, shape_info): + if self.is_audio_type: + return x.transpose(1, 2).contiguous() + else: + t, h, w = shape_info + c = self.out_channels + pt, ph, pw = self.patch_size + + x = x.reshape(x.shape[0], t, h, w, c, pt, ph, pw) + x = torch.einsum("nthwcopq->nctohpwq", x) + x = x.reshape(x.shape[0], c, t * pt, h * ph, w * pw) + return x diff --git a/unison/models/transformers/modules/activation_layers.py b/unison/models/transformers/modules/activation_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..b3f3fa95d91191c85bd24180aaec3e6b5c7d590f --- /dev/null +++ b/unison/models/transformers/modules/activation_layers.py @@ -0,0 +1,39 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch.nn as nn + + +def get_activation_layer(act_type): + """get activation layer + + Args: + act_type (str): the activation type + + Returns: + torch.nn.functional: the activation layer + """ + if act_type == "gelu": + return lambda: nn.GELU() + elif act_type == "gelu_tanh": + # Approximate `tanh` requires torch >= 1.13 + return lambda: nn.GELU(approximate="tanh") + elif act_type == "relu": + return nn.ReLU + elif act_type == "silu": + return nn.SiLU + else: + raise ValueError(f"Unknown activation type: {act_type}") diff --git a/unison/models/transformers/modules/attention.py b/unison/models/transformers/modules/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..8dd5b78a16f4e14dafde48980737b290e75c1e28 --- /dev/null +++ b/unison/models/transformers/modules/attention.py @@ -0,0 +1,299 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import einops +import torch +from typing import Optional +from loguru import logger +import numpy as np +import torch.nn.functional as F + +from unison.commons.parallel_states import get_parallel_state +from unison.utils.communications import ( + all_gather, + all_to_all_4D, +) +from unison.utils.flash_attn_no_pad import ( + flash_attn_no_pad, + flash_attn_no_pad_v3, +) +from unison.commons import maybe_fallback_attn_mode + +try: + from torch.nn.attention.flex_attention import flex_attention + + flex_attention = torch.compile(flex_attention, dynamic=False) + torch._dynamo.config.cache_size_limit = 192 + torch._dynamo.config.accumulated_cache_size_limit = 192 + flex_mask_cache = {} +except Exception: + logger.warning("Could not load Sliding Tile Attention of FlexAttn.") + +from unison.models.transformers.modules.ssta_attention import ssta_3d_attention +from unison.commons.infer_state import get_infer_state + + + +@torch.compiler.disable +def attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + drop_rate: float = 0.0, + attn_mask: Optional[torch.Tensor] = None, + causal: bool = False, + attn_mode: str = "flash", +) -> torch.Tensor: + """ + Compute attention using flash_attn_no_pad or torch scaled_dot_product_attention. + + Args: + q: Query tensor of shape [B, L, H, D] + k: Key tensor of shape [B, L, H, D] + v: Value tensor of shape [B, L, H, D] + drop_rate: Dropout rate for attention weights. + attn_mask: Optional attention mask of shape [B, L]. + causal: Whether to apply causal masking. + attn_mode: Attention mode, either "flash" or "torch". Defaults to "flash". + + Returns: + Output tensor after attention of shape [B, L, H*D] + """ + attn_mode = maybe_fallback_attn_mode(attn_mode) + + if attn_mode == "torch": + # transpose q,k,v dim to fit scaled_dot_product_attention + query = q.transpose(1, 2) # B * H * L * D + key = k.transpose(1, 2) # B * H * L * D + value = v.transpose(1, 2) # B * H * L * D + + if attn_mask is not None: + if attn_mask.dtype != torch.bool and attn_mask.dtype in [torch.int64, torch.int32]: + assert attn_mask.max() <= 1 and attn_mask.min() >= 0, f'Integer attention mask must be between 0 and 1 for torch attention.' + attn_mask = attn_mask.to(torch.bool) + elif attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(query.dtype) + raise NotImplementedError(f'Float attention mask is not implemented for torch attention.') + attn_mask1 = einops.rearrange(attn_mask, 'b l -> b 1 l 1') + attn_mask2 = einops.rearrange(attn_mask1, 'b 1 l 1 -> b 1 1 l') + attn_mask = attn_mask1 & attn_mask2 + + x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal) + + # transpose back + x = x.transpose(1, 2) # B * L * H * D + b, s, h, d = x.shape + out = x.reshape(b, s, -1) + return out + else: + # flash mode (default) + qkv = torch.stack([q, k, v], dim=2) + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.bool() + x = flash_attn_no_pad(qkv, attn_mask, causal=causal, dropout_p=drop_rate, softmax_scale=None) + b, s, a, d = x.shape + out = x.reshape(b, s, -1) + return out + + +@torch.compiler.disable +def parallel_attention(q, k, v, img_q_len, img_kv_len, + attn_mode=None, text_mask=None, + attn_param=None, + block_idx=None, + ): + return sequence_parallel_attention(q, k, v, img_q_len, img_kv_len, attn_mode, text_mask, attn_param=attn_param, block_idx=block_idx) + + +def sequence_parallel_attention(q, k, v, + img_q_len, img_kv_len, + attn_mode=None, text_mask=None, + attn_param=None, + block_idx=None, + ): + assert attn_mode is not None + query, encoder_query = q + key, encoder_key = k + value, encoder_value = v + + parallel_dims = get_parallel_state() + enable_sp = parallel_dims.sp_enabled + + if enable_sp: + sp_group = parallel_dims.sp_group + sp_size = parallel_dims.sp + sp_rank = parallel_dims.sp_rank + + if enable_sp: + # batch_size, seq_len, attn_heads, head_dim + query = all_to_all_4D(query, sp_group, scatter_dim=2, gather_dim=1) + key = all_to_all_4D(key, sp_group, scatter_dim=2, gather_dim=1) + value = all_to_all_4D(value, sp_group, scatter_dim=2, gather_dim=1) + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // sp_size + return encoder_state.narrow( + dim, sp_rank * local_heads, local_heads + ) + + encoder_query = shrink_head(encoder_query, dim=2) + encoder_key = shrink_head(encoder_key, dim=2) + encoder_value = shrink_head(encoder_value, dim=2) + + sequence_length = query.size(1) + encoder_sequence_length = encoder_query.size(1) + + attn_mode = maybe_fallback_attn_mode(attn_mode, get_infer_state(), block_idx) + + if attn_mode == "sageattn": + from sageattention import sageattn + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + hidden_states = sageattn(query, key, value, tensor_layout="NHD", is_causal=False) + elif attn_mode == "torch": + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + if text_mask is not None: + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + else: + attn_mask = None + + if attn_mask is not None: + if attn_mask.dtype != torch.bool and attn_mask.dtype in [torch.int64, torch.int32]: + assert attn_mask.max() <= 1 and attn_mask.min() >= 0, f'Integer attention mask must be between 0 and 1 for torch attention.' + attn_mask = attn_mask.to(torch.bool) + elif attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(query.dtype) + raise NotImplementedError(f'Float attention mask is not implemented for torch attention.') + + # transpose q,k,v dim to fit scaled_dot_product_attention + query = query.transpose(1, 2) # B * Head_num * length * dim + key = key.transpose(1, 2) # B * Head_num * length * dim + value = value.transpose(1, 2) # B * Head_num * length * dim + + def score_mod(score, b, h, q_idx, kv_idx): + return torch.where(attn_mask[b, q_idx] & attn_mask[b, kv_idx], score, float('-inf')) + + hidden_states = flex_attention(query, key, value, score_mod=score_mod) + + # transpose back + hidden_states = hidden_states.transpose(1, 2) + + elif attn_mode == "flash2": + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + # B, S, 3, H, D + qkv = torch.stack([query, key, value], dim=2) + + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) + + elif attn_mode == "flash3": + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) + # B, S, 3, H, D + qkv = torch.stack([query, key, value], dim=2) + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + hidden_states = flash_attn_no_pad_v3(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) + + elif attn_mode == "flex-block-attn": + sparse_type = attn_param["attn_sparse_type"] # sta/block_attn/ssta + ssta_threshold = attn_param["ssta_threshold"] + ssta_lambda = attn_param["ssta_lambda"] + ssta_sampling_type = attn_param["ssta_sampling_type"] + ssta_adaptive_pool = attn_param["ssta_adaptive_pool"] + + attn_pad_type = attn_param["attn_pad_type"] # repeat/zero + attn_use_text_mask = attn_param["attn_use_text_mask"] + attn_mask_share_within_head = attn_param["attn_mask_share_within_head"] + + ssta_topk = attn_param["ssta_topk"] + thw = attn_param["thw"] + tile_size = attn_param["tile_size"] + win_size = attn_param["win_size"][0].copy() + + def get_image_tile(tile_size): + block_size = np.prod(tile_size) + if block_size == 384: + tile_size = (1, 16, 24) + elif block_size == 128: + tile_size = (1, 16, 8) + elif block_size == 64: + tile_size = (1, 8, 8) + elif block_size == 16: + tile_size = (1, 4, 4) + else: + raise ValueError(f"Error tile_size {tile_size}, only support in [16, 64, 128, 384]") + return tile_size + + if thw[0] == 1: + tile_size = get_image_tile(tile_size) + win_size = [1, 1, 1] + elif thw[0] <= 31: # 16fps: 5 * 16 / 4 + 1 = 21; 24fps: 5 * 24 / 4 + 1 = 31 + ssta_topk = ssta_topk // 2 + + # Concatenate and permute query, key, value to (B, H, S, D) + query = torch.cat([query, encoder_query], dim=1).permute(0, 2, 1, 3) + key = torch.cat([key, encoder_key], dim=1).permute(0, 2, 1, 3) + value = torch.cat([value, encoder_value], dim=1).permute(0, 2, 1, 3) + + assert ( + query.shape[-1] == 128 + ), "The last dimension of query, key and value must be 128 for flex-block-attn." + + hidden_states = ssta_3d_attention( + query, + key, + value, + thw, + topk=ssta_topk, + tile_thw=tile_size, + kernel_thw=win_size, + text_len=encoder_sequence_length, + sparse_type=sparse_type, + threshold=ssta_threshold, + lambda_=ssta_lambda, + pad_type=attn_pad_type, + text_mask=text_mask if attn_use_text_mask else None, + sampling_type=ssta_sampling_type, + adaptive_pool=ssta_adaptive_pool, + mask_share_within_head=attn_mask_share_within_head, + ) + hidden_states, sparse_ratio = hidden_states + hidden_states = hidden_states.permute(0, 2, 1, 3) + + else: + raise NotImplementedError( + f'Unsupported attention mode: {attn_mode}. Only torch, flash, flash3, sageattn and flex-block-attn are supported.' + ) + + + if enable_sp: + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes((sequence_length, encoder_sequence_length), dim=1) + hidden_states = all_to_all_4D(hidden_states, sp_group, scatter_dim=1, gather_dim=2) + encoder_hidden_states = all_gather(encoder_hidden_states, dim=2, group=sp_group).contiguous() + hidden_states = hidden_states.to(query.dtype) + encoder_hidden_states = encoder_hidden_states.to(query.dtype) + hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + b, s, a, d = hidden_states.shape + hidden_states = hidden_states.reshape(b, s, -1) + + return hidden_states diff --git a/unison/models/transformers/modules/embed_layers.py b/unison/models/transformers/modules/embed_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..9a6fbdac3473d2e36dfb13e7f1c165bbe7566d1c --- /dev/null +++ b/unison/models/transformers/modules/embed_layers.py @@ -0,0 +1,463 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Tuple, Union +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models import ModelMixin +from math import pi +from unison.commons import to_2tuple, to_3tuple + + +class ChannelLastConv1d(nn.Module): + """ + Conv1d that consumes [B, L, C] and returns [B, L, C]. + Used for Audio branch. + """ + def __init__(self, in_channels, out_channels, kernel_size, padding=0, stride=1, bias=True): + super().__init__() + self.conv = nn.Conv1d( + in_channels, out_channels, kernel_size=kernel_size, + padding=padding, stride=stride, bias=bias, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Input: [B, L, C] + x = x.transpose(1, 2) # [B, L, C] -> [B, C, L] + x = self.conv(x) + x = x.transpose(1, 2) # [B, C, L] -> [B, L, C] + return x + +class ConvMLP(nn.Module): + """ + SwiGLU style ConvMLP used by MMAudio/Ovi style embedding. + """ + def __init__( + self, + dim: int, + hidden_dim: int, + multiple_of: int = 256, + kernel_size: int = 3, + padding: int = 1, + ): + super().__init__() + hidden_dim = int(2 * hidden_dim / 3) + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + + self.w1 = ChannelLastConv1d(dim, hidden_dim, kernel_size=kernel_size, padding=padding, bias=False) + self.w2 = ChannelLastConv1d(hidden_dim, dim, kernel_size=kernel_size, padding=padding, bias=False) + self.w3 = ChannelLastConv1d(dim, hidden_dim, kernel_size=kernel_size, padding=padding, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + +class UniversalPatchEmbed(ModelMixin, ConfigMixin): + """ + Universal Patch Embedding for HunyuanVideo 1.5 Architecture. + Supports both Video (3D) and Audio (1D w/ Context). + + Audio patch embed types (audio_patch_type): + "conv_mlp" — Original: Conv1d + SwiGLU ConvMLP (~132M params, heavy) + "mlp" — Pointwise MLP: Linear-SiLU-Linear (~2.4M params) + "conv_lite" — Small Conv + Linear: Conv1d(k=3)-SiLU-Linear (~2.5M params) + """ + + @register_to_config + def __init__( + self, + patch_size=(1, 2, 2), + in_chans=48, + embed_dim=1280, + is_reshape_temporal_channels=False, + concat_condition=False, + norm_layer=None, + flatten=True, + bias=True, + is_audio=False, + audio_kernel_size=7, + audio_padding=3, + audio_patch_type="conv_mlp", + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + + self.is_audio = is_audio + self.flatten = flatten + self.audio_patch_type = audio_patch_type + self.patch_size = to_3tuple(patch_size) if not is_audio else patch_size + + orig_in_chans = in_chans + if concat_condition: + if is_reshape_temporal_channels: + in_chans = in_chans + in_chans//2 + 1 + else: + in_chans = in_chans * 2 + 1 + + self.in_chans = in_chans + self.orig_in_chans = orig_in_chans + + if is_audio: + if audio_patch_type == "conv_mlp": + self.proj = nn.Sequential( + ChannelLastConv1d(in_chans, embed_dim, kernel_size=audio_kernel_size, padding=audio_padding, bias=bias), + nn.SiLU(), + ConvMLP(embed_dim, embed_dim * 4, kernel_size=audio_kernel_size, padding=audio_padding), + ) + first_layer = self.proj[0].conv + elif audio_patch_type == "mlp": + self.proj = nn.Sequential( + nn.Linear(in_chans, embed_dim, bias=bias), + nn.SiLU(), + nn.Linear(embed_dim, embed_dim, bias=bias), + ) + first_layer = self.proj[0] + elif audio_patch_type == "conv_lite": + self.proj = nn.Sequential( + ChannelLastConv1d(in_chans, embed_dim, kernel_size=3, padding=1, bias=bias), + nn.SiLU(), + nn.Linear(embed_dim, embed_dim, bias=bias), + ) + first_layer = self.proj[0].conv + else: + raise ValueError(f"Unknown audio_patch_type: {audio_patch_type}") + + self._init_audio_proj(first_layer, orig_in_chans, concat_condition, bias) + + else: + self.patch_size = to_3tuple(patch_size) + self.proj = nn.Conv3d( + in_chans, + embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=bias, + **factory_kwargs, + ) + + nn.init.xavier_uniform_(self.proj.weight[:, :orig_in_chans].view(self.proj.weight[:, :orig_in_chans].size(0), -1)) + if concat_condition: + nn.init.zeros_(self.proj.weight[:, orig_in_chans:].view(self.proj.weight[:, orig_in_chans:].size(0), -1)) + if bias: + nn.init.zeros_(self.proj.bias) + + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + @staticmethod + def _init_audio_proj(first_layer, orig_in_chans, concat_condition, bias): + """Xavier init + zero-init concat channels for any audio proj first layer.""" + if isinstance(first_layer, nn.Conv1d): + nn.init.xavier_uniform_(first_layer.weight) + if bias and first_layer.bias is not None: + nn.init.zeros_(first_layer.bias) + if concat_condition: + with torch.no_grad(): + first_layer.weight[:, orig_in_chans:, :] = 0.0 + elif isinstance(first_layer, nn.Linear): + nn.init.xavier_uniform_(first_layer.weight) + if bias and first_layer.bias is not None: + nn.init.zeros_(first_layer.bias) + if concat_condition: + with torch.no_grad(): + first_layer.weight[:, orig_in_chans:] = 0.0 + + def forward(self, x): + if self.is_audio: + if x.dim() == 5: + x = x.squeeze(2).squeeze(2) + elif x.dim() == 4: + x = x.squeeze(2) + + # [B, C, L] -> [B, L, C] + x = x.transpose(1, 2) + x = self.proj(x) + + else: + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) + + x = self.norm(x) + return x + +class PatchEmbed(nn.Module): + """2D Image to Patch Embedding + + Image to Patch Embedding using Conv2d + + A convolution based approach to patchifying a 2D image w/ embedding projection. + + Based on the impl in https://github.com/google-research/vision_transformer + + Hacked together by / Copyright 2020 Ross Wightman + + Remove the _assert function in forward function to be compatible with multi-resolution images. + """ + + def __init__( + self, + patch_size=16, + in_chans=3, + embed_dim=768, + is_reshape_temporal_channels=False, + concat_condition=True, + norm_layer=None, + flatten=True, + bias=True, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + self.flatten = flatten + + # Only support concat mode (multitask mask training) + orig_in_chans = in_chans + if concat_condition: + if is_reshape_temporal_channels: + in_chans = in_chans + in_chans//2 + 1 + else: + in_chans = in_chans * 2 + 1 + + self.proj = nn.Conv3d( + in_chans, + embed_dim, + kernel_size=patch_size, + stride=patch_size, + bias=bias, + **factory_kwargs, + ) + + nn.init.xavier_uniform_(self.proj.weight[:, :orig_in_chans].view(self.proj.weight[:, :orig_in_chans].size(0), -1)) + # Special initialization for concat mode + nn.init.zeros_(self.proj.weight[:, orig_in_chans:].view(self.proj.weight[:, orig_in_chans:].size(0), -1)) + + if bias: + nn.init.zeros_(self.proj.bias) + + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCHW -> BNC + x = self.norm(x) + return x + + +class TextProjection(nn.Module): + """ + Projects text embeddings. Also handles dropout for classifier-free guidance. + + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.linear_1 = nn.Linear( + in_features=in_channels, + out_features=hidden_size, + bias=True, + **factory_kwargs + ) + self.act_1 = act_layer() + self.linear_2 = nn.Linear( + in_features=hidden_size, + out_features=hidden_size, + bias=True, + **factory_kwargs + ) + + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + + +class VisionProjection(torch.nn.Module): + + def __init__(self, input_dim, output_dim): + super().__init__() + + self.proj = torch.nn.Sequential( + torch.nn.LayerNorm(input_dim), + torch.nn.Linear(input_dim, input_dim), + torch.nn.GELU(), + torch.nn.Linear(input_dim, output_dim), + torch.nn.LayerNorm(output_dim) + ) + + + def forward(self, vision_embeds): + return self.proj(vision_embeds) + +class ClipVisionProjection(nn.Module): + def __init__(self, in_channels, out_channels): + super().__init__() + self.up = nn.Linear(in_channels, out_channels * 3) + self.down = nn.Linear(out_channels * 3, out_channels) + torch.nn.init.zeros_(self.down.weight) + torch.nn.init.zeros_(self.down.bias) + + def forward(self, x): + projected_x = self.down(nn.functional.silu(self.up(x))) + return projected_x + +def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): the dimension of the output. + max_period (int): controls the minimum frequency of the embeddings. + + Returns: + embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. + + .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__( + self, + hidden_size, + act_layer, + frequency_embedding_size=256, + max_period=10000, + out_size=None, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + if out_size is None: + out_size = hidden_size + + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True, **factory_kwargs), + act_layer(), + nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.normal_(self.mlp[2].weight, std=0.02) + + def forward(self, t): + t_freq = timestep_embedding( + t, self.frequency_embedding_size, self.max_period + ).type(self.mlp[0].weight.dtype) + t_emb = self.mlp(t_freq) + return t_emb + +class StableAudioPositionalEmbedding(nn.Module): + """Used for continuous time + Adapted from Stable Audio Open. + """ + + def __init__(self, dim: int): + super().__init__() + assert (dim % 2) == 0 + half_dim = dim // 2 + self.weights = nn.Parameter(torch.randn(half_dim)) + + def forward(self, times: torch.Tensor) -> torch.Tensor: + times = times[..., None] + freqs = times * self.weights[None] * 2 * pi + fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1) + fouriered = torch.cat((times, fouriered), dim=-1) + return fouriered + +class DurationEmbedder(nn.Module): + """ + A simple linear projection model to map numbers to a latent space. + + Code is adapted from + https://github.com/Stability-AI/stable-audio-tools + + Args: + number_embedding_dim (`int`): + Dimensionality of the number embeddings. + min_value (`int`): + The minimum value of the seconds number conditioning modules. + max_value (`int`): + The maximum value of the seconds number conditioning modules + internal_dim (`int`): + Dimensionality of the intermediate number hidden states. + """ + + def __init__( + self, + number_embedding_dim, + min_value, + max_value, + internal_dim= 256, + ): + super().__init__() + self.time_positional_embedding = nn.Sequential( + StableAudioPositionalEmbedding(internal_dim), + nn.Linear(in_features=internal_dim + 1, out_features=number_embedding_dim), + ) + + self.number_embedding_dim = number_embedding_dim + self.min_value = min_value + self.max_value = max_value + self.dtype = torch.float32 + + def forward( + self, + floats: torch.Tensor, + ): + floats = floats.clamp(self.min_value, self.max_value) + + normalized_floats = (floats - self.min_value) / ( + self.max_value - self.min_value + ) + + # Cast floats to same type as embedder + embedder_dtype = next(self.time_positional_embedding.parameters()).dtype + normalized_floats = normalized_floats.to(embedder_dtype) + + embedding = self.time_positional_embedding(normalized_floats) + float_embeds = embedding.view(-1, 1, self.number_embedding_dim) + + return float_embeds \ No newline at end of file diff --git a/unison/models/transformers/modules/mlp_layers.py b/unison/models/transformers/modules/mlp_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..a633f5fb647f4bf59713e74ee85037c242a4cfb5 --- /dev/null +++ b/unison/models/transformers/modules/mlp_layers.py @@ -0,0 +1,137 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +# Modified from timm library: +# https://github.com/huggingface/pytorch-image-models/blob/648aaa41233ba83eb38faf5ba9d415d574823241/timm/layers/mlp.py#L13 + +from functools import partial + +import torch +import torch.nn as nn + +from unison.commons import to_2tuple +from .modulate_layers import modulate + + +class MLP(nn.Module): + """MLP as used in Vision Transformer, MLP-Mixer and related networks""" + + def __init__( + self, + in_channels, + hidden_channels=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + drop=0.0, + use_conv=False, + device=None, + dtype=None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + out_features = out_features or in_channels + hidden_channels = hidden_channels or in_channels + bias = to_2tuple(bias) + drop_probs = to_2tuple(drop) + linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear + + self.fc1 = linear_layer(in_channels, hidden_channels, bias=bias[0], **factory_kwargs) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.norm = norm_layer(hidden_channels, **factory_kwargs) if norm_layer is not None else nn.Identity() + self.fc2 = linear_layer(hidden_channels, out_features, bias=bias[1], **factory_kwargs) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.norm(x) + x = self.fc2(x) + x = self.drop2(x) + return x + + +class LinearWarpforSingle(nn.Module): + def __init__(self, in_dim: int, out_dim: int, bias=False, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.fc = nn.Linear(in_dim, out_dim, bias=bias, **factory_kwargs) + + def forward(self, x, y): + input = torch.cat([x.contiguous(), y.contiguous()], dim=2).contiguous() + return self.fc(input) + + +# +class MLPEmbedder(nn.Module): + """copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py""" + + def __init__(self, in_dim: int, hidden_dim: int, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True, **factory_kwargs) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True, **factory_kwargs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class FinalLayer(nn.Module): + """The final layer of DiT.""" + + def __init__(self, hidden_size, patch_size, out_channels, act_layer, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + # Just use LayerNorm for the final layer + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + if isinstance(patch_size, int): + self.linear = nn.Linear( + hidden_size, + patch_size * patch_size * out_channels, + bias=True, + **factory_kwargs, + ) + else: + out_size = ( + patch_size[0] * patch_size[1] * patch_size[2] if len(patch_size) == 3 else patch_size[0] * patch_size[1] + ) * out_channels + self.linear = nn.Linear( + hidden_size, + out_size, + bias=True, + ) + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + # Here we don't distinguish between the modulate types. Just use the simple one. + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift=shift, scale=scale) + x = self.linear(x) + return x diff --git a/unison/models/transformers/modules/modulate_layers.py b/unison/models/transformers/modules/modulate_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..614424bda53166fa5e425335073eed2c1c51fe11 --- /dev/null +++ b/unison/models/transformers/modules/modulate_layers.py @@ -0,0 +1,92 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Callable + +import torch +import torch.nn as nn + + +class ModulateDiT(nn.Module): + """Modulation layer for DiT.""" + + def __init__( + self, + hidden_size: int, + factor: int, + act_layer: Callable, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.act = act_layer() + self.linear = nn.Linear(hidden_size, factor * hidden_size, bias=True, **factory_kwargs) + # Zero-initialize the modulation + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(self.act(x)) + + +def modulate(x, shift=None, scale=None): + """modulate by shift and scale + + Args: + x (torch.Tensor): input tensor. + shift (torch.Tensor, optional): shift tensor. Defaults to None. + scale (torch.Tensor, optional): scale tensor. Defaults to None. + + Returns: + torch.Tensor: the output tensor after modulate. + """ + if scale is None and shift is None: + return x + elif shift is None: + return x * (1 + scale.unsqueeze(1)) + elif scale is None: + return x + shift.unsqueeze(1) + else: + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +def apply_gate(x, gate=None, tanh=False): + """AI is creating summary for apply_gate + + Args: + x (torch.Tensor): input tensor. + gate (torch.Tensor, optional): gate tensor. Defaults to None. + tanh (bool, optional): whether to use tanh function. Defaults to False. + + Returns: + torch.Tensor: the output tensor after apply gate. + """ + if gate is None: + return x + if tanh: + return x * gate.unsqueeze(1).tanh() + else: + return x * gate.unsqueeze(1) + + +def ckpt_wrapper(module): + def ckpt_forward(*inputs): + outputs = module(*inputs) + return outputs + + return ckpt_forward + diff --git a/unison/models/transformers/modules/norm_layers.py b/unison/models/transformers/modules/norm_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..1d02aa8d220c225581e4028062784f0104ec8598 --- /dev/null +++ b/unison/models/transformers/modules/norm_layers.py @@ -0,0 +1,97 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch +import torch.nn as nn + + +class RMSNorm(nn.Module): + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x): + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def reset_parameters(self): + if hasattr(self, "weight"): + self.weight.fill_(1) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +def get_norm_layer(norm_layer): + """ + Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") diff --git a/unison/models/transformers/modules/posemb_layers.py b/unison/models/transformers/modules/posemb_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..2fc8f80bc08668de8b0131bd233af30894654404 --- /dev/null +++ b/unison/models/transformers/modules/posemb_layers.py @@ -0,0 +1,413 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch +from typing import Union, Tuple, List + + +def _to_tuple(x, dim=2): + if isinstance(x, int): + return (x,) * dim + elif len(x) == dim: + return x + else: + raise ValueError(f"Expected length {dim} or int, but got {x}") + + +def get_meshgrid_nd(start, *args, dim=2): + """ + Get n-D meshgrid with start, stop and num. + + Args: + start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop, + step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num + should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in + n-tuples. + *args: See above. + dim (int): Dimension of the meshgrid. Defaults to 2. + + Returns: + grid (np.ndarray): [dim, ...] + """ + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0,) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = [stop[i] - start[i] for i in range(dim)] + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [W, H, D] + grid = torch.stack(grid, dim=0) # [dim, W, H, D] + + return grid + + +################################################################################# +# Rotary Positional Embedding Functions # +################################################################################# +# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80 + + +def reshape_for_broadcast( + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], + x: torch.Tensor, + head_first=False, +): + """ + Reshape frequency tensor for broadcasting it with another tensor. + + This function reshapes the frequency tensor to have the same shape as the target tensor 'x' + for the purpose of broadcasting the frequency tensor during element-wise operations. + + Notes: + When using FlashMHAModified, head_first should be False. + When using Attention, head_first should be True. + + Args: + freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped. + x (torch.Tensor): Target tensor for broadcasting compatibility. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + torch.Tensor: Reshaped frequency tensor. + + Raises: + AssertionError: If the frequency tensor doesn't match the expected shape. + AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions. + """ + ndim = x.ndim + assert 0 <= 1 < ndim + + if isinstance(freqs_cis, tuple): + # freqs_cis: (cos, sin) in real space + if head_first: + assert freqs_cis[0].shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [ + d if i == ndim - 2 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + else: + assert freqs_cis[0].shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape) + else: + # freqs_cis: values in complex space + if head_first: + assert freqs_cis.shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [ + d if i == ndim - 2 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + else: + assert freqs_cis.shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return freqs_cis.view(*shape) + + +def rotate_half(x): + x_real, x_imag = ( + x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1) + ) # [B, S, H, D//2] + return torch.stack([-x_imag, x_real], dim=-1).flatten(3) + + +def apply_rotary_emb( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + head_first: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. + + This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided + frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor + is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are + returned as real tensors. + + Args: + xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D] + xk (torch.Tensor): Key tensor to apply rotary embeddings. [B, S, H, D] + freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. + + """ + xk_out = None + if isinstance(freqs_cis, tuple): + cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first) # [S, D] + cos, sin = cos.to(xq.device), sin.to(xq.device) + # real * cos - imag * sin + # imag * cos + real * sin + xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq) + xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk) + else: + # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex) + xq_ = torch.view_as_complex( + xq.float().reshape(*xq.shape[:-1], -1, 2) + ) # [B, S, H, D//2] + freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to( + xq.device + ) # [S, D//2] --> [1, S, 1, D//2] + # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin) + # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq) + xk_ = torch.view_as_complex( + xk.float().reshape(*xk.shape[:-1], -1, 2) + ) # [B, S, H, D//2] + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk) + + return xq_out, xk_out + + +def get_nd_rotary_pos_embed( + rope_dim_list, + start, + *args, + theta=10000.0, + use_real=False, + theta_rescale_factor: Union[float, List[float]] = 1.0, + interpolation_factor: Union[float, List[float]] = 1.0, +): + """ + This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure. + + Args: + rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n. + sum(rope_dim_list) should equal to head_dim of attention layer. + start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start, + args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. + *args: See above. + theta (float): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers. + Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real + part and an imaginary part separately. + theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0. + + Returns: + pos_embed (torch.Tensor): [HW, D/2] + """ + + grid = get_meshgrid_nd( + start, *args, dim=len(rope_dim_list) + ) # [3, W, H, D] / [2, W, H] + + if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float): + theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list) + elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1: + theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list) + assert len(theta_rescale_factor) == len( + rope_dim_list + ), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + + if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float): + interpolation_factor = [interpolation_factor] * len(rope_dim_list) + elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1: + interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list) + assert len(interpolation_factor) == len( + rope_dim_list + ), "len(interpolation_factor) should equal to len(rope_dim_list)" + + # use 1/ndim of dimensions to encode grid_axis + embs = [] + for i in range(len(rope_dim_list)): + emb = get_1d_rotary_pos_embed( + rope_dim_list[i], + grid[i].reshape(-1), + theta, + use_real=use_real, + theta_rescale_factor=theta_rescale_factor[i], + interpolation_factor=interpolation_factor[i], + ) # 2 x [WHD, rope_dim_list[i]] + embs.append(emb) + + if use_real: + cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2) + sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2) + return cos, sin + else: + emb = torch.cat(embs, dim=1) # (WHD, D/2) + return emb + + +def get_1d_rotary_pos_embed( + dim: int, + pos: Union[torch.FloatTensor, int], + theta: float = 10000.0, + use_real: bool = False, + theta_rescale_factor: float = 1.0, + interpolation_factor: float = 1.0, +) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """ + Precompute the frequency tensor for complex exponential (cis) with given dimensions. + (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.) + + This function calculates a frequency tensor with complex exponential using the given dimension 'dim' + and the end index 'end'. The 'theta' parameter scales the frequencies. + The returned tensor contains complex values in complex64 data type. + + Args: + dim (int): Dimension of the frequency tensor. + pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar + theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool, optional): If True, return real part and imaginary part separately. + Otherwise, return complex numbers. + theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0. + + Returns: + freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2] + freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D] + """ + if isinstance(pos, int): + pos = torch.arange(pos).float() + + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + if theta_rescale_factor != 1.0: + theta *= theta_rescale_factor ** (dim / (dim - 2)) + + freqs = 1.0 / ( + theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim) + ) # [D/2] + # assert interpolation_factor == 1.0, f"interpolation_factor: {interpolation_factor}" + freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2] + if use_real: + freqs_cos = freqs.cos().repeat_interleave(2, dim=1) # [S, D] + freqs_sin = freqs.sin().repeat_interleave(2, dim=1) # [S, D] + return freqs_cos, freqs_sin + else: + freqs_cis = torch.polar( + torch.ones_like(freqs), freqs + ) # complex64 # [S, D/2] + return freqs_cis + +def get_audio_rotary_pos_embed( + rope_dim_list: List[int], # [T_dim, H_dim, W_dim] same as Video + length: int, # Audio sequence length L + theta: float = 10000.0, + freqs_scaling: float = 1.0, # Linear Scaling for FPS alignment + use_real: bool = True, +): + """ + """ + + # 1. Compute Active Part (Time Dim) + t_dim = rope_dim_list[0] + assert t_dim % 2 == 0 + + # Construct position index [0, 1, ..., L-1] + pos = torch.arange(length).float() + + # Compute frequency index [0, 2/d, 4/d, ...] + freqs_idx = torch.arange(0, t_dim, 2).float() / t_dim + + # Compute base frequency + freqs = 1.0 / (theta ** freqs_idx) + + # Linear Scaling: Align Video/Audio time rate + freqs = freqs * freqs_scaling + + # Generate angles: [L, t_dim/2] + angles_t = torch.outer(pos, freqs) + + # 2. Compute Passive Part (Space Dims) + # This part is not rotated (Pass-through) -> set angle to 0 + + rest_dim = sum(rope_dim_list[1:]) # H_dim + W_dim + if rest_dim > 0: + assert rest_dim % 2 == 0 + # Pad 0 angle -> [L, rest_dim/2] + angles_rest = torch.zeros(length, rest_dim // 2) + + # Concat -> [L, total_dim/2] + angles = torch.cat([angles_t, angles_rest], dim=1) + else: + angles = angles_t + + # 3. Generate Output + if use_real: + # Need [L, D] format, cos/sin need to be repeated + cos = angles.cos().repeat_interleave(2, dim=1) + sin = angles.sin().repeat_interleave(2, dim=1) + return cos, sin + else: + return torch.polar(torch.ones_like(angles), angles) + + +def get_audio_rotary_pos_embed_with_offset( + rope_dim_list: List[int], + length: int, + offset: int = 0, + theta: float = 10000.0, + freqs_scaling: float = 1.0, + use_real: bool = True, +): + """Same as get_audio_rotary_pos_embed but positions start from `offset`.""" + t_dim = rope_dim_list[0] + assert t_dim % 2 == 0 + + pos = torch.arange(offset, offset + length).float() + + freqs_idx = torch.arange(0, t_dim, 2).float() / t_dim + freqs = 1.0 / (theta ** freqs_idx) + freqs = freqs * freqs_scaling + angles_t = torch.outer(pos, freqs) + + rest_dim = sum(rope_dim_list[1:]) + if rest_dim > 0: + assert rest_dim % 2 == 0 + angles_rest = torch.zeros(length, rest_dim // 2) + angles = torch.cat([angles_t, angles_rest], dim=1) + else: + angles = angles_t + + if use_real: + cos = angles.cos().repeat_interleave(2, dim=1) + sin = angles.sin().repeat_interleave(2, dim=1) + return cos, sin + else: + return torch.polar(torch.ones_like(angles), angles) \ No newline at end of file diff --git a/unison/models/transformers/modules/ssta_attention.py b/unison/models/transformers/modules/ssta_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..438f3eedeadd1d3812186b8d7cf4b5573b63c794 --- /dev/null +++ b/unison/models/transformers/modules/ssta_attention.py @@ -0,0 +1,700 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import math +import torch +import numpy as np +from einops import rearrange +from functools import lru_cache + +def tile(x, canvas_thw, tile_thw, sp_size=1): + r"""Rearrange tensor into tiles for block-based attention. + + Args: + x: Input tensor with shape (b, head, s, d) where s = t * h * w + canvas_thw: Tuple of (t, h, w) representing temporal, height, width dimensions + tile_thw: Tuple of (tile_t, tile_h, tile_w) representing tile dimensions + sp_size: Spatial size parameter, defaults to 1 + + Returns: + Rearranged tensor organized by tiles + """ + b, h, s, d = x.shape + t, h, w = canvas_thw + assert t * h * w == s, f"t:{t} * h:{h} * w:{w} == s:{s}" + + tile_t_dim, tile_h_dim, tile_w_dim = tile_thw + n_t = int(t/tile_t_dim) + n_h = int(h/tile_h_dim) + n_w = int(w/tile_w_dim) + x = rearrange(x, "b head (sp t h w) d -> b head (t sp h w) d", sp=sp_size, t=t // sp_size, h=h, w=w) + return rearrange(x, + "b h (n_t ts_t n_h ts_h n_w ts_w) d -> b h (n_t n_h n_w ts_t ts_h ts_w) d", + n_t=n_t, + n_h=n_h, + n_w=n_w, + ts_t=tile_t_dim, + ts_h=tile_h_dim, + ts_w=tile_w_dim) + +def untile(x, canvas_thw, tile_thw, sp_size=1): + r"""Reverse the tiling operation to restore original tensor layout. + + Args: + x: Tiled tensor + canvas_thw: Tuple of (t, h, w) representing temporal, height, width dimensions + tile_thw: Tuple of (tile_t, tile_h, tile_w) representing tile dimensions + sp_size: Spatial size parameter, defaults to 1 + + Returns: + Restored tensor with original layout + """ + t, h, w = canvas_thw + + tile_t_dim, tile_h_dim, tile_w_dim = tile_thw + n_t = int(t/tile_t_dim) + n_h = int(h/tile_h_dim) + n_w = int(w/tile_w_dim) + + x = rearrange(x, + "b h (n_t n_h n_w ts_t ts_h ts_w) d -> b h (n_t ts_t n_h ts_h n_w ts_w) d", + n_t=n_t, + n_h=n_h, + n_w=n_w, + ts_t=tile_t_dim, + ts_h=tile_h_dim, + ts_w=tile_w_dim) + return rearrange(x, "b head (t sp h w) d -> b head (sp t h w) d", sp=sp_size, t=t // sp_size, h=h, w=w) + +def get_tile_t_h_w(tile_id, tile_thw_dim): + """Extract temporal, height, and width indices from a flattened tile ID.""" + tile_t_dim, tile_h_dim, tile_w_dim = tile_thw_dim + tile_t = tile_id // (tile_h_dim * tile_w_dim) + tile_h = (tile_id % (tile_h_dim * tile_w_dim)) // tile_w_dim + tile_w = tile_id % tile_w_dim + + return tile_t, tile_h, tile_w +def importance_sampling(q, k, topk, threshold=0.0, lambda_=0.9, adaptive_pool=None): + r"""Select top-k blocks based on importance scores considering both similarity and redundancy. + + Args: + q: Query tensor with shape (B, H, S, D) + k: Key tensor with shape (B, H, K, D) + topk: Number of top blocks to select + threshold: Threshold parameter (not implemented) + lambda_: Weight factor balancing similarity and redundancy + adaptive_pool: Adaptive pooling parameter (unused) + + Returns: + top_block_indices: Indices of selected blocks with shape (B, H, S, topk) + """ + if threshold > 0.0: + raise NotImplementedError("importance_sampling with threshold not implemented") + + q = q / q.norm(dim=-1, keepdim=True) + k = k / k.norm(dim=-1, keepdim=True) + gate_similarity = torch.einsum("bhsd,bhkd->bhsk", q, k) + gate_similarity = (gate_similarity + 1.0) / 2.0 + gate_unique = torch.einsum("bhsd,bhkd->bhsk", k, k) + gate_unique = (gate_unique + 1.0) / 2.0 + + B, H, K_num, D = k.shape + diag_indices = torch.arange(K_num, device=k.device) + gate_unique[:, :, diag_indices, diag_indices] = torch.nan + + mean_redundancy = torch.nanmean(gate_unique, dim=-2, keepdim=True) + + importance_scores = lambda_ * gate_similarity - (1 - lambda_) * mean_redundancy + + topk = min(topk, importance_scores.size(-1)) + _, top_block_indices = importance_scores.topk(k=topk, dim=-1, sorted=False) + return top_block_indices + +def similarity_sampling(q, k, topk, + threshold=0.0, + block_num=None, + adaptive_pool=None, + temperature=0.01): + r"""Select top-k blocks based on similarity scores between query and key averages. + + Args: + q: Query tensor with shape (B, H, S, D) + k: Key tensor with shape (B, H, K, D) + topk: Number of top blocks to select + threshold: Cumulative score threshold for dynamic topk selection + block_num: Total number of blocks (unused) + adaptive_pool: Adaptive pooling parameter (unused) + temperature: Temperature scaling for softmax + + Returns: + top_block_indices: Indices of selected blocks with shape (B, H, S, topk) + """ + if threshold > 0.0: + gate = torch.einsum("bhsd,bhkd->bhsk", q, k) + gate = gate / temperature + gate_ = torch.softmax(gate, dim=-1) + sorted_gate, sorted_indices = torch.sort(gate_, dim=-1, descending=True) + cum_scores = torch.cumsum(sorted_gate, dim=-1) + above_threshold = cum_scores >= threshold + has_any_above = above_threshold.any(dim=-1, keepdim=True) + # Find the first position exceeding threshold, add 1 to get the number of blocks to select + dynamic_topk = above_threshold.int().argmax(dim=-1, keepdim=True) + 1 + dynamic_topk = torch.where(has_any_above, dynamic_topk, + torch.full_like(dynamic_topk, topk)) + # Generate fixed-shape indices and pad with top1 indices for parts not meeting threshold + # Limit minimum k value + dynamic_topk = torch.clamp(dynamic_topk, min=8, max=topk) + indices = torch.arange(gate.size(-1), device=gate.device).expand(gate.size()) + mask = (indices < dynamic_topk).int() + # Use top1 to pad to topk to ensure same size + top_block_indices = torch.gather(sorted_indices, -1, indices) * mask + (1 - mask) * sorted_indices[..., 0:1] + else: + gate = torch.einsum("bhsd,bhkd->bhsk", q, k) + topk = min(topk, gate.size(-1)) + _, top_block_indices = gate.topk(k=topk, dim=-1, sorted=False) + return top_block_indices + +def create_moba_3d_mask(q, + k, + canvas_thw, + topk, + tile_thw, + kernel_thw, + text_block_num=0, + add_text_mask=False, + threshold=0.0, + lambda_=None, + mask_share_within_head=True, + q_block_avg_pool=True, + adaptive_pool=None, + sampling_type=None): + r"""Create MOBA (Mixture of Block Attention) 3D mask for sparse attention. + + Args: + q: Query tensor + k: Key tensor + canvas_thw: Canvas dimensions (t, h, w) + topk: Number of top blocks to attend to + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_block_num: Number of text blocks + add_text_mask: Whether to add text mask + threshold: Threshold for similarity sampling + lambda_: Weight factor for importance sampling + mask_share_within_head: Whether to share mask across heads + q_block_avg_pool: Whether to apply average pooling to query blocks + adaptive_pool: Adaptive pooling size + sampling_type: Type of sampling ("similarity" or "importance") + + Returns: + moba_3d_mask: 3D attention mask with shape (num_heads, block_num, block_num) + """ + seq_len = q.size(2) + block_size = np.prod(tile_thw) + block_num = int(seq_len / block_size) + + image_shape = canvas_thw + block_shape = tile_thw + batch_size, num_heads, seq_len, head_dim = k.shape + num_blocks_t = np.ceil(image_shape[0] / block_shape[0]).astype(int) + num_blocks_h = np.ceil(image_shape[1] / block_shape[1]).astype(int) + num_blocks_w = np.ceil(image_shape[2] / block_shape[2]).astype(int) + + def get_block_avg_feat(x: torch.Tensor, adaptive_pool=None, pooling_type="avg") -> torch.Tensor: + x_block_means = x.view(batch_size, num_heads, num_blocks_t, num_blocks_h, num_blocks_w, + block_shape[0], block_shape[1], block_shape[2], head_dim) + if adaptive_pool is not None: + x_block_means = x_block_means.view(-1, block_shape[0], block_shape[1], block_shape[2], head_dim) + x_block_means = x_block_means.permute(0, 4, 1, 2, 3) + if pooling_type == "avg": + global_avg = torch.nn.functional.adaptive_avg_pool3d(x_block_means, adaptive_pool) + global_avg = global_avg.permute(0, 2, 3, 4, 1) + global_avg = global_avg.reshape(batch_size, num_heads, -1, head_dim * adaptive_pool[0] * adaptive_pool[1] * adaptive_pool[2]) + x_block_means = global_avg + elif pooling_type == "max": + max_pool = torch.nn.functional.adaptive_max_pool3d(x_block_means, adaptive_pool) + max_pool = max_pool.permute(0, 2, 3, 4, 1) + max_pool = max_pool.reshape(batch_size, num_heads, -1, head_dim * adaptive_pool[0] * adaptive_pool[1] * adaptive_pool[2]) + x_block_means = max_pool + elif pooling_type == "mix": + global_avg = torch.nn.functional.adaptive_avg_pool3d(x_block_means, (1, 1, 1)) + global_avg = global_avg / (global_avg.norm(dim=1, keepdim=True) + 1e-8) + global_avg = global_avg.permute(0, 2, 3, 4, 1) + global_avg = global_avg.reshape(batch_size, num_heads, -1, head_dim) + + max_pool = torch.nn.functional.adaptive_max_pool3d(x_block_means, adaptive_pool) + max_pool = max_pool / (max_pool.norm(dim=1, keepdim=True) + 1e-8) + max_pool = max_pool.permute(0, 2, 3, 4, 1) + max_pool = max_pool.reshape(batch_size, num_heads, -1, head_dim * adaptive_pool[0] * adaptive_pool[1] * adaptive_pool[2]) + x_block_means = torch.cat([max_pool, global_avg], dim=-1) + else: + raise ValueError(f"pooling_type={pooling_type} is not Supported") + else: + x_block_means = x_block_means.mean(dim=(-2, -3, -4)).view(batch_size, num_heads, -1, head_dim) + return x_block_means + + if (sampling_type == "similarity" and threshold > 0.0) and adaptive_pool is None: + adaptive_pool = (2,2,2) + + k_block_means = get_block_avg_feat(k, adaptive_pool) + + if q_block_avg_pool: + q = get_block_avg_feat(q, adaptive_pool) + q = q.type(torch.float32) + k_block_means = k_block_means.type(torch.float32) # float logit for better gate logit perception + + if mask_share_within_head: + q = q.mean(dim=1, keepdim=True) + k_block_means = k_block_means.mean(dim=1, keepdim=True) + + if sampling_type == "similarity": + top_block_indices = similarity_sampling(q, k_block_means, topk, threshold, + block_num=block_num, + adaptive_pool=adaptive_pool) + elif sampling_type == "importance": + top_block_indices = importance_sampling(q, + k_block_means, + topk, + threshold, + lambda_=lambda_, + adaptive_pool=adaptive_pool) + else: + raise NotImplementedError(f"sampling_type={sampling_type} is not Supported") + + q = q.type_as(k) + # Keep head dimension to dynamically build mask + assert top_block_indices.size(0) == 1, "top_block_indices batch size must be 1" + top_block_indices = top_block_indices.squeeze(0) # Remove batch dimension [h, s_block, topk] + + # Create 3D mask template (heads, block_num, block_num) + gate_idx_mask = torch.zeros( + (top_block_indices.size(0), block_num, block_num), + dtype=torch.bool, + device=q.device + ) + + # Iterate over each attention head using dim=0 + for head_idx in range(top_block_indices.size(0)): + # Build index matrix for each head + head_mask = torch.zeros( + (block_num, block_num), + dtype=torch.bool, + device=q.device + ) + # Fill indices for current head + head_mask.scatter_( + dim=-1, + index=top_block_indices[head_idx], + value=True + ) + gate_idx_mask[head_idx] = head_mask + + if text_block_num > 0: + pad_block_num = block_num + text_block_num + moba_3d_mask = torch.full( + (gate_idx_mask.size(0), pad_block_num, pad_block_num), + False, + dtype=torch.bool, + device=gate_idx_mask.device + ) + moba_3d_mask[:, :block_num, :block_num] = gate_idx_mask + if add_text_mask: + moba_3d_mask[:, :, -text_block_num:] = True + moba_3d_mask[:, -text_block_num:, :] = True + else: + moba_3d_mask = gate_idx_mask + + return moba_3d_mask + +@lru_cache(maxsize=4096) +def create_sta_3d_mask_optimize(canvas_thw, tile_thw, kernel_thw): + r"""Create optimized STA (Spatio-Temporal Attention) 3D mask using vectorized operations. + + Args: + canvas_thw: String representation of canvas dimensions "t_h_w" + tile_thw: String representation of tile dimensions "t_h_w" + kernel_thw: String representation of kernel dimensions "t_h_w" + + Returns: + block_mask: Boolean mask tensor with shape (block_num, block_num) + """ + canvas_thw = tuple(map(int, canvas_thw.split('_'))) + seq_len = np.prod(canvas_thw) + tile_thw = tuple(map(int, tile_thw.split('_'))) + kernel_thw = tuple(map(int, kernel_thw.split('_'))) + + kernel_t, kernel_h, kernel_w = kernel_thw + + block_size = np.prod(tile_thw) + block_num = int(seq_len / block_size) + + block_mask = np.full((block_num, block_num), False, dtype=bool) + tile_thw_num=(canvas_thw[0] // tile_thw[0], canvas_thw[1] // tile_thw[1], canvas_thw[2] // tile_thw[2]) + + i_indices = np.arange(block_num) + j_indices = np.arange(block_num) + + i_grid, j_grid = np.meshgrid(i_indices, j_indices, indexing='ij') + + q_t_tile = i_grid // (tile_thw_num[1] * tile_thw_num[2]) + q_h_tile = (i_grid % (tile_thw_num[1] * tile_thw_num[2])) // tile_thw_num[2] + q_w_tile = i_grid % tile_thw_num[2] + + kv_t_tile = j_grid // (tile_thw_num[1] * tile_thw_num[2]) + kv_h_tile = (j_grid % (tile_thw_num[1] * tile_thw_num[2])) // tile_thw_num[2] + kv_w_tile = j_grid % tile_thw_num[2] + + kernel_center_t = np.clip(q_t_tile, kernel_t // 2, (tile_thw_num[0] - 1) - kernel_t // 2) + kernel_center_h = np.clip(q_h_tile, kernel_h // 2, (tile_thw_num[1] - 1) - kernel_h // 2) + kernel_center_w = np.clip(q_w_tile, kernel_w // 2, (tile_thw_num[2] - 1) - kernel_w // 2) + + time_mask = np.abs(kernel_center_t - kv_t_tile) <= kernel_t // 2 + hori_mask = np.abs(kernel_center_h - kv_h_tile) <= kernel_h // 2 + vert_mask = np.abs(kernel_center_w - kv_w_tile) <= kernel_w // 2 + + block_mask = time_mask & hori_mask & vert_mask + + block_mask = torch.tensor(block_mask, dtype=torch.bool) + return block_mask + +@torch.no_grad() +def create_sta_3d_mask(canvas_thw, tile_thw, kernel_thw, text_block_num=0): + r"""Create STA (Spatio-Temporal Attention) 3D mask. + + Args: + canvas_thw: Canvas dimensions (t, h, w) + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_block_num: Number of text blocks to pad + + Returns: + sta_mask: Boolean mask tensor with optional text block padding + """ + block_mask = create_sta_3d_mask_optimize("_".join([str(x) for x in canvas_thw]), "_".join([str(x) for x in tile_thw]), "_".join([str(x) for x in kernel_thw])) + sta_mask = None + block_num = block_mask.size(0) + if text_block_num > 0: + pad_block_num = block_num + text_block_num + sta_mask = torch.full( + (pad_block_num, pad_block_num), + False, + dtype=torch.bool, + ) + sta_mask[:block_num, :block_num] = block_mask + sta_mask[:, -text_block_num:] = True + sta_mask[-text_block_num:, :] = True + else: + sta_mask = block_mask + return sta_mask + +@torch.no_grad() +def create_ssta_3d_mask(q, + k, + canvas_thw, + topk, + tile_thw, + kernel_thw, + text_block_num=0, + threshold=0.0, + lambda_=None, + text_mask=None, + mask_share_within_head=True, + adaptive_pool=None, + sampling_type=None): + r"""Create SSTA (Sparse Spatio-Temporal Attention) 3D mask combining STA and MOBA masks. + + Args: + q: Query tensor + k: Key tensor + canvas_thw: Canvas dimensions (t, h, w) + topk: Number of top blocks to attend to + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_block_num: Number of text blocks + threshold: Threshold for similarity sampling + lambda_: Weight factor for importance sampling + text_mask: Optional text mask tensor + mask_share_within_head: Whether to share mask across heads + adaptive_pool: Adaptive pooling size + sampling_type: Type of sampling ("similarity" or "importance") + + Returns: + ssta_3d_mask: Combined 3D attention mask + """ + sta_3d_mask = create_sta_3d_mask(canvas_thw, tile_thw, kernel_thw, text_block_num) + sta_3d_mask = sta_3d_mask.to(q.device) + + moba_3d_mask = create_moba_3d_mask(q, k, canvas_thw, topk, tile_thw, kernel_thw, text_block_num, + threshold=threshold, + lambda_=lambda_, + mask_share_within_head=mask_share_within_head, + adaptive_pool=adaptive_pool, + sampling_type=sampling_type) + ssta_3d_mask = torch.logical_or(sta_3d_mask.unsqueeze(0), moba_3d_mask) + assert len(ssta_3d_mask.size()) == 3, "ssta_3d_mask should be 3D" + + # set text block mask + if text_mask is not None: + block_size = np.prod(tile_thw) + seq_len = q.size(2) + block_num = int(seq_len / block_size) + text_mask_index = torch.ceil(text_mask.sum() / block_size).long() + text_mask_index = text_mask_index.clamp(min=1).item() + assert ssta_3d_mask.shape[-1] == ssta_3d_mask.shape[-2], "ssta_3d_mask should be square" + + pad_start_index = block_num + text_mask_index + ssta_3d_mask[:, pad_start_index:, :] = False + ssta_3d_mask[:, :, pad_start_index:] = False + eye_mask = torch.eye(ssta_3d_mask.shape[1]-pad_start_index, dtype=torch.bool, device=ssta_3d_mask.device).unsqueeze(0) + ssta_3d_mask[:, pad_start_index:, pad_start_index:] = ssta_3d_mask[:, pad_start_index:, pad_start_index:] | eye_mask + return ssta_3d_mask + +def ssta_3d_attention(all_q, + all_k, + all_v, + canvas_thw, + topk=1, + tile_thw=(6, 8, 8), + kernel_thw=(1, 1, 1), + text_len=0, + sparse_type='ssta', + threshold=0.0, + lambda_=None, + pad_type="zero", + text_mask=None, + mask_share_within_head=True, + sampling_type=None, + adaptive_pool=None): + r"""Sparse Spatio-Temporal Attention (SSTA) 3D attention mechanism. + + Args: + all_q: Query tensor with shape (B, H, S, D) + all_k: Key tensor with shape (B, H, S, D) + all_v: Value tensor with shape (B, H, S, D) + canvas_thw: Canvas dimensions (t, h, w) + topk: Number of top blocks to attend to + tile_thw: Tile dimensions + kernel_thw: Kernel dimensions + text_len: Length of text sequence + sparse_type: Type of sparse attention ('sta', 'block_attn', or 'ssta') + threshold: Threshold for similarity sampling + lambda_: Weight factor for importance sampling + pad_type: Padding type ("zero" or "repeat") + text_mask: Optional text mask tensor + mask_share_within_head: Whether to share mask across heads + sampling_type: Type of sampling ("similarity" or "importance") + adaptive_pool: Adaptive pooling size + + Returns: + tuple: (output tensor, sparse_ratio) + - output: Attention output with shape (B, H, S, D) + - sparse_ratio: Ratio of non-zero attention weights + """ + try: + from flex_block_attn import flex_block_attn_func + except ImportError as e: + raise Exception("Could not load flex-block-attn. Please install the flex-block-attn package. You can install it via 'pip install flex-block-attn'.") from e + assert pad_type in ["zero", "repeat"] + assert sampling_type in ["similarity", "importance"] + assert (lambda_ is not None and sampling_type=="importance") or sampling_type == "similarity" + + if text_len > 0: + image_q = all_q[:, :, :-text_len, :] + image_k = all_k[:, :, :-text_len, :] + image_v = all_v[:, :, :-text_len, :] + + text_q = all_q[:, :, -text_len:, :] + text_k = all_k[:, :, -text_len:, :] + text_v = all_v[:, :, -text_len:, :] + else: + image_q = all_q + image_k = all_k + image_v = all_v + + b, hd, s, d = image_q.shape + t, h, w = canvas_thw + assert t * h * w == s, f"t:{t} * h:{h} * w:{w} != s:{s}" + tile_t, tile_h, tile_w = tile_thw + block_size = np.prod(tile_thw) + + need_pad = False + if t % tile_t != 0 or h % tile_h !=0 or w % tile_w != 0: + need_pad = True + pad_image_q = image_q.reshape(b, hd, t, h, w, d) + pad_image_k = image_k.reshape(b, hd, t, h, w, d) + pad_image_v = image_v.reshape(b, hd, t, h, w, d) + + pad_t = 0 if t % tile_t == 0 else tile_t - t % tile_t + if pad_t > 0: + t = t + pad_t + repeat_q = pad_image_q[:, :, -1:, :, :, :].expand(-1, -1, pad_t, -1, -1, -1) + repeat_k = pad_image_k[:, :, -1:, :, :, :].expand(-1, -1, pad_t, -1, -1, -1) + repeat_v = pad_image_v[:, :, -1:, :, :, :].expand(-1, -1, pad_t, -1, -1, -1) + if pad_type == "zero": + repeat_q = torch.zeros_like(repeat_q) + repeat_k = torch.zeros_like(repeat_k) + repeat_v = torch.zeros_like(repeat_v) + pad_image_q = torch.cat([pad_image_q, repeat_q], dim=2) + pad_image_k = torch.cat([pad_image_k, repeat_k], dim=2) + pad_image_v = torch.cat([pad_image_v, repeat_v], dim=2) + + pad_h = 0 if h % tile_h == 0 else tile_h - h % tile_h + if pad_h > 0: + h = h + pad_h + repeat_q = pad_image_q[:, :, :, -1:, :, :].expand(-1, -1, -1, pad_h, -1, -1) + repeat_k = pad_image_k[:, :, :, -1:, :, :].expand(-1, -1, -1, pad_h, -1, -1) + repeat_v = pad_image_v[:, :, :, -1:, :, :].expand(-1, -1, -1, pad_h, -1, -1) + if pad_type == "zero": + repeat_q = torch.zeros_like(repeat_q) + repeat_k = torch.zeros_like(repeat_k) + repeat_v = torch.zeros_like(repeat_v) + pad_image_q = torch.cat([pad_image_q, repeat_q], dim=3) + pad_image_k = torch.cat([pad_image_k, repeat_k], dim=3) + pad_image_v = torch.cat([pad_image_v, repeat_v], dim=3) + + pad_w = 0 if w % tile_w == 0 else tile_w - w % tile_w + if pad_w > 0: + w = w + pad_w + repeat_q = pad_image_q[:, :, :, :, -1:, :].expand(-1, -1, -1, -1, pad_w, -1) + repeat_k = pad_image_k[:, :, :, :, -1:, :].expand(-1, -1, -1, -1, pad_w, -1) + repeat_v = pad_image_v[:, :, :, :, -1:, :].expand(-1, -1, -1, -1, pad_w, -1) + if pad_type == "zero": + repeat_q = torch.zeros_like(repeat_q) + repeat_k = torch.zeros_like(repeat_k) + repeat_v = torch.zeros_like(repeat_v) + pad_image_q = torch.cat([pad_image_q, repeat_q], dim=4) + pad_image_k = torch.cat([pad_image_k, repeat_k], dim=4) + pad_image_v = torch.cat([pad_image_v, repeat_v], dim=4) + + image_q = pad_image_q.reshape(b, hd, -1, d) + image_k = pad_image_k.reshape(b, hd, -1, d) + image_v = pad_image_v.reshape(b, hd, -1, d) + + canvas_thw = (t, h, w) + + need_pad_text = False + text_block_num = math.ceil(text_len / block_size) + text_target_size = text_block_num * block_size + if text_len % block_size > 0: + need_pad_text = True + text_pad_size = text_target_size - text_len + + pad_text_q = text_q[:, :, -1, :].unsqueeze(2).expand(-1, -1, text_pad_size, -1) + pad_text_k = text_k[:, :, -1, :].unsqueeze(2).expand(-1, -1, text_pad_size, -1) + pad_text_v = text_v[:, :, -1, :].unsqueeze(2).expand(-1, -1, text_pad_size, -1) + + text_q = torch.cat([text_q, pad_text_q], dim=2) + text_k = torch.cat([text_k, pad_text_k], dim=2) + text_v = torch.cat([text_v, pad_text_v], dim=2) + + image_q = tile(image_q, canvas_thw, tile_thw) + image_k = tile(image_k, canvas_thw, tile_thw) + image_v = tile(image_v, canvas_thw, tile_thw) + + if text_len > 0: + q = torch.cat([image_q, text_q], dim=2) + k = torch.cat([image_k, text_k], dim=2) + v = torch.cat([image_v, text_v], dim=2) + else: + q = image_q + k = image_k + v = image_v + + if sparse_type == 'sta': + assert text_mask is None, "text_mask do not support in sta sparse_type" + block_mask = create_sta_3d_mask(canvas_thw, tile_thw, kernel_thw, text_block_num).to(q.device) + o = flex_block_attn_func(q, k, v, block_size, block_size, block_mask) + elif sparse_type == 'block_attn': + image_q_list = torch.split(image_q, 1, dim=0) + image_k_list = torch.split(image_k, 1, dim=0) + + mask_list = [] + for i in range(b): + block_mask = create_moba_3d_mask(image_q_list[i], + image_k_list[i], + canvas_thw=canvas_thw, + topk=topk, + tile_thw=tile_thw, + kernel_thw=kernel_thw, + text_block_num=text_block_num, + add_text_mask=True, + lambda_=lambda_, + threshold=threshold, + mask_share_within_head=mask_share_within_head, + adaptive_pool=adaptive_pool, + sampling_type=sampling_type) + mask_list.append(block_mask) + block_mask = torch.stack(mask_list, dim=0) + o = flex_block_attn_func(q, k, v, block_size, block_size, block_mask) + + elif sparse_type == 'ssta': + image_q_list = torch.split(image_q, 1, dim=0) + image_k_list = torch.split(image_k, 1, dim=0) + + mask_list = [] + for i in range(b): + block_mask = create_ssta_3d_mask( + image_q_list[i], + image_k_list[i], + canvas_thw=canvas_thw, + tile_thw=tile_thw, + kernel_thw=kernel_thw, + text_block_num=text_block_num, + topk=topk, + threshold=threshold, + lambda_=lambda_, + text_mask=text_mask[i] if text_mask is not None else None, + mask_share_within_head=mask_share_within_head, + adaptive_pool=adaptive_pool, + sampling_type=sampling_type + ) + mask_list.append(block_mask) + block_mask = torch.stack(mask_list, dim=0) + if mask_share_within_head: + block_mask = block_mask.unsqueeze(1) # [b, 1, s_block, s_block] + o = flex_block_attn_func(q, k, v, block_size, block_size, block_mask) + else: + raise Exception(f"unsupported sparse_type:{sparse_type}") + sparse_ratio = block_mask.float().mean().cpu().item() + + if text_len > 0: + image_o = o[:, :, :-text_target_size, :] + if need_pad_text: + text_o = o[:, :, -text_target_size : -text_pad_size, :] + else: + text_o = o[:, :, -text_target_size:, :] + else: + image_o = o + + image_o = untile(image_o, canvas_thw, tile_thw) + + if need_pad: + # Remove padding from output + unpad_image_o = image_o.reshape(b, hd, t, h, w, d) + if pad_t > 0: + unpad_image_o = unpad_image_o[:, :, :-pad_t, :, :, :] + if pad_h > 0: + unpad_image_o = unpad_image_o[:, :, :, :-pad_h, :, :] + if pad_w > 0: + unpad_image_o = unpad_image_o[:, :, :, :, :-pad_w, :] + image_o = unpad_image_o.reshape(b, hd, -1, d) + + if text_len > 0: + o = torch.cat([image_o, text_o], dim=2) + else: + o = image_o + + return o, sparse_ratio diff --git a/unison/models/transformers/modules/token_refiner.py b/unison/models/transformers/modules/token_refiner.py new file mode 100644 index 0000000000000000000000000000000000000000..69307315e6ec3f5f527f21ff7c633f0a2486ece3 --- /dev/null +++ b/unison/models/transformers/modules/token_refiner.py @@ -0,0 +1,283 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Optional + +import torch +import torch.nn as nn +from einops import rearrange + +from .activation_layers import get_activation_layer +from .embed_layers import TextProjection, TimestepEmbedder +from .mlp_layers import MLP +from .modulate_layers import apply_gate +from .norm_layers import get_norm_layer +from .attention import attention + + + + +class IndividualTokenRefinerBlock(nn.Module): + """ + A single block for token refinement with self-attention and MLP. + + Args: + hidden_size: Hidden dimension size. + heads_num: Number of attention heads. + mlp_width_ratio: Expansion ratio for MLP hidden size. + mlp_drop_rate: Dropout rate for MLP. + act_type: Activation function type. + qk_norm: Whether to use QK normalization. + qk_norm_type: Type of QK normalization. + qkv_bias: Whether to use bias in QKV projections. + dtype: Optional torch dtype. + device: Optional torch device. + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs) + self.self_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.self_attn_q_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.self_attn_k_norm = ( + qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity() + ) + self.self_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs) + act_layer = get_activation_layer(act_type) + self.mlp = MLP( + in_channels=hidden_size, + hidden_channels=mlp_hidden_dim, + act_layer=act_layer, + drop=mlp_drop_rate, + **factory_kwargs, + ) + + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, # timestep_aware_representations + context_aware_representations + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Forward pass for IndividualTokenRefinerBlock. + + Args: + x: Input tensor of shape [B, L, C]. + c: Conditioning tensor of shape [B, C]. + attn_mask: Optional attention mask of shape [B, L]. + + Returns: + Refined tensor of shape [B, L, C]. + """ + gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1) + norm_x = self.norm1(x) + qkv = self.self_attn_qkv(norm_x) + q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) + q = self.self_attn_q_norm(q).to(v) + k = self.self_attn_k_norm(k).to(v) + attn = attention(q, k, v, attn_mask=attn_mask) + x = x + apply_gate(self.self_attn_proj(attn), gate_msa) + x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp) + return x + + +class IndividualTokenRefiner(nn.Module): + """ + Stacks multiple IndividualTokenRefinerBlock modules. + + Args: + hidden_size: Hidden dimension size. + heads_num: Number of attention heads. + depth: Number of blocks. + mlp_width_ratio: Expansion ratio for MLP hidden size. + mlp_drop_rate: Dropout rate for MLP. + act_type: Activation function type. + qk_norm: Whether to use QK normalization. + qk_norm_type: Type of QK normalization. + qkv_bias: Whether to use bias in QKV projections. + dtype: Optional torch dtype. + device: Optional torch device. + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + depth: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.blocks = nn.ModuleList( + [ + IndividualTokenRefinerBlock( + hidden_size=hidden_size, + heads_num=heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) + for _ in range(depth) + ] + ) + + def forward( + self, + x: torch.Tensor, + c: torch.LongTensor, + mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Forward pass for IndividualTokenRefiner. + + Args: + x: Input tensor of shape [B, L, C]. + c: Conditioning tensor of shape [B, C]. + mask: Optional mask tensor of shape [B, L]. + + Returns: + Refined tensor of shape [B, L, C]. + """ + if mask is not None: + mask = mask.clone().bool() + mask[:, 0] = True # Prevent attention weights from becoming NaN + for block in self.blocks: + x = block(x, c, mask) + return x + + +class SingleTokenRefiner(nn.Module): + """ + Single token refiner block for LLM text embedding refinement. + + Args: + in_channels: Input feature dimension. + hidden_size: Hidden dimension size. + heads_num: Number of attention heads. + depth: Number of blocks. + mlp_width_ratio: Expansion ratio for MLP hidden size. + mlp_drop_rate: Dropout rate for MLP. + act_type: Activation function type. + qk_norm: Whether to use QK normalization. + qk_norm_type: Type of QK normalization. + qkv_bias: Whether to use bias in QKV projections. + dtype: Optional torch dtype. + device: Optional torch device. + """ + + def __init__( + self, + in_channels: int, + hidden_size: int, + heads_num: int, + depth: int, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True, **factory_kwargs) + act_layer = get_activation_layer(act_type) + self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs) + self.c_embedder = TextProjection(in_channels, hidden_size, act_layer, **factory_kwargs) + self.individual_token_refiner = IndividualTokenRefiner( + hidden_size=hidden_size, + heads_num=heads_num, + depth=depth, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) + + def forward( + self, + x: torch.Tensor, + t: torch.LongTensor, + mask: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + """ + Forward pass for SingleTokenRefiner. + + Args: + x: Input tensor of shape [B, L, in_channels]. + t: Timestep tensor of shape [B]. + mask: Optional mask tensor of shape [B, L]. + + Returns: + Refined tensor of shape [B, L, hidden_size]. + """ + timestep_aware_representations = self.t_embedder(t) + if mask is None: + context_aware_representations = x.mean(dim=1) + else: + mask_float = mask.float().unsqueeze(-1) # [B, L, 1] + context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1) + context_aware_representations = self.c_embedder(context_aware_representations) + c = timestep_aware_representations + context_aware_representations + x = self.input_embedder(x) + x = self.individual_token_refiner(x, c, mask) + return x diff --git a/unison/models/transformers/modules/upsample.py b/unison/models/transformers/modules/upsample.py new file mode 100644 index 0000000000000000000000000000000000000000..b931aea8a0a167bcbd94f4d76b23cec3598ea7ce --- /dev/null +++ b/unison/models/transformers/modules/upsample.py @@ -0,0 +1,169 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor +from diffusers.models import ModelMixin +from diffusers.configuration_utils import ConfigMixin, register_to_config + +from unison.models.autoencoders.hunyuanvideo_15_vae import ( + CausalConv3d, + ResnetBlock, + RMS_norm, + forward_with_checkpointing, + swish, +) + + +class UpsamplerType(Enum): + LEARNED = "learned" + FIXED = "fixed" + NONE = "none" + LEARNED_FIXED = "learned_fixed" + + +@dataclass +class UpsamplerConfig: + load_from: str + enable: bool = False + hidden_channels: int = 128 + num_blocks: int = 16 + model_type: UpsamplerType = UpsamplerType.NONE + version: str = "720p" + + +class SRResidualCausalBlock3D(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.block = nn.Sequential( + CausalConv3d(channels, channels, kernel_size=3), + nn.SiLU(inplace=True), + CausalConv3d(channels, channels, kernel_size=3), + nn.SiLU(inplace=True), + CausalConv3d(channels, channels, kernel_size=3), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.block(x) + + +class SRTo720pUpsampler(ModelMixin, ConfigMixin): + + @register_to_config + def __init__( + self, + in_channels: int, + out_channels: int, + hidden_channels: int | None = None, + num_blocks: int = 6, + global_residual: bool = False, + ): + super().__init__() + if hidden_channels is None: + hidden_channels = 64 + self.in_conv = CausalConv3d(in_channels, hidden_channels, kernel_size=3) + self.blocks = nn.ModuleList([SRResidualCausalBlock3D(hidden_channels) for _ in range(num_blocks)]) + self.out_conv = CausalConv3d(hidden_channels, out_channels, kernel_size=3) + self.global_residual = bool(global_residual) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + y = self.in_conv(x) + for blk in self.blocks: + y = blk(y) + y = self.out_conv(y) + if self.global_residual and (y.shape == residual.shape): + y = y + residual + return y + + +class SRTo1080pUpsampler(ModelMixin, ConfigMixin): + + @register_to_config + def __init__( + self, + z_channels: int, + out_channels: int, + block_out_channels: tuple[int, ...], + num_res_blocks: int = 2, + is_residual: bool = False, + ): + super().__init__() + self.num_res_blocks = num_res_blocks + self.block_out_channels = block_out_channels + self.z_channels = z_channels + + block_in = block_out_channels[0] + self.conv_in = CausalConv3d(z_channels, block_in, kernel_size=3) + + self.up = nn.ModuleList() + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + + self.up.append(up) + + self.norm_out = RMS_norm(block_in, images=False) + self.conv_out = CausalConv3d(block_in, out_channels, kernel_size=3) + + self.gradient_checkpointing = False + self.is_residual = is_residual + + def forward(self, z: Tensor, target_shape: Sequence[int] = None) -> Tensor: + """ + Args: + z: (B, C, T, H, W) + target_shape: (H, W) + """ + use_checkpointing = bool(self.training and self.gradient_checkpointing) + if target_shape is not None and z.shape[-2:] != target_shape: + bsz = z.shape[0] + z = rearrange(z, "b c f h w -> (b f) c h w") + z = F.interpolate(z, size=target_shape, mode="bilinear", align_corners=False) + z = rearrange(z, "(b f) c h w -> b c f h w", b=bsz) + + # z to block_in + repeats = self.block_out_channels[0] // (self.z_channels) + h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1) + + # upsampling + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks + 1): + h = forward_with_checkpointing( + self.up[i_level].block[i_block], + h, + use_checkpointing=use_checkpointing, + ) + if hasattr(self.up[i_level], "upsample"): + h = forward_with_checkpointing(self.up[i_level].upsample, h, use_checkpointing=use_checkpointing) + + # end + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h diff --git a/unison/pipelines/__init__.py b/unison/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a98854003ddedbd3c39e7f7b493a2e7451f32957 --- /dev/null +++ b/unison/pipelines/__init__.py @@ -0,0 +1 @@ +"""UNISON training and inference pipelines.""" diff --git a/unison/pipelines/infer.py b/unison/pipelines/infer.py new file mode 100644 index 0000000000000000000000000000000000000000..bbb768f6c4860de32c5d7fe795191aa6db957443 --- /dev/null +++ b/unison/pipelines/infer.py @@ -0,0 +1,607 @@ +#!/usr/bin/env python +# coding: utf-8 +""" +Inference pipeline for UNISON with channel-cat + Qwen2.5-Omni deep fusion. + +Channel-Cat approach: +- ref audio is concatenated into the source channel (not as separate ref_latents). +- For zero-shot TTS: source = [ref_latent | zeros], mask = [2 | 0]. + The model sees ref via channel concatenation and learns to generate target audio + conditioned on the ref portion. +- Optional inpainting trick: at each ODE step, clamp the ref region in latent space + so the ref portion stays faithful to the original encoding. + +Supports three task modes: generation / editing / zeroshotts. +""" + +import os +import re +from typing import List, Optional + +import torch +import torch.nn.functional as F +import soundfile as sf +import torchaudio +import torchaudio.functional as AF +from diffusers import FlowMatchEulerDiscreteScheduler +from diffusers.training_utils import EMAModel +from safetensors.torch import load_file as safe_load_file +from tqdm.auto import tqdm +import yaml +import logging + +from unison.models.transformers.backbone import UnisonBackbone +from unison.models.text_encoders.omni_encoder import ( + QwenOmniThinkerExtractor, + resolve_omni_model_path, +) + + +# ------------------------- +# Defaults / constants +# ------------------------- +MAX_AUDIO_DURATION = 10.0 +AUDIO_FPS = 31.25 +TARGET_FRAMES = int(MAX_AUDIO_DURATION * AUDIO_FPS) +DEFAULT_TARGET_SAMPLE_RATE = 16000 +REF_DURATION = 3.0 # seconds +TRAINING_TARGET_RMS = 0.1 # must match GPUAudioProcessorChannelCat.target_rms + + +# ------------------------- +# Logging filter +# ------------------------- +class _OmniFilter(logging.Filter): + def filter(self, record): + if record.levelno != logging.WARNING: + return True + return "System prompt modified" not in record.getMessage() + +logging.getLogger().addFilter(_OmniFilter()) + + +# ------------------------- +# Model initializers +# ------------------------- +def init_text_hidden_extractor( + text_encoder_type: str, + omni_model_path: str, + text_encoder_model_path: Optional[str], + dit_depth: int, + device, + dtype=torch.bfloat16, + omni_last_layer_idx=-1, +): + """Load Qwen2.5-Omni as the frozen text-feature extractor.""" + omni_path = resolve_omni_model_path( + text_encoder_type, + text_encoder_model_path or omni_model_path, + ) + print(f" Omni preset={text_encoder_type} path={omni_path}") + extractor = QwenOmniThinkerExtractor( + model_path=omni_path, + dit_depth=dit_depth, + select_mode="interval", + device=device, + dtype=dtype, + omni_last_layer_idx=omni_last_layer_idx, + ) + extractor.eval() + extractor.requires_grad_(False) + return extractor + + +def sync_omni_dim_with_text_encoder(model_config, text_extractor): + """Align model_config['omni_dim'] with loaded text encoder hidden_size.""" + backbone = getattr(text_extractor, "text_backbone", None) + if backbone is None: + print("[WARN] Text encoder has no text_backbone; skip omni_dim sync.") + return + cfg = getattr(backbone, "config", None) + hidden = getattr(cfg, "hidden_size", None) if cfg is not None else None + if hidden is None: + print("[WARN] Could not read text encoder hidden_size; skip omni_dim sync.") + return + prev = model_config.get("omni_dim") + if prev != hidden: + print( + f"[WARN] model_config omni_dim={prev} != text encoder hidden_size={hidden}; " + f"updating omni_dim to {hidden}." + ) + model_config["omni_dim"] = hidden + + +# ------------------------- +# Sampling (generation / editing / zero-shot TTS via channel-cat) +# ------------------------- +@torch.no_grad() +def sample_latents( + model, + scheduler, + omni_extractor, + prompts: List[str], + source_latents=None, + masks=None, + num_inference_steps: int = 30, + guidance_scale: float = 3.5, + device=torch.device("cuda"), + target_frames: int = None, + latent_shape: tuple = None, + inpaint_mask=None, + inpaint_source=None, + noise_init=None, +): + """ + Sample latents with channel-cat conditioning. + + Audio latent shape: [B, C, T] (C = VAE latent channels) + Generation: source_latents=None, masks=None → zeros are used + Editing: source_latents=[B,C,T], masks=[B,1,T] (mask value 1) + ZS-TTS: source_latents=[B,C,T] (ref | zeros), masks=[B,1,T] (2=ref region, 0=target) + + Optional inpainting trick (ZS-TTS only): + inpaint_mask: [B, C, T] bool, True at ref positions + inpaint_source: [B, C, T] clean ref latent (already scaled by vae_scale_factor) + noise_init: [B, C, T] initial noise tensor + At each ODE step the ref region is clamped to keep it faithful to the encoding. + """ + model.eval() + batch_size = len(prompts) + is_edit = source_latents is not None and masks is not None + C = model.config.in_channels + is_5d = latent_shape is not None and len(latent_shape) == 3 + + # Omni features + omni_emb_cond, omni_mask_cond, omni_last_emb_cond = omni_extractor(prompts) + omni_emb_cond = [e.to(device) for e in omni_emb_cond] + omni_mask_cond = omni_mask_cond.to(device) + omni_last_emb_cond = omni_last_emb_cond.to(device) + + if is_edit: + source_latents = source_latents.to(device=device, dtype=torch.bfloat16) + masks = masks.to(device=device, dtype=torch.bfloat16) + if noise_init is not None: + noise = noise_init.to(device=device, dtype=torch.bfloat16) + else: + noise = torch.randn_like(source_latents) + x = torch.cat([noise, source_latents, masks], dim=1) + else: + if is_5d: + lt, lh, lw = latent_shape + noise = torch.randn(batch_size, C, lt, lh, lw, device=device, dtype=torch.bfloat16) + zero_src = torch.zeros_like(noise) + zero_mask = torch.zeros(batch_size, 1, lt, lh, lw, device=device, dtype=torch.bfloat16) + else: + T = target_frames if target_frames is not None else TARGET_FRAMES + noise = torch.randn(batch_size, C, T, device=device, dtype=torch.bfloat16) + zero_src = torch.zeros_like(noise) + zero_mask = torch.zeros(batch_size, 1, T, device=device, dtype=torch.bfloat16) + x = torch.cat([noise, zero_src, zero_mask], dim=1) + + # Move inpainting tensors to device + if inpaint_mask is not None: + inpaint_mask = inpaint_mask.to(device=device) + if inpaint_source is not None: + inpaint_source = inpaint_source.to(device=device, dtype=torch.bfloat16) + if noise_init is not None: + noise_init = noise_init.to(device=device, dtype=torch.bfloat16) + + scheduler.set_timesteps(num_inference_steps, device=device) + do_cfg = guidance_scale is not None and guidance_scale > 1.0 + + noise_slice = slice(None, C) + + if do_cfg: + omni_emb_uncond = [torch.zeros_like(e) for e in omni_emb_cond] + omni_mask_uncond = torch.zeros_like(omni_mask_cond) + omni_last_emb_uncond = torch.zeros_like(omni_last_emb_cond) + + timesteps = scheduler.timesteps + for step_idx, t in enumerate(tqdm(timesteps, desc="Sampling (CFG)")): + t_in = (t.float() / scheduler.config.num_train_timesteps).repeat(batch_size).to(device=device, dtype=torch.bfloat16) + + out_unc = model(x=x, t=t_in, duration=None, + omni_emb_list=omni_emb_uncond, omni_last_emb=omni_last_emb_uncond, + omni_mask=omni_mask_uncond) + pred_unc = out_unc[0] + + out_cond = model(x=x, t=t_in, duration=None, + omni_emb_list=omni_emb_cond, omni_last_emb=omni_last_emb_cond, + omni_mask=omni_mask_cond) + pred_cond = out_cond[0] + + pred = pred_unc + guidance_scale * (pred_cond - pred_unc) + x[:, noise_slice] = scheduler.step(model_output=pred, timestep=t, sample=x[:, noise_slice]).prev_sample + + # Inpainting trick: clamp ref region after each ODE step + if inpaint_mask is not None and noise_init is not None and inpaint_source is not None: + # Compute sigma for the NEXT step (after the update we just did) + if step_idx + 1 < len(timesteps): + t_next = timesteps[step_idx + 1] + sigma_next = t_next.float() / scheduler.config.num_train_timesteps + else: + sigma_next = 0.0 + # clamp_val = (1 - sigma_next) * clean + sigma_next * noise + clamp_val = (1.0 - sigma_next) * inpaint_source + sigma_next * noise_init + for dim_idx in range(C): + x[:, dim_idx][inpaint_mask[:, dim_idx]] = clamp_val[:, dim_idx][inpaint_mask[:, dim_idx]] + else: + timesteps = scheduler.timesteps + for step_idx, t in enumerate(tqdm(timesteps, desc="Sampling")): + t_in = (t.float() / scheduler.config.num_train_timesteps).repeat(batch_size).to(device=device, dtype=torch.bfloat16) + out = model(x=x, t=t_in, duration=None, + omni_emb_list=omni_emb_cond, omni_last_emb=omni_last_emb_cond, + omni_mask=omni_mask_cond) + pred = out[0] + + x[:, noise_slice] = scheduler.step(model_output=pred, timestep=t, sample=x[:, noise_slice]).prev_sample + + # Inpainting trick: clamp ref region after each ODE step + if inpaint_mask is not None and noise_init is not None and inpaint_source is not None: + if step_idx + 1 < len(timesteps): + t_next = timesteps[step_idx + 1] + sigma_next = t_next.float() / scheduler.config.num_train_timesteps + else: + sigma_next = 0.0 + clamp_val = (1.0 - sigma_next) * inpaint_source + sigma_next * noise_init + for dim_idx in range(C): + x[:, dim_idx][inpaint_mask[:, dim_idx]] = clamp_val[:, dim_idx][inpaint_mask[:, dim_idx]] + + return x[:, noise_slice] + + +# ------------------------- +# Decode & save +# ------------------------- +@torch.no_grad() +def decode_and_save(audio_vae, latents, durations, output_paths, sample_rate=None): + if sample_rate is None: + sample_rate = DEFAULT_TARGET_SAMPLE_RATE + device = next(audio_vae.parameters()).device + dtype = getattr(audio_vae, 'dtype', torch.bfloat16) + latents = latents.to(device=device, dtype=dtype) + audio = audio_vae.wrapped_decode(latents).cpu() + + for i, path in enumerate(output_paths): + dur = min(float(durations[i]), MAX_AUDIO_DURATION) + max_smp = int(dur * sample_rate) + wav = audio[i] + if wav.dim() > 1: + wav = wav.mean(0) + wav = wav[..., :max_smp] + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + sf.write(path, wav.float().numpy(), sample_rate) + print(f" Saved: {path} ({dur:.1f}s)") + + +@torch.no_grad() +def decode_and_save_full(audio_vae, latents_full, ref_samples, + durations, output_paths, sample_rate=None, + fade_in_ms: float = 20.0, + max_duration: float = None, + ref_wav_raw: torch.Tensor = None, + ref_output_paths=None): + """Decode the FULL latent (ref+target), then crop target in waveform space. + + This avoids boundary artifacts that occur when cropping in latent space + (the decoder's convolutional receptive field needs left-side context). + + The crop point is snapped to the nearest VAE hop boundary so the cut + aligns with a latent frame edge, then a short fade-in suppresses any + residual ref tail bleed from the decoder's receptive field. + + If ref_wav_raw is provided (the truncated ref waveform before VAE encode), + it is saved to ref_output_paths alongside the generated target files. This + lets you verify exactly what the model heard as the reference. + """ + if sample_rate is None: + sample_rate = DEFAULT_TARGET_SAMPLE_RATE + device = next(audio_vae.parameters()).device + dtype = getattr(audio_vae, 'dtype', torch.bfloat16) + latents_full = latents_full.to(device=device, dtype=dtype) + audio_full = audio_vae.wrapped_decode(latents_full).cpu() + + wav_total = audio_full.shape[-1] + lat_total = latents_full.shape[-1] + vae_hop = wav_total // lat_total if lat_total > 0 else 1 + # Snap crop point to nearest VAE hop boundary + crop_start = round(ref_samples / vae_hop) * vae_hop + crop_start_s = crop_start / sample_rate + + save_cap = MAX_AUDIO_DURATION if max_duration is None else float(max_duration) + + for i, path in enumerate(output_paths): + dur = min(float(durations[i]), save_cap) + max_smp = int(dur * sample_rate) + wav = audio_full[i] + if wav.dim() > 1: + wav = wav.mean(0) + target_wav = wav[crop_start: crop_start + max_smp].clone() + fade_len = min(int(fade_in_ms * sample_rate / 1000), target_wav.shape[-1]) + if fade_len > 1: + target_wav[:fade_len] *= torch.linspace(0.0, 1.0, fade_len) + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + sf.write(path, target_wav.float().numpy(), sample_rate) + print(f" Saved (target): {path} ({dur:.1f}s, crop@{crop_start} = {crop_start_s:.3f}s, fade={fade_in_ms}ms)") + + # Optionally save the truncated ref waveform used as conditioning + if ref_wav_raw is not None and ref_output_paths is not None and i < len(ref_output_paths): + ref_path = ref_output_paths[i] + ref_out = ref_wav_raw.squeeze().cpu().float() + if ref_out.dim() == 0: + ref_out = ref_out.unsqueeze(0) + os.makedirs(os.path.dirname(ref_path) or ".", exist_ok=True) + sf.write(ref_path, ref_out.numpy(), sample_rate) + print(f" Saved (ref): {ref_path} ({ref_out.shape[-1] / sample_rate:.3f}s)") + + +# ------------------------- +# Audio utilities +# ------------------------- +def _rms_normalize(wav: torch.Tensor, target_rms: float = TRAINING_TARGET_RMS) -> torch.Tensor: + """Normalize waveform RMS to match training preprocessing.""" + rms = torch.sqrt(torch.mean(wav ** 2)) + if rms > 0: + wav = wav * (target_rms / (rms + 1e-8)) + return wav + + +def load_source_audio(audio_path, target_sr=DEFAULT_TARGET_SAMPLE_RATE, + target_length=None, max_duration=MAX_AUDIO_DURATION, device=None): + device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") + wav, sr = torchaudio.load(audio_path) + if wav.shape[0] > 1: + wav = wav.mean(0, keepdim=True) + if sr != target_sr: + wav = AF.resample(wav, sr, target_sr) + wav = _rms_normalize(wav) + L = target_length or int(max_duration * target_sr) + if wav.shape[-1] > L: + wav = wav[..., :L] + elif wav.shape[-1] < L: + wav = F.pad(wav, (0, L - wav.shape[-1])) + return wav.to(device) + + +_whisper_model = None + +def get_whisper_model(model_size="base"): + global _whisper_model + if _whisper_model is None: + try: + import whisper + except ImportError: + raise ImportError( + "openai-whisper is required for auto-transcription. " + "Install with: pip install openai-whisper" + ) + print(f" Loading Whisper ({model_size}) for auto-transcription...") + _whisper_model = whisper.load_model(model_size) + return _whisper_model + + +def transcribe_ref_audio(wav_tensor, sr, whisper_size="base"): + """Auto-transcribe a (possibly truncated) waveform tensor using Whisper. + + Args: + wav_tensor: [1, T] or [T] waveform on any device / sample rate + sr: sample rate of wav_tensor + """ + model = get_whisper_model(whisper_size) + wav_16k = wav_tensor.squeeze().cpu().float() + if sr != 16000: + wav_16k = AF.resample(wav_16k.unsqueeze(0), sr, 16000).squeeze(0) + audio_np = wav_16k.numpy() + result = model.transcribe(audio_np, fp16=torch.cuda.is_available()) + text = result["text"].strip() + print(f" [ASR] ref_text = \"{text}\"") + return text + + +# --- Silence-aware ref cut + tail padding helpers -------------------------- +# +# Why these exist: training does a hard cut at a random sample for ref/target +# split, so the model has learned to *continue the on-going acoustic state* at +# the cut point. At inference, if the user feeds a hard-cut ref that ends mid +# word / mid syllable, the model faithfully extends that ref tail into the +# target generation (ref-tail bleed into output). Two cheap fixes: +# 1. Snap the cut backward to the nearest silent frame (energy-based VAD). +# 2. Append a short silence pad after the ref so the model sees a calm +# starting state for target generation. + +def _find_silence_cut( + wav: torch.Tensor, + sr: int, + target_len: int, + search_back_s: float = 0.5, + win_ms: float = 20.0, + threshold_db: float = -40.0, +) -> int: + """Snap target_len backward to the last silent frame within search_back_s. + + Returns the cut sample index in [target_len - search_back_s*sr, target_len]. + Falls back to target_len if no silent frame is found in the search window. + """ + if wav.dim() == 2: + mono = wav.mean(dim=0) + else: + mono = wav + + total = mono.shape[-1] + if target_len <= 0 or total == 0: + return max(0, target_len) + if total <= target_len: + return total + + win = max(int(win_ms * sr / 1000), 1) + hop = max(win // 2, 1) + search_back = max(int(search_back_s * sr), win) + search_start = max(0, target_len - search_back) + search_end = target_len + if search_end - search_start < win: + return target_len + + seg = mono[search_start:search_end].float() + frames = seg.unfold(-1, win, hop) # [n_frames, win] + rms = frames.pow(2).mean(-1).sqrt() + rms_db = 20.0 * torch.log10(rms.clamp(min=1e-10)) + is_silence = rms_db < threshold_db + if not bool(is_silence.any()): + return target_len + + silence_idx = torch.where(is_silence)[0] + last = int(silence_idx[-1].item()) + cut = search_start + last * hop + win // 2 + return max(1, min(cut, target_len)) + + +# --- Trailing-punctuation strip + ref/target join helpers ------------------ +# +# Whisper transcription always emits punctuation; for a ref hard-cut mid +# sentence that produces things like "Hello there.", which (a) tells the +# model the speaker has finished a thought and (b) double-punctuates when we +# concatenate with target text. Strip the trailing punct and re-insert a +# soft "continue" comma between ref and target. + +_TRAILING_PUNCT_CHARS = ( + " \t\r\n" + ".,;:!?\"'`-—…" + "。，；：！？、…—" + "\u3002\uff0c\uff1b\uff1a\uff01\uff1f\u3001" + "“”‘’《》「」『』" +) +_TRAILING_PUNCT_RE = re.compile( + "[" + re.escape(_TRAILING_PUNCT_CHARS) + "]+$" +) +_CJK_RE = re.compile(r"[\u3400-\u9fff\uff00-\uffef]") + + +def _strip_trailing_punct(text: str) -> str: + if not text: + return "" + return _TRAILING_PUNCT_RE.sub("", text).rstrip() + + +def _has_cjk(text: str) -> bool: + return bool(_CJK_RE.search(text or "")) + + +def join_ref_target_text(ref_text: str, target_text: str) -> str: + """Strip ref's trailing punctuation, then insert a continuation comma + between ref_text and target_text (Chinese 「，」 if either side is CJK, + else English ", ").""" + ref_text = _strip_trailing_punct(ref_text or "") + target_text = (target_text or "").strip() + if not ref_text: + return target_text + if not target_text: + return ref_text + sep = "，" if (_has_cjk(ref_text) or _has_cjk(target_text)) else ", " + return f"{ref_text}{sep}{target_text}" + + +def load_ref_audio( + audio_path, + target_sr=DEFAULT_TARGET_SAMPLE_RATE, + max_ref_duration=None, + device=None, + silence_search_s: float = 0.5, + silence_threshold_db: float = -40.0, + tail_pad_s: float = 0.1, +): + """Load reference audio, optionally truncate to max_ref_duration. + + If silence_search_s > 0 and max_ref_duration is set, the cut sample is + snapped backward to the last silent frame within + [target_len - silence_search_s*sr, target_len]. After RMS-normalizing + the cut waveform, ``tail_pad_s`` of true silence is appended so the + model starts target generation from a calm state instead of continuing + the ref's last formant. + """ + device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") + wav, sr = torchaudio.load(audio_path) + if wav.shape[0] > 1: + wav = wav.mean(0, keepdim=True) + if sr != target_sr: + wav = AF.resample(wav, sr, target_sr) + + if max_ref_duration is not None: + ref_len = int(max_ref_duration * target_sr) + if wav.shape[-1] > ref_len: + if silence_search_s > 0: + cut = _find_silence_cut( + wav, target_sr, ref_len, + search_back_s=silence_search_s, + threshold_db=silence_threshold_db, + ) + cut_dur = cut / target_sr + if cut < ref_len: + print( + f" [ref-cut] snapped {max_ref_duration:.2f}s -> " + f"{cut_dur:.2f}s at silence (search_back={silence_search_s:.2f}s, " + f"thresh={silence_threshold_db:.0f}dB)" + ) + wav = wav[..., :cut] + else: + wav = wav[..., :ref_len] + + wav = _rms_normalize(wav) + + if tail_pad_s > 0: + pad_len = int(tail_pad_s * target_sr) + if pad_len > 0: + wav = F.pad(wav, (0, pad_len)) + + return wav.to(device) + + +def make_edit_mask(waveform_length, device): + return torch.ones(1, 1, waveform_length, device=device, dtype=torch.float32) + + +def downsample_mask(mask_wav, latent_length): + T = mask_wav.shape[-1] + if T != latent_length: + m = F.adaptive_avg_pool1d(mask_wav.float(), latent_length) + return (m > 0.5).float() + return mask_wav.float() + + +def _load_model(args, device, model_config): + mcfg = dict(model_config) + mcfg.pop("omni_last_layer_idx", None) + mcfg["concat_condition"] = True + + model = UnisonBackbone(**mcfg) + + ckpt_path = args.model_ckpt + if os.path.isdir(ckpt_path): + for name in ("ema_model.pt", "model.safetensors", "pytorch_model.bin"): + candidate = os.path.join(ckpt_path, name) + if os.path.exists(candidate): + ckpt_path = candidate + break + + print(f"Loading checkpoint: {ckpt_path}") + if ckpt_path.endswith(".safetensors"): + sd = safe_load_file(ckpt_path, device="cpu") + else: + sd = torch.load(ckpt_path, map_location="cpu") + + is_ema = "shadow_params" in sd or any(k.startswith("shadow_params") for k in sd) + if is_ema: + print(">>> EMA wrapper detected - unwrapping...") + ema = EMAModel(model.parameters(), decay=0.9999, + model_cls=UnisonBackbone, model_config=mcfg) + ema.load_state_dict(sd) + ema.copy_to(model.parameters()) + print(">>> EMA weights injected.") + else: + ignore = {'decay','num_updates','min_decay','update_after_step', + 'use_ema_warmup','inv_gamma','power','shadow_params'} + model.load_state_dict({k: v for k, v in sd.items() if k not in ignore}, strict=False) + + model.to(device=device, dtype=torch.bfloat16).eval() + return model + + diff --git a/unison/utils/communications.py b/unison/utils/communications.py new file mode 100644 index 0000000000000000000000000000000000000000..9d63ccae0eaecaa42b3c2610d887fa41a8fc2ddf --- /dev/null +++ b/unison/utils/communications.py @@ -0,0 +1,294 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from typing import Any, Tuple +import torch +import torch.distributed as dist + +from torch.nn import functional as F + + +def broadcast(input_: torch.Tensor, group: dist.ProcessGroup): + src = dist.get_global_rank(group, 0) + dist.broadcast(input_, src=src, group=group) + + +def _all_to_all_4D( + input: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None +) -> torch.tensor: + """ + all-to-all for QKV + + Args: + input (torch.tensor): a tensor sharded along dim scatter dim + scatter_idx (int): default 1 + gather_idx (int): default 2 + group : torch process group + + Returns: + torch.tensor: resharded tensor (bs, seqlen/P, hc, hs) + """ + assert ( + input.dim() == 4 + ), f"input must be 4D tensor, got {input.dim()} and shape {input.shape}" + + seq_world_size = dist.get_world_size(group) + + if scatter_idx == 2 and gather_idx == 1: + + seq_lens = [None] * seq_world_size + dist.all_gather_object(seq_lens, input.shape[1], group) + # uneven + if seq_lens[-1] != seq_lens[0] : + assert seq_lens[0] > seq_lens[-1] + gap = seq_lens[0] - seq_lens[-1] + if dist.get_group_rank(group, dist.get_rank()) == seq_world_size - 1: + assert input.shape[1] == seq_lens[-1] + input = F.pad(input, (0, 0, 0, 0, 0, gap)) + else: + gap = 0 + + + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs) + bs, shard_seqlen, hc, hs = input.shape + seqlen = shard_seqlen * seq_world_size + assert hc % seq_world_size == 0, f'Invalid Head size: {hc}, which should be divisible by spsize {seq_world_size}' + shard_hc = hc // seq_world_size + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen/P, hc, hs) -reshape-> (bs, seq_len/P, P, hc/P, hs) -transpose(0,2)-> (P, seq_len/P, bs, hc/P, hs) + input_t = ( + input.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs) + .transpose(0, 2) + .contiguous() + ) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, seq_len/P, bs, hc/P, hs) scatter seqlen -all2all-> (P, seq_len/P, bs, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + else: + output = input_t + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(seqlen, bs, shard_hc, hs) + + # (seq_len, bs, hc/P, hs) -reshape-> (bs, seq_len, hc/P, hs) + output = output.transpose(0, 1).contiguous().reshape(bs, seqlen, shard_hc, hs) + if gap > 0: + output = output[:, :-gap] + + return output + + elif scatter_idx == 1 and gather_idx == 2: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs) + bs, seqlen, shard_hc, hs = input.shape + + hc = shard_hc * seq_world_size + if seqlen % seq_world_size != 0: + new_seqlen = (seqlen // seq_world_size + 1) * seq_world_size + gap = new_seqlen - seqlen + input = F.pad(input, (0, 0, 0, 0, 0, gap)) + bs, seqlen, shard_hc, hs = input.shape + else: + gap = 0 + + + assert seqlen % seq_world_size == 0 + + shard_seqlen = seqlen // seq_world_size + seq_world_size = dist.get_world_size(group) + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen, hc/P, hs) -reshape-> (bs, P, seq_len/P, hc/P, hs) -transpose(0, 3)-> (hc/P, P, seqlen/P, bs, hs) -transpose(0, 1) -> (P, hc/P, seqlen/P, bs, hs) + input_t = ( + input.reshape(bs, seq_world_size, shard_seqlen, shard_hc, hs) + .transpose(0, 3) + .transpose(0, 1) + .contiguous() + .reshape(seq_world_size, shard_hc, shard_seqlen, bs, hs) + ) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, bs x hc/P, seqlen/P, hs) scatter seqlen -all2all-> (P, bs x seq_len/P, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + else: + output = input_t + + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(hc, shard_seqlen, bs, hs) + + # (hc, seqlen/N, bs, hs) -tranpose(0,2)-> (bs, seqlen/N, hc, hs) + output = output.transpose(0, 2).contiguous().reshape(bs, shard_seqlen, hc, hs) + + if gap > 0 and dist.get_group_rank(group, dist.get_rank()) == seq_world_size - 1: + output = output[:, :-gap] + + return output + else: + raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2") + + +class SeqAllToAll4D(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + group: dist.ProcessGroup, + input: torch.Tensor, + scatter_idx: int, + gather_idx: int, + ) -> torch.Tensor: + ctx.group = group + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + + return _all_to_all_4D(input, scatter_idx, gather_idx, group=group) + + @staticmethod + def backward(ctx: Any, *grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None, None]: + return ( + None, + SeqAllToAll4D.apply( + ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx + ), + None, + None, + ) + + +def all_to_all_4D( + input_: torch.Tensor, group: dist.ProcessGroup, scatter_dim: int = 2, gather_dim: int = 1, +): + return SeqAllToAll4D.apply(group, input_, scatter_dim, gather_dim) + + +def _all_to_all( + input_: torch.Tensor, + world_size: int, + group: dist.ProcessGroup, + scatter_dim: int, + gather_dim: int, +): + input_list = [ + t.contiguous() for t in torch.tensor_split(input_, world_size, scatter_dim) + ] + output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)] + dist.all_to_all(output_list, input_list, group=group) + return torch.cat(output_list, dim=gather_dim).contiguous() + + +class _AllToAll(torch.autograd.Function): + """All-to-all communication. + + Args: + input_: input matrix + process_group: communication group + scatter_dim: scatter dimension + gather_dim: gather dimension + """ + + @staticmethod + def forward(ctx, input_, process_group, scatter_dim, gather_dim): + ctx.process_group = process_group + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + ctx.world_size = dist.get_world_size(process_group) + output = _all_to_all( + input_, ctx.world_size, process_group, scatter_dim, gather_dim + ) + return output + + @staticmethod + def backward(ctx, grad_output): + grad_output = _all_to_all( + grad_output, + ctx.world_size, + ctx.process_group, + ctx.gather_dim, + ctx.scatter_dim, + ) + return ( + grad_output, + None, + None, + None, + ) + + +def all_to_all( + input_: torch.Tensor, group: dist.ProcessGroup, scatter_dim: int = 2, gather_dim: int = 1 +): + return _AllToAll.apply(input_, group, scatter_dim, gather_dim) + + +class _AllGather(torch.autograd.Function): + """All-gather communication with autograd support. + + Args: + input_: input tensor + dim: dimension along which to concatenate + """ + + @staticmethod + def forward(ctx, input_, dim, group): + ctx.dim = dim + ctx.group = group + world_size = dist.get_world_size(group) + input_size = list(input_.size()) + + sizes = [None] * world_size + dist.all_gather_object(sizes, input_.shape, group) + + ctx.input_size = input_size[dim] + + tensor_list = [torch.empty(sizes[i], dtype=input_.dtype, device=input_.device) for i in range(world_size)] + input_ = input_.contiguous() + dist.all_gather(tensor_list, input_, group=group) + + output = torch.cat(tensor_list, dim=dim) + return output + + @staticmethod + def backward(ctx, grad_output): + group = ctx.group + world_size = dist.get_world_size(group) + global_rank = dist.get_rank() + rank = dist.get_group_rank(group, global_rank) + dim = ctx.dim + input_size = ctx.input_size + + sizes = [None] * world_size + dist.all_gather_object(sizes, input_size, group=group) + + grad_input_list = torch.split(grad_output, sizes, dim=dim) + grad_input = grad_input_list[rank] + + return grad_input, None, None + + +def all_gather(input_: torch.Tensor, dim: int = 1, group=None): + """Performs an all-gather operation on the input tensor along the specified dimension. + + Args: + input_ (torch.Tensor): Input tensor of shape [B, H, S, D]. + dim (int, optional): Dimension along which to concatenate. Defaults to 1. + + Returns: + torch.Tensor: Output tensor after all-gather operation, concatenated along 'dim'. + """ + return _AllGather.apply(input_, dim, group) diff --git a/unison/utils/diffusers_compat.py b/unison/utils/diffusers_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..04357bdbd6694ede669517182b0c2fc8a6f0a3c5 --- /dev/null +++ b/unison/utils/diffusers_compat.py @@ -0,0 +1,14 @@ +"""Compatibility shims for diffusers + transformers version mismatches.""" + + +def patch_transformers_deepspeed() -> None: + """Patch transformers.deepspeed for diffusers EMAModel. + + Must be called AFTER importing diffusers (diffusers reloads transformers). + """ + import transformers + + if not hasattr(transformers, "deepspeed"): + import transformers.integrations.deepspeed as deepspeed + + transformers.deepspeed = deepspeed diff --git a/unison/utils/flash_attn_no_pad.py b/unison/utils/flash_attn_no_pad.py new file mode 100644 index 0000000000000000000000000000000000000000..6a6fcbbfb7efedda229d330966f59d3db2dfd385 --- /dev/null +++ b/unison/utils/flash_attn_no_pad.py @@ -0,0 +1,92 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from einops import rearrange + + +def flash_attn_no_pad( + qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None, deterministic=False +): + from flash_attn import flash_attn_varlen_qkvpacked_func + from flash_attn.bert_padding import pad_input, unpad_input + batch_size = qkv.shape[0] + seqlen = qkv.shape[1] + nheads = qkv.shape[-2] + x = rearrange(qkv, "b s three h d -> b s (three h d)") + x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input( + x, key_padding_mask + ) + + x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads) + output_unpad = flash_attn_varlen_qkvpacked_func( + x_unpad, + cu_seqlens, + max_s, + dropout_p, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic, + ) + output = rearrange( + pad_input( + rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen + ), + "b s (h d) -> b s h d", + h=nheads, + ) + return output + +def flash_attn_no_pad_v3( + qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None, deterministic=False +): + from flash_attn import flash_attn_varlen_qkvpacked_func + from flash_attn.bert_padding import pad_input, unpad_input + from flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func_v3 + + if flash_attn_varlen_func_v3 is None: + raise ImportError("FlashAttention V3 backend not available") + + batch_size, seqlen, _, nheads, head_dim = qkv.shape + query, key, value = qkv.unbind(dim=2) + + query_unpad, indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input( + rearrange(query, "b s h d -> b s (h d)"), key_padding_mask + ) + key_unpad, _, cu_seqlens_k, _, _ = unpad_input( + rearrange(key, "b s h d -> b s (h d)"), key_padding_mask + ) + value_unpad, _, _, _, _ = unpad_input( + rearrange(value, "b s h d -> b s (h d)"), key_padding_mask + ) + + query_unpad = rearrange(query_unpad, "nnz (h d) -> nnz h d", h=nheads) + key_unpad = rearrange(key_unpad, "nnz (h d) -> nnz h d", h=nheads) + value_unpad = rearrange(value_unpad, "nnz (h d) -> nnz h d", h=nheads) + + output_unpad = flash_attn_varlen_func_v3( + query_unpad, key_unpad, value_unpad, + cu_seqlens_q, cu_seqlens_k, + max_seqlen_q, max_seqlen_q, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic + ) + + output = rearrange( + pad_input(rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen), + "b s (h d) -> b s h d", h=nheads + ) + return output diff --git a/unison/utils/infer_utils.py b/unison/utils/infer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..940deaa59b59c3074ebb1c32df479ea393c65c95 --- /dev/null +++ b/unison/utils/infer_utils.py @@ -0,0 +1,31 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +import torch + +from unison.commons.infer_state import get_infer_state + +def torch_compile_wrapper(): + """Return a decorator that lazily applies torch.compile when enabled in infer_state.""" + def decorator(func): + def wrapper(*args, **kwargs): + if get_infer_state() and get_infer_state().enable_torch_compile: + compiled_func = torch.compile(func) + return compiled_func(*args, **kwargs) + else: + return func(*args, **kwargs) + return wrapper + return decorator