inference.py

import numpy as np
# import onnxruntime as ort
import ztu_somemodelruntime_rknnlite2 as ort
import sentencepiece as spm
import soundfile as sf

ort.set_default_logger_verbosity(0)

def load_onnx_model(model_path):
    """加载ONNX模型"""
    return ort.InferenceSession(
        model_path,
        providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
    )

class SimpleT5Tokenizer:
    def __init__(self, model_path, max_length=128):
        """初始化tokenizer
        
        Args:
            model_path: sentencepiece模型路径
            max_length: 序列最大长度,默认128
        """
        self.sp = spm.SentencePieceProcessor()
        self.sp.Load(model_path)
        
        # T5特殊token的ID
        self.pad_token_id = 0
        self.eos_token_id = 1
        self.max_length = max_length
        
    def __call__(self, texts, padding=True, truncation=True, max_length=None, return_tensors="np"):
        """处理文本序列
        
        Args:
            texts: 文本或文本列表
            padding: 是否padding
            truncation: 是否截断
            max_length: 可选,覆盖默认max_length
            return_tensors: 返回类型(只支持"np")
            
        Returns:
            dict: 包含input_ids和attention_mask
        """
        if isinstance(texts, str):
            texts = [texts]
            
        max_len = max_length if max_length is not None else self.max_length
            
        # 分词并转换为ID
        input_ids = []
        attention_mask = []
        for text in texts:
            ids = self.sp.EncodeAsIds(text)
            
            # 截断处理(预留EOS token位置)
            if truncation and len(ids) > max_len - 1:
                ids = ids[:max_len-1]
            ids.append(self.eos_token_id)
            
            # 创建attention mask
            mask = [1] * len(ids)
            
            # Padding处理
            if padding:
                pad_length = max_len - len(ids)
                ids.extend([self.pad_token_id] * pad_length)
                mask.extend([0] * pad_length)
                
            input_ids.append(ids)
            attention_mask.append(mask)
                
        # 转换为numpy array
        input_ids = np.array(input_ids, dtype=np.int64)
        attention_mask = np.array(attention_mask, dtype=np.int64)
            
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask
        }

def encode_text(prompt, negative_prompt, tokenizer, text_encoder_onnx, guidance_scale=None):
    """编码文本,同时处理条件和无条件文本
    
    Args:
        prompt: 文本提示
        tokenizer: T5 tokenizer
        text_encoder_onnx: T5 ONNX模型
        guidance_scale: 引导系数
    """
    if not isinstance(prompt, list):
        prompt = [prompt]
    
    if guidance_scale is not None and guidance_scale > 1.0:
        # 同时处理条件和无条件文本
        all_prompts = [negative_prompt] + prompt
        batch = tokenizer(
            all_prompts,
            padding=True,
            truncation=True,
            return_tensors="np"
        )
        
        # ONNX推理
        all_hidden_states = text_encoder_onnx.run(
            ['last_hidden_state'],
            {
                'input_ids': batch['input_ids'].astype(np.int64),
                'attention_mask': batch['attention_mask'].astype(np.int64)
            }
        )[0]
            
        # 分离无条件和条件结果
        uncond_hidden_states = all_hidden_states[0:1]
        cond_hidden_states = all_hidden_states[1:]
        uncond_mask = batch['attention_mask'][0:1]
        cond_mask = batch['attention_mask'][1:]
        
        return (uncond_hidden_states, uncond_mask), (cond_hidden_states, cond_mask)
    else:
        # 只处理条件文本
        batch = tokenizer(
            prompt,
            padding=True,
            truncation=True,
            return_tensors="np"
        )
        
        # ONNX推理
        hidden_states = text_encoder_onnx.run(
            ['last_hidden_state'],
            {
                'input_ids': batch['input_ids'].astype(np.int64),
                'attention_mask': batch['attention_mask'].astype(np.int64)
            }
        )[0]
            
        return hidden_states, batch['attention_mask']

def retrieve_timesteps(scheduler, num_inference_steps, device, timesteps=None, sigmas=None):
    """获取timesteps"""
    if sigmas is not None:
        scheduler.set_timesteps(sigmas=sigmas)
        timesteps = scheduler.timesteps
        num_inference_steps = len(timesteps)
    else:
        scheduler.set_timesteps(num_inference_steps)
        timesteps = scheduler.timesteps
    return timesteps, num_inference_steps

# 添加一个简单的FlowMatchScheduler类
class SimpleFlowMatchScheduler:
    def __init__(self, num_train_timesteps=1000, shift=1.0):
        """初始化scheduler
        
        Args:
            num_train_timesteps: 训练步数
            shift: 时间步偏移量
        """
        # 生成线性timesteps
        timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
        
        # 计算sigmas
        sigmas = timesteps / num_train_timesteps
        sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
        
        # 添加终止sigma
        self.sigmas = np.append(sigmas, 0.0)
        self.timesteps = sigmas * num_train_timesteps
        self.step_index = None
        
    def set_timesteps(self, num_inference_steps):
        """设置推理时的timesteps
        
        Args:
            num_inference_steps: 推理步数
        """
        timesteps = np.linspace(1, len(self.timesteps), num_inference_steps, dtype=np.float32)[::-1].copy()
        sigmas = timesteps / len(self.timesteps)
        self.sigmas = np.append(sigmas, 0.0)
        self.timesteps = sigmas * len(self.timesteps)
        self.step_index = 0
        
    def step(self, model_output, timestep, sample):
        """执行一步euler更新
        
        Args:
            model_output: 模型输出
            timestep: 当前时间步
            sample: 当前样本
            
        Returns:
            prev_sample: 更新后的样本
        """
        sigma = self.sigmas[self.step_index] 
        sigma_next = self.sigmas[self.step_index + 1]
        
        # euler更新
        prev_sample = sample + (sigma_next - sigma) * model_output
        
        self.step_index += 1
        return prev_sample

def generate_audio_onnx(
    prompt="",
    negative_prompt="",
    duration=10,
    steps=50,
    guidance_scale=4.5,
    onnx_dir="./onnx_models",
    output_path="output_onnx.wav",
    seed=None
):
    if seed is not None:
        np.random.seed(seed)
    
    # 加载tokenizer和ONNX模型,设置固定长度
    tokenizer = SimpleT5Tokenizer(f"{onnx_dir}/spiece.model", max_length=63)
    text_encoder_onnx = load_onnx_model(f"{onnx_dir}/text_encoder_nf4.onnx")
    
    # 加载其他ONNX模型
    vae_decoder = load_onnx_model(f"{onnx_dir}/vae_decoder.onnx")
    duration_embedder = load_onnx_model(f"{onnx_dir}/duration_embedder.onnx")
    transformer = load_onnx_model(f"{onnx_dir}/transformer.onnx")
    proj_layer = load_onnx_model(f"{onnx_dir}/proj.onnx")
    
    # 1. duration embedding
    duration_input = np.array([[duration]], dtype=np.float32)
    print(f"[Shape] duration输入: {duration_input.shape}")
    
    duration_hidden_states = duration_embedder.run(
        ['embedding'],
        {'duration': duration_input}
    )[0]
    print(f"[Shape] duration embedding: {duration_hidden_states.shape}")

    if guidance_scale > 1.0:
        duration_hidden_states = np.concatenate([duration_hidden_states] * 2, axis=0)
        print(f"[Shape] 复制后的duration embedding: {duration_hidden_states.shape}")
    
    # 2. text encoder
    if guidance_scale > 1.0:
        (uncond_hidden_states, uncond_mask), (cond_hidden_states, cond_mask) = encode_text(
            prompt, negative_prompt, tokenizer, text_encoder_onnx, guidance_scale=guidance_scale
        )
        print(cond_hidden_states)
        encoder_hidden_states = np.concatenate([uncond_hidden_states, cond_hidden_states])
        attention_mask = np.concatenate([uncond_mask, cond_mask])
    else:
        encoder_hidden_states, attention_mask = encode_text(
            prompt, tokenizer, text_encoder_onnx
        )
    
    # 3. pooled_text
    boolean_encoder_mask = (attention_mask == 1)
    mask_expanded = boolean_encoder_mask[..., None].repeat(encoder_hidden_states.shape[-1], axis=-1)
    masked_data = np.where(mask_expanded, encoder_hidden_states, np.nan)
    pooled = np.nanmean(masked_data, axis=1)
    
    # 使用projection层处理
    pooled_text = proj_layer.run(
        ['projected'],
        {'text_embedding': pooled.astype(np.float32)}
    )[0]
    
    # 4. 合并duration和text特征
    encoder_hidden_states = np.concatenate(
        [encoder_hidden_states, duration_hidden_states],
        axis=1
    )
    
    # 5. 创建其他输入
    txt_ids = np.zeros((1, encoder_hidden_states.shape[1], 3), dtype=np.int64)
    img_ids = np.tile(
        np.arange(645, dtype=np.int64)[None, :, None],
        (1, 1, 3)
    )
    
    # 6. scheduler
    scheduler = SimpleFlowMatchScheduler(num_train_timesteps=1000)
    scheduler.set_timesteps(steps)
    
    # 初始化latents
    latents = np.random.randn(1, 645, 64).astype(np.float32)
    
    # 7. 生成循环
    for i in range(steps):
        # Transformer前向传播
        noise_pred = transformer.run(
            ['output'],
            {
                'hidden_states': latents,
                'timestep': np.array([scheduler.timesteps[i]/1000], dtype=np.float32),
                'pooled_text': pooled_text,
                'encoder_hidden_states': encoder_hidden_states, 
                'txt_ids': txt_ids,
                'img_ids': img_ids
            }
        )[0]
        
        if i == 0:  # 只在第一步打印
            print(f"[Shape] noise预测输出: {noise_pred.shape}")
            
        # 应用classifier free guidance
        if guidance_scale > 1.0:
            noise_pred_uncond, noise_pred_text = noise_pred[0:1], noise_pred[1:2]
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
            
        # 使用scheduler更新latents
        latents = scheduler.step(noise_pred, scheduler.timesteps[i], latents)
        
        if i % 10 == 0:
            print(f"生成进度: {i}/{steps}")
            
    # 8. VAE解码前的处理
    latents = latents / scheduler.sigmas[0]
    latents = np.transpose(latents, (0, 2, 1))
    
    # 9. VAE解码
    wave = vae_decoder.run(['audio'], {'latent': latents})[0]
    
    # 10. 裁剪
    sample_rate = 44100
    waveform_end = int(duration * sample_rate)
    wave = wave[:, :, :waveform_end]
    print(f"[Shape] 裁剪后的最终波形: {wave.shape}")
    
    # 11. 保存音频
    wave = wave[0]  # 移除batch维度
    sf.write(output_path, wave.T, sample_rate)  # soundfile需要(samples, channels)格式
    
    return wave

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="测试ONNX模型推理")
    parser.add_argument("--prompt", type=str, default="What does the fox say?", help="文本提示")
    parser.add_argument("--negative_prompt", type=str, default="", help="负文本提示")
    parser.add_argument("--onnx_dir", type=str, default=".", help="ONNX模型目录")
    parser.add_argument("--duration", type=float, default=10.0, help="生成音频时长（秒）")
    parser.add_argument("--steps", type=int, default=30, help="推理步数")
    parser.add_argument("--guidance_scale", type=float, default=4.5, help="引导系数")
    parser.add_argument("--output", type=str, default="output_onnx.wav", help="输出音频路径")
    parser.add_argument("--seed", type=int, default=42, help="随机种子")
    
    args = parser.parse_args()
    
    # 生成音频
    wave = generate_audio_onnx(
        # prompt="What does the fox say?",
        # prompt="Never gonna give you up, never gonna let you down",
        # prompt="Electonic music, future house style",
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        duration=args.duration,
        steps=args.steps,
        guidance_scale=args.guidance_scale,
        onnx_dir=args.onnx_dir,
        output_path=args.output,
        seed=args.seed
    )
    
    print(f"生成的音频shape为: {wave.shape}")
    print(f"音频已保存到: {args.output}")