import torch import os # for offline loading path from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked from diffusers_helper.memory import DynamicSwapInstaller from .base_generator import BaseModelGenerator class OriginalModelGenerator(BaseModelGenerator): """ Model generator for the Original HunyuanVideo model. """ def __init__(self, **kwargs): """ Initialize the Original model generator. """ super().__init__(**kwargs) self.model_name = "Original" self.model_path = 'lllyasviel/FramePackI2V_HY' self.model_repo_id_for_cache = "models--lllyasviel--FramePackI2V_HY" def get_model_name(self): """ Get the name of the model. """ return self.model_name def load_model(self): """ Load the Original transformer model. If offline mode is True, attempts to load from a local snapshot. """ print(f"Loading {self.model_name} Transformer...") path_to_load = self.model_path # Initialize with the default path if self.offline: path_to_load = self._get_offline_load_path() # Calls the method in BaseModelGenerator # Create the transformer model self.transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained( path_to_load, torch_dtype=torch.bfloat16 ).cpu() # Configure the model self.transformer.eval() self.transformer.to(dtype=torch.bfloat16) self.transformer.requires_grad_(False) # Set up dynamic swap if not in high VRAM mode if not self.high_vram: DynamicSwapInstaller.install_model(self.transformer, device=self.gpu) else: # In high VRAM mode, move the entire model to GPU self.transformer.to(device=self.gpu) print(f"{self.model_name} Transformer Loaded from {path_to_load}.") return self.transformer def prepare_history_latents(self, height, width): """ Prepare the history latents tensor for the Original model. Args: height: The height of the image width: The width of the image Returns: The initialized history latents tensor """ return torch.zeros( size=(1, 16, 1 + 2 + 16, height // 8, width // 8), dtype=torch.float32 ).cpu() def get_latent_paddings(self, total_latent_sections): """ Get the latent paddings for the Original model. Args: total_latent_sections: The total number of latent sections Returns: A list of latent paddings """ # Original model uses reversed latent paddings if total_latent_sections > 4: return [3] + [2] * (total_latent_sections - 3) + [1, 0] else: return list(reversed(range(total_latent_sections))) def prepare_indices(self, latent_padding_size, latent_window_size): """ Prepare the indices for the Original model. Args: latent_padding_size: The size of the latent padding latent_window_size: The size of the latent window Returns: A tuple of (clean_latent_indices, latent_indices, clean_latent_2x_indices, clean_latent_4x_indices) """ indices = torch.arange(0, sum([1, latent_padding_size, latent_window_size, 1, 2, 16])).unsqueeze(0) clean_latent_indices_pre, blank_indices, latent_indices, clean_latent_indices_post, clean_latent_2x_indices, clean_latent_4x_indices = indices.split([1, latent_padding_size, latent_window_size, 1, 2, 16], dim=1) clean_latent_indices = torch.cat([clean_latent_indices_pre, clean_latent_indices_post], dim=1) return clean_latent_indices, latent_indices, clean_latent_2x_indices, clean_latent_4x_indices def prepare_clean_latents(self, start_latent, history_latents): """ Prepare the clean latents for the Original model. Args: start_latent: The start latent history_latents: The history latents Returns: A tuple of (clean_latents, clean_latents_2x, clean_latents_4x) """ clean_latents_pre = start_latent.to(history_latents) clean_latents_post, clean_latents_2x, clean_latents_4x = history_latents[:, :, :1 + 2 + 16, :, :].split([1, 2, 16], dim=2) clean_latents = torch.cat([clean_latents_pre, clean_latents_post], dim=2) return clean_latents, clean_latents_2x, clean_latents_4x def update_history_latents(self, history_latents, generated_latents): """ Update the history latents with the generated latents for the Original model. Args: history_latents: The history latents generated_latents: The generated latents Returns: The updated history latents """ # For Original model, we prepend the generated latents 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 Original model. Args: history_latents: The history latents total_generated_latent_frames: The total number of generated latent frames Returns: The real history latents """ 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 Original 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 Original model, current_pixels is first, history_pixels is second return soft_append_bcthw(current_pixels, history_pixels, overlapped_frames) def get_section_latent_frames(self, latent_window_size, is_last_section): """ Get the number of section latent frames for the Original model. Args: latent_window_size: The size of the latent window is_last_section: Whether this is the last section Returns: The number of section latent frames """ return latent_window_size * 2 def get_current_pixels(self, real_history_latents, section_latent_frames, vae): """ Get the current pixels for the Original 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 """ from diffusers_helper.hunyuan import vae_decode 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 Original model. Args: total_generated_latent_frames: The total number of generated latent frames current_pos: The current position in seconds original_pos: The original position in seconds current_prompt: The current prompt Returns: The formatted position description """ return (f'Total generated frames: {int(max(0, total_generated_latent_frames * 4 - 3))}, ' f'Video length: {max(0, (total_generated_latent_frames * 4 - 3) / 30):.2f} seconds (FPS-30). ' f'Current position: {current_pos:.2f}s (original: {original_pos:.2f}s). ' f'using prompt: {current_prompt[:256]}...')