local_inference.py

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