wan / appold.py
68ed8100's picture
1
afbf058
Raw
History Blame Contribute Delete
9.33 kB
import os
import gc
import torch
import gradio as gr
from diffusers import WanAnimatePipeline
from diffusers.utils import load_image, load_video, export_to_video
# ============================================================
# CONFIG
# ============================================================
MODEL_ID = "Wan-AI/Wan2.2-Animate-14B-Diffusers"
OUTPUT_DIR = "/tmp/outputs"
os.makedirs(OUTPUT_DIR, exist_ok=True)
# ============================================================
# PRESETS
# ============================================================
PRESETS = {
"Realistic": {
"prompt": (
"A highly realistic video of the reference character performing "
"the movements from the input motion video. Preserve identity, "
"facial appearance, hairstyle and clothing. Natural body motion, "
"realistic lighting and physics."
),
"steps": 20,
"guidance": 1.0,
},
"Cinematic": {
"prompt": (
"A cinematic photorealistic video of the character performing "
"the exact movements and actions from the reference motion video. "
"Natural facial expressions, realistic skin, detailed clothing, "
"cinematic lighting, shallow depth of field, professional camera."
),
"steps": 20,
"guidance": 1.0,
},
"Portrait": {
"prompt": (
"A photorealistic portrait video of the reference character. "
"Preserve the character's identity and facial features while "
"accurately following the body and facial motion from the input video. "
"Natural expression, realistic skin and cinematic portrait lighting."
),
"steps": 20,
"guidance": 1.0,
},
"Anime": {
"prompt": (
"An anime-style cinematic video featuring the reference character, "
"accurately reproducing the movement and performance from the input "
"motion video. Consistent character identity, expressive animation, "
"detailed anime background."
),
"steps": 20,
"guidance": 1.0,
},
}
# ============================================================
# LOAD MODEL
# ============================================================
print("Loading Wan2.2 Animate...")
if not torch.cuda.is_available():
raise RuntimeError(
"CUDA GPU is required. "
"Use a GPU-enabled Hugging Face Space."
)
print("GPU:", torch.cuda.get_device_name(0))
pipe = WanAnimatePipeline.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
)
# Important for limited VRAM
pipe.enable_model_cpu_offload()
print("Model loaded.")
# ============================================================
# GENERATION
# ============================================================
def generate_video(
prompt,
reference_image,
motion_video,
preset,
seed,
inference_steps,
guidance_scale,
progress=gr.Progress(),
):
if reference_image is None:
raise gr.Error("Please upload a reference image.")
if motion_video is None:
raise gr.Error("Please upload a motion video.")
# --------------------------------------------------------
# PRESET
# --------------------------------------------------------
preset_config = PRESETS[preset]
if not prompt or not prompt.strip():
prompt = preset_config["prompt"]
# --------------------------------------------------------
# LOAD INPUTS
# --------------------------------------------------------
progress(0.1, desc="Loading reference image...")
image = load_image(reference_image)
progress(0.2, desc="Loading motion video...")
motion = load_video(motion_video)
# --------------------------------------------------------
# CONDITIONING
#
# For the basic workflow we use the motion video as
# both pose and face conditioning.
#
# For best results, replace this with Wan Animate's
# official preprocessing pipeline.
# --------------------------------------------------------
pose_video = motion
face_video = motion
# --------------------------------------------------------
# GENERATION
# --------------------------------------------------------
progress(0.3, desc="Generating video...")
generator = torch.Generator(
device="cuda"
).manual_seed(int(seed))
with torch.inference_mode():
result = pipe(
image=image,
pose_video=pose_video,
face_video=face_video,
prompt=prompt,
mode="animate",
segment_frame_length=77,
prev_segment_conditioning_frames=1,
guidance_scale=float(guidance_scale),
num_inference_steps=int(inference_steps),
generator=generator,
).frames[0]
progress(0.9, desc="Encoding output video...")
# --------------------------------------------------------
# SAVE
# --------------------------------------------------------
output_path = os.path.join(
OUTPUT_DIR,
f"generated_{int(seed)}.mp4"
)
export_to_video(
result,
output_path,
fps=30,
)
# --------------------------------------------------------
# CLEANUP
# --------------------------------------------------------
gc.collect()
torch.cuda.empty_cache()
progress(1.0, desc="Done")
return output_path
# ============================================================
# PRESET HANDLER
# ============================================================
def update_preset(preset):
config = PRESETS[preset]
return (
config["prompt"],
config["steps"],
config["guidance"],
)
# ============================================================
# GRADIO UI
# ============================================================
with gr.Blocks(
title="Wan2.2 Character Animation"
) as demo:
gr.Markdown(
"""
# 🎬 Wan2.2 Character Animation
Generate a video using:
**Reference Photo + Motion Video + Prompt → Generated Video**
The reference image provides the character identity.
The motion video provides the movement.
"""
)
with gr.Row():
# ----------------------------------------------------
# LEFT
# ----------------------------------------------------
with gr.Column():
prompt = gr.Textbox(
label="Prompt",
placeholder=(
"Describe the generated video..."
),
lines=5,
)
preset = gr.Dropdown(
choices=list(PRESETS.keys()),
value="Realistic",
label="Style Preset",
)
reference_image = gr.Image(
label="Reference Photo / Character",
type="filepath",
)
motion_video = gr.Video(
label="Motion Video",
sources=["upload"],
)
# ----------------------------------------------------
# RIGHT
# ----------------------------------------------------
with gr.Column():
output_video = gr.Video(
label="Generated Video",
autoplay=True,
)
generate_button = gr.Button(
"🎬 Generate Video",
variant="primary",
)
# ========================================================
# ADVANCED SETTINGS
# ========================================================
with gr.Accordion(
"Advanced Settings",
open=False,
):
seed = gr.Number(
label="Seed",
value=42,
precision=0,
)
inference_steps = gr.Slider(
minimum=5,
maximum=50,
value=20,
step=1,
label="Inference Steps",
)
guidance_scale = gr.Slider(
minimum=0.5,
maximum=5.0,
value=1.0,
step=0.1,
label="Guidance Scale",
)
# ========================================================
# PRESET EVENT
# ========================================================
preset.change(
fn=update_preset,
inputs=[preset],
outputs=[
prompt,
inference_steps,
guidance_scale,
],
)
# ========================================================
# GENERATE EVENT
# ========================================================
generate_button.click(
fn=generate_video,
inputs=[
prompt,
reference_image,
motion_video,
preset,
seed,
inference_steps,
guidance_scale,
],
outputs=output_video,
)
# ============================================================
# LAUNCH
# ============================================================
if __name__ == "__main__":
demo.queue(
max_size=10,
default_concurrency_limit=1,
)
demo.launch()