diff --git a/.gitattributes b/.gitattributes index de46872..34d8762 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1,2 @@ -assets/video.mp4 filter=lfs diff=lfs merge=lfs -text +**/*.png filter=lfs diff=lfs merge=lfs -text +**/*.mp4 filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..82f9275 --- /dev/null +++ b/.gitignore @@ -0,0 +1,162 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/latest/usage/project/#working-with-version-control +.pdm.toml +.pdm-python +.pdm-build/ + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..7c0a6f7 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,28 @@ +# general formatting +repos: +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.6.0 + hooks: + - id: check-yaml + - id: end-of-file-fixer + - id: trailing-whitespace + - id: requirements-txt-fixer + - id: check-added-large-files + - id: check-toml +# Isort (import sorting) +- repo: https://github.com/PyCQA/isort + rev: 5.13.2 + hooks: + - id: isort + args: [--profile, black] +# Black (code formatting) +- repo: https://github.com/psf/black + rev: 24.4.2 # Replace by any tag/version: https://github.com/psf/black/tags + hooks: + - id: black + language_version: python3.10 +# flake8 linter +- repo: https://github.com/charliermarsh/ruff-pre-commit + rev: 'v0.4.4' + hooks: + - id: ruff diff --git a/__init__.py b/__init__.py index e5d0e39..94fc28f 100644 --- a/__init__.py +++ b/__init__.py @@ -1,8 +1,7 @@ +from .loader_node import LTXVLoader # noqa: F401 from .nodes_registry import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS - -from .transformer import LTXVModelConfigurator, LTXVShiftSigmas -from .t5_encoder import LTXVCLIPModelLoader -from .loader_node import LTXVLoader +from .t5_encoder import LTXVCLIPModelLoader # noqa: F401 +from .transformer import LTXVModelConfigurator, LTXVShiftSigmas # noqa: F401 # Export so that ComfyUI can pick them up. -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/assets/baloons.png b/assets/baloons.png index 65fca41..1e5b422 100644 Binary files a/assets/baloons.png and b/assets/baloons.png differ diff --git a/assets/ltxvideo-i2v.json b/assets/ltxvideo-i2v.json index f9b6092..eb9f43b 100644 --- a/assets/ltxvideo-i2v.json +++ b/assets/ltxvideo-i2v.json @@ -816,4 +816,4 @@ } }, "version": 0.4 -} \ No newline at end of file +} diff --git a/assets/ltxvideo-i2v.png b/assets/ltxvideo-i2v.png index eea1053..d315d24 100644 Binary files a/assets/ltxvideo-i2v.png and b/assets/ltxvideo-i2v.png differ diff --git a/assets/ltxvideo-t2v.json b/assets/ltxvideo-t2v.json index 6b8fe1d..2b1cef5 100644 --- a/assets/ltxvideo-t2v.json +++ b/assets/ltxvideo-t2v.json @@ -770,4 +770,4 @@ } }, "version": 0.4 -} \ No newline at end of file +} diff --git a/assets/ltxvideo-t2v.png b/assets/ltxvideo-t2v.png index 9d4ca60..e132ab3 100644 Binary files a/assets/ltxvideo-t2v.png and b/assets/ltxvideo-t2v.png differ diff --git a/assets/ltxvideo-v2v.json b/assets/ltxvideo-v2v.json index 7252170..5ea4706 100644 --- a/assets/ltxvideo-v2v.json +++ b/assets/ltxvideo-v2v.json @@ -911,4 +911,4 @@ } }, "version": 0.4 -} \ No newline at end of file +} diff --git a/assets/ltxvideo-v2v.png b/assets/ltxvideo-v2v.png index 78c7c76..45aefb2 100644 Binary files a/assets/ltxvideo-v2v.png and b/assets/ltxvideo-v2v.png differ diff --git a/img2vid.py b/img2vid.py index 343779f..0bcfa60 100644 --- a/img2vid.py +++ b/img2vid.py @@ -1,14 +1,11 @@ +import comfy.latent_formats import comfy.model_base import comfy.model_management -import comfy.supported_models_base -import comfy.utils import comfy.model_patcher import comfy.sd -import comfy.latent_formats -import nodes +import comfy.supported_models_base +import comfy.utils import torch -from .nodes_registry import comfy_node - from ltx_video.models.autoencoders.vae_encode import get_vae_size_scale_factor diff --git a/loader_node.py b/loader_node.py index 496e7bb..1c6cd6e 100644 --- a/loader_node.py +++ b/loader_node.py @@ -1,22 +1,20 @@ from pathlib import Path -import safetensors.torch -import torch import comfy import comfy.model_management import comfy.model_patcher import folder_paths - - -from ltx_video.models.transformers.transformer3d import Transformer3DModel -from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder -from ltx_video.models.autoencoders.vae_encode import get_vae_size_scale_factor +import safetensors.torch +import torch +from ltx_video.models.autoencoders.causal_video_autoencoder import ( + CausalVideoAutoencoder, +) from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier +from ltx_video.models.transformers.transformer3d import Transformer3DModel -from .model import LTXVTransformer3D, LTXVModel, LTXVModelConfig -from .vae import LTXVVAE - +from .model import LTXVModel, LTXVModelConfig, LTXVTransformer3D from .nodes_registry import comfy_node +from .vae import LTXVVAE @comfy_node(name="LTXVLoader") @@ -53,7 +51,11 @@ class LTXVLoader: num_latent_channels = vae.first_stage_model.config.latent_channels model = self._load_unet( - load_device, offload_device, weights, num_latent_channels, dtype=dtype_map[dtype] + load_device, + offload_device, + weights, + num_latent_channels, + dtype=dtype_map[dtype], ) return (model, vae) @@ -94,7 +96,9 @@ class LTXVLoader: ) return vae - def _load_unet(self, load_device, offload_device, weights, num_latent_channels, dtype): + def _load_unet( + self, load_device, offload_device, weights, num_latent_channels, dtype + ): config = { "_class_name": "Transformer3DModel", "_diffusers_version": "0.25.1", diff --git a/model.py b/model.py index 555d26b..44e86c3 100644 --- a/model.py +++ b/model.py @@ -1,20 +1,18 @@ +import math from contextlib import nullcontext -import torch -import torch.nn as nn + +import comfy.latent_formats import comfy.model_base import comfy.model_management +import comfy.model_patcher +import comfy.model_sampling +import comfy.sd import comfy.supported_models_base import comfy.utils -import comfy.model_patcher -import comfy.sd -import comfy.latent_formats -import comfy.model_sampling -from comfy.ldm.modules.diffusionmodules.util import make_beta_schedule -import math - -from ltx_video.models.transformers.transformer3d import Transformer3DModel +import torch +import torch.nn as nn from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier -from ltx_video.models.autoencoders.vae_encode import get_vae_size_scale_factor +from ltx_video.models.transformers.transformer3d import Transformer3DModel class LTXVModelConfig: @@ -109,14 +107,15 @@ class LTXVTransformer3D(nn.Module): latent_patchified = self.patchifier.patchify(latent) context_mask = (context != 0).any(dim=2).to(self.transformer.dtype) - l = latent_patchified if mixed_precision: context_manager = torch.autocast("cuda", dtype=torch.bfloat16) else: context_manager = nullcontext() with context_manager: noise_pred = self.transformer( - l.to(self.transformer.dtype).to(self.transformer.device), + latent_patchified.to(self.transformer.dtype).to( + self.transformer.device + ), indices_grid.to(self.transformer.device), encoder_hidden_states=context.to(self.transformer.device), encoder_attention_mask=context_mask.to(self.transformer.device).to( @@ -145,6 +144,8 @@ class LTXVTransformer3DWrapper(nn.Module): indices_grid, ): super().__init__() + self.generator = torch.Generator(transformer.transformer.device).manual_seed(42) + self.indices_grid = indices_grid self.dtype = transformer.dtype self.wrapped_transformer = transformer @@ -157,12 +158,23 @@ class LTXVTransformer3DWrapper(nn.Module): def forward(self, x, timesteps, context, img_hw=None, aspect_ratio=None, **kwargs): transformer_options = kwargs.get("transformer_options", {}) mixed_precision = transformer_options.get("mixed_precision", False) + noise_scale = transformer_options.get("noise_scale", 0.15) mask = self.conditioning_mask_patchified.to(x.device) ndim_mask = mask.ndimension() expanded_timesteps = timesteps.view(timesteps.size(0), *([1] * (ndim_mask - 1))) timesteps_masked = expanded_timesteps * (1 - mask) + + noise = torch.randn(size=x.shape, device=x.device, generator=self.generator) + + timesteps_unsqeezed = timesteps + for _ in range(x.dim() - timesteps.dim()): + timesteps_unsqeezed = timesteps_unsqeezed.unsqueeze(-1) + latent = ( + x + noise_scale * noise * (timesteps_unsqeezed**2) + ) * self.conditioning_mask + x * (1 - self.conditioning_mask) + result = self.wrapped_transformer.forward( - x, + latent, timesteps_masked, context, self.indices_grid, diff --git a/nodes_registry.py b/nodes_registry.py index 598c077..3e4b6de 100644 --- a/nodes_registry.py +++ b/nodes_registry.py @@ -1,8 +1,6 @@ import re from typing import Callable, Optional, Type -from easy_nodes import initialize_easy_nodes - NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} @@ -11,8 +9,6 @@ EXPERIMENTAL_DISPLAY_NAME_PREFIX = "(Experimental 🧪)" DEPRECATED_DISPLAY_NAME_PREFIX = "(Deprecated 🚫)" DEFAULT_CATEGORY_NAME = "Lightricks" -initialize_easy_nodes(default_category=DEFAULT_CATEGORY_NAME, auto_register=True) - def register_node(node_class: Type, name: str, description: str) -> None: """ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..cf5e632 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "comfyui-ltxvideo" +version = "0.1.0" +description = "Custom nodes for LTX-Video support in ComfyUI" +authors = [ + { name = "Andrew Kvochko", email = "akvochko@lightricks.com" } +] +requires-python = ">=3.10" +readme = "README.md" +license = { file = "LICENSE" } +dependencies = [ + "ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9", +] + +[project.optional-dependencies] +dev = [ + "pre-commit>=4.0.1", + "pytest>=8.0.0", + "websocket-client==1.6.1", + "scikit-image==0.24.0" +] diff --git a/requirements.txt b/requirements.txt index 2dd40c8..b21e0f8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1 @@ -comfyui-easynodes -ltx-video@git+https://github.com/Lightricks/LTX-Video.git@ltx-video-0.9 -huggingface_hub==0.25.2 +ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9 diff --git a/t5_encoder.py b/t5_encoder.py index 9ba6542..571b850 100644 --- a/t5_encoder.py +++ b/t5_encoder.py @@ -1,19 +1,20 @@ -import torch +from pathlib import Path + import comfy.clip_model +import comfy.latent_formats import comfy.model_base import comfy.model_management +import comfy.model_patcher +import comfy.sd import comfy.sd1_clip import comfy.supported_models_base import comfy.utils -import comfy.model_patcher -import comfy.sd -import comfy.latent_formats import folder_paths -from pathlib import Path -from .nodes_registry import comfy_node - +import torch from transformers import T5EncoderModel, T5Tokenizer +from .nodes_registry import comfy_node + class LTXVTokenizer(comfy.sd1_clip.SDTokenizer): def __init__(self, tokenizer_path: str): @@ -84,7 +85,7 @@ class LTXVTextEncoderModel(torch.nn.Module): [[w[1] for w in token_weight_pairs_t5]], device=self.t5xxl.device, ) - self.to(self.t5xxl.device) # comfyui skips loading some weights to gpu + self.to(self.t5xxl.device) # comfyui skips loading some weights to gpu out = self.t5xxl(text_input_ids, attention_mask=prompt_attention_mask)[0] out = out * prompt_attention_mask.unsqueeze(2) return out, None, {"attention_mask": prompt_attention_mask} @@ -119,10 +120,14 @@ def ltxv_tokenizer(tokenizer_path): class LTXVCLIPModelLoader: @classmethod def INPUT_TYPES(s): - return {"required": {"clip_path": ( - folder_paths.get_filename_list("text_encoders"), - {"tooltip": "The name of the text encoder model to load."}, - )}} + return { + "required": { + "clip_path": ( + folder_paths.get_filename_list("text_encoders"), + {"tooltip": "The name of the text encoder model to load."}, + ) + } + } RETURN_TYPES = ("CLIP",) RETURN_NAMES = ("clip",) diff --git a/transformer.py b/transformer.py index f389a3d..805e6c3 100644 --- a/transformer.py +++ b/transformer.py @@ -1,29 +1,19 @@ -import torch +import math + +import comfy.latent_formats import comfy.model_base import comfy.model_management -import comfy.supported_models_base -import comfy.utils import comfy.model_patcher import comfy.sd -import comfy.latent_formats -from comfy.ldm.modules.diffusionmodules.util import make_beta_schedule -import folder_paths -import math -from pathlib import Path -from .nodes_registry import comfy_node - -from ltx_video.models.transformers.transformer3d import Transformer3DModel -from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier +import comfy.supported_models_base +import comfy.utils +import torch from ltx_video.models.autoencoders.vae_encode import get_vae_size_scale_factor +from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier -from .model import ( - LTXVTransformer3DWrapper, - LTXVTransformer3D, - LTXVModel, - LTXVModelConfig, - LTXVSampling, -) from .img2vid import encode_media_conditioning +from .model import LTXVSampling, LTXVTransformer3DWrapper +from .nodes_registry import comfy_node def get_normal_shift( @@ -85,12 +75,36 @@ class LTXVModelConfigurator: ), "width": ("INT", {"default": 768, "min": 1, "max": 10000}), "height": ("INT", {"default": 512, "min": 1, "max": 10000}), - "frames_number": ("INT", {"default": 65, "min": 9, "max": 257, "step": 8, "tooltip": "Must be equal to N * 8 + 1"}), + "frames_number": ( + "INT", + { + "default": 65, + "min": 9, + "max": 257, + "step": 8, + "tooltip": "Must be equal to N * 8 + 1", + }, + ), "frame_rate": ("INT", {"default": 25, "min": 1, "max": 60}), "batch": ("INT", {"default": 1, "min": 1, "max": 60}), "mixed_precision": ("BOOLEAN", {"default": True}), + "noise_scale": ( + "FLOAT", + { + "default": 0.15, + "min": 0.0, + "max": 100, + "step": 0.01, + "tooltip": "Amount of noise to apply on conditioning image latent.", + }, + ), + }, + "optional": { + "conditioning": ( + "IMAGE", + {"tooltip": "Optional conditioning image or video."}, + ) }, - "optional": {"conditioning": ("IMAGE",{"tooltip": "Optional conditioning image or video."})}, } RETURN_TYPES = ("MODEL", "LATENT", "FLOAT") @@ -160,10 +174,10 @@ class LTXVModelConfigurator: frame_rate, batch, mixed_precision, + noise_scale, conditioning=None, ): load_device = comfy.model_management.get_torch_device() - offload_device = comfy.model_management.unet_offload_device() if preset != "Custom": preset = preset.split("|") width, height = map(int, preset[0].strip().split("x")) @@ -189,7 +203,7 @@ class LTXVModelConfigurator: patchifier = SymmetricPatchifier(1) _, vae_scale_factor, _ = get_vae_size_scale_factor(vae.first_stage_model) - use_rope = model.model.diffusion_model.transformer.use_rope + use_rope = model.get_model_object("diffusion_model").transformer.use_rope indices_grid = self.indices_grid( latent_shape, vae_scale_factor, @@ -199,26 +213,23 @@ class LTXVModelConfigurator: load_device, ) wrapper = LTXVTransformer3DWrapper( - transformer=model.model.diffusion_model, + transformer=model.get_model_object("diffusion_model"), patchifier=patchifier, conditioning_mask=conditioning_mask, indices_grid=indices_grid, ) - model.model.diffusion_model.to(load_device) - new_model = LTXVModel( - LTXVModelConfig(vae.first_stage_model.config.latent_channels, model.model.model_config.manual_cast_dtype), - model_type=comfy.model_base.ModelType.FLOW, - device=comfy.model_management.get_torch_device(), - ) - new_model.model_sampling = LTXVSampling(wrapper.conditioning_mask) - new_model.diffusion_model = wrapper - patcher = comfy.model_patcher.ModelPatcher( - new_model, load_device, offload_device + patcher = model.clone() + patcher.add_object_patch("diffusion_model", wrapper) + patcher.add_object_patch( + "model_sampling", LTXVSampling(wrapper.conditioning_mask) ) patcher.model_options.setdefault("transformer_options", {})[ "mixed_precision" ] = mixed_precision + patcher.model_options.setdefault("transformer_options", {})["noise_scale"] = ( + noise_scale if conditioning is not None else 0.0 + ) num_latent_patches = latent_shape[2] * latent_shape[3] * latent_shape[4] return (patcher, {"samples": latent}, get_normal_shift(num_latent_patches)) @@ -232,15 +243,21 @@ class LTXVShiftSigmas: "required": { "sigmas": ("SIGMAS",), "sigma_shift": ("FLOAT", {"default": 1.820833333}), - "stretch": ("BOOLEAN", { - "default": True, - "tooltip": "Stretch the sigmas to be in the range [terminal, 1]." - }), + "stretch": ( + "BOOLEAN", + { + "default": True, + "tooltip": "Stretch the sigmas to be in the range [terminal, 1].", + }, + ), "terminal": ( "FLOAT", { - "default": 0.1, "min": 0.0, "max": 0.99, "step": 0.01, - "tooltip": "The terminal value of the sigmas after stretching." + "default": 0.1, + "min": 0.0, + "max": 0.99, + "step": 0.01, + "tooltip": "The terminal value of the sigmas after stretching.", }, ), } @@ -250,7 +267,9 @@ class LTXVShiftSigmas: CATEGORY = "lightricks/LTXV" FUNCTION = "shift_sigmas" - DESCRIPTION = "Transforms sigmas to values where the model can focus on denoising high noise." + DESCRIPTION = ( + "Transforms sigmas to values where the model can focus on denoising high noise." + ) def shift_sigmas(self, sigmas, sigma_shift, stretch, terminal): power = 1 diff --git a/vae.py b/vae.py index 0a39c20..6adfd05 100644 --- a/vae.py +++ b/vae.py @@ -1,24 +1,17 @@ -import torch +import comfy.latent_formats import comfy.model_base import comfy.model_management -import comfy.supported_models_base -import comfy.utils import comfy.model_patcher import comfy.sd -import comfy.latent_formats -import folder_paths -from pathlib import Path -from .nodes_registry import comfy_node - -from ltx_video.models.autoencoders.video_autoencoder import VideoAutoencoder -from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder +import comfy.supported_models_base +import comfy.utils +import torch +from diffusers.image_processor import VaeImageProcessor from ltx_video.models.autoencoders.vae_encode import ( get_vae_size_scale_factor, vae_decode, vae_encode, ) -from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier -from diffusers.image_processor import VaeImageProcessor class LTXVVAE(comfy.sd.VAE):