import os import gc import sys import time import cv2 import librosa import numpy as np import torch import torchvision import PIL from PIL import Image, ImageFile import moviepy as mpy import soundfile as sf from tqdm import tqdm from moviepy import AudioFileClip from pathlib import Path # Import the xfuser mock first to register it in sys.modules import xfuser from diffsynth import save_video from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig from diffsynth.models.model_manager import ModelManager FPS = 30 ImageFile.LOAD_TRUNCATED_IMAGES = True def get_music_base_feature(music_path, output_path, fps=30): hop_length = 512 sr = fps * hop_length data, 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).T chroma = 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.0 start_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.0 audio_feature = np.concatenate( [envelope[:, None], mfcc, chroma, peak_onehot[:, None], beat_onehot[:, None]], axis=-1, ) np.save(output_path, audio_feature) return audio_feature def 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) / FPS idx = 0 t = 0 while t + 0.2 < total_duration: start_time = t end_time = t + duration if end_time >= total_duration: end_time = total_duration sliced_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 += duration idx += 1 def get_music_features(music_folder): dirs = [f for f in sorted(os.listdir(music_folder)) if f.endswith('.wav')] for idx, name in 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.size scale = 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 ) target_image = np.ones((target_height, target_width, 3), dtype=np.uint8) * 127 tl_x = (target_width - resized_width) // 2 tl_y = (target_height - resized_height) // 2 br_x = tl_x + resized_width br_y = tl_y + resized_height target_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: break frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frames.append(frame) cap.release() if not frames: frames = [np.zeros((height, width, 3), dtype=np.uint8)] N = len(frames) seg_num = int(np.ceil(total_frames / 149)) frame_interval_num = float(total_frames) / N keyframes_mask_list = [] for i in range(seg_num): mask = np.zeros(149, dtype=np.int32) if i != seg_num - 1: cnt = 0 while (cnt * frame_interval_num < 149 - frame_interval_num): index = int(np.ceil(frame_interval_num * cnt)) mask[index] = 1 cnt += 1 else: end_index = total_frames - 149 * i - 1 mask[end_index] = 1 cnt = 0 while (cnt * frame_interval_num < end_index - frame_interval_num): index = int(np.ceil(frame_interval_num * cnt)) mask[index] = 1 cnt += 1 keyframes_mask_list.append(mask) keyframes_list = [] index = 0 for 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_idx = min(index, N - 1) frame = Image.fromarray(frames[frame_idx].astype('uint8')) frame, _ = crop_and_resize(frame, target_height=height, target_width=width) keyframes[j] = frame.copy() index += 1 keyframes_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] = 1 return keyframes_list, keyframes_mask_list _GLOBAL_PIPE = None def init_pipeline(model_file_name, local_model_path="./models"): pipe = WanVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", model_configs=[ ModelConfig(model_id="Wan-AI/Wan-Dancer-14B", origin_file_pattern=model_file_name, offload_device="cpu"), ModelConfig(model_id="Wan-AI/Wan-Dancer-14B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth", offload_device="cpu"), ModelConfig(model_id="Wan-AI/Wan-Dancer-14B", origin_file_pattern="Wan2.1_VAE.pth", offload_device="cpu"), ModelConfig(model_id="Wan-AI/Wan-Dancer-14B", origin_file_pattern="models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth", offload_device="cpu"), ], tokenizer_config=ModelConfig(model_id="Wan-AI/Wan-Dancer-14B", origin_file_pattern="google/umt5-xxl/"), skip_download=True, redirect_common_files=False, use_usp=False, dit_model_type=1, enable_music_inject=True, enable_refimage=True, enable_global=True, enable_dynamicfps=True, enable_unimodel=True ) pipe.enable_vram_management() return pipe def get_or_init_pipeline(local_model_path="./models"): global _GLOBAL_PIPE if _GLOBAL_PIPE is None: print("Initializing base WanVideoPipeline with global_model...") pipe = init_pipeline("global_model.safetensors", local_model_path=local_model_path) pipe.global_dit = pipe.dit local_model_file = os.path.join(local_model_path, "Wan-AI/Wan-Dancer-14B", "local_model.safetensors") if os.path.exists(local_model_file): print("Loading local_model.safetensors DIT backbone...") model_manager = ModelManager() model_manager.load_model_dit( local_model_file, device="cpu", torch_dtype=torch.bfloat16, enable_music_inject=True, music_inject_layers=[0, 4, 8, 12, 16, 20, 24, 27], dit_model_type=1, enable_videojam=False, enable_double=False, use_usp=False, enable_refimage=True, enable_refface=False, enable_global=True, enable_dynamicfps=True, enable_unimodel=True ) pipe.local_dit = model_manager.fetch_model("wan_video_dit") pipe.dit = pipe.local_dit pipe.enable_vram_management() else: pipe.local_dit = pipe.global_dit pipe.dit = pipe.global_dit _GLOBAL_PIPE = pipe return _GLOBAL_PIPE def gen_global_video(pipe, image_path, music_feature_path, prompt, output_video_path, seed=0, height=1280, width=720, num_inference_steps=24, cfg_scale=5): negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" img = Image.open(image_path) img, (tl_x, tl_y, br_x, br_y) = crop_and_resize(img, target_width=width, target_height=height) music_feature = np.load(music_feature_path) music_feature = torch.from_numpy(music_feature).to(dtype=torch.bfloat16, device='cuda') input_fps = 30.0 / int(music_feature.shape[0] / 149.0 + 0.5) input_fps = float("{:.4f}".format(input_fps)) prompt += f"帧率是{input_fps}" mask = np.zeros(149, dtype=np.int32) mask[0] = 1 keyframes = np.zeros((149, height, width, 3), dtype=np.uint8) keyframes[mask == 1] = np.array(img, dtype=np.uint8) keyframes = [Image.fromarray(i.astype("uint8")) for i in keyframes] keyframes_mask = torch.tensor(mask).to(torch.int32) video = pipe( prompt=prompt, negative_prompt=negative_prompt, input_image=None, num_inference_steps=num_inference_steps, seed=seed, tiled=True, height=height, width=width, enable_music_inject=True, music_feature=music_feature, num_frames=149, interp_mode="bilinear", enable_refimage=True, refimage=img, enable_global=True, keyframes=keyframes, keyframes_mask=keyframes_mask, enable_dynamicfps=True, input_fps=input_fps, enable_vae_decode_framewise=True, enable_skip_layer=True, enable_unimodel=True, sigma_shift=5, cfg_scale=cfg_scale, ) tmp_video_path = output_video_path[:-4] + "_tmp.mp4" save_video(video, tmp_video_path, fps=8, quality=5) # Crop video clip = mpy.VideoFileClip(tmp_video_path) croper = mpy.video.fx.Crop(x1=tl_x, y1=tl_y, x2=br_x, y2=br_y) clip = croper.apply(clip) clip.write_videofile(output_video_path, codec="libx264", audio_codec="aac") try: clip.close() os.remove(tmp_video_path) except Exception: pass def gen_local_video_segment(pipe, music_path, music_feature_path, prompt, output_video_path, seed=0, height=1280, width=720, keyframes=None, keyframes_mask=None, num_inference_steps=24, cfg_scale=5, refimage_path=None): negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" refimage = Image.open(refimage_path) refimage, (tl_x, tl_y, br_x, br_y) = crop_and_resize(refimage, target_width=width, target_height=height) music_feature = np.load(music_feature_path) music_feature = torch.from_numpy(music_feature).to(dtype=torch.bfloat16, device='cuda') video = pipe( prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, seed=seed, tiled=True, height=height, width=width, enable_music_inject=True, music_feature=music_feature, num_frames=149, interp_mode="bilinear", enable_refimage=True, refimage=refimage, keyframes=keyframes, keyframes_mask=keyframes_mask, enable_dynamicfps=True, input_fps=30, enable_skip_layer=True, sigma_shift=5, cfg_scale=cfg_scale, ) save_video(video, output_video_path, fps=FPS, quality=5) # Crop video and add music clip = mpy.VideoFileClip(output_video_path) croper = mpy.video.fx.Crop(x1=tl_x, y1=tl_y, x2=br_x, y2=br_y) clip = croper.apply(clip) clip.audio = mpy.AudioFileClip(music_path) final_output = output_video_path[:-4] + "_music.mp4" clip.write_videofile(final_output, codec='libx264', audio_codec='aac') try: clip.close() except Exception: pass return final_output def generate_dance_video(image_path, music_path, genre, output_folder, local_model_path="./models", seed=0, height=1280, width=720, steps=24, cfg=5, progress=None): # Set prompts prompts_global = { "chinese_classical": "一个人正在跳舞，舞蹈种类是古典舞。", "k_pop": "一个人正在跳舞，舞蹈种类是韩舞。", "street": "一个人正在跳舞，舞蹈种类是街舞。", "tap": "一个人正在跳舞，舞蹈种类是踢踏舞。", "latin": "一个人正在跳舞，舞蹈种类是拉丁舞。", } prompts_local = { "chinese_classical": "一个人正在跳舞，舞蹈种类是古典舞,图像清晰程度高,人物动作平均幅度中等,人物动作最大幅度中等。", "k_pop": "一个人正在跳舞，舞蹈种类是韩舞,图像清晰程度高,人物动作平均幅度中等,人物动作最大幅度中等。", "street": "一个人正在跳舞，舞蹈种类是街舞,图像清晰程度高,人物动作平均幅度中等,人物动作最大幅度中等。", "tap": "一个人正在跳舞，舞蹈种类是踢踏舞,图像清晰程度高,人物动作平均幅度高,人物动作最大幅度高。", "latin": "一个人正在跳舞，舞蹈种类是拉丁舞,图像清晰程度高,人物动作平均幅度高,人物动作最大幅度中等。", } genre = genre.strip().lower() p_global = prompts_global.get(genre, prompts_global["chinese_classical"]) p_local = prompts_local.get(genre, prompts_local["chinese_classical"]) timestamp = str(int(time.time())) temp_dir = os.path.join(output_folder, f"tmp_{timestamp}") os.makedirs(temp_dir, exist_ok=True) # 1. Encode Music Features if progress is not None: progress(0, desc="Preprocessing audio...") else: print("Preprocessing audio...") global_music_npy = os.path.join(temp_dir, "global_music_feature.npy") get_music_base_feature(music_path, global_music_npy, fps=30) # 2. Initialize and Run Global Model if progress is not None: progress(0.1, desc="Preparing global model pipeline...") else: print("Preparing global model pipeline...") pipe = get_or_init_pipeline(local_model_path=local_model_path) pipe.dit = pipe.global_dit if progress is not None: progress(0.2, desc="Generating global video sketch...") else: print("Generating global video sketch...") global_video_mp4 = os.path.join(temp_dir, "global_sketch.mp4") gen_global_video( pipe=pipe, image_path=image_path, music_feature_path=global_music_npy, prompt=p_global, output_video_path=global_video_mp4, seed=seed, height=height, width=width, num_inference_steps=steps, cfg_scale=cfg ) # 3. Slice Music & Slice Global Video into segments if progress is not None: progress(0.4, desc="Slicing global sketch and music...") else: print("Slicing global sketch and music...") audio = AudioFileClip(music_path) total_duration = audio.duration total_frames = int(total_duration * FPS) keyframes_list, keyframes_mask_list = process_global_video_firstlastframe( global_video_mp4, height, width, total_frames ) # Replace first frame with reference image input_image = Image.open(image_path) input_image_resized, _ = crop_and_resize(input_image, target_height=height, target_width=width) keyframes_list[0][0] = input_image_resized # Slice music chunks get_music_clip_149f(music_path, temp_dir) get_music_features(temp_dir) # 4. Prepare Local Model Pipeline if progress is not None: progress(0.5, desc="Switching to local model pipeline...") else: print("Switching to local model pipeline...") pipe.dit = pipe.local_dit # 5. Run Local Refinement on each segment dirs = [f for f in sorted(os.listdir(temp_dir)) if f.endswith('.wav')] while len(keyframes_list) < len(dirs): keyframes_list.append(keyframes_list[-1] if keyframes_list else [[Image.new("RGB", (width, height))] * 149]) while len(keyframes_mask_list) < len(dirs): keyframes_mask_list.append(keyframes_mask_list[-1] if keyframes_mask_list else [np.zeros(149, dtype=np.int32)]) video_paths = [] for idx, name in enumerate(dirs): seg_music_path = os.path.join(temp_dir, name) seg_npy_path = os.path.join(temp_dir, name[:-4] + '_librosa_feature.npy') seg_seed = idx * 10 + seed seg_output_path = os.path.join(temp_dir, name[:-4] + f"_seg_{idx}.mp4") if progress is not None: progress(0.5 + 0.4 * (idx / len(dirs)), desc=f"Refining dance segment {idx+1}/{len(dirs)}...") else: print(f"Refining dance segment {idx+1}/{len(dirs)}...") final_seg = gen_local_video_segment( pipe=pipe, music_path=seg_music_path, music_feature_path=seg_npy_path, prompt=p_local + ", 帧率是30fps。", output_video_path=seg_output_path, seed=seg_seed, height=height, width=width, keyframes=keyframes_list[idx], keyframes_mask=keyframes_mask_list[idx], num_inference_steps=steps, cfg_scale=cfg, refimage_path=image_path ) video_paths.append(final_seg) pipe.load_models_to_device([]) torch.cuda.empty_cache() # 6. Combine segments and audio if progress is not None: progress(0.9, desc="Merging segments into final video...") else: print("Merging segments into final video...") output_video_path = os.path.join(output_folder, f"wan_dancer_{timestamp}.mp4") clips = [mpy.VideoFileClip(vp) for vp in video_paths] final_clip = mpy.concatenate_videoclips(clips, method="compose") final_clip.audio = mpy.AudioFileClip(music_path) # Trim to match music duration exactly (minus small margin) final_clip = final_clip[:total_duration-0.2] final_clip.write_videofile(output_video_path, codec='libx264', audio_codec='aac', fps=FPS) # Close clips for c in clips: c.close() final_clip.close() # Clean up temp folder try: import shutil shutil.rmtree(temp_dir) except Exception: pass return output_video_path from mutagen import File as MutagenFile from pydub import AudioSegment def get_audio_duration_seconds(file_path: str) -> float: if not file_path or not os.path.isfile(file_path): return 0.0 try: audio_seg = AudioSegment.from_file(file_path) return float(audio_seg.duration_seconds) except Exception as e: print(f"Error reading audio duration with pydub: {e}") try: audio = MutagenFile(file_path) if audio is not None and audio.info is not None: return float(audio.info.length) except Exception as e2: print(f"Error reading audio duration with mutagen: {e2}") return 0.0 def truncate_audio(file_path: str, max_seconds: float = 30, start_second: float = 0.0) -> str: if not file_path or not os.path.isfile(file_path): return file_path duration = get_audio_duration_seconds(file_path) if duration <= 0: return file_path try: ext = Path(file_path).suffix.lower() fmt = "mp3" if ext == ".mp3" else "wav" audio_seg = AudioSegment.from_file(file_path) start_ms = int(start_second * 1000) end_ms = start_ms + int(max_seconds * 1000) # Clamp start_ms to length of audio if start_ms >= len(audio_seg): start_ms = 0 end_ms = int(max_seconds * 1000) truncated = audio_seg[start_ms:end_ms] import tempfile cache_dir = os.environ.get("TASK_CACHE_DIR", "./task_cache") os.makedirs(cache_dir, exist_ok=True) tmp = tempfile.NamedTemporaryFile( delete=False, suffix=ext or ".mp3", dir=cache_dir ) tmp.close() truncated.export(tmp.name, format=fmt) return tmp.name except Exception as e: print(f"Error truncating audio: {e}") return file_path