Delete demo.py

demo.py

@@ -1,135 +0,0 @@
-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()