import os import sys current_dir = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, current_dir) import torch import traceback import einops import safetensors.torch as sf import numpy as np import argparse import math import time from PIL import Image from diffusers import AutoencoderKLHunyuanVideo from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer from diffusers_helper.hunyuan import encode_prompt_conds, vae_decode, vae_encode, vae_decode_fake from diffusers_helper.utils import save_bcthw_as_mp4, crop_or_pad_yield_mask, soft_append_bcthw, resize_and_center_crop, state_dict_weighted_merge, state_dict_offset_merge, generate_timestamp from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan from diffusers_helper.memory import cpu, gpu, get_cuda_free_memory_gb, move_model_to_device_with_memory_preservation, offload_model_from_device_for_memory_preservation, fake_diffusers_current_device, DynamicSwapInstaller, unload_complete_models, load_model_as_complete from transformers import SiglipImageProcessor, SiglipVisionModel from diffusers_helper.clip_vision import hf_clip_vision_encode from diffusers_helper.bucket_tools import find_nearest_bucket import hashlib import random import string import torchvision from torchvision.transforms.functional import to_pil_image import comfy.utils from PIL import Image import folder_paths class Kiki_FramePack: @classmethod def INPUT_TYPES(s): return { "required": { "ref_image": ("IMAGE", ), "prompt": ("STRING", {"multiline": True}), # "n_prompt": ("STRING", {"multiline": True}), "total_second_length": ("INT", {"default": 5, "min": 1, "max": 120, "step": 1}), "seed": ("INT", {"default": 3407}), "steps": ("INT", {"default": 25, "min": 1, "max": 100, "step": 1}), "use_teacache": ("BOOLEAN", {"default": True}), "upscale": ("FLOAT", {"default": 1.2, "min": 0.1, "max": 2.0, "step": 0.1, "description": "Resolution scaling factor. 1.0 = original size, >1.0 = upscale, <1.0 = downscale"}), }, "optional": { "end_image": ("IMAGE", ), }, } RETURN_TYPES = ("IMAGE", "FLOAT") RETURN_NAMES = ("frames", "fps") CATEGORY = "Runninghub/FramePack" FUNCTION = "run" TITLE = 'RunningHub FramePack' OUTPUT_NODE = True def __init__(self): self.high_vram = False self.frames = None self.fps = None hunyuan_root = os.path.join(folder_paths.models_dir, 'HunyuanVideo') flux_redux_bfl_root = os.path.join(folder_paths.models_dir, 'flux_redux_bfl') framePackI2V_root = os.path.join(folder_paths.models_dir, 'FramePackI2V_HY') self.text_encoder = LlamaModel.from_pretrained(hunyuan_root, subfolder='text_encoder', torch_dtype=torch.float16).cpu() self.text_encoder_2 = CLIPTextModel.from_pretrained(hunyuan_root, subfolder='text_encoder_2', torch_dtype=torch.float16).cpu() self.tokenizer = LlamaTokenizerFast.from_pretrained(hunyuan_root, subfolder='tokenizer') self.tokenizer_2 = CLIPTokenizer.from_pretrained(hunyuan_root, subfolder='tokenizer_2') self.vae = AutoencoderKLHunyuanVideo.from_pretrained(hunyuan_root, subfolder='vae', torch_dtype=torch.float16).cpu() self.feature_extractor = SiglipImageProcessor.from_pretrained(flux_redux_bfl_root, subfolder='feature_extractor') self.image_encoder = SiglipVisionModel.from_pretrained(flux_redux_bfl_root, subfolder='image_encoder', torch_dtype=torch.float16).cpu() self.transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(framePackI2V_root, torch_dtype=torch.bfloat16).cpu() self.vae.eval() self.text_encoder.eval() self.text_encoder_2.eval() self.image_encoder.eval() self.transformer.eval() if not self.high_vram: self.vae.enable_slicing() self.vae.enable_tiling() self.transformer.high_quality_fp32_output_for_inference = True print('transformer.high_quality_fp32_output_for_inference = True') self.transformer.to(dtype=torch.bfloat16) self.vae.to(dtype=torch.float16) self.image_encoder.to(dtype=torch.float16) self.text_encoder.to(dtype=torch.float16) self.text_encoder_2.to(dtype=torch.float16) self.vae.requires_grad_(False) self.text_encoder.requires_grad_(False) self.text_encoder_2.requires_grad_(False) self.image_encoder.requires_grad_(False) self.transformer.requires_grad_(False) if not self.high_vram: # DynamicSwapInstaller is same as huggingface's enable_sequential_offload but 3x faster DynamicSwapInstaller.install_model(self.transformer, device=gpu) DynamicSwapInstaller.install_model(self.text_encoder, device=gpu) def strict_align(self, h, w, scale): raw_h = h * scale raw_w = w * scale aligned_h = int(round(raw_h / 64)) * 64 aligned_w = int(round(raw_w / 64)) * 64 assert (aligned_h % 64 == 0) and (aligned_w % 64 == 0), "尺寸必须是64的倍数" assert (aligned_h//8) % 8 == 0 and (aligned_w//8) % 8 == 0, "潜在空间需要8的倍数" return aligned_h, aligned_w def preprocess_image(self, image): if image is None: return None image_np = 255. * image[0].cpu().numpy() image = Image.fromarray(np.clip(image_np, 0, 255).astype(np.uint8)).convert("RGB") input_image = np.array(image) return input_image def run(self, **kwargs): try: image = kwargs['ref_image'] end_image = kwargs.get('end_image', None) # Use get with None as default image_np = self.preprocess_image(image) end_image_np = self.preprocess_image(end_image) if end_image is not None else None prompt = kwargs['prompt'] seed = kwargs['seed'] total_second_length = kwargs['total_second_length'] steps = kwargs['steps'] use_teacache = kwargs['use_teacache'] upscale = kwargs['upscale'] random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=16)) video_path = os.path.join(folder_paths.get_output_directory(), f'{random_str}.mp4') self.pbar = comfy.utils.ProgressBar(steps * total_second_length) self.exec(input_image=image_np, end_image=end_image_np, prompt=prompt, seed=seed, total_second_length=total_second_length, video_path=video_path, steps=steps, use_teacache=use_teacache, scale=upscale) if os.path.exists(video_path): self.fps = self.get_fps_with_torchvision(video_path) self.frames = self.extract_frames_as_pil(video_path) print(f'{video_path}:{self.fps} {len(self.frames)}') else: self.frames = [] self.fps = 0.0 except Exception as e: print(f"Error in run: {str(e)}") traceback.print_exc() self.frames = [] self.fps = 0.0 return (self.frames, self.fps) @torch.no_grad() def exec(self, input_image, video_path, end_image=None, prompt="The girl dances gracefully, with clear movements, full of charm.", n_prompt="", seed=31337, total_second_length=5, latent_window_size=9, steps=25, cfg=1, gs=32, rs=0, gpu_memory_preservation=6, use_teacache=True, scale=1.0): total_latent_sections = (total_second_length * 30) / (latent_window_size * 4) total_latent_sections = int(max(round(total_latent_sections), 1)) try: # Clean GPU if not self.high_vram: unload_complete_models( self.text_encoder, self.text_encoder_2, self.image_encoder, self.vae, self.transformer ) # Text encoding print('Text encoding') if not self.high_vram: fake_diffusers_current_device(self.text_encoder, gpu) load_model_as_complete(self.text_encoder_2, target_device=gpu) llama_vec, clip_l_pooler = encode_prompt_conds(prompt, self.text_encoder, self.text_encoder_2, self.tokenizer, self.tokenizer_2) if cfg == 1: llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler) else: llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, self.text_encoder, self.text_encoder_2, self.tokenizer, self.tokenizer_2) llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512) llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512) # Processing input image (start frame) print('Processing start frame ...') H, W, C = input_image.shape height, width = find_nearest_bucket(H, W, resolution=640) print(f"Resized height: {height}, Resized width: {width}") height, width = self.strict_align(height, width, scale) print(f"After Resized height: {height}, Resized width: {width}") input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height) input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1 input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None] # Processing end image if provided has_end_image = end_image is not None end_image_np = None end_image_pt = None if has_end_image: print('Processing end frame ...') H_end, W_end, C_end = end_image.shape end_image_np = resize_and_center_crop(end_image, target_width=width, target_height=height) end_image_pt = torch.from_numpy(end_image_np).float() / 127.5 - 1 end_image_pt = end_image_pt.permute(2, 0, 1)[None, :, None] # VAE encoding print('VAE encoding ...') if not self.high_vram: load_model_as_complete(self.vae, target_device=gpu) start_latent = vae_encode(input_image_pt, self.vae) end_latent = None if has_end_image: end_latent = vae_encode(end_image_pt, self.vae) # CLIP Vision print('CLIP Vision encoding ...') if not self.high_vram: load_model_as_complete(self.image_encoder, target_device=gpu) # Start image encoding image_encoder_output = hf_clip_vision_encode(input_image_np, self.feature_extractor, self.image_encoder) image_encoder_last_hidden_state = image_encoder_output.last_hidden_state # End image encoding if available if has_end_image: end_image_encoder_output = hf_clip_vision_encode(end_image_np, self.feature_extractor, self.image_encoder) end_image_encoder_last_hidden_state = end_image_encoder_output.last_hidden_state # Use a simple average of embeddings - exactly like in the original code image_encoder_last_hidden_state = (image_encoder_last_hidden_state + end_image_encoder_last_hidden_state) / 2 # Dtype llama_vec = llama_vec.to(self.transformer.dtype) llama_vec_n = llama_vec_n.to(self.transformer.dtype) clip_l_pooler = clip_l_pooler.to(self.transformer.dtype) clip_l_pooler_n = clip_l_pooler_n.to(self.transformer.dtype) image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(self.transformer.dtype) print('Start Sample') rnd = torch.Generator("cpu").manual_seed(seed) num_frames = latent_window_size * 4 - 3 history_latents = torch.zeros(size=(1, 16, 1 + 2 + 16, height // 8, width // 8), dtype=torch.float32).cpu() history_pixels = None total_generated_latent_frames = 0 latent_paddings = list(reversed(range(total_latent_sections))) if total_latent_sections > 4: latent_paddings = [3] + [2] * (total_latent_sections - 3) + [1, 0] for i, latent_padding in enumerate(latent_paddings): is_last_section = latent_padding == 0 is_first_section = latent_padding == latent_paddings[0] # Use the original method latent_padding_size = latent_padding * latent_window_size print(f'latent_padding_size = {latent_padding_size}, is_last_section = {is_last_section}, is_first_section = {is_first_section}') 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) # Always use start_latent for the first position (exactly like in the original code) clean_latents_pre = start_latent.to(history_latents) # For the second position, use history clean_latents_post, clean_latents_2x, clean_latents_4x = history_latents[:, :, :1 + 2 + 16, :, :].split([1, 2, 16], dim=2) # Create clean_latents first clean_latents = torch.cat([clean_latents_pre, clean_latents_post], dim=2) # Then if we have end_image and this is the first section, override clean_latents_post with end_latent if has_end_image and is_first_section: clean_latents_post = end_latent.to(history_latents) clean_latents = torch.cat([clean_latents_pre, clean_latents_post], dim=2) if not self.high_vram: unload_complete_models() move_model_to_device_with_memory_preservation(self.transformer, target_device=gpu, preserved_memory_gb=gpu_memory_preservation) if use_teacache: self.transformer.initialize_teacache(enable_teacache=True, num_steps=steps) else: self.transformer.initialize_teacache(enable_teacache=False) def callback(d): self.update(1) return generated_latents = sample_hunyuan( transformer=self.transformer, sampler='unipc', width=width, height=height, frames=num_frames, real_guidance_scale=cfg, distilled_guidance_scale=gs, guidance_rescale=rs, num_inference_steps=steps, generator=rnd, prompt_embeds=llama_vec, prompt_embeds_mask=llama_attention_mask, prompt_poolers=clip_l_pooler, negative_prompt_embeds=llama_vec_n, negative_prompt_embeds_mask=llama_attention_mask_n, negative_prompt_poolers=clip_l_pooler_n, device=gpu, dtype=torch.bfloat16, image_embeddings=image_encoder_last_hidden_state, latent_indices=latent_indices, clean_latents=clean_latents, clean_latent_indices=clean_latent_indices, clean_latents_2x=clean_latents_2x, clean_latent_2x_indices=clean_latent_2x_indices, clean_latents_4x=clean_latents_4x, clean_latent_4x_indices=clean_latent_4x_indices, callback=callback, ) # For the last section, add start_latent back to the beginning - just like in the original if is_last_section: generated_latents = torch.cat([start_latent.to(generated_latents), generated_latents], dim=2) # Accumulate generated frames total_generated_latent_frames += int(generated_latents.shape[2]) history_latents = torch.cat([generated_latents.to(history_latents), history_latents], dim=2) if not self.high_vram: offload_model_from_device_for_memory_preservation(self.transformer, target_device=gpu, preserved_memory_gb=8) load_model_as_complete(self.vae, target_device=gpu) # Only decode up to the total number of frames we've generated real_history_latents = history_latents[:, :, :total_generated_latent_frames, :, :] # Decode latents to pixels if history_pixels is None: history_pixels = vae_decode(real_history_latents, self.vae).cpu() else: # For appending new frames to existing ones section_latent_frames = (latent_window_size * 2 + 1) if is_last_section else (latent_window_size * 2) overlapped_frames = latent_window_size * 4 - 3 current_pixels = vae_decode(real_history_latents[:, :, :section_latent_frames], self.vae).cpu() history_pixels = soft_append_bcthw(current_pixels, history_pixels, overlapped_frames) if not self.high_vram: unload_complete_models() # If this is the last section, save the video if is_last_section: save_bcthw_as_mp4(history_pixels, video_path, fps=30) break except Exception as e: print(f"Error in exec: {str(e)}") traceback.print_exc() finally: unload_complete_models() def update(self, in_progress): self.pbar.update(in_progress) def extract_frames_as_pil(self, video_path): video, _, _ = torchvision.io.read_video(video_path, pts_unit='sec') # (T, H, W, C) frames = [to_pil_image(frame.permute(2, 0, 1)) for frame in video] frames = [torch.from_numpy(np.array(frame).astype(np.float32) / 255.0) for frame in frames] return frames def get_fps_with_torchvision(self, video_path): _, _, info = torchvision.io.read_video(video_path, pts_unit='sec') return info['video_fps'] # --- Start of Kiki_FramePack_F1 Class --- class Kiki_FramePack_F1: @classmethod def INPUT_TYPES(s): return { "required": { "ref_image": ("IMAGE", ), "prompt": ("STRING", {"multiline": True}), "total_second_length": ("INT", {"default": 5, "min": 1, "max": 120, "step": 1}), "fps": ("INT", {"default": 30, "min": 1, "max": 60, "step": 1}), "seed": ("INT", {"default": 3407}), "steps": ("INT", {"default": 25, "min": 1, "max": 100, "step": 1}), "gs": ("FLOAT", {"default": 10.0, "min": 1.0, "max": 32.0, "step": 0.1, "round": 0.01, "label": "Distilled CFG Scale"}), "use_teacache": ("BOOLEAN", {"default": True}), "upscale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 2.0, "step": 0.1, "description": "Resolution scaling factor."}), }, "optional": { "n_prompt": ("STRING", {"multiline": True, "default": ""}), } } RETURN_TYPES = ("IMAGE", "FLOAT") RETURN_NAMES = ("frames", "fps") CATEGORY = "Runninghub/FramePack" FUNCTION = "run_f1" TITLE = 'RunningHub FramePack F1' OUTPUT_NODE = True def __init__(self): self.high_vram = False self.frames = None self.fps = None hunyuan_root = os.path.join(folder_paths.models_dir, 'HunyuanVideo') flux_redux_bfl_root = os.path.join(folder_paths.models_dir, 'flux_redux_bfl') framePackF1_root = os.path.join(folder_paths.models_dir, 'FramePackF1_HY') if not os.path.isdir(framePackF1_root): print(f"Warning: FramePack F1 model directory not found at {framePackF1_root}") self.text_encoder = LlamaModel.from_pretrained(hunyuan_root, subfolder='text_encoder', torch_dtype=torch.float16).cpu() self.text_encoder_2 = CLIPTextModel.from_pretrained(hunyuan_root, subfolder='text_encoder_2', torch_dtype=torch.float16).cpu() self.tokenizer = LlamaTokenizerFast.from_pretrained(hunyuan_root, subfolder='tokenizer') self.tokenizer_2 = CLIPTokenizer.from_pretrained(hunyuan_root, subfolder='tokenizer_2') self.vae = AutoencoderKLHunyuanVideo.from_pretrained(hunyuan_root, subfolder='vae', torch_dtype=torch.float16).cpu() self.feature_extractor = SiglipImageProcessor.from_pretrained(flux_redux_bfl_root, subfolder='feature_extractor') self.image_encoder = SiglipVisionModel.from_pretrained(flux_redux_bfl_root, subfolder='image_encoder', torch_dtype=torch.float16).cpu() try: self.transformer_f1 = HunyuanVideoTransformer3DModelPacked.from_pretrained(framePackF1_root, torch_dtype=torch.bfloat16).cpu() except Exception as e: print(f"Error loading FramePack F1 transformer model from {framePackF1_root}: {e}") print("Please ensure the F1 model weights (e.g., transformer.safetensors) are correctly placed in the directory.") self.transformer_f1 = None self.vae.eval() self.text_encoder.eval() self.text_encoder_2.eval() self.image_encoder.eval() if self.transformer_f1: self.transformer_f1.eval() if not self.high_vram: self.vae.enable_slicing() self.vae.enable_tiling() if self.transformer_f1: self.transformer_f1.high_quality_fp32_output_for_inference = True print('F1 transformer.high_quality_fp32_output_for_inference = True') self.transformer_f1.to(dtype=torch.bfloat16) self.transformer_f1.requires_grad_(False) if not self.high_vram: DynamicSwapInstaller.install_model(self.transformer_f1, device=gpu) self.vae.to(dtype=torch.float16) self.image_encoder.to(dtype=torch.float16) self.text_encoder.to(dtype=torch.float16) self.text_encoder_2.to(dtype=torch.float16) self.vae.requires_grad_(False) self.text_encoder.requires_grad_(False) self.text_encoder_2.requires_grad_(False) self.image_encoder.requires_grad_(False) if not self.high_vram: DynamicSwapInstaller.install_model(self.text_encoder, device=gpu) def strict_align(self, h, w, scale): raw_h = h * scale raw_w = w * scale aligned_h = int(round(raw_h / 64)) * 64 aligned_w = int(round(raw_w / 64)) * 64 assert (aligned_h % 64 == 0) and (aligned_w % 64 == 0), "尺寸必须是64的倍数" assert (aligned_h//8) % 8 == 0 and (aligned_w//8) % 8 == 0, "潜在空间需要8的倍数" return aligned_h, aligned_w def preprocess_image(self, image): if image is None: return None if image.dim() == 4 and image.shape[0] == 1: img_tensor = image[0] else: img_tensor = image print(f"Warning: Unexpected input image tensor shape: {image.shape}. Assuming HWC.") image_np = 255. * img_tensor.cpu().numpy() image = Image.fromarray(np.clip(image_np, 0, 255).astype(np.uint8)).convert("RGB") input_image = np.array(image) return input_image def run_f1(self, **kwargs): if not self.transformer_f1: print("Error: Kiki_FramePack_F1 cannot run because the transformer model failed to load.") return (torch.empty((0, 1, 1, 3), dtype=torch.float32), 0.0) try: image = kwargs['ref_image'] image_np = self.preprocess_image(image) prompt = kwargs['prompt'] n_prompt = kwargs.get('n_prompt', "") seed = kwargs['seed'] total_second_length = kwargs['total_second_length'] fps = kwargs['fps'] steps = kwargs['steps'] gs = kwargs['gs'] use_teacache = kwargs['use_teacache'] upscale = kwargs['upscale'] cfg = 1.0 rs = 0.0 latent_window_size = 9 random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=16)) video_path = os.path.join(folder_paths.get_output_directory(), f'{random_str}_f1.mp4') # --- Initialize Progress Bar (Aligned with demo's section calc) --- # Use demo's calculation for total_latent_sections, assuming 30fps basis for consistency total_latent_sections = int(max(round((total_second_length * 30) / (latent_window_size * 4)), 1)) total_progress_steps = total_latent_sections * steps self.pbar = comfy.utils.ProgressBar(total_progress_steps) # Call exec_f1, passing latent_window_size as well self.exec_f1(input_image=image_np, prompt=prompt, n_prompt=n_prompt, seed=seed, total_second_length=total_second_length, video_path=video_path, fps=fps, steps=steps, gs=gs, cfg=cfg, rs=rs, latent_window_size=latent_window_size, # Pass latent_window_size use_teacache=use_teacache, scale=upscale, gpu_memory_preservation=6) if os.path.exists(video_path): self.fps = float(fps) self.frames = self.extract_frames_to_tensor(video_path) print(f'F1 Video saved: {video_path} | FPS: {self.fps} | Frames: {self.frames.shape[0] if self.frames is not None else 0}') else: self.frames = torch.empty((0, 1, 1, 3), dtype=torch.float32) self.fps = 0.0 print(f'F1 Video generation failed or file not found: {video_path}') except Exception as e: print(f"Error in run_f1: {str(e)}") traceback.print_exc() self.frames = torch.empty((0, 1, 1, 3), dtype=torch.float32) self.fps = 0.0 return (self.frames, self.fps) @torch.no_grad() def exec_f1(self, input_image, video_path, prompt, n_prompt, seed, total_second_length, fps, steps, gs, cfg, rs, latent_window_size, # Receive latent_window_size use_teacache, scale, gpu_memory_preservation=6): print("--- Starting Kiki_FramePack_F1 exec_f1 (Aligned with Demo Logic) ---") print(f"Params: seed={seed}, length={total_second_length}s@{fps}fps, steps={steps}, gs={gs}, cfg={cfg}, rs={rs}, lws={latent_window_size}") vae_time_stride = 4 # --- Use Demo's total_latent_sections calculation --- total_latent_sections = int(max(round((total_second_length * 30) / (latent_window_size * 4)), 1)) print(f"Total generation sections (Demo calc): {total_latent_sections}") # --- Calculate target frames needed (still useful for trimming) --- target_pixel_frames = int(round(total_second_length * fps)) try: # --- 1. Initialization & Setup --- torch.manual_seed(seed) rnd = torch.Generator("cpu").manual_seed(seed) # ... (Unload models if needed) ... # --- 2. Encoding Inputs --- print('Encoding text prompts...') if not self.high_vram: fake_diffusers_current_device(self.text_encoder, gpu) load_model_as_complete(self.text_encoder_2, target_device=gpu) llama_vec, clip_l_pooler = encode_prompt_conds(prompt, self.text_encoder, self.text_encoder_2, self.tokenizer, self.tokenizer_2) llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, self.text_encoder, self.text_encoder_2, self.tokenizer, self.tokenizer_2) llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512) llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512) print('Processing reference image...') H, W, C = input_image.shape if scale == 1.0: height, width = find_nearest_bucket(H, W, resolution=640) height, width = self.strict_align(height, width, 1.0) else: height, width = self.strict_align(H, W, scale) print(f"Target dimensions: {width}x{height}") input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height) input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1 input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None] print('VAE encoding reference image...') if not self.high_vram: load_model_as_complete(self.vae, target_device=gpu) start_latent = vae_encode(input_image_pt.to(self.vae.device, dtype=self.vae.dtype), self.vae) print(f"Start latent shape: {start_latent.shape}") print('CLIP Vision encoding reference image...') if not self.high_vram: load_model_as_complete(self.image_encoder, target_device=gpu) image_encoder_output = hf_clip_vision_encode(input_image_np, self.feature_extractor, self.image_encoder.to(gpu)) image_embeddings = image_encoder_output.last_hidden_state transformer_dtype = self.transformer_f1.dtype start_latent = start_latent.to(transformer_dtype).cpu() # --- 3. Diffusion Loop (Aligned with Demo) --- print(f'Starting diffusion loop for {total_latent_sections} sections...') latent_channels = start_latent.shape[1] latent_height = start_latent.shape[-2] latent_width = start_latent.shape[-1] history_context_size = 16 + 2 + 1 # --- Initialize history_latents like demo --- # Start with zeros matching context size history_latents = torch.zeros(size=(1, latent_channels, history_context_size, latent_height, latent_width), dtype=torch.float32).cpu() # Use float32 like demo? # Immediately add start_latent history_latents = torch.cat([history_latents, start_latent.to(history_latents.dtype)], dim=2) total_generated_latent_frames = 1 # Account for start_latent history_pixels = None # ... (Progress bar callback setup) ... current_section_step = 0 total_progress_steps = total_latent_sections * steps def callback_f1(d): # ... (Update pbar logic remains the same) ... nonlocal current_section_step step_in_section = d['i'] current_total_step = current_section_step * steps + step_in_section + 1 if hasattr(self, 'pbar') and self.pbar: self.pbar.update_absolute(current_total_step, total_progress_steps) # Calculate frames generated per step based on demo frames_per_latent_window = latent_window_size * 4 - 3 for section_index in range(total_latent_sections): section_start_time = time.time() print(f'Generating section {section_index + 1} / {total_latent_sections}') current_section_step = section_index # ... (Load transformer if needed) ... # --- Prepare context and indices (same as before, uses history_latents) --- indices = torch.arange(0, sum([1, 16, 2, 1, latent_window_size])).unsqueeze(0) clean_latent_indices_start, clean_latent_4x_indices, clean_latent_2x_indices, clean_latent_1x_indices, latent_indices = indices.split([1, 16, 2, 1, latent_window_size], dim=1) clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1) # Get history context from the *end* of the current history_latents # No padding needed here because history starts with context + start_latent history_context = history_latents[:, :, -history_context_size:, :, :] clean_latents_4x, clean_latents_2x, clean_latents_1x = history_context.split([16, 2, 1], dim=2) clean_latents = torch.cat([start_latent.cpu(), clean_latents_1x.cpu()], dim=2) # --- Prepare sample_kwargs (same as before) --- sample_kwargs = dict( transformer=self.transformer_f1, sampler='unipc', width=width, height=height, frames=frames_per_latent_window, # Use demo's frame count real_guidance_scale=cfg, distilled_guidance_scale=gs, guidance_rescale=rs, num_inference_steps=steps, generator=rnd, # --- Add missing positive prompt embeddings & ENSURE DTYPE --- prompt_embeds=llama_vec.to(gpu, dtype=transformer_dtype), prompt_embeds_mask=llama_attention_mask.to(gpu), # Mask dtype usually okay # --- Existing embeddings/poolers & ENSURE DTYPE --- prompt_poolers=clip_l_pooler.to(gpu, dtype=transformer_dtype), negative_prompt_embeds=llama_vec_n.to(gpu, dtype=transformer_dtype), negative_prompt_embeds_mask=llama_attention_mask_n.to(gpu), # Mask dtype usually okay negative_prompt_poolers=clip_l_pooler_n.to(gpu, dtype=transformer_dtype), device=gpu, # Device is already GPU dtype=transformer_dtype, # Explicitly passing transformer's dtype image_embeddings=image_embeddings.to(gpu, dtype=transformer_dtype), latent_indices=latent_indices.to(gpu), # Indices dtype usually okay clean_latents=clean_latents.to(gpu, dtype=transformer_dtype), # Ensure correct dtype clean_latent_indices=clean_latent_indices.to(gpu), # Indices dtype usually okay clean_latents_2x=clean_latents_2x.to(gpu, dtype=transformer_dtype), # Ensure correct dtype clean_latent_2x_indices=clean_latent_2x_indices.to(gpu), # Indices dtype usually okay clean_latents_4x=clean_latents_4x.to(gpu, dtype=transformer_dtype), # Ensure correct dtype clean_latent_4x_indices=clean_latent_4x_indices.to(gpu), # Indices dtype usually okay callback=callback_f1, ) # ... (Initialize teacache) ... if hasattr(self.transformer_f1, 'initialize_teacache'): self.transformer_f1.initialize_teacache(enable_teacache=use_teacache, num_steps=steps) # --- Call sample_hunyuan --- generated_latents = sample_hunyuan(**sample_kwargs) generated_latents = generated_latents.to(cpu, dtype=torch.float32) print(f" Sampled latent section shape: {generated_latents.shape}") # --- Update history_latents (Aligned with Demo: Always append) --- total_generated_latent_frames += int(generated_latents.shape[2]) history_latents = torch.cat([history_latents, generated_latents.to(history_latents.dtype)], dim=2) # --- Decode and append pixels (Aligned with Demo) --- if not self.high_vram: offload_model_from_device_for_memory_preservation(self.transformer_f1, target_device=gpu, preserved_memory_gb=8) load_model_as_complete(self.vae, target_device=gpu) else: if self.vae.device != gpu: self.vae.to(gpu) # Calculate the slice of history to decode based on total generated frames real_history_latents = history_latents[:, :, -total_generated_latent_frames:, :, :] # Use actual generated frames if history_pixels is None: # First time: decode the current relevant history history_pixels = vae_decode(real_history_latents.to(gpu, dtype=self.vae.dtype), self.vae).cpu() print(f" Decoded initial pixels. Shape: {history_pixels.shape}") else: # Subsequent times: decode only the part needed for smooth append section_latent_frames = latent_window_size * 2 overlapped_frames = latent_window_size * 4 - 3 # Use demo's overlap calculation # Decode the relevant tail end of the history latents current_latents_to_decode = real_history_latents[:, :, -section_latent_frames:, :, :] current_pixels = vae_decode(current_latents_to_decode.to(gpu, dtype=self.vae.dtype), self.vae).cpu() # Append smoothly using demo's overlap value history_pixels = soft_append_bcthw(history_pixels, current_pixels, overlapped_frames) print(f" Appended pixels. New history shape: {history_pixels.shape}") # ... (Unload VAE if needed) ... if not self.high_vram: unload_complete_models(self.vae) section_end_time = time.time() print(f" Section {section_index + 1} took {section_end_time - section_start_time:.2f} seconds.") # --- 4. Final Saving (Aligned with Demo, keeping variable fps) --- print('Saving final video...') if history_pixels is None or history_pixels.shape[2] == 0: raise ValueError("No pixel frames were generated or decoded.") if history_pixels.shape[2] > target_pixel_frames: print(f"Trimming final video from {history_pixels.shape[2]} to {target_pixel_frames} frames.") history_pixels = history_pixels[:,:,:target_pixel_frames,:,:] save_bcthw_as_mp4( history_pixels, video_path, fps=fps, # Keep user FPS for now # crf=18 # Omit crf until utils.py is confirmed synced ) print(f"Final video saved to: {video_path}") except Exception as e: print(f"Error during Kiki_FramePack_F1 execution: {str(e)}") traceback.print_exc() if os.path.exists(video_path): try: os.remove(video_path) except OSError: pass if hasattr(self, 'pbar') and self.pbar: self.pbar.update_absolute(total_progress_steps, total_progress_steps) raise finally: print('Cleaning up models...') unload_complete_models( self.text_encoder, self.text_encoder_2, self.image_encoder, self.vae, self.transformer_f1 ) torch.cuda.empty_cache() print("--- Finished Kiki_FramePack_F1 exec_f1 (Aligned with Demo Logic) ---") def extract_frames_to_tensor(self, video_path): try: video_tensor, _, metadata = torchvision.io.read_video(video_path, pts_unit='sec', output_format='TCHW') video_tensor = video_tensor.permute(0, 2, 3, 1) video_tensor = video_tensor.float() / 255.0 print(f"Extracted video tensor shape: {video_tensor.shape}") return video_tensor except Exception as e: print(f"Error extracting frames using torchvision.io.read_video: {e}") traceback.print_exc() return torch.empty((0, 1, 1, 3), dtype=torch.float32) def get_fps_with_torchvision(self, video_path): try: _, _, metadata = torchvision.io.read_video(video_path, pts_unit='sec') fps = metadata.get('video_fps', 30.0) return float(fps) except Exception as e: print(f"Error reading FPS using torchvision.io.read_video: {e}") traceback.print_exc() return 30.0 # NODE CLASS MAPPINGS NODE_CLASS_MAPPINGS = { "RunningHub_FramePack": Kiki_FramePack, "RunningHub_FramePack_F1": Kiki_FramePack_F1 } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "RunningHub_FramePack": Kiki_FramePack.TITLE, "RunningHub_FramePack_F1": Kiki_FramePack_F1.TITLE }