diff --git a/README.md b/README.md index 7e4af4e..e7b0fb3 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,15 @@ ComfyUI-LTXVideo is a collection of custom nodes for ComfyUI, designed to provid The model itself is supported in the core ComfyUI [code](https://github.com/comfyanonymous/ComfyUI/tree/master/comfy/ldm/lightricks). The main LTXVideo repository can be found [here](https://github.com/Lightricks/LTX-Video). +# ⭐ 9.06.2025 – LTXVideo VAE Patcher, Mask manipulation and Q8 LoRA loader nodes. ⭐ +1. **LTXV Patcher VAE**
The new node improves VAE decoding performance by reducing runtime and cutting memory consumption by up to 50%. This allows generation of higher-resolution outputs on consumer-grade GPUs with limited VRAM, without needing to load the VAE partially or decode in tiles.
+⚠️ On *Windows*, you may need to add the paths to the *MSVC compiler (cl.exe)* and *ninja.exe* to your system environment PATH variable.
+2. **LTXV Preprocess Masks**
+Preprocesses masks for use with the LTXVideo model's latent masking. It validates mask dimensions based on VAE downscaling, supports optional inversion, handles the first frame mask separately, combines temporal masks via max pooling, applies morphological operations to grow or shrink masks, and clamps values to ensure correct opacity. The result is a set of masks ready for latent-space masking. +3. **LTXV Q8 Lora Model Loader**
+Applying LoRA to an FP8-quantized model requires special handling to preserve output quality. It's crucial to apply LoRA weights using the correct precision, as the current LoRA implementation in ComfyUI does so in a non-optimal way. This node addresses that limitation by ensuring LoRA weights are applied properly, resulting in significantly better quality. If you're working with an FP8 LTXV model, using this node guarantees that LoRA behaves as expected and delivers the intended effect. + + # ⭐ 14.05.2025 – LTXVideo 13B 0.9.7 Distilled Release ⭐ ### 🚀 What's New in LTXVideo 13B 0.9.7 Distilled diff --git a/__init__.py b/__init__.py index 7fbcbc5..076fdc8 100644 --- a/__init__.py +++ b/__init__.py @@ -5,6 +5,7 @@ from .guide import LTXVAddGuideAdvanced from .latent_adain import LTXVAdainLatent from .latent_upsampler import LTXVLatentUpsampler from .latents import LTXVSelectLatents, LTXVSetVideoLatentNoiseMasks +from .masks import LTXVPreprocessMasks from .nodes_registry import NODE_CLASS_MAPPINGS as RUNTIME_NODE_CLASS_MAPPINGS from .nodes_registry import ( NODE_DISPLAY_NAME_MAPPINGS as RUNTIME_NODE_DISPLAY_NAME_MAPPINGS, @@ -22,6 +23,7 @@ from .stg import ( from .tiled_sampler import LTXVTiledSampler from .tricks import NODE_CLASS_MAPPINGS as TRICKS_NODE_CLASS_MAPPINGS from .tricks import NODE_DISPLAY_NAME_MAPPINGS as TRICKS_NODE_DISPLAY_NAME_MAPPINGS +from .vae_patcher.vae_patcher import LTXVPatcherVAE # Static node mappings, required for ComfyUI-Manager mapping to work NODE_CLASS_MAPPINGS = { @@ -32,7 +34,9 @@ NODE_CLASS_MAPPINGS = { "LTXVApplySTG": LTXVApplySTG, "LTXVBaseSampler": LTXVBaseSampler, "LTXVFilmGrain": LTXVFilmGrain, + "LTXVPreprocessMasks": LTXVPreprocessMasks, "LTXVLatentUpsampler": LTXVLatentUpsampler, + "LTXVPatcherVAE": LTXVPatcherVAE, "LTXVPromptEnhancer": LTXVPromptEnhancer, "LTXVPromptEnhancerLoader": LTXVPromptEnhancerLoader, "LTXQ8Patch": LTXVQ8Patch, diff --git a/masks.py b/masks.py new file mode 100644 index 0000000..7660695 --- /dev/null +++ b/masks.py @@ -0,0 +1,212 @@ +from typing import Any + +import torch +from comfy.model_management import is_device_cpu +from comfy_extras.nodes_mask import GrowMask + +from .nodes_registry import comfy_node + + +@comfy_node(name="LTXVPreprocessMasks") +class LTXVPreprocessMasks: + """ + Preprocesses masks for use with the LTXVideo model's latent masking. + + This node handles temporal mask processing by: + 1. Validating mask dimensions against VAE downscaling factors + 2. Optionally inverting masks + 3. Handling the first frame mask separately (can be ignored) + 4. Combining temporal masks using max pooling + 5. Growing/shrinking masks with morphological operations + 6. Clamping mask values to ensure proper opacity + + The output is a set of masks ready for latent-space masking operations. + """ + + @classmethod + def INPUT_TYPES(cls: "LTXVPreprocessMasks") -> dict[str, Any]: # noqa: N802 + return { + "required": { + "masks": ("MASK",), + "vae": ("VAE",), + "invert_input_masks": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Invert the input masks before processing." + "Useful for masking vs unmasking operations.", + }, + ), + "ignore_first_mask": ( + "BOOLEAN", + { + "default": True, + "tooltip": "Zero out the first mask, typically used when" + "it corresponds to the conditioning frame.", + }, + ), + "pooling_method": ( + ["max", "mean", "min"], + { + "default": "max", + "tooltip": "Method to combine temporal masks. Max preserves strongest values, " + "mean averages them, min takes weakest values.", + }, + ), + "grow_mask": ( + "INT", + { + "default": 0, + "min": -16384, + "max": 16384, + "step": 1, + "tooltip": "Pixels to grow (positive) or shrink (negative) the mask." + "Uses morphological operations.", + }, + ), + "tapered_corners": ( + "BOOLEAN", + { + "default": True, + "tooltip": "Use tapered corners in morphological operations for smoother mask edges.", + }, + ), + "clamp_min": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Minimum value to clamp masks to, ensuring minimum opacity.", + }, + ), + "clamp_max": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Maximum value to clamp masks to.", + }, + ), + }, + } + + RETURN_TYPES = ("MASK",) + FUNCTION = "preprocess_masks" + CATEGORY = "Lightricks/mask_operations" + DESCRIPTION = ( + "Preprocess masks to be used for masking latents in the LTXVideo model." + ) + + def preprocess_masks( + self, + masks: torch.Tensor, + vae, # noqa: ANN001 + pooling_method: str, + grow_mask: int, + tapered_corners: bool, + clamp_min: float, + clamp_max: float, + ignore_first_mask: bool, + invert_input_masks: bool, + ) -> tuple[torch.Tensor]: + """ + Preprocess masks for temporal latent masking. + + Args: + masks: Input masks tensor of shape (batch_size, H, W) where batch_size = 1 + (num_frames - 1) + The first mask is for the conditioning frame, remaining masks are for generated frames + vae: VAE model object containing downscale_index_formula attribute + pooling_method: Method to combine temporal masks ('max', 'mean', 'min') + grow_mask: Pixels to grow (positive) or shrink (negative) the mask + tapered_corners: Whether to use tapered corners in morphological operations + clamp_min: Minimum value to clamp masks to (ensures minimum opacity) + clamp_max: Maximum value to clamp masks to + ignore_first_mask: Whether to zero out the first mask (often used for conditioning) + invert_input_masks: Whether to invert all input masks (1 - mask) + + Returns: + Tuple containing the processed masks tensor + + Raises: + ValueError: If mask dimensions are not compatible with VAE downscaling factors + """ + + if masks.ndim != 3: + raise ValueError("Masks must be of shape (batch_size, H, W).") + + batch_size, H, W = masks.shape + + # Get VAE downscaling factors for temporal and spatial dimensions + time_scale_factor, width_scale_factor, height_scale_factor = ( + vae.downscale_index_formula + ) + + # Validate that mask dimensions are compatible with VAE downscaling + # The batch size minus 1 (for the first frame) must be divisible by time_scale_factor + if (batch_size - 1) % time_scale_factor != 0: + raise ValueError( + f"Masks batch size must have a multiple of {time_scale_factor} masks + 1." + ) + + if H % height_scale_factor != 0: + raise ValueError( + f"Masks height must have a multiple of {height_scale_factor}." + ) + + if W % width_scale_factor != 0: + raise ValueError( + f"Masks width must have a multiple of {width_scale_factor}." + ) + + # Optionally invert all masks (useful for masking vs unmasking operations) + if invert_input_masks: + masks = 1.0 - masks + + # Handle first mask separately - often used as conditioning frame + if ignore_first_mask: + first_mask = masks[0:1, :, :] * 0 # Zero out the first mask + else: + first_mask = masks[0:1, :, :] + + # Process remaining masks by grouping them according to time_scale_factor + # and taking the maximum value across each temporal group + remaining = masks[1:].view(-1, time_scale_factor, H, W) + + # Apply selected pooling method + pooling_methods = { + "max": lambda x: x.max(dim=1)[0], + "mean": lambda x: x.mean(dim=1), + "min": lambda x: x.min(dim=1)[0], + } + + combined = pooling_methods[pooling_method](remaining) + + # Apply morphological operations to grow or shrink masks + if grow_mask != 0: + grow_mask_node = GrowMask() + + # GrowMask node is not supported on GPU, so we need to move the combined tensor to CPU + # if it is not already on CPU + input_device = combined.device + if not is_device_cpu(input_device): + combined = combined.cpu() + + (combined,) = grow_mask_node.expand_mask( + mask=combined, expand=grow_mask, tapered_corners=tapered_corners + ) + + # Move the combined tensor back to the original device + if not is_device_cpu(input_device): + combined = combined.to(input_device) + + # Combine first mask with processed remaining masks and clamp values + # Clamping ensures masks have proper opacity values + processed_masks = torch.cat( + [first_mask, torch.clamp(combined, clamp_min, clamp_max)], dim=0 + ) + + return (processed_masks,) diff --git a/pyproject.toml b/pyproject.toml index 84fa532..1749530 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,14 +12,15 @@ dependencies = [ "diffusers", "huggingface_hub>=0.25.2", "transformers[timm]>=4.45.0", - "einops" + "einops", + "ninja~=1.11.1.4" ] [project.optional-dependencies] internal = [ - "ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.7", + "ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.1", "av>=10.0.0", - "q8-kernels==0.0.5" + "q8-kernels==0.0.6" ] dev = [ "pre-commit>=4.0.1", diff --git a/q8_nodes.py b/q8_nodes.py index c243fc0..748b506 100644 --- a/q8_nodes.py +++ b/q8_nodes.py @@ -1,4 +1,13 @@ +import comfy +import comfy.model_detection +import comfy.model_management +import comfy.model_patcher +import comfy.utils +import folder_paths +import torch + try: + from q8_kernels.functional.ops import hadamard_transform from q8_kernels.integration.patch_transformer import ( patch_comfyui_native_transformer, patch_comfyui_transformer, @@ -53,3 +62,73 @@ class LTXVQ8Patch: patcher(transformer, use_fp8_attention, True) m.add_object_patch(transformer_key, transformer) return (m,) + + +def idendity_quant_fn(x, t): + return x.to(dtype=t) + + +@comfy_node(name="LTXV Q8 Lora Model Loader") +class LoraLoaderQ8ModelOnly: + def __init__(self): + self.loaded_lora = None + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "lora_name": (folder_paths.get_filename_list("loras"),), + "strength_model": ( + "FLOAT", + {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}, + ), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "load_lora_model_only" + + def load_lora(self, model, lora_name, strength_model): + quant_fn = hadamard_transform + is_patched_transformer = getattr( + model.get_model_object("diffusion_model"), "is_patched", False + ) + if not is_patched_transformer or not Q8_AVAILABLE: + raise ValueError( + "LTXV Q8 Patcher is not applied to the model. Please use LTXQ8Patch node before loading lora or install q8_kernels." + ) + + if strength_model == 0: + return model + + lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) + lora = None + if self.loaded_lora is not None: + if self.loaded_lora[0] == lora_path: + lora = self.loaded_lora[1] + else: + self.loaded_lora = None + + if lora is None: + lora = comfy.utils.load_torch_file(lora_path, safe_load=True) + new_lora = {} + for k in lora: + device = lora[k].device + if lora[k].ndim == 2: + if "lora_A" in k: + new_lora[k] = quant_fn( + lora[k].to(device="cuda", dtype=torch.bfloat16), + out_type=torch.bfloat16, + ).to(device) + else: + new_lora[k] = lora[k] + self.loaded_lora = (lora_path, new_lora) + + model_lora, _ = comfy.sd.load_lora_for_models( + model, None, new_lora, strength_model, 0 + ) + return model_lora + + def load_lora_model_only(self, model, lora_name, strength_model): + return (self.load_lora(model, lora_name, strength_model),) diff --git a/requirements.txt b/requirements.txt index 3110f2f..28a4ee0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ diffusers einops huggingface_hub>=0.25.2 +ninja~=1.11.1.4 transformers[timm]>=4.45.0 diff --git a/vae_patcher.py b/vae_patcher.py new file mode 100644 index 0000000..8c36eca --- /dev/null +++ b/vae_patcher.py @@ -0,0 +1,38 @@ +from comfy.model_management import dtype_size + +from .nodes_registry import comfy_node + + +@comfy_node(name="LTXVPatcherVAE") +class LTXVPatcherVAE: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "vae": ("VAE",), + } + } + + RETURN_TYPES = ("VAE",) + FUNCTION = "patch" + CATEGORY = "lightricks/LTXV" + TITLE = "LTXV VAE Patcher" + + def patch(self, vae): + from .vae_patcher.patch import patch_vae + + vae_model = vae.first_stage_model + vae.memory_used_decode = ( + lambda shape, dtype: shape[1] + * shape[2] + * shape[3] + * shape[4] + * 8 + * 8 + * 8 + * 3 + * dtype_size(dtype) + ) + + patch_vae(vae_model, patch_block=4) + return (vae,) diff --git a/vae_patcher/__init__.py b/vae_patcher/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/vae_patcher/forwards.py b/vae_patcher/forwards.py new file mode 100644 index 0000000..711fde1 --- /dev/null +++ b/vae_patcher/forwards.py @@ -0,0 +1,273 @@ +from typing import Optional + +import numpy as np +import torch +from einops import rearrange + +from .kernels.ops import add_inplace, pixel_norm_inplace + + +def find_split_size(F: int) -> int: + s = 16 + while F % s < 3: + s -= 1 + return s + + +def inplace_patch_conv_2(x, conv: torch.nn.Conv3d): + + assert conv.kernel_size == (3, 3, 3) + assert conv.stride == (1, 1, 1) + assert conv.padding == (0, 1, 1) + assert conv.dilation == (1, 1, 1) + assert conv.groups == 1 + + x[:, :, 0, :, :] = x[:, :, 1, :, :].clone() + x[:, :, -1, :, :] = x[:, :, -2, :, :].clone() + + if x.shape[2] > 16: + split_size = find_split_size(x.shape[2]) + num_splits = (x.shape[2] + split_size - 1) // split_size + else: + split_size = x.shape[2] - 1 + num_splits = 1 + + out_channels = conv.out_channels + in_channels = conv.in_channels + x_buffers = torch.empty( + x.shape[0], + x.shape[1], + 1, + x.shape[3], + x.shape[4], + device=x.device, + dtype=x.dtype, + ) + o_buffers = torch.empty( + x.shape[0], + x.shape[1], + 1, + x.shape[3], + x.shape[4], + device=x.device, + dtype=x.dtype, + ) + # 0 case + x_buffers[:, :, 0, :, :] = x[:, :, split_size - 1, :, :].clone() + x[:, :out_channels, 1:split_size, :, :] = conv( + x[:, :in_channels, : split_size + 1, :, :] + ) + + for i in range(1, num_splits): + curr_frame_start = i * split_size + curr_frame_end = min((i + 1) * split_size, x.shape[2] - 1) + o_buffers[:, :, 0, :, :] = x[:, :, curr_frame_start - 1, :, :].clone() + x[:, :, curr_frame_start - 1, :, :] = x_buffers[:, :, 0, :, :].clone() + x_buffers[:, :, 0, :, :] = x[:, :, curr_frame_end - 1, :, :].clone() + x[:, :out_channels, curr_frame_start:curr_frame_end, :, :] = conv( + x[:, :in_channels, curr_frame_start - 1 : curr_frame_end + 1, :, :] + ) + x[:, :, curr_frame_start - 1, :, :] = o_buffers[:, :, 0, :, :].clone() + + +def res_block_forward( + self, + hidden_states: torch.FloatTensor, + causal: bool = True, + timestep: Optional[torch.Tensor] = None, +) -> torch.FloatTensor: + + batch_size = hidden_states.shape[0] + ada_values = self.scale_shift_table[None, ..., None, None, None].to( + timestep.device + ) + timestep.reshape( + batch_size, + 4, + -1, + timestep.shape[-3], + timestep.shape[-2], + timestep.shape[-1], + ) + shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1) + pixel_norm_inplace(hidden_states, scale1, shift1, 1e-9) + inplace_patch_conv_2(hidden_states, self.conv1.conv) + pixel_norm_inplace(hidden_states, scale2, shift2, 1e-9) + inplace_patch_conv_2(hidden_states, self.conv2.conv) + + +def block_forward( + self, + hidden_states: torch.FloatTensor, + causal: bool = True, + timestep: Optional[torch.Tensor] = None, +) -> torch.FloatTensor: + timestep_embed = None + if self.timestep_conditioning: + assert ( + timestep is not None + ), "should pass timestep with timestep_conditioning=True" + batch_size = hidden_states.shape[0] + timestep_embed = self.time_embedder( + timestep=timestep.flatten(), + resolution=None, + aspect_ratio=None, + batch_size=batch_size, + hidden_dtype=hidden_states.dtype, + ) + timestep_embed = timestep_embed.view( + batch_size, timestep_embed.shape[-1], 1, 1, 1 + ) + workspace = torch.empty( + hidden_states.shape[0], + hidden_states.shape[1], + hidden_states.shape[2] + 2, + hidden_states.shape[3], + hidden_states.shape[4], + device=hidden_states.device, + dtype=hidden_states.dtype, + memory_format=torch.channels_last_3d, + ) + for resnet in self.res_blocks: + workspace[:, :, 1:-1].copy_(hidden_states) + resnet(workspace, causal=causal, timestep=timestep_embed) + add_inplace(hidden_states, workspace, 1) + del workspace + torch.cuda.empty_cache() + return hidden_states + + +def unpatchify(x, patch_size_hw, patch_size_t=1): + if patch_size_hw == 1 and patch_size_t == 1: + return x + + if x.dim() == 4: + x = rearrange( + x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw + ) + elif x.dim() == 5: + x = rearrange( + x, + "b (c p r q) f h w -> b c (f p) (h q) (w r)", + p=patch_size_t, + q=patch_size_hw, + r=patch_size_hw, + ) + + return x + + +def upsample_forward(self, x, causal: bool = True): + if self.residual: + # Reshape and duplicate the input to match the output shape + x_in = rearrange( + x, + "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + num_repeat = np.prod(self.stride) // self.out_channels_reduction_factor + x_in = x_in.repeat(1, num_repeat, 1, 1, 1) + + workspace = torch.empty( + x.shape[0], + self.conv.conv.out_channels, + x.shape[2] + 2, + x.shape[3], + x.shape[4], + device=x.device, + dtype=x.dtype, + ) + workspace[:, : x.shape[1], 1:-1, :, :].copy_(x) + del x + torch.cuda.empty_cache() + inplace_patch_conv_2(workspace, self.conv.conv) + x = workspace[:, : self.conv.conv.out_channels, 1:-1, :, :] + + x = rearrange( + x, + "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + if self.residual: + x.add_(x_in) + del x_in + torch.cuda.empty_cache() + return x[:, :, 1:, :, :] + + +def decoder_forward( + self, + sample: torch.FloatTensor, + timestep: Optional[torch.Tensor] = None, +) -> torch.FloatTensor: + r"""The forward method of the `Decoder` class.""" + batch_size = sample.shape[0] + + sample = self.conv_in(sample, causal=self.causal) + + upscale_dtype = next(iter(self.up_blocks.parameters())).dtype + sample = sample.to(upscale_dtype) + + if self.timestep_conditioning: + assert ( + timestep is not None + ), "should pass timestep with timestep_conditioning=True" + scaled_timestep = (timestep * self.timestep_scale_multiplier).to(sample.device) + + for up_block in self.up_blocks: + if ( + self.timestep_conditioning + and up_block.__class__.__name__ == "UNetMidBlock3D" + ): + sample = up_block(sample, causal=self.causal, timestep=scaled_timestep) + else: + sample = up_block(sample, causal=self.causal) + + if self.timestep_conditioning: + embedded_timestep = self.last_time_embedder( + timestep=scaled_timestep.flatten(), + resolution=None, + aspect_ratio=None, + batch_size=sample.shape[0], + hidden_dtype=sample.dtype, + ) + embedded_timestep = embedded_timestep.view( + batch_size, embedded_timestep.shape[-1], 1, 1, 1 + ) + ada_values = self.last_scale_shift_table[None, ..., None, None, None].to( + embedded_timestep.device + ) + embedded_timestep.reshape( + batch_size, + 2, + -1, + embedded_timestep.shape[-3], + embedded_timestep.shape[-2], + embedded_timestep.shape[-1], + ) + shift, scale = ada_values.unbind(dim=1) + + workspace = torch.empty( + sample.shape[0], + sample.shape[1], + sample.shape[2] + 2, + sample.shape[3], + sample.shape[4], + device=sample.device, + dtype=sample.dtype, + memory_format=torch.channels_last_3d, + ) + workspace[:, :, 1:-1, :, :].copy_(sample) + del sample + torch.cuda.empty_cache() + + pixel_norm_inplace(workspace, scale, shift, 1e-9) + inplace_patch_conv_2(workspace, self.conv_out.conv) + + return unpatchify( + workspace[:, : self.conv_out.conv.out_channels, 1:-1], + patch_size_hw=self.patch_size, + patch_size_t=1, + ) diff --git a/vae_patcher/kernels/__init__.py b/vae_patcher/kernels/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/vae_patcher/kernels/add_inplace.cpp b/vae_patcher/kernels/add_inplace.cpp new file mode 100644 index 0000000..22cc80b --- /dev/null +++ b/vae_patcher/kernels/add_inplace.cpp @@ -0,0 +1,31 @@ +#include +#include +#include +#include + +#include + +void inplace_add_cuda( + at::Tensor& x, + at::Tensor& input2, + int offset, + cudaStream_t stream +); + + + +void inplace_add(at::Tensor &x, at::Tensor &workspace, int offset) { + at::cuda::CUDAGuard device_guard{x.get_device()}; + auto stream = at::cuda::getCurrentCUDAStream().stream(); + inplace_add_cuda( + x, + workspace, + offset, + stream + ); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("inplace_add", &inplace_add, + "inplace_add"); +} diff --git a/vae_patcher/kernels/add_inplace_cuda.cu b/vae_patcher/kernels/add_inplace_cuda.cu new file mode 100644 index 0000000..21518ee --- /dev/null +++ b/vae_patcher/kernels/add_inplace_cuda.cu @@ -0,0 +1,91 @@ +#include +#include // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK + +#include +#include +#include +#include +#include + +#include +#include + +using bfloat16 = at::BFloat16; +using accessor = at::PackedTensorAccessor64; +constexpr int BLOCK_ROWS = 8; + +template +__global__ void add_tensors_tiled(const accessor_t tensor1, + const accessor_t tensor2, + accessor_t result, + int offset +) { + int B = tensor1.size(0); + int C = tensor1.size(1); + int T = tensor1.size(2); + int H = tensor1.size(3); + int W = tensor1.size(4); + + const int outer_size = B * C; + const int inner_dim2 = H * W; + const int hw_ind = blockIdx.z * blockDim.x + threadIdx.x; + const int h = hw_ind / W; + const int w = hw_ind % W; + if (h >= H || w >= W) { + return; + } + const int t = blockIdx.y; + const int block_bc = blockIdx.x * BLOCK_ROWS; + + const int b = block_bc / C; + const int c = block_bc % C; + + for (int i = 0; i < BLOCK_ROWS; i++) { + result[b][c + i][t][h][w] = tensor1[b][c + i][t][h][w] + tensor2[b][c + i][t + offset][h][w]; + } +} + +void inplace_add_cuda( + at::Tensor& x, + at::Tensor& workspace, + int offset, + cudaStream_t stream +) { + + int B = x.size(0); + int C = x.size(1); + int T = x.size(2); + int H = x.size(3); + int W = x.size(4); + + int outer_size = B * (C / BLOCK_ROWS); + int inner_dim2 = H * W; + const int max_threads_per_block = 1024; + int threads_per_block_x = (inner_dim2 > max_threads_per_block) ? max_threads_per_block : inner_dim2; + int threads_per_block_y = (T > max_threads_per_block / threads_per_block_x) ? max_threads_per_block / threads_per_block_x : T; + + dim3 threadsPerBlock(threads_per_block_x, threads_per_block_y); + int max_v = (threadsPerBlock.x * threadsPerBlock.y); + int inner_dim = ((inner_dim2 + max_v - 1) / max_v); + + dim3 blocksPerGrid(outer_size, T, inner_dim); + if (x.numel() >= std::numeric_limits::max() || + workspace.numel() >= std::numeric_limits::max()) { + add_tensors_tiled<<>>( + x.packed_accessor64(), + workspace.packed_accessor64(), + x.packed_accessor64(), + offset + ); + } else { + add_tensors_tiled<<>>( + x.packed_accessor32(), + workspace.packed_accessor32(), + x.packed_accessor32(), + offset + ); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void inplace_add_cuda(at::Tensor&, at::Tensor&, int, cudaStream_t); diff --git a/vae_patcher/kernels/ops.py b/vae_patcher/kernels/ops.py new file mode 100644 index 0000000..aa812b5 --- /dev/null +++ b/vae_patcher/kernels/ops.py @@ -0,0 +1,28 @@ +import os + +from torch.utils.cpp_extension import load + +pixel_norm = load( + name="pixel_norm", + sources=[ + os.path.join(os.path.dirname(__file__), "pixel_norm.cpp"), + os.path.join(os.path.dirname(__file__), "pixel_norm_cuda.cu"), + ], +) + + +def pixel_norm_inplace(x, scale, shift, eps=1e-5): + return pixel_norm.pixel_norm_inplace(x, scale, shift, eps) + + +inplace_add = load( + name="inplace_add", + sources=[ + os.path.join(os.path.dirname(__file__), "add_inplace.cpp"), + os.path.join(os.path.dirname(__file__), "add_inplace_cuda.cu"), + ], +) + + +def add_inplace(x, workspace, offset): + return inplace_add.inplace_add(x, workspace, offset) diff --git a/vae_patcher/kernels/pixel_norm.cpp b/vae_patcher/kernels/pixel_norm.cpp new file mode 100644 index 0000000..783f5cf --- /dev/null +++ b/vae_patcher/kernels/pixel_norm.cpp @@ -0,0 +1,43 @@ +#include +#include +#include +#include +#include + +template +void pixel_norm_cuda( + const void* input, + const void* scale, + const void* shift, + void* output, + const float eps, + const int batch_size, + const int num_channels, + const int num_pixels, + cudaStream_t stream +); + +void pixel_norm_inplace(at::Tensor &x, at::Tensor &scale, at::Tensor &shift, float eps) { + int num_pixels = x.size(2)*x.size(3)*x.size(4); + int batch_size = x.size(0); + int num_channels = x.size(1); + at::cuda::CUDAGuard device_guard{(char)x.get_device()}; + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + pixel_norm_cuda( + x.data_ptr(), + scale.data_ptr(), + shift.data_ptr(), + x.data_ptr(), + eps, + batch_size, + num_channels, + num_pixels, + stream + ); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("pixel_norm_inplace", &pixel_norm_inplace, + "pixel norm inplace"); +} diff --git a/vae_patcher/kernels/pixel_norm_cuda.cu b/vae_patcher/kernels/pixel_norm_cuda.cu new file mode 100644 index 0000000..ac5212e --- /dev/null +++ b/vae_patcher/kernels/pixel_norm_cuda.cu @@ -0,0 +1,161 @@ + +#include +#include // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK + +#include +#include +#include + +template +struct pixel_norm_kernel_traits { + static constexpr int CHANNELS = CHANNELS_; + static constexpr int NUM_WARPS = NUM_WARPS_; + using input_t = input_t_; + using copy_t = uint4; + + static constexpr int NUM_THREADS = NUM_WARPS * 32; + static constexpr int NUM_ELEMS_PER_COPY = sizeof(copy_t) / sizeof(input_t); + static constexpr int NUM_THREADS_PER_PIXEL = CHANNELS / NUM_ELEMS_PER_COPY; + static constexpr int NUM_WARPS_PER_PIXEL = NUM_THREADS_PER_PIXEL / 32; + static constexpr int PIXELS_PER_BLOCK = NUM_THREADS / NUM_THREADS_PER_PIXEL; +}; + +template +inline __device__ void load_input(const input_t *x, float x_vals[NUM_ELEMS]) {; + input_t x_vals_load[NUM_ELEMS] = {0}; + reinterpret_cast(x_vals_load)[0] = reinterpret_cast(x)[threadIdx.x % COPY_THREADS]; + for (int i = 0; i < NUM_ELEMS; i++) { + x_vals[i] = float(x_vals_load[i]); + } +} + +inline __device__ float silu(float x) { + return x / (1.0f + expf(-x)); +} + +template +__global__ void pixel_norm_kernel( + const typename kernel_traits::input_t* input, + const typename kernel_traits::input_t* scale, + const typename kernel_traits::input_t* shift, + typename kernel_traits::input_t* output, + const float eps, + const int num_pixels, + const int batch_size +) { + extern __shared__ float smem_[]; + + const int tid = threadIdx.x; + const int pixel_id = blockIdx.x * kernel_traits::PIXELS_PER_BLOCK; + const int batch_id = pixel_id / num_pixels; + const int total_pixels = num_pixels * batch_size; + const int processing_pixels = pixel_id + tid / kernel_traits::NUM_THREADS_PER_PIXEL; + + if (processing_pixels >= total_pixels) { + return; + } + + float x_vals[kernel_traits::NUM_ELEMS_PER_COPY]; + float sum = 0.0f; + + const typename kernel_traits::input_t* x_ptr = input + static_cast(pixel_id) * static_cast(kernel_traits::CHANNELS); + typename kernel_traits::input_t* output_ptr = output + static_cast(pixel_id) * static_cast(kernel_traits::CHANNELS); + + const typename kernel_traits::input_t* scale_ptr = scale + batch_id * kernel_traits::CHANNELS; + const typename kernel_traits::input_t* shift_ptr = shift + batch_id * kernel_traits::CHANNELS; + + float scale_vals[kernel_traits::NUM_ELEMS_PER_COPY]; + float shift_vals[kernel_traits::NUM_ELEMS_PER_COPY]; + + load_input(x_ptr, x_vals); + load_input(scale_ptr, scale_vals); + load_input(shift_ptr, shift_vals); + + for (int i = 0; i < kernel_traits::NUM_ELEMS_PER_COPY; i++) { + sum += x_vals[i] * x_vals[i]; + } + + for (int i = 1; i < kernel_traits::NUM_THREADS_PER_PIXEL; i *= 2) { + sum += __shfl_xor_sync(uint32_t(-1), sum, i); + } + if constexpr (kernel_traits::NUM_WARPS_PER_PIXEL > 1) { + + smem_[tid / 32] = sum; + __syncthreads(); + sum = 0.0f; + for (int i = 0; i < kernel_traits::NUM_WARPS_PER_PIXEL; i++) { + sum += smem_[i]; + } + } + + sum = sum / kernel_traits::CHANNELS; + + float inv_sqrt_sum = rsqrtf(sum + eps); + #pragma unroll + for (int i = 0; i < kernel_traits::NUM_ELEMS_PER_COPY; i++) { + x_vals[i] *= inv_sqrt_sum; + x_vals[i] = (1.0f + scale_vals[i]) * x_vals[i] + shift_vals[i]; + x_vals[i] = silu(x_vals[i]); + } + typename kernel_traits::input_t out_vals_store[kernel_traits::NUM_ELEMS_PER_COPY]; + for (int i = 0; i < kernel_traits::NUM_ELEMS_PER_COPY; i++) { + out_vals_store[i] = static_cast(x_vals[i]); + } + reinterpret_cast(output_ptr)[tid] = reinterpret_cast(out_vals_store)[0]; +} + +template +void pixel_norm_launch( + const void* input, + const void* scale, + const void* shift, + void* output, + const float eps, + const int batch_size, + const int num_pixels, + cudaStream_t stream +) { + + using kernel_traits = pixel_norm_kernel_traits; + int BLOCK_SIZE = (num_pixels*batch_size + kernel_traits::PIXELS_PER_BLOCK - 1) / kernel_traits::PIXELS_PER_BLOCK; + + dim3 block(kernel_traits::NUM_THREADS); + dim3 grid(BLOCK_SIZE); + + auto kernel = &pixel_norm_kernel; + kernel<<>>( + (const typename kernel_traits::input_t*)input, + (const typename kernel_traits::input_t*)scale, + (const typename kernel_traits::input_t*)shift, + (typename kernel_traits::input_t*)output, + eps, + num_pixels, + batch_size + ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void pixel_norm_cuda( + const void* input, + const void* scale, + const void* shift, + void* output, + const float eps, + const int batch_size, + const int num_channels, + const int num_pixels, + cudaStream_t stream +) { + if (num_channels == 128) { + pixel_norm_launch<2, 128, input_t, input_t>(input, scale, shift, output, eps, batch_size, num_pixels, stream); + } else if (num_channels == 256) { + pixel_norm_launch<4, 256, input_t, input_t>(input, scale, shift, output, eps, batch_size, num_pixels, stream); + } else if (num_channels == 512) { + pixel_norm_launch<4, 512, input_t, input_t>(input, scale, shift, output, eps, batch_size, num_pixels, stream); + } else if (num_channels == 1024) { + pixel_norm_launch<8, 1024, input_t, input_t>(input, scale, shift, output, eps, batch_size, num_pixels, stream); + } +} + +template void pixel_norm_cuda(const void*, const void*, const void*, void*, float, int, int, int, cudaStream_t); diff --git a/vae_patcher/patch.py b/vae_patcher/patch.py new file mode 100644 index 0000000..28ed289 --- /dev/null +++ b/vae_patcher/patch.py @@ -0,0 +1,35 @@ +import types + +import torch + +from .forwards import ( + block_forward, + decoder_forward, + res_block_forward, + upsample_forward, +) + + +def patch_vae(vae_model, patch_block=4): + + vae_model.decoder.forward = types.MethodType(decoder_forward, vae_model.decoder) + vae_model.decoder.conv_out.conv.weight.data = ( + vae_model.decoder.conv_out.conv.weight.data.to( + memory_format=torch.channels_last_3d + ) + ) + + for i, block in enumerate(reversed(vae_model.decoder.up_blocks)): + if i >= patch_block: + break + if block.__class__.__name__ == "UNetMidBlock3D": + block.forward = types.MethodType(block_forward, block) + for name, param in block.named_parameters(): + if "conv1" in name or "conv2" in name: + if "weight" in name: + param = param.to(memory_format=torch.channels_last_3d) + for res_block in block.res_blocks: + if res_block.__class__.__name__ == "ResnetBlock3D": + res_block.forward = types.MethodType(res_block_forward, res_block) + elif block.__class__.__name__ == "DepthToSpaceUpsample": + block.forward = types.MethodType(upsample_forward, block) diff --git a/vae_patcher/vae_patcher.py b/vae_patcher/vae_patcher.py new file mode 100644 index 0000000..7f75d35 --- /dev/null +++ b/vae_patcher/vae_patcher.py @@ -0,0 +1,38 @@ +from comfy.model_management import dtype_size + +from ..nodes_registry import comfy_node + + +@comfy_node(name="LTXVPatcherVAE") +class LTXVPatcherVAE: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "vae": ("VAE",), + } + } + + RETURN_TYPES = ("VAE",) + FUNCTION = "patch" + CATEGORY = "lightricks/LTXV" + TITLE = "LTXV VAE Patcher" + + def patch(self, vae): + from .patch import patch_vae + + vae_model = vae.first_stage_model + vae.memory_used_decode = ( + lambda shape, dtype: shape[1] + * shape[2] + * shape[3] + * shape[4] + * 8 + * 8 + * 8 + * 3 + * dtype_size(dtype) + ) + + patch_vae(vae_model, patch_block=4) + return (vae,)