From 431a8acde9e369899ca598fc6072e2700fd22a31 Mon Sep 17 00:00:00 2001 From: facok <128763816+facok@users.noreply.github.com> Date: Wed, 8 Jan 2025 00:42:41 +0800 Subject: [PATCH] Add files via upload --- __init__.py | 3 + example_workflow.json | 820 ++++++++++++++++++++++++++++++++++++++++++ nodes.py | 275 ++++++++++++++ 3 files changed, 1098 insertions(+) create mode 100644 __init__.py create mode 100644 example_workflow.json create mode 100644 nodes.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ceb0e38 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/example_workflow.json b/example_workflow.json new file mode 100644 index 0000000..6a862d4 --- /dev/null +++ b/example_workflow.json @@ -0,0 +1,820 @@ +{ + "last_node_id": 99, + "last_link_id": 285, + "nodes": [ + { + "id": 67, + "type": "ModelSamplingSD3", + "pos": [ + 360, + 0 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 284 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 195 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ModelSamplingSD3" + }, + "widgets_values": [ + 7 + ] + }, + { + "id": 11, + "type": "DualCLIPLoader", + "pos": [ + 0, + 270 + ], + "size": [ + 350, + 106 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 205 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "DualCLIPLoader" + }, + "widgets_values": [ + "clip_l.safetensors", + "llava_llama3_fp8_scaled.safe.safetensors", + "hunyuan_video", + "default" + ] + }, + { + "id": 26, + "type": "FluxGuidance", + "pos": [ + 986.8018188476562, + -19.807586669921875 + ], + "size": [ + 317.4000244140625, + 58 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 175 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 129 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "FluxGuidance" + }, + "widgets_values": [ + 6 + ], + "color": "#233", + "bgcolor": "#355" + }, + { + "id": 87, + "type": "VAEDecodeTiled", + "pos": [ + 1259.380126953125, + 174.9122772216797 + ], + "size": [ + 210, + 150 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 282 + }, + { + "name": "vae", + "type": "VAE", + "link": 237 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 238, + 239 + ] + } + ], + "properties": { + "Node name for S&R": "VAEDecodeTiled" + }, + "widgets_values": [ + 256, + 64, + 64, + 8 + ] + }, + { + "id": 10, + "type": "VAELoader", + "pos": [ + 0, + 420 + ], + "size": [ + 350, + 60 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 237 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "hunyuan_video_vae_fp32.safetensors" + ] + }, + { + "id": 83, + "type": "VHS_VideoCombine", + "pos": [ + 1503.88720703125, + 29.917325973510742 + ], + "size": [ + 494.6632385253906, + 1191.845703125 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 238 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null, + "shape": 7 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null, + "shape": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 16, + "loop_count": 0, + "filename_prefix": "mix1/HunyuanVideo", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 18, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "HunyuanVideo_00291.mp4", + "subfolder": "mix1", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 16, + "workflow": "HunyuanVideo_00291.png", + "fullpath": "D:\\ComfyUI_windows_portable\\ComfyUI\\output\\mix1\\HunyuanVideo_00291.mp4" + }, + "muted": false + } + } + }, + { + "id": 85, + "type": "SaveImage", + "pos": [ + 2250.943603515625, + 268.50616455078125 + ], + "size": [ + 531.666748046875, + 889.0310668945312 + ], + "flags": {}, + "order": 15, + "mode": 4, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 239 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "SaveImage" + }, + "widgets_values": [ + "mix1/HunyuanVideo" + ] + }, + { + "id": 22, + "type": "BasicGuider", + "pos": [ + 600, + 0 + ], + "size": [ + 222.3482666015625, + 46 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 195, + "slot_index": 0 + }, + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 129, + "slot_index": 1 + } + ], + "outputs": [ + { + "name": "GUIDER", + "type": "GUIDER", + "links": [ + 278 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "BasicGuider" + }, + "widgets_values": [] + }, + { + "id": 17, + "type": "BasicScheduler", + "pos": [ + 477.8186340332031, + 865.2354736328125 + ], + "size": [ + 315, + 106 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 285, + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "SIGMAS", + "type": "SIGMAS", + "links": [ + 280 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "BasicScheduler" + }, + "widgets_values": [ + "simple", + 30, + 1 + ] + }, + { + "id": 16, + "type": "KSamplerSelect", + "pos": [ + 475.5769958496094, + 750.35205078125 + ], + "size": [ + 315, + 58 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "SAMPLER", + "type": "SAMPLER", + "links": [ + 279 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "KSamplerSelect" + }, + "widgets_values": [ + "euler" + ] + }, + { + "id": 25, + "type": "RandomNoise", + "pos": [ + 479, + 618 + ], + "size": [ + 315, + 82 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "NOISE", + "type": "NOISE", + "links": [ + 277 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "RandomNoise" + }, + "widgets_values": [ + 499675609572506, + "randomize" + ], + "color": "#2a363b", + "bgcolor": "#3f5159" + }, + { + "id": 98, + "type": "TeaCacheHunyuanVideoSampler", + "pos": [ + 975.8233032226562, + 172.5495147705078 + ], + "size": [ + 243.60000610351562, + 138 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "noise", + "type": "NOISE", + "link": 277 + }, + { + "name": "guider", + "type": "GUIDER", + "link": 278 + }, + { + "name": "sampler", + "type": "SAMPLER", + "link": 279 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 280 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 281 + } + ], + "outputs": [ + { + "name": "output", + "type": "LATENT", + "links": [ + 282 + ] + }, + { + "name": "denoised_output", + "type": "LATENT", + "links": null + } + ], + "properties": { + "Node name for S&R": "TeaCacheHunyuanVideoSampler" + }, + "widgets_values": [ + "Fast (1.6x)" + ] + }, + { + "id": 44, + "type": "CLIPTextEncode", + "pos": [ + 440.9069519042969, + 155.72299194335938 + ], + "size": [ + 422.84503173828125, + 164.31304931640625 + ], + "flags": { + "collapsed": false + }, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 205 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 175 + ], + "slot_index": 0 + } + ], + "title": "CLIP Text Encode (Positive Prompt)", + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "yor forger wearing a black dress, the background is the bed, the photo adopts panoramic composition, half-body, showing the sexy atmosphere of modern urban women.,Crane Dolly Shot, Combining crane and dolly movements for complex, dynamic shots." + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 12, + "type": "UNETLoader", + "pos": [ + -44.083003997802734, + 98.69648742675781 + ], + "size": [ + 350, + 82 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 283 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "UNETLoader" + }, + "widgets_values": [ + "hunyuan_video_720_cfgdistill_bf16.safetensors", + "default" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 99, + "type": "LoraLoaderModelOnly", + "pos": [ + 970.1587524414062, + 400.90899658203125 + ], + "size": [ + 315, + 82 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 283 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 284, + 285 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "LoraLoaderModelOnly" + }, + "widgets_values": [ + "hy\\Yor_hyvid_v1.safetensors", + 1 + ] + }, + { + "id": 45, + "type": "EmptyHunyuanLatentVideo", + "pos": [ + 475.540771484375, + 432.673583984375 + ], + "size": [ + 315, + 130 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 281 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EmptyHunyuanLatentVideo" + }, + "widgets_values": [ + 720, + 1280, + 49, + 1 + ] + } + ], + "links": [ + [ + 129, + 26, + 0, + 22, + 1, + "CONDITIONING" + ], + [ + 175, + 44, + 0, + 26, + 0, + "CONDITIONING" + ], + [ + 195, + 67, + 0, + 22, + 0, + "MODEL" + ], + [ + 205, + 11, + 0, + 44, + 0, + "CLIP" + ], + [ + 237, + 10, + 0, + 87, + 1, + "VAE" + ], + [ + 238, + 87, + 0, + 83, + 0, + "IMAGE" + ], + [ + 239, + 87, + 0, + 85, + 0, + "IMAGE" + ], + [ + 277, + 25, + 0, + 98, + 0, + "NOISE" + ], + [ + 278, + 22, + 0, + 98, + 1, + "GUIDER" + ], + [ + 279, + 16, + 0, + 98, + 2, + "SAMPLER" + ], + [ + 280, + 17, + 0, + 98, + 3, + "SIGMAS" + ], + [ + 281, + 45, + 0, + 98, + 4, + "LATENT" + ], + [ + 282, + 98, + 0, + 87, + 0, + "LATENT" + ], + [ + 283, + 12, + 0, + 99, + 0, + "MODEL" + ], + [ + 284, + 99, + 0, + 67, + 0, + "MODEL" + ], + [ + 285, + 99, + 0, + 17, + 0, + "MODEL" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.8769226950000013, + "offset": [ + -243.18507049080947, + 79.75379974398345 + ] + }, + "groupNodes": {}, + "ue_links": [] + }, + "version": 0.4 +} \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..463f2f9 --- /dev/null +++ b/nodes.py @@ -0,0 +1,275 @@ +import os +import torch +import numpy as np +import comfy.model_management as mm +import comfy.utils +import latent_preview +from typing import Any, List, Tuple, Optional, Union, Dict + +def modulate(x, shift, scale): + """Modulate layer implementation for HunyuanVideo""" + try: + # Ensure consistent data types + shift = shift.to(dtype=x.dtype, device=x.device) + scale = scale.to(dtype=x.dtype, device=x.device) + + # Reshape shift and scale to match x dimensions + B = x.shape[0] # batch size + + # Adjust shift and scale based on x dimensions + if len(x.shape) == 3: # [B, L, D] + shift = shift.view(B, 1, -1) # [B, 1, D] + scale = scale.view(B, 1, -1) # [B, 1, D] + # Expand to match x length + shift = shift.expand(-1, x.shape[1], -1) # [B, L, D] + scale = scale.expand(-1, x.shape[1], -1) # [B, L, D] + elif len(x.shape) == 5: # [B, C, T, H, W] + shift = shift.view(B, -1, 1, 1, 1) # [B, C, 1, 1, 1] + scale = scale.view(B, -1, 1, 1, 1) # [B, C, 1, 1, 1] + # Expand to match x temporal and spatial dimensions + shift = shift.expand(-1, -1, x.shape[2], x.shape[3], x.shape[4]) # [B, C, T, H, W] + scale = scale.expand(-1, -1, x.shape[2], x.shape[3], x.shape[4]) # [B, C, T, H, W] + else: + raise ValueError(f"Unsupported input shape: {x.shape}") + + # Step-by-step calculation to reduce memory usage + result = x.mul_(1 + scale) # in-place operation + result.add_(shift) # in-place operation + + return result + + except Exception as e: + raise RuntimeError(f"Modulation failed: {str(e)}") + +class TeaCacheHunyuanVideoSampler: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "noise": ("NOISE",), + "guider": ("GUIDER",), + "sampler": ("SAMPLER",), + "sigmas": ("SIGMAS",), + "latent_image": ("LATENT",), + "speedup": (["Original (1x)", "Fast (1.6x)", "Faster (2.1x)"], { + "default": "Fast (1.6x)", + "tooltip": "Control TeaCache speed/quality trade-off:\nOriginal: Base quality\nFast: 1.6x speedup\nFaster: 2.1x speedup" + }), + } + } + + RETURN_TYPES = ("LATENT", "LATENT") + RETURN_NAMES = ("output", "denoised_output") + FUNCTION = "sample" + CATEGORY = "sampling/custom_sampling" + + def teacache_forward(self, transformer, x, timestep, context=None, y=None, guidance=None, attention_mask=None, control=None, transformer_options={}, **kwargs): + """TeaCache forward implementation""" + should_calc = True + + if transformer.enable_teacache: + try: + # Get input dimensions + B, C, T, H, W = x.shape + + # Ensure consistent input dtypes + try: + x = x.to(dtype=transformer.dtype) + timestep = timestep.to(dtype=transformer.dtype) + if context is not None: + context = context.to(dtype=transformer.dtype) + if guidance is not None: + guidance = guidance.to(dtype=transformer.dtype) + except Exception as e: + raise RuntimeError(f"Failed to convert input dtypes: {str(e)}") + + # Prepare modulation vectors + try: + # HunyuanVideo uses timestep_embedding for time step encoding + time_emb = comfy.ldm.flux.layers.timestep_embedding(timestep, 256, time_factor=1.0).to(x.dtype) + vec = transformer.time_in(time_emb) # [B, hidden_size] + + # text modulation - HunyuanVideo uses vector_in to process y instead of context + if y is not None: + if not hasattr(transformer, 'params') or not hasattr(transformer.params, 'vec_in_dim'): + raise AttributeError("Transformer missing required attributes: params.vec_in_dim") + vec = vec + transformer.vector_in(y[:, :transformer.params.vec_in_dim]) + + # guidance modulation + if guidance is not None and transformer.params.guidance_embed: + guidance_emb = comfy.ldm.flux.layers.timestep_embedding(guidance, 256).to(x.dtype) + guidance_vec = transformer.guidance_in(guidance_emb) + vec = vec + guidance_vec + + except Exception as e: + raise RuntimeError(f"Failed to prepare modulation vector: {str(e)}") + + # Embed image + try: + img = transformer.img_in(x) + except Exception as e: + raise RuntimeError(f"Failed to embed image: {str(e)}") + + if transformer.enable_teacache: + try: + # Use in-place operation to reduce memory usage + inp = img.clone() + vec_ = vec.clone() + + # Get modulation parameters + modulation_output = transformer.double_blocks[0].img_mod(vec_) + + # Process modulation output + if isinstance(modulation_output, tuple): + if len(modulation_output) >= 2: + mod_shift = modulation_output[0] + mod_scale = modulation_output[1] + if hasattr(mod_shift, 'shift') and hasattr(mod_scale, 'scale'): + img_mod1_shift = mod_shift.shift + img_mod1_scale = mod_scale.scale + else: + img_mod1_shift = mod_shift + img_mod1_scale = mod_scale + else: + raise ValueError(f"Tuple too short, expected at least 2 elements, got {len(modulation_output)}") + elif hasattr(modulation_output, 'shift') and hasattr(modulation_output, 'scale'): + img_mod1_shift = modulation_output.shift + img_mod1_scale = modulation_output.scale + elif hasattr(modulation_output, 'chunk'): + chunks = modulation_output.chunk(6, dim=-1) + img_mod1_shift = chunks[0] + img_mod1_scale = chunks[1] + else: + raise ValueError(f"Unsupported modulation output format: {type(modulation_output)}") + + # Ensure get is tensor + if not isinstance(img_mod1_shift, torch.Tensor) or not isinstance(img_mod1_scale, torch.Tensor): + raise ValueError(f"Failed to get tensor values for shift and scale") + + # Apply normalization and modulation + normed_inp = transformer.double_blocks[0].img_norm1(inp) + del inp # Release memory + + modulated_inp = modulate(normed_inp, shift=img_mod1_shift, scale=img_mod1_scale) + del normed_inp # Release memory + + # Calculate relative L1 distance and decide whether to calculate + if transformer.cnt == 0 or transformer.cnt == transformer.num_steps-1: + should_calc = True + transformer.accumulated_rel_l1_distance = 0 + else: + try: + coefficients = [7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02] + rescale_func = np.poly1d(coefficients) + rel_l1 = ((modulated_inp-transformer.previous_modulated_input).abs().mean() / + transformer.previous_modulated_input.abs().mean()).cpu().item() + transformer.accumulated_rel_l1_distance += rescale_func(rel_l1) + + if transformer.accumulated_rel_l1_distance < transformer.rel_l1_thresh: + should_calc = False + else: + should_calc = True + transformer.accumulated_rel_l1_distance = 0 + except Exception as e: + should_calc = True + + transformer.previous_modulated_input = modulated_inp + transformer.cnt += 1 + + except Exception as e: + should_calc = True + + except Exception as e: + should_calc = True + + # If calculation is needed, call the original forward function + if should_calc: + try: + out = transformer.original_forward(x, timestep, context, y, guidance, + attention_mask=attention_mask, + control=control, + transformer_options=transformer_options, + **kwargs) + transformer.previous_residual = out + return out + except Exception as e: + raise + else: + # If calculation is not needed, return the previous result + return transformer.previous_residual + + def sample(self, noise, guider, sampler, sigmas, latent_image, speedup): + """Sampling implementation""" + device = mm.get_torch_device() + + # Convert options to specific thresholds + thresh_map = { + "Original (1x)": 0.0, + "Fast (1.6x)": 0.1, + "Faster (2.1x)": 0.15 + } + actual_thresh = thresh_map[speedup] + + try: + # Get transformer + transformer = guider.model_patcher.model.diffusion_model + + # Initialize TeaCache state + transformer.enable_teacache = True + transformer.cnt = 0 + transformer.num_steps = len(sigmas) - 1 + transformer.rel_l1_thresh = actual_thresh + transformer.accumulated_rel_l1_distance = 0 + transformer.previous_modulated_input = None + transformer.previous_residual = None + + latent = latent_image + latent_image = latent["samples"] + latent = latent.copy() + + noise_mask = None + if "noise_mask" in latent: + noise_mask = latent["noise_mask"] + + # Save original forward function + transformer.original_forward = transformer.forward + + # Replace with TeaCache forward + transformer.forward = lambda *args, **kwargs: self.teacache_forward(transformer, *args, **kwargs) + + try: + x0_output = {} + callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output) + + disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED + samples = guider.sample(noise.generate_noise(latent), latent_image, sampler, sigmas, + denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, + seed=noise.seed) + samples = samples.to(mm.intermediate_device()) + + finally: + # Restore original forward function + transformer.forward = transformer.original_forward + delattr(transformer, 'original_forward') + transformer.enable_teacache = False + + out = latent.copy() + out["samples"] = samples + if "x0" in x0_output: + out_denoised = latent.copy() + out_denoised["samples"] = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu()) + else: + out_denoised = out + + return (out, out_denoised) + + except Exception as e: + raise + +NODE_CLASS_MAPPINGS = { + "TeaCacheHunyuanVideoSampler": TeaCacheHunyuanVideoSampler +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "TeaCacheHunyuanVideoSampler": "TeaCache HunyuanVideo Sampler" +} \ No newline at end of file