Yuan-ManX-ComfyUI-Matrix-Game/nodes.py

import argparse
import os
import glob
from typing import Tuple, Dict, List, Optional
import torch
import numpy as np
from PIL import Image
import imageio
from diffusers.utils import load_image
from diffusers.video_processor import VideoProcessor
from einops import rearrange
from safetensors.torch import load_file as safe_load
from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict

from matrixgame.sample.pipeline_matrixgame import MatrixGameVideoPipeline
from matrixgame.model_variants import get_dit
from matrixgame.vae_variants import get_vae
from matrixgame.encoder_variants import get_text_enc
from matrixgame.model_variants.matrixgame_dit_src import MGVideoDiffusionTransformerI2V
from matrixgame.sample.flow_matching_scheduler_matrixgame import FlowMatchDiscreteScheduler
from tools.visualize import process_video
from condtions import Bench_actions_76
from teacache_forward import teacache_forward


class VideoGenerator:
    """Main class for video generation using MatrixGame model."""
    
    def __init__(self, dit_path, vae_path, textenc_path, image_path, mouse_icon_path, mouse_scale, mouse_rotation, fps, 
                 output_path, video_length, guidance_scale, inference_steps, shift, num_pre_frames, num_steps, 
                 rel_l1_thresh, resolution, bfloat16, max_images, gpu_id):
        """
        Initialize the video generator with configuration parameters.
        """

        # Model paths
        self.dit_path = dit_path
        self.vae_path = vae_path
        self.textenc_path = textenc_path

        # Basic parameters
        self.image_path = image_path
        self.output_path = output_path
        self.video_length = video_length
        self.guidance_scale = guidance_scale

        # Mouse icon parameters
        self.mouse_icon_path = mouse_icon_path
        self.mouse_scale = mouse_scale
        self.mouse_rotation = mouse_rotation
        self.fps = fps
                     
        # Inference parameters             
        self.inference_steps = inference_steps
        self.shift = shift
        self.num_pre_frames = num_pre_frames
        self.num_steps = num_steps
        self.rel_l1_thresh = rel_l1_thresh

        # Resolution and format
        self.resolution = resolution
        self.bfloat16 = bfloat16
        self.max_images = max_images

        # GPU settings
        self.gpu_id = gpu_id
        
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.scheduler = FlowMatchDiscreteScheduler(
            shift=self.shift,
            reverse=True,
            solver="euler"
        )
        
        # Initialize models
        self._init_models()
        
        # Teacache settings
        self._setup_teacache()
    
    def _init_models(self) -> None:
        """Initialize all required models (VAE, text encoder, transformer)."""
        # Initialize VAE
        vae_path = self.vae_path 
        self.vae = get_vae("matrixgame", vae_path, torch.float16)
        self.vae.requires_grad_(False)
        self.vae.eval()
        self.vae.enable_tiling()
        
        # Initialize DIT (Transformer)
        dit = MGVideoDiffusionTransformerI2V.from_pretrained(self.dit_path)
        dit.requires_grad_(False)
        dit.eval()
        
        # Initialize text encoder
        textenc_path = self.textenc_path
        weight_dtype = torch.bfloat16 if self.bfloat16 else torch.float32
        self.text_enc = get_text_enc('matrixgame', textenc_path, weight_dtype=weight_dtype, i2v_type='refiner')
        
        # Move models to devices
        self.pipeline = MatrixGameVideoPipeline(
            vae=self.vae.vae,
            text_encoder=self.text_enc,
            transformer=dit,
            scheduler=self.scheduler,
        ).to(weight_dtype).to(self.device)
    
    def _setup_teacache(self) -> None:
        """Configure teacache for the transformer."""
        self.pipeline.transformer.__class__.enable_teacache = True
        self.pipeline.transformer.__class__.cnt = 0
        self.pipeline.transformer.__class__.num_steps = self.num_steps  # should be aligned with infer_steps
        self.pipeline.transformer.__class__.accumulated_rel_l1_distance = 0
        self.pipeline.transformer.__class__.rel_l1_thresh = self.rel_l1_thresh
        self.pipeline.transformer.__class__.previous_modulated_input = None
        self.pipeline.transformer.__class__.previous_residual = None
        self.pipeline.transformer.__class__.forward = teacache_forward
    
    def _resize_and_crop_image(self, image: Image.Image, target_size: Tuple[int, int]) -> Image.Image:
        """
        Resize and crop image to target dimensions while maintaining aspect ratio.
        
        Args:
            image: Input PIL image
            target_size: Target (width, height) tuple
            
        Returns:
            Resized and cropped PIL image
        """
        w, h = image.size
        tw, th = target_size
        
        if h / w > th / tw:
            new_w = int(w)
            new_h = int(new_w * th / tw)
        else:
            new_h = int(h)
            new_w = int(new_h * tw / th)
        
        left = (w - new_w) / 2
        top = (h - new_h) / 2
        right = (w + new_w) / 2
        bottom = (h + new_h) / 2
        
        return image.crop((left, top, right, bottom))
    
    def _load_images(self, root_dir: str) -> List[str]:
        """
        Load image paths from directory with specified extensions.
        
        Args:
            root_dir: Root directory to search for images
            
        Returns:
            List of image file paths
        """
        image_extensions = ('*.png', '*.jpg', '*.jpeg')
        image_paths = []
        
        for ext in image_extensions:
            image_paths.extend(glob.glob(os.path.join(root_dir, '**', ext), recursive=True))
            
        return image_paths[:self.max_images] if hasattr(self, 'max_images') else image_paths
    
    def _process_condition(self, condition: Dict, image_path: str) -> None:
        """
        Process a single condition and generate video.
        
        Args:
            condition: Condition dictionary containing action and conditions
            image_path: Path to input image
        """
        # Prepare conditions
        keyboard_condition = torch.tensor(condition['keyboard_condition'], dtype=torch.float32).unsqueeze(0)
        mouse_condition = torch.tensor(condition['mouse_condition'], dtype=torch.float32).unsqueeze(0)
        
        # Move to device
        keyboard_condition = keyboard_condition.to(torch.bfloat16 if self.bfloat16 else torch.float16).to(self.device)
        mouse_condition = mouse_condition.to(torch.bfloat16 if self.bfloat16 else torch.float16).to(self.device)
        
        # Load and preprocess image
        image = Image.open(image_path).convert("RGB")
        new_width, new_height = self.resolution
        initial_image = self._resize_and_crop_image(image, (new_width, new_height))
        semantic_image = initial_image
        vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        video_processor = VideoProcessor(vae_scale_factor=vae_scale_factor)
        initial_image = video_processor.preprocess(initial_image, height=new_height, width=new_width)
        
        if self.num_pre_frames > 0:
            past_frames = initial_image.repeat(self.num_pre_frames, 1, 1, 1)
            initial_image = torch.cat([initial_image, past_frames], dim=0)
        
        # Generate video
        with torch.no_grad():
            video = self.pipeline(
                height=new_height,
                width=new_width,
                video_length=self.video_length,
                mouse_condition=mouse_condition,
                keyboard_condition=keyboard_condition,
                initial_image=initial_image,
                num_inference_steps=self.inference_steps,
                guidance_scale=self.guidance_scale,
                embedded_guidance_scale=None,
                data_type="video",
                vae_ver='884-16c-hy',
                enable_tiling=True,
                generator=torch.Generator(device="cuda").manual_seed(42),
                i2v_type='refiner',
                semantic_images=semantic_image
            ).videos[0]
        
        # Save video
        img_tensors = rearrange(video.permute(1, 0, 2, 3) * 255, 't c h w -> t h w c').contiguous()
        img_tensors = img_tensors.cpu().numpy().astype(np.uint8)
        
        config = (
            keyboard_condition[0].float().cpu().numpy(),
            mouse_condition[0].float().cpu().numpy()
        )
        
        action_name = condition['action_name']
        output_filename = f"{os.path.splitext(os.path.basename(image_path))[0]}_{action_name}.mp4"
        output_path = os.path.join(self.output_path, output_filename)
        
        process_video(
            img_tensors,
            output_path,
            config,
            mouse_icon_path=self.mouse_icon_path,
            mouse_scale=self.mouse_scale,
            mouse_rotation=self.mouse_rotation,
            fps=self.fps
        )
            

    
    def generate_videos(self) -> None:
        """Main method to generate videos for all conditions."""
        # Create output directory
        os.makedirs(self.output_path, exist_ok=True)
        
        # Load conditions
        conditions = Bench_actions_76()
        print(f"Found {len(conditions)} conditions to process")
        
        # Load sample images
        root_dir = self.image_path
        image_paths = self._load_images(root_dir)
        
        if not image_paths:
            print("No images found in the specified directory")
            return
        
        # Process each condition
        for idx, condition in enumerate(conditions):
            for image_path in image_paths:
                print(f"Processing condition {idx+1}/{len(conditions)} with image {os.path.basename(image_path)}")
                self._process_condition(condition, image_path)


class LoadDiTModel:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "model_path": ("STRING", {"multiline": False, "default": "models/matrixgame/dit"}),
        }}

    RETURN_TYPES = ("DIT",)
    RETURN_NAMES = ("dit_path",)
    FUNCTION = "load_dit"
    CATEGORY = "Matrix-Game"

    def load_dit(self, model_path):
        dit_path = model_path
        return (dit_path,)
    

class LoadVAEModel:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "vae_path": ("STRING", {"default": "models/matrixgame/vae"}),
            }
        }

    RETURN_TYPES = ("VAE",)
    RETURN_NAMES = ("vae_path",)
    FUNCTION = "load_vae"
    CATEGORY = "Matrix-Game"

    def load_vae(self, vae_path):
        vae_path = vae_path
        return (vae_path,)


class LoadTextEncoderModel:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "textencoder_path": ("STRING", {"default": "models/matrixgame"}),
            }
        }

    RETURN_TYPES = ("TEXTENCODER",)
    RETURN_NAMES = ("textencoder",)
    FUNCTION = "load_textencoder"
    CATEGORY = "Matrix-Game"

    def load_textencoder(self, textencoder_path):
        textenc_path = textencoder_path
        return (textenc_path,)


class LoadGameImage:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "game_image_path": ("STRING", {"default": "initial_image/"}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image_path",)
    FUNCTION = "load_image"
    CATEGORY = "Matrix-Game"

    def load_image(self, game_image_path):
        image_path = game_image_path
        return (image_path,)


class LoadMouseIcon:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mouse_icon_path": ("STRING", {"default": "models/matrixgame/assets/mouse.png"}),
                "mouse_scale": ("FLOAT", {"default": 0.1}),
                "mouse_rotation": ("FLOAT", {"default": -20}),
                "fps": ("INT", {"default": 16}),
            }
        }

    RETURN_TYPES = ("MOUSEICON", "MOUSESCALE", "MOUSEROTATION", "FPS")
    RETURN_NAMES = ("mouse_icon_path", "mouse_scale", "mouse_rotation", "fps")
    FUNCTION = "load_mouse_icon"
    CATEGORY = "Matrix-Game"

    def load_mouse_icon(self, mouse_icon_path, mouse_scale, mouse_rotation, fps):
        mouse_icon_path = mouse_icon_path
        mouse_scale = mouse_scale
        mouse_rotation = mouse_rotation
        fps = fps
        return (mouse_icon_path, mouse_scale, mouse_rotation, mouse_rotation, fps)


class GameVideoGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "dit_path": ("DIT",),
                "vae_path": ("VAE",),
                "textenc_path": ("TEXTENCODER",),
                "image_path": ("IMAGE",),
                "mouse_icon_path": ("MOUSEICON",),
                "mouse_scale": ("MOUSESCALE",),
                "mouse_rotation": ("MOUSEROTATION",),
                "fps": ("FPS",),
                "output_path": ("OUTPUT",),
                "video_length": ("INT", {"default": 65}),
                "guidance_scale": ("FLOAT", {"default": 6}),
                "inference_steps": ("INT", {"default": 50}),
                "shift": ("INT", {"default": 15.0}),
                "num_pre_frames": ("INT", {"default": 5}),
                "num_steps": ("INT", {"default": 50}),
                "rel_l1_thresh": ("FLOAT", {"default": 0.075}),
                "resolution_h": ("INT", {"default": 720}),
                "resolution_w": ("INT", {"default": 1280}),
                "bfloat16": ("STRING", {"default": "store_true"}),
                "max_images": ("INT", {"default": 3}),
                "gpu_id": ("STRING", {"default": "0"}),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "generate_game_videos"
    CATEGORY = "Matrix-Game"

    def generate_game_videos(self, dit_path, vae_path, textenc_path, image_path, mouse_icon_path, mouse_scale, mouse_rotation, fps, 
                             output_path, video_length, guidance_scale, inference_steps, shift, num_pre_frames, num_steps, 
                             rel_l1_thresh, resolution_h, resolution_w, bfloat16, max_images, gpu_id):
        
        resolution = [resolution_w, resolution_h]
                   
        # Set GPU device if specified
        if gpu_id and torch.cuda.is_available():
            os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
        
        generator = VideoGenerator(dit_path, vae_path, textenc_path, image_path, mouse_icon_path, mouse_scale, mouse_rotation, fps, 
                             output_path, video_length, guidance_scale, inference_steps, shift, num_pre_frames, num_steps, 
                             rel_l1_thresh, resolution, bfloat16, max_images, gpu_id)
                                 
        generator.generate_videos()
        return ()


class MatrixGameOutput:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "output_path": ("STRING", {"default": "./test"}),
            }
        }

    RETURN_TYPES = ("OUTPUT",)
    RETURN_NAMES = ("output_path",)
    FUNCTION = "load_output_path"
    CATEGORY = "Matrix-Game"

    def load_output_path(self, output_path):
        output_path = output_path
        return (output_path,)