jhqxxx
523 lines
21 KiB
Rust
523 lines
21 KiB
Rust
use crate::{
|
|
models::voxcpm::{
|
|
audio_vae::AudioVAE,
|
|
config::VoxCPMConfig,
|
|
minicpm4::MiniCPMModel,
|
|
model::{ScalarQuantizationLayer, UnifiedCFM, VoxCPMLocDiT, VoxCPMLocEnc},
|
|
},
|
|
utils::tensor_utils::masked_scatter_dim0,
|
|
};
|
|
use anyhow::Result;
|
|
use candle_core::{D, DType, Device, IndexOp, Tensor};
|
|
use candle_nn::{Linear, Module, VarBuilder, linear, linear_no_bias};
|
|
use rocket::async_stream::stream;
|
|
use rocket::futures::Stream;
|
|
|
|
pub struct VoxCPMModelRefact {
|
|
config: VoxCPMConfig,
|
|
patch_size: usize,
|
|
latent_dim: usize,
|
|
decode_patch_len: usize,
|
|
// audio_start_token: u32,
|
|
// // audio_end_token: u32,
|
|
// ref_audio_start_token: u32,
|
|
// ref_audio_end_token: u32,
|
|
// chunk_size: usize,
|
|
// sample_rate: usize,
|
|
base_lm: MiniCPMModel,
|
|
residual_lm: MiniCPMModel,
|
|
feat_encoder: VoxCPMLocEnc,
|
|
feat_decoder: UnifiedCFM,
|
|
fsq_layer: ScalarQuantizationLayer,
|
|
enc_to_lm_proj: Linear,
|
|
lm_to_dit_proj: Linear,
|
|
res_to_dit_proj: Linear,
|
|
fusion_concat_proj: Option<Linear>,
|
|
stop_proj: Linear,
|
|
stop_head: Linear,
|
|
device: Device,
|
|
dtype: DType,
|
|
}
|
|
|
|
impl VoxCPMModelRefact {
|
|
pub fn new(
|
|
vb: VarBuilder,
|
|
config: VoxCPMConfig,
|
|
latent_dim: usize,
|
|
decode_chunk_size: usize,
|
|
) -> Result<Self> {
|
|
let base_lm = MiniCPMModel::new(vb.pp("base_lm"), config.lm_config.clone())?;
|
|
let mut residual_lm_config = config.lm_config.clone();
|
|
residual_lm_config.num_hidden_layers = config.residual_lm_num_layers;
|
|
residual_lm_config.vocab_size = 0;
|
|
residual_lm_config.no_rope = config.residual_lm_no_rope;
|
|
let residual_lm = MiniCPMModel::new(vb.pp("residual_lm"), residual_lm_config)?;
|
|
let mut encoder_config = config.lm_config.clone();
|
|
encoder_config.hidden_size = config.encoder_config.hidden_dim;
|
|
encoder_config.intermediate_size = config.encoder_config.ffn_dim;
|
|
encoder_config.num_attention_heads = config.encoder_config.num_heads;
|
|
encoder_config.num_hidden_layers = config.encoder_config.num_layers;
|
|
encoder_config.kv_channels = config.encoder_config.kv_channels;
|
|
encoder_config.vocab_size = 0;
|
|
let feat_encoder =
|
|
VoxCPMLocEnc::new(vb.pp("feat_encoder"), encoder_config, config.feat_dim)?;
|
|
|
|
let mut decoder_config = config.lm_config.clone();
|
|
decoder_config.hidden_size = config.dit_config.hidden_dim;
|
|
decoder_config.intermediate_size = config.dit_config.ffn_dim;
|
|
decoder_config.num_attention_heads = config.dit_config.num_heads;
|
|
decoder_config.num_hidden_layers = config.dit_config.num_layers;
|
|
decoder_config.kv_channels = config.dit_config.kv_channels;
|
|
decoder_config.vocab_size = 0;
|
|
let estimator = VoxCPMLocDiT::new(
|
|
vb.pp("feat_decoder.estimator"),
|
|
decoder_config,
|
|
config.feat_dim,
|
|
)?;
|
|
let feat_decoder = UnifiedCFM::new(
|
|
config.feat_dim,
|
|
config.dit_config.cfm_config.clone(),
|
|
estimator,
|
|
false,
|
|
// config.architecture.clone(),
|
|
)?;
|
|
let fsq_layer = ScalarQuantizationLayer::new(
|
|
vb.pp("fsq_layer"),
|
|
config.lm_config.hidden_size,
|
|
config.lm_config.hidden_size,
|
|
config.scalar_quantization_latent_dim,
|
|
config.scalar_quantization_scale,
|
|
)?;
|
|
let enc_to_lm_proj = linear(
|
|
config.encoder_config.hidden_dim,
|
|
config.lm_config.hidden_size,
|
|
vb.pp("enc_to_lm_proj"),
|
|
)?;
|
|
let lm_to_dit_proj = linear(
|
|
config.lm_config.hidden_size,
|
|
config.dit_config.hidden_dim,
|
|
vb.pp("lm_to_dit_proj"),
|
|
)?;
|
|
let res_to_dit_proj = linear(
|
|
config.lm_config.hidden_size,
|
|
config.dit_config.hidden_dim,
|
|
vb.pp("res_to_dit_proj"),
|
|
)?;
|
|
|
|
let fusion_concat_proj = if config.architecture.to_lowercase().eq("voxcpm2") {
|
|
Some(linear(
|
|
config.lm_config.hidden_size * 2,
|
|
config.lm_config.hidden_size,
|
|
vb.pp("fusion_concat_proj"),
|
|
)?)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let stop_proj = linear(
|
|
config.lm_config.hidden_size,
|
|
config.lm_config.hidden_size,
|
|
vb.pp("stop_proj"),
|
|
)?;
|
|
let stop_head = linear_no_bias(config.lm_config.hidden_size, 2, vb.pp("stop_head"))?;
|
|
|
|
let patch_size = config.patch_size;
|
|
let decode_patch_len = patch_size * decode_chunk_size;
|
|
Ok(Self {
|
|
config,
|
|
patch_size,
|
|
latent_dim,
|
|
decode_patch_len,
|
|
// audio_start_token: 101,
|
|
// // audio_end_token: 102,
|
|
// ref_audio_start_token: 103,
|
|
// ref_audio_end_token: 104,
|
|
// chunk_size: audio_vae.chunk_size,
|
|
// sample_rate: audio_vae.sample_rate,
|
|
// tokenizer,
|
|
// audio_vae,
|
|
base_lm,
|
|
residual_lm,
|
|
feat_encoder,
|
|
feat_decoder,
|
|
fsq_layer,
|
|
enc_to_lm_proj,
|
|
lm_to_dit_proj,
|
|
res_to_dit_proj,
|
|
fusion_concat_proj,
|
|
stop_proj,
|
|
stop_head,
|
|
device: vb.device().clone(),
|
|
dtype: vb.dtype(),
|
|
})
|
|
}
|
|
|
|
pub fn inference(
|
|
&mut self,
|
|
text: &Tensor,
|
|
audio_feat: Option<&Tensor>,
|
|
audio_mask: Option<&Tensor>,
|
|
min_len: usize,
|
|
max_len: usize,
|
|
inference_timesteps: usize,
|
|
cfg_value: f64,
|
|
audio_vae: &AudioVAE,
|
|
) -> Result<Tensor> {
|
|
let (b, t) = text.dims2()?;
|
|
let scale_emb = if self.config.lm_config.use_mup {
|
|
self.config.lm_config.scale_emb
|
|
} else {
|
|
1.0
|
|
};
|
|
let text_embed = self
|
|
.base_lm
|
|
.embed_tokens
|
|
.as_ref()
|
|
.unwrap()
|
|
.forward(text)?
|
|
.affine(scale_emb as f64, 0.0)?;
|
|
let (combined_embed, mut prefix_feat_cond, feat_embed) = if let Some(audio_feat) =
|
|
audio_feat
|
|
&& let Some(audio_mask) = audio_mask
|
|
{
|
|
let audio_feat = audio_feat.to_dtype(self.dtype)?;
|
|
let audio_t = audio_feat.dim(1)?;
|
|
let feat_embed = self.feat_encoder.forward(&audio_feat)?; // [b, audio_t, h_feat]
|
|
let feat_embed = self.enc_to_lm_proj.forward(&feat_embed)?.squeeze(0)?;
|
|
let embeds = masked_scatter_dim0(&text_embed, &feat_embed, audio_mask)?;
|
|
let prefix_feat_cond = audio_feat.i((.., audio_t - 1, ..))?;
|
|
(embeds, prefix_feat_cond, Some(feat_embed))
|
|
} else {
|
|
let prefix_feat_cond = Tensor::zeros(
|
|
(b, self.patch_size, self.latent_dim),
|
|
self.dtype,
|
|
&self.device,
|
|
)?;
|
|
(text_embed, prefix_feat_cond, None)
|
|
};
|
|
let mut pred_feat_seq = Vec::new();
|
|
// if feat_mask.i((1, t-1))?.to_scalar::<f32>()? == 0.0 {
|
|
// // TODO for stream
|
|
// }
|
|
let mut position_id = 0;
|
|
let mut seq_len = t;
|
|
let enc_outputs = self
|
|
.base_lm
|
|
.forward_with_cache(&combined_embed, position_id)?;
|
|
|
|
let (mut lm_hidden, input_embeds) = if let Some(_) = audio_feat
|
|
&& let Some(audio_mask) = audio_mask
|
|
&& let Some(feat_embed) = feat_embed
|
|
{
|
|
let fsq_emb = self.fsq_layer.forward(&enc_outputs)?;
|
|
let audio_mask_broadcast = audio_mask
|
|
.unsqueeze(D::Minus1)?
|
|
.broadcast_as(fsq_emb.shape())?;
|
|
let enc_outputs = audio_mask_broadcast.where_cond(&fsq_emb, &enc_outputs)?;
|
|
let lm_hidden = enc_outputs.i((.., t - 1, ..))?;
|
|
let input_embeds = if let Some(fusion) = &self.fusion_concat_proj {
|
|
let feat = enc_outputs.zeros_like()?;
|
|
let feat = masked_scatter_dim0(&feat, &feat_embed, audio_mask)?;
|
|
let concat = Tensor::cat(&[&enc_outputs, &feat], D::Minus1)?;
|
|
fusion.forward(&concat)?
|
|
} else {
|
|
let feat = enc_outputs.zeros_like()?;
|
|
let feat = masked_scatter_dim0(&feat, &feat_embed, audio_mask)?;
|
|
enc_outputs.add(&feat)?
|
|
};
|
|
(lm_hidden, input_embeds)
|
|
} else {
|
|
let lm_hidden = enc_outputs.i((.., t - 1, ..))?;
|
|
let input_embeds = if let Some(fusion) = &self.fusion_concat_proj {
|
|
let feat = enc_outputs.zeros_like()?;
|
|
let concat = Tensor::cat(&[&enc_outputs, &feat], D::Minus1)?;
|
|
fusion.forward(&concat)?
|
|
} else {
|
|
enc_outputs
|
|
};
|
|
(lm_hidden, input_embeds)
|
|
};
|
|
let residual_enc_outputs = self
|
|
.residual_lm
|
|
.forward_with_cache(&input_embeds, position_id)?;
|
|
let mut residual_hidden = residual_enc_outputs.i((.., t - 1, ..))?;
|
|
for i in 0..max_len {
|
|
let dit_hidden_1 = self.lm_to_dit_proj.forward(&lm_hidden)?; // [b, h_dit]
|
|
let dit_hidden_2 = self.res_to_dit_proj.forward(&residual_hidden)?; // [b, h_dit]
|
|
// let dit_hidden = dit_hidden_1.add(&dit_hidden_2)?;
|
|
let dit_hidden = if self.fusion_concat_proj.is_some() {
|
|
Tensor::cat(&[&dit_hidden_1, &dit_hidden_2], D::Minus1)?
|
|
} else {
|
|
dit_hidden_1.add(&dit_hidden_2)?
|
|
};
|
|
let cond = prefix_feat_cond.transpose(1, 2)?.contiguous()?;
|
|
let pred_feat = self
|
|
.feat_decoder
|
|
.forward(
|
|
&dit_hidden,
|
|
inference_timesteps,
|
|
self.patch_size,
|
|
&cond,
|
|
1.0,
|
|
cfg_value,
|
|
1.0,
|
|
true,
|
|
)?
|
|
.transpose(1, 2)?; // [b, p, d]
|
|
let curr_embed = self.feat_encoder.forward(&pred_feat.unsqueeze(1)?)?; // [b, 1, c]
|
|
let curr_embed = self.enc_to_lm_proj.forward(&curr_embed)?;
|
|
pred_feat_seq.push(pred_feat.unsqueeze(1)?);
|
|
|
|
prefix_feat_cond = pred_feat;
|
|
let stop_flag = self.stop_proj.forward(&lm_hidden)?.silu()?;
|
|
let stop_flag = self
|
|
.stop_head
|
|
.forward(&stop_flag)?
|
|
.argmax(D::Minus1)?
|
|
.i(0)?
|
|
.to_scalar::<u32>()?;
|
|
if i > min_len && stop_flag == 1 {
|
|
break;
|
|
}
|
|
position_id += seq_len;
|
|
seq_len = 1;
|
|
lm_hidden = self
|
|
.base_lm
|
|
.forward_with_cache(&curr_embed.i((.., 0, ..))?, position_id)?
|
|
.squeeze(1)?;
|
|
lm_hidden = self.fsq_layer.forward(&lm_hidden)?;
|
|
let curr_residual_input = if let Some(fusion) = &self.fusion_concat_proj {
|
|
let curr_embed = curr_embed.i((.., 0, ..))?;
|
|
let concat = Tensor::cat(&[&lm_hidden, &curr_embed], D::Minus1)?;
|
|
fusion.forward(&concat)?
|
|
} else {
|
|
lm_hidden.add(&curr_embed.i((.., 0, ..))?)?
|
|
};
|
|
residual_hidden = self
|
|
.residual_lm
|
|
.forward_with_cache(&curr_residual_input, position_id)?
|
|
.squeeze(1)?;
|
|
}
|
|
let pred_seq = Tensor::cat(&pred_feat_seq, 1)?; // (b, t, p, d)
|
|
let (b, _, _, d) = pred_seq.dims4()?;
|
|
let feat_pred = pred_seq
|
|
.permute((0, 3, 1, 2))?
|
|
.reshape((b, d, ()))?
|
|
.contiguous()?;
|
|
self.clear_kv_cache();
|
|
|
|
let decode_audio = audio_vae
|
|
.decode(&feat_pred.to_dtype(DType::F32)?, None)?
|
|
.squeeze(1)?;
|
|
let decode_audio_len = decode_audio.dim(D::Minus1)? - 640 - 640;
|
|
let decode_audio = decode_audio.narrow(D::Minus1, 640, decode_audio_len)?;
|
|
Ok(decode_audio)
|
|
}
|
|
|
|
pub fn inference_stream(
|
|
&mut self,
|
|
text: Tensor,
|
|
audio_feat: Option<Tensor>,
|
|
audio_mask: Option<Tensor>,
|
|
min_len: usize,
|
|
max_len: usize,
|
|
inference_timesteps: usize,
|
|
cfg_value: f64,
|
|
audio_vae: &AudioVAE,
|
|
) -> Result<impl Stream<Item = Result<Tensor, anyhow::Error>>> {
|
|
let (b, t) = text.dims2()?;
|
|
let scale_emb = if self.config.lm_config.use_mup {
|
|
self.config.lm_config.scale_emb
|
|
} else {
|
|
1.0
|
|
};
|
|
let text_embed = self
|
|
.base_lm
|
|
.embed_tokens
|
|
.as_ref()
|
|
.unwrap()
|
|
.forward(&text)?
|
|
.affine(scale_emb as f64, 0.0)?;
|
|
let (combined_embed, mut prefix_feat_cond, feat_embed) = if let Some(audio_feat) =
|
|
&audio_feat
|
|
&& let Some(audio_mask) = &audio_mask
|
|
{
|
|
let audio_feat = audio_feat.to_dtype(self.dtype)?;
|
|
let audio_t = audio_feat.dim(1)?;
|
|
let feat_embed = self.feat_encoder.forward(&audio_feat)?; // [b, audio_t, h_feat]
|
|
let feat_embed = self.enc_to_lm_proj.forward(&feat_embed)?.squeeze(0)?;
|
|
let embeds = masked_scatter_dim0(&text_embed, &feat_embed, audio_mask)?;
|
|
let prefix_feat_cond = audio_feat.i((.., audio_t - 1, ..))?;
|
|
(embeds, prefix_feat_cond, Some(feat_embed))
|
|
} else {
|
|
let prefix_feat_cond = Tensor::zeros(
|
|
(b, self.patch_size, self.latent_dim),
|
|
self.dtype,
|
|
&self.device,
|
|
)?;
|
|
(text_embed, prefix_feat_cond, None)
|
|
};
|
|
|
|
let streaming_prefix_len = 4usize;
|
|
// 流式处理固定4个结果使得VAE decode结果正常
|
|
let mut pred_feat_seq = Vec::with_capacity(streaming_prefix_len);
|
|
// 流式处理添加prompt块
|
|
if let Some(audio_feat) = &audio_feat
|
|
&& let Some(audio_mask) = &audio_mask
|
|
{
|
|
let audio_len = audio_mask.dim(1)?;
|
|
if audio_mask
|
|
.narrow(1, audio_len - 1, 1)?
|
|
.squeeze(0)?
|
|
.squeeze(0)?
|
|
.to_scalar::<u32>()?
|
|
== 1
|
|
{
|
|
let audio_len = audio_mask.sum_all()?.to_scalar::<u32>()? as usize;
|
|
let context_len = audio_len.min(streaming_prefix_len - 1);
|
|
let start = audio_feat.dim(1)? - context_len;
|
|
let last_feat = audio_feat
|
|
.narrow(1, start, context_len)?
|
|
.to_dtype(self.dtype)?;
|
|
pred_feat_seq.push(last_feat);
|
|
}
|
|
}
|
|
let mut position_id = 0;
|
|
let mut seq_len = t;
|
|
let enc_outputs = self
|
|
.base_lm
|
|
.forward_with_cache(&combined_embed, position_id)?;
|
|
|
|
let (mut lm_hidden, input_embeds) = if let Some(_) = &audio_feat
|
|
&& let Some(audio_mask) = &audio_mask
|
|
&& let Some(feat_embed) = feat_embed
|
|
{
|
|
let fsq_emb = self.fsq_layer.forward(&enc_outputs)?;
|
|
let audio_mask_broadcast = audio_mask
|
|
.unsqueeze(D::Minus1)?
|
|
.broadcast_as(fsq_emb.shape())?;
|
|
let enc_outputs = audio_mask_broadcast.where_cond(&fsq_emb, &enc_outputs)?;
|
|
let lm_hidden = enc_outputs.i((.., t - 1, ..))?;
|
|
let input_embeds = if let Some(fusion) = &self.fusion_concat_proj {
|
|
let feat = enc_outputs.zeros_like()?;
|
|
let feat = masked_scatter_dim0(&feat, &feat_embed, audio_mask)?;
|
|
let concat = Tensor::cat(&[&enc_outputs, &feat], D::Minus1)?;
|
|
fusion.forward(&concat)?
|
|
} else {
|
|
let feat = enc_outputs.zeros_like()?;
|
|
let feat = masked_scatter_dim0(&feat, &feat_embed, audio_mask)?;
|
|
enc_outputs.add(&feat)?
|
|
};
|
|
(lm_hidden, input_embeds)
|
|
} else {
|
|
let lm_hidden = enc_outputs.i((.., t - 1, ..))?;
|
|
let input_embeds = if let Some(fusion) = &self.fusion_concat_proj {
|
|
let feat = enc_outputs.zeros_like()?;
|
|
let concat = Tensor::cat(&[&enc_outputs, &feat], D::Minus1)?;
|
|
fusion.forward(&concat)?
|
|
} else {
|
|
enc_outputs
|
|
};
|
|
(lm_hidden, input_embeds)
|
|
};
|
|
let residual_enc_outputs = self
|
|
.residual_lm
|
|
.forward_with_cache(&input_embeds, position_id)?;
|
|
let mut residual_hidden = residual_enc_outputs.i((.., t - 1, ..))?;
|
|
let stream = stream! {
|
|
let mut first_flag = true;
|
|
for i in 0..max_len {
|
|
let dit_hidden_1 = self.lm_to_dit_proj.forward(&lm_hidden)?; // [b, h_dit]
|
|
let dit_hidden_2 = self.res_to_dit_proj.forward(&residual_hidden)?; // [b, h_dit]
|
|
// let dit_hidden = dit_hidden_1.add(&dit_hidden_2)?;
|
|
let dit_hidden = if self.fusion_concat_proj.is_some() {
|
|
Tensor::cat(&[&dit_hidden_1, &dit_hidden_2], D::Minus1)?
|
|
} else {
|
|
dit_hidden_1.add(&dit_hidden_2)?
|
|
};
|
|
let cond = prefix_feat_cond.transpose(1, 2)?.contiguous()?;
|
|
let pred_feat = self
|
|
.feat_decoder
|
|
.forward(
|
|
&dit_hidden,
|
|
inference_timesteps,
|
|
self.patch_size,
|
|
&cond,
|
|
1.0,
|
|
cfg_value,
|
|
1.0,
|
|
true,
|
|
)?
|
|
.transpose(1, 2)?; // [b, p, d]
|
|
let curr_embed = self.feat_encoder.forward(&pred_feat.unsqueeze(1)?)?; // [b, 1, c]
|
|
let curr_embed = self.enc_to_lm_proj.forward(&curr_embed)?;
|
|
// 保持容量不超过最大值
|
|
if pred_feat_seq.len() == streaming_prefix_len {
|
|
pred_feat_seq.remove(0);
|
|
}
|
|
pred_feat_seq.push(pred_feat.unsqueeze(1)?);
|
|
prefix_feat_cond = pred_feat;
|
|
let stop_flag = self.stop_proj.forward(&lm_hidden)?.silu()?;
|
|
let stop_flag = self
|
|
.stop_head
|
|
.forward(&stop_flag)?
|
|
.argmax(D::Minus1)?
|
|
.i(0)?
|
|
.to_scalar::<u32>()?;
|
|
|
|
let pred_feat_chunk = Tensor::cat(&pred_feat_seq, 1)?;
|
|
let (b, _, _, d) = pred_feat_chunk.dims4()?;
|
|
let feat_pred = pred_feat_chunk.permute((0, 3, 1, 2))?
|
|
.reshape((b, d, ()))?
|
|
.contiguous()?;
|
|
let mut decode_audio = audio_vae
|
|
.decode(&feat_pred.to_dtype(DType::F32)?, None)?
|
|
.squeeze(1)?;
|
|
// 只取当前帧结果
|
|
let decode_start = decode_audio.dim(D::Minus1)? - self.decode_patch_len;
|
|
decode_audio = decode_audio.narrow(D::Minus1, decode_start, self.decode_patch_len)?;
|
|
if i > min_len && stop_flag == 1 {
|
|
// // 最后一段容易有噪音，暂时不返回
|
|
// let decode_audio_len = decode_audio.dim(D::Minus1)? - 640;
|
|
// decode_audio = decode_audio.narrow(D::Minus1, 0, decode_audio_len)?;
|
|
// yield Ok(decode_audio);
|
|
break;
|
|
}
|
|
if first_flag {
|
|
// 去除初始噪音
|
|
let decode_audio_len = decode_audio.dim(D::Minus1)? - 1280;
|
|
decode_audio = decode_audio.narrow(D::Minus1, 1280, decode_audio_len)?;
|
|
first_flag = false;
|
|
}
|
|
yield Ok(decode_audio);
|
|
position_id += seq_len;
|
|
seq_len = 1;
|
|
lm_hidden = self
|
|
.base_lm
|
|
.forward_with_cache(&curr_embed.i((.., 0, ..))?, position_id)?
|
|
.squeeze(1)?;
|
|
lm_hidden = self.fsq_layer.forward(&lm_hidden)?;
|
|
let curr_residual_input = if let Some(fusion) = &self.fusion_concat_proj {
|
|
let curr_embed = curr_embed.i((.., 0, ..))?;
|
|
let concat = Tensor::cat(&[&lm_hidden, &curr_embed], D::Minus1)?;
|
|
fusion.forward(&concat)?
|
|
} else {
|
|
lm_hidden.add(&curr_embed.i((.., 0, ..))?)?
|
|
};
|
|
residual_hidden = self
|
|
.residual_lm
|
|
.forward_with_cache(&curr_residual_input, position_id)?
|
|
.squeeze(1)?;
|
|
|
|
}
|
|
self.clear_kv_cache();
|
|
};
|
|
Ok(stream)
|
|
}
|
|
|
|
pub fn clear_kv_cache(&mut self) {
|
|
self.base_lm.clear_kv_cache();
|
|
self.residual_lm.clear_kv_cache();
|
|
}
|
|
}
|