inference_blockwise.py

from typing import List

import torch

from inference import (
    KVCache,
    _concat_kv_caches,
    _multiply_kv_cache,
    _temporal_score_rescale,
)
from model import EchoDiT


@torch.inference_mode()
def sample_blockwise_euler_cfg_independent_guidances(
    model: EchoDiT,
    speaker_latent: torch.Tensor,
    speaker_mask: torch.Tensor,
    text_input_ids: torch.Tensor,
    text_mask: torch.Tensor,
    rng_seed: int,
    block_sizes: List[int],
    num_steps: int,
    cfg_scale_text: float,
    cfg_scale_speaker: float,
    cfg_min_t: float,
    cfg_max_t: float,
    truncation_factor: float | None,
    rescale_k: float | None,
    rescale_sigma: float | None,
    speaker_kv_scale: float | None,
    speaker_kv_max_layers: int | None,
    speaker_kv_min_t: float | None,
    continuation_latent: torch.Tensor | None = None,
) -> torch.Tensor:

    INIT_SCALE = 0.999  # so that we can apply rescale to first step

    device, dtype = model.device, model.dtype
    batch_size = text_input_ids.shape[0]

    rng = torch.Generator(device=device).manual_seed(rng_seed)

    t_schedule = torch.linspace(1., 0., num_steps + 1, device=device) * INIT_SCALE

    text_mask_uncond = torch.zeros_like(text_mask)
    speaker_mask_uncond = torch.zeros_like(speaker_mask)

    kv_text_cond = model.get_kv_cache_text(text_input_ids, text_mask)
    kv_speaker_cond = model.get_kv_cache_speaker(speaker_latent.to(dtype))

    # masks prevent decoder from attending to unconds:
    kv_text_full = _concat_kv_caches(kv_text_cond, kv_text_cond, kv_text_cond)
    kv_speaker_full = _concat_kv_caches(kv_speaker_cond, kv_speaker_cond, kv_speaker_cond)

    full_text_mask = torch.cat([text_mask, text_mask_uncond, text_mask], dim=0)
    full_speaker_mask = torch.cat([speaker_mask, speaker_mask, speaker_mask_uncond], dim=0)

    prefix_latent = torch.zeros((batch_size, sum(block_sizes) , 80), device=device, dtype=torch.float32)

    start_pos = 0
    if continuation_latent is not None:
        continuation_len = continuation_latent.shape[1]
        prefix_latent = torch.cat([continuation_latent, prefix_latent], dim=1)
        start_pos = continuation_len

    for block_size in block_sizes:
        if speaker_kv_scale is not None:
            _multiply_kv_cache(kv_speaker_cond, speaker_kv_scale, speaker_kv_max_layers)
            kv_speaker_full = _concat_kv_caches(kv_speaker_cond, kv_speaker_cond, kv_speaker_cond)

        full_prefix_latent = torch.cat([prefix_latent, prefix_latent, prefix_latent], dim=0)
        kv_latent_full = model.get_kv_cache_latent(full_prefix_latent.to(dtype))
        kv_latent_cond = [(k[:batch_size], v[:batch_size]) for k, v in kv_latent_full]

        x_t = torch.randn((batch_size, block_size, 80), device=device, dtype=torch.float32, generator=rng)
        if truncation_factor is not None:
            x_t = x_t * truncation_factor

        for i in range(num_steps):
            t, t_next = t_schedule[i], t_schedule[i + 1]

            has_cfg = ((t >= cfg_min_t) * (t <= cfg_max_t)).item()

            if has_cfg:
                v_cond, v_uncond_text, v_uncond_speaker = model(
                    x=torch.cat([x_t, x_t, x_t], dim=0).to(dtype),
                    t=(torch.ones((batch_size * 3,), device=device) * t).to(dtype),
                    text_mask=full_text_mask,
                    speaker_mask=full_speaker_mask,
                    start_pos=start_pos,
                    kv_cache_text=kv_text_full,
                    kv_cache_speaker=kv_speaker_full,
                    kv_cache_latent=kv_latent_full,
                ).float().chunk(3, dim=0)
                v_pred = v_cond + cfg_scale_text * (v_cond - v_uncond_text) + cfg_scale_speaker * (v_cond - v_uncond_speaker)
            else:
                v_pred = model(
                    x=x_t.to(dtype),
                    t=(torch.ones((batch_size,), device=device) * t).to(dtype),
                    text_mask=text_mask,
                    speaker_mask=speaker_mask,
                    start_pos=start_pos,
                    kv_cache_text=kv_text_cond,
                    kv_cache_speaker=kv_speaker_cond,
                    kv_cache_latent=kv_latent_cond,
                ).float()

            # optional temporal score rescaling: https://arxiv.org/pdf/2510.01184
            if rescale_k is not None and rescale_sigma is not None:
                v_pred = _temporal_score_rescale(v_pred, x_t, t, rescale_k, rescale_sigma)

            # optional kv speaker scaling
            if speaker_kv_scale is not None and t_next < speaker_kv_min_t and t >= speaker_kv_min_t:
                _multiply_kv_cache(kv_speaker_cond, 1. / speaker_kv_scale, speaker_kv_max_layers)
                kv_speaker_full = _concat_kv_caches(kv_speaker_cond, kv_speaker_cond, kv_speaker_cond)

            x_t = x_t + v_pred * (t_next - t)

        prefix_latent[:, start_pos:start_pos + block_size] = x_t
        start_pos += block_size

    return prefix_latent


if __name__ == "__main__":
    import torchaudio
    from inference import (
        load_model_from_hf,
        load_fish_ae_from_hf,
        load_pca_state_from_hf,
        load_audio,
        get_text_input_ids_and_mask,
        get_speaker_latent_and_mask,
        ae_encode,
        ae_decode,
        crop_audio_to_flattening_point,
    )

    model = load_model_from_hf()
    fish_ae = load_fish_ae_from_hf()
    pca_state = load_pca_state_from_hf()


    # example 1, generate 320 in three blocks

    speaker_audio_path = "/path/to/speaker/audio.wav"
    speaker_audio = load_audio(speaker_audio_path).cuda()
    speaker_latent, speaker_mask = get_speaker_latent_and_mask(fish_ae, pca_state, speaker_audio)

    text = "[S1] Alright, I'm going to demo this new model called Echo TTS."
    text_input_ids, text_mask = get_text_input_ids_and_mask([text], max_length=None, device="cuda")

    latent_out = sample_blockwise_euler_cfg_independent_guidances(
        model=model,
        speaker_latent=speaker_latent,
        speaker_mask=speaker_mask,
        text_input_ids=text_input_ids,
        text_mask=text_mask,
        rng_seed=0,
        block_sizes=[128, 128, 64], # (sums to 320, so will be ~15 seconds; supports up to 640)
        num_steps=40,
        cfg_scale_text=3.0,
        cfg_scale_speaker=5.0,
        cfg_min_t=0.5,
        cfg_max_t=1.0,
        truncation_factor=0.8,
        rescale_k=None,
        rescale_sigma=None,
        speaker_kv_scale=None,
        speaker_kv_max_layers=None,
        speaker_kv_min_t=None,
    )
    audio_out = ae_decode(fish_ae, pca_state, latent_out)
    audio_out = crop_audio_to_flattening_point(audio_out, latent_out[0])
    torchaudio.save("output_blockwise.wav", audio_out[0].cpu(), 44100)



    # ___________________________________________________________
    # example 2: with continuation latent (use same speaker audio as first example, generate from partial output of first example)

    continuation_audio_path = "output_blockwise.wav" # can be any path
    continuation_audio = load_audio(continuation_audio_path).cuda()
    continuation_latent, continuation_mask = get_speaker_latent_and_mask(fish_ae, pca_state, continuation_audio)

    continuation_latent = continuation_latent[:, :continuation_mask.sum()]

    text = "[S1] Alright, I'm going to demo this new model called Echo TTS, and now, we're going to continue from the audio we already generated and add some more text." 
    # NOTE this MUST include the text from the continuation prefix. can use https://huggingface.co/jordand/whisper-d-v1a to get in-distribution transcription automatically.

    text_input_ids, text_mask = get_text_input_ids_and_mask([text], max_length=None, device="cuda")

    continuation_block_sizes = [256] # (generate up to 12 more seconds)
    # NOTE: these do not include the continuation latent length, so sum(block_sizes) + continuation_latent.shape[1] should be < 640 (to be in-distribution with training data)

    latent_out_continued = sample_blockwise_euler_cfg_independent_guidances(
        model=model,
        speaker_latent=speaker_latent,
        speaker_mask=speaker_mask,
        text_input_ids=text_input_ids,
        text_mask=text_mask,
        rng_seed=0,
        block_sizes=continuation_block_sizes,
        num_steps=40,
        cfg_scale_text=3.0,
        cfg_scale_speaker=3.0,
        cfg_min_t=0.5,
        cfg_max_t=1.0,
        truncation_factor=0.8,
        rescale_k=None,
        rescale_sigma=None,
        speaker_kv_scale=None,
        speaker_kv_max_layers=None,
        speaker_kv_min_t=None,
        continuation_latent=continuation_latent,
    )
    audio_out_continued = ae_decode(fish_ae, pca_state, latent_out_continued)
    audio_out_continued = crop_audio_to_flattening_point(audio_out_continued, latent_out_continued[0])
    torchaudio.save("output_blockwise_continued.wav", audio_out_continued[0].cpu(), 44100)