cera 0.2.6

//! Audio-aware generation loop with text ↔ audio modality switching.

use anyhow::Result;

use crate::kv_cache::InferenceState;
use crate::model::Model;
use crate::model::audio_decoder::{
    AudioDecoderWeights, AudioGpu, DepthformerState, DetokenizerState, DetokenizerWeights,
    detokenize_to_spectrum, embed_audio_token, istft_to_pcm, sample_audio_frame,
};
use crate::sampler::{Sampler, SamplerConfig};
use crate::time::{Duration, Instant};
use crate::tokenizer::BpeTokenizer;

/// Audio generation configuration.
pub struct AudioGenerateConfig {
    pub max_tokens: usize,
    pub sampler: SamplerConfig,
    /// Audio sampling temperature (0.0 = greedy, >0 = stochastic).
    pub audio_temperature: f32,
    /// Audio top-k for stochastic sampling.
    pub audio_top_k: usize,
    /// Generation mode.
    pub mode: AudioMode,
    /// Use GPU for depthformer (code sampling). Disabled by default because
    /// GEMV accumulation order differences can produce different codes.
    pub gpu_depthformer: bool,
}

#[derive(Clone, Copy)]
pub enum AudioMode {
    /// All text first, then audio when <|audio_start|> (128) is emitted.
    Sequential,
    /// Alternate: 6 text tokens, 12 audio frames, repeat.
    Interleaved,
}

/// Result of audio generation.
pub struct AudioGenerateResult {
    pub text_tokens: usize,
    pub audio_frames: usize,
    pub audio_samples: usize,
    pub elapsed_secs: f64,
    pub depthformer_secs: f64,
    pub detokenizer_secs: f64,
}

/// Special token IDs for modality control.
const TOKEN_AUDIO_START: u32 = 128;
const TOKEN_TEXT_END: u32 = 130;
const AUDIO_END_CODE: i32 = 2048;

#[derive(PartialEq)]
enum Modality {
    Text,
    Audio,
}

// ---------------------------------------------------------------------------
// AudioOutputDecoder
// ---------------------------------------------------------------------------

/// Per-frame outcome from [`AudioOutputDecoder::decode_frame`].
enum FrameOutcome {
    /// Codes sampled + detokenized; `audio_embedding` is the feedback
    /// embedding the caller should pass back through the main LLM.
    Codes { audio_embedding: Vec<f32> },
    /// Audio stream terminated (`codes[0] == AUDIO_END_CODE`). No
    /// spectrum produced for this frame; caller should return control
    /// to the text modality per its mode's exit convention.
    End,
}

/// Owns the audio output-decoder state and exposes per-frame operations.
/// Extracted from `generate_audio` so the same per-frame logic is shared
/// between the Sequential and Interleaved paths and so future decoder
/// variants can slot in behind a common interface.
///
/// External `generate_audio` signature is unchanged; this is a purely
/// internal refactor.
struct AudioOutputDecoder<'a> {
    weights: &'a AudioDecoderWeights,
    detok_weights: &'a DetokenizerWeights,
    gpu: Option<&'a dyn AudioGpu>,
    df_state: DepthformerState,
    detok_state: DetokenizerState,
    /// Accumulated spectrum across the entire generate_audio call.
    /// Flushed by `finish()` via a single ISTFT pass — per-frame ISTFT
    /// would produce discontinuities from fresh overlap buffers.
    all_spectrum: Vec<f32>,
    audio_frames: usize,
    time_depthformer: Duration,
    time_detokenizer: Duration,
    audio_temperature: f32,
    audio_top_k: usize,
    /// Precomputed: `gpu.is_some() && config.gpu_depthformer`.
    use_gpu_df: bool,
}

impl<'a> AudioOutputDecoder<'a> {
    fn new(
        weights: &'a AudioDecoderWeights,
        detok_weights: &'a DetokenizerWeights,
        gpu: Option<&'a dyn AudioGpu>,
        audio_temperature: f32,
        audio_top_k: usize,
        gpu_depthformer: bool,
    ) -> Self {
        // One-shot GPU reset so repeated generate_audio calls don't leak
        // state across invocations.
        if let Some(g) = gpu {
            g.reset_detokenizer();
            g.reset_depthformer();
        }
        let df_state = DepthformerState::new(&weights.depthformer_config);
        let detok_state = DetokenizerState::new(&detok_weights.config);
        Self {
            weights,
            detok_weights,
            gpu,
            df_state,
            detok_state,
            all_spectrum: Vec::new(),
            audio_frames: 0,
            time_depthformer: Duration::ZERO,
            time_detokenizer: Duration::ZERO,
            audio_temperature,
            audio_top_k,
            use_gpu_df: gpu_depthformer && gpu.is_some(),
        }
    }

    /// Sample one audio frame via depthformer, detect end-of-stream,
    /// detokenize into spectrum, and produce the feedback embedding the
    /// caller feeds back through the main LLM.
    ///
    /// `embed` is the main LLM's hidden state / embedding to condition
    /// this frame on (the audio_start token embedding on the first
    /// frame, or the prior frame's feedback embedding afterward).
    fn decode_frame(&mut self, embed: &[f32]) -> FrameOutcome {
        let t0 = Instant::now();
        let codes = if self.use_gpu_df {
            let g = self
                .gpu
                .expect("use_gpu_df implies gpu is Some (set at construction)");
            g.sample_audio_frame(embed, self.audio_temperature, self.audio_top_k)
        } else {
            sample_audio_frame(
                self.weights,
                &mut self.df_state,
                embed,
                self.audio_temperature,
                self.audio_top_k,
            )
        };
        self.time_depthformer += t0.elapsed();

        if codes[0] == AUDIO_END_CODE {
            return FrameOutcome::End;
        }

        let t1 = Instant::now();
        let spectrum = if let Some(g) = self.gpu {
            g.detokenize_to_spectrum(self.detok_weights, &codes)
        } else {
            detokenize_to_spectrum(
                self.detok_weights,
                self.weights,
                &mut self.detok_state,
                &codes,
            )
        };
        self.time_detokenizer += t1.elapsed();
        self.all_spectrum.extend_from_slice(&spectrum);
        self.audio_frames += 1;

        let audio_embedding = embed_audio_token(self.weights, &codes);
        FrameOutcome::Codes { audio_embedding }
    }

    /// Drain the accumulated spectrum through a single ISTFT pass and
    /// emit the resulting PCM via `sink`. Returns the PCM sample count.
    /// A single end-of-generation ISTFT (rather than per-frame) avoids
    /// discontinuities from fresh overlap buffers.
    fn finish(&mut self, mut sink: impl FnMut(&[f32], u32)) -> usize {
        if self.all_spectrum.is_empty() {
            return 0;
        }
        let pcm = istft_to_pcm(
            &self.all_spectrum,
            self.detok_weights.config.n_fft,
            self.detok_weights.config.hop_length,
        );
        if pcm.is_empty() {
            return 0;
        }
        let n = pcm.len();
        sink(&pcm, self.detok_weights.config.sample_rate as u32);
        n
    }
}

// ---------------------------------------------------------------------------
// generate_audio
// ---------------------------------------------------------------------------

/// Generate text + audio from a model with vocoder.
///
/// `gpu`: optional GPU backend for depthformer + detokenizer acceleration.
#[allow(unused_assignments, clippy::too_many_arguments)]
pub fn generate_audio(
    model: &dyn Model,
    decoder_weights: &AudioDecoderWeights,
    detok_weights: &DetokenizerWeights,
    tokenizer: &BpeTokenizer,
    prompt_tokens: &[u32],
    config: &AudioGenerateConfig,
    gpu: Option<&dyn AudioGpu>,
    mut text_callback: impl FnMut(&str),
    mut audio_callback: impl FnMut(&[f32], u32),
) -> Result<AudioGenerateResult> {
    anyhow::ensure!(!prompt_tokens.is_empty(), "prompt_tokens must not be empty");

    let model_config = model.config();
    let mut state = InferenceState::from_config(model_config)?;
    let mut sampler = Sampler::new(config.sampler.clone());
    let mut decoder = AudioOutputDecoder::new(
        decoder_weights,
        detok_weights,
        gpu,
        config.audio_temperature,
        config.audio_top_k,
        config.gpu_depthformer,
    );

    let start = Instant::now();

    // Prefill.
    let mut logits = model.forward_prefill(prompt_tokens, 0, &mut state);

    let mut modality = Modality::Text;
    let mut generated = 0usize;
    let mut text_tokens = 0usize;
    let mut pos = prompt_tokens.len();

    // Interleaved mode counters.
    let mut modality_budget = match config.mode {
        AudioMode::Interleaved => 6, // start with 6 text tokens
        AudioMode::Sequential => usize::MAX,
    };
    let mut text_done = false;

    let mut next_token = sampler.sample(&mut logits);

    // Track consecutive audio segments after text_done to detect trailing garbage.
    // When the model finishes text (text_done) but doesn't emit audio_end cleanly,
    // we cap the number of trailing audio segments to avoid infinite generation.
    let mut trailing_audio_segments: usize = 0;
    const MAX_TRAILING_AUDIO_SEGMENTS: usize = 3;

    'outer: loop {
        if generated >= config.max_tokens || pos >= model_config.max_seq_len {
            break;
        }

        if modality == Modality::Text {
            // Check for EOG.
            if tokenizer.eos_token() == Some(next_token) {
                break;
            }

            // Sequential mode: switch on audio_start token.
            if next_token == TOKEN_AUDIO_START {
                modality = Modality::Audio;
                modality_budget = match config.mode {
                    AudioMode::Interleaved => 12,
                    AudioMode::Sequential => usize::MAX,
                };
                continue;
            }

            if next_token == TOKEN_TEXT_END {
                text_done = true;
            }

            // Emit text token.
            if next_token != TOKEN_TEXT_END {
                let piece = tokenizer.decode(&[next_token]);
                text_callback(&piece);
                text_tokens += 1;
            }

            generated += 1;
            modality_budget = modality_budget.saturating_sub(1);

            if generated >= config.max_tokens {
                break;
            }

            // Interleaved: check budget AFTER consuming the current token.
            // When budget hits 0, use forward_embedding on this token to
            // extract the audio embedding. This matches the reference which
            // extracts from the decode of the LAST text token.
            if matches!(config.mode, AudioMode::Interleaved) && (modality_budget == 0 || text_done)
            {
                if text_done {
                    trailing_audio_segments += 1;
                    if trailing_audio_segments > MAX_TRAILING_AUDIO_SEGMENTS {
                        break;
                    }
                }

                let mut emb = model.forward_embedding(&[next_token], pos, &mut state);
                pos += 1;

                modality = Modality::Audio;
                modality_budget = 12;

                // Run audio loop with this embedding.
                loop {
                    let outcome = decoder.decode_frame(&emb);
                    let audio_emb = match outcome {
                        FrameOutcome::End => {
                            text_done = true;
                            break;
                        }
                        FrameOutcome::Codes { audio_embedding } => audio_embedding,
                    };
                    modality_budget = modality_budget.saturating_sub(1);

                    if generated >= config.max_tokens || pos >= model_config.max_seq_len {
                        break;
                    }
                    if modality_budget == 0 && !text_done {
                        // Switch back to text. The reference transitions by
                        // decoding the last audio code embedding and sampling
                        // text from those logits (not by injecting TEXT_END).
                        logits = model.forward_from_embedding(&audio_emb, pos, &mut state);
                        next_token = sampler.sample(&mut logits);
                        pos += 1;
                        break;
                    }

                    emb = model.forward_hidden_from_embedding(&audio_emb, pos, &mut state);
                    pos += 1;
                    generated += 1;
                }

                // Switch back to text.
                modality = Modality::Text;
                modality_budget = 6;
                continue;
            }

            // Normal text: forward and sample next token.
            logits = model.forward(&[next_token], pos, &mut state);
            next_token = sampler.sample(&mut logits);
            pos += 1;
        } else {
            // Sequential audio mode: embedding from the audio_start token.
            let mut emb = model.forward_embedding(&[next_token], pos, &mut state);
            // The output norm naturally produces the right scale (~0.14 RMS)
            // when the hidden state has the activation outlier at channel 1455.
            pos += 1;
            generated += 1;

            loop {
                let outcome = decoder.decode_frame(&emb);
                let audio_emb = match outcome {
                    FrameOutcome::End => match config.mode {
                        AudioMode::Sequential => {
                            // Sequential TTS: audio is the final output.
                            // Returning to text + forwarding TEXT_END +
                            // resampling produces runaway garbage tokens (the
                            // model has nothing useful left to say) until
                            // max_tokens caps. Exit cleanly.
                            break 'outer;
                        }
                        AudioMode::Interleaved => {
                            // Interleaved: transition back to text. The
                            // trailing-audio-segments cap
                            // (MAX_TRAILING_AUDIO_SEGMENTS) bounds runaway
                            // post-audio cycles in the text branch above.
                            // Reachable here when the model emits
                            // TOKEN_AUDIO_START explicitly (line ~273) — the
                            // text-branch's budget-driven Interleaved block
                            // has its own audio loop that handles End inline.
                            modality = Modality::Text;
                            text_done = true;
                            modality_budget = 6;
                            logits = model.forward(&[TOKEN_TEXT_END], pos, &mut state);
                            next_token = sampler.sample(&mut logits);
                            pos += 1;
                            break;
                        }
                    },
                    FrameOutcome::Codes { audio_embedding } => audio_embedding,
                };
                modality_budget = modality_budget.saturating_sub(1);

                // Feed codes back as embedding → next hidden state.
                emb = model.forward_hidden_from_embedding(&audio_emb, pos, &mut state);
                pos += 1;
                generated += 1;

                if generated >= config.max_tokens || pos >= model_config.max_seq_len {
                    break;
                }
                if matches!(config.mode, AudioMode::Interleaved)
                    && modality_budget == 0
                    && !text_done
                {
                    modality = Modality::Text;
                    modality_budget = 6;
                    logits = model.forward_from_embedding(&audio_emb, pos, &mut state);
                    next_token = sampler.sample(&mut logits);
                    pos += 1;
                    break;
                }
            }
        }
    }

    // Batch ISTFT: all accumulated spectrum → PCM in one pass with proper overlap.
    let audio_samples = decoder.finish(&mut audio_callback);

    Ok(AudioGenerateResult {
        text_tokens,
        audio_frames: decoder.audio_frames,
        audio_samples,
        elapsed_secs: start.elapsed().as_secs_f64(),
        depthformer_secs: decoder.time_depthformer.as_secs_f64(),
        detokenizer_secs: decoder.time_detokenizer.as_secs_f64(),
    })
}