app.py

import torch
import psutil
import argparse
import os
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.utils import load_image
from transformers import AutoTokenizer, Wav2Vec2Model, Wav2Vec2Processor
from omegaconf import OmegaConf
from wan.models.cache_utils import get_teacache_coefficients
from wan.models.wan_fantasy_transformer3d_1B import WanTransformer3DFantasyModel
from wan.models.wan_text_encoder import WanT5EncoderModel
from wan.models.wan_vae import AutoencoderKLWan
from wan.models.wan_image_encoder import CLIPModel
from wan.pipeline.wan_inference_long_pipeline import WanI2VTalkingInferenceLongPipeline
from wan.utils.fp8_optimization import replace_parameters_by_name, convert_weight_dtype_wrapper, convert_model_weight_to_float8
from wan.utils.utils import get_image_to_video_latent, save_videos_grid
import numpy as np
import librosa
import datetime
import random
import math
import subprocess
from huggingface_hub import snapshot_download
import requests
import shutil

# --- 全域設定 ---
if torch.cuda.is_available():
    device = "cuda"
    if torch.cuda.get_device_capability()[0] >= 8:
        dtype = torch.bfloat16
    else:
        dtype = torch.float16
else:
    device = "cpu"
    dtype = torch.float32

def filter_kwargs(cls, kwargs):
    """過濾掉不屬於類別建構函式的關鍵字參數"""
    import inspect
    sig = inspect.signature(cls.__init__)
    valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
    filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
    return filtered_kwargs

def resolve_path(user_path, repo_root):
    """
    以正確的優先級解析檔案路徑。
    1. 優先檢查本地路徑（絕對或相對）。
    2. 如果找不到，則嘗試從 HF 快取目錄中尋找。
    """
    # 優先檢查本地路徑
    if os.path.exists(user_path):
        print(f"找到本地檔案： {os.path.abspath(user_path)}")
        return os.path.abspath(user_path)
    
    # 其次，嘗試從 HF 快取目錄中尋找
    potential_repo_path = os.path.join(repo_root, user_path)
    if os.path.exists(potential_repo_path):
        print(f"在 Hugging Face 快取目錄中找到檔案： {potential_repo_path}")
        return potential_repo_path
        
    return None

def setup_models(repo_root, model_version):
    """載入所有必要的模型和設定"""
    pretrained_model_name_or_path = os.path.join(repo_root, "Wan2.1-Fun-V1.1-1.3B-InP")
    pretrained_wav2vec_path = os.path.join(repo_root, "wav2vec2-base-960h")

    # <<< 核心修正：對設定檔使用與輸入檔案相同的路徑解析邏輯 >>>
    config_relative_path = "deepspeed_config/wan2.1/wan_civitai.yaml"
    config_path = resolve_path(config_relative_path, repo_root)
    
    if not config_path:
        raise FileNotFoundError(f"設定檔 '{config_relative_path}' 在當前目錄或 HF 快取中都找不到。請確保該檔案存在。")
    # <<< 修正結束 >>>

    print(f"正在從 {config_path} 載入設定...")
    config = OmegaConf.load(config_path)
    sampler_name = "Flow"
    
    print("正在載入 Tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(os.path.join(pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')))
    
    print("正在載入 Text Encoder...")
    text_encoder = WanT5EncoderModel.from_pretrained(
        os.path.join(pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
        additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
        low_cpu_mem_usage=True,
        torch_dtype=dtype,
    ).eval()
    
    print("正在載入 VAE...")
    vae = AutoencoderKLWan.from_pretrained(
        os.path.join(pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')),
        additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
    )
    
    print("正在載入 Wav2Vec...")
    wav2vec_processor = Wav2Vec2Processor.from_pretrained(pretrained_wav2vec_path)
    wav2vec = Wav2Vec2Model.from_pretrained(pretrained_wav2vec_path).to("cpu")
    
    print("正在載入 CLIP Image Encoder...")
    clip_image_encoder = CLIPModel.from_pretrained(os.path.join(pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder'))).eval()
    
    print("正在載入 Transformer 3D 基礎模型...")
    transformer3d = WanTransformer3DFantasyModel.from_pretrained(
        os.path.join(pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
        transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
        low_cpu_mem_usage=False,
        torch_dtype=dtype,
    )

    if model_version == "square":
        transformer_path = os.path.join(repo_root, "StableAvatar-1.3B", "transformer3d-square.pt")
    else: # rec_vec
        transformer_path = os.path.join(repo_root, "StableAvatar-1.3B", "transformer3d-rec-vec.pt")

    if os.path.exists(transformer_path):
        print(f"正在從 {transformer_path} 載入 StableAvatar 權重...")
        state_dict = torch.load(transformer_path, map_location="cpu")
        state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
        m, u = transformer3d.load_state_dict(state_dict, strict=False)
        print(f"StableAvatar 權重載入成功。 Missing keys: {len(m)}; Unexpected keys: {len(u)}")
    else:
        raise FileNotFoundError(f"找不到 StableAvatar 權重檔案：{transformer_path}。請確保模型已完整下載。")

    scheduler_class = { "Flow": FlowMatchEulerDiscreteScheduler }[sampler_name]
    scheduler = scheduler_class(**filter_kwargs(scheduler_class, OmegaConf.to_container(config['scheduler_kwargs'])))
    
    print("正在建立 Pipeline...")
    pipeline = WanI2VTalkingInferenceLongPipeline(
        tokenizer=tokenizer, text_encoder=text_encoder, vae=vae,
        transformer=transformer3d, clip_image_encoder=clip_image_encoder,
        scheduler=scheduler, wav2vec_processor=wav2vec_processor, wav2vec=wav2vec,
    )
    
    return pipeline, transformer3d, vae

def run_inference(
    pipeline, transformer3d, vae, image_path, audio_path, prompt,
    negative_prompt, seed, output_filename, gpu_memory_mode="model_cpu_offload",
    width=512, height=512, num_inference_steps=50, fps=25, **kwargs
):
    """執行推理以生成影片。"""
    if seed < 0:
        seed = random.randint(0, np.iinfo(np.int32).max)
    print(f"使用的種子： {seed}")

    if gpu_memory_mode == "sequential_cpu_offload":
        pipeline.enable_sequential_cpu_offload(device=device)
    elif gpu_memory_mode == "model_cpu_offload":
        pipeline.enable_model_cpu_offload(device=device)
    else:
        pipeline.to(device=device)

    with torch.no_grad():
        print("正在準備輸入資料...")
        video_length = 81
        input_video, input_video_mask, clip_image = get_image_to_video_latent(image_path, None, video_length=video_length, sample_size=[height, width])
        
        sr = 16000
        vocal_input, _ = librosa.load(audio_path, sr=sr)
        
        print("Pipeline 執行中... 這可能需要一些時間。")
        sample = pipeline(
            prompt, num_frames=video_length, negative_prompt=negative_prompt,
            width=width, height=height, guidance_scale=6.0,
            generator=torch.Generator().manual_seed(seed), num_inference_steps=num_inference_steps,
            video=input_video, mask_video=input_video_mask, clip_image=clip_image,
            text_guide_scale=3.0, audio_guide_scale=5.0, vocal_input_values=vocal_input,
            motion_frame=25, fps=fps, sr=sr, cond_file_path=image_path,
            overlap_window_length=10, seed=seed, overlapping_weight_scheme="uniform",
        ).videos
        
        print("正在儲存影片...")
        os.makedirs("outputs", exist_ok=True)
        video_path = os.path.join("outputs", f"{output_filename}.mp4")
        save_videos_grid(sample, video_path, fps=fps)
        
        output_video_with_audio = os.path.join("outputs", f"{output_filename}_audio.mp4")
        
        print("正在將音訊合併到影片中...")
        subprocess.run([
            "ffmpeg", "-y", "-loglevel", "quiet", "-i", video_path, "-i", audio_path,
            "-c:v", "copy", "-c:a", "aac", "-strict", "experimental",
            output_video_with_audio
        ], check=True)
        
        os.remove(video_path)
        
    print(f"✅ 生成完成！影片已儲存至： {output_video_with_audio}")
    return output_video_with_audio, seed

def main():
    parser = argparse.ArgumentParser(description="StableAvatar 命令列推理工具")
    parser.add_argument('--prompt', type=str, default="a beautiful woman is talking, masterpiece, best quality", help='正面提示詞')
    parser.add_argument('--input_image', type=str, default="example_case/case-6/reference.png", help='輸入圖片的路徑')
    parser.add_argument('--input_audio', type=str, default="example_case/case-6/audio.wav", help='輸入音訊的路徑')
    parser.add_argument('--seed', type=int, default=42, help='隨機種子，-1 表示隨機')
    parser.add_argument('--negative_prompt', type=str, default="vivid color, static, blur details, text, style, painting, picture, still, gray, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, malformed, deformed, bad anatomy, fused fingers, still image, messy background, many people in the background, walking backwards", help='負面提示詞')
    parser.add_argument('--width', type=int, default=512, help='影片寬度')
    parser.add_argument('--height', type=int, default=512, help='影片高度')
    parser.add_argument('--num_inference_steps', type=int, default=50, help='推理步數')
    parser.add_argument('--fps', type=int, default=25, help='影片幀率')
    parser.add_argument('--gpu_memory_mode', type=str, default="model_cpu_offload", choices=["Normal", "model_cpu_offload"], help='GPU 記憶體優化模式')
    parser.add_argument('--model_version', type=str, default="square", choices=["square", "rec_vec"], help='StableAvatar 模型版本')
    args = parser.parse_args()

    print("--- 步驟 1: 正在檢查並下載模型與設定檔 ---")
    # 確保所有需要的檔案都被下載，以作為本地找不到檔案時的後備
    repo_root = snapshot_download(
        repo_id="FrancisRing/StableAvatar",
        allow_patterns=[
            "StableAvatar-1.3B/*",
            "Wan2.1-Fun-V1.1-1.3B-InP/*",
            "wav2vec2-base-960h/*",
            "deepspeed_config/**",
            "example_case/**" # 也下載範例，以防使用者直接執行預設參數
        ],
    )
    print("模型檔案已準備就緒。")

    print("\n--- 步驟 2: 正在解析輸入檔案路徑 ---")
    final_image_path = resolve_path(args.input_image, repo_root)
    if not final_image_path:
        print(f"錯誤：無法找到圖片檔案 {args.input_image}")
        return

    final_audio_path = resolve_path(args.input_audio, repo_root)
    if not final_audio_path:
        print(f"錯誤：無法找到音訊檔案 {args.input_audio}")
        return

    print("\n--- 步驟 3: 正在載入模型 ---")
    pipeline, transformer3d, vae = setup_models(repo_root, args.model_version)
    print("模型載入完成。")
    
    print("\n--- 步驟 4: 開始執行推理 ---")
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    run_inference(
        pipeline=pipeline, transformer3d=transformer3d, vae=vae,
        image_path=final_image_path, audio_path=final_audio_path,
        prompt=args.prompt, negative_prompt=args.negative_prompt,
        seed=args.seed, output_filename=f"output_{timestamp}",
        gpu_memory_mode=args.gpu_memory_mode, width=args.width,
        height=args.height, num_inference_steps=args.num_inference_steps,
        fps=args.fps
    )

if __name__ == "__main__":
    main()