import dataclasses
import logging
from pathlib import Path
from typing import Optional
import torch
from colorlog import ColoredFormatter
from torchvision.transforms import v2
from mmaudio.data.av_utils import VideoInfo, read_frames, reencode_with_audio
from mmaudio.model.flow_matching import FlowMatching
from mmaudio.model.networks import MMAudio
from mmaudio.model.sequence_config import (CONFIG_16K, CONFIG_44K, SequenceConfig)
from mmaudio.model.utils.features_utils import FeaturesUtils
from mmaudio.utils.download_utils import download_model_if_needed
log = logging.getLogger()
@dataclasses.dataclass
class ModelConfig:
model_name: str
model_path: Path
vae_path: Path
bigvgan_16k_path: Optional[Path]
mode: str
synchformer_ckpt: Path = Path('./ext_weights/synchformer_state_dict.pth')
@property
def seq_cfg(self) -> SequenceConfig:
if self.mode == '16k':
return CONFIG_16K
elif self.mode == '44k':
return CONFIG_44K
def download_if_needed(self):
download_model_if_needed(self.model_path)
download_model_if_needed(self.vae_path)
if self.bigvgan_16k_path is not None:
download_model_if_needed(self.bigvgan_16k_path)
download_model_if_needed(self.synchformer_ckpt)
small_16k = ModelConfig(model_name='small_16k',
model_path=Path('./weights/mmaudio_small_16k.pth'),
vae_path=Path('./ext_weights/v1-16.pth'),
bigvgan_16k_path=Path('./ext_weights/best_netG.pt'),
mode='16k')
small_44k = ModelConfig(model_name='small_44k',
model_path=Path('./weights/mmaudio_small_44k.pth'),
vae_path=Path('./ext_weights/v1-44.pth'),
bigvgan_16k_path=None,
mode='44k')
medium_44k = ModelConfig(model_name='medium_44k',
model_path=Path('./weights/mmaudio_medium_44k.pth'),
vae_path=Path('./ext_weights/v1-44.pth'),
bigvgan_16k_path=None,
mode='44k')
large_44k = ModelConfig(model_name='large_44k',
model_path=Path('./weights/mmaudio_large_44k.pth'),
vae_path=Path('./ext_weights/v1-44.pth'),
bigvgan_16k_path=None,
mode='44k')
large_44k_v2 = ModelConfig(model_name='large_44k_v2',
model_path=Path('./weights/mmaudio_large_44k_v2.pth'),
vae_path=Path('./ext_weights/v1-44.pth'),
bigvgan_16k_path=None,
mode='44k')
all_model_cfg: dict[str, ModelConfig] = {
'small_16k': small_16k,
'small_44k': small_44k,
'medium_44k': medium_44k,
'large_44k': large_44k,
'large_44k_v2': large_44k_v2,
}
def generate(
clip_video: Optional[torch.Tensor],
sync_video: Optional[torch.Tensor],
text: Optional[list[str]],
*,
negative_text: Optional[list[str]] = None,
feature_utils: FeaturesUtils,
net: MMAudio,
fm: FlowMatching,
rng: torch.Generator,
cfg_strength: float,
clip_batch_size_multiplier: int = 40,
sync_batch_size_multiplier: int = 40,
) -> torch.Tensor:
device = feature_utils.device
dtype = feature_utils.dtype
bs = len(text)
if clip_video is not None:
clip_video = clip_video.to(device, dtype, non_blocking=True)
clip_features = feature_utils.encode_video_with_clip(clip_video,
batch_size=bs *
clip_batch_size_multiplier)
else:
clip_features = net.get_empty_clip_sequence(bs)
if sync_video is not None:
sync_video = sync_video.to(device, dtype, non_blocking=True)
sync_features = feature_utils.encode_video_with_sync(sync_video,
batch_size=bs *
sync_batch_size_multiplier)
else:
sync_features = net.get_empty_sync_sequence(bs)
if text is not None:
text_features = feature_utils.encode_text(text)
else:
text_features = net.get_empty_string_sequence(bs)
if negative_text is not None:
assert len(negative_text) == bs
negative_text_features = feature_utils.encode_text(negative_text)
else:
negative_text_features = net.get_empty_string_sequence(bs)
x0 = torch.randn(bs,
net.latent_seq_len,
net.latent_dim,
device=device,
dtype=dtype,
generator=rng)
preprocessed_conditions = net.preprocess_conditions(clip_features, sync_features, text_features)
empty_conditions = net.get_empty_conditions(
bs, negative_text_features=negative_text_features if negative_text is not None else None)
cfg_ode_wrapper = lambda t, x: net.ode_wrapper(t, x, preprocessed_conditions, empty_conditions,
cfg_strength)
x1 = fm.to_data(cfg_ode_wrapper, x0)
x1 = net.unnormalize(x1)
spec = feature_utils.decode(x1)
audio = feature_utils.vocode(spec)
return audio
LOGFORMAT = " %(log_color)s%(levelname)-8s%(reset)s | %(log_color)s%(message)s%(reset)s"
def setup_eval_logging(log_level: int = logging.INFO):
logging.root.setLevel(log_level)
formatter = ColoredFormatter(LOGFORMAT)
stream = logging.StreamHandler()
stream.setLevel(log_level)
stream.setFormatter(formatter)
log = logging.getLogger()
log.setLevel(log_level)
log.addHandler(stream)
def load_video(video_path: Path, duration_sec: float, load_all_frames: bool = True) -> VideoInfo:
_CLIP_SIZE = 384
_CLIP_FPS = 8.0
_SYNC_SIZE = 224
_SYNC_FPS = 25.0
clip_transform = v2.Compose([
v2.Resize((_CLIP_SIZE, _CLIP_SIZE), interpolation=v2.InterpolationMode.BICUBIC),
v2.ToImage(),
v2.ToDtype(torch.float32, scale=True),
])
sync_transform = v2.Compose([
v2.Resize(_SYNC_SIZE, interpolation=v2.InterpolationMode.BICUBIC),
v2.CenterCrop(_SYNC_SIZE),
v2.ToImage(),
v2.ToDtype(torch.float32, scale=True),
v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
output_frames, all_frames, orig_fps = read_frames(video_path,
list_of_fps=[_CLIP_FPS, _SYNC_FPS],
start_sec=0,
end_sec=duration_sec,
need_all_frames=load_all_frames)
clip_chunk, sync_chunk = output_frames
clip_chunk = torch.from_numpy(clip_chunk).permute(0, 3, 1, 2)
sync_chunk = torch.from_numpy(sync_chunk).permute(0, 3, 1, 2)
clip_frames = clip_transform(clip_chunk)
sync_frames = sync_transform(sync_chunk)
clip_length_sec = clip_frames.shape[0] / _CLIP_FPS
sync_length_sec = sync_frames.shape[0] / _SYNC_FPS
if clip_length_sec < duration_sec:
log.warning(f'Clip video is too short: {clip_length_sec:.2f} < {duration_sec:.2f}')
log.warning(f'Truncating to {clip_length_sec:.2f} sec')
duration_sec = clip_length_sec
if sync_length_sec < duration_sec:
log.warning(f'Sync video is too short: {sync_length_sec:.2f} < {duration_sec:.2f}')
log.warning(f'Truncating to {sync_length_sec:.2f} sec')
duration_sec = sync_length_sec
clip_frames = clip_frames[:int(_CLIP_FPS * duration_sec)]
sync_frames = sync_frames[:int(_SYNC_FPS * duration_sec)]
video_info = VideoInfo(
duration_sec=duration_sec,
fps=orig_fps,
clip_frames=clip_frames,
sync_frames=sync_frames,
all_frames=all_frames if load_all_frames else None,
)
return video_info
def make_video(video_info: VideoInfo, output_path: Path, audio: torch.Tensor, sampling_rate: int):
reencode_with_audio(video_info, output_path, audio, sampling_rate)