import torch
from fastvideo.models.hunyuan.diffusion.pipelines.pipeline_hunyuan_video import HunyuanVideoPipeline
from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer
from fastvideo.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D


transformer=HYVideoDiffusionTransformer.from_pretrained(
    'data/hyvideo-diffusers',
    torch_dtype=torch.bfloat16,
    subfolder='transformer'
)
vae=AutoencoderKLCausal3D.from_pretrained(
    'data/hyvideo-diffusers',
    torch_dtype=torch.float16,
    subfolder='vae'
)

pipe = HunyuanVideoPipeline.from_pretrained(
    'data/hyvideo-diffusers',
    transformer=transformer,
    vae=vae
)
pipe = pipe.to('cuda')
pipe.vae.enable_tiling()

prompt = "Close-up, A little girl wearing a red hoodie in winter strikes a match. The sky is dark, there is a layer of snow on the ground, and it is still snowing lightly. The flame of the match flickers, illuminating the girl's face intermittently."

result = pipe(
    prompt,
    height=512,
    width=512,
    video_length=29,
)

import PIL.Image
from diffusers.utils import export_to_video

output = result.videos[0].permute(1, 2, 3, 0).detach().cpu().numpy()
output = (output * 255).clip(0, 255).astype("uint8")
output = [PIL.Image.fromarray(x) for x in output]

export_to_video(output, "output.mp4", fps=24)