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]}...')