04 - 视频生成实战¶

# 方法1：使用HuggingFace Hub下载（推荐）
from huggingface_hub import snapshot_download

# CogVideoX-5B（约20GB）
snapshot_download(
    repo_id="THUDM/CogVideoX-5b",
    local_dir="./models/CogVideoX-5b",
    # 国内用户可设置镜像
    # endpoint="https://hf-mirror.com"
)

# CogVideoX-2B（约10GB，显存不够时使用）
snapshot_download(
    repo_id="THUDM/CogVideoX-2b",
    local_dir="./models/CogVideoX-2b",
)

import torch
import diffusers

print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA可用: {torch.cuda.is_available()}")
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"显存: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f}GB")
print(f"Diffusers版本: {diffusers.__version__}")

"""
CogVideoX 文本到视频 — 完整推理流程
"""
import torch
from diffusers import CogVideoXPipeline
from diffusers.utils import export_to_video

# === 1. 加载模型 ===
model_path = "THUDM/CogVideoX-5b"  # 或本地路径 "./models/CogVideoX-5b"

pipe = CogVideoXPipeline.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,  # CogVideoX推荐bfloat16
)

# === 2. 显存优化 ===
pipe.enable_model_cpu_offload()   # 按需加载到GPU
pipe.vae.enable_tiling()          # VAE分块解码
pipe.vae.enable_slicing()         # VAE切片处理

# === 3. 生成视频 ===
prompt = """
A majestic eagle soaring over snow-capped mountains during golden hour.
The camera slowly pans from left to right, revealing a vast landscape
with a crystal-clear lake below. Cinematic, 4K quality.
"""

# 生成参数
video_frames = pipe(
    prompt=prompt,
    num_videos_per_prompt=1,
    num_inference_steps=50,       # 采样步数（默认50）
    num_frames=49,                # 帧数（CogVideoX: 49帧=6秒@8fps）
    guidance_scale=6.0,           # CFG强度
    generator=torch.Generator(device="cpu").manual_seed(42),
).frames[0]

# === 4. 保存视频 ===
export_to_video(video_frames, "text2video_output.mp4", fps=8)
print(f"视频已保存: text2video_output.mp4")
print(f"帧数: {len(video_frames)}, 分辨率: {video_frames[0].size}")

"""
CogVideoX 批量生成与参数探索
"""
import torch
import os
from diffusers import CogVideoXPipeline
from diffusers.utils import export_to_video

pipe = CogVideoXPipeline.from_pretrained(
    "THUDM/CogVideoX-2b",      # 使用2B版本降低显存需求
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()
pipe.vae.enable_tiling()

# 不同提示词和参数的生成实验
experiments = [
    {
        "prompt": "Ocean waves crashing on a tropical beach at sunset, slow motion",
        "guidance_scale": 6.0,
        "num_inference_steps": 50,
        "seed": 42,
    },
    {
        "prompt": "A cat sitting on a windowsill watching rain outside, cozy atmosphere",
        "guidance_scale": 7.0,
        "num_inference_steps": 50,
        "seed": 123,
    },
    {
        "prompt": "Timelapse of a flower blooming in a garden, macro photography",
        "guidance_scale": 5.0,
        "num_inference_steps": 30,  # 测试较少步数
        "seed": 456,
    },
]

os.makedirs("outputs", exist_ok=True)

for i, exp in enumerate(experiments):  # enumerate同时获取索引和元素
    print(f"\n[{i+1}/{len(experiments)}] 生成: {exp['prompt'][:50]}...")  # 切片操作，取前n个元素

    video = pipe(
        prompt=exp["prompt"],
        num_frames=49,
        guidance_scale=exp["guidance_scale"],
        num_inference_steps=exp["num_inference_steps"],
        generator=torch.Generator("cpu").manual_seed(exp["seed"]),
    ).frames[0]

    output_path = f"outputs/exp_{i+1}_cfg{exp['guidance_scale']}_steps{exp['num_inference_steps']}.mp4"
    export_to_video(video, output_path, fps=8)
    print(f"  已保存: {output_path}")

# === CogVideoX的Prompt最佳实践 ===

# ✅ 好的prompt：详细、结构化
good_prompt = """
A young woman with long black hair walks through a Japanese garden in autumn.
Cherry blossom petals fall gently around her.
The camera follows her from behind, slowly revealing a traditional wooden bridge.
Warm afternoon sunlight filters through the red maple leaves.
Cinematic, shallow depth of field, 4K, Film grain.
"""

# ❌ 差的prompt：太短、缺乏细节
bad_prompt = "woman walking in garden"

# 📌 Prompt结构建议：
# 1. 主体描述（who/what）
# 2. 场景环境（where）
# 3. 动作/运动（action）
# 4. 摄像机运动（camera movement）
# 5. 光照/氛围（lighting/mood）
# 6. 质量标签（quality tags）

"""
CogVideoX 图像到视频 — 从单张图片生成动态视频
"""
import torch
from diffusers import CogVideoXImageToVideoPipeline
from diffusers.utils import export_to_video, load_image

# 加载Image-to-Video模型
pipe = CogVideoXImageToVideoPipeline.from_pretrained(
    "THUDM/CogVideoX-5b-I2V",
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()
pipe.vae.enable_tiling()

# 加载输入图像（建议720×480或480×720）
image = load_image("input_scene.png").resize((720, 480))

prompt = "The scene comes alive with gentle wind blowing through the trees, birds flying across the sky"

video = pipe(
    prompt=prompt,
    image=image,
    num_frames=49,
    num_inference_steps=50,
    guidance_scale=6.0,
    generator=torch.Generator("cpu").manual_seed(42),
).frames[0]

export_to_video(video, "img2video_output.mp4", fps=8)
print("图像到视频生成完成!")

"""
Open-Sora 文本到视频推理
注意：需要在Open-Sora仓库目录下运行
"""
import torch
from opensora.utils.inference_utils import (
    load_model,
    get_save_path_name,
    save_sample,
)

# 加载模型配置
config = {
    "model_type": "STDiT2-XL/2",
    "resolution": "480p",
    "num_frames": 51,
    "fps": 24,
}

# 初始化模型
model, vae, text_encoder, scheduler = load_model(
    model_path="hpcai-tech/OpenSora-STDiT-v3",
    config=config,
)

# 生成
prompt = "Aerial view of a coastal city with waves crashing on the shore, golden hour lighting"

samples = model.generate(
    prompt=prompt,
    num_frames=config["num_frames"],
    resolution=config["resolution"],
    num_sampling_steps=30,
    cfg_scale=7.0,
)

# 保存
save_sample(samples, "opensora_output.mp4", fps=24)

# === 显存优化方案 ===

# 方案1: 降低精度
pipe = CogVideoXPipeline.from_pretrained(model_path, torch_dtype=torch.float16)

# 方案2: CPU Offload + VAE优化（最推荐）
pipe.enable_model_cpu_offload()
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()

# 方案3: 使用更小的模型
# CogVideoX: 5B → 2B
# Open-Sora: 选择较小的STDiT变体

# 方案4: 减少帧数和分辨率
video = pipe(prompt=prompt, num_frames=25, height=320, width=480)

# 方案5: 序列并行 (多GPU)
# CogVideoX支持在多张GPU上分布式推理
pipe.enable_sequential_cpu_offload()  # 逐层卸载，比enable_model_cpu_offload更激进

# === 常见质量问题与解决方案 ===

# 问题1: 运动不自然/抖动
# 解决: 增加采样步数，适当降低CFG
video = pipe(prompt=prompt, num_inference_steps=50, guidance_scale=5.0)

# 问题2: 时间不一致（物体闪烁/变形）
# 解决: 使用更长的负面提示词
negative = "flickering, morphing, inconsistent, blurry, distorted, low quality"

# 问题3: 内容与文本不匹配
# 解决: 改进prompt，使用更精确的描述
# 加入: 主语、动作、场景、摄像机运动、光照、风格

# 问题4: 生成速度慢
# 解决: 减少步数 + 使用torch.compile
pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=True)

"""
视频后处理工具
"""
import subprocess
import os

def upscale_video(input_path, output_path, scale=2):
    """使用ffmpeg进行视频超分辨率（双三次插值）"""
    cmd = f'ffmpeg -i {input_path} -vf "scale=iw*{scale}:ih*{scale}:flags=bicubic" -c:v libx264 -crf 18 {output_path}'
    subprocess.run(cmd, shell=True)

def interpolate_frames(input_path, output_path, target_fps=24):
    """帧插值提高帧率"""
    cmd = f'ffmpeg -i {input_path} -filter:v "minterpolate=fps={target_fps}:mi_mode=mci" -c:v libx264 -crf 18 {output_path}'
    subprocess.run(cmd, shell=True)

def add_audio(video_path, audio_path, output_path):
    """为生成的视频添加背景音乐"""
    cmd = f'ffmpeg -i {video_path} -i {audio_path} -c:v copy -c:a aac -shortest {output_path}'
    subprocess.run(cmd, shell=True)

def create_loop(input_path, output_path, loops=3):
    """创建循环视频"""
    cmd = f'ffmpeg -stream_loop {loops-1} -i {input_path} -c copy {output_path}'
    subprocess.run(cmd, shell=True)

# 使用示例
# upscale_video("output_480p.mp4", "output_960p.mp4", scale=2)
# interpolate_frames("output_8fps.mp4", "output_24fps.mp4", target_fps=24)

"""
完整视频生成项目 — 从prompt到最终输出
"""
import torch
import os
import time
from diffusers import CogVideoXPipeline
from diffusers.utils import export_to_video

def generate_video(
    prompt: str,
    output_dir: str = "project_outputs",
    model_id: str = "THUDM/CogVideoX-2b",
    num_frames: int = 49,
    num_steps: int = 50,
    guidance_scale: float = 6.0,
    seed: int = 42,
):
    """端到端视频生成函数"""
    os.makedirs(output_dir, exist_ok=True)

    print("=" * 60)
    print(f"视频生成项目")
    print(f"模型: {model_id}")
    print(f"Prompt: {prompt[:80]}...")
    print("=" * 60)

    # 1. 加载模型
    print("\n[1/4] 加载模型...")
    t0 = time.time()
    pipe = CogVideoXPipeline.from_pretrained(
        model_id, torch_dtype=torch.bfloat16
    )
    pipe.enable_model_cpu_offload()
    pipe.vae.enable_tiling()
    print(f"  模型加载耗时: {time.time()-t0:.1f}s")

    # 2. 生成视频
    print("\n[2/4] 生成视频...")
    t1 = time.time()
    result = pipe(
        prompt=prompt,
        num_frames=num_frames,
        num_inference_steps=num_steps,
        guidance_scale=guidance_scale,
        generator=torch.Generator("cpu").manual_seed(seed),
    )
    gen_time = time.time() - t1
    print(f"  生成耗时: {gen_time:.1f}s")

    # 3. 保存结果
    print("\n[3/4] 保存视频...")
    video_frames = result.frames[0]
    output_path = os.path.join(output_dir, f"video_seed{seed}.mp4")
    export_to_video(video_frames, output_path, fps=8)

    # 4. 输出摘要
    print("\n[4/4] 生成摘要:")
    print(f"  输出路径: {output_path}")
    print(f"  帧数: {len(video_frames)}")
    print(f"  分辨率: {video_frames[0].size}")
    print(f"  时长: {len(video_frames)/8:.1f}秒")
    print(f"  总耗时: {time.time()-t0:.1f}s")

    return output_path

if __name__ == "__main__":
    prompts = [
        "A time-lapse of clouds moving over a mountain range, dramatic lighting, 4K cinematic",
        "Close-up of coffee being poured into a white cup, steam rising, warm morning light",
    ]

    for i, prompt in enumerate(prompts):
        print(f"\n{'='*60}")
        print(f"生成第 {i+1}/{len(prompts)} 个视频")
        generate_video(prompt=prompt, seed=i*100)