Automated PR - 2024-12-03

.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

.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/

.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

__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"]

assets/ltxvideo-i2v.json

@@ -816,4 +816,4 @@
     }
   },
   "version": 0.4
-}
+}

assets/ltxvideo-t2v.json

@@ -770,4 +770,4 @@
     }
   },
   "version": 0.4
-}
+}

assets/ltxvideo-v2v.json

@@ -911,4 +911,4 @@
     }
   },
   "version": 0.4
-}
+}

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
 
 

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",

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,

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:
     """

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"
+]

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

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",)

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

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):