backend/models/audio_generator.py

# audio_generator.py placeholder
"""
WHAT: This module defines the Audio Generator, the fourth stage of the Multiverse AI pipeline.
WHY: A compelling video requires a soundscape. This model takes the specialized audio prompt 
     (generated by our LLM in stage 1) and synthesizes an ambient music track or sound effect 
     to accompany the visual scene.
HOW: It implements the `BaseModel` interface. It uses Hugging Face's `transformers` library to 
     load Meta's `MusicGen` model. It converts the generated audio tensors into raw WAV bytes 
     so our file_manager can save it directly to disk.
"""

import gc
import io
import torch
import scipy.io.wavfile
from typing import Dict, Union

# Import specific classes for MusicGen from transformers
from transformers import AutoProcessor, MusicgenForConditionalGeneration

# Import our base interface and global configurations
from .base import BaseModel
from ..config import MODEL_IDS, DEVICE, MOCK_INFERENCE, FORCE_CPU_INFERENCE

class AudioGenerator(BaseModel):
    """
    Wrapper for the Text-to-Audio model (MusicGen-Small).
    Generates background audio/music based on the expanded prompt and returns raw WAV bytes.
    """

    def __init__(self):
        self.model = None
        self.processor = None
        self.model_id = MODEL_IDS["audio_generation"]

    def initialize(self) -> None:
        # If running on CPU and FORCE_CPU_INFERENCE is False, bypass the heavy model load to protect system RAM.
        if MOCK_INFERENCE or (DEVICE == "cpu" and not FORCE_CPU_INFERENCE):
            print("[AudioGenerator] Running in MOCK mode or CPU bypass active. Bypassing audio model load.")
            return

        try:
            torch_dtype = torch.float16 if DEVICE == "cuda" else torch.float32

            # The processor handles converting text into token IDs that the model understands
            self.processor = AutoProcessor.from_pretrained(self.model_id)
            
            # The core generative model
            self.model = MusicgenForConditionalGeneration.from_pretrained(
                self.model_id,
                torch_dtype=torch_dtype
            )
            
            self.model.to(DEVICE)
        except Exception as e:
            print(f"[AudioGenerator Warning] Failed to load local weights: {e}. Degrading gracefully to MOCK fallback.")
            self.model = None

    def generate(self, **kwargs) -> bytes:
        """
        WHAT: Synthesizes audio based on the text prompt and serializes it into a WAV format byte stream.
        WHY: The output of this function needs to be written to a file system and eventually 
             mixed with the video. Returning standard WAV bytes makes file handling trivial.
        HOW: 
             1. Extracts the `audio_prompt`.
             2. Tokenizes the text.
             3. Runs model.generate() (MusicGen generates ~50 tokens per second of audio).
             4. Extracts the raw waveform array.
             5. Uses `scipy.io.wavfile` to encode the array into WAV bytes in memory.
        """
        # Extract the specific, optimized audio prompt from the dict passed by the orchestrator
        prompt_data: Union[Dict[str, str], str] = kwargs.get("prompt", "")
        
        if isinstance(prompt_data, dict):
            actual_prompt = prompt_data.get("audio_prompt", "")
        else:
            actual_prompt = str(prompt_data)

        if not actual_prompt:
            raise ValueError("AudioGenerator requires a valid 'prompt' in kwargs.")

        if MOCK_INFERENCE or self.model is None:
            print("[AudioGenerator] Generating mock soundscape audio...")
            import numpy as np
            sr = 16000  # 16kHz
            duration = 2.0  # 2 seconds
            t = np.linspace(0, duration, int(sr * duration), endpoint=False)
            # A warm 220Hz bass sine wave tone
            audio_data = 0.5 * np.sin(2 * np.pi * 220 * t)
            wav_io = io.BytesIO()
            scipy.io.wavfile.write(wav_io, rate=sr, data=audio_data.astype(np.float32))
            return wav_io.getvalue()

        try:
            # Process the text input
            inputs = self.processor(
                text=[actual_prompt],
                padding=True,
                return_tensors="pt"
            )
            
            # Move inputs to the same device as the model
            inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

            # Generate audio tokens. 
            # MusicGen yields ~1 second of audio per 50 tokens. 
            # max_new_tokens=256 creates roughly 5 seconds of audio, ideal for our short video clip.
            audio_values = self.model.generate(
                **inputs, 
                max_new_tokens=256
            )

            # The output is a PyTorch tensor of shape (batch, channels, sequence_length)
            # We extract the sequence and convert it to a NumPy array for SciPy.
            sampling_rate = self.model.config.audio_encoder.sampling_rate
            audio_data = audio_values[0, 0].cpu().numpy()

            # Convert the raw numpy array into a WAV file formatted byte string in-memory
            wav_io = io.BytesIO()
            scipy.io.wavfile.write(wav_io, rate=sampling_rate, data=audio_data)
            
            return wav_io.getvalue()
        except Exception as e:
            raise RuntimeError(f"Audio generation failed: {str(e)}") from e

    def cleanup(self) -> None:
        """
        WHAT: Unloads MusicGen from memory and flushes the GPU cache.
        WHY: VRAM MANAGEMENT. Once the audio bytes are created, we no longer need this model. 
             We must aggressively clear it out so the final Video Generation model (which is 
             usually the heaviest of all) has completely empty VRAM to operate within.
        HOW: Deletes Python references to both the model and processor, forces garbage collection, 
             and calls PyTorch's empty_cache() utility.
        """
        if MOCK_INFERENCE:
            return

        if self.model is not None:
            del self.model
            self.model = None
            
        if self.processor is not None:
            del self.processor
            self.processor = None
            
        # Force Python to run garbage collection
        gc.collect()
        
        # If using a GPU, force PyTorch to release the freed memory back to the OS
        if DEVICE == "cuda":
            torch.cuda.empty_cache()