#!/usr/bin/env python3
"""
Wan2.1 图像到视频 CLI 脚本。
用法：python generate_i2v.py --model Wan-AI/Wan2.1-I2V-14B-480P \
--image input.jpg --prompt "camera slowly zooms out"
"""
import argparse
批处理处理
"""通过 Ollama 使用 LLaVA 分析图像"""
import torch
from diffusers import WanImageToVideoPipeline
from diffusers.utils import load_image, export_to_video
from PIL import Image
def parse_args():
parser = argparse.ArgumentParser(description="Wan2.1 图像到视频生成器")
parser.add_argument(
"--model",
type=str,
default="Wan-AI/Wan2.1-I2V-14B-480P",
help="来自 Hugging Face 的模型 ID（默认：Wan-AI/Wan2.1-I2V-14B-480P）",
)
parser.add_argument(
"--image",
type=str,
required=True,
help="输入图像路径（JPEG 或 PNG）",
)
parser.add_argument(
"--prompt",
type=str,
required=True,
help='描述期望运动的文本提示（例如 "camera slowly zooms out"）',
)
parser.add_argument(
"--negative-prompt",
type=str,
default="模糊、低质量、变形、抖动的运动、伪影",
help="用于避免不想要伪影的负面提示",
)
parser.add_argument(
"--frames",
type=int,
default=49,
help="要生成的视频帧数（默认：49，最大：81）",
)
parser.add_argument(
"--steps",
type=int,
default=50,
help="扩散步骤数（默认：50）",
)
parser.add_argument(
"--guidance",
type=float,
default=7.0,
help="无分类器引导尺度（默认：7.0）",
)
parser.add_argument(
"--seed",
type=int,
default=-1,
help="用于可重复性的随机种子（-1 = 随机）",
)
parser.add_argument(
"--fps",
type=int,
default=16,
help="输出视频帧率（默认：16）",
)
parser.add_argument(
"--output",
type=str,
default="output_i2v.mp4",
help="输出视频文件路径（默认：output_i2v.mp4）",
)
parser.add_argument(
"--height",
type=int,
default=480,
help="输出视频高度（像素）（默认：480）",
)
parser.add_argument(
"--width",
type=int,
default=854,
help="输出视频宽度（像素）（默认：854）",
)
parser.add_argument(
"--cpu-offload",
action="store_true",
default=True,
help="启用模型 CPU 卸载以节省显存（默认：True）",
)
parser.add_argument(
"--vae-tiling",
action="store_true",
default=False,
help="启用 VAE 平铺以生成高分辨率输出",
)
return parser.parse_args()
def load_and_resize_image(image_path: str, width: int, height: int) -> Image.Image:
"""从路径加载图像并调整为目标尺寸。"""
if not os.path.exists(image_path):
print(f"[ERROR] 未找到图像：{image_path}", file=sys.stderr)
sys.exit(1)
img = Image.open(image_path).convert("RGB")
original_size = img.size
img = img.resize((width, height), Image.LANCZOS)
print(f"[INFO] 已加载图像：{image_path} ({original_size[0]}x{original_size[1]}) → 调整为 {width}x{height}")
return img
def load_pipeline(model_id: str, cpu_offload: bool, vae_tiling: bool):
"""使用内存优化加载 Wan I2V 管道。"""
print(f"[INFO] 正在加载模型：{model_id}")
print(f"[INFO] CUDA 可用：{torch.cuda.is_available()}")
if torch.cuda.is_available():
vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
print(f"[INFO] GPU：{torch.cuda.get_device_name(0)} ({vram_gb:.1f} GB 显存)")
if vram_gb < 23:
print("[WARN] 检测到显存少于 24GB — 启用 --cpu-offload 或使用 1.3B 模型")
pipe = WanImageToVideoPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16,
)
if cpu_offload:
print("[INFO] 启用模型 CPU 卸载")
pipe.enable_model_cpu_offload()
else:
pipe.to("cuda")
if vae_tiling:
print("[INFO] 为高分辨率生成启用 VAE 平铺")
pipe.enable_vae_tiling()
return pipe
def generate_video(pipe, args) -> None:
"""运行 I2V 管道并保存输出视频。"""
image = load_and_resize_image(args.image, args.width, args.height)
generator = None
if args.seed >= 0:
generator = torch.Generator("cuda").manual_seed(args.seed)
print(f"[INFO] 使用种子：{args.seed}")
else:
print("[INFO] 使用随机种子")
print(f"[INFO] 正在生成 {args.frames} 帧，分辨率为 {args.width}x{args.height}")
print(f"[INFO] 步数：{args.steps} | 引导：{args.guidance} | FPS：{args.fps}")
print(f"[INFO] 提示词：{args.prompt}")
output = pipe(
prompt=args.prompt,
negative_prompt=args.negative_prompt,
image=image,
num_frames=args.frames,
height=args.height,
width=args.width,
num_inference_steps=args.steps,
guidance_scale=args.guidance,
generator=generator,
)
export_to_video(output.frames[0], args.output, fps=args.fps)
print(f"[INFO] 视频已保存到：{os.path.abspath(args.output)}")
duration = args.frames / args.fps
print(f"[INFO] 时长：{duration:.1f}s，{args.fps}fps（{args.frames} 帧）")
def main():
args = parse_args()
if not torch.cuda.is_available():
print("[ERROR] 未找到 CUDA GPU。Wan2.1-I2V-14B 需要支持 CUDA 的 GPU。", file=sys.stderr)
sys.exit(1)
pipe = load_pipeline(args.model, args.cpu_offload, args.vae_tiling)
generate_video(pipe, args)
print("[DONE] 图像到视频生成完成！")
if __name__ == "__main__":
main()
