import torch
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers import HunyuanVideoTransformer3DModel, HunyuanVideoPipeline
from diffusers.utils import export_to_video

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
# model = "tencent/HunyuanVideo"
model = "hunyuanvideo-community/HunyuanVideo"
transformer_8bit = HunyuanVideoTransformer3DModel.from_pretrained(
    model,
    subfolder="transformer",
    quantization_config=quant_config,
    torch_dtype=torch.float16,
    # device_map="balanced",
)

pipeline = HunyuanVideoPipeline.from_pretrained(
    "hunyuanvideo-community/HunyuanVideo",
    transformer=transformer_8bit,
    torch_dtype=torch.float16,
    device_map="balanced",
)

prompt = "A cat walks on the grass, realistic style."
video = pipeline(prompt=prompt, num_frames=61, num_inference_steps=3).frames[0]
export_to_video(video, "cat.mp4", fps=15)