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 moviepy.editor import VideoFileClip
from huggingface_hub import snapshot_download
import shutil
import requests
import uuid

# Device and dtype setup
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 load_transformer_model(model_version, repo_root):
    transformer_path = os.path.join(repo_root, "StableAvatar-1.3B", f"transformer3d-{model_version}.pt")
    print(f"Loading model: {transformer_path}")
    if os.path.exists(transformer_path):
        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"Model loaded successfully: {transformer_path}")
        print(f"Missing keys: {len(m)}; Unexpected keys: {len(u)}")
        return transformer3d
    else:
        print(f"Error: Model file does not exist: {transformer_path}")
        return None

def download_file(url, local_path):
    """Download file from URL to local path"""
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()
        with open(local_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        return local_path
    except Exception as e:
        print(f"Error downloading file from {url}: {e}")
        return None

def prepare_input_file(input_path, file_type="image"):
    """Handle local or remote file inputs"""
    if input_path.startswith("http://") or input_path.startswith("https://"):
        ext = ".png" if file_type == "image" else ".wav"
        local_path = os.path.join("temp", f"{uuid.uuid4()}{ext}")
        os.makedirs("temp", exist_ok=True)
        return download_file(input_path, local_path)
    elif os.path.exists(input_path):
        return input_path
    else:
        print(f"Error: {file_type.capitalize()} file {input_path} does not exist")
        return None

# Initialize model paths
REPO_ID = "FrancisRing/StableAvatar"
repo_root = snapshot_download(
    repo_id=REPO_ID,
    allow_patterns=[
        "StableAvatar-1.3B/*",
        "Wan2.1-Fun-V1.1-1.3B-InP/*",
        "wav2vec2-base-960h/*",
        "assets/**",
        "Kim_Vocal_2.onnx",
    ],
)
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")
audio_separator_model_file = os.path.join(repo_root, "Kim_Vocal_2.onnx")

# Load configuration and models
config = OmegaConf.load("deepspeed_config/wan2.1/wan_civitai.yaml")
sampler_name = "Flow"
clip_sample_n_frames = 81

tokenizer = AutoTokenizer.from_pretrained(
    os.path.join(pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer'))
)
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()
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']),
)
wav2vec_processor = Wav2Vec2Processor.from_pretrained(pretrained_wav2vec_path)
wav2vec = Wav2Vec2Model.from_pretrained(pretrained_wav2vec_path).to("cpu")
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()
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,
)

# Load default transformer model
load_transformer_model("square", repo_root)

# Initialize scheduler and pipeline
scheduler_dict = {"Flow": FlowMatchEulerDiscreteScheduler}
Choosen_Scheduler = scheduler_dict[sampler_name]
scheduler = Choosen_Scheduler(
    **filter_kwargs(Choosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
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,
)

def generate(
    GPU_memory_mode="model_cpu_offload",
    teacache_threshold=0,
    num_skip_start_steps=5,
    image_path=None,
    audio_path=None,
    prompt="",
    negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
    width=512,
    height=512,
    guidance_scale=6.0,
    num_inference_steps=50,
    text_guide_scale=3.0,
    audio_guide_scale=5.0,
    motion_frame=25,
    fps=25,
    overlap_window_length=10,
    seed_param=42,
    overlapping_weight_scheme="uniform"
):
    global pipeline, transformer3d
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    if seed_param < 0:
        seed = random.randint(0, np.iinfo(np.int32).max)
    else:
        seed = seed_param

    # Handle input files
    image_path = prepare_input_file(image_path, "image")
    audio_path = prepare_input_file(audio_path, "audio")
    if not image_path or not audio_path:
        return None, None, "Error: Invalid input file paths"

    # Configure pipeline based on GPU memory mode
    if GPU_memory_mode == "sequential_cpu_offload":
        replace_parameters_by_name(transformer3d, ["modulation"], device=device)
        transformer3d.freqs = transformer3d.freqs.to(device=device)
        pipeline.enable_sequential_cpu_offload(device=device)
    elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
        convert_model_weight_to_float8(transformer3d, exclude_module_name=["modulation"])
        convert_weight_dtype_wrapper(transformer3d, dtype)
        pipeline.enable_model_cpu_offload(device=device)
    elif GPU_memory_mode == "model_cpu_offload":
        pipeline.enable_model_cpu_offload(device=device)
    else:
        pipeline.to(device=device)

    # Enable TeaCache if specified
    if teacache_threshold > 0:
        coefficients = get_teacache_coefficients(pretrained_model_name_or_path)
        pipeline.transformer.enable_teacache(
            coefficients,
            num_inference_steps,
            teacache_threshold,
            num_skip_start_steps=num_skip_start_steps,
        )

    # Perform inference
    with torch.no_grad():
        video_length = int((clip_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if clip_sample_n_frames != 1 else 1
        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, sample_rate = librosa.load(audio_path, sr=sr)
        sample = pipeline(
            prompt,
            num_frames=video_length,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            guidance_scale=guidance_scale,
            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=text_guide_scale,
            audio_guide_scale=audio_guide_scale,
            vocal_input_values=vocal_input,
            motion_frame=motion_frame,
            fps=fps,
            sr=sr,
            cond_file_path=image_path,
            overlap_window_length=overlap_window_length,
            seed=seed,
            overlapping_weight_scheme=overlapping_weight_scheme,
        ).videos
        os.makedirs("outputs", exist_ok=True)
        video_path = os.path.join("outputs", f"{timestamp}.mp4")
        save_videos_grid(sample, video_path, fps=fps)
        output_video_with_audio = os.path.join("outputs", f"{timestamp}_audio.mp4")
        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)

    return output_video_with_audio, seed, f"Generated outputs/{timestamp}.mp4"

def main():
    parser = argparse.ArgumentParser(description="StableAvatar Inference Script")
    parser.add_argument("--prompt", type=str, default="", help="Text prompt for generation")
    parser.add_argument("--seed", type=int, default=42, help="Random seed, -1 for random")
    parser.add_argument("--input_image", type=str, required=True, help="Path or URL to input image (e.g., ./image.png or https://example.com/image.png)")
    parser.add_argument("--input_audio", type=str, required=True, help="Path or URL to input audio (e.g., ./audio.wav or https://example.com/audio.wav)")
    parser.add_argument("--GPU_memory_mode", type=str, default="model_cpu_offload", choices=["Normal", "model_cpu_offload", "model_cpu_offload_and_qfloat8", "sequential_cpu_offload"], help="GPU memory mode")
    parser.add_argument("--teacache_threshold", type=float, default=0, help="TeaCache threshold, 0 to disable")
    parser.add_argument("--num_skip_start_steps", type=int, default=5, help="Number of start steps to skip")
    parser.add_argument("--negative_prompt", type=str, default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", help="Negative prompt")
    parser.add_argument("--width", type=int, default=512, help="Output video width")
    parser.add_argument("--height", type=int, default=512, help="Output video height")
    parser.add_argument("--guidance_scale", type=float, default=6.0, help="Guidance scale")
    parser.add_argument("--num_inference_steps", type=int, default=50, help="Number of inference steps")
    parser.add_argument("--text_guide_scale", type=float, default=3.0, help="Text guidance scale")
    parser.add_argument("--audio_guide_scale", type=float, default=5.0, help="Audio guidance scale")
    parser.add_argument("--motion_frame", type=int, default=25, help="Motion frame")
    parser.add_argument("--fps", type=int, default=25, help="Frames per second")
    parser.add_argument("--overlap_window_length", type=int, default=10, help="Overlap window length")
    parser.add_argument("--overlapping_weight_scheme", type=str, default="uniform", choices=["uniform", "log"], help="Overlapping weight scheme")

    args = parser.parse_args()

    video_path, seed, message = generate(
        GPU_memory_mode=args.GPU_memory_mode,
        teacache_threshold=args.teacache_threshold,
        num_skip_start_steps=args.num_skip_start_steps,
        image_path=args.input_image,
        audio_path=args.input_audio,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        width=args.width,
        height=args.height,
        guidance_scale=args.guidance_scale,
        num_inference_steps=args.num_inference_steps,
        text_guide_scale=args.text_guide_scale,
        audio_guide_scale=args.audio_guide_scale,
        motion_frame=args.motion_frame,
        fps=args.fps,
        overlap_window_length=args.overlap_window_length,
        seed_param=args.seed,
        overlapping_weight_scheme=args.overlapping_weight_scheme
    )

    if video_path:
        print(f"{message}\nSeed: {seed}")
    else:
        print("Generation failed.")

if __name__ == "__main__":
    main()