CANN/cann-recipes-infer：在昇腾Atlas A2/A3环境上适配SANA-Video模型的推理

在昇腾Atlas A2/A3环境上适配SANA-Video模型的推理【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-inferSANA-Video模型是一个多模态视频生成模型提供了文生视频功能。本项目旨在提供 SANA-Video 模型的 Atlas A2/A3 适配版本为开发者开展相关 NPU 迁移工作提供参考。本项目基于NPU主要完成以下优化点具体内容可至NPU SANA-Video模型推理优化实践查看转换1×1 conv2d为matmul计算提升性能调换时序conv2d的hw轴提升性能支持NPU npu_rotary_mul融合算子支持NPU npu_rms_norm融合算子支持NPU npu_fusion_attention融合算子执行样例本样例支持昇腾Atlas A2/A3环境的单机单卡推理和单机多卡DP推理。使用一站式平台的用户可直接跳转 「一站式平台的快速启动」章节。CANN环境准备1.安装CANN软件包本样例的编译执行依赖CANN开发套件包cann-toolkit与CANN二进制算子包cann-kernels支持的CANN软件版本为CANN 8.5.0。请从软件包下载地址下载Ascend-cann-toolkit_{version}_linux-{arch}.run和Ascend-cann-{soc}-ops_{version}_linux-{arch}.run软件包并参考CANN安装文档进行安装。2.安装Ascend Extension for PyTorchtorch_npuAscend Extension for PyTorchtorch_npu为支撑PyTorch框架运行在NPU上的适配插件本样例支持的Ascend Extension for PyTorch版本为torch-npu 7.1.0.post17支持的Torch版本为torch 2.6.0请从Ascend Extension for PyTorch插件下载并安装torch与torch_npu安装包。建议在conda环境中安装conda create -n sana python3.10 conda activate sana依赖安装本仓库依赖 SANA-Video 的开源代码。 首先在项目目录拉取SANA-Video源代码git clone https://github.com/NVlabs/Sana.git进入到源代码目录cd Sana切换到指定版本git checkout 08c656c3回到原目录cd ..拉取本仓库代码git clone https://gitcode.com/cann/cann-recipes-infer.git将Sana仓库的代码以非覆盖模式复制到本项目目录下cp -rn Sana/* cann-recipes-infer/models/sana-video/说明models/sana-video现在只保留 NPU patch 层、入口脚本与文档绝大部分diffusion/源码来自上游 Sana。推理入口会在导入上游模块前自动执行patches.apply_all()patches/__init__.py中包含NPU适配 patch包括patch_triton_rms_norm_import()和patch_wan_rotary_npu()patches/npu_patches.py中包含优化相关 patch并通过apply_npu_optimization_patches()统一应用。安装依赖cd cann-recipes-infer/models/sana-video pip install -e . pip uninstall -y opencv-python pip install opencv-python-headless4.8.0.76说明依赖中的mmengine会引入opencv-python本样例需要opencv-python-headless4.8.0.76因此需切换为 headless 版本。编译安装mmcv库(注意1.x分支才有Registry模块)git clone -b 1.x https://github.com/open-mmlab/mmcv.git cd mmcv MMCV_WITH_OPS1 FORCE_NPU1 python setup.py install cd ..模型权重普通infer.sh默认使用本地离线权重路径运行前需准备以下文件或目录并按实际存放位置修改config/2b_480p_single.yaml权重类型YAML 字段默认相对路径说明DiT / transformermodel_args.model_path./SANA-Video_2B_480p/checkpoints/SANA_Video_2B_480p.pthSANA-Video 主模型权重VAEmodel_args.vae.vae_pretrained./SANA-Video_2B_480p/vae/Wan2.1_VAE.pthWanVAE 权重文本编码器model_args.text_encoder.text_encoder_name./gemma-2-2b-it本地 Gemma HuggingFace 目录需包含config.json、tokenizer 文件和模型权重权重可从 SANA-Video_2B_480p 和 gemma-2-2b-it 下载后放置到上述目录。infer_platform.sh面向一站式平台场景会使用config/2b_480p_single_platform.yaml和脚本内的WEIGHTS_DIR自动下载/定位权重。快速启动本样例通过bash infer.sh拉起推理参数集中在config/*.yaml维护启动器使用accelerate。1. 配置 CANN 环境变量执行推理前先完成 CANN 环境变量配置参考 Ascend 社区source /usr/local/Ascend/ascend-toolkit/set_env.sh注意PYTHONPATH不会也不应全局设置项目根路径mm_function.sh会在accelerate子进程中按需注入避免影响 CANN TBE 编译器的内置 Python 解释器。2. 选择推理配置config/目录下已预置以下 YAMLYAML 文件卡数启动器适用场景2b_480p_single.yaml1acceleratemixed_precision: bf16SANA-Video 2B 480p 单卡文生视频本地离线权重2b_480p_single_a8w8.yaml1acceleratemixed_precision: bf16SANA-Video 2B 480p 做mxfp_a8w8单卡文生视频2b_480p_single_a4w4.yaml1acceleratemixed_precision: bf16SANA-Video 2B 480p 做mxfp_a4w4单卡文生视频(部分回退a8w8)2b_480p_single_platform.yaml1python 直接拉起一站式平台专用模板由infer_platform.sh读取并生成临时本地权重配置YAML 中的关键字段launcher: accelerate launcher_args: mixed_precision: bf16 # 透传 accelerate launch --mixed_precisionbf16 env_vars: DISABLE_XFORMERS: 1 # 屏蔽 xformers走 NPU 融合算子 HF_HUB_OFFLINE: 1 # 普通 infer.sh 默认离线运行 TRANSFORMERS_OFFLINE: 1 HF_DATASETS_OFFLINE: 1 model_args: config: configs/sana_video_config/Sana_2000M_480px_AdamW_fsdp.yaml model_path: ./checkpoints/SANA_Video_2B_480p.pth vae.vae_pretrained: ./vae/Wan2.1_VAE.pth text_encoder.text_encoder_name: ./gemma-2-2b-it txt_file: asset/samples/video_prompts_samples.txt cfg_scale: 6 motion_score: 30 flow_shift: 8 work_dir: output/sana_t2v_video_results model.fp32_attention: False # 用字符串 Falseargparse 侧按字面量接收 quant_type: bf16 # 8bit量化:a8w8 4bit量化:a4w43. 修改推理输入根据需要修改 YAML 中model_args下的字段YAML 字段含义config推理使用的配置文件上游 Sana 的 YAMLmodel_path推理使用的 DiT / transformer 权重文件路径vae.vae_pretrained推理使用的 VAE 权重文件路径text_encoder.text_encoder_name推理使用的文本编码器普通离线配置中填写本地 Gemma HuggingFace 目录txt_file推理使用的文本提示词文件cfg_scale提示词对齐强度motion_score运动强度分数flow_shift流偏移参数用于调整扩散模型去噪时间步work_dir生成视频输出路径model.fp32_attentionattn 是否使用 fp32 精度推理时置False以提升性能quant_type量化配置可选 bf16 / a8w8 / a4w4透传规则适用于所有字段YAML 类型命令行效果字符串/数字--key value布尔true--keyflag布尔false忽略列表--key v1 v2 …加引号字符串False--key False用于 pyrallis/Hydra 类解析器4. 拉起推理bash infer.sh拉起过程会自动解析 YAML 中的world_size、master_port、entry_script、env_vars、launcher_args、model_args设置通用 NPU 优化环境变量PYTORCH_NPU_ALLOC_CONF、TASK_QUEUE_ENABLE、CPU_AFFINITY_CONF、TOKENIZERS_PARALLELISM及 HCCL 通信配置在res/YYYYMMDD/model_name/下自动创建日志目录并将推理输出tee到log_timestamp.log调用accelerate launch --num_processesworld_size --num_machines1 --main_process_portmaster_port --mixed_precisionbf16 inference_video_scripts/inference_sana_video.py model_args。如需切换到多卡 DP可将 YAML 中的world_size调整为目标卡数mm_function.sh会自动转换为accelerate --num_processes。一站式平台的快速启动本章节面向使用一站式平台的用户平台已预置完整的 CANN 环境按以下步骤即可在单卡上完成 SANA-Video 文生视频推理。1. 安装 Miniconda 并创建 Python 环境若当前环境未预装 conda在用户目录下安装 Miniconda无需 root 权限wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh -O /tmp/miniconda.sh bash /tmp/miniconda.sh -b -p $HOME/miniconda3 source $HOME/miniconda3/etc/profile.d/conda.sh echo source $HOME/miniconda3/etc/profile.d/conda.sh ~/.bashrc # 永久生效说明x86_64 机器将命令中的aarch64替换为x86_64。创建并激活 Python 环境conda create -n sana python3.10 -y --override-channels -c conda-forge conda activate sana2. 修改infer_platform.sh中的两处路径打开infer_platform.sh按本机实际情况修改「User configuration」段的两个路径WEIGHTS_DIR/mnt/workspace/gitCode/cann/models # 权重与 HF cache 的存放目录 CANN_SET_ENV/home/developer/Ascend/ascend-toolkit/set_env.sh # 平台 CANN 的 set_env.sh3. 一键拉起推理cd cann-recipes-infer/models/sana-video bash infer_platform.sh前置条件已完成 §1conda 环境已激活。脚本本身会依次完成source CANN 环境脚本 → 安装 torch / torch_npu → 合入上游 Sana 源码 → 安装项目依赖与 NPU mmcv → 通过 hf-mirror 下载权重 → 生成指向本地权重的临时 YAML → 调用 python 启动单卡推理。各步骤均幂等已完成的环节会自动跳过。说明infer_platform.sh不直接读取普通离线配置2b_480p_single.yaml而是读取2b_480p_single_platform.yaml作为平台模板再按WEIGHTS_DIR生成临时 YAML。普通infer.sh与一站式平台启动互不影响。生成结果位于output/sana_t2v_video_results/vis/*.mp4一条 prompt 对应一个 mp4。性能数据本样例在Atlas A2/A3的推理性能如下表所示| 规格| 单步时延(s) | Diffusion sampling总时长(s) | |--|--| --| | 480p81f | 2.75 | 132 |本样例在 Atlas A5 的推理性能如下表所示 |量化方式 | 规格| 单步时延(s) | Diffusion sampling总时长(s) | 整网耗时(s) | |--|--|--|--| -- | |bf16| 480p81f | 1.67 | 83.5 | 90.44 | |8bit| 480p81f | 1.46 | 73 | 79.38 | |4bit| 480p81f | 1.43 | 71.5 | 77.92 |【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考