358 lines
16 KiB
Python
358 lines
16 KiB
Python
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
|
|
|
|
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}),
|
|
},
|
|
}
|
|
|
|
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
|
|
|
|
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 preprocess_image(self, image):
|
|
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):
|
|
image = kwargs['ref_image']
|
|
image_np = self.preprocess_image(image)
|
|
prompt = kwargs['prompt']
|
|
seed = kwargs['seed']
|
|
total_second_length = kwargs['total_second_length']
|
|
steps = kwargs['steps']
|
|
use_teacache = kwargs['use_teacache']
|
|
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, prompt=prompt, seed=seed, total_second_length=total_second_length, video_path=video_path, steps=steps, use_teacache=use_teacache)
|
|
if os.path.exists(video_path):
|
|
fps = self.get_fps_with_torchvision(video_path)
|
|
frames = self.extract_frames_as_pil(video_path)
|
|
# os.remove(video_path)
|
|
print(f'{video_path}:{fps} {len(frames)}')
|
|
|
|
return (frames, fps)
|
|
|
|
@torch.no_grad()
|
|
def exec(self, input_image, video_path,
|
|
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):
|
|
|
|
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) # since we only encode one text - that is one model move and one encode, offload is same time consumption since it is also one load and one encode.
|
|
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
|
|
|
|
print('Image processing ...')
|
|
|
|
H, W, C = input_image.shape
|
|
height, width = find_nearest_bucket(H, W, resolution=640)
|
|
print(f"Resized height: {height}, Resized width: {width}") # Print resized dimensions
|
|
scale = 1.2
|
|
def strict_align(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
|
|
height, width = strict_align(height, width, scale)
|
|
print(f"After Resized height: {height}, Resized width: {width}") # Print resized dimensions
|
|
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]
|
|
|
|
# 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)
|
|
|
|
# CLIP Vision
|
|
|
|
print('CLIP Vision encoding ...')
|
|
|
|
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)
|
|
image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
|
|
|
|
# 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 = reversed(range(total_latent_sections))
|
|
|
|
if total_latent_sections > 4:
|
|
# In theory the latent_paddings should follow the above sequence, but it seems that duplicating some
|
|
# items looks better than expanding it when total_latent_sections > 4
|
|
# One can try to remove below trick and just
|
|
# use `latent_paddings = list(reversed(range(total_latent_sections)))` to compare
|
|
latent_paddings = [3] + [2] * (total_latent_sections - 3) + [1, 0]
|
|
|
|
for latent_padding in latent_paddings:
|
|
is_last_section = latent_padding == 0
|
|
latent_padding_size = latent_padding * latent_window_size
|
|
|
|
# if stream.input_queue.top() == 'end':
|
|
# stream.output_queue.push(('end', None))
|
|
# return
|
|
|
|
print(f'latent_padding_size = {latent_padding_size}, is_last_section = {is_last_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)
|
|
|
|
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)
|
|
|
|
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,
|
|
# shift=3.0,
|
|
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,
|
|
)
|
|
|
|
if is_last_section:
|
|
generated_latents = torch.cat([start_latent.to(generated_latents), generated_latents], dim=2)
|
|
|
|
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)
|
|
|
|
real_history_latents = history_latents[:, :, :total_generated_latent_frames, :, :]
|
|
|
|
if history_pixels is None:
|
|
history_pixels = vae_decode(real_history_latents, self.vae).cpu()
|
|
else:
|
|
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 is_last_section:
|
|
save_bcthw_as_mp4(history_pixels, video_path, fps=30)
|
|
|
|
# print('video fps:', get_fps_with_torchvision('final.mp4'))
|
|
# print(f'Decoded. Current latent shape {real_history_latents.shape}; pixel shape {history_pixels.shape}')
|
|
except Exception as e:
|
|
print(e)
|
|
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']
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RunningHub_FramePack": Kiki_FramePack,
|
|
}
|