基于K100_AI四卡全离线部署的舞蹈编排系统设计与实现

一、引言

近年来，生成式人工智能在视频创作领域展现出巨大潜力，尤其是基于扩散模型的文本到视频生成技术，能够根据简单的文本提示和参考图像生成连贯、逼真的动态内容。然而，针对舞蹈这类强节奏、高精度姿态控制的任务，通用视频生成模型往往难以捕捉音乐节拍与动作韵律的同步关系。为此，研究者提出了Wan-Dancer-14B等专用模型，通过引入音乐特征注入和姿态控制机制，实现音乐驱动的舞蹈视频生成。但在实际部署中，该类模型参数量大（14B级）、计算资源需求高，且对推理速度和显存管理有严格要求，传统单卡环境难以满足实时或高效生产需求。这篇内容以K100_AI四卡GPU服务器，设计并实现了一套全离线部署的舞蹈编排系统。系统采用分布式推理框架，利用torchrun进行多卡并行，结合序列并行（USP）技术，有效提升了模型推理效率。整个流程分为全局关键帧生成和局部高清细化两个阶段，通过音乐特征提取、关键帧采样和动态帧率控制，实现了从静态参考图到完整舞蹈视频的全自动化生成。该系统完全运行于本地环境，无需联网，保障了数据隐私与稳定性，适用于影视预可视化、虚拟人舞蹈创作等实际应用场景。

二、方案设计

系统整体采用“先全局后局部”的两阶段生成策略。第一阶段，基于全局模型（global_model.safetensors）生成低分辨率的关键帧序列，该序列捕捉舞蹈的宏观姿态变化和音乐节拍对齐；第二阶段，利用局部模型（local_model.safetensors）以全局视频为条件，逐段生成高分辨率（1280×720）的视频片段，最终拼接为完整作品。两阶段均采用相同的底层架构：WanVideoPipeline集成T5文本编码器、VAE、CLIP视觉编码器以及DiT（Diffusion Transformer）主干网络，并开启音乐注入、参考图像、动态帧率及统一模型（unimodel）等增强功能。分布式方案上，系统通过环境变量MASTER_ADDR、MASTER_PORT和WORLD_SIZE控制，启动4个进程分别绑定GPU，利用序列并行将长序列拆分为多个子序列，降低单卡显存压力。音乐处理方面，使用librosa提取包络、MFCC、色度特征和节拍/峰值热编码，形成多维特征向量，与视频帧率对齐后注入扩散过程。关键帧控制模块通过掩码（mask）标记指定帧，确保生成视频首帧为输入图像，且中间帧与全局视频采样帧一致，从而保证动作连贯性。最后，所有片段视频裁剪、添加音轨并拼接，输出带音乐的完整舞蹈视频。

本系统的代码源自Wan-Dancer-14B大模型官方示例代码，代码下载链接：https://github.com/Wan-Video/Wan-Dancer，其代码原始实现基于NVIDIA CUDA环境。为适配国产算力平台，本文对代码进行了针对性改造，使其能够在ROCM异构计算环境下稳定运行，主要包括对torch.version.cuda的兼容性处理以及对torch.cuda.get_arch_list的接口适配，确保依赖库的版本检查与架构识别在ROCM环境中正常通过。两阶段生成流程（全局关键帧规划与局部时序细化）均遵循Wan-Dancer官方提出的分层生成框架。

如果你在研究中使用了官方代码或框架，请引用：

@article{wan-dancer-2026,

title={Wan-Dancer: A Hierarchical Framework for Minute-scale Coherent Music-to-Dance Generation},

author={Mingyang Huang, Peng Zhang, Li Hu, Guangyuan Wang, Bang Zhang},

website={https://humanaigc.github.io/wan-dancer/},

url={https://arxiv.org/abs/2607.09581},

year={2026}

}

程序工作流程图如下：

三、实施方法及代码

3.1硬件环境

本方案的硬件平台为一台H3C服务器，配置如下：

需占用四张K100_AI显卡

3.2软件栈

本方案为了方便离线部署，软件栈基于Docker容器化技术构建，使用经过海光DCU适配的镜像：

镜像：

harbor.sourcefind.cn:5443/dcu/admin/base/custom:vllm0.15.1-ubuntu22.04-dtk26.04-0130-py3.10-20260220

在该镜像里用四卡部署Wan-Dancer-14B大模型。

3.3软件环境：

采用python虚拟环境部署（docker容器里执行）：

python -m venv venv_wan_dancersource venv_wan_dancer/bin/activatepip listPackageVersion---------------------- --------------------------------accelerate 1.14.0annotated-doc0.0.4anyio4.14.2audioread3.1.0beautifulsoup4 4.15.0blinker1.9.0certifi2026.6.17cffi 2.1.0charset-normalizer 3.4.9click8.4.2cloudpickle3.1.2controlnet_aux 0.0.7cuda-bindings13.3.1cuda-pathfinder1.5.6cuda-toolkit 13.0.3.0decorator5.3.1diffusers0.34.0DistVAE0.0.0b5einops 0.8.2exceptiongroup 1.3.1filelock 3.29.7flash_attn 2.8.3+das.opt1.dtk2604.torch271Flask3.1.3fsspec 2026.6.0ftfy 6.3.1h110.16.0hf-xet 1.5.1httpcore 1.0.9httpx0.28.1huggingface_hub0.36.2idna 3.18ImageIO2.37.3imageio-ffmpeg 0.6.0importlib_metadata 9.0.0iniconfig2.3.0itsdangerous 2.2.0Jinja2 3.1.6joblib 1.5.3lazy-loader0.5librosa0.11.0llvmlite 0.48.0loguru 0.7.3markdown-it-py 4.2.0MarkupSafe 3.0.3mdurl0.1.2modelscope 1.38.1modelscope-hub 0.1.7moviepy2.2.1mpmath 1.3.0msgpack1.2.1networkx 3.4.2ninja1.13.0numba0.66.0numpy1.26.4nvidia-cublas13.1.1.3nvidia-cuda-cupti13.0.85nvidia-cuda-nvrtc13.0.88nvidia-cuda-runtime13.0.96nvidia-cudnn-cu139.20.0.48nvidia-cufft 12.0.0.61nvidia-cufile1.15.1.6nvidia-curand10.4.0.35nvidia-cusolver12.0.4.66nvidia-cusparse12.6.3.3nvidia-cusparselt-cu13 0.8.1nvidia-ml-py 13.610.43nvidia-nccl-cu13 2.29.7nvidia-nvjitlink 13.3.33nvidia-nvshmem-cu133.4.5nvidia-nvtx13.0.85opencv-python5.0.0.93optimum-quanto 0.2.7packaging26.2pandas 2.3.3pillow 11.3.0pip22.0.2platformdirs 4.10.0pluggy 1.6.0pooch1.9.0proglog0.1.12protobuf 7.35.1psutil 7.2.2pycparser3.0Pygments 2.20.0pynvml 13.0.1pytest 9.1.1python-dateutil2.9.0.post0python-dotenv1.2.2pytz 2026.2PyYAML 6.0.3regex2026.7.10requests 2.34.2rich 15.0.0safetensors0.8.0scikit-image 0.25.2scikit-learn 1.7.2scipy1.15.3sentencepiece0.2.2setuptools 83.0.0shellingham1.5.4six1.17.0soundfile0.14.0soupsieve2.8.4soxr 1.1.0sympy1.14.0threadpoolctl3.6.0tifffile 2025.5.10timm 1.0.28tokenizers 0.20.3tomli2.4.1torch2.7.1+das.opt1.dtk2604torchvision0.22.0+das.opt1.dtk2604.torch271tqdm 4.68.4transformers 4.46.2triton 3.7.1typer0.26.8typing_extensions4.16.0tzdata 2026.3urllib32.7.0wcwidth0.8.2Werkzeug 3.1.8wheel0.47.0xfuser 0.4.0yunchang 0.5.0zipp 4.1.0

注意事项：

安装的torch系列及相关模块版本如下：

pip list|grep torchflash_attn 2.8.3+das.opt1.dtk2604.torch271torch2.7.1+das.opt1.dtk2604torchvision0.22.0+das.opt1.dtk2604.torch271

3.4软件项目：

系统代码参照https://github.com/Wan-Video/Wan-Dancer，为了适应ROCM环境以及离线部署，代码做了相应修改如下：

gen_video_global-chinese.sh代码（生成全局关键帧视频）：# Activate your venv, and run the script in the project root dirseed=0image_path='gen_video/ref_image/1001.jpg'prompt_path='gen_video/prompt/古典舞_local.txt'music_path='gen_video/music/ChineseClassicDance.WAV'output_folder="outputs/global_video/"timestamp=$(date "+%Y%m%d_%H%M%S_%N")num_inference_steps=48cfg_scale=5main_process_ip=${MASTER_ADDR:-localhost} main_process_port=${MASTER_PORT:-8899}machine_rank=${RANK:-0}num_machines=${WORLD_SIZE:-1}gpu_count=4num_processes=$((num_machines * gpu_count))torchrun --nproc_per_node=${gpu_count} --master_addr=${main_process_ip} --master_port=${main_process_port} --nnodes=${num_machines} --node_rank=${machine_rank} gen_video/gen_video_global.py --seed $seed --image_path $image_path --prompt_path $prompt_path --music_path $music_path --output_folder $output_folder --timestamp ${timestamp} --num_inference_steps $num_inference_steps --cfg_scale $cfg_scalegen_video/gen_video_global.py代码（生成全局关键帧视频）：import librosaimport osimport torchif torch.cuda.is_available() and torch.version.cuda is None:# 1. 假装有一个 CUDA 版本，使 xfuser 的版本检查通过torch.version.cuda = "12.4"# 2. 修补 get_arch_list，返回 CUDA 格式的架构列表，避免 _extract_arch_version 索引错误import torch.cudaoriginal_get_arch_list = torch.cuda.get_arch_listdef fake_get_arch_list():# 返回一个 CUDA 风格的架构（例如 sm_70），这样 split("_")[1] 就能正常取出数字return ["sm_120"]torch.cuda.get_arch_list = fake_get_arch_listimport torchvisionimport torch.distributed as distimport numpy as npimport moviepy as mpyfrom loguru import loggerimport PILfrom PIL import Image, ImageFilefrom diffsynth import save_videofrom diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfigFPS = 30def parse_args():import argparseparser = argparse.ArgumentParser(description="input parameters")parser.add_argument("--image_path",type=str,default="",required=True,help="image_path",)parser.add_argument("--prompt_path",type=str,default="",required=True,help="prompt_path",)parser.add_argument("--music_path",type=str,default="",required=True,help="music_path",)parser.add_argument("--seed",type=int,default=0,required=True,help="seed",)parser.add_argument("--dit_model_path",type=str,default="Wan-Dancer-14B/wan_dancer_model/global_model.safetensors",required=False,help="dit model path",)parser.add_argument("--height",type=int,default=1280,required=False,help="height",)parser.add_argument("--width",type=int,default=720,required=False,help="width",)parser.add_argument("--music_inject_layers",type=str,default="0, 4, 8, 12, 16, 20, 24, 27",help="music inject layers",)parser.add_argument("--num_frames",type=int,default=149,help="num frames",)parser.add_argument("--output_folder",type=str,default="outputs/global_video",help="output folder",)parser.add_argument("--timestamp",type=str,default="",help="timestamp",)parser.add_argument("--sigma_shift",type=int,default=5,help="sigma shift",)parser.add_argument("--num_inference_steps",type=int,default=48,help="number inference steps",)parser.add_argument("--cfg_scale",type=int,default=5,help="cfg scale",)args = parser.parse_args()return argsdef init_dit_model(args):world_size = int(os.environ.get("WORLD_SIZE", 1))assert world_size == 4, "WORLD_SIZE must be 8"## usp configulysses_degree = world_sizesequence_parallel_degree = world_sizering_degree = sequence_parallel_degree // ulysses_degreedata_parallel_degree = 1usp_config = {"data_parallel_degree": data_parallel_degree,"sequence_parallel_degree": sequence_parallel_degree,"ring_degree": ring_degree,"ulysses_degree": ulysses_degree,}## load modelspipe = WanVideoPipeline.from_pretrained(torch_dtype=torch.bfloat16,device="cuda",model_configs=[ModelConfig(path="/home/models/Wan-Dancer-main/Wan-Dancer-14B/global_model.safetensors",model_id="/home/models/Wan-Dancer-main/Wan-Dancer-14B",origin_file_pattern="global_model.safetensors",offload_device="cuda",),ModelConfig( path="/home/models/Wan-Dancer-main/Wan-Dancer-14B/models_t5_umt5-xxl-enc-bf16.pth",model_id="/home/models/Wan-Dancer-main/Wan-Dancer-14B",origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth",offload_device="cuda",),ModelConfig(path="/home/models/Wan-Dancer-main/Wan-Dancer-14B/Wan2.1_VAE.pth",model_id="/home/models/Wan-Dancer-main/Wan-Dancer-14B",origin_file_pattern="Wan2.1_VAE.pth",offload_device="cuda",),ModelConfig(path="/home/models/Wan-Dancer-main/Wan-Dancer-14B/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth",model_id="/home/models/Wan-Dancer-main/Wan-Dancer-14B",origin_file_pattern="models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth",offload_device="cuda",),],tokenizer_config=ModelConfig(path="/home/models/Wan-Dancer-main/Wan-Dancer-14B/google/umt5-xxl/",model_id="/home/models/Wan-Dancer-main/Wan-Dancer-14B", origin_file_pattern="google/umt5-xxl/"),skip_download=True,redirect_common_files=False,use_usp=True,usp_config=usp_config,dit_model_type=1,# 1 for our trained modelenable_music_inject=True,enable_refimage=True,enable_global=True,enable_dynamicfps=True,enable_unimodel=True,)pipe.enable_vram_management()return pipedef get_music_base_feature(music_path, output_path, fps=30):hop_length = 512sr = fps * hop_lengthdata, sr = librosa.load(music_path, sr=sr)sr = 22050envelope = librosa.onset.onset_strength(y=data, sr=sr)mfcc = librosa.feature.mfcc(y=data, sr=sr, n_mfcc=20).Tchroma = librosa.feature.chroma_cens(y=data, sr=sr, hop_length=hop_length, n_chroma=12).Tpeak_idxs = librosa.onset.onset_detect(onset_envelope=envelope.flatten(), sr=sr, hop_length=hop_length)peak_onehot = np.zeros_like(envelope, dtype=np.float32)peak_onehot[peak_idxs] = 1.0start_bpm = librosa.beat.tempo(y=librosa.load(music_path)[0])[0]_, beat_idxs = librosa.beat.beat_track(onset_envelope=envelope,sr=sr,hop_length=hop_length,start_bpm=start_bpm,tightness=100,)beat_onehot = np.zeros_like(envelope, dtype=np.float32)beat_onehot[beat_idxs] = 1.0audio_feature = np.concatenate([envelope[:, None], mfcc, chroma, peak_onehot[:, None], beat_onehot[:, None]],axis=-1,)logger.info(f"audio feature shape: {audio_feature.shape} for {music_path}")np.save(output_path, audio_feature)def crop_and_resize(image: PIL.Image.Image, target_width=720, target_height=1280):width, height = image.sizescale = min(target_width / width, target_height / height)resized_height = round(height * scale)resized_width = round(width * scale)image = torchvision.transforms.functional.resize(image,(resized_height, resized_width),interpolation=torchvision.transforms.InterpolationMode.BILINEAR,)## pad 127 to target sizetarget_image = np.ones((target_height, target_width, 3), dtype=np.uint8) * 127tl_x = (target_width - resized_width) // 2tl_y = (target_height - resized_height) // 2br_x = tl_x + resized_widthbr_y = tl_y + resized_heighttarget_image[tl_y:br_y, tl_x:br_x, :] = np.array(image, dtype=np.uint8)image = Image.fromarray(target_image)return image, (tl_x, tl_y, br_x, br_y)def gen_video_single(pipe, prompt, input_config):negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"img_path = input_config["img_path"]ImageFile.LOAD_TRUNCATED_IMAGES = Trueimg = Image.open(img_path)target_height = input_config["target_height"]target_width = input_config["target_width"]## resize input imageimg, (tl_x, tl_y, br_x, br_y) = crop_and_resize(img, target_width=target_width, target_height=target_height)## refimageenable_refimage = input_config.get("enable_refimage", False)if enable_refimage:refimage_path = input_config["refimage_path"]refimage = Image.open(refimage_path)refimage, (tl_x, tl_y, br_x, br_y) = crop_and_resize(refimage, target_width=target_width, target_height=target_height)else:refimage = None## music featureenable_music_inject = input_config["enable_music_inject"]music_feature = input_config["music_feature"]## input parametersseed = input_config["seed"]num_frames = input_config["num_frames"]interp_mode = input_config["interp_mode"]enable_global = input_config.get("enable_global", False)enable_dynamicfps = input_config.get("enable_dynamicfps", False)input_fps = input_config.get("input_fps", 30.0)enable_vae_decode_framewise = input_config.get("enable_vae_decode_framewise", False)enable_skip_layer = input_config.get("enable_skip_layer", False)enable_unimodel = input_config.get("enable_unimodel", False)sigma_shift = input_config.get("sigma_shift", 5)num_inference_steps = input_config.get("num_inference_steps", 48)cfg_scale = input_config.get("cfg_scale", 5)## maskmask = np.zeros(num_frames, dtype=np.int32)mask[0] = 1keyframes = np.zeros((num_frames, target_height, target_width, 3), dtype=np.uint8)keyframes[mask == 1] = np.array(img, dtype=np.uint8)keyframes = [Image.fromarray(img.astype("uint8")) if isinstance(img, np.ndarray) else imgfor img in keyframes]mask = torch.tensor(mask).to(torch.int32)keyframes_mask = mask## generated videovideo = pipe(prompt=prompt,negative_prompt=negative_prompt,input_image=None,num_inference_steps=num_inference_steps,seed=seed,tiled=True,height=target_height,width=target_width,enable_music_inject=enable_music_inject,music_feature=music_feature,num_frames=num_frames,interp_mode=interp_mode,enable_refimage=enable_refimage,refimage=refimage,enable_global=enable_global,keyframes=keyframes,keyframes_mask=keyframes_mask,enable_dynamicfps=enable_dynamicfps,input_fps=input_fps,enable_vae_decode_framewise=enable_vae_decode_framewise,enable_skip_layer=enable_skip_layer,enable_unimodel=enable_unimodel,sigma_shift=sigma_shift,cfg_scale=cfg_scale,)if dist.get_rank() == 0:save_video_path = input_config["save_video_path"]tmp_video_path = save_video_path[:-4] + "_tmp.mp4"if enable_vae_decode_framewise:fps = 8# 8fps for vae decode framewisesave_video(video, tmp_video_path, fps=fps, quality=5)else:save_video(video, tmp_video_path, fps=FPS, quality=5)## crop generated videovideo = mpy.VideoFileClip(tmp_video_path)croper = mpy.video.fx.Crop(x1=tl_x, y1=tl_y, x2=br_x, y2=br_y)video = croper.apply(video)video.write_videofile(save_video_path, codec="libx264", audio_codec="aac")os.remove(tmp_video_path)def gen_video(pipe,music_feature_path,img_path,prompt,output_video_path,seed=0,max_pixels=1280 * 720,height=1280,width=720,num_frames=149,enable_refimage=False,refimage_path=None,enable_global=False,enable_dynamicfps=False,enable_vae_decode_framewise=False,enable_skip_layer=False,enable_unimodel=False,sigma_shift=5,num_inference_steps=48,cfg_scale=5,):## input parametersinput_config = {}input_config["img_path"] = img_pathinput_config["enable_music_inject"] = Truemusic_feature = np.load(music_feature_path)music_feature = torch.from_numpy(music_feature).to(dtype=torch.bfloat16, device="cuda")input_config["music_feature"] = music_featureinput_config["max_pixels"] = max_pixelsinput_config["save_video_path"] = output_video_pathinput_config["seed"] = seedinput_config["target_height"] = heightinput_config["target_width"] = widthinput_config["num_frames"] = num_framesinput_config["interp_mode"] = "bilinear"input_config["enable_refimage"] = enable_refimageinput_config["refimage_path"] = refimage_pathinput_config["enable_global"] = enable_globalinput_fps = 30.0 / int(music_feature.shape[0] / 149.0 + 0.5)input_fps = "{:.4f}".format(input_fps)input_config["input_fps"] = float(input_fps)logger.info(f"input fps: {input_fps}")## update promptprompt += f"帧率是{input_fps}"logger.info(f"prompt: {prompt}")input_config["enable_dynamicfps"] = enable_dynamicfpsinput_config["enable_vae_decode_framewise"] = enable_vae_decode_framewiseinput_config["enable_skip_layer"] = enable_skip_layerinput_config["enable_unimodel"] = enable_unimodelinput_config["sigma_shift"] = sigma_shiftinput_config["num_inference_steps"] = num_inference_stepsinput_config["cfg_scale"] = cfg_scale## gen videogen_video_single(pipe, prompt, input_config)def main():args = parse_args()prompt_path = args.prompt_path## promptwith open(prompt_path, "r") as f:prompt = f.read().strip()## init dit modelpipe = init_dit_model(args)final_name = (args.image_path.split("/")[-1].split(".")[0]+ "_"+ args.music_path.split("/")[-1].split(".")[0])time_name = args.timestampmusic_folder = "outputs/tmp_results/" + final_name + "_" + str(time_name)os.makedirs(music_folder, exist_ok=True)## encode musicoriginal_music_path = args.music_pathmusic_feature_path = os.path.join(music_folder, final_name + "_librosa_feature.npy")if dist.get_rank() == 0:get_music_base_feature(original_music_path, music_feature_path, fps=30)dist.barrier(device_ids=[dist.get_rank()])## generate global videoseed = args.seedimg_path = args.image_pathrefimage_path = args.image_pathoutput_video_folder = args.output_folderos.makedirs(output_video_folder, exist_ok=True)output_video_path = os.path.join(output_video_folder,final_name + "_seed" + str(seed) + "_" + str(time_name) + ".mp4",)gen_video(pipe,music_feature_path,img_path,prompt,output_video_path,seed=seed,height=args.height,width=args.width,num_frames=args.num_frames,enable_refimage=True,refimage_path=refimage_path,enable_global=True,enable_dynamicfps=True,enable_vae_decode_framewise=True,enable_skip_layer=True,enable_unimodel=True,sigma_shift=args.sigma_shift,num_inference_steps=args.num_inference_steps,cfg_scale=args.cfg_scale,)if __name__ == "__main__":main()gen_video_local.sh（生成最终高分辨率视频）：# Activate your venv, and run the script in the project root dirseed=0image_path='gen_video/ref_image/1001.jpg'prompt_path='gen_video/prompt/古典舞_local.txt'music_path='gen_video/music/ChineseClassicDance.WAV'output_folder="outputs/final_video/"timestamp=$(date "+%Y%m%d_%H%M%S_%N")num_inference_steps=24cfg_scale=5global_video_path='outputs/global_video/1001_ChineseClassicDance_seed0_20260716_123658_713049149.mp4'main_process_ip=${MASTER_ADDR:-localhost} main_process_port=${MASTER_PORT:-8089}machine_rank=${RANK:-0}num_machines=${WORLD_SIZE:-1}gpu_count=4num_processes=$((num_machines * gpu_count))torchrun --nproc_per_node=${gpu_count} --master_addr=${main_process_ip} --master_port=${main_process_port} --nnodes=${num_machines} --node_rank=${machine_rank} gen_video/gen_video_local.py --seed $seed --image_path $image_path --prompt_path $prompt_path --music_path $music_path --output_folder $output_folder --timestamp ${timestamp} --num_inference_steps $num_inference_steps --cfg_scale $cfg_scale --global_video_path $global_video_path gen_video/gen_video_local.py（生成最终高分辨率视频）：import cv2import librosaimport osimport timeimport torchif torch.cuda.is_available() and torch.version.cuda is None:# 1. 假装有一个 CUDA 版本，使 xfuser 的版本检查通过torch.version.cuda = "12.4"# 2. 修补 get_arch_list，返回 CUDA 格式的架构列表，避免 _extract_arch_version 索引错误import torch.cudaoriginal_get_arch_list = torch.cuda.get_arch_listdef fake_get_arch_list():# 返回一个 CUDA 风格的架构（例如 sm_70），这样 split("_")[1] 就能正常取出数字return ["sm_120"]torch.cuda.get_arch_list = fake_get_arch_listimport torchvisionimport torch.distributed as distimport numpy as npimport moviepy as mpyfrom loguru import loggerfrom tqdm import tqdmimport soundfile as sfimport PILfrom PIL import Image, ImageFilefrom moviepy import AudioFileClipfrom diffsynth import save_videofrom diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfigFPS = 30LAYER = 66def parse_args():import argparseparser = argparse.ArgumentParser(description="input parameters")parser.add_argument("--image_path",type=str,default='',required=True,help="image_path",)parser.add_argument("--prompt_path",type=str,default='',required=True,help="prompt_path",)parser.add_argument("--music_path",type=str,default='',required=True,help="music_path",)parser.add_argument("--seed",type=int,default=0,required=True,help="seed",)parser.add_argument("--enable_music_inject", default=False, action="store_true", help="Whether to inject music.")parser.add_argument("--height",type=int,default=1280,required=False,help="height",)parser.add_argument("--width",type=int,default=720,required=False,help="width",)parser.add_argument("--music_inject_layers",type=str,default="0, 4, 8, 12, 16, 20, 24, 27",help="music inject layers")parser.add_argument("--enable_refimage", default=False, action="store_true", help="enable refimage.")parser.add_argument("--num_frames",type=int,default=149,help="num frames",)parser.add_argument("--output_folder",type=str,default='outputs/local_video',help="output folder",)parser.add_argument("--global_video_path",type=str,default='',required=True,help="global video path",)parser.add_argument("--timestamp",type=str,default='',required=True,help="timestamp",)parser.add_argument("--enable_dynamicfps", default=False, action="store_true", help="enable dynamicfps.")parser.add_argument("--enable_skip_layer", default=False, action="store_true", help="enable skip dit layer.")parser.add_argument("--sigma_shift",type=int,default=5,help="sigma_shift",)parser.add_argument("--num_inference_steps",type=int,default=48,help="num_inference_steps",)parser.add_argument("--cfg_scale",type=int,default=5,help="cfg scale",)args = parser.parse_args()return argsdef get_music_clip_149f(original_music_path, target_music_folder):audio = AudioFileClip(original_music_path)total_duration = audio.duration audio, sr = librosa.load(original_music_path, sr=None)duration = float(149) / FPSidx = 0t = 0while t + 0.2 < total_duration: start_time = tend_time = t + durationif end_time >= total_duration:end_time = total_durationsliced_audio = audio[int(start_time * sr):int(end_time * sr)]timestamp = time.time()save_path = os.path.join(target_music_folder, str(idx).zfill(3) + '_' + str(timestamp).replace('.', '') + '.wav')sf.write(save_path, sliced_audio, sr)t += durationidx += 1def get_music_base_feature(music_path, output_path, fps=30):hop_length = 512sr = fps * hop_lengthdata, sr = librosa.load(music_path, sr=sr)sr = 22050 envelope = librosa.onset.onset_strength(y=data, sr=sr)mfcc = librosa.feature.mfcc(y=data, sr=sr, n_mfcc=20).Tchroma = librosa.feature.chroma_cens(y=data, sr=sr, hop_length=hop_length, n_chroma=12).T peak_idxs = librosa.onset.onset_detect(onset_envelope=envelope.flatten(), sr=sr, hop_length=hop_length)peak_onehot = np.zeros_like(envelope, dtype=np.float32)peak_onehot[peak_idxs] = 1.0start_bpm = librosa.beat.tempo(y=librosa.load(music_path)[0])[0]_, beat_idxs = librosa.beat.beat_track(onset_envelope=envelope,sr=sr,hop_length=hop_length,start_bpm=start_bpm,tightness=100,)beat_onehot = np.zeros_like(envelope, dtype=np.float32)beat_onehot[beat_idxs] = 1.0audio_feature = np.concatenate([envelope[:, None], mfcc, chroma, peak_onehot[:, None], beat_onehot[:, None]],axis=-1,)logger.info(f'audio feature shape: {audio_feature.shape} for {music_path}')np.save(output_path, audio_feature)def get_music_features(music_folder):dirs = [f for f in sorted(os.listdir(music_folder)) if f.endswith('.wav')]for idx, name in tqdm(enumerate(dirs)):music_path = os.path.join(music_folder, name)output_path = os.path.join(music_folder, name.replace('.wav', '_librosa_feature.npy'))if os.path.exists(output_path) is False:get_music_base_feature(music_path, output_path)def crop_and_resize(image: PIL.Image.Image, target_width=720, target_height=1280):width, height = image.sizescale = min(target_width / width, target_height / height)resized_height = round(height * scale)resized_width = round(width * scale)image = torchvision.transforms.functional.resize(image,(resized_height, resized_width),interpolation=torchvision.transforms.InterpolationMode.BILINEAR)## pad 127 to target sizetarget_image = np.ones((target_height, target_width, 3), dtype=np.uint8) * 127tl_x = (target_width - resized_width) // 2tl_y = (target_height - resized_height) // 2br_x = tl_x + resized_widthbr_y = tl_y + resized_heighttarget_image[tl_y: br_y, tl_x: br_x, :] = np.array(image, dtype=np.uint8)image = Image.fromarray(target_image)return image, (tl_x, tl_y, br_x, br_y)def process_global_video_firstlastframe(video_path, height, width, total_frames):cap = cv2.VideoCapture(video_path)frames = []while True:ret, frame = cap.read()if not ret:breakframe = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)frames.append(frame)cap.release()N = len(frames) logger.info(f'global video frame number: {N}')seg_num = int(np.ceil(total_frames / 149)) # 149 cannot be changegframe_interval_num = float(total_frames) / Nkeyframes_list = []keyframes_mask_list = []for i in range(seg_num): mask = np.zeros(149, dtype=np.int32)if i != seg_num - 1:cnt = 0while (cnt * frame_interval_num < 149 - frame_interval_num):index = int(np.ceil(frame_interval_num * cnt))mask[index] = 1cnt += 1else:end_index = total_frames - 149 * i - 1mask[end_index] = 1cnt = 0while (cnt * frame_interval_num < end_index - frame_interval_num):index = int(np.ceil(frame_interval_num * cnt))mask[index] = 1cnt += 1keyframes_mask_list.append(mask)sum = 0for mask in keyframes_mask_list:sum += np.sum(mask)## fill keyframes_listindex = 0for mask in keyframes_mask_list:keyframes = np.zeros((149, height, width, 3), dtype=np.uint8)keyframes = [Image.fromarray(img.astype('uint8')) for img in keyframes]for j in range(len(mask)):if mask[j] == 1:frame = Image.fromarray(frames[index].astype('uint8'))frame, _ = crop_and_resize(frame, target_height=height, target_width=width)keyframes[j] = frame.copy()index += 1keyframes_list.append(keyframes)for i in range(len(keyframes_list) - 1):keyframes_list[i][-1] = keyframes_list[i + 1][0]keyframes_mask_list[i][-1] = 1return keyframes_list, keyframes_mask_listdef gen_video_single(pipe, prompt, input_config):negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"target_height = input_config['target_height']target_width = input_config['target_width']## refimageenable_refimage = input_config.get('enable_refimage', False)if enable_refimage:refimage_path = input_config['refimage_path']refimage = Image.open(refimage_path)refimage, (tl_x, tl_y, br_x, br_y) = crop_and_resize(refimage, target_width=target_width, target_height=target_height) # resizeelse:refimage = None## music featureenable_music_inject = input_config['enable_music_inject']music_feature = input_config['music_feature']## generated videoseed = input_config['seed']num_frames = input_config['num_frames']interp_mode = input_config['interp_mode']keyframes = input_config['keyframes']keyframes_mask = input_config['keyframes_mask']enable_dynamicfps = input_config.get('enable_dynamicfps', False)input_fps = input_config.get('input_fps', 30.0)enable_skip_layer = input_config.get('enable_skip_layer', False)sigma_shift = input_config.get('sigma_shift', 5)num_inference_steps = input_config.get('num_inference_steps', 48)cfg_scale = input_config.get('cfg_scale', 5)video = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, seed=seed, tiled=True, height=target_height, width=target_width,enable_music_inject=enable_music_inject,music_feature=music_feature,num_frames=num_frames,interp_mode=interp_mode,enable_refimage=enable_refimage,refimage=refimage,keyframes=keyframes,keyframes_mask=keyframes_mask,enable_dynamicfps=enable_dynamicfps,input_fps=input_fps,enable_skip_layer=enable_skip_layer,sigma_shift=sigma_shift,cfg_scale=cfg_scale)if dist.get_rank() == 0:save_video_path = input_config['save_video_path']save_video(video, save_video_path, fps=FPS, quality=5)## crop generated videovideo = mpy.VideoFileClip(save_video_path)croper = mpy.video.fx.Crop(x1=tl_x, y1=tl_y, x2=br_x, y2=br_y)video = croper.apply(video)## add musicmusic_path = input_config['music_path']video.audio = mpy.AudioFileClip(music_path)save_video_path = save_video_path[:-4] + "_music.mp4"video.write_videofile(save_video_path, codec='libx264', audio_codec='aac')dist.barrier(device_ids=[dist.get_rank()])def init_dit_model(args):world_size = int(os.environ.get("WORLD_SIZE", 1)) ## usp configulysses_degree = world_sizesequence_parallel_degree = world_sizering_degree=sequence_parallel_degree // ulysses_degreedata_parallel_degree = 1usp_config = {'data_parallel_degree': data_parallel_degree,'sequence_parallel_degree': sequence_parallel_degree,'ring_degree': ring_degree,'ulysses_degree': ulysses_degree}## load modelspipe = WanVideoPipeline.from_pretrained(torch_dtype=torch.bfloat16,device="cuda",model_configs=[ModelConfig(path="/home/models/Wan-Dancer-main/Wan-Dancer-14B/local_model.safetensors",model_id="/home/models/Wan-Dancer-main/Wan-Dancer-14B/", origin_file_pattern="local_model.safetensors", offload_device="cuda"),ModelConfig(path="/home/models/Wan-Dancer-main/Wan-Dancer-14B/models_t5_umt5-xxl-enc-bf16.pth",model_id="/home/models/Wan-Dancer-main/Wan-Dancer-14B/", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth", offload_device="cuda"),ModelConfig(path="/home/models/Wan-Dancer-main/Wan-Dancer-14B/Wan2.1_VAE.pth",model_id="/home/models/Wan-Dancer-main/Wan-Dancer-14B/", origin_file_pattern="Wan2.1_VAE.pth", offload_device="cuda"),ModelConfig(path="/home/models/Wan-Dancer-main/Wan-Dancer-14B/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth",model_id="/home/models/Wan-Dancer-main/Wan-Dancer-14B/", origin_file_pattern="models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth", offload_device="cuda"),],tokenizer_config=ModelConfig(path="/home/models/Wan-Dancer-main/Wan-Dancer-14B/google/umt5-xxl/",model_id="/home/models/Wan-Dancer-main/Wan-Dancer-14B/",origin_file_pattern="google/umt5-xxl/"),skip_download=True,redirect_common_files=False,use_usp=True,usp_config=usp_config,dit_model_type=1, # 1 for our trained modelenable_music_inject=True,enable_refimage=True,enable_global=True,enable_dynamicfps=True,enable_unimodel=True)pipe.enable_vram_management()return pipedef gen_video(pipe, music_path, music_feature_path, prompt, output_video_path, seed=0, max_pixels=1280*720, height=1280, width=720,num_frames=81, enable_refimage=False, refimage_path=None,keyframes=None, keyframes_mask=None, enable_dynamicfps=False, enable_skip_layer=False, sigma_shift=5, num_inference_steps=48):## input parametersinput_config = {}input_config['music_path'] = music_pathinput_config['enable_music_inject'] = Truemusic_feature = np.load(music_feature_path)music_feature = torch.from_numpy(music_feature).to(dtype=torch.bfloat16, device='cuda')input_config['music_feature'] = music_featureinput_config['max_pixels'] = max_pixelsinput_config['save_video_path'] = output_video_pathinput_config['seed'] = seedinput_config['target_height'] = heightinput_config['target_width'] = widthinput_config['num_frames'] = num_framesinput_config['interp_mode'] = 'bilinear'input_config['enable_refimage'] = enable_refimageinput_config['refimage_path'] = refimage_pathinput_config['keyframes'] = keyframesinput_config['keyframes_mask'] = keyframes_maskinput_config['input_fps'] = 30input_config['enable_dynamicfps'] = enable_dynamicfpsinput_config['enable_skip_layer'] = enable_skip_layerinput_config['sigma_shift'] = sigma_shiftinput_config['num_inference_steps'] = num_inference_steps# gen videogen_video_single(pipe, prompt, input_config)def process():args = parse_args()final_name = args.image_path.split('/')[-1].split('.')[0] + '_' + args.music_path.split('/')[-1].split('.')[0]time_name = args.timestampmusic_folder = 'outputs/tmp_results/' + final_name + '_' + str(time_name)os.makedirs(music_folder, exist_ok=True)## 0. process global videovideo_path = args.global_video_pathheight = args.heightwidth = args.widthaudio = AudioFileClip(args.music_path)# load the audio filetotal_duration = audio.duration# total duration of the audio filetotal_frames = int(total_duration * FPS)logger.info(f'total frames: {total_frames}')keyframes_list, keyframes_mask_list = process_global_video_firstlastframe(video_path, height, width, total_frames)# replace the first frame of the first segment with the input imageinput_image = Image.open(args.image_path)video_h, video_w = keyframes_list[0][0].size[1], keyframes_list[0][0].size[0]if min(video_h, video_w) < 512: input_image_resized, _ = crop_and_resize(input_image, target_height=height, target_width=width)keyframes_list[0][0] = input_image_resized# save keyframe videosfor i in range(len(keyframes_list)):save_keyframes_path = os.path.join(music_folder, f'keyframes_{str(i).zfill(2)}.mp4')save_video(keyframes_list[i], save_keyframes_path, fps=FPS, quality=5)## 1. init modelprompt_path = args.prompt_path# promptwith open(prompt_path, 'r') as f:prompt = f.read().strip()prompt += ', 帧率是30fps。'logger.info(f'---- prompt: {prompt} ----')pipe = init_dit_model(args)## 2. slice music to 5s segments；original_music_path = args.music_pathtarget_music_root_folder = music_folderif dist.get_rank() == 0:get_music_clip_149f(original_music_path, target_music_root_folder)dist.barrier(device_ids=[dist.get_rank()])## 3. encode musicif dist.get_rank() == 0:get_music_features(music_folder)dist.barrier(device_ids=[dist.get_rank()])## 4. generate videovideo_paths = []dirs = [f for f in sorted(os.listdir(music_folder)) if f.endswith('.wav')]for idx, name in tqdm(enumerate(dirs)):music_path = os.path.join(music_folder, name)music_feature_path = os.path.join(music_folder, name[:-4] + '_librosa_feature.npy')refimage_path = args.image_path seed = idx * 10 + args.seedoutput_video_path = os.path.join(music_folder, name[:-4] + "_seed" + str(seed) + ".mp4")gen_video(pipe, music_path, music_feature_path, prompt, output_video_path, seed=seed,height=args.height, width=args.width, num_frames=args.num_frames,enable_refimage=True, refimage_path=refimage_path,keyframes=keyframes_list[idx], keyframes_mask=keyframes_mask_list[idx],enable_dynamicfps=True,enable_skip_layer=True,sigma_shift=args.sigma_shift,num_inference_steps=args.num_inference_steps)video_paths.append(output_video_path[:-4] + "_music.mp4")## 5. combine with musicmusic_path = args.music_pathmusic = mpy.AudioFileClip(music_path)total_duration = music.durationoutput_video_folder = args.output_folderos.makedirs(output_video_folder, exist_ok=True)seed = args.seedif len(video_paths) > 0:output_video_path = os.path.join(output_video_folder, final_name + '_' + str(time_name) + "_seed" + str(seed) + '.mp4')clips = [mpy.VideoFileClip(vp) for vp in video_paths]final_clip = mpy.concatenate_videoclips(clips, method="compose")final_clip.audio = musicfinal_clip = final_clip[:total_duration-0.2] final_clip.write_videofile(output_video_path, codec='libx264', audio_codec='aac', fps=FPS)logger.info(f"Final video saved to {output_video_path}")else:logger.warning("No video files found to concatenate.")def main():process()if __name__ == '__main__':main()

注意关键代码修改（用于适应ROCM环境）：

import torchif torch.cuda.is_available() and torch.version.cuda is None:# 1. 假装有一个 CUDA 版本，使 xfuser 的版本检查通过torch.version.cuda = "12.4"# 2. 修补 get_arch_list，返回 CUDA 格式的架构列表，避免 _extract_arch_version 索引错误import torch.cudaoriginal_get_arch_list = torch.cuda.get_arch_listdef fake_get_arch_list():# 返回一个 CUDA 风格的架构（例如 sm_70），这样 split("_")[1] 就能正常取出数字return ["sm_120"]torch.cuda.get_arch_list = fake_get_arch_list

3.5模型下载与准备

Wan-Dancer-14B（舞蹈视频生成大模型）：

下载链接：https://modelscope.cn/models/Wan-AI/Wan-Dancer-14B/files

四、运行测试

1、生成全局关键帧视频:

export PYTORCH_ROCM_ARCH="gfx928"

./gen_video_global-chinese.sh

2、生成最终高分辨率视频：

export PYTORCH_ROCM_ARCH="gfx928"

./gen_video_local.sh

四卡K100_AI运行生成30秒视频，整个过程耗时约8小时。（据网友反馈，四卡H20运行生成30秒视频，整个过程耗时6小时）

程序用到的图片（官方示例中提供的图片）：

五、舞蹈视频生成效果：

Wan-Dancer-14B生成视频效果