modules/generators/video_generator.py

import torch
import os
import numpy as np
import math
import decord
from tqdm import tqdm
import pathlib
from PIL import Image

from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
from diffusers_helper.memory import DynamicSwapInstaller
from diffusers_helper.utils import resize_and_center_crop
from diffusers_helper.bucket_tools import find_nearest_bucket
from diffusers_helper.hunyuan import vae_encode, vae_decode
from .video_base_generator import VideoBaseModelGenerator

class VideoModelGenerator(VideoBaseModelGenerator):
    """
    Generator for the Video (backward) extension of the Original HunyuanVideo model.
    These generators accept video input instead of a single image.
    """
    
    def __init__(self, **kwargs):
        """
        Initialize the Video model generator.
        """
        super().__init__(**kwargs)
        self.model_name = "Video"
        self.model_path = 'lllyasviel/FramePackI2V_HY'  # Same as Original
        self.model_repo_id_for_cache = "models--lllyasviel--FramePackI2V_HY"
    
    def get_latent_paddings(self, total_latent_sections):
        """
        Get the latent paddings for the Video model.
        
        Args:
            total_latent_sections: The total number of latent sections
            
        Returns:
            A list of latent paddings
        """
        # Video model uses reversed latent paddings like Original
        if total_latent_sections > 4:
            return [3] + [2] * (total_latent_sections - 3) + [1, 0]
        else:
            return list(reversed(range(total_latent_sections)))

    def video_prepare_clean_latents_and_indices(self, end_frame_output_dimensions_latent, end_frame_weight, end_clip_embedding, end_of_input_video_embedding, latent_paddings, latent_padding, latent_padding_size, latent_window_size, video_latents, history_latents, num_cleaned_frames=5):
        """
        Combined method to prepare clean latents and indices for the Video model.
        
        Args:
            Work in progress - better not to pass in latent_paddings and latent_padding.
            num_cleaned_frames: Number of context frames to use from the video (adherence to video)
            
        Returns:
            A tuple of (clean_latent_indices, latent_indices, clean_latent_2x_indices, clean_latent_4x_indices, clean_latents, clean_latents_2x, clean_latents_4x)
        """
        # Get num_cleaned_frames from job_params if available, otherwise use default value of 5
        num_clean_frames = num_cleaned_frames if num_cleaned_frames is not None else 5


        # HACK SOME STUFF IN THAT SHOULD NOT BE HERE
        # Placeholders for end frame processing
        # Colin, I'm only leaving them for the moment in case you want separate models for
        # Video-backward and Video-backward-Endframe.
        # end_latent = None
        # end_of_input_video_embedding = None # Placeholder for end frame's CLIP embedding. SEE: 20250507 pftq: Process end frame if provided
        # end_clip_embedding = None # Placeholders for end frame processing. SEE: 20250507 pftq: Process end frame if provided
        # end_frame_weight = 0.0 # Placeholders for end frame processing. SEE: 20250507 pftq: Process end frame if provided
        # HACK MORE STUFF IN THAT PROBABLY SHOULD BE ARGUMENTS OR OTHWISE MADE AVAILABLE
        end_of_input_video_latent = video_latents[:, :, -1:] # Last frame of the input video (produced by video_encode in the PR)
        is_start_of_video = latent_padding == 0 # This refers to the start of the *generated* video part
        is_end_of_video = latent_padding == latent_paddings[0] # This refers to the end of the *generated* video part (closest to input video) (better not to pass in latent_paddings[])
        # End of HACK STUFF

        # Dynamic frame allocation for context frames (clean latents)
        # This determines which frames from history_latents are used as input for the transformer.
        available_frames = video_latents.shape[2] if is_start_of_video else history_latents.shape[2] # Use input video frames for first segment, else previously generated history
        effective_clean_frames = max(0, num_clean_frames - 1) if num_clean_frames > 1 else 1
        if is_start_of_video:
            effective_clean_frames = 1 # Avoid jumpcuts if input video is too different

        clean_latent_pre_frames = effective_clean_frames 
        num_2x_frames = min(2, max(1, available_frames - clean_latent_pre_frames - 1)) if available_frames > clean_latent_pre_frames + 1 else 1
        num_4x_frames = min(16, max(1, available_frames - clean_latent_pre_frames - num_2x_frames)) if available_frames > clean_latent_pre_frames + num_2x_frames else 1
        total_context_frames = num_2x_frames + num_4x_frames
        total_context_frames = min(total_context_frames, available_frames - clean_latent_pre_frames)

        # Prepare indices for the transformer's input (these define the *relative positions* of different frame types in the input tensor)
        # The total length is the sum of various frame types:
        # clean_latent_pre_frames: frames before the blank/generated section
        # latent_padding_size: blank frames before the generated section (for backward generation)
        # latent_window_size: the new frames to be generated
        # post_frames: frames after the generated section
        # num_2x_frames, num_4x_frames: frames for lower resolution context
        # 20250511 pftq: Dynamically adjust post_frames based on clean_latents_post
        post_frames = 1 if is_end_of_video and end_frame_output_dimensions_latent is not None else effective_clean_frames  # 20250511 pftq: Single frame for end_latent, otherwise padding causes still image
        indices = torch.arange(0, clean_latent_pre_frames + latent_padding_size + latent_window_size + post_frames + num_2x_frames + num_4x_frames).unsqueeze(0)
        clean_latent_indices_pre, blank_indices, latent_indices, clean_latent_indices_post, clean_latent_2x_indices, clean_latent_4x_indices = indices.split(
            [clean_latent_pre_frames, latent_padding_size, latent_window_size, post_frames, num_2x_frames, num_4x_frames], dim=1
        )
        clean_latent_indices = torch.cat([clean_latent_indices_pre, clean_latent_indices_post], dim=1) # Combined indices for 1x clean latents

        # Prepare the *actual latent data* for the transformer's context inputs
        # These are extracted from history_latents (or video_latents for the first segment)
        context_frames = history_latents[:, :, -(total_context_frames + clean_latent_pre_frames):-clean_latent_pre_frames, :, :] if total_context_frames > 0 else history_latents[:, :, :1, :, :]
        # clean_latents_4x: 4x downsampled context frames. From history_latents (or video_latents).
        # clean_latents_2x: 2x downsampled context frames. From history_latents (or video_latents).
        split_sizes = [num_4x_frames, num_2x_frames]
        split_sizes = [s for s in split_sizes if s > 0]
        if split_sizes and context_frames.shape[2] >= sum(split_sizes):
            splits = context_frames.split(split_sizes, dim=2)
            split_idx = 0
            clean_latents_4x = splits[split_idx] if num_4x_frames > 0 else history_latents[:, :, :1, :, :]
            split_idx += 1 if num_4x_frames > 0 else 0
            clean_latents_2x = splits[split_idx] if num_2x_frames > 0 and split_idx < len(splits) else history_latents[:, :, :1, :, :]
        else:
            clean_latents_4x = clean_latents_2x = history_latents[:, :, :1, :, :]

        # clean_latents_pre: Latents from the *end* of the input video (if is_start_of_video), or previously generated history.
        # Its purpose is to provide a smooth transition *from* the input video.
        clean_latents_pre = video_latents[:, :, -min(effective_clean_frames, video_latents.shape[2]):].to(history_latents)
        
        # clean_latents_post: Latents from the *beginning* of the previously generated video segments.
        # Its purpose is to provide a smooth transition *to* the existing generated content.
        clean_latents_post = history_latents[:, :, :min(effective_clean_frames, history_latents.shape[2]), :, :]

        # Special handling for the end frame:
        # If it's the very first segment being generated (is_end_of_video in terms of generation order),
        # and an end_latent was provided, force clean_latents_post to be that end_latent.
        if is_end_of_video:
            clean_latents_post = torch.zeros_like(end_of_input_video_latent).to(history_latents) # Initialize to zero
        
        # RT_BORG: end_of_input_video_embedding and end_clip_embedding shouldn't need to be checked, since they should
        # always be provided if end_latent is provided. But bulletproofing before the release since test time will be short.
        if end_frame_output_dimensions_latent is not None and end_of_input_video_embedding is not None and end_clip_embedding is not None:
            # image_encoder_last_hidden_state: Weighted average of CLIP embedding of first input frame and end frame's CLIP embedding
            # This guides the overall content to transition towards the end frame.
            image_encoder_last_hidden_state = (1 - end_frame_weight) * end_of_input_video_embedding + end_clip_embedding * end_frame_weight
            image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(self.transformer.dtype)
            
            if is_end_of_video:
                # For the very first generated segment, the "post" part is the end_latent itself.
                clean_latents_post = end_frame_output_dimensions_latent.to(history_latents)[:, :, :1, :, :] # Ensure single frame
            
        # Pad clean_latents_pre/post if they have fewer frames than specified by clean_latent_pre_frames/post_frames
        if clean_latents_pre.shape[2] < clean_latent_pre_frames:
            clean_latents_pre = clean_latents_pre.repeat(1, 1, math.ceil(clean_latent_pre_frames / clean_latents_pre.shape[2]), 1, 1)[:,:,:clean_latent_pre_frames]
        if clean_latents_post.shape[2] < post_frames:
            clean_latents_post = clean_latents_post.repeat(1, 1, math.ceil(post_frames / clean_latents_post.shape[2]), 1, 1)[:,:,:post_frames]
            
        # clean_latents: Concatenation of pre and post clean latents. These are the 1x resolution context frames.
        clean_latents = torch.cat([clean_latents_pre, clean_latents_post], dim=2)

        return clean_latent_indices, latent_indices, clean_latent_2x_indices, clean_latent_4x_indices, clean_latents, clean_latents_2x, clean_latents_4x
    
    def update_history_latents(self, history_latents, generated_latents):
        """
        Backward Generation: Update the history latents with the generated latents for the Video model.
        
        Args:
            history_latents: The history latents
            generated_latents: The generated latents
            
        Returns:
            The updated history latents
        """
        # For Video model, we prepend the generated latents to the front of history latents
        # This matches the original implementation in video-example.py
        # It generates new sections backwards in time, chunk by chunk
        return torch.cat([generated_latents.to(history_latents), history_latents], dim=2)
    
    def get_real_history_latents(self, history_latents, total_generated_latent_frames):
        """
        Get the real history latents for the backward Video model. For Video, this is the first
        `total_generated_latent_frames` frames of the history latents.
        
        Args:
            history_latents: The history latents
            total_generated_latent_frames: The total number of generated latent frames
            
        Returns:
            The real history latents
        """
        # Generated frames at the front.
        return history_latents[:, :, :total_generated_latent_frames, :, :]
    
    def update_history_pixels(self, history_pixels, current_pixels, overlapped_frames):
        """
        Update the history pixels with the current pixels for the Video model.
        
        Args:
            history_pixels: The history pixels
            current_pixels: The current pixels
            overlapped_frames: The number of overlapped frames
            
        Returns:
            The updated history pixels
        """
        from diffusers_helper.utils import soft_append_bcthw
        # For Video model, we prepend the current pixels to the history pixels
        # This matches the original implementation in video-example.py
        return soft_append_bcthw(current_pixels, history_pixels, overlapped_frames)
    
    def get_current_pixels(self, real_history_latents, section_latent_frames, vae):
        """
        Get the current pixels for the Video model.
        
        Args:
            real_history_latents: The real history latents
            section_latent_frames: The number of section latent frames
            vae: The VAE model
            
        Returns:
            The current pixels
        """
        # For backward Video mode, current pixels are at the front of history.
        return vae_decode(real_history_latents[:, :, :section_latent_frames], vae).cpu()
    
    def format_position_description(self, total_generated_latent_frames, current_pos, original_pos, current_prompt):
        """
        Format the position description for the Video model.
        
        Args:
            total_generated_latent_frames: The total number of generated latent frames
            current_pos: The current position in seconds (includes input video time)
            original_pos: The original position in seconds
            current_prompt: The current prompt
            
        Returns:
            The formatted position description
        """
        # For Video model, current_pos already includes the input video time
        # We just need to display the total generated frames and the current position
        return (f'Total generated frames: {int(max(0, total_generated_latent_frames * 4 - 3))}, '
                f'Generated video length: {max(0, (total_generated_latent_frames * 4 - 3) / 30):.2f} seconds (FPS-30). '
                f'Current position: {current_pos:.2f}s (remaining: {original_pos:.2f}s). '
                f'using prompt: {current_prompt[:256]}...')