Spaces:
Runtime error
Runtime error
File size: 5,215 Bytes
dbac20f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
import logging
from argparse import ArgumentParser
from pathlib import Path
import torch
import torchaudio
from mmaudio.eval_utils import (ModelConfig, all_model_cfg, generate,
load_video, make_video, setup_eval_logging)
from mmaudio.model.flow_matching import FlowMatching
from mmaudio.model.networks import MMAudio, get_my_mmaudio
from mmaudio.model.utils.features_utils import FeaturesUtils
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
log = logging.getLogger()
@torch.inference_mode()
def main():
setup_eval_logging()
parser = ArgumentParser()
parser.add_argument('--variant',
type=str,
default='large_44k_v2',
help='small_16k, small_44k, medium_44k, large_44k, large_44k_v2')
parser.add_argument('--video', type=Path, help='Path to the video file')
parser.add_argument('--prompt', type=str, help='Input prompt', default='')
parser.add_argument('--negative_prompt', type=str, help='Negative prompt', default='')
parser.add_argument('--duration', type=float, default=8.0)
parser.add_argument('--cfg_strength', type=float, default=4.5)
parser.add_argument('--num_steps', type=int, default=25)
parser.add_argument('--mask_away_clip', action='store_true')
parser.add_argument('--output', type=Path, help='Output directory', default='./output')
parser.add_argument('--seed', type=int, help='Random seed', default=42)
parser.add_argument('--skip_video_composite', action='store_true')
parser.add_argument('--full_precision', action='store_true')
args = parser.parse_args()
if args.variant not in all_model_cfg:
raise ValueError(f'Unknown model variant: {args.variant}')
model: ModelConfig = all_model_cfg[args.variant]
model.download_if_needed()
seq_cfg = model.seq_cfg
if args.video:
video_path: Path = Path(args.video).expanduser()
else:
video_path = None
prompt: str = args.prompt
negative_prompt: str = args.negative_prompt
output_dir: str = args.output.expanduser()
seed: int = args.seed
num_steps: int = args.num_steps
duration: float = args.duration
cfg_strength: float = args.cfg_strength
skip_video_composite: bool = args.skip_video_composite
mask_away_clip: bool = args.mask_away_clip
device = 'cuda'
dtype = torch.float32 if args.full_precision else torch.bfloat16
output_dir.mkdir(parents=True, exist_ok=True)
# load a pretrained model
net: MMAudio = get_my_mmaudio(model.model_name).to(device, dtype).eval()
net.load_weights(torch.load(model.model_path, map_location=device, weights_only=True))
log.info(f'Loaded weights from {model.model_path}')
# misc setup
rng = torch.Generator(device=device)
rng.manual_seed(seed)
fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=num_steps)
feature_utils = FeaturesUtils(tod_vae_ckpt=model.vae_path,
synchformer_ckpt=model.synchformer_ckpt,
enable_conditions=True,
mode=model.mode,
bigvgan_vocoder_ckpt=model.bigvgan_16k_path)
feature_utils = feature_utils.to(device, dtype).eval()
if video_path is not None:
log.info(f'Using video {video_path}')
clip_frames, sync_frames, duration = load_video(video_path, duration)
if mask_away_clip:
clip_frames = None
else:
clip_frames = clip_frames.unsqueeze(0)
sync_frames = sync_frames.unsqueeze(0)
else:
log.info('No video provided -- text-to-audio mode')
clip_frames = sync_frames = None
seq_cfg.duration = duration
net.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
log.info(f'Prompt: {prompt}')
log.info(f'Negative prompt: {negative_prompt}')
audios = generate(clip_frames,
sync_frames, [prompt],
negative_text=[negative_prompt],
feature_utils=feature_utils,
net=net,
fm=fm,
rng=rng,
cfg_strength=cfg_strength)
audio = audios.float().cpu()[0]
if video_path is not None:
save_path = output_dir / f'{video_path.stem}.flac'
else:
safe_filename = prompt.replace(' ', '_').replace('/', '_').replace('.', '')
save_path = output_dir / f'{safe_filename}.flac'
torchaudio.save(save_path, audio, seq_cfg.sampling_rate)
log.info(f'Audio saved to {save_path}')
if video_path is not None and not skip_video_composite:
video_save_path = output_dir / f'{video_path.stem}.mp4'
make_video(video_path,
video_save_path,
audio,
sampling_rate=seq_cfg.sampling_rate,
duration_sec=seq_cfg.duration)
log.info(f'Video saved to {output_dir / video_save_path}')
log.info('Memory usage: %.2f GB', torch.cuda.max_memory_allocated() / (2**30))
if __name__ == '__main__':
main()
|