import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import torch
from modular_pipeline import MatrixGameWanModularPipeline
from modular_blocks import ACTION2VIDEO_BLOCKS
from diffusers.modular_pipelines import SequentialPipelineBlocks
from diffusers import AutoModel
from diffusers.utils import load_image, export_to_video

blocks = SequentialPipelineBlocks.from_blocks_dict(ACTION2VIDEO_BLOCKS.copy())
pipe = MatrixGameWanModularPipeline(blocks)
pipe.load_components(torch_dtype=torch.bfloat16)
pipe.load_components(["vae"], repo="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
pipe.load_components(["scheduler"], repo="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="scheduler")

transformer = AutoModel.from_pretrained("./transformer", trust_remote_code=True, torch_dtype=torch.bfloat16)
pipe.transformer = transformer
pipe.to("cuda")

image = load_image("/home/dhruv/matrix-game-workspace/Matrix-Game/Matrix-Game-2/demo_images/universal/0000.png")
output = pipe(image=image, num_frames=141)
export_to_video(output.values['videos'][0], "modular-matrix-game.mp4")