import warnings warnings.filterwarnings("ignore", category=FutureWarning) import logging from argparse import ArgumentParser from pathlib import Path import torch import torchaudio from hydra import compose, initialize from resonate.eval_utils import generate_fm, setup_eval_logging from resonate.model.flow_matching import FlowMatching from resonate.model.networks import FluxAudio, get_model from resonate.model.utils.features_utils import FeaturesUtils from resonate.model.sequence_config import CONFIG_16K, CONFIG_44K from torchaudio.transforms import Resample torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True from tqdm import tqdm log = logging.getLogger() @torch.inference_mode() def main(): setup_eval_logging() parser = ArgumentParser() parser.add_argument('--config_name', type=str, required=True, help='config file name under config/ (e.g., train_config_online_feature_umt5.yaml)') 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=9.975) parser.add_argument('--cfg_strength', type=float, default=4.5) parser.add_argument('--num_steps', type=int, default=25) 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('--full_precision', action='store_true') parser.add_argument('--model_path', type=str, help='Path of trained model') parser.add_argument('--debug', action='store_true') parser.add_argument('--ref_audio_path', type=str, required=False) args = parser.parse_args() if args.debug: import debugpy debugpy.listen(6666) print("Waiting for debugger attach (rank 0)...") debugpy.wait_for_client() with initialize(version_base="1.3.2", config_path="config"): cfg = compose(config_name=args.config_name) if cfg.audio_sample_rate == 16000: seq_cfg = CONFIG_16K elif cfg.audio_sample_rate == 44100: seq_cfg = CONFIG_44K else: raise ValueError(f'Invalid audio sample rate: {cfg.audio_sample_rate}') 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 device = 'cpu' if torch.cuda.is_available(): device = 'cuda' elif torch.backends.mps.is_available(): device = 'mps' else: log.warning('CUDA/MPS are not available, running on CPU') dtype = torch.float32 if args.full_precision else torch.bfloat16 output_dir.mkdir(parents=True, exist_ok=True) # load a pretrained model with train-style params (only overwrite model-loading related fields) use_rope = cfg.get('use_rope', True) text_dim = cfg.get('text_dim', None) text_c_dim = cfg.get('text_c_dim', None) net: FluxAudio = get_model(cfg.model, use_rope=use_rope, text_dim=text_dim, text_c_dim=text_c_dim).to(device, dtype).eval() net.load_weights(torch.load(args.model_path, map_location=device, weights_only=True)) log.info(f'Loaded weights from {args.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) encoder_name = cfg.get('text_encoder_name', 'flan-t5') if cfg.audio_sample_rate == 16000: feature_utils = FeaturesUtils(tod_vae_ckpt=cfg.get('vae_16k_ckpt'), enable_conditions=True, encoder_name=encoder_name, mode='16k', bigvgan_vocoder_ckpt=cfg.get('bigvgan_vocoder_ckpt'), need_vae_encoder=True) elif cfg.audio_sample_rate == 44100: feature_utils = FeaturesUtils(tod_vae_ckpt=cfg.get('vae_44k_ckpt'), enable_conditions=True, encoder_name=encoder_name, mode='44k', need_vae_encoder=True) else: raise ValueError(f'Invalid audio sample rate: {cfg.audio_sample_rate}') feature_utils = feature_utils.to(device, dtype).eval() seq_cfg.duration = duration net.update_seq_lengths(seq_cfg.latent_seq_len) log.info(f'Updated seq_cfg latent_seq_len: {seq_cfg.latent_seq_len}') # prompts: str = [args.prompt] if args.prompt != "": prompts = [args.prompt] else: prompts = [ # ---------- Sound scene prompts ---------- "Light rain taps steadily against the pavement while distant cars hiss past on the wet road. Occasional footsteps splash through puddles, accompanied by the low hum of city traffic.", "Birds chirp energetically from different directions as leaves rustle softly in a mild breeze. Somewhere deeper in the forest, a woodpecker taps rhythmically against a tree trunk.", "Espresso machines hiss and steam, cups clink against saucers, and quiet conversations overlap with soft background music. A barista calls out drink orders above the ambient chatter.", "Metal clanks sharply as tools collide, followed by the constant rumble of heavy machinery. Occasional warning beeps and shouted instructions cut through the industrial noise.", "Waves crash and retreat in a steady rhythm, mixing with the cries of seagulls overhead. The wind carries the faint sound of water splashing against nearby rocks.", "The room is mostly silent, broken only by the gentle whir of a ceiling fan and the occasional creak of furniture. From outside, a distant siren fades slowly into the night.", "Announcements echo through the hall as trains arrive and depart. Rolling suitcases, hurried footsteps, and overlapping voices create a constant, restless background noise.", "Deep bass pulses through the air while sharp synthetic melodies cut in and out. The rhythm builds steadily, accompanied by crowd cheers and reverberating echoes.", "Oil sizzles loudly in a pan as vegetables are tossed and chopped. The clatter of utensils and the soft bubbling of boiling water fill the space.", "Low thunder rumbles in the distance, growing louder with each passing moment. Wind howls through trees as the first heavy raindrops strike the ground.", # ---------- Music generation prompts ---------- "A slow ambient electronic track with warm synthesizer pads, minimal rhythm, and a calm, floating atmosphere.", "An energetic electronic dance track driven by punchy kick drums, crisp hi-hats, and a powerful bassline.", "A cinematic orchestral piece featuring rising strings, deep brass, and dramatic percussion for an epic battle scene.", "A soft piano solo with gentle reverb, slow tempo, and an emotional, reflective mood.", "A lo-fi hip hop beat with vinyl crackle, mellow chords, relaxed drums, and a nostalgic late-night feeling.", "A dark industrial techno track with distorted synths, mechanical rhythms, and an intense, aggressive energy.", "An acoustic folk song with fingerpicked guitar, light percussion, and a warm, intimate atmosphere.", "A futuristic synthwave track inspired by the 1980s, featuring analog synth leads, steady arpeggios, and retro drum machines.", "A jazz trio performance with upright bass, brushed drums, and expressive piano improvisation.", "A minimal drone music piece built on sustained tones, subtle texture changes, and a tense, immersive soundscape." ] for prompt in tqdm(prompts): log.info(f'Prompt: {prompt}') log.info(f'Negative prompt: {negative_prompt}') audios = generate_fm([prompt], negative_text=[negative_prompt], feature_utils=feature_utils, net=net, fm=fm, rng=rng, cfg_strength=cfg_strength) audio = audios.float().cpu()[0] safe_filename = prompt.replace(' ', '_').replace('/', '_').replace('.', '') safe_filename = safe_filename[:200] save_path = output_dir / f'{safe_filename}--numsteps{num_steps}--seed{args.seed}--duration{args.duration}.wav' torchaudio.save(save_path, audio, seq_cfg.sampling_rate) log.info(f'Audio saved to {save_path}') log.info('Memory usage: %.2f GB', torch.cuda.max_memory_allocated() / (2**30)) if __name__ == '__main__': main()