import os import torch import math from tqdm import tqdm import sys import logging from pathlib import Path 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 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 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 # Import original function for fallback from diffusers.loaders.lora_conversion_utils import _convert_hunyuan_video_lora_to_diffusers 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) 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) # Move model directly to device transformer = transformer.to(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) # Load state dict directly to the target device sd = load_torch_file(model_path, device=transformer_load_device, safe_load=True) model_weight_dtype = sd['single_transformer_blocks.0.attn.to_k.weight'].dtype # Create model on the target device directly with proper dtype 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 # Store the original dtype for later use with LoRA after_lora_dtype = dtype if lora is not None: # When LoRA is present, we need to use base_dtype initially dtype = base_dtype # Convert dtype for state dict if needed when LoRA is present # Only convert if dtypes actually differ to avoid unnecessary memory operations if dtype != model_weight_dtype: print(f"Converting model weights from {model_weight_dtype} to {dtype} for LoRA compatibility") for key in sd.keys(): if any(keyword in key for keyword in params_to_keep): sd[key] = sd[key].to(base_dtype) else: sd[key] = sd[key].to(dtype) else: # When no LoRA, only convert if model dtype differs from target dtype # AND if quantization is disabled (otherwise we handle it later) if quantization == "disabled" and dtype != model_weight_dtype: print(f"Converting model weights from {model_weight_dtype} to {dtype}") for key in sd.keys(): if any(keyword in key for keyword in params_to_keep): sd[key] = sd[key].to(base_dtype) else: sd[key] = sd[key].to(dtype) # Load state dict all at once print(f"Loading transformer state dict to {transformer_load_device}") info = transformer.load_state_dict(sd, strict=True) print(f"Loaded state dict: {info}") # Move model to device with proper dtype if quantization != "disabled" and lora is None: # For quantized models without LoRA, convert to target dtype directly transformer = transformer.to(transformer_load_device).to(after_lora_dtype) else: # For non-quantized or LoRA models, just move to device transformer = transformer.to(transformer_load_device) 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" 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) # Move model to device if it's not already there if transformer_load_device != device: transformer = transformer.to(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, ) 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}, NODE_CLASS_MAPPINGS = { "DownloadAndLoadFramePackModel": DownloadAndLoadFramePackModel, "FramePackSampler": FramePackSampler, "FramePackTorchCompileSettings": FramePackTorchCompileSettings, "FramePackFindNearestBucket": FramePackFindNearestBucket, "LoadFramePackModel": LoadFramePackModel, "FramePackLoraSelect": FramePackLoraSelect, } NODE_DISPLAY_NAME_MAPPINGS = { "DownloadAndLoadFramePackModel": "(Down)Load FramePackModel", "FramePackSampler": "FramePackSampler", "FramePackTorchCompileSettings": "Torch Compile Settings", "FramePackFindNearestBucket": "Find Nearest Bucket", "LoadFramePackModel": "Load FramePackModel", "FramePackLoraSelect": "Select Lora", }