Upload 13 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+transformer.rknn filter=lfs diff=lfs merge=lfs -text
+vae_decoder.rknn filter=lfs diff=lfs merge=lfs -text

convert_rknn.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+import datetime
+import argparse
+from rknn.api import RKNN
+from sys import exit
+
+AUDIO_LENGTH = 645 # 音频长度, 645为10秒
+TEXT_LENGTH = 64 # 文本长度(token)
+
+# 模型配置
+MODELS = {
+    'transformer': 'transformer.onnx',
+    'vae_decoder': 'vae_decoder.onnx',
+}
+
+SHAPES = {
+    'transformer': [
+        [
+            [1, AUDIO_LENGTH, 64],  # hidden_states
+            [1,],  # timestep
+            [2, 1024],  # pooled_text
+            [2, TEXT_LENGTH, 1024],  # encoder_hidden_states
+            [1, TEXT_LENGTH, 3],  # txt_ids
+            [1, AUDIO_LENGTH, 3],  # img_ids
+        ],
+    ],
+    'vae_decoder': [
+        [
+            [1, 64, AUDIO_LENGTH],
+        ],
+    ],
+}
+
+QUANTIZE=False
+detailed_performance_log = True
+
+def convert_model(model_type):
+    """转换指定类型的模型到RKNN格式"""
+    if model_type not in MODELS:
+        print(f"错误: 不支持的模型类型 {model_type}")
+        return False
+        
+    onnx_model = MODELS[model_type]
+    rknn_model = onnx_model.replace(".onnx",".rknn")
+    
+    timedate_iso = datetime.datetime.now().isoformat()
+    
+    rknn = RKNN(verbose=True)
+    rknn.config(
+        quantized_dtype='w8a8',
+        quantized_algorithm='normal',
+        quantized_method='channel',
+        quantized_hybrid_level=0,
+        target_platform='rk3588',
+        quant_img_RGB2BGR = False,
+        float_dtype='float16',
+        optimization_level=3,
+        custom_string=f"converted at {timedate_iso}",
+        remove_weight=False,
+        compress_weight=False,
+        inputs_yuv_fmt=None,
+        single_core_mode=False,
+        dynamic_input=SHAPES[model_type],
+        model_pruning=False,
+        op_target=None,
+        quantize_weight=False,
+        remove_reshape=False,
+        sparse_infer=False,
+        enable_flash_attention=False,
+        #  disable_rules=['convert_gemm_by_exmatmul']
+    )
+
+    print(f"开始转换 {model_type} 模型...")
+    ret = rknn.load_onnx(model=onnx_model)
+    if ret != 0:
+        print("加载ONNX模型失败")
+        return False
+        
+    ret = rknn.build(do_quantization=False, rknn_batch_size=None)
+    if ret != 0:
+        print("构建RKNN模型失败")
+        return False
+        
+    ret = rknn.export_rknn(rknn_model)
+    if ret != 0:
+        print("导出RKNN模型失败")
+        return False
+        
+    print(f"成功转换模型: {rknn_model}")
+    return True
+
+def main():
+    parser = argparse.ArgumentParser(description='转换ONNX模型到RKNN格式')
+    parser.add_argument('model_type', nargs='?', default='all',
+                      choices=['all', 'transformer', 'vae_decoder'],
+                      help='要转换的模型类型 (默认: all)')
+    
+    args = parser.parse_args()
+    
+    if args.model_type == 'all':
+        # 转换所有模型
+        for model_type in MODELS.keys():
+            if not convert_model(model_type):
+                print(f"转换 {model_type} 失败")
+    else:
+        # 转换指定模型
+        if not convert_model(args.model_type):
+            print(f"转换 {args.model_type} 失败")
+
+if __name__ == '__main__':
+    main()
+
+

duration_embedder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b2bd04d4bbd075e7c663711e55b6d09c68bbd35a772587ae46d8339599e03e3
+size 1061046

export_onnx.py

@@ -0,0 +1,234 @@
+import torch
+import torch.nn as nn
+from diffusers import AutoencoderOobleck
+from diffusers import FluxTransformer2DModel
+from tangoflux import TangoFluxInference
+from tangoflux.model import DurationEmbedder, TangoFlux
+
+def export_vae_encoder(vae, save_path, batch_size=1, audio_length=441000):
+    """导出VAE编码器到ONNX格式
+    
+    Args:
+        vae: AutoencoderOobleck实例
+        save_path: 保存路径
+        batch_size: batch大小
+        audio_length: 音频长度(默认10秒,44100Hz采样率)
+    """
+    vae.eval()
+    
+    # 创建dummy input - 注意这里是双声道音频
+    dummy_input = torch.randn(batch_size, 2, audio_length)
+    
+    # 创建一个包装类来处理forward调用
+    class VAEEncoderWrapper(nn.Module):
+        def __init__(self, vae):
+            super().__init__()
+            self.vae = vae
+            
+        def forward(self, audio):
+            return self.vae.encode(audio).latent_dist.sample()
+    
+    wrapper = VAEEncoderWrapper(vae)
+    
+    # 导出encoder部分
+    torch.onnx.export(
+        wrapper,
+        dummy_input,
+        save_path,
+        input_names=['audio'],
+        output_names=['latent'],
+        dynamic_axes={
+            'audio': {0: 'batch_size', 2: 'audio_length'},
+            'latent': {0: 'batch_size', 2: 'latent_length'}
+        },
+        opset_version=17
+    )
+
+def export_vae_decoder(vae, save_path, batch_size=1, latent_length=645):
+    """导出VAE解码器到ONNX格式
+    
+    Args:
+        vae: AutoencoderOobleck实例
+        save_path: 保存路径
+        batch_size: batch大小
+        latent_length: 潜在向量长度
+    """
+    vae.eval()
+    
+    # 创建dummy input
+    dummy_input = torch.randn(batch_size, 64, latent_length)
+    
+    # 创建一个包装类来处理forward调用
+    class VAEDecoderWrapper(nn.Module):
+        def __init__(self, vae):
+            super().__init__()
+            self.vae = vae
+            
+        def forward(self, latent):
+            return self.vae.decode(latent).sample
+    
+    wrapper = VAEDecoderWrapper(vae)
+    
+    # 导出decoder部分
+    torch.onnx.export(
+        wrapper,
+        dummy_input,
+        save_path,
+        input_names=['latent'],
+        output_names=['audio'],
+        dynamic_axes={
+            'latent': {0: 'batch_size', 2: 'latent_length'},
+            'audio': {0: 'batch_size', 2: 'audio_length'}
+        },
+        opset_version=17
+    )
+
+def export_duration_embedder(duration_embedder, save_path, batch_size=1):
+    """导出Duration Embedder到ONNX格式
+    
+    Args:
+        duration_embedder: DurationEmbedder实例
+        save_path: 保存路径
+        batch_size: batch大小
+    """
+    duration_embedder.eval()
+    
+    # 创建dummy input - 注意这里是标量值
+    dummy_input = torch.tensor([[10.0]], dtype=torch.float32)  # 10秒
+    
+    # 导出
+    torch.onnx.export(
+        duration_embedder,
+        dummy_input,
+        save_path,
+        input_names=['duration'],
+        output_names=['embedding'],
+        dynamic_axes={
+            'duration': {0: 'batch_size'},
+            'embedding': {0: 'batch_size'}
+        },
+        opset_version=17
+    )
+
+def export_flux_transformer(transformer, save_path, batch_size=1, seq_length=645):
+    """导出FluxTransformer2D到ONNX格式
+    
+    Args:
+        transformer: FluxTransformer2DModel实例
+        save_path: 保存路径
+        batch_size: batch大小
+        seq_length: 序列长度
+    """
+    transformer.eval()
+    
+    # 创建dummy inputs - 注意所有输入的形状
+    hidden_states = torch.randn(batch_size, seq_length, 64)  # [B, S, C]
+    timestep = torch.tensor([0.5])  # [1]
+    pooled_text = torch.randn(batch_size, 1024)  # [B, D]
+    encoder_hidden_states = torch.randn(batch_size, 64, 1024)  # [B, L, D]
+    txt_ids = torch.zeros(batch_size, 64, 3).to(torch.int64) # [B, L, 3]
+    img_ids = torch.arange(seq_length).unsqueeze(0).unsqueeze(-1).repeat(batch_size, 1, 3).to(torch.int64)  # [B, S, 3]
+    
+    # 创建一个包装类来处理forward调用
+    class TransformerWrapper(nn.Module):
+        def __init__(self, transformer):
+            super().__init__()
+            self.transformer = transformer
+            
+        def forward(self, hidden_states, timestep, pooled_text, encoder_hidden_states, txt_ids, img_ids):
+            return self.transformer(
+                hidden_states=hidden_states,
+                timestep=timestep,
+                guidance=None,
+                pooled_projections=pooled_text,
+                encoder_hidden_states=encoder_hidden_states,
+                txt_ids=txt_ids,
+                img_ids=img_ids,
+                return_dict=False
+            )[0]
+    
+    wrapper = TransformerWrapper(transformer)
+    
+    # 导出
+    torch.onnx.export(
+        wrapper,
+        (hidden_states, timestep, pooled_text, encoder_hidden_states, txt_ids, img_ids),
+        save_path,
+        input_names=['hidden_states', 'timestep', 'pooled_text', 'encoder_hidden_states', 'txt_ids', 'img_ids'],
+        output_names=['output'],
+        dynamic_axes={
+            'hidden_states': {0: 'batch_size', 1: 'sequence_length'},
+            'pooled_text': {0: 'batch_size'},
+            'encoder_hidden_states': {0: 'batch_size', 1: 'text_length'},
+            'txt_ids': {0: 'batch_size', 1: 'text_length'},
+            'img_ids': {0: 'batch_size', 1: 'sequence_length'}
+        },
+        opset_version=17
+    )
+
+def export_proj_layer(proj_layer, save_path, batch_size=1):
+    """导出projection层到ONNX格式
+    
+    Args:
+        proj_layer: 投影层(fc层)实例
+        save_path: 保存路径
+        batch_size: batch大小
+    """
+    proj_layer.eval()
+    
+    # 创建dummy input - 使用T5的hidden size
+    dummy_input = torch.randn(batch_size, 1024)  # T5-large hidden size
+    
+    # 导出
+    torch.onnx.export(
+        proj_layer,
+        dummy_input,
+        save_path,
+        input_names=['text_embedding'],
+        output_names=['projected'],
+        dynamic_axes={
+            'text_embedding': {0: 'batch_size'},
+            'projected': {0: 'batch_size'}
+        },
+        opset_version=17
+    )
+
+def export_all(model_path, output_dir):
+    """导出所有组件到ONNX格式
+    
+    Args:
+        model_path: TangoFlux模型路径
+        output_dir: 输出目录
+    """
+    import os
+    
+    # 加载模型
+    model = TangoFluxInference(name=model_path, device="cpu")
+    
+    # 创建输出目录
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # 导出VAE
+    export_vae_encoder(model.vae, f"{output_dir}/vae_encoder.onnx")
+    export_vae_decoder(model.vae, f"{output_dir}/vae_decoder.onnx")
+    
+    # 导出Duration Embedder
+    export_duration_embedder(model.model.duration_emebdder, f"{output_dir}/duration_embedder.onnx")
+    
+    # 导出Transformer
+    export_flux_transformer(model.model.transformer, f"{output_dir}/transformer.onnx")
+    
+    # 导出Projection层
+    export_proj_layer(model.model.fc, f"{output_dir}/proj.onnx")
+    
+    print(f"所有模型已导出到: {output_dir}")
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="导出TangoFlux模型到ONNX格式")
+    parser.add_argument("--model_path", type=str, required=True, help="TangoFlux模型路径")
+    parser.add_argument("--output_dir", type=str, required=True, help="输出目录")
+    
+    args = parser.parse_args()
+    export_all(args.model_path, args.output_dir) 

inference.py

@@ -0,0 +1,365 @@
+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}") 

proj.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af92c1c262e6a217a75ef59c922304eb90770f9a67a6253c9c477fbe3fa9eba8
+size 4198734

spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d60acb128cf7b7f2536e8f38a5b18a05535c9e14c7a355904270e15b0945ea86
+size 791656

text_encoder_bnb4.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6334ec5d0eeaba54449c24ec0784c3e224db6c483a968c0fca055e001b80e39
+size 305592280

transformer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9079229578ab2b683c271f9f585fddabcfeb588191d9c02c597f0aa4b6a383b
+size 2068637351

transformer.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dffd321f320c949a0cffc9b3bf92b371fccaebb5f25826710c8b89d84184d2c7
+size 1118028281

vae_decoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25cdc6d16f896906df2cea9374b2746842ef808563e6daeb7b48b2eb6360a4a2
+size 312595968

vae_decoder.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c9782440915ffe1698717576b7889fc7d387a9c62b4874ff83337b8473b5049
+size 352599027

vae_encoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:846b8fd27f2f6309954fb1420066b56123bfe34c3adcc100809c578011179980
+size 312074746

ztu_somemodelruntime_rknnlite2.py

@@ -0,0 +1,535 @@
+# 模块级常量和函数
+from rknnlite.api import RKNNLite
+import numpy as np
+import os
+import warnings
+import logging
+from typing import List, Dict, Union, Optional
+
+try:
+    import onnxruntime as ort
+    HAS_ORT = True
+except ImportError:
+    HAS_ORT = False
+    warnings.warn("onnxruntime未安装,只能使用RKNN后端", ImportWarning)
+
+# 配置日志
+logger = logging.getLogger("somemodelruntime_rknnlite2")
+logger.setLevel(logging.ERROR)  # 默认只输出错误信息
+if not logger.handlers:
+    handler = logging.StreamHandler()
+    handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
+    logger.addHandler(handler)
+
+# ONNX Runtime日志级别到Python logging级别的映射
+_LOGGING_LEVEL_MAP = {
+    0: logging.DEBUG,    # Verbose
+    1: logging.INFO,     # Info
+    2: logging.WARNING,  # Warning
+    3: logging.ERROR,    # Error
+    4: logging.CRITICAL  # Fatal
+}
+
+def set_default_logger_severity(level: int) -> None:
+    """
+    Sets the default logging severity. 0:Verbose, 1:Info, 2:Warning, 3:Error, 4:Fatal
+    
+    Args:
+        level: 日志级别(0-4)
+    """
+    if level not in _LOGGING_LEVEL_MAP:
+        raise ValueError(f"无效的日志级别: {level}, 应该是0-4之间的整数")
+    logger.setLevel(_LOGGING_LEVEL_MAP[level])
+
+def set_default_logger_verbosity(level: int) -> None:
+    """
+    Sets the default logging verbosity level. To activate the verbose log, 
+    you need to set the default logging severity to 0:Verbose level.
+    
+    Args:
+        level: 日志级别(0-4)
+    """
+    set_default_logger_severity(level)
+
+# NPU核心模式常量
+NPU_CORE_AUTO = 0      # 自动选择
+NPU_CORE_0 = 1        # 使用核心0
+NPU_CORE_1 = 2        # 使用核心1
+NPU_CORE_2 = 4        # 使用核心2
+NPU_CORE_0_1 = 3      # 使用核心0和1
+NPU_CORE_0_1_2 = 7    # 使用所有核心
+NPU_CORE_ALL = 0xffff # 使用所有核心
+
+# RKNN tensor type到numpy dtype的映射
+RKNN_DTYPE_MAP = {
+    0: np.float32,  # RKNN_TENSOR_FLOAT32
+    1: np.float16,  # RKNN_TENSOR_FLOAT16
+    2: np.int8,     # RKNN_TENSOR_INT8
+    3: np.uint8,    # RKNN_TENSOR_UINT8
+    4: np.int16,    # RKNN_TENSOR_INT16
+    5: np.uint16,   # RKNN_TENSOR_UINT16
+    6: np.int32,    # RKNN_TENSOR_INT32
+    7: np.uint32,   # RKNN_TENSOR_UINT32
+    8: np.int64,    # RKNN_TENSOR_INT64
+    9: bool,        # RKNN_TENSOR_BOOL
+    10: np.int8,    # RKNN_TENSOR_INT4 (用int8表示)
+}
+
+def get_available_providers() -> List[str]:
+    """
+    获取可用的设备提供者列表(为保持接口兼容性的占位函数)
+    
+    Returns:
+        list: 可用的设备提供者列表,总是返回["CPUExecutionProvider"]
+    """
+    return ["CPUExecutionProvider"]
+
+def get_version_info() -> Dict[str, str]:
+    """
+    获取版本信息
+    
+    Returns:
+        dict: 包含API和驱动版本信息的字典
+    """
+    runtime = RKNNLite()
+    version = runtime.get_sdk_version()
+    return {
+        "api_version": version.split('\n')[2].split(': ')[1].split(' ')[0],
+        "driver_version": version.split('\n')[3].split(': ')[1]
+    }
+
+class IOTensor:
+    """输入/输出张量的信息封装类"""
+    def __init__(self, name, shape, type=None):
+        self.name = name.decode() if isinstance(name, bytes) else name
+        self.shape = shape
+        self.type = type
+
+    def __str__(self):
+        return f"IOTensor(name='{self.name}', shape={self.shape}, type={self.type})"
+
+class SessionOptions:
+    """会话选项类"""
+    def __init__(self):
+        self.async_mode = False  # 是否使用异步模式
+        self.core_mask = 0      # NPU核心选择
+        self.perf_debug = False # 是否启用性能分析
+
+class InferenceSession:
+    """
+    RKNNLite运行时封装类,API风格类似ONNX Runtime
+    """
+    
+    def __new__(cls, model_path: str, verbose: bool = False, sess_options: Optional[SessionOptions] = None, fallback: bool = True, **kwargs):
+        """
+        创建运行时实例
+        
+        Args:
+            model_path: 模型文件路径(.rknn或.onnx)
+            verbose: 是否打印详细日志
+            sess_options: 会话选项
+            fallback: 是否自动加载同名.rknn文件
+            **kwargs: 其他初始化参数
+        """
+        # 只在verbose=True时开启详细日志
+        if verbose:
+            set_default_logger_severity(0)
+            
+        if not os.path.exists(model_path):
+            logger.error(f"模型文件不存在: {model_path}")
+            raise FileNotFoundError(f"模型文件不存在: {model_path}")
+            
+        # 检查是否是ONNX文件
+        is_onnx = model_path.lower().endswith('.onnx')
+        
+        if is_onnx and fallback:
+            # 尝试查找对应的RKNN文件
+            rknn_path = os.path.splitext(model_path)[0] + '.rknn'
+            if os.path.exists(rknn_path):
+                logger.info(f"找到对应的RKNN模型,将使用RKNN: {rknn_path}")
+                # 创建RKNN运行时实例
+                instance = super().__new__(cls)
+                instance.model_path = rknn_path
+                return instance
+                
+        if is_onnx:
+            # 使用ONNX Runtime
+            logger.info(f"使用ONNX Runtime加载模型: {model_path}")
+            if not HAS_ORT:
+                raise RuntimeError("未安装onnxruntime,无法加载ONNX模型")
+            return ort.InferenceSession(model_path, sess_options=sess_options, **kwargs)
+            
+        # 创建RKNN运行时实例
+        instance = super().__new__(cls)
+        instance.model_path = model_path
+        return instance
+
+    def __init__(self, model_path: str, verbose: bool = False, sess_options: Optional[SessionOptions] = None, fallback: bool = True, **kwargs):
+        """
+        初始化RKNN运行时
+        
+        Args:
+            model_path: 模型文件路径(.rknn或.onnx)
+            verbose: 是否打印详细日志
+            sess_options: 会话选项
+            fallback: 是否自动加载同名.rknn文件
+            **kwargs: 其他初始化参数
+        """
+        # 如果是ONNX模型,__init__不会被调用
+        if not hasattr(self, 'model_path'):  # 如果是ONNX Runtime实例
+            return
+            
+        self.runtime = RKNNLite(verbose=verbose)
+        
+        # 加载模型
+        logger.debug(f"正在加载模型: {self.model_path}")
+        ret = self.runtime.load_rknn(self.model_path)
+        if ret != 0:
+            logger.error(f"加载RKNN模型失败: {self.model_path}")
+            raise RuntimeError(f'加载RKNN模型失败: {self.model_path}')
+        logger.debug("模型加载成功")
+            
+        # 应用会话选项
+        options = sess_options or SessionOptions()
+        
+        # 初始化运行时
+        logger.debug("正在初始化运行时环境")
+        ret = self.runtime.init_runtime(
+            async_mode=options.async_mode,
+            core_mask=options.core_mask
+        )
+        if ret != 0:
+            logger.error("初始化运行时环境失败")
+            raise RuntimeError('初始化运行时环境失败')
+        logger.debug("运行时环境初始化成功")
+
+        # 获取输入输出信息
+        self._init_io_info()
+        
+        # 保存选项
+        self.options = options
+
+    def get_performance_info(self) -> Dict[str, float]:
+        """
+        获取性能信息
+        
+        Returns:
+            dict: 包含性能信息的字典
+        """
+        if not self.options.perf_debug:
+            raise RuntimeError("性能分析未启用,请在SessionOptions中设置perf_debug=True")
+            
+        perf = self.runtime.rknn_runtime.get_run_perf()
+        return {
+            "run_duration": perf.run_duration / 1000.0  # 转换为毫秒
+        }
+
+    def set_core_mask(self, core_mask: int) -> None:
+        """
+        设置NPU核心使用模式
+        
+        Args:
+            core_mask: NPU核心掩码,使用NPU_CORE_*常量
+        """
+        ret = self.runtime.rknn_runtime.set_core_mask(core_mask)
+        if ret != 0:
+            raise RuntimeError("设置NPU核心模式失败")
+
+    def _convert_nhwc_to_nchw(self, shape):
+        """将NHWC格式的shape转换为NCHW格式"""
+        if len(shape) == 4:
+            # NHWC -> NCHW
+            n, h, w, c = shape
+            return [n, c, h, w]
+        return shape
+        
+    def _init_io_info(self):
+        """初始化模型的输入输出信息"""
+        runtime = self.runtime.rknn_runtime
+        
+        # 获取输入输出数量
+        n_input, n_output = runtime.get_in_out_num()
+        
+        # 获取输入信息
+        self.input_tensors = []
+        for i in range(n_input):
+            attr = runtime.get_tensor_attr(i)
+            shape = [attr.dims[j] for j in range(attr.n_dims)]
+            # 对四维输入进行NHWC到NCHW的转换
+            shape = self._convert_nhwc_to_nchw(shape)
+            # 获取dtype
+            dtype = RKNN_DTYPE_MAP.get(attr.type, None)
+            tensor = IOTensor(attr.name, shape, dtype)
+            self.input_tensors.append(tensor)
+            
+        # 获取输出信息
+        self.output_tensors = []
+        for i in range(n_output):
+            attr = runtime.get_tensor_attr(i, is_output=True)
+            shape = runtime.get_output_shape(i)
+            # 获取dtype
+            dtype = RKNN_DTYPE_MAP.get(attr.type, None)
+            tensor = IOTensor(attr.name, shape, dtype)
+            self.output_tensors.append(tensor)
+        
+    def get_inputs(self):
+        """
+        获取模型输入信息
+        
+        Returns:
+            list: 包含输入信息的列表
+        """
+        return self.input_tensors
+        
+    def get_outputs(self):
+        """
+        获取模型输出信息
+        
+        Returns:
+            list: 包含输出信息的列表
+        """
+        return self.output_tensors
+        
+    def run(self, output_names=None, input_feed=None, data_format="nchw", **kwargs):
+        """
+        执行模型推理
+        
+        Args:
+            output_names: 输出节点名称列表,指定需要返回哪些输出
+            input_feed: 输入数据字典或列表
+            data_format: 输入数据格式,"nchw"或"nhwc"
+            **kwargs: 其他运行时参数
+            
+        Returns:
+            list: 模型输出结果列表,如果指定了output_names则只返回指定的输出
+        """
+        if input_feed is None:
+            logger.error("input_feed不能为None")
+            raise ValueError("input_feed不能为None")
+            
+        # 准备输入数据
+        if isinstance(input_feed, dict):
+            # 如果是字典,按照模型输入顺序排列
+            inputs = []
+            input_map = {tensor.name: i for i, tensor in enumerate(self.input_tensors)}
+            for tensor in self.input_tensors:
+                if tensor.name not in input_feed:
+                    raise ValueError(f"缺少输入: {tensor.name}")
+                inputs.append(input_feed[tensor.name])
+        elif isinstance(input_feed, (list, tuple)):
+            # 如果是列表,确保长度匹配
+            if len(input_feed) != len(self.input_tensors):
+                raise ValueError(f"输入数量不匹配: 期望{len(self.input_tensors)}, 实际{len(input_feed)}")
+            inputs = list(input_feed)
+        else:
+            logger.error("input_feed必须是字典或列表类型")
+            raise ValueError("input_feed必须是字典或列表类型")
+            
+        # 执行推理
+        try:
+            logger.debug("开始执行推理")
+            all_outputs = self.runtime.inference(inputs=inputs, data_format=data_format)
+            
+            # 如果没有指定output_names,返回所有输出
+            if output_names is None:
+                return all_outputs
+                
+            # 获取指定的输出
+            output_map = {tensor.name: i for i, tensor in enumerate(self.output_tensors)}
+            selected_outputs = []
+            for name in output_names:
+                if name not in output_map:
+                    raise ValueError(f"未找到输出节点: {name}")
+                selected_outputs.append(all_outputs[output_map[name]])
+                    
+            return selected_outputs
+            
+        except Exception as e:
+            logger.error(f"推理执行失败: {str(e)}")
+            raise RuntimeError(f"推理执行失败: {str(e)}")
+        
+    def close(self):
+        """
+        关闭会话,释放资源
+        """
+        if self.runtime is not None:
+            logger.info("正在释放运行时资源")
+            self.runtime.release()
+            self.runtime = None
+            
+    def __enter__(self):
+        return self
+        
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def end_profiling(self) -> Optional[str]:
+        """
+        结束性能分析的存根方法
+        
+        Returns:
+            Optional[str]: None
+        """
+        warnings.warn("end_profiling()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return None
+        
+    def get_profiling_start_time_ns(self) -> int:
+        """
+        获取性能分析开始时间的存根方法
+        
+        Returns:
+            int: 0
+        """
+        warnings.warn("get_profiling_start_time_ns()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return 0
+        
+    def get_modelmeta(self) -> Dict[str, str]:
+        """
+        获取模型元数据的存根方法
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_modelmeta()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+        
+    def get_session_options(self) -> SessionOptions:
+        """
+        获取会话选项
+        
+        Returns:
+            SessionOptions: 当前会话选项
+        """
+        return self.options
+        
+    def get_providers(self) -> List[str]:
+        """
+        获取当前使用的providers的存根方法
+        
+        Returns:
+            List[str]: ["CPUExecutionProvider"]
+        """
+        warnings.warn("get_providers()是存根方法,始终返回CPUExecutionProvider", RuntimeWarning, stacklevel=2)
+        return ["CPUExecutionProvider"]
+        
+    def get_provider_options(self) -> Dict[str, Dict[str, str]]:
+        """
+        获取provider选项的存根方法
+        
+        Returns:
+            Dict[str, Dict[str, str]]: 空字典
+        """
+        warnings.warn("get_provider_options()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {} 
+
+    def get_session_config(self) -> Dict[str, str]:
+        """
+        获取会话配置的存根方法
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_session_config()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def get_session_state(self) -> Dict[str, str]:
+        """
+        获取会话状态的存根方法
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_session_state()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def set_session_config(self, config: Dict[str, str]) -> None:
+        """
+        设置会话配置的存根方法
+        
+        Args:
+            config: 会话配置字典
+        """
+        warnings.warn("set_session_config()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def get_memory_info(self) -> Dict[str, int]:
+        """
+        获取内存使用信息的存根方法
+        
+        Returns:
+            Dict[str, int]: 空字典
+        """
+        warnings.warn("get_memory_info()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def set_memory_pattern(self, enable: bool) -> None:
+        """
+        设置内存模式的存根方法
+        
+        Args:
+            enable: 是否启用内存模式
+        """
+        warnings.warn("set_memory_pattern()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def disable_memory_pattern(self) -> None:
+        """
+        禁用内存模式的存根方法
+        """
+        warnings.warn("disable_memory_pattern()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def get_optimization_level(self) -> int:
+        """
+        获取优化级别的存根方法
+        
+        Returns:
+            int: 0
+        """
+        warnings.warn("get_optimization_level()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return 0
+
+    def set_optimization_level(self, level: int) -> None:
+        """
+        设置优化级别的存根方法
+        
+        Args:
+            level: 优化级别
+        """
+        warnings.warn("set_optimization_level()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def get_model_metadata(self) -> Dict[str, str]:
+        """
+        获取模型元数据的存根方法(与get_modelmeta不同的接口)
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_model_metadata()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def get_model_path(self) -> str:
+        """
+        获取模型路径
+        
+        Returns:
+            str: 模型文件路径
+        """
+        return self.model_path
+
+    def get_input_type_info(self) -> List[Dict[str, str]]:
+        """
+        获取输入类型信息的存根方法
+        
+        Returns:
+            List[Dict[str, str]]: 空列表
+        """
+        warnings.warn("get_input_type_info()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return []
+
+    def get_output_type_info(self) -> List[Dict[str, str]]:
+        """
+        获取输出类型信息的存根方法
+        
+        Returns:
+            List[Dict[str, str]]: 空列表
+        """
+        warnings.warn("get_output_type_info()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return []