inference.py

from dataclasses import dataclass
from typing import Callable, List, Tuple
import torch
import safetensors.torch as st
from huggingface_hub import hf_hub_download

from model import EchoDiT
from autoencoder import build_ae, DAC

import torchaudio
from torchcodec.decoders import AudioDecoder

# from samplers import Sampler

SampleFn = Callable[
    [EchoDiT, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int],
    torch.Tensor
]

@dataclass
class SpeakerCache:
    """Cached speaker encoder output and per-layer KV projections for reuse across text prompts."""
    speaker_state: torch.Tensor          # (1, 640, 1280) - speaker encoder output
    speaker_mask: torch.Tensor           # (1, 640) - patched mask
    speaker_kvs: List[Tuple[torch.Tensor, torch.Tensor]]  # 24 layers × (k, v)
### Loading

def load_model_from_hf(repo_id: str = 'jordand/echo-tts-base', device: str = 'cuda', dtype: torch.dtype | None = torch.bfloat16, compile: bool = False, token: str | None = None) -> EchoDiT:
    with torch.device('meta'):
        model = EchoDiT(
            latent_size=80, model_size=2048, num_layers=24, num_heads=16,
            intermediate_size=5888, norm_eps=1e-5, max_seq_len=640,
            text_vocab_size=256, text_model_size=1280, text_num_layers=14,
            text_num_heads=10, text_intermediate_size=3328, text_max_seq_len=768,
            speaker_patch_size=4, speaker_model_size=1280, speaker_num_layers=14,
            speaker_num_heads=10, speaker_intermediate_size=3328,
            speaker_max_patched_seq_len=640, timestep_embed_size=512, adaln_rank=256,
        )
    w_path = hf_hub_download(repo_id, 'pytorch_model.safetensors', token=token)
    
    # Load to CPU first
    state = st.load_file(w_path, device='cpu')
    
    # Convert dtype on CPU if needed
    if dtype is not None:
        state = {k: v.to(dtype=dtype) for k, v in state.items()}
    
    # Now move to device
    state = {k: v.to(device=device) for k, v in state.items()}
    
    model.load_state_dict(state, strict=False, assign=True)
    model = model.eval()

    if compile:
        model = torch.compile(model)
        model.get_kv_cache = torch.compile(model.get_kv_cache)

    return model
    
def load_fish_ae_from_hf(repo_id: str = 'jordand/fish-s1-dac-min', device: str = 'cuda', dtype: torch.dtype | None = torch.float32, compile: bool = False, token: str | None = None) -> DAC:
    # have not tested lower precisions with fish AE yet
   
    with torch.device('meta'):
        fish_ae = build_ae()

    w_path = hf_hub_download(repo_id, 'pytorch_model.safetensors', token=token)

    # Load to CPU first (required for ZeroGPU compatibility)
    state = st.load_file(w_path, device='cpu')

    # Convert dtype if needed
    if dtype is not None:
        state = {k: v.to(dtype=dtype) for k, v in state.items()}

    # Move to device
    state = {k: v.to(device=device) for k, v in state.items()}
    fish_ae.load_state_dict(state, strict=False, assign=True)

    fish_ae = fish_ae.eval().to(device)

    if compile:
        fish_ae.encoder = torch.compile(fish_ae.encoder)
        fish_ae.decoder = torch.compile(fish_ae.decoder)
    
    return fish_ae


@dataclass
class PCAState:
    pca_components: torch.Tensor
    pca_mean: torch.Tensor
    latent_scale: float

def load_pca_state_from_hf(repo_id: str = 'jordand/echo-tts-base', device: str = 'cuda', filename: str = 'pca_state.safetensors', token: str | None = None) -> PCAState:
    p_path = hf_hub_download(repo_id, filename, token=token)
    # Load to CPU first (required for ZeroGPU compatibility)
    t = st.load_file(p_path, device='cpu')
    # Move tensors to device
    t = {k: v.to(device) for k, v in t.items()}
    return PCAState(
        pca_components=t["pca_components"],
        pca_mean=t["pca_mean"],
        latent_scale=float(t["latent_scale"].item()),
    )

### default load audio

def load_audio(path: str) -> torch.Tensor:

    decoder = AudioDecoder(path)
    sr = decoder.metadata.sample_rate
    audio = decoder.get_samples_played_in_range(0, 120)
    audio = audio.data.mean(dim=0).unsqueeze(0)
    audio = torchaudio.functional.resample(audio, sr, 44_100)
    audio = audio / torch.maximum(audio.abs().max(), torch.tensor(1.))
    # TODO is this better than clipping? should we target a specific energy level?
    return audio



### Text helpers

def tokenizer_encode(text: str, append_bos: bool = True, normalize: bool = True) -> torch.Tensor:

    if normalize:
        text = text.replace('…', '...')
        text = text.replace('“', '"')
        text = text.replace('”', '"')
        text = text.replace('’', "'")
        text = text.replace('\n', " ")
        text = text.replace(':', ',')
        text = text.replace(';', ',')

    b = list(text.encode('utf-8'))
    if append_bos:
        b.insert(0, 0)
    return torch.tensor(b)

def get_text_input_ids_and_mask(text_arr: List[str], max_length: int | None, device: str | None = None) -> tuple[torch.Tensor, torch.Tensor]:
    batch_size = len(text_arr)
    if max_length is None:
        max_length = max(len(tokenizer_encode(text)) for text in text_arr) # obviously bad...

    tokens = torch.zeros((batch_size, max_length), dtype=torch.int32)
    mask = torch.zeros((batch_size, max_length), dtype=torch.bool)
    
    for i, text in enumerate(text_arr):
        encoded = tokenizer_encode(text)
        length = min(len(encoded), max_length)
        tokens[i, :length] = encoded[:length]
        mask[i, :length] = 1

    if device is not None:
        tokens = tokens.to(device)
        mask = mask.to(device)

    return tokens, mask


### Autoencoder Inference

@torch.inference_mode()
def ae_encode(fish_ae: DAC, pca_state: PCAState, audio: torch.Tensor) -> torch.Tensor:
    assert audio.ndim == 3 and audio.shape[1] == 1 # (b, 1, length)
    z_q = fish_ae.encode_zq(audio).float()
    z_q = (z_q.transpose(1, 2) - pca_state.pca_mean) @ pca_state.pca_components.T
    z_q = z_q * pca_state.latent_scale
    return z_q

@torch.inference_mode()
def ae_decode(fish_ae: DAC, pca_state: PCAState, z_q: torch.Tensor) -> torch.Tensor:
    z_q = (z_q / pca_state.latent_scale) @ pca_state.pca_components + pca_state.pca_mean
    return fish_ae.decode_zq(z_q.transpose(1, 2).to(fish_ae.dtype)).float()

@torch.inference_mode()
def ae_reconstruct(fish_ae: DAC, pca_state: PCAState, audio: torch.Tensor) -> torch.Tensor:
    # (audio is (b, 1, length))
    z_q = ae_encode(fish_ae, pca_state, audio.to(fish_ae.dtype))
    return ae_decode(fish_ae, pca_state, z_q)


@torch.inference_mode()
def get_speaker_latent_and_mask(
    fish_ae: DAC,
    pca_state: PCAState,
    audio: torch.Tensor, # (1, length)
    max_speaker_latent_len: int = 2560, # pretrained max length
    audio_chunk_size: int = 640 * 2048 # (~30 seconds, 1/4 max speaker condition size)
) -> tuple[torch.Tensor, torch.Tensor]:

    # gets speaker latent and mask from audio, computes in chunks and concatenates (similar to pretraining setup)
    
    AE_DOWNSAMPLE_FACTOR = 2048
    max_audio_len = max_speaker_latent_len * AE_DOWNSAMPLE_FACTOR
    
    assert audio.ndim == 2 and audio.shape[0] == 1  # (1, length)
    audio = audio[:, :max_audio_len]
    audio_len = audio.shape[1]

    latent_arr = []
    
    for i in range(0, audio_len, audio_chunk_size):
        audio_chunk = audio[:, i:i + audio_chunk_size]
        if audio_chunk.shape[1] < audio_chunk_size:
            audio_chunk = torch.nn.functional.pad(audio_chunk, (0, audio_chunk_size - audio_chunk.shape[1]))

        latent_chunk = ae_encode(fish_ae, pca_state, audio_chunk.unsqueeze(0))
        latent_arr.append(latent_chunk)
    
    speaker_latent = torch.cat(latent_arr, dim=1)
    
    actual_latent_len = audio_len // AE_DOWNSAMPLE_FACTOR
    speaker_mask = (torch.arange(speaker_latent.shape[1], device=speaker_latent.device) < actual_latent_len).unsqueeze(0)
    
    if speaker_latent.shape[1] < max_speaker_latent_len:
        speaker_latent = torch.nn.functional.pad(speaker_latent, (0, 0, 0, max_speaker_latent_len - speaker_latent.shape[1]))
        speaker_mask = torch.nn.functional.pad(speaker_mask, (0, max_speaker_latent_len - speaker_mask.shape[1]))
    
    return speaker_latent, speaker_mask


### Speaker cache functions for batch generation

@torch.inference_mode()
def build_speaker_cache(
    model: EchoDiT,
    fish_ae: DAC,
    pca_state: PCAState,
    speaker_audio: torch.Tensor,
    pad_to_max_speaker_latent_len: int = 2560
) -> SpeakerCache:
    """Build speaker cache: encode speaker audio once and project into all layer KVs."""
    speaker_latent, speaker_mask = get_speaker_latent_and_mask(
        fish_ae, pca_state,
        speaker_audio.to(fish_ae.dtype),
        max_speaker_latent_len=pad_to_max_speaker_latent_len
    )
    speaker_state = model.speaker_encoder(speaker_latent.to(model.dtype))
    speaker_state = model.speaker_norm(speaker_state)
    speaker_mask_patched = speaker_mask[..., ::model.speaker_patch_size]

    # Project speaker_state into per-layer K/V and apply k_norm
    speaker_kvs = []
    for block in model.blocks:
        attn = block.attention
        bs = speaker_state.shape[0]
        xk = attn.wk_speaker(speaker_state).reshape(bs, speaker_state.shape[1], attn.num_heads, -1)
        xv = attn.wv_speaker(speaker_state).reshape(bs, speaker_state.shape[1], attn.num_heads, -1)
        xk = attn.k_norm(xk)  # Apply k_norm before caching
        speaker_kvs.append((xk, xv))

    return SpeakerCache(
        speaker_state=speaker_state,
        speaker_mask=speaker_mask_patched,
        speaker_kvs=speaker_kvs
    )


@torch.inference_mode()
def get_zero_speaker_cache(model: EchoDiT, device: str = 'cuda') -> SpeakerCache:
    """Compute speaker cache for zero speaker (all-zeros latent) - computed once at startup."""
    speaker_latent_zero = torch.zeros((1, 2560, 80), device=device, dtype=model.dtype)
    speaker_state_zero = model.speaker_encoder(speaker_latent_zero)
    speaker_state_zero = model.speaker_norm(speaker_state_zero)
    speaker_mask_zero = torch.zeros((1, 640), device=device, dtype=torch.bool)

    speaker_kvs = []
    for block in model.blocks:
        attn = block.attention
        bs = 1
        xk = attn.wk_speaker(speaker_state_zero).reshape(bs, speaker_state_zero.shape[1], attn.num_heads, -1)
        xv = attn.wv_speaker(speaker_state_zero).reshape(bs, speaker_state_zero.shape[1], attn.num_heads, -1)
        xk = attn.k_norm(xk)  # Apply k_norm before caching
        speaker_kvs.append((xk, xv))

    return SpeakerCache(
        speaker_state=speaker_state_zero,
        speaker_mask=speaker_mask_zero,
        speaker_kvs=speaker_kvs
    )


### Full sample pipeline

def find_flattening_point(data, target_value=0.0, window_size=20, std_threshold=0.05):
    padded_data = torch.cat([data, torch.zeros(window_size, *data.shape[1:], device=data.device, dtype=data.dtype)])
    for i in range(len(padded_data) - window_size):
        window = padded_data[i:i + window_size]
        if window.std() < std_threshold and abs(window.mean() - target_value) < 0.1:
            return i
    return len(data)


@torch.inference_mode()
def sample_pipeline(
    model: EchoDiT,
    fish_ae: DAC,
    pca_state: PCAState,
    sample_fn: SampleFn,
    text_prompt: str,
    speaker_audio: torch.Tensor | None,
    rng_seed: int,
    pad_to_max_speaker_latent_len: int | None = 2560,
    pad_to_max_text_seq_len: int | None = 768,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

    MAX_SPEAKER_LATENT_LEN = 2560
    MAX_TEXT_SEQ_LEN = 768

    device, dtype = model.device, model.dtype

    text_input_ids, text_mask = get_text_input_ids_and_mask([text_prompt], min(pad_to_max_text_seq_len or MAX_TEXT_SEQ_LEN, MAX_TEXT_SEQ_LEN), device=device)

    # print('initial text input ids length: ', text_input_ids.shape[1])
    # torch.cuda.synchronize()

    # import time
    
    # t0 = time.time()

    if speaker_audio is None:
        # No speaker prompt - use zero speaker latent and mask
        speaker_latent = torch.zeros((1, pad_to_max_speaker_latent_len if pad_to_max_speaker_latent_len else MAX_SPEAKER_LATENT_LEN, 80), device=device, dtype=dtype)
        speaker_mask = torch.zeros((1, pad_to_max_speaker_latent_len if pad_to_max_speaker_latent_len else MAX_SPEAKER_LATENT_LEN), device=device, dtype=torch.bool)
        # print("Using zero speaker latent and mask (no speaker prompt)")
    else:
        speaker_latent, speaker_mask = get_speaker_latent_and_mask(
            fish_ae, 
            pca_state, 
            speaker_audio.to(fish_ae.dtype), 
            max_speaker_latent_len=pad_to_max_speaker_latent_len if pad_to_max_speaker_latent_len else MAX_SPEAKER_LATENT_LEN
        )
        speaker_latent = speaker_latent.to(device)
        speaker_mask = speaker_mask.to(device)
        
        # print('speaker latent shape: ', speaker_latent.shape)
        # print('speaker mask shape: ', speaker_mask.shape)

    # torch.cuda.synchronize()
    # t1 = time.time()
    # print(f"Time taken encode: {t1 - t0} seconds")

    latent_out = sample_fn(model, speaker_latent, speaker_mask, text_input_ids, text_mask, rng_seed)

    # torch.cuda.synchronize()
    # t2 = time.time()

    # print(f"Time taken sample: {t2 - t1} seconds")

    audio_out = ae_decode(fish_ae, pca_state, latent_out)
    # torch.cuda.synchronize()
    # t3 = time.time()
    # print(f"Time taken decode: {t3 - t2} seconds")

    flattening_point = find_flattening_point(latent_out[0])
    audio_out = audio_out[..., :flattening_point * 2048]

    # print(f"\nTime taken total: {t3 - t0} seconds")

    # peak_mem = torch.cuda.max_memory_allocated()
    # print(f"Peak memory: {peak_mem / 1024**2:.2f} MB")
    # print(torch.cuda.memory_summary(abbreviated=True))

    return audio_out


@torch.inference_mode()
def build_kv_cache_from_speaker_cache(
    model: EchoDiT,
    speaker_cache: SpeakerCache,
    zero_speaker_cache: SpeakerCache,
    text_input_ids: torch.Tensor,
    text_mask: torch.Tensor,
    guidance_mode: str = "independent"
) -> tuple[List[Tuple[torch.Tensor, torch.Tensor]], torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Build full KV cache for CFG from cached speaker KVs + fresh text encoding.

    Returns: (kv_cache, combined_text_mask, combined_speaker_mask, text_mask, text_mask_uncond)
    The last two are for alternating mode which needs per-step switching.
    """
    device = model.device

    # Text encoding - the only variable cost per text
    text_state = model.text_encoder(text_input_ids, text_mask)
    text_state = model.text_norm(text_state)

    # Unconditional text (BOS only)
    text_input_ids_uncond = torch.zeros_like(text_input_ids)
    text_mask_uncond = torch.zeros_like(text_mask)
    text_mask_uncond[:, 0] = True
    text_state_uncond = model.text_encoder(text_input_ids_uncond, text_mask_uncond)
    text_state_uncond = model.text_norm(text_state_uncond)

    # Assemble CFG batch based on guidance mode
    kv_cache = []
    for i, block in enumerate(model.blocks):
        attn = block.attention
        bs = text_state.shape[0]

        def _text_kv(ts):
            xk = attn.wk_text(ts).reshape(bs, ts.shape[1], attn.num_heads, -1)
            xv = attn.wv_text(ts).reshape(bs, ts.shape[1], attn.num_heads, -1)
            xk = attn.k_norm(xk)  # Apply k_norm to text KVs
            return xk, xv

        xk_txt_cond, xv_txt_cond = _text_kv(text_state)
        xk_txt_uncond, xv_txt_uncond = _text_kv(text_state_uncond)

        # Speaker KVs are already k_norm'd in the cache
        xk_sp_cond, xv_sp_cond = speaker_cache.speaker_kvs[i]
        xk_sp_zero, xv_sp_zero = zero_speaker_cache.speaker_kvs[i]

        if guidance_mode in ["independent", "apg-independent"]:
            # 3× batch: [cond, text_uncond, speaker_uncond]
            xk = torch.cat([
                torch.cat([xk_txt_cond, xk_sp_cond], dim=1),    # batch 0: cond
                torch.cat([xk_txt_uncond, xk_sp_cond], dim=1),  # batch 1: text-uncond
                torch.cat([xk_txt_cond, xk_sp_zero], dim=1),    # batch 2: speaker-uncond
            ], dim=0)
            xv = torch.cat([
                torch.cat([xv_txt_cond, xv_sp_cond], dim=1),
                torch.cat([xv_txt_uncond, xv_sp_cond], dim=1),
                torch.cat([xv_txt_cond, xv_sp_zero], dim=1),
            ], dim=0)
        elif guidance_mode == "alternating":
            # 2× batch: [cond, cond] - note: using cond text in both, will switch masks per step
            # This follows the original implementation which has known limitations (see samplers.py line 337):
            # Uses text_cond with alternating masks rather than true text_uncond encoding.
            # Works because text encoder is bidirectional (not causal), but is "technically incorrect"
            # per original author. In practice: "might not make much of a difference."
            # No presets use alternating mode. Could be improved by building separate KV caches per step,
            # but that would defeat the speaker cache optimization.
            xk = torch.cat([
                torch.cat([xk_txt_cond, xk_sp_cond], dim=1),     # batch 0: cond
                torch.cat([xk_txt_cond, xk_sp_zero], dim=1),     # batch 1: text_cond + speaker_zero (masks will alternate)
            ], dim=0)
            xv = torch.cat([
                torch.cat([xv_txt_cond, xv_sp_cond], dim=1),
                torch.cat([xv_txt_cond, xv_sp_zero], dim=1),
            ], dim=0)
        else:  # joint-unconditional
            # 2× batch: [cond, joint_uncond]
            xk = torch.cat([
                torch.cat([xk_txt_cond, xk_sp_cond], dim=1),     # batch 0: cond
                torch.cat([xk_txt_uncond, xk_sp_zero], dim=1),   # batch 1: joint uncond
            ], dim=0)
            xv = torch.cat([
                torch.cat([xv_txt_cond, xv_sp_cond], dim=1),
                torch.cat([xv_txt_uncond, xv_sp_zero], dim=1),
            ], dim=0)

        # k_norm already applied to both text and speaker KVs before concatenation
        kv_cache.append((xk, xv))

    # Combined masks based on mode
    # When using prebuilt kv_cache, we need to return a single combined cross-attention mask
    # (text + speaker concatenated), not separate masks
    if guidance_mode in ["independent", "apg-independent"]:
        text_masks = torch.cat([text_mask, text_mask_uncond, text_mask], dim=0)
        speaker_masks = torch.cat([
            speaker_cache.speaker_mask,
            speaker_cache.speaker_mask,
            zero_speaker_cache.speaker_mask
        ], dim=0)
    else:  # alternating or joint
        text_masks = torch.cat([text_mask, text_mask_uncond], dim=0)
        speaker_masks = torch.cat([
            speaker_cache.speaker_mask,
            zero_speaker_cache.speaker_mask
        ], dim=0)

    # Concatenate text and speaker masks into single cross-attention mask
    combined_cross_mask = torch.cat([text_masks, speaker_masks], dim=1)
    # Return empty mask for speaker_mask parameter (model will concat [self, cross, empty])
    empty_mask = torch.zeros((combined_cross_mask.shape[0], 0), dtype=torch.bool, device=device)

    return kv_cache, combined_cross_mask, empty_mask, text_mask, text_mask_uncond