modules/audio_detokenizer/semantic_fm_prefix_streaming.py

import yaml
import logging
import time

import os
import torch

from modules.audio_detokenizer.flow_matching.ode_wrapper import StreamingODEWrapperForPrefix
from modules.audio_detokenizer.flow_matching.model import DiTPrefix
from modules.audio_detokenizer.flow_matching.scheduler import StreamingFlowMatchingScheduler


logger = logging.getLogger(__name__)


class StreamingSemanticFMWrapper:
    def __init__(self, speech_model: DiTPrefix, max_kv_cache_tokens=900, max_prompt_chunk=2,
                 use_cfg=True, use_cfg_rescale=True, cfg_init=1.5, cfg_scale=7.5, cfg_schedule="linear", cfg_token_id=0, 
                 normalize_mel=False, mel_mean=None, mel_std=None, device: torch.device = torch.device("cpu")) -> None:
        
        self.dtype = torch.bfloat16
        self.speech_model = speech_model.to(device).to(self.dtype)
        self.speech_model = self.speech_model.eval()
        self.device = device
        self.normalize_mel = normalize_mel
        self.mel_mean = mel_mean
        self.mel_std = mel_std

        self.use_cfg = use_cfg
        self.use_cfg_rescale = use_cfg_rescale
        self.cfg_init = cfg_init
        self.cfg_scale = cfg_scale
        self.cfg_schedule = cfg_schedule
        
        self.incremental_state = {}
        self.condition_cache = {"previous_seqlen": 0}

        logger.info(f">>> SemanticFMWrapper initialized with use_cfg={use_cfg}, use_cfg_rescale={use_cfg_rescale}, cfg_init={cfg_init}, cfg_scale={cfg_scale}, cfg_schedule={cfg_schedule}")

        self.scheduler = StreamingFlowMatchingScheduler()
        self.ode_wrapper = StreamingODEWrapperForPrefix(net=self.speech_model, x_mask=None, x_cond=None,
                                      use_cfg=use_cfg, use_cfg_rescale=use_cfg_rescale, cfg_init=cfg_init, cfg_scale=cfg_scale, cfg_schedule=cfg_schedule, cfg_token_id=cfg_token_id)
    
        self.max_kv_cache_tokens = max_kv_cache_tokens
        self.max_prompt_chunk = max_prompt_chunk
        self.reserve_kv_cache_tokens = 0

    @torch.inference_mode()
    def infer_chunk(self, xt_chunk, semantic_tokens_chunk, start_position_id, 
                    cache = None, look_ahead_tokens=0,
                    ode_steps=15, verbose=False, ode_solver="neural_ode_euler"):
        """
            semantic_tokens: [T_1], torch.LongTensor
            xt: [T_2, 80], torch.Tensor, DO NOT normalize it outside
            ode_steps: int, number of ode steps, default 15
            verbose: bool, default False
            ode_solver: str, ode solver, expected in ("neural_ode_euler", "naive_euler"), default "neural_ode_euler"
        """
        bs = 1

        self.scheduler.set_timesteps(ode_steps)

        semantic_tokens_chunk = semantic_tokens_chunk.unsqueeze(0).to(self.device)
        xt_chunk = xt_chunk.unsqueeze(0).to(self.device).to(self.dtype)

        t_span = torch.linspace(0, 1, self.scheduler.timesteps)

        x_mask = torch.zeros(bs, xt_chunk.shape[1], device=self.device).bool()
        
        cache_ret = self.ode_wrapper.set_conditions(x_mask=x_mask, x_cond=semantic_tokens_chunk, start_position_id=start_position_id, cache=self.condition_cache)

        if verbose:
            t_start = time.time()
        if ode_solver == "neural_ode_euler":
            x_t = self.scheduler.sample_by_neuralode(self.ode_wrapper, time_steps=t_span, xt=xt_chunk, verbose=False)
        elif ode_solver == "naive_euler":
            x_t = self.scheduler.sample(ode_wrapper=self.ode_wrapper, time_steps=t_span, xt=xt_chunk, verbose=False)
        else:
            raise NotImplementedError("ode_solver should be in ('neural_ode_euler', 'naive_euler')")
        
        if look_ahead_tokens > 0:
            semantic_tokens_left = semantic_tokens_chunk.view(-1)[-look_ahead_tokens:]
            cache["semantic_token"] = semantic_tokens_left
            x_t_ret = x_t[:, :-look_ahead_tokens, :]
        else:
            x_t_ret = x_t

        if look_ahead_tokens > 0:
            x_mask = torch.zeros(bs, xt_chunk.shape[1] - look_ahead_tokens, device=self.device).bool()
            self.condition_cache = self.ode_wrapper.set_conditions(x_mask=x_mask, x_cond=semantic_tokens_chunk[:, :-look_ahead_tokens], start_position_id=start_position_id, cache=self.condition_cache)
            self.ode_wrapper(torch.Tensor([0.999]).to(x_t_ret.device), x_t_ret)
        else:
            self.condition_cache = cache_ret

        if verbose:
            t_end = time.time()
            logger.info(f"[ODE Chunk] Time cost: {t_end - t_start}")

        if self.normalize_mel:
            x_t_ret = x_t_ret * self.mel_std + self.mel_mean
        return x_t_ret.squeeze(0)


    @torch.inference_mode()
    def infer_mel(self, semantic_tokens, ode_steps=15, chunk_size=150, verbose=False, ode_solver="neural_ode_euler"):
        """
            semantic_tokens: [T_1], torch.LongTensor
            prompt: [T_2, 80], torch.Tensor, DO NOT normalize it outside
            prompt_semantic_tokens, [T_2], torch.LongTensor
            ode_steps: int, number of ode steps, default 15
            verbose: bool, default False
            ode_solver: str, ode solver, expected in ("neural_ode_euler", "naive_euler"), default "neural_ode_euler"
        """
        assert semantic_tokens.dim() == 1

        x_t = torch.randn(semantic_tokens.shape[0], 80).to(self.device).to(self.dtype)

        seq_len = semantic_tokens.shape[0]

        num_chunks = seq_len // chunk_size
        if seq_len % chunk_size != 0:
            num_chunks += 1

        x_pred_collect = []

        if verbose:
            t_start = time.time()

        for chunk_id in range(num_chunks):
            start = chunk_id * chunk_size
            end = min(start + chunk_size, seq_len)
            semantic_tokens_chunk = semantic_tokens[start:end]
            x_t_chunk = x_t[start:end, :]

            x_pred = self.infer_chunk(xt_chunk=x_t_chunk, semantic_tokens_chunk=semantic_tokens_chunk, start_position_id=self.start_position_id,
                                      ode_steps=ode_steps, verbose=verbose, ode_solver=ode_solver)
            self.start_position_id += end - start
            self.update_incremental_state()

            x_pred_collect.append(x_pred)

        if verbose:
            t_end = time.time()
            logger.info(f"[ODE] Time cost: {t_end - t_start}")
        
        x_pred = torch.cat(x_pred_collect, dim=0)

        return x_pred
    
    def clear_all_states(self):
        self.start_position_id = 0
        self.condition_cache = {"previous_seqlen": 0}
        self.ode_wrapper.clear_all_states()
    
    def state_dict(self):
        return {
            "start_position_id": self.start_position_id,
            "ode_wrapper": self.ode_wrapper.state_dict(),
            "condition_cache": self.condition_cache
        }
    
    def load_state_dict(self, state_dict):
        if state_dict is not None:
            self.start_position_id = state_dict["start_position_id"]
            self.ode_wrapper.load_state_dict(state_dict["ode_wrapper"])
            self.condition_cache = state_dict["condition_cache"]
    
    def update_incremental_state(self):
        self.ode_wrapper.update_incremental_state(reserve_kv_cache_tokens=0, max_kv_cache_tokens=self.max_kv_cache_tokens, condition_cache=self.condition_cache)
    
    @torch.inference_mode()
    def prefill(self, mel, semantic_token, chunk_size=150, verbose=False):
        """
            mel: [T, 80], torch.Tensor
            semantic_token: [T], torch.LongTensor
            chunk_size: int, default 150
        """
        assert mel.dim() == 2
        assert semantic_token.dim() == 1
        assert semantic_token.shape[0] == mel.shape[0], "Semantic token and mel shape mismatch"
        seq_len = mel.shape[0]
        num_chunks = min(seq_len // chunk_size, self.max_prompt_chunk)
        start_pos = seq_len - num_chunks * chunk_size
        
        res_mel = mel[:start_pos, :]
        res_semantic_token = semantic_token[:start_pos]
        self.prefill_chunk(res_mel, res_semantic_token, start_position_id=self.start_position_id)
        self.start_position_id += start_pos
        self.update_incremental_state()
        self.reserve_kv_cache_tokens += self.ode_wrapper.kv_cache_tokens

        if verbose:
            logger.info("Prefilling prompt with {} chunks".format(num_chunks))
            start_time = time.time()

        for chunk_id in range(num_chunks):
            start = start_pos + chunk_id * chunk_size
            end = start + chunk_size
            mel_chunk = mel[start:end, :]
            semantic_token_chunk = semantic_token[start:end]

            self.prefill_chunk(mel_chunk, semantic_token_chunk, start_position_id=self.start_position_id)
            self.start_position_id += end - start
            
            self.update_incremental_state()
            self.reserve_kv_cache_tokens += self.ode_wrapper.kv_cache_tokens
        
        
        if verbose:
            logger.info("Prefilling done in {:.2f} seconds".format(time.time() - start_time))
    
    def prefill_chunk(self, mel_chunk, semantic_tokens_chunk, start_position_id=0):
        """
            mel_chunk: [T, 80], torch.Tensor, T is the chunk size
            semantic_tokens_chunk: [T], torch.LongTensor
            start_position_id: int, default 0
        """
        bs = 1

        semantic_tokens_chunk = semantic_tokens_chunk.unsqueeze(0).to(self.device)
        mel_chunk = mel_chunk.unsqueeze(0).to(self.device).to(self.dtype)

        if self.normalize_mel:
            mel_chunk = (mel_chunk - self.mel_mean) / self.mel_std

        x_mask = torch.zeros(bs, mel_chunk.shape[1], device=self.device).bool()
        
        self.condition_cache = self.ode_wrapper.set_conditions(x_mask=x_mask, x_cond=semantic_tokens_chunk, start_position_id=start_position_id, cache=self.condition_cache)

        x_t = torch.Tensor([0.999]).to(self.device)

        self.ode_wrapper(x_t, mel_chunk)

        
    @classmethod
    def from_pretrained(cls, model_config, ckpt_path, device, max_prompt_chunk=2, max_kv_cache_tokens=900, use_cfg=True, use_cfg_rescale=True, cfg_init=1.5, cfg_scale=7.5, cfg_schedule="linear"):

        # open yaml file
        with open(model_config, 'r') as f:
            config = yaml.safe_load(f)
        model_config = config["model"]["dit"]
        dit = DiTPrefix(
            input_size=model_config["input_size"],
            semantic_vocab_size=model_config["semantic_vocab_size"] + 1,
            hidden_size=model_config["hidden_size"],
            depth=model_config["depth"],
            num_heads=model_config["num_heads"],
            mlp_ratio=model_config["mlp_ratio"],
            ffn_type=model_config.get("ffn_type", "conv1d_conv1d"),
            ffn_gated_glu=model_config.get("ffn_gated_glu", True),
            ffn_act_layer=model_config.get("ffn_act_layer", "gelu"),
            ffn_conv_kernel_size=model_config.get("ffn_conv_kernel_size", 5),

            use_rope=model_config.get("use_rope", False),
            rope_params=model_config.get("rope_params", { "max_position_embeddings": 4096,"rope_base": 10000,"rope_interpolation_factor": 1 }),

            position_embedding_type=model_config["position_embedding_type"],
            max_seq_len=model_config["max_seq_len"],
            output_size=model_config["input_size"],
            prompt_cfg_dropout=0
        )
        cfg_semantic_token_id = model_config["semantic_vocab_size"]
        
        # load state_dict
        state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True)["state_dict"]
        speech_model_params = {k.replace("speech_model.", ""): v for k, v in state_dict.items() if "speech_model" in k}
        dit.load_state_dict(speech_model_params, strict=True)
        logger.info(f">>> Loaded checkpoint from {ckpt_path}")

        return cls(speech_model=dit, device=device, normalize_mel=config["normalize_mel"], mel_mean=config["mel_mean"], mel_std=config["mel_std"], max_prompt_chunk=max_prompt_chunk, max_kv_cache_tokens=max_kv_cache_tokens,
                   use_cfg=use_cfg, use_cfg_rescale=use_cfg_rescale, cfg_init=cfg_init, cfg_scale=cfg_scale, cfg_schedule=cfg_schedule, cfg_token_id=cfg_semantic_token_id)