|
---
|
|
license: apache-2.0
|
|
---
|
|
|
|
```
|
|
import torch
|
|
from transformers import AutoTokenizer, UMT5EncoderModel
|
|
from diffusers import AutoencoderKLWan, WanPipeline, WanTransformer3DModel, FlowMatchEulerDiscreteScheduler
|
|
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
|
|
from diffusers.utils import export_to_video
|
|
from torchvision import transforms
|
|
import os
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
from pathlib import Path
|
|
import json
|
|
from safetensors.torch import safe_open
|
|
|
|
device = "cuda"
|
|
seed = 0
|
|
|
|
# TODO: impl AutoencoderKLWan
|
|
vae = vae.from_pretrained("StevenZhang/Wan2.1-VAE_Diff")
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
vae = vae.to(device)
|
|
|
|
# TODO: impl FlowDPMSolverMultistepScheduler
|
|
scheduler = UniPCMultistepScheduler(prediction_type='flow_prediction', use_flow_sigmas=True, num_train_timesteps=1000, flow_shift=1.0)
|
|
|
|
text_encoder = UMT5EncoderModel.from_pretrained("google/umt5-xxl", torch_dtype=torch.bfloat16)
|
|
tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl")
|
|
|
|
# 14B
|
|
transformer = WanTransformer3DModel.from_pretrained('StevenZhang/Wan2.1-T2V-14B-Diff', torch_dtype=torch.bfloat16)
|
|
# transformer = WanTransformer3DModel.from_pretrained('StevenZhang/Wan2.1-T2V-1.3B-Diff', torch_dtype=torch.bfloat16)
|
|
|
|
components = {
|
|
"transformer": transformer,
|
|
"vae": vae,
|
|
"scheduler": scheduler,
|
|
"text_encoder": text_encoder,
|
|
"tokenizer": tokenizer,
|
|
}
|
|
pipe = WanPipeline(**components)
|
|
|
|
pipe.to(device)
|
|
|
|
negative_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
|
|
|
|
generator = torch.Generator(device=device).manual_seed(seed)
|
|
inputs = {
|
|
"prompt": "两只拟人化的猫咪身穿舒适的拳击装备，戴着鲜艳的手套，在聚光灯照射的舞台上激烈对战",
|
|
"negative_prompt": negative_prompt, # TODO
|
|
"generator": generator,
|
|
"num_inference_steps": 50,
|
|
"flow_shift": 3.0,
|
|
"guidance_scale": 5.0,
|
|
"height": 480,
|
|
"width": 832,
|
|
"num_frames": 81,
|
|
"max_sequence_length": 512,
|
|
"output_type": "np"
|
|
}
|
|
|
|
video = pipe(**inputs).frames[0]
|
|
|
|
print(video.shape)
|
|
|
|
export_to_video(video, "output.mp4", fps=16)
|
|
```
|
|
|
