ACE-Step Captioner 在昇腾 NPU 上的适配

pip install torch_npu transformers accelerate modelscope qwen-omni-utils soundfile librosa safetensors

from modelscope import snapshot_download
model_dir = snapshot_download('ACE-Step/acestep-captioner', cache_dir='./model_cache')

python inference.py \
    --model_path ./model_cache/ACE-Step/acestep-captioner \
    --audio_path ./test_audio.wav \
    --device npu:0

import os
import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu

from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
from qwen_omni_utils import process_mm_info

os.environ["TORCH_NPU_ALLOC_CONF"] = "expandable_segments:True"

model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
    "ACE-Step/acestep-captioner",
    torch_dtype=torch.bfloat16,
    device_map="npu:0",
)
processor = Qwen2_5OmniProcessor.from_pretrained("ACE-Step/acestep-captioner")

conversation = [
    {"role": "user", "content": [
        {"type": "text", "text": "*Task* Describe this audio in detail"},
        {"type": "audio", "audio": "path/to/audio.wav"},
    ]},
]

text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
audios, images, videos = process_mm_info(conversation, use_audio_in_video=False)
inputs = processor(text=text, audio=audios, images=images, videos=videos,
                   return_tensors="pt", padding=True, use_audio_in_video=False)
inputs = inputs.to(model.device).to(model.dtype)

text_ids = model.generate(**inputs, use_audio_in_video=False, return_audio=False, max_new_tokens=512)
result = processor.batch_decode(text_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
print(result)

A single, sustained, and heavily distorted sawtooth synthesizer note creates a harsh, buzzing texture. The sound is static and monolithic, with a gritty, abrasive character and a very long, abrupt decay that ends in silence.

A single, sustained, and heavily distorted sawtooth synthesizer note creates a harsh, buzzing texture. The sound is static and monolithic, with a gritty, abrasive quality and a very long, abrupt decay that ends in silence.

NPU available: 4 x Ascend910B2
Loading model from: /home/openmind/acestep-captioner/model_cache/ACE-Step/acestep-captioner
Target device: npu:0
Loading checkpoint shards: 100%|██████████| 5/5 [00:02<00:00,  2.26it/s]
Model loaded: 10.73B parameters
Model dtype: torch.bfloat16

Input audio: /home/openmind/acestep-captioner/test_audio.wav
Running inference...
Inference completed in 13.38s (48 tokens, 3.6 tok/s)

============================================================
Caption Result:
============================================================
system
You are a helpful assistant.
user
*Task* Describe this audio in detail
assistant
A single, sustained, and heavily distorted sawtooth synthesizer note creates a harsh, buzzing texture. The sound is static and monolithic, with a gritty, abrasive character and a very long, abrupt decay that ends in silence.
============================================================