#!/usr/bin/env python3 import os import json import threading import time import numpy as np import torch from tqdm import tqdm from suno_utils.audio import Audio from suno_utils.diffusion import generation as diffusion_gen from suno_utils.tasks.upsample_engine import UpsampleEngine, Request from suno_utils.tasks.dac_vae_fixed_25hz import decode_stream_to_full_audio, encode def setup_distributed(): """Setup distributed processing from SLURM environment variables.""" procid = int(os.environ.get("SLURM_PROCID", 0)) localid = int(os.environ.get("SLURM_LOCALID", 0)) world_size = int(os.environ.get("SLURM_JOB_NUM_NODES", 1)) * int( os.environ.get("SLURM_NTASKS_PER_NODE", 1) ) assert world_size > 0, "WORLD_SIZE is 0" print(f"PROCID: {procid}, LOCALID: {localid}, WORLD_SIZE: {world_size}") # Set CUDA device torch.cuda.set_device(localid) return procid, localid, world_size def load_cover_metadata(metadata_file): """Load cover metadata from JSONL file.""" print(f"Loading cover metadata from {metadata_file}") cover_paths = [] with open(metadata_file, "r") as f: for line in tqdm(f, desc="Loading cover metadata"): if not line.strip(): continue record = json.loads(line) cover_paths.append(record.get("local_filepath")) print(f"Found {len(cover_paths):,} cover tracks with audio files") return cover_paths def init_models(): """Initialize diffusion and upsample models.""" print("Starting model initialization...") start_time = time.time() diffusion_gen.preload_models( dit_model_filepath="/app/suno/checkpoints/2025-05-26_22-49-34_s8646/last_ckpt_infer.pt", codec_filepath="s3://suno-data/minz/models/dac_vae_tuned_25hz.pth", compile=True, ) model_load_time = time.time() - start_time print(f"Model loading took {model_load_time:.2f}s") engine = UpsampleEngine(min_chunk_size=25 * 30) total_time = time.time() - start_time print(f"Total model initialization took {total_time:.2f}s") return engine def gen_stem( audio: Audio, engine: UpsampleEngine, stem_type_cfg_scale=1.0, tags="extract [split_karaoke]", steps=4, seed=3, codec_scale_factor=0.4, scale_ctx_vector=True, noise_ctx_level=0.0, infill_prefix_latents=None, infill_suffix_latents=None, ): """Generate stems from audio using diffusion model.""" print(f" Starting stem generation (steps={steps})...") start_time = time.time() # VAE encoding encode_start = time.time() vae = encode(audio) encode_time = time.time() - encode_start print(f" VAE encoding took {encode_time:.2f}s, shape: {vae.shape}") gen_cfg = diffusion_gen.DiffusionGenerationConfig( lyrics=tags, steps=steps, seed=seed, codec_scale_factor=codec_scale_factor, scale_ctx_vector=scale_ctx_vector, noise_ctx_level=noise_ctx_level, text_cfg_coef=stem_type_cfg_scale, infill_prefix_latents=infill_prefix_latents, infill_suffix_latents=infill_suffix_latents, drop_semantic_tokens=True, ) request = Request( id="dummy", generation_config=gen_cfg, tokens=np.zeros((vae.shape[0], 1)), input_tokens_finished=True, stem_ctx_latents=vae, ) # Diffusion generation diffusion_start = time.time() result = engine.run_request(request, tqdm_enabled=False) vae_latents = torch.concat(result.vae_latents) diffusion_time = time.time() - diffusion_start print(f" Diffusion generation took {diffusion_time:.2f}s, output shape: {vae_latents.shape}") # Decoding stems decode_start = time.time() audios = [] for i in tqdm(range(vae_latents.shape[1]), desc="Decoding stems", disable=True): audios.append(decode_stream_to_full_audio(vae_latents[:, i], n_stride_tokens=25 * 10)) decode_time = time.time() - decode_start print(f" Decoding {vae_latents.shape[1]} stems took {decode_time:.2f}s") total_time = time.time() - start_time print(f" Total stem generation took {total_time:.2f}s") return audios def write_stem(args): """Write a single stem to disk.""" out_root, category, stem = args if stem.loudness < -45: return None out_path = f"{out_root}/{category}.opus" os.makedirs(os.path.dirname(out_path), exist_ok=True) stem.write_opus(out_path) return category def process_id(input_path, engine, out_stem_dir, categories): """Process a single audio file and generate stems.""" name = input_path.split("/")[-1].split(".")[0] output_path = f"{out_stem_dir}/{name}" print(f"Processing {name}") total_start = time.time() try: # Load audio load_start = time.time() audio = Audio.from_file(input_path, n_channels=2) load_time = time.time() - load_start print(f" Audio loading took {load_time:.2f}s, duration: {audio.duration_s:.2f}s") # Generate stems stems = gen_stem(audio, engine, steps=8) print(f" Generated {len(stems)} stems") # Write stems to disk write_start = time.time() write_args = [(output_path, category, stem) for category, stem in zip(categories, stems)] # Use threading to write stems in parallel results = [None] * len(write_args) threads = [] def write_stem_thread(i, args): results[i] = write_stem(args) for i, args in enumerate(write_args): thread = threading.Thread(target=write_stem_thread, args=(i, args)) threads.append(thread) thread.start() for thread in threads: thread.join() write_time = time.time() - write_start print(f" Writing stems took {write_time:.2f}s") # Filter out None results found_categories = [result for result in results if result is not None] total_time = time.time() - total_start print(f" Total processing time: {total_time:.2f}s") print(f" Successfully wrote stems: {found_categories}") return found_categories except Exception as e: print(f"Error processing {name}: {e}") return [] def main(): """Main processing function.""" # Setup distributed processing procid, localid, world_size = setup_distributed() # Output directory for stems out_stem_dir = "/app2/suno/data/cover_stems" os.makedirs(out_stem_dir, exist_ok=True) # Load cover metadata cover_metadata_file = "/home/victor/neon/sunoData/src/sunodata/gpt/metas_covers.jsonl" cover_paths = load_cover_metadata(cover_metadata_file) # Initialize models engine = init_models() # Stem categories categories = [ "Vocals", "Backing_Vocals", "Drums", "Bass", "Guitar", "Keyboard", "Percussion", "Strings", "Synth", "FX", "Brass", "Woodwinds", ] # Distribute work across processes total_files = len(cover_paths) files_per_process = total_files // world_size start_idx = procid * files_per_process end_idx = start_idx + files_per_process if procid < world_size - 1 else total_files my_files = cover_paths[start_idx:end_idx] print(f"Process {procid} will handle {len(my_files)} files (indices {start_idx}:{end_idx})") # Process files assigned to this process for i, input_path in enumerate(tqdm(my_files, desc=f"Process {procid}", mininterval=120)): if not os.path.exists(input_path): print(f"File not found: {input_path}") continue result = process_id(input_path, engine, out_stem_dir, categories) if __name__ == "__main__": main()