diff --git a/nodes_one.py b/nodes_one.py index fe0e157..ccdd429 100644 --- a/nodes_one.py +++ b/nodes_one.py @@ -1,926 +1,22 @@ -import os -import torch -import math -from tqdm import tqdm -import sys import logging -from pathlib import Path +import math + +import torch logger = logging.getLogger(__name__) -from accelerate import init_empty_weights -from accelerate.utils import set_module_tensor_to_device - -import folder_paths -import comfy.model_management as mm -from comfy.utils import load_torch_file, ProgressBar, common_upscale -import comfy.model_base import comfy.latent_formats -from comfy.cli_args import args, LatentPreviewMethod +import comfy.model_base +import comfy.model_management as mm +from comfy.utils import common_upscale -from .utils import log - -script_directory = os.path.dirname(os.path.abspath(__file__)) -vae_scaling_factor = 0.476986 - -from .diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModel -from .diffusers_helper.memory import DynamicSwapInstaller, move_model_to_device_with_memory_preservation +from .diffusers_helper.memory import move_model_to_device_with_memory_preservation from .diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan from .diffusers_helper.utils import crop_or_pad_yield_mask -from .diffusers_helper.bucket_tools import find_nearest_bucket +from .nodes import HyVideoModel, HyVideoModelConfig -# Import original function for fallback -from diffusers.loaders.lora_conversion_utils import _convert_hunyuan_video_lora_to_diffusers +vae_scaling_factor = 0.476986 -def patched_convert_hunyuan_video_lora(original_state_dict): - """Patched version that handles problematic tensors during conversion""" - try: - # Filter out problematic tensors first - filtered_state_dict = {} - for key, value in original_state_dict.items(): - if isinstance(value, torch.Tensor): - if value.dim() == 0: - print(f"Skipping 0-dimensional tensor: {key}") - continue - filtered_state_dict[key] = value - else: - print(f"Skipping non-tensor value: {key}") - - print(f"After filtering: {len(filtered_state_dict)} valid keys") - - # First try the original conversion - try: - from diffusers.loaders.lora_conversion_utils import _convert_hunyuan_video_lora_to_diffusers - result = _convert_hunyuan_video_lora_to_diffusers(filtered_state_dict) - print("Successfully converted LoRA weights") - return result - except Exception as e: - print(f"Error in standard conversion: {e}") - print("Falling back to custom conversion") - - # If standard conversion fails, use the custom implementation - # [Insert the custom conversion code here that was previously unreachable] - # Include all the remapper functions and conversion logic - - # Return the result of the custom conversion - return converted_state_dict - - except Exception as e: - print(f"LoRA conversion failed: {str(e)}") - # Return empty state dict as fallback - return {} - - except Exception as e: - print(f"LoRA conversion failed: {str(e)}") - # Return empty state dict as fallback - return {} - - def remap_img_attn_qkv_(key, state_dict): - try: - weight = state_dict.pop(key) - - # Add dimension check - if weight.dim() == 0: - logger.warning(f"Invalid tensor dimensions for {key}: scalar tensor. Skipping.") - return - - if "lora_A" in key: - state_dict[key.replace("img_attn_qkv", "attn.to_q")] = weight - state_dict[key.replace("img_attn_qkv", "attn.to_k")] = weight - state_dict[key.replace("img_attn_qkv", "attn.to_v")] = weight - else: - # Ensure tensor is properly sized before chunking - if weight.dim() == 0 or weight.size(0) < 3: - logger.warning(f"Invalid tensor size for {key}: {weight.shape}. Using equal splits.") - # Create minimal placeholders - if weight.dim() > 0 and weight.size(0) > 0: - to_q = weight[:1] - to_k = weight[:1] if weight.size(0) == 1 else weight[1:2] - to_v = weight[:1] if weight.size(0) <= 2 else weight[2:3] - else: - # For zero-dim tensors, create basic ones - to_q = torch.ones(1, dtype=weight.dtype, device=weight.device) - to_k = torch.ones(1, dtype=weight.dtype, device=weight.device) - to_v = torch.ones(1, dtype=weight.dtype, device=weight.device) - else: - to_q, to_k, to_v = weight.chunk(3, dim=0) - - state_dict[key.replace("img_attn_qkv", "attn.to_q")] = to_q - state_dict[key.replace("img_attn_qkv", "attn.to_k")] = to_k - state_dict[key.replace("img_attn_qkv", "attn.to_v")] = to_v - except Exception as e: - logger.warning(f"Error processing key {key}: {str(e)}. Skipping.") - # Just skip the problematic key - - def remap_txt_attn_qkv_(key, state_dict): - try: - weight = state_dict.pop(key) - - # Add dimension check - if weight.dim() == 0: - logger.warning(f"Invalid tensor dimensions for {key}: scalar tensor. Skipping.") - return - - if "lora_A" in key: - state_dict[key.replace("txt_attn_qkv", "attn.add_q_proj")] = weight - state_dict[key.replace("txt_attn_qkv", "attn.add_k_proj")] = weight - state_dict[key.replace("txt_attn_qkv", "attn.add_v_proj")] = weight - else: - # Ensure tensor is properly sized before chunking - if weight.dim() == 0 or weight.size(0) < 3: - logger.warning(f"Invalid tensor size for {key}: {weight.shape}. Using equal splits.") - # Create minimal placeholders - if weight.dim() > 0 and weight.size(0) > 0: - to_q = weight[:1] - to_k = weight[:1] if weight.size(0) == 1 else weight[1:2] - to_v = weight[:1] if weight.size(0) <= 2 else weight[2:3] - else: - # For zero-dim tensors, create basic ones - to_q = torch.ones(1, dtype=weight.dtype, device=weight.device) - to_k = torch.ones(1, dtype=weight.dtype, device=weight.device) - to_v = torch.ones(1, dtype=weight.dtype, device=weight.device) - else: - to_q, to_k, to_v = weight.chunk(3, dim=0) - - state_dict[key.replace("txt_attn_qkv", "attn.add_q_proj")] = to_q - state_dict[key.replace("txt_attn_qkv", "attn.add_k_proj")] = to_k - state_dict[key.replace("txt_attn_qkv", "attn.add_v_proj")] = to_v - except Exception as e: - logger.warning(f"Error processing key {key}: {str(e)}. Skipping.") - # Just skip the problematic key - - def remap_txt_in_(key, state_dict): - def rename_key(key): - new_key = key.replace("individual_token_refiner.blocks", "token_refiner.refiner_blocks") - new_key = new_key.replace("adaLN_modulation.1", "norm_out.linear") - new_key = new_key.replace("txt_in", "context_embedder") - new_key = new_key.replace("t_embedder.mlp.0", "time_text_embed.timestep_embedder.linear_1") - new_key = new_key.replace("t_embedder.mlp.2", "time_text_embed.timestep_embedder.linear_2") - new_key = new_key.replace("c_embedder", "time_text_embed.text_embedder") - new_key = new_key.replace("mlp", "ff") - return new_key - - try: - if "self_attn_qkv" in key: - weight = state_dict.pop(key) - # Ensure tensor is at least 1D before chunking - if weight.dim() == 0 or weight.size(0) < 3: - logger.warning(f"Invalid tensor dimensions for {key}: {weight.shape}. Skipping.") - return - - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_q"))] = to_q - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_k"))] = to_k - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_v"))] = to_v - else: - state_dict[rename_key(key)] = state_dict.pop(key) - except Exception as e: - logger.warning(f"Error processing key {key}: {str(e)}") - # Skip if we can't process this key properly - if key in state_dict: - state_dict.pop(key) - - def remap_single_transformer_blocks_(key, state_dict): - try: - hidden_size = 3072 - - if "linear1.lora_A.weight" in key or "linear1.lora_B.weight" in key: - linear1_weight = state_dict.pop(key) - if "lora_A" in key: - new_key = key.replace("single_blocks", "single_transformer_blocks") - if new_key.endswith(".linear1.lora_A.weight"): - new_key = new_key[:-len(".linear1.lora_A.weight")] - state_dict[f"{new_key}.attn.to_q.lora_A.weight"] = linear1_weight - state_dict[f"{new_key}.attn.to_k.lora_A.weight"] = linear1_weight - state_dict[f"{new_key}.attn.to_v.lora_A.weight"] = linear1_weight - state_dict[f"{new_key}.proj_mlp.lora_A.weight"] = linear1_weight - else: - # Ensure tensor size is sufficient for splitting - if linear1_weight.dim() == 0 or linear1_weight.size(0) < 3 * hidden_size: - logger.warning(f"Invalid tensor size for {key}: {linear1_weight.shape}. Skipping splitting.") - return - - split_size = (hidden_size, hidden_size, hidden_size, linear1_weight.size(0) - 3 * hidden_size) - q, k, v, mlp = torch.split(linear1_weight, split_size, dim=0) - new_key = key.replace("single_blocks", "single_transformer_blocks") - if new_key.endswith(".linear1.lora_B.weight"): - new_key = new_key[:-len(".linear1.lora_B.weight")] - state_dict[f"{new_key}.attn.to_q.lora_B.weight"] = q - state_dict[f"{new_key}.attn.to_k.lora_B.weight"] = k - state_dict[f"{new_key}.attn.to_v.lora_B.weight"] = v - state_dict[f"{new_key}.proj_mlp.lora_B.weight"] = mlp - - elif "linear1.lora_A.bias" in key or "linear1.lora_B.bias" in key: - linear1_bias = state_dict.pop(key) - if "lora_A" in key: - new_key = key.replace("single_blocks", "single_transformer_blocks") - if new_key.endswith(".linear1.lora_A.bias"): - new_key = new_key[:-len(".linear1.lora_A.bias")] - state_dict[f"{new_key}.attn.to_q.lora_A.bias"] = linear1_bias - state_dict[f"{new_key}.attn.to_k.lora_A.bias"] = linear1_bias - state_dict[f"{new_key}.attn.to_v.lora_A.bias"] = linear1_bias - state_dict[f"{new_key}.proj_mlp.lora_A.bias"] = linear1_bias - else: - # Ensure tensor size is sufficient for splitting - if linear1_bias.dim() == 0 or linear1_bias.size(0) < 3 * hidden_size: - logger.warning(f"Invalid tensor size for {key}: {linear1_bias.shape}. Skipping splitting.") - return - - split_size = (hidden_size, hidden_size, hidden_size, linear1_bias.size(0) - 3 * hidden_size) - q_bias, k_bias, v_bias, mlp_bias = torch.split(linear1_bias, split_size, dim=0) - new_key = key.replace("single_blocks", "single_transformer_blocks") - if new_key.endswith(".linear1.lora_B.bias"): - new_key = new_key[:-len(".linear1.lora_B.bias")] - state_dict[f"{new_key}.attn.to_q.lora_B.bias"] = q_bias - state_dict[f"{new_key}.attn.to_k.lora_B.bias"] = k_bias - state_dict[f"{new_key}.attn.to_v.lora_B.bias"] = v_bias - state_dict[f"{new_key}.proj_mlp.lora_B.bias"] = mlp_bias - - else: - new_key = key.replace("single_blocks", "single_transformer_blocks") - new_key = new_key.replace("linear2", "proj_out") - new_key = new_key.replace("q_norm", "attn.norm_q") - new_key = new_key.replace("k_norm", "attn.norm_k") - state_dict[new_key] = state_dict.pop(key) - except Exception as e: - logger.warning(f"Error processing key {key}: {str(e)}") - # Skip if we can't process this key properly - if key in state_dict: - state_dict.pop(key) - - TRANSFORMER_KEYS_RENAME_DICT = { - "img_in": "x_embedder", - "time_in.mlp.0": "time_text_embed.timestep_embedder.linear_1", - "time_in.mlp.2": "time_text_embed.timestep_embedder.linear_2", - "guidance_in.mlp.0": "time_text_embed.guidance_embedder.linear_1", - "guidance_in.mlp.2": "time_text_embed.guidance_embedder.linear_2", - "vector_in.in_layer": "time_text_embed.text_embedder.linear_1", - "vector_in.out_layer": "time_text_embed.text_embedder.linear_2", - "double_blocks": "transformer_blocks", - "img_attn_q_norm": "attn.norm_q", - "img_attn_k_norm": "attn.norm_k", - "img_attn_proj": "attn.to_out.0", - "txt_attn_q_norm": "attn.norm_added_q", - "txt_attn_k_norm": "attn.norm_added_k", - "txt_attn_proj": "attn.to_add_out", - "img_mod.linear": "norm1.linear", - "img_norm1": "norm1.norm", - "img_norm2": "norm2", - "img_mlp": "ff", - "txt_mod.linear": "norm1_context.linear", - "txt_norm1": "norm1.norm", - "txt_norm2": "norm2_context", - "txt_mlp": "ff_context", - "self_attn_proj": "attn.to_out.0", - "modulation.linear": "norm.linear", - "pre_norm": "norm.norm", - "final_layer.norm_final": "norm_out.norm", - "final_layer.linear": "proj_out", - "fc1": "net.0.proj", - "fc2": "net.2", - "input_embedder": "proj_in", - } - - TRANSFORMER_SPECIAL_KEYS_REMAP = { - "txt_in": remap_txt_in_, - "img_attn_qkv": remap_img_attn_qkv_, - "txt_attn_qkv": remap_txt_attn_qkv_, - "single_blocks": remap_single_transformer_blocks_, - "final_layer.adaLN_modulation.1": remap_norm_scale_shift_, - } - - # Some folks attempt to make their state dict compatible with diffusers by adding "transformer." prefix to all keys - # and use their custom code. To make sure both "original" and "attempted diffusers" loras work as expected, we make - # sure that both follow the same initial format by stripping off the "transformer." prefix. - for key in list(converted_state_dict.keys()): - try: - if key.startswith("transformer."): - converted_state_dict[key[len("transformer.") :]] = converted_state_dict.pop(key) - if key.startswith("diffusion_model."): - converted_state_dict[key[len("diffusion_model.") :]] = converted_state_dict.pop(key) - except Exception as e: - logger.warning(f"Error processing key {key}: {str(e)}") - # Skip if we can't process this key properly - - # Rename and remap the state dict keys - for key in list(converted_state_dict.keys()): - try: - new_key = key[:] - for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - converted_state_dict[new_key] = converted_state_dict.pop(key) - except Exception as e: - logger.warning(f"Error processing key {key}: {str(e)}") - # Skip if we can't process this key properly - - for key in list(converted_state_dict.keys()): - try: - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, converted_state_dict) - except Exception as e: - logger.warning(f"Error processing key {key}: {str(e)}") - # Skip if we can't process this key properly - - # Add back the "transformer." prefix - for key in list(converted_state_dict.keys()): - try: - converted_state_dict[f"transformer.{key}"] = converted_state_dict.pop(key) - except Exception as e: - logger.warning(f"Error processing key {key}: {str(e)}") - # Skip if we can't process this key properly - - return converted_state_dict - except Exception as e: - logger.error(f"LoRA conversion failed: {str(e)}") - # Return empty state dict as fallback - return {} - - -class HyVideoModel(comfy.model_base.BaseModel): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.pipeline = {} - self.load_device = mm.get_torch_device() - - def __getitem__(self, k): - return self.pipeline[k] - - def __setitem__(self, k, v): - self.pipeline[k] = v - - -class HyVideoModelConfig: - def __init__(self, dtype): - self.unet_config = {} - self.unet_extra_config = {} - self.latent_format = comfy.latent_formats.HunyuanVideo - self.latent_format.latent_channels = 16 - self.manual_cast_dtype = dtype - self.sampling_settings = {"multiplier": 1.0} - self.memory_usage_factor = 2.0 - self.unet_config["disable_unet_model_creation"] = True - -class FramePackTorchCompileSettings: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "backend": (["inductor","cudagraphs"], {"default": "inductor"}), - "fullgraph": ("BOOLEAN", {"default": False, "tooltip": "Enable full graph mode"}), - "mode": (["default", "max-autotune", "max-autotune-no-cudagraphs", "reduce-overhead"], {"default": "default"}), - "dynamic": ("BOOLEAN", {"default": False, "tooltip": "Enable dynamic mode"}), - "dynamo_cache_size_limit": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 1, "tooltip": "torch._dynamo.config.cache_size_limit"}), - "compile_single_blocks": ("BOOLEAN", {"default": True, "tooltip": "Enable single block compilation"}), - "compile_double_blocks": ("BOOLEAN", {"default": True, "tooltip": "Enable double block compilation"}), - }, - } - RETURN_TYPES = ("FRAMEPACKCOMPILEARGS",) - RETURN_NAMES = ("torch_compile_args",) - FUNCTION = "loadmodel" - CATEGORY = "HunyuanVideoWrapper" - DESCRIPTION = "torch.compile settings, when connected to the model loader, torch.compile of the selected layers is attempted. Requires Triton and torch 2.5.0 is recommended" - - def loadmodel(self, backend, fullgraph, mode, dynamic, dynamo_cache_size_limit, compile_single_blocks, compile_double_blocks): - - compile_args = { - "backend": backend, - "fullgraph": fullgraph, - "mode": mode, - "dynamic": dynamic, - "dynamo_cache_size_limit": dynamo_cache_size_limit, - "compile_single_blocks": compile_single_blocks, - "compile_double_blocks": compile_double_blocks - } - - return (compile_args, ) - -#region Model loading -class DownloadAndLoadFramePackModel: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "model": (["lllyasviel/FramePackI2V_HY"],), - - "base_precision": (["fp32", "bf16", "fp16"], {"default": "bf16"}), - "quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_e5m2'], {"default": 'disabled', "tooltip": "optional quantization method"}), - }, - "optional": { - "attention_mode": ([ - "sdpa", - "flash_attn", - "sageattn", - ], {"default": "sdpa"}), - "compile_args": ("FRAMEPACKCOMPILEARGS", ), - } - } - - RETURN_TYPES = ("FramePackMODEL",) - RETURN_NAMES = ("model", ) - FUNCTION = "loadmodel" - CATEGORY = "FramePackWrapper" - - def loadmodel(self, model, base_precision, quantization, - compile_args=None, attention_mode="sdpa"): - - base_dtype = {"fp8_e4m3fn": torch.float8_e4m3fn, "fp8_e4m3fn_fast": torch.float8_e4m3fn, "bf16": torch.bfloat16, "fp16": torch.float16, "fp16_fast": torch.float16, "fp32": torch.float32}[base_precision] - - device = mm.get_torch_device() - - model_path = os.path.join(folder_paths.models_dir, "diffusers", "lllyasviel", "FramePackI2V_HY") - if not os.path.exists(model_path): - print(f"Downloading clip model to: {model_path}") - from huggingface_hub import snapshot_download - snapshot_download( - repo_id=model, - local_dir=model_path, - local_dir_use_symlinks=False, - ) - - transformer = HunyuanVideoTransformer3DModel.from_pretrained(model_path, torch_dtype=base_dtype, attention_mode=attention_mode).cpu() - params_to_keep = {"norm", "bias", "time_in", "vector_in", "guidance_in", "txt_in", "img_in"} - if quantization == 'fp8_e4m3fn' or quantization == 'fp8_e4m3fn_fast': - transformer = transformer.to(torch.float8_e4m3fn) - if quantization == "fp8_e4m3fn_fast": - from .fp8_optimization import convert_fp8_linear - convert_fp8_linear(transformer, base_dtype, params_to_keep=params_to_keep) - elif quantization == 'fp8_e5m2': - transformer = transformer.to(torch.float8_e5m2) - else: - transformer = transformer.to(base_dtype) - - DynamicSwapInstaller.install_model(transformer, device=device) - - if compile_args is not None: - if compile_args["compile_single_blocks"]: - for i, block in enumerate(transformer.single_transformer_blocks): - transformer.single_transformer_blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"]) - if compile_args["compile_double_blocks"]: - for i, block in enumerate(transformer.transformer_blocks): - transformer.transformer_blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"]) - - #transformer = torch.compile(transformer, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"]) - - pipe = { - "transformer": transformer.eval(), - "dtype": base_dtype, - } - return (pipe, ) - -class FramePackLoraSelect: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "lora": (folder_paths.get_filename_list("loras"), - {"tooltip": "LORA models are expected to be in ComfyUI/models/loras with .safetensors extension"}), - "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}), - "fuse_lora": ("BOOLEAN", {"default": True, "tooltip": "Fuse the LORA model with the base model. This is recommended for better performance."}), - }, - "optional": { - "prev_lora":("FPLORA", {"default": None, "tooltip": "For loading multiple LoRAs"}), - } - } - - RETURN_TYPES = ("FPLORA",) - RETURN_NAMES = ("lora", ) - FUNCTION = "getlorapath" - CATEGORY = "FramePackWrapper" - DESCRIPTION = "Select a LoRA model from ComfyUI/models/loras" - - def getlorapath(self, lora, strength, prev_lora=None, fuse_lora=True): - loras_list = [] - - lora = { - "path": folder_paths.get_full_path("loras", lora), - "strength": strength, - "name": lora.split(".")[0], - "fuse_lora": fuse_lora, - } - if prev_lora is not None: - loras_list.extend(prev_lora) - - loras_list.append(lora) - return (loras_list,) - -class LoadFramePackModel: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder",}), - - "base_precision": (["fp32", "bf16", "fp16"], {"default": "bf16"}), - "quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_e5m2'], {"default": 'disabled', "tooltip": "optional quantization method"}), - "load_device": (["main_device", "offload_device"], {"default": "cuda", "tooltip": "Initialize the model on the main device or offload device"}), - }, - "optional": { - "attention_mode": ([ - "sdpa", - "flash_attn", - "sageattn", - ], {"default": "sdpa"}), - "compile_args": ("FRAMEPACKCOMPILEARGS", ), - "lora": ("FPLORA", {"default": None, "tooltip": "LORA model to load"}), - } - } - - RETURN_TYPES = ("FramePackMODEL",) - RETURN_NAMES = ("model", ) - FUNCTION = "loadmodel" - CATEGORY = "FramePackWrapper" - - def loadmodel(self, model, base_precision, quantization, - compile_args=None, attention_mode="sdpa", lora=None, load_device="main_device"): - base_dtype = {"fp8_e4m3fn": torch.float8_e4m3fn, "fp8_e4m3fn_fast": torch.float8_e4m3fn, "bf16": torch.bfloat16, "fp16": torch.float16, "fp16_fast": torch.float16, "fp32": torch.float32}[base_precision] - device = mm.get_torch_device() - offload_device = mm.unet_offload_device() - if load_device == "main_device": - transformer_load_device = device - else: - transformer_load_device = offload_device - model_path = folder_paths.get_full_path_or_raise("diffusion_models", model) - model_config_path = os.path.join(script_directory, "transformer_config.json") - import json - with open(model_config_path, "r") as f: - config = json.load(f) - sd = load_torch_file(model_path, device=offload_device, safe_load=True) - model_weight_dtype = sd['single_transformer_blocks.0.attn.to_k.weight'].dtype - with init_empty_weights(): - transformer = HunyuanVideoTransformer3DModel(**config, attention_mode=attention_mode) - params_to_keep = {"norm", "bias", "time_in", "vector_in", "guidance_in", "txt_in", "img_in"} - if quantization == "fp8_e4m3fn" or quantization == "fp8_e4m3fn_fast" or quantization == "fp8_scaled": - dtype = torch.float8_e4m3fn - elif quantization == "fp8_e5m2": - dtype = torch.float8_e5m2 - else: - dtype = base_dtype - if lora is not None: - after_lora_dtype = dtype - dtype = base_dtype - print("Using accelerate to load and assign model weights to device...") - param_count = sum(1 for _ in transformer.named_parameters()) - for name, param in tqdm(transformer.named_parameters(), - desc=f"Loading transformer parameters to {transformer_load_device}", - total=param_count, - leave=True): - dtype_to_use = base_dtype if any(keyword in name for keyword in params_to_keep) else dtype - set_module_tensor_to_device(transformer, name, device=transformer_load_device, dtype=dtype_to_use, value=sd[name]) - - if lora is not None: - adapter_list = [] - adapter_weights = [] - - for l in lora: - fuse = True if l["fuse_lora"] else False - lora_sd = load_torch_file(l["path"]) - - if "lora_unet_single_transformer_blocks_0_attn_to_k.lora_up.weight" in lora_sd: - from .utils import convert_to_diffusers - lora_sd = convert_to_diffusers("lora_unet_", lora_sd) - - if not "transformer.single_transformer_blocks.0.attn.to_k.lora_A.weight" in lora_sd: - log.info(f"Converting LoRA weights from {l['path']} to diffusers format...") - # Make a copy of the original state dict to avoid modifying it - state_dict_copy = {} - - # Remove scalar (0-dimensional) tensors that cause problems - for key, value in lora_sd.items(): - if isinstance(value, torch.Tensor): - if value.dim() == 0: - print(f"Skipping 0-dimensional tensor: {key}") - continue - state_dict_copy[key] = value - else: - print(f"Skipping non-tensor value: {key}") - - print(f"After filtering: {len(state_dict_copy)} valid keys") - - # Try the original conversion with the filtered state dict - try: - from diffusers.loaders.lora_conversion_utils import _convert_hunyuan_video_lora_to_diffusers - lora_sd = _convert_hunyuan_video_lora_to_diffusers(state_dict_copy) - print("Successfully converted LoRA weights") - except Exception as e: - print(f"Error in standard conversion: {e}") - # Fall back to empty dict if conversion fails - print("Conversion failed, returning empty state dict") - lora_sd = {} - - lora_rank = None - for key, val in lora_sd.items(): - if "lora_B" in key or "lora_up" in key: - lora_rank = val.shape[1] - break - if lora_rank is not None: - log.info(f"Merging rank {lora_rank} LoRA weights from {l['path']} with strength {l['strength']}") - adapter_name = l['path'].split("/")[-1].split(".")[0] - adapter_weight = l['strength'] - transformer.load_lora_adapter(lora_sd, weight_name=l['path'].split("/")[-1], lora_rank=lora_rank, adapter_name=adapter_name) - - adapter_list.append(adapter_name) - adapter_weights.append(adapter_weight) - - del lora_sd - mm.soft_empty_cache() - if adapter_list: - transformer.set_adapters(adapter_list, weights=adapter_weights) - if fuse: - if model_weight_dtype not in [torch.float32, torch.float16, torch.bfloat16]: - raise ValueError("Fusing LoRA doesn't work well with fp8 model weights. Please use a bf16 model file, or disable LoRA fusing.") - lora_scale = 1 - transformer.fuse_lora(lora_scale=lora_scale) - transformer.delete_adapters(adapter_list) - - if quantization == "fp8_e4m3fn" or quantization == "fp8_e4m3fn_fast" or quantization == "fp8_e5m2": - params_to_keep = {"norm", "bias", "time_in", "vector_in", "guidance_in", "txt_in", "img_in"} - for name, param in transformer.named_parameters(): - # Make sure to not cast the LoRA weights to fp8. - if not any(keyword in name for keyword in params_to_keep) and not 'lora' in name: - param.data = param.data.to(after_lora_dtype) - - if quantization == "fp8_e4m3fn_fast": - from .fp8_optimization import convert_fp8_linear - convert_fp8_linear(transformer, base_dtype, params_to_keep=params_to_keep) - - DynamicSwapInstaller.install_model(transformer, device=device) - - if compile_args is not None: - if compile_args["compile_single_blocks"]: - for i, block in enumerate(transformer.single_transformer_blocks): - transformer.single_transformer_blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"]) - if compile_args["compile_double_blocks"]: - for i, block in enumerate(transformer.transformer_blocks): - transformer.transformer_blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"]) - - pipe = { - "transformer": transformer.eval(), - "dtype": base_dtype, - } - return (pipe, ) - -class FramePackFindNearestBucket: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "image": ("IMAGE", {"tooltip": "Image to resize"}), - "base_resolution": ("INT", {"default": 640, "min": 64, "max": 2048, "step": 16, "tooltip": "Width of the image to encode"}), - }, - } - - RETURN_TYPES = ("INT", "INT", ) - RETURN_NAMES = ("width","height",) - FUNCTION = "process" - CATEGORY = "FramePackWrapper" - DESCRIPTION = "Finds the closes resolution bucket as defined in the orignal code" - - def process(self, image, base_resolution): - - H, W = image.shape[1], image.shape[2] - - new_height, new_width = find_nearest_bucket(H, W, resolution=base_resolution) - - return (new_width, new_height, ) - -# 元のFramePackSamplerを保持（動画生成用） -class FramePackSampler: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "model": ("FramePackMODEL",), - "positive": ("CONDITIONING",), - "negative": ("CONDITIONING",), - "start_latent": ("LATENT", {"tooltip": "init Latents to use for image2video"} ), - "steps": ("INT", {"default": 30, "min": 1}), - "use_teacache": ("BOOLEAN", {"default": True, "tooltip": "Use teacache for faster sampling."}), - "teacache_rel_l1_thresh": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The threshold for the relative L1 loss."}), - "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 30.0, "step": 0.01}), - "guidance_scale": ("FLOAT", {"default": 10.0, "min": 0.0, "max": 32.0, "step": 0.01}), - "shift": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01}), - "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "latent_window_size": ("INT", {"default": 9, "min": 1, "max": 33, "step": 1, "tooltip": "The size of the latent window to use for sampling."}), - "total_second_length": ("FLOAT", {"default": 5, "min": 1, "max": 120, "step": 0.1, "tooltip": "The total length of the video in seconds."}), - "gpu_memory_preservation": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 128.0, "step": 0.1, "tooltip": "The amount of GPU memory to preserve."}), - "sampler": (["unipc_bh1", "unipc_bh2"], - { - "default": 'unipc_bh1' - }), - }, - "optional": { - "image_embeds": ("CLIP_VISION_OUTPUT", ), - "end_latent": ("LATENT", {"tooltip": "end Latents to use for image2video"} ), - "end_image_embeds": ("CLIP_VISION_OUTPUT", {"tooltip": "end Image's clip embeds"} ), - "embed_interpolation": (["disabled", "weighted_average", "linear"], {"default": 'disabled', "tooltip": "Image embedding interpolation type. If linear, will smoothly interpolate with time, else it'll be weighted average with the specified weight."}), - "start_embed_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Weighted average constant for image embed interpolation. If end image is not set, the embed's strength won't be affected"}), - "initial_samples": ("LATENT", {"tooltip": "init Latents to use for video2video"} ), - "denoise_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - } - } - - RETURN_TYPES = ("LATENT", ) - RETURN_NAMES = ("samples",) - FUNCTION = "process" - CATEGORY = "FramePackWrapper" - - def process(self, model, shift, positive, negative, latent_window_size, use_teacache, total_second_length, teacache_rel_l1_thresh, steps, cfg, - guidance_scale, seed, sampler, gpu_memory_preservation, start_latent=None, image_embeds=None, end_latent=None, end_image_embeds=None, embed_interpolation="linear", start_embed_strength=1.0, initial_samples=None, denoise_strength=1.0): - total_latent_sections = (total_second_length * 30) / (latent_window_size * 4) - total_latent_sections = int(max(round(total_latent_sections), 1)) - print("total_latent_sections: ", total_latent_sections) - - transformer = model["transformer"] - base_dtype = model["dtype"] - - device = mm.get_torch_device() - offload_device = mm.unet_offload_device() - - mm.unload_all_models() - mm.cleanup_models() - mm.soft_empty_cache() - - if start_latent is not None: - start_latent = start_latent["samples"] * vae_scaling_factor - if initial_samples is not None: - initial_samples = initial_samples["samples"] * vae_scaling_factor - if end_latent is not None: - end_latent = end_latent["samples"] * vae_scaling_factor - has_end_image = end_latent is not None - print("start_latent", start_latent.shape) - B, C, T, H, W = start_latent.shape - - if image_embeds is not None: - start_image_encoder_last_hidden_state = image_embeds["last_hidden_state"].to(device, base_dtype) - - if has_end_image: - assert end_image_embeds is not None - end_image_encoder_last_hidden_state = end_image_embeds["last_hidden_state"].to(device, base_dtype) - else: - if image_embeds is not None: - end_image_encoder_last_hidden_state = torch.zeros_like(start_image_encoder_last_hidden_state) - - llama_vec = positive[0][0].to(device, base_dtype) - clip_l_pooler = positive[0][1]["pooled_output"].to(device, base_dtype) - - if not math.isclose(cfg, 1.0): - llama_vec_n = negative[0][0].to(device, base_dtype) - clip_l_pooler_n = negative[0][1]["pooled_output"].to(device, base_dtype) - else: - llama_vec_n = torch.zeros_like(llama_vec, device=device) - clip_l_pooler_n = torch.zeros_like(clip_l_pooler, device=device) - - 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) - - # Sampling - rnd = torch.Generator("cpu").manual_seed(seed) - - num_frames = latent_window_size * 4 - 3 - - history_latents = torch.zeros(size=(1, 16, 1 + 2 + 16, H, W), dtype=torch.float32).cpu() - - total_generated_latent_frames = 0 - - latent_paddings_list = list(reversed(range(total_latent_sections))) - latent_paddings = latent_paddings_list.copy() # Create a copy for iteration - - comfy_model = HyVideoModel( - HyVideoModelConfig(base_dtype), - model_type=comfy.model_base.ModelType.FLOW, - device=device, - ) - - patcher = comfy.model_patcher.ModelPatcher(comfy_model, device, torch.device("cpu")) - from latent_preview import prepare_callback - callback = prepare_callback(patcher, steps) - - move_model_to_device_with_memory_preservation(transformer, target_device=device, preserved_memory_gb=gpu_memory_preservation) - - 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] - latent_paddings_list = latent_paddings.copy() - - for i, latent_padding in enumerate(latent_paddings): - print(f"latent_padding: {latent_padding}") - is_last_section = latent_padding == 0 - is_first_section = latent_padding == latent_paddings[0] - latent_padding_size = latent_padding * latent_window_size - - if image_embeds is not None: - if embed_interpolation != "disabled": - if embed_interpolation == "linear": - if total_latent_sections <= 1: - frac = 1.0 # Handle case with only one section - else: - frac = 1 - i / (total_latent_sections - 1) # going backwards - else: - frac = start_embed_strength if has_end_image else 1.0 - - image_encoder_last_hidden_state = start_image_encoder_last_hidden_state * frac + (1 - frac) * end_image_encoder_last_hidden_state - else: - image_encoder_last_hidden_state = start_image_encoder_last_hidden_state * start_embed_strength - else: - image_encoder_last_hidden_state = None - - print(f'latent_padding_size = {latent_padding_size}, is_last_section = {is_last_section}, is_first_section = {is_first_section}') - - start_latent_frames = T # 0 or 1 - indices = torch.arange(0, sum([start_latent_frames, 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([start_latent_frames, 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) - - # Use end image latent for the first section if provided - 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) - - #vid2vid WIP - if initial_samples is not None: - total_length = initial_samples.shape[2] - - # Get the max padding value for normalization - max_padding = max(latent_paddings_list) - - if is_last_section: - # Last section should capture the end of the sequence - start_idx = max(0, total_length - latent_window_size) - else: - # Calculate windows that distribute more evenly across the sequence - # This normalizes the padding values to create appropriate spacing - if max_padding > 0: # Avoid division by zero - progress = (max_padding - latent_padding) / max_padding - start_idx = int(progress * max(0, total_length - latent_window_size)) - else: - start_idx = 0 - - end_idx = min(start_idx + latent_window_size, total_length) - print(f"start_idx: {start_idx}, end_idx: {end_idx}, total_length: {total_length}") - input_init_latents = initial_samples[:, :, start_idx:end_idx, :, :].to(device) - - if use_teacache: - transformer.initialize_teacache(enable_teacache=True, num_steps=steps, rel_l1_thresh=teacache_rel_l1_thresh) - else: - transformer.initialize_teacache(enable_teacache=False) - - with torch.autocast(device_type=mm.get_autocast_device(device), dtype=base_dtype, enabled=True): - generated_latents = sample_hunyuan( - transformer=transformer, - sampler=sampler, - initial_latent=input_init_latents if initial_samples is not None else None, - strength=denoise_strength, - width=W * 8, - height=H * 8, - frames=num_frames, - real_guidance_scale=cfg, - distilled_guidance_scale=guidance_scale, - guidance_rescale=0, - shift=shift if shift != 0 else None, - 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=device, - dtype=base_dtype, - 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) - - real_history_latents = history_latents[:, :, :total_generated_latent_frames, :, :] - - if is_last_section: - break - - transformer.to(offload_device) - mm.soft_empty_cache() - - return {"samples": real_history_latents / vae_scaling_factor}, class FramePackSingleFrameSampler: @classmethod @@ -930,43 +26,149 @@ class FramePackSingleFrameSampler: "model": ("FramePackMODEL",), "positive": ("CONDITIONING",), "negative": ("CONDITIONING",), - "start_latent": ("LATENT", {"tooltip": "init Latents to use for image2image"} ), + "start_latent": ( + "LATENT", + {"tooltip": "init Latents to use for image2image"}, + ), "steps": ("INT", {"default": 30, "min": 1}), - "use_teacache": ("BOOLEAN", {"default": True, "tooltip": "Use teacache for faster sampling."}), - "teacache_rel_l1_thresh": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The threshold for the relative L1 loss."}), - "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 30.0, "step": 0.01}), - "guidance_scale": ("FLOAT", {"default": 10.0, "min": 0.0, "max": 32.0, "step": 0.01}), - "shift": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01}), - "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "latent_window_size": ("INT", {"default": 9, "min": 1, "max": 33, "step": 1, "tooltip": "The size of the latent window to use for sampling."}), - "gpu_memory_preservation": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 128.0, "step": 0.1, "tooltip": "The amount of GPU memory to preserve."}), - "sampler": (["unipc_bh1", "unipc_bh2"], {"default": 'unipc_bh1'}), - "use_kisekaeichi": ("BOOLEAN", {"default": False, "tooltip": "Enable Kisekaeichi mode for style transfer"}), + "use_teacache": ( + "BOOLEAN", + {"default": True, "tooltip": "Use teacache for faster sampling."}, + ), + "teacache_rel_l1_thresh": ( + "FLOAT", + { + "default": 0.15, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "The threshold for the relative L1 loss.", + }, + ), + "cfg": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 30.0, "step": 0.01}, + ), + "guidance_scale": ( + "FLOAT", + {"default": 10.0, "min": 0.0, "max": 32.0, "step": 0.01}, + ), + "shift": ( + "FLOAT", + {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01}, + ), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + "latent_window_size": ( + "INT", + { + "default": 9, + "min": 1, + "max": 33, + "step": 1, + "tooltip": "The size of the latent window to use for sampling.", + }, + ), + "gpu_memory_preservation": ( + "FLOAT", + { + "default": 6.0, + "min": 0.0, + "max": 128.0, + "step": 0.1, + "tooltip": "The amount of GPU memory to preserve.", + }, + ), + "sampler": (["unipc_bh1", "unipc_bh2"], {"default": "unipc_bh1"}), + "use_kisekaeichi": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Enable Kisekaeichi mode for style transfer", + }, + ), }, "optional": { - "image_embeds": ("CLIP_VISION_OUTPUT", ), - "initial_samples": ("LATENT", {"tooltip": "init Latents to use for image2image variation"} ), - "denoise_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - "reference_latent": ("LATENT", {"tooltip": "Reference image latent for kisekaeichi mode"}), - "reference_image_embeds": ("CLIP_VISION_OUTPUT", {"tooltip": "Reference image CLIP embeds for kisekaeichi mode"}), - "target_index": ("INT", {"default": 1, "min": 0, "max": 8, "step": 1, "tooltip": "Target index for kisekaeichi mode (recommended: 1)"}), - "history_index": ("INT", {"default": 13, "min": 0, "max": 16, "step": 1, "tooltip": "History index for kisekaeichi mode (recommended: 13)"}), - "input_mask": ("MASK", {"tooltip": "Input mask for selective application"}), - "reference_mask": ("MASK", {"tooltip": "Reference mask for selective features"}), - } + "image_embeds": ("CLIP_VISION_OUTPUT",), + "initial_samples": ( + "LATENT", + {"tooltip": "init Latents to use for image2image variation"}, + ), + "denoise_strength": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, + ), + "reference_latent": ( + "LATENT", + {"tooltip": "Reference image latent for kisekaeichi mode"}, + ), + "reference_image_embeds": ( + "CLIP_VISION_OUTPUT", + {"tooltip": "Reference image CLIP embeds for kisekaeichi mode"}, + ), + "target_index": ( + "INT", + { + "default": 1, + "min": 0, + "max": 8, + "step": 1, + "tooltip": "Target index for kisekaeichi mode (recommended: 1)", + }, + ), + "history_index": ( + "INT", + { + "default": 13, + "min": 0, + "max": 16, + "step": 1, + "tooltip": "History index for kisekaeichi mode (recommended: 13)", + }, + ), + "input_mask": ( + "MASK", + {"tooltip": "Input mask for selective application"}, + ), + "reference_mask": ( + "MASK", + {"tooltip": "Reference mask for selective features"}, + ), + }, } - RETURN_TYPES = ("LATENT", ) + RETURN_TYPES = ("LATENT",) RETURN_NAMES = ("samples",) FUNCTION = "process" CATEGORY = "FramePackWrapper" DESCRIPTION = "Single frame sampler with Kisekaeichi (style transfer) support" - def process(self, model, shift, positive, negative, latent_window_size, use_teacache, teacache_rel_l1_thresh, steps, cfg, - guidance_scale, seed, sampler, gpu_memory_preservation, start_latent=None, image_embeds=None, - initial_samples=None, denoise_strength=1.0, use_kisekaeichi=False, reference_latent=None, - reference_image_embeds=None, target_index=1, history_index=13, input_mask=None, reference_mask=None): - + def process( + self, + model, + shift, + positive, + negative, + latent_window_size, + use_teacache, + teacache_rel_l1_thresh, + steps, + cfg, + guidance_scale, + seed, + sampler, + gpu_memory_preservation, + start_latent=None, + image_embeds=None, + initial_samples=None, + denoise_strength=1.0, + use_kisekaeichi=False, + reference_latent=None, + reference_image_embeds=None, + target_index=1, + history_index=13, + input_mask=None, + reference_mask=None, + ): print("=== 1フレーム推論モード ===") if use_kisekaeichi: print("Kisekaeichi（着せ替え）モード有効") @@ -989,18 +191,22 @@ class FramePackSingleFrameSampler: if use_kisekaeichi and reference_latent is not None: reference_latent = reference_latent["samples"] * vae_scaling_factor print(f"参照画像latent: {reference_latent.shape}") - + print("start_latent", start_latent.shape) B, C, T, H, W = start_latent.shape # 画像エンベッディング処理 if image_embeds is not None: - start_image_encoder_last_hidden_state = image_embeds["last_hidden_state"].to(device, base_dtype) + start_image_encoder_last_hidden_state = image_embeds[ + "last_hidden_state" + ].to(device, base_dtype) else: start_image_encoder_last_hidden_state = None if use_kisekaeichi and reference_image_embeds is not None: - reference_image_encoder_last_hidden_state = reference_image_embeds["last_hidden_state"].to(device, base_dtype) + reference_image_encoder_last_hidden_state = reference_image_embeds[ + "last_hidden_state" + ].to(device, base_dtype) print("参照画像のCLIP embeddingを設定しました") else: reference_image_encoder_last_hidden_state = None @@ -1017,179 +223,247 @@ class FramePackSingleFrameSampler: clip_l_pooler_n = torch.zeros_like(clip_l_pooler, device=device) 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) + llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask( + llama_vec_n, length=512 + ) # シード設定 rnd = torch.Generator("cpu").manual_seed(seed) # === 一つ目のコードに完全準拠した設定 === - + # 1フレームモード固定設定 sample_num_frames = 1 total_latent_sections = 1 latent_padding = 0 latent_padding_size = latent_padding * latent_window_size # 0 - + # 一つ目のコードと同じインデックス構造 - indices = torch.arange(0, sum([1, latent_padding_size, latent_window_size, 1, 2, 16])).unsqueeze(0) + indices = torch.arange( + 0, sum([1, latent_padding_size, latent_window_size, 1, 2, 16]) + ).unsqueeze(0) split_sizes = [1, latent_padding_size, latent_window_size, 1, 2, 16] - + # latent_padding_sizeが0の場合の分割（一つ目のコードと完全同一） if latent_padding_size == 0: clean_latent_indices_pre = indices[:, 0:1] - latent_indices = indices[:, 1:1+latent_window_size] - clean_latent_indices_post = indices[:, 1+latent_window_size:2+latent_window_size] - clean_latent_2x_indices = indices[:, 2+latent_window_size:4+latent_window_size] - clean_latent_4x_indices = indices[:, 4+latent_window_size:20+latent_window_size] + latent_indices = indices[:, 1 : 1 + latent_window_size] + clean_latent_indices_post = indices[ + :, 1 + latent_window_size : 2 + latent_window_size + ] + clean_latent_2x_indices = indices[ + :, 2 + latent_window_size : 4 + latent_window_size + ] + clean_latent_4x_indices = indices[ + :, 4 + latent_window_size : 20 + latent_window_size + ] blank_indices = torch.empty((1, 0), dtype=torch.long) else: - clean_latent_indices_pre, blank_indices, latent_indices, clean_latent_indices_post, clean_latent_2x_indices, clean_latent_4x_indices = indices.split(split_sizes, dim=1) + ( + clean_latent_indices_pre, + blank_indices, + latent_indices, + clean_latent_indices_post, + clean_latent_2x_indices, + clean_latent_4x_indices, + ) = indices.split(split_sizes, dim=1) # 一つ目のコードの重要な処理：1フレームモード時のone_frame_inference処理 if sample_num_frames == 1: if use_kisekaeichi and reference_latent is not None: print("=== Kisekaeichi モード設定（完全版） ===") - + # 一つ目のコードのone_frame_inference処理を完全再現 one_frame_inference = set() one_frame_inference.add(f"target_index={target_index}") one_frame_inference.add(f"history_index={history_index}") - + # 公式実装に従った処理 latent_indices = indices[:, -1:] # デフォルトは最後のフレーム - + # パラメータ解析と処理（一つ目のコードと同じ） for one_frame_param in one_frame_inference: if one_frame_param.startswith("target_index="): target_idx = int(one_frame_param.split("=")[1]) latent_indices[:, 0] = target_idx print(f"latent_indices設定: target_index={target_idx}") - + elif one_frame_param.startswith("history_index="): history_idx = int(one_frame_param.split("=")[1]) clean_latent_indices_post[:, 0] = history_idx - print(f"clean_latent_indices_post設定: history_index={history_idx}") + print( + f"clean_latent_indices_post設定: history_index={history_idx}" + ) # history_latentsのダミーを作成（一つ目のコードと同じ構造） - history_latents = torch.zeros(size=(1, 16, 1 + 2 + 16, H, W), dtype=torch.float32, device='cpu') - + history_latents = torch.zeros( + size=(1, 16, 1 + 2 + 16, H, W), dtype=torch.float32, device="cpu" + ) + # clean_latents_preの設定（入力画像） - clean_latents_pre = start_latent.to(history_latents.dtype).to(history_latents.device) + clean_latents_pre = start_latent.to(history_latents.dtype).to( + history_latents.device + ) if len(clean_latents_pre.shape) < 5: clean_latents_pre = clean_latents_pre.unsqueeze(2) - + # マスクの適用（入力画像） if input_mask is not None: print("入力画像マスクを適用中...") try: height_latent, width_latent = clean_latents_pre.shape[-2:] - + if isinstance(input_mask, torch.Tensor): input_mask_tensor = input_mask else: input_mask_tensor = torch.from_numpy(input_mask) - - input_mask_resized = common_upscale(input_mask_tensor.unsqueeze(0).unsqueeze(0), - width_latent, height_latent, - "bilinear", "center").squeeze(0).squeeze(0) - input_mask_resized = input_mask_resized.to(clean_latents_pre.device)[None, None, None, :, :] + + input_mask_resized = ( + common_upscale( + input_mask_tensor.unsqueeze(0).unsqueeze(0), + width_latent, + height_latent, + "bilinear", + "center", + ) + .squeeze(0) + .squeeze(0) + ) + input_mask_resized = input_mask_resized.to( + clean_latents_pre.device + )[None, None, None, :, :] clean_latents_pre = clean_latents_pre * input_mask_resized print("入力画像マスクを適用しました") except Exception as e: print(f"入力マスク適用エラー: {e}") - + # clean_latents_postの設定（参照画像） - clean_latents_post = reference_latent[:, :, 0:1, :, :].to(history_latents.dtype).to(history_latents.device) - + clean_latents_post = ( + reference_latent[:, :, 0:1, :, :] + .to(history_latents.dtype) + .to(history_latents.device) + ) + # マスクの適用（参照画像） if reference_mask is not None: print("参照画像マスクを適用中...") try: height_latent, width_latent = clean_latents_post.shape[-2:] - + if isinstance(reference_mask, torch.Tensor): reference_mask_tensor = reference_mask else: reference_mask_tensor = torch.from_numpy(reference_mask) - - reference_mask_resized = common_upscale(reference_mask_tensor.unsqueeze(0).unsqueeze(0), - width_latent, height_latent, - "bilinear", "center").squeeze(0).squeeze(0) - reference_mask_resized = reference_mask_resized.to(clean_latents_post.device)[None, None, None, :, :] + + reference_mask_resized = ( + common_upscale( + reference_mask_tensor.unsqueeze(0).unsqueeze(0), + width_latent, + height_latent, + "bilinear", + "center", + ) + .squeeze(0) + .squeeze(0) + ) + reference_mask_resized = reference_mask_resized.to( + clean_latents_post.device + )[None, None, None, :, :] clean_latents_post = clean_latents_post * reference_mask_resized print("参照画像マスクを適用しました") except Exception as e: print(f"参照マスク適用エラー: {e}") - + # clean_latentsを結合 - clean_latents = torch.cat([clean_latents_pre, clean_latents_post], dim=2) - + clean_latents = torch.cat( + [clean_latents_pre, clean_latents_post], dim=2 + ) + # clean_latent_indicesを結合 - clean_latent_indices = torch.cat([clean_latent_indices_pre, clean_latent_indices_post], dim=1) - + clean_latent_indices = torch.cat( + [clean_latent_indices_pre, clean_latent_indices_post], dim=1 + ) + clean_latents_2x_param = None clean_latents_4x_param = None clean_latent_2x_indices = None clean_latent_4x_indices = None - + # 2x, 4x latentsの設定も無効化 clean_latents_2x = None clean_latents_4x = None - + print("Kisekaeichi: 2x/4xインデックスを無効化しました") - + # 画像エンベッディングの処理（両方を活用） - if reference_image_encoder_last_hidden_state is not None and start_image_encoder_last_hidden_state is not None: + if ( + reference_image_encoder_last_hidden_state is not None + and start_image_encoder_last_hidden_state is not None + ): # 重み付き平均で統合 ref_weight = 0.3 # 参照画像の重み input_weight = 1.0 - ref_weight image_encoder_last_hidden_state = ( - start_image_encoder_last_hidden_state * input_weight + - reference_image_encoder_last_hidden_state * ref_weight + start_image_encoder_last_hidden_state * input_weight + + reference_image_encoder_last_hidden_state * ref_weight + ) + print( + f"画像エンベッディングを統合 (入力:{input_weight:.2f}, 参照:{ref_weight:.2f})" ) - print(f"画像エンベッディングを統合 (入力:{input_weight:.2f}, 参照:{ref_weight:.2f})") elif reference_image_encoder_last_hidden_state is not None: - image_encoder_last_hidden_state = reference_image_encoder_last_hidden_state + image_encoder_last_hidden_state = ( + reference_image_encoder_last_hidden_state + ) else: - image_encoder_last_hidden_state = start_image_encoder_last_hidden_state - + image_encoder_last_hidden_state = ( + start_image_encoder_last_hidden_state + ) + print(f"Kisekaeichi設定完了:") print(f" - clean_latents.shape: {clean_latents.shape} (入力+参照)") print(f" - latent_indices: {latent_indices}") print(f" - clean_latent_indices: {clean_latent_indices}") print(f" - sample_num_frames: {sample_num_frames}") print(f" - 2x/4x無効化: True") - + else: # 通常モード（参照画像なし） all_indices = torch.arange(0, latent_window_size).unsqueeze(0) latent_indices = all_indices[:, -1:] - + clean_latents_pre = start_latent.to(torch.float32).cpu() if len(clean_latents_pre.shape) < 5: clean_latents_pre = clean_latents_pre.unsqueeze(2) - + clean_latents_post = torch.zeros_like(clean_latents_pre) - clean_latents = torch.cat([clean_latents_pre, clean_latents_post], dim=2) - clean_latent_indices = torch.cat([clean_latent_indices_pre, clean_latent_indices_post], dim=1) - + clean_latents = torch.cat( + [clean_latents_pre, clean_latents_post], dim=2 + ) + clean_latent_indices = torch.cat( + [clean_latent_indices_pre, clean_latent_indices_post], dim=1 + ) + # 通常モードでのインデックス調整 - clean_latent_indices = torch.tensor([[0]], dtype=clean_latent_indices.dtype, device=clean_latent_indices.device) + clean_latent_indices = torch.tensor( + [[0]], + dtype=clean_latent_indices.dtype, + device=clean_latent_indices.device, + ) clean_latents = clean_latents[:, :, :1, :, :] - + clean_latents_2x_param = None clean_latents_4x_param = None clean_latent_2x_indices = None clean_latent_4x_indices = None - + # 2x, 4x latentsの設定も無効化 clean_latents_2x = None clean_latents_4x = None - + print("Kisekaeichi: 2x/4xインデックスを無効化しました") - + image_encoder_last_hidden_state = start_image_encoder_last_hidden_state - + print("通常モード設定:") print(f" - clean_latents.shape: {clean_latents.shape}") print(f" - latent_indices: {latent_indices}") @@ -1207,16 +481,27 @@ class FramePackSingleFrameSampler: model_type=comfy.model_base.ModelType.FLOW, device=device, ) - patcher = comfy.model_patcher.ModelPatcher(comfy_model, device, torch.device("cpu")) + patcher = comfy.model_patcher.ModelPatcher( + comfy_model, device, torch.device("cpu") + ) from latent_preview import prepare_callback + callback = prepare_callback(patcher, steps) # モデルをGPUに移動 - move_model_to_device_with_memory_preservation(transformer, target_device=device, preserved_memory_gb=gpu_memory_preservation) + move_model_to_device_with_memory_preservation( + transformer, + target_device=device, + preserved_memory_gb=gpu_memory_preservation, + ) # TeaCacheの設定 if use_teacache: - transformer.initialize_teacache(enable_teacache=True, num_steps=steps, rel_l1_thresh=teacache_rel_l1_thresh) + transformer.initialize_teacache( + enable_teacache=True, + num_steps=steps, + rel_l1_thresh=teacache_rel_l1_thresh, + ) else: transformer.initialize_teacache(enable_teacache=False) @@ -1225,8 +510,10 @@ class FramePackSingleFrameSampler: print(f"clean_latents使用フレーム数: {clean_latents.shape[2]}") print(f"clean_latent_2x_indices: {clean_latent_2x_indices}") print(f"clean_latent_4x_indices: {clean_latent_4x_indices}") - - with torch.autocast(device_type=mm.get_autocast_device(device), dtype=base_dtype, enabled=True): + + with torch.autocast( + device_type=mm.get_autocast_device(device), dtype=base_dtype, enabled=True + ): generated_latents = sample_hunyuan( transformer=transformer, sampler=sampler, @@ -1265,27 +552,20 @@ class FramePackSingleFrameSampler: mm.soft_empty_cache() # 出力処理 - mode_info = "Kisekaeichi（latent+エンベッディング）" if use_kisekaeichi else "通常" - print(f"=== 1フレーム生成完了 ({mode_info}モード): {generated_latents.shape} ===") - - return {"samples": generated_latents / vae_scaling_factor}, + mode_info = ( + "Kisekaeichi（latent+エンベッディング）" if use_kisekaeichi else "通常" + ) + print( + f"=== 1フレーム生成完了 ({mode_info}モード): {generated_latents.shape} ===" + ) + + return ({"samples": generated_latents / vae_scaling_factor},) + NODE_CLASS_MAPPINGS = { - "DownloadAndLoadFramePackModel": DownloadAndLoadFramePackModel, - "FramePackSampler": FramePackSampler, - "FramePackSingleFrameSampler": FramePackSingleFrameSampler, # 新しい1フレーム用サンプラー - "FramePackTorchCompileSettings": FramePackTorchCompileSettings, - "FramePackFindNearestBucket": FramePackFindNearestBucket, - "LoadFramePackModel": LoadFramePackModel, - "FramePackLoraSelect": FramePackLoraSelect, - } + "FramePackSingleFrameSampler": FramePackSingleFrameSampler, +} NODE_DISPLAY_NAME_MAPPINGS = { - "DownloadAndLoadFramePackModel": "(Down)Load FramePackModel", - "FramePackSampler": "FramePackSampler", - "FramePackSingleFrameSampler": "FramePack Single Frame Sampler", # 新しいノード - "FramePackTorchCompileSettings": "Torch Compile Settings", - "FramePackFindNearestBucket": "Find Nearest Bucket", - "LoadFramePackModel": "Load FramePackModel", - "FramePackLoraSelect": "Select Lora", - } + "FramePackSingleFrameSampler": "FramePack Single Frame Sampler", +}