Refactor audio generation scripts to streamline processing and enhance functionality

This commit introduces significant improvements to the audio generation workflow in both `audiogen_medium.py` and `stable_audio.py`. Key changes include:
- Removal of redundant seed extraction logic and integration of a new `process_audio_generations` function to handle audio generation in a more organized manner.
- Consolidation of argument preparation for audio generation into a dedicated `prepare_args` function, improving code clarity and maintainability.
- Enhanced user feedback during the audio generation process, ensuring clearer communication of the actions being performed.

These modifications optimize the audio generation process, improve code organization, and enhance the overall user experience.

audio/audiogen_medium.py

@@ -4,18 +4,10 @@
 import sys
 import os
 import torch
-import torchaudio
 import random
-import multiprocessing as mp
 from audiocraft.models import AudioGen
 from audiocraft.data.audio import audio_write
-
-def get_seed_from_filename(filename):
-    """Extract seed from filename like '12345.wav'"""
-    try:
-        return int(filename.split('.')[0])
-    except:
-        return None
+from utils.audio_utils import process_audio_generations
 
 def generate_audio(args):
     description, seed, prompt_dir = args
@@ -37,61 +29,14 @@ def generate_audio(args):
     # Will save with loudness normalization at -14 db LUFS
     audio_write(file_path, wav[0].cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True)
 
-if __name__ == '__main__':
-    # Set start method to spawn for CUDA multiprocessing
-    mp.set_start_method('spawn')
-    
-    descriptions = sys.argv[1:]
-    if not descriptions:
-        print('At least one prompt should be provided')
-        sys.exit(1)
-
-    # Base output directory
-    base_output_dir = 'generated_audio'
-    os.makedirs(base_output_dir, exist_ok=True)
-
-    # Generate 25 variations for each prompt
-    num_variations = 25
-    num_processes = 3  # Number of parallel models to run
-    seed_range = (0, 1000000)  # Use seeds between 0 and 1,000,000
-
-    for description in descriptions:
-        # Create a safe folder name from the description
-        folder_name = description.replace(' ', '_').replace('/', '_').replace('\\', '_')
-        folder_name = ''.join(c for c in folder_name if c.isalnum() or c in '_-')
-        prompt_dir = os.path.join(base_output_dir, folder_name)
-        os.makedirs(prompt_dir, exist_ok=True)
-        
-        print(f"\nGenerating variations for prompt: {description}")
-        print(f"Saving in directory: {prompt_dir}")
-        
-        # Get existing seeds
-        existing_seeds = set()
-        for filename in os.listdir(prompt_dir):
-            if filename.endswith('.wav'):
-                seed = get_seed_from_filename(filename)
-                if seed is not None:
-                    existing_seeds.add(seed)
-        
-        if len(existing_seeds) >= num_variations:
-            print(f"All {num_variations} variations already exist in {prompt_dir}, skipping...")
-            continue
-        
-        # Generate new random seeds that haven't been used yet
-        needed_variations = num_variations - len(existing_seeds)
-        new_seeds = set()
-        while len(new_seeds) < needed_variations:
-            seed = random.randint(*seed_range)
-            if seed not in existing_seeds and seed not in new_seeds:
-                new_seeds.add(seed)
-        
-        print(f"Generating {needed_variations} new variations using {num_processes} parallel processes...")
-        print(f"Using seeds: {sorted(new_seeds)}")
-        
-        # Prepare arguments for parallel processing
-        args_list = [(description, seed, prompt_dir) for seed in new_seeds]
-        
-        # Use multiprocessing to distribute the work
-        with mp.Pool(processes=num_processes) as pool:
-            pool.map(generate_audio, args_list)
+def prepare_args(description, seeds, prompt_dir):
+    """Prepare arguments for the generate_audio function"""
+    return [(description, seed, prompt_dir) for seed in seeds]
 
+if __name__ == '__main__':
+    process_audio_generations(
+        descriptions=sys.argv[1:],
+        model_name='audiogen',
+        generate_fn=generate_audio,
+        prepare_args_fn=prepare_args
+    )

audio/stable_audio.py

@@ -6,15 +6,11 @@ import os
 import torch
 import soundfile as sf
 import random
-import multiprocessing as mp
 from diffusers import StableAudioPipeline
+from utils.audio_utils import process_audio_generations
 
-def get_seed_from_filename(filename):
-    """Extract seed from filename like '12345.wav'"""
-    try:
-        return int(filename.split('.')[0])
-    except:
-        return None
+# Default negative prompt
+NEGATIVE_PROMPT = "Low quality, noise, distortion, low fidelity"
 
 def generate_audio(args):
     description, negative_prompt, seed, prompt_dir = args
@@ -46,63 +42,14 @@ def generate_audio(args):
     print(f"Saving audio to: {file_path}")
     sf.write(file_path, output, pipe.vae.sampling_rate)
 
-if __name__ == '__main__':
-    # Set start method to spawn for CUDA multiprocessing
-    mp.set_start_method('spawn')
-    
-    descriptions = sys.argv[1:]
-    if not descriptions:
-        print('At least one prompt should be provided')
-        sys.exit(1)
-
-    # Default negative prompt
-    negative_prompt = "Low quality, noise, distortion, low fidelity"
-
-    # Base output directory
-    base_output_dir = 'generated_audio/sa'
-    os.makedirs(base_output_dir, exist_ok=True)
+def prepare_args(description, seeds, prompt_dir):
+    """Prepare arguments for the generate_audio function"""
+    return [(description, NEGATIVE_PROMPT, seed, prompt_dir) for seed in seeds]
 
-    # Generate 25 variations for each prompt
-    num_variations = 25
-    num_processes = 3  # Number of parallel models to run
-    seed_range = (0, 1000000)  # Use seeds between 0 and 1,000,000
-
-    for description in descriptions:
-        # Create a safe folder name from the description
-        folder_name = description.replace(' ', '_').replace('/', '_').replace('\\', '_')
-        folder_name = ''.join(c for c in folder_name if c.isalnum() or c in '_-')
-        prompt_dir = os.path.join(base_output_dir, folder_name)
-        os.makedirs(prompt_dir, exist_ok=True)
-        
-        print(f"\nGenerating variations for prompt: {description}")
-        print(f"Saving in directory: {prompt_dir}")
-        
-        # Get existing seeds
-        existing_seeds = set()
-        for filename in os.listdir(prompt_dir):
-            if filename.endswith('.wav'):
-                seed = get_seed_from_filename(filename)
-                if seed is not None:
-                    existing_seeds.add(seed)
-        
-        if len(existing_seeds) >= num_variations:
-            print(f"All {num_variations} variations already exist in {prompt_dir}, skipping...")
-            continue
-        
-        # Generate new random seeds that haven't been used yet
-        needed_variations = num_variations - len(existing_seeds)
-        new_seeds = set()
-        while len(new_seeds) < needed_variations:
-            seed = random.randint(*seed_range)
-            if seed not in existing_seeds and seed not in new_seeds:
-                new_seeds.add(seed)
-        
-        print(f"Generating {needed_variations} new variations using {num_processes} parallel processes...")
-        print(f"Using seeds: {sorted(new_seeds)}")
-        
-        # Prepare arguments for parallel processing
-        args_list = [(description, negative_prompt, seed, prompt_dir) for seed in new_seeds]
-        
-        # Use multiprocessing to distribute the work
-        with mp.Pool(processes=num_processes) as pool:
-            pool.map(generate_audio, args_list) 
+if __name__ == '__main__':
+    process_audio_generations(
+        descriptions=sys.argv[1:],
+        model_name='sa',
+        generate_fn=generate_audio,
+        prepare_args_fn=prepare_args
+    ) 

audio/tango_audio.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import sys
+import os
+import random
+import soundfile as sf
+from tango import Tango
+from utils.audio_utils import process_audio_generations
+
+def generate_audio(args):
+    description, seed, prompt_dir = args
+    wav_path = os.path.join(prompt_dir, f"{seed}.wav")
+    if os.path.exists(wav_path):
+        print(f"Skipping seed {seed} - file already exists")
+        return
+        
+    # Initialize model for this process
+    tango = Tango("declare-lab/tango")
+    
+    # Set random seed for reproducibility
+    random.seed(seed)
+    
+    # Generate audio
+    audio = tango.generate(description)
+    
+    # Save the audio
+    file_path = os.path.join(prompt_dir, f"{seed}.wav")
+    print(f"Saving audio to: {file_path}")
+    sf.write(file_path, audio, samplerate=16000)
+
+def prepare_args(description, seeds, prompt_dir):
+    """Prepare arguments for the generate_audio function"""
+    return [(description, seed, prompt_dir) for seed in seeds]
+
+if __name__ == '__main__':
+    process_audio_generations(
+        descriptions=sys.argv[1:],
+        model_name='tango',
+        generate_fn=generate_audio,
+        prepare_args_fn=prepare_args
+    )

caption/jtp2.py

@@ -447,5 +447,3 @@ def create_tags(threshold):
 
 if __name__ == "__main__":
     process_directory(args.directory, args.threshold, args.cpu, args.no_grad)
-
-

caption/wdv3.py

@@ -395,4 +395,3 @@ if __name__ == "__main__":
         print(f"Available models: {list(MODEL_REPO_MAP.keys())}")
         raise ValueError(f"Unknown model name '{opts.model}'")
     main(opts)
-

utils/audio_utils.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import os
+import sys
+import random
+import multiprocessing as mp
+
+def get_seed_from_filename(filename):
+    """Extract seed from filename like '12345.wav'"""
+    try:
+        return int(filename.split('.')[0])
+    except:
+        return None
+
+def setup_generation_dir(base_output_dir, description):
+    """Setup and return the directory for a given prompt"""
+    os.makedirs(base_output_dir, exist_ok=True)
+    
+    # Create a safe folder name from the description
+    folder_name = description.replace(' ', '_').replace('/', '_').replace('\\', '_')
+    folder_name = ''.join(c for c in folder_name if c.isalnum() or c in '_-')
+    prompt_dir = os.path.join(base_output_dir, folder_name)
+    os.makedirs(prompt_dir, exist_ok=True)
+    return prompt_dir
+
+def get_existing_seeds(prompt_dir):
+    """Get set of seeds from existing wav files in directory"""
+    existing_seeds = set()
+    for filename in os.listdir(prompt_dir):
+        if filename.endswith('.wav'):
+            seed = get_seed_from_filename(filename)
+            if seed is not None:
+                existing_seeds.add(seed)
+    return existing_seeds
+
+def generate_new_seeds(needed_variations, existing_seeds, seed_range=(0, 1000000)):
+    """Generate new unique random seeds"""
+    new_seeds = set()
+    while len(new_seeds) < needed_variations:
+        seed = random.randint(*seed_range)
+        if seed not in existing_seeds and seed not in new_seeds:
+            new_seeds.add(seed)
+    return new_seeds
+
+def process_audio_generations(descriptions, model_name, generate_fn, prepare_args_fn, num_variations=25, num_processes=3):
+    """
+    Shared logic for processing audio generations across different models.
+    
+    Args:
+        descriptions: List of text prompts to generate audio for
+        model_name: Name of the model (used for output directory)
+        generate_fn: Function that generates a single audio sample
+        prepare_args_fn: Function that prepares arguments for generate_fn
+        num_variations: Number of variations to generate per prompt
+        num_processes: Number of parallel processes to use
+    """
+    # Set start method for multiprocessing
+    mp.set_start_method('spawn', force=True)
+    
+    if not descriptions:
+        print('At least one prompt should be provided')
+        sys.exit(1)
+
+    # Base output directory
+    base_output_dir = f'generated_audio/{model_name}'
+    
+    for description in descriptions:
+        prompt_dir = setup_generation_dir(base_output_dir, description)
+        print(f"\nGenerating variations for prompt: {description}")
+        print(f"Saving in directory: {prompt_dir}")
+        
+        # Get existing seeds and check if we need to generate more
+        existing_seeds = get_existing_seeds(prompt_dir)
+        if len(existing_seeds) >= num_variations:
+            print(f"All {num_variations} variations already exist in {prompt_dir}, skipping...")
+            continue
+        
+        # Generate new random seeds that haven't been used yet
+        needed_variations = num_variations - len(existing_seeds)
+        new_seeds = generate_new_seeds(needed_variations, existing_seeds)
+        
+        print(f"Generating {needed_variations} new variations using {num_processes} parallel processes...")
+        print(f"Using seeds: {sorted(new_seeds)}")
+        
+        # Prepare arguments for parallel processing
+        args_list = prepare_args_fn(description, new_seeds, prompt_dir)
+        
+        # Use multiprocessing to distribute the work
+        with mp.Pool(processes=num_processes) as pool:
+            pool.map(generate_fn, args_list)