nodes/Lightricks-ComfyUI-LTXVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Automated PR - 2024-12-19

Browse Source
This commit is contained in:
sync-bot
2024-12-19 17:34:13 +00:00
parent 366c43d754
commit 5c5dbecdfa
19 changed files with 2793 additions and 2012 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+13 -14
View File
@@ -1,6 +1,13 @@
# ComfyUI-LTXVideo
ComfyUI-LTXVideo is a collection of custom nodes for ComfyUI designed to integrate the LTXVideo diffusion model. These nodes enable workflows for text-to-video, image-to-video, and video-to-video generation. The main LTXVideo repository can be found [here](https://github.com/Lightricks/LTX-Video).
## 19.12.2024 ⭐ Update ⭐
1. Improved model - removes "strobing texture" artifacts and generates better motion. Download from [here](https://huggingface.co/Lightricks/LTX-Video/resolve/main/ltx-video-2b-v0.9.1.safetensors).
2. STG support
3. Integrated image degradation system for improved motion generation.
3. Additional initial latent optional input to chain latents for high res generation.
4. Image captioning in image to video [flow](assets/ltxvideo-i2v.json).
## Installation
Installation via [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager) is preferred. Simply search for `ComfyUI-LTXVideo` in the list of nodes and follow installation instructions.
@@ -20,27 +27,19 @@ For portable ComfyUI installations, run
### Models
1. Download [ltx-video-2b-v0.9.safetensors](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltx-video-2b-v0.9.safetensors) from Hugging Face and place it under `models/checkpoints`.
2. [Install git-lfs](https://docs.github.com/en/repositories/working-with-files/managing-large-files/installing-git-large-file-storage) and clone the text encoder model to `models/text_encoders`:
```bash
cd models/text_encoders && git clone https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS
```
1. Download [ltx-video-2b-v0.9.1.safetensors](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltx-video-2b-v0.9.1.safetensors) from Hugging Face and place it under `models/checkpoints`.
2. Install one of the t5 text encoders, for example [google_t5-v1_1-xxl_encoderonly](https://huggingface.co/mcmonkey/google_t5-v1_1-xxl_encoderonly/tree/main). You can install it using ComfyUI Model Manager.
## Example workflows
Note that to run the example workflows, you need to have [ComfyUI-VideoHelperSuite](https://github.com/kosinkadink/ComfyUI-VideoHelperSuite) installed.
### Text-to-video
[Download workflow](assets/ltxvideo-t2v.json)
![workflow](assets/ltxvideo-t2v.png)
Note that to run the example workflows, you need to have some additional custom nodes, like [ComfyUI-VideoHelperSuite](https://github.com/kosinkadink/ComfyUI-VideoHelperSuite) and others, installed. You can do it by pressing "Install Missing Custom Nodes" button in ComfyUI Manager.
### Image-to-video
[Download workflow](assets/ltxvideo-i2v.json)
![workflow](assets/ltxvideo-i2v.png)
### Video-to-video
### Text-to-video
[Download workflow](assets/ltxvideo-v2v.json)
![workflow](assets/ltxvideo-v2v.png)
[Download workflow](assets/ltxvideo-t2v.json)
![workflow](assets/ltxvideo-t2v.png)
__init__.py
+1
View File
@@ -1,5 +1,6 @@
from .loader_node import LTXVLoader # noqa: F401
from .nodes_registry import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
from .stg import LTXVApplySTG, STGGuiderNode # noqa: F401
from .t5_encoder import LTXVCLIPModelLoader # noqa: F401
from .transformer import LTXVModelConfigurator, LTXVShiftSigmas # noqa: F401
assets/baloons.png
-3
View File
@@ -1,3 +0,0 @@
version https://git-lfs.github.com/spec/v1
oid sha256:57f8f573779ebb37ca2c2bd1377e430d4e9abce0e728cf72228687b72c455244
size 699352
assets/ltxvideo-i2v.json
+1235 -455
View File
File diff suppressed because it is too large Load Diff
assets/ltxvideo-i2v.png
+2 -2
View File
@@ -1,3 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:f247a5e042f13afcf0e6a88ea94e4187b6eb079d904ef6f54b7aeda47710c7ac
size 711224
oid sha256:6ec16f1372ece0cd63395edf29d9997599580df3703d7fd30025fd0870ad5532
size 1205996
assets/ltxvideo-t2v.json
+995 -455
View File
File diff suppressed because it is too large Load Diff
assets/ltxvideo-t2v.png
+2 -2
View File
@@ -1,3 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:156f19cbe09b2510a4c86803e7dc07e510643f2ff6d709bfd821cf0cceb59e8c
size 494899
oid sha256:1e508763410d009f524ed22c1c420e798904afe92098130a4d155d8318ef6b3a
size 1112804
assets/ltxvideo-v2v.json
-914
View File
@@ -1,914 +0,0 @@
{
"last_node_id": 110,
"last_link_id": 271,
"nodes": [
{
"id": 105,
"type": "LTXVCLIPModelLoader",
"pos": [
-458,
-122
],
"size": [
314.85308837890625,
60.962398529052734
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "clip",
"type": "CLIP",
"links": [
264,
265
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "LTXVCLIPModelLoader"
},
"widgets_values": [
"PixArt-XL-2-1024-MS/text_encoder/model-00001-of-00002.safetensors"
]
},
{
"id": 87,
"type": "CLIPTextEncode",
"pos": [
-27,
-245
],
"size": [
400,
200
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 264
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
215
],
"slot_index": 0,
"shape": 3
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"A group of colorful hot air balloons take off at dawn in Cappadocia, Turkey.\nDozens of balloons in various bright colors and patterns slowly rise into the pink and orange sky. Below them, the unique landscape of Cappadocia unfolds, with its distinctive \"fairy chimneys\" - tall, cone-shaped rock formations scattered across the valley. The rising sun casts long shadows across the terrain, highlighting the otherworldly topography."
]
},
{
"id": 36,
"type": "SamplerCustomAdvanced",
"pos": [
1107,
-259
],
"size": [
355.20001220703125,
106
],
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "noise",
"type": "NOISE",
"link": 72
},
{
"name": "guider",
"type": "GUIDER",
"link": 87,
"slot_index": 1
},
{
"name": "sampler",
"type": "SAMPLER",
"link": 70
},
{
"name": "sigmas",
"type": "SIGMAS",
"link": 270
},
{
"name": "latent_image",
"type": "LATENT",
"link": 258
}
],
"outputs": [
{
"name": "output",
"type": "LATENT",
"links": [
78
],
"slot_index": 0,
"shape": 3
},
{
"name": "denoised_output",
"type": "LATENT",
"links": [],
"slot_index": 1,
"shape": 3
}
],
"properties": {
"Node name for S&R": "SamplerCustomAdvanced"
},
"widgets_values": []
},
{
"id": 14,
"type": "KSamplerSelect",
"pos": [
768,
-132
],
"size": [
315,
58
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "SAMPLER",
"type": "SAMPLER",
"links": [
70
],
"slot_index": 0,
"shape": 3
}
],
"properties": {
"Node name for S&R": "KSamplerSelect"
},
"widgets_values": [
"euler"
]
},
{
"id": 81,
"type": "BasicScheduler",
"pos": [
409,
-10
],
"size": [
315,
106
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 261
}
],
"outputs": [
{
"name": "SIGMAS",
"type": "SIGMAS",
"links": [
262
],
"slot_index": 0,
"shape": 3
}
],
"properties": {
"Node name for S&R": "BasicScheduler"
},
"widgets_values": [
"normal",
20,
1
]
},
{
"id": 8,
"type": "VAEDecode",
"pos": [
1491,
-252
],
"size": [
210,
46
],
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 78
},
{
"name": "vae",
"type": "VAE",
"link": 256,
"slot_index": 1
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
121
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "VAEDecode"
},
"widgets_values": []
},
{
"id": 103,
"type": "LTXVModelConfigurator",
"pos": [
-22,
257
],
"size": [
378,
242
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 254
},
{
"name": "vae",
"type": "VAE",
"link": 255
},
{
"name": "conditioning",
"type": "IMAGE",
"link": 271,
"shape": 7
}
],
"outputs": [
{
"name": "model",
"type": "MODEL",
"links": [
260,
261
],
"slot_index": 0
},
{
"name": "latent",
"type": "LATENT",
"links": [
258
],
"slot_index": 1
},
{
"name": "sigma_shift",
"type": "FLOAT",
"links": [
259
],
"slot_index": 2
}
],
"properties": {
"Node name for S&R": "LTXVModelConfigurator"
},
"widgets_values": [
"Custom",
768,
512,
65,
25,
1,
true
]
},
{
"id": 60,
"type": "VHS_VideoCombine",
"pos": [
1197,
-104
],
"size": [
315,
520.6666870117188
],
"flags": {},
"order": 14,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 121
},
{
"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,
"shape": 3
}
],
"properties": {
"Node name for S&R": "VHS_VideoCombine"
},
"widgets_values": {
"frame_rate": 25,
"loop_count": 0,
"filename_prefix": "LTXVideo",
"format": "video/h264-mp4",
"pix_fmt": "yuv420p",
"crf": 19,
"save_metadata": true,
"pingpong": false,
"save_output": true,
"videopreview": {
"hidden": false,
"paused": false,
"params": {
"filename": "LTXVideo_00002.mp4",
"subfolder": "",
"type": "output",
"format": "video/h264-mp4",
"frame_rate": 25
}
}
}
},
{
"id": 88,
"type": "CLIPTextEncode",
"pos": [
-29,
5
],
"size": [
400,
200
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 265
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
214
],
"slot_index": 0,
"shape": 3
}
],
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"worst quality, inconsistent motion, blurry, jittery, distorted, watermarks"
]
},
{
"id": 104,
"type": "LTXVShiftSigmas",
"pos": [
759,
-19
],
"size": [
315,
106
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "sigmas",
"type": "SIGMAS",
"link": 262
},
{
"name": "sigma_shift",
"type": "FLOAT",
"link": 259,
"widget": {
"name": "sigma_shift"
}
}
],
"outputs": [
{
"name": "SIGMAS",
"type": "SIGMAS",
"links": [
269
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "LTXVShiftSigmas"
},
"widgets_values": [
1.820833333,
true,
0.1
]
},
{
"id": 109,
"type": "SplitSigmas",
"pos": [
761,
137
],
"size": [
315,
78
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "sigmas",
"type": "SIGMAS",
"link": 269
}
],
"outputs": [
{
"name": "high_sigmas",
"type": "SIGMAS",
"links": null
},
{
"name": "low_sigmas",
"type": "SIGMAS",
"links": [
270
],
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "SplitSigmas"
},
"widgets_values": [
1
]
},
{
"id": 42,
"type": "CFGGuider",
"pos": [
766,
265
],
"size": [
315,
98
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 260
},
{
"name": "positive",
"type": "CONDITIONING",
"link": 215
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 214
}
],
"outputs": [
{
"name": "GUIDER",
"type": "GUIDER",
"links": [
87
],
"slot_index": 0,
"shape": 3
}
],
"properties": {
"Node name for S&R": "CFGGuider"
},
"widgets_values": [
4
]
},
{
"id": 37,
"type": "RandomNoise",
"pos": [
765,
-259
],
"size": [
315,
82
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "NOISE",
"type": "NOISE",
"links": [
72
],
"slot_index": 0,
"shape": 3
}
],
"properties": {
"Node name for S&R": "RandomNoise"
},
"widgets_values": [
43,
"fixed"
]
},
{
"id": 110,
"type": "VHS_LoadVideo",
"pos": [
-442,
-10
],
"size": [
283.6130065917969,
473.7097473144531
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "meta_batch",
"type": "VHS_BatchManager",
"link": null,
"shape": 7
},
{
"name": "vae",
"type": "VAE",
"link": null,
"shape": 7
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
271
],
"slot_index": 0
},
{
"name": "frame_count",
"type": "INT",
"links": null
},
{
"name": "audio",
"type": "AUDIO",
"links": null
},
{
"name": "video_info",
"type": "VHS_VIDEOINFO",
"links": null
}
],
"properties": {
"Node name for S&R": "VHS_LoadVideo"
},
"widgets_values": {
"video": "video.mp4",
"force_rate": 0,
"force_size": "Disabled",
"custom_width": 512,
"custom_height": 512,
"frame_load_cap": 0,
"skip_first_frames": 0,
"select_every_nth": 1,
"choose video to upload": "image",
"videopreview": {
"hidden": false,
"paused": false,
"params": {
"force_rate": 0,
"frame_load_cap": 0,
"skip_first_frames": 0,
"select_every_nth": 1,
"filename": "video.mp4",
"type": "input",
"format": "video/mp4"
},
"muted": false
}
}
},
{
"id": 102,
"type": "LTXVLoader",
"pos": [
-461,
-249
],
"size": [
315,
78
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "model",
"type": "MODEL",
"links": [
254
],
"slot_index": 0
},
{
"name": "vae",
"type": "VAE",
"links": [
255,
256
],
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "LTXVLoader"
},
"widgets_values": [
"ltx-video-2b-v0.9.safetensors"
]
}
],
"links": [
[
70,
14,
0,
36,
2,
"SAMPLER"
],
[
72,
37,
0,
36,
0,
"NOISE"
],
[
78,
36,
0,
8,
0,
"LATENT"
],
[
87,
42,
0,
36,
1,
"GUIDER"
],
[
121,
8,
0,
60,
0,
"IMAGE"
],
[
214,
88,
0,
42,
2,
"CONDITIONING"
],
[
215,
87,
0,
42,
1,
"CONDITIONING"
],
[
254,
102,
0,
103,
0,
"MODEL"
],
[
255,
102,
1,
103,
1,
"VAE"
],
[
256,
102,
1,
8,
1,
"VAE"
],
[
258,
103,
1,
36,
4,
"LATENT"
],
[
259,
103,
2,
104,
1,
"FLOAT"
],
[
260,
103,
0,
42,
0,
"MODEL"
],
[
261,
103,
0,
81,
0,
"MODEL"
],
[
262,
81,
0,
104,
0,
"SIGMAS"
],
[
264,
105,
0,
87,
0,
"CLIP"
],
[
265,
105,
0,
88,
0,
"CLIP"
],
[
269,
104,
0,
109,
0,
"SIGMAS"
],
[
270,
109,
1,
36,
3,
"SIGMAS"
],
[
271,
110,
0,
103,
2,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 0.6115909044841471,
"offset": [
595.203978662415,
513.6446771351451
]
}
},
"version": 0.4
}
assets/ltxvideo-v2v.png
-3
View File
@@ -1,3 +0,0 @@
version https://git-lfs.github.com/spec/v1
oid sha256:22d3a1732d18e5aaf2c3fd8263e271cdd40b6e0384eaef1082eac6cc39c0cd60
size 708563
assets/moto.png
+3
View File
@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:8681ff86edb43d0a12472175d5503a590910b24d5ec5d5d5b78eed079960eac0
size 998115
assets/video.mp4
-3
View File
@@ -1,3 +0,0 @@
version https://git-lfs.github.com/spec/v1
oid sha256:610240226d1f9f7e89e6d585176d952787d79523403d1c9a6c2feaba85e8ed35
size 1538809
img2vid.py
+66 -7
View File
@@ -1,3 +1,6 @@
import io
import av
import comfy.latent_formats
import comfy.model_base
import comfy.model_management
@@ -5,10 +8,52 @@ import comfy.model_patcher
import comfy.sd
import comfy.supported_models_base
import comfy.utils
import numpy as np
import torch
from ltx_video.models.autoencoders.vae_encode import get_vae_size_scale_factor
def encode_single_frame(output_file, image_array: np.ndarray, crf):
container = av.open(output_file, "w", format="mp4")
try:
stream = container.add_stream(
"h264", rate=1, options={"crf": str(crf), "preset": "veryfast"}
)
stream.height = image_array.shape[0]
stream.width = image_array.shape[1]
av_frame = av.VideoFrame.from_ndarray(image_array, format="rgb24").reformat(
format="yuv420p"
)
container.mux(stream.encode(av_frame))
container.mux(stream.encode())
finally:
container.close()
def decode_single_frame(video_file):
container = av.open(video_file)
try:
stream = next(s for s in container.streams if s.type == "video")
frame = next(container.decode(stream))
finally:
container.close()
return frame.to_ndarray(format="rgb24")
def videofy(image: torch.Tensor, crf=29):
if crf == 0:
return image
image_array = (image * 255.0).byte().cpu().numpy()
with io.BytesIO() as output_file:
encode_single_frame(output_file, image_array, crf)
video_bytes = output_file.getvalue()
with io.BytesIO(video_bytes) as video_file:
image_array = decode_single_frame(video_file)
tensor = torch.tensor(image_array, dtype=image.dtype, device=image.device) / 255.0
return tensor
def pad_tensor(tensor, target_len):
dim = 2
repeat_factor = target_len - tensor.shape[dim] # Ceiling division
@@ -17,24 +62,38 @@ def pad_tensor(tensor, target_len):
return torch.cat([tensor, padding], dim=dim)
def encode_media_conditioning(init_media, vae, width, height, frames_number):
def encode_media_conditioning(
init_media, vae, width, height, frames_number, image_compression, initial_latent
):
pixels = comfy.utils.common_upscale(
init_media.movedim(-1, 1), width, height, "bilinear", "center"
init_media.movedim(-1, 1), width, height, "bilinear", ""
).movedim(1, -1)
encode_pixels = pixels[:, :, :, :3]
init_latents = vae.encode(encode_pixels).float()
if image_compression > 0:
for i in range(encode_pixels.shape[0]):
image = videofy(encode_pixels[i], image_compression)
encode_pixels[i] = image
encoded_latents = vae.encode(encode_pixels).float()
video_scale_factor, _, _ = get_vae_size_scale_factor(vae.first_stage_model)
video_scale_factor = video_scale_factor if frames_number > 1 else 1
target_len = (frames_number // video_scale_factor) + 1
init_latents = init_latents[:, :, :target_len]
encoded_latents = encoded_latents[:, :, :target_len]
if initial_latent is None:
initial_latent = encoded_latents
else:
if encoded_latents.shape[2] > initial_latent.shape[2]:
initial_latent = pad_tensor(initial_latent, encoded_latents.shape[2])
initial_latent[:, :, : encoded_latents.shape[2], ...] = encoded_latents
init_image_frame_number = init_media.shape[0]
if init_image_frame_number == 1:
result = pad_tensor(init_latents, target_len)
result = pad_tensor(initial_latent, target_len)
elif init_image_frame_number % 8 != 1:
result = pad_tensor(init_latents, target_len)
result = pad_tensor(initial_latent, target_len)
else:
result = init_latents
result = initial_latent
return result
loader_node.py
+67 -44
View File
@@ -1,3 +1,4 @@
import json
from pathlib import Path
import comfy
@@ -11,6 +12,7 @@ from ltx_video.models.autoencoders.causal_video_autoencoder import (
)
from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier
from ltx_video.models.transformers.transformer3d import Transformer3DModel
from safetensors import safe_open
from .model import LTXVModel, LTXVModelConfig, LTXVTransformer3D
from .nodes_registry import comfy_node
@@ -27,7 +29,7 @@ class LTXVLoader:
folder_paths.get_filename_list("checkpoints"),
{"tooltip": "The name of the checkpoint (model) to load."},
),
"dtype": (["bfloat16", "float32"], {"default": "float32"}),
"dtype": (["bfloat16", "float32"], {"default": "bfloat16"}),
}
}
@@ -45,9 +47,20 @@ class LTXVLoader:
ckpt_path = Path(folder_paths.get_full_path("checkpoints", ckpt_name))
weights = safetensors.torch.load_file(ckpt_path, device=str(load_device))
vae_config = None
unet_config = None
with safe_open(ckpt_path, framework="pt", device="cpu") as f:
metadata = f.metadata()
if metadata is not None:
config_metadata = metadata.get("config", None)
if config_metadata is not None:
config_metadata = json.loads(config_metadata)
vae_config = config_metadata.get("vae", None)
unet_config = config_metadata.get("transformer", None)
vae = self._load_vae(weights)
weights = safetensors.torch.load_file(ckpt_path, device="cpu")
vae = self._load_vae(weights, vae_config)
num_latent_channels = vae.first_stage_model.config.latent_channels
model = self._load_unet(
@@ -56,14 +69,13 @@ class LTXVLoader:
weights,
num_latent_channels,
dtype=dtype_map[dtype],
config=unet_config,
)
return (model, vae)
def _load_vae(self, weights):
vae_prefix = "vae."
vae = LTXVVAE.from_config_and_state_dict(
vae_class=CausalVideoAutoencoder,
config={
def _load_vae(self, weights, config=None):
if config is None:
config = {
"_class_name": "CausalVideoAutoencoder",
"dims": 3,
"in_channels": 3,
@@ -87,7 +99,11 @@ class LTXVLoader:
"latent_log_var": "uniform",
"use_quant_conv": False,
"causal_decoder": False,
},
}
vae_prefix = "vae."
vae = LTXVVAE.from_config_and_state_dict(
vae_class=CausalVideoAutoencoder,
config=config,
state_dict={
key.removeprefix(vae_prefix): value
for key, value in weights.items()
@@ -97,40 +113,47 @@ class LTXVLoader:
return vae
def _load_unet(
self, load_device, offload_device, weights, num_latent_channels, dtype
self,
load_device,
offload_device,
weights,
num_latent_channels,
dtype,
config=None,
):
config = {
"_class_name": "Transformer3DModel",
"_diffusers_version": "0.25.1",
"_name_or_path": "PixArt-alpha/PixArt-XL-2-256x256",
"activation_fn": "gelu-approximate",
"attention_bias": True,
"attention_head_dim": 64,
"attention_type": "default",
"caption_channels": 4096,
"cross_attention_dim": 2048,
"double_self_attention": False,
"dropout": 0.0,
"in_channels": 128,
"norm_elementwise_affine": False,
"norm_eps": 1e-06,
"norm_num_groups": 32,
"num_attention_heads": 32,
"num_embeds_ada_norm": 1000,
"num_layers": 28,
"num_vector_embeds": None,
"only_cross_attention": False,
"out_channels": 128,
"project_to_2d_pos": True,
"upcast_attention": False,
"use_linear_projection": False,
"qk_norm": "rms_norm",
"standardization_norm": "rms_norm",
"positional_embedding_type": "rope",
"positional_embedding_theta": 10000.0,
"positional_embedding_max_pos": [20, 2048, 2048],
"timestep_scale_multiplier": 1000,
}
if config is None:
config = {
"_class_name": "Transformer3DModel",
"_diffusers_version": "0.25.1",
"_name_or_path": "PixArt-alpha/PixArt-XL-2-256x256",
"activation_fn": "gelu-approximate",
"attention_bias": True,
"attention_head_dim": 64,
"attention_type": "default",
"caption_channels": 4096,
"cross_attention_dim": 2048,
"double_self_attention": False,
"dropout": 0.0,
"in_channels": 128,
"norm_elementwise_affine": False,
"norm_eps": 1e-06,
"norm_num_groups": 32,
"num_attention_heads": 32,
"num_embeds_ada_norm": 1000,
"num_layers": 28,
"num_vector_embeds": None,
"only_cross_attention": False,
"out_channels": 128,
"project_to_2d_pos": True,
"upcast_attention": False,
"use_linear_projection": False,
"qk_norm": "rms_norm",
"standardization_norm": "rms_norm",
"positional_embedding_type": "rope",
"positional_embedding_theta": 10000.0,
"positional_embedding_max_pos": [20, 2048, 2048],
"timestep_scale_multiplier": 1000,
}
transformer = Transformer3DModel.from_config(config)
unet_prefix = "model.diffusion_model."
@@ -142,8 +165,8 @@ class LTXVLoader:
}
)
transformer.to(dtype).to(load_device).eval()
diffusion_model = LTXVTransformer3D(transformer, SymmetricPatchifier(1))
patchifier = SymmetricPatchifier(1)
diffusion_model = LTXVTransformer3D(transformer, patchifier, None, None, None)
model = LTXVModel(
LTXVModelConfig(num_latent_channels, dtype=dtype),
model_type=comfy.model_base.ModelType.FLOW,
model.py
+81 -45
View File
@@ -29,24 +29,25 @@ class LTXVModelConfig:
class LTXVSampling(torch.nn.Module, comfy.model_sampling.CONST):
def __init__(self, condition_mask):
def __init__(self, condition_mask, guiding_latent=None):
super().__init__()
self.condition_mask = condition_mask
self.guiding_latent = guiding_latent
self.set_parameters(shift=1.0, multiplier=1)
def set_parameters(self, shift=1.0, timesteps=1000, multiplier=1000):
self.shift = shift
self.multiplier = multiplier
ts = self.sigma((torch.arange(1, timesteps + 1, 1) / timesteps) * multiplier)
ts = self.sigma((torch.arange(0, timesteps + 1, 1) / timesteps) * multiplier)
self.register_buffer("sigmas", ts)
@property
def sigma_min(self):
return 0.0
return self.sigmas[0]
@property
def sigma_max(self):
return 1.0
return self.sigmas[-1]
def timestep(self, sigma):
return sigma * self.multiplier
@@ -61,6 +62,14 @@ class LTXVSampling(torch.nn.Module, comfy.model_sampling.CONST):
return 0.0
return 1.0 - percent
def calculate_input(self, sigma, noise):
if self.guiding_latent is not None:
noise = (
noise * (1 - self.condition_mask)
+ self.guiding_latent * self.condition_mask
)
return noise
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
self.condition_mask = self.condition_mask.to(latent_image.device)
scaled = latent_image * (1 - sigma) + noise * sigma
@@ -72,7 +81,15 @@ class LTXVSampling(torch.nn.Module, comfy.model_sampling.CONST):
sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
result = model_input - model_output * sigma
# In order to d * dT to be zero in euler step, we need to set result equal to input in first latent frame.
result = result * (1 - self.condition_mask) + model_input * self.condition_mask
if self.guiding_latent is not None:
result = (
result * (1 - self.condition_mask)
+ self.guiding_latent * self.condition_mask
)
else:
result = (
result * (1 - self.condition_mask) + model_input * self.condition_mask
)
return result
@@ -84,19 +101,53 @@ class LTXVModel(comfy.model_base.BaseModel):
class LTXVTransformer3D(nn.Module):
def __init__(
self, transformer: Transformer3DModel, patchifier: SymmetricPatchifier
self,
transformer: Transformer3DModel,
patchifier: SymmetricPatchifier,
conditioning_mask,
latent_frame_rate,
vae_scale_factor,
):
super().__init__()
self.transformer = transformer
self.dtype = transformer.dtype
self.patchifier = patchifier
def forward(
self.dtype = transformer.dtype
self.transformer = transformer
self.patchifier = patchifier
self.conditioning_mask = conditioning_mask
self.latent_frame_rate = latent_frame_rate
self.vae_scale_factor = vae_scale_factor
def indices_grid(
self,
latent_shape,
device,
):
use_rope = self.transformer.use_rope
scale_grid = (
(1 / self.latent_frame_rate, self.vae_scale_factor, self.vae_scale_factor)
if use_rope
else None
)
indices_grid = self.patchifier.get_grid(
orig_num_frames=latent_shape[2],
orig_height=latent_shape[3],
orig_width=latent_shape[4],
batch_size=latent_shape[0],
scale_grid=scale_grid,
device=device,
)
return indices_grid
def wrapped_transformer(
self,
latent,
timesteps,
context,
indices_grid,
skip_layer_mask=None,
skip_layer_strategy=None,
img_hw=None,
aspect_ratio=None,
mixed_precision=True,
@@ -122,6 +173,8 @@ class LTXVTransformer3D(nn.Module):
torch.int64
),
timestep=timesteps,
skip_layer_mask=skip_layer_mask,
skip_layer_strategy=skip_layer_strategy,
return_dict=False,
)[0]
@@ -134,50 +187,33 @@ class LTXVTransformer3D(nn.Module):
)
return result
class LTXVTransformer3DWrapper(nn.Module):
def __init__(
self,
transformer: LTXVTransformer3D,
patchifier: SymmetricPatchifier,
conditioning_mask,
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
self.patchifier = patchifier
self.conditioning_mask = conditioning_mask
self.conditioning_mask_patchified = patchifier.patchify(
self.conditioning_mask
).squeeze(-1)
def forward(self, x, timesteps, context, img_hw=None, aspect_ratio=None, **kwargs):
transformer_options = kwargs.get("transformer_options", {})
ptb_index = transformer_options.get("ptb_index", None)
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)
cond_or_uncond = transformer_options.get("cond_or_uncond", [])
skip_block_list = transformer_options.get("skip_block_list", [])
skip_layer_strategy = transformer_options.get("skip_layer_strategy", None)
mask = self.patchifier.patchify(self.conditioning_mask).squeeze(-1).to(x.device)
ndim_mask = mask.ndimension()
expanded_timesteps = timesteps.view(timesteps.size(0), *([1] * (ndim_mask - 1)))
timesteps_masked = expanded_timesteps * (1 - mask)
skip_layer_mask = None
if ptb_index is not None and ptb_index in cond_or_uncond:
skip_layer_mask = self.transformer.create_skip_layer_mask(
skip_block_list,
1,
len(cond_or_uncond),
len(cond_or_uncond) - 1 - cond_or_uncond.index(ptb_index),
)
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(
latent,
result = self.wrapped_transformer(
x,
timesteps_masked,
context,
self.indices_grid,
indices_grid=self.indices_grid(x.shape, x.device),
mixed_precision=mixed_precision,
skip_layer_mask=skip_layer_mask,
skip_layer_strategy=skip_layer_strategy,
)
return result
pyproject.toml
+2 -1
View File
@@ -9,7 +9,8 @@ requires-python = ">=3.10"
readme = "README.md"
license = { file = "LICENSE" }
dependencies = [
"ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9",
"ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.1",
"av>=10.0.0"
]
[project.optional-dependencies]
requirements.txt
+2 -1
View File
@@ -1 +1,2 @@
ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9
av>=10.0.0
ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.1
stg.py
+196
View File
@@ -0,0 +1,196 @@
import comfy.samplers
import comfy.utils
import torch
from comfy.model_patcher import ModelPatcher
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
from .nodes_registry import comfy_node
def stg(
noise_pred_pos,
noise_pred_neg,
noise_pred_pertubed,
cfg_scale,
stg_scale,
rescale_scale,
):
noise_pred = (
noise_pred_neg
+ cfg_scale * (noise_pred_pos - noise_pred_neg)
+ stg_scale * (noise_pred_pos - noise_pred_pertubed)
)
if rescale_scale != 0:
factor = noise_pred_pos.std() / noise_pred.std()
factor = rescale_scale * factor + (1 - rescale_scale)
noise_pred = noise_pred * factor
return noise_pred
@comfy_node(name="LTXVApplySTG")
class LTXVApplySTG:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": (
"MODEL",
{"tooltip": "The model to apply the STG to."},
),
"stg_mode": (["attention", "residual"],),
"block_indices": (
"STRING",
{
"default": "14, 19",
"tooltip": "Comma-separated indices of the blocks to apply the STG to.",
},
),
}
}
FUNCTION = "apply_stg"
RETURN_TYPES = ("MODEL",)
RETURN_NAMES = ("model",)
CATEGORY = "lightricks/LTXV"
def apply_stg(self, model: ModelPatcher, stg_mode: str, block_indices: str):
skip_block_list = [int(i.strip()) for i in block_indices.split(",")]
stg_mode = (
SkipLayerStrategy.Attention
if stg_mode == "attention"
else SkipLayerStrategy.Residual
)
new_model = model.clone()
new_model.model_options["transformer_options"]["skip_layer_strategy"] = stg_mode
if "skip_block_list" in new_model.model_options["transformer_options"]:
skip_block_list.extend(
new_model.model_options["transformer_options"]["skip_block_list"]
)
new_model.model_options["transformer_options"][
"skip_block_list"
] = skip_block_list
return (new_model,)
class STGGuider(comfy.samplers.CFGGuider):
def set_conds(self, positive, negative):
self.inner_set_conds(
{"positive": positive, "negative": negative, "perturbed": positive}
)
def set_cfg(self, cfg, stg_scale, rescale_scale: float = None):
self.cfg = cfg
self.stg_scale = stg_scale
self.rescale_scale = rescale_scale
def predict_noise(
self,
x: torch.Tensor,
timestep: torch.Tensor,
model_options: dict = {},
seed=None,
):
# in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
# but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly
perturbed_cond = self.conds.get("perturbed", None)
positive_cond = self.conds.get("positive", None)
negative_cond = self.conds.get("negative", None)
noise_pred_neg = 0
# no similar optimization for stg=0, use CFG guider instead.
if self.cfg > 1:
model_options["transformer_options"]["ptb_index"] = 2
(noise_pred_perturbed, noise_pred_pos, noise_pred_neg) = (
comfy.samplers.calc_cond_batch(
self.inner_model,
[perturbed_cond, positive_cond, negative_cond],
x,
timestep,
model_options,
)
)
else:
model_options["transformer_options"]["ptb_index"] = 1
(noise_pred_perturbed, noise_pred_pos) = comfy.samplers.calc_cond_batch(
self.inner_model,
[perturbed_cond, positive_cond],
x,
timestep,
model_options,
)
stg_result = stg(
noise_pred_pos,
noise_pred_neg,
noise_pred_perturbed,
self.cfg,
self.stg_scale,
self.rescale_scale,
)
# normally this would be done in cfg_function, but we skipped
# that for efficiency: we can compute the noise predictions in
# a single call to calc_cond_batch() (rather than two)
# so we replicate the hook here
for fn in model_options.get("sampler_post_cfg_function", []):
args = {
"denoised": stg_result,
"cond": positive_cond,
"uncond": negative_cond,
"model": self.inner_model,
"uncond_denoised": noise_pred_neg,
"cond_denoised": noise_pred_pos,
"sigma": timestep,
"model_options": model_options,
"input": x,
# not in the original call in samplers.py:cfg_function, but made available for future hooks
"perturbed_cond": positive_cond,
"perturbed_cond_denoised": noise_pred_perturbed,
}
stg_result = fn(args)
return stg_result
@comfy_node(name="STGGuider")
class STGGuiderNode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"cfg": (
"FLOAT",
{
"default": 1.0,
"min": 0.0,
"max": 100.0,
"step": 0.1,
"round": 0.01,
},
),
"stg": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01},
),
"rescale": (
"FLOAT",
{"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01},
),
}
}
RETURN_TYPES = ("GUIDER",)
FUNCTION = "get_guider"
CATEGORY = "lightricks/LTXV"
def get_guider(self, model, positive, negative, cfg, stg, rescale):
guider = STGGuider(model)
guider.set_conds(positive, negative)
guider.set_cfg(cfg, stg, rescale)
return (guider,)
transformer.py
+36 -59
View File
@@ -9,10 +9,9 @@ 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 .img2vid import encode_media_conditioning
from .model import LTXVSampling, LTXVTransformer3DWrapper
from .model import LTXVSampling
from .nodes_registry import comfy_node
@@ -88,14 +87,13 @@ class LTXVModelConfigurator:
"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",
"img_compression": (
"INT",
{
"default": 0.15,
"min": 0.0,
"default": 29,
"min": 0,
"max": 100,
"step": 0.01,
"tooltip": "Amount of noise to apply on conditioning image latent.",
"tooltip": "Amount of compression to apply on conditioning image.",
},
),
},
@@ -103,7 +101,13 @@ class LTXVModelConfigurator:
"conditioning": (
"IMAGE",
{"tooltip": "Optional conditioning image or video."},
)
),
"initial_latent": (
"LATENT",
{
"tooltip": "initial latent that is combined with conditioning if given"
},
),
},
}
@@ -137,32 +141,6 @@ class LTXVModelConfigurator:
]
return latent_shape, latent_frame_rate
def indices_grid(
self,
latent_shape,
vae_scale_factor,
latent_frame_rate,
patchifier,
use_rope,
device,
):
scale_grid = (
(1 / latent_frame_rate, vae_scale_factor, vae_scale_factor)
if use_rope # self.wrapped_transformer.transformer.use_rope
else None
)
indices_grid = patchifier.get_grid(
orig_num_frames=latent_shape[2],
orig_height=latent_shape[3],
orig_width=latent_shape[4],
batch_size=latent_shape[0],
scale_grid=scale_grid,
device=device,
)
return indices_grid
def configure_sizes(
self,
model,
@@ -174,8 +152,9 @@ class LTXVModelConfigurator:
frame_rate,
batch,
mixed_precision,
noise_scale,
img_compression,
conditioning=None,
initial_latent=None,
):
load_device = comfy.model_management.get_torch_device()
if preset != "Custom":
@@ -193,43 +172,41 @@ class LTXVModelConfigurator:
latent_shape[4],
]
conditioning_mask = torch.zeros(mask_shape, device=load_device)
initial_latent = (
None
if initial_latent is None
else initial_latent["samples"].to(load_device)
)
guiding_latent = None
if conditioning is not None:
latent = encode_media_conditioning(
conditioning, vae, width, height, frames_number
conditioning,
vae,
width,
height,
frames_number,
image_compression=img_compression,
initial_latent=initial_latent,
)
conditioning_mask[:, :, 0] = 1.0
guiding_latent = latent[:, :, :1, ...]
else:
latent = torch.zeros(latent_shape, dtype=torch.float32)
latent = torch.zeros(latent_shape, dtype=torch.float32, device=load_device)
if initial_latent is not None:
latent[:, :, : initial_latent.shape[2], ...] = initial_latent
patchifier = SymmetricPatchifier(1)
_, vae_scale_factor, _ = get_vae_size_scale_factor(vae.first_stage_model)
use_rope = model.get_model_object("diffusion_model").transformer.use_rope
indices_grid = self.indices_grid(
latent_shape,
vae_scale_factor,
latent_frame_rate,
patchifier,
use_rope,
load_device,
)
wrapper = LTXVTransformer3DWrapper(
transformer=model.get_model_object("diffusion_model"),
patchifier=patchifier,
conditioning_mask=conditioning_mask,
indices_grid=indices_grid,
)
patcher = model.clone()
patcher.add_object_patch("diffusion_model", wrapper)
patcher.add_object_patch("diffusion_model.conditioning_mask", conditioning_mask)
patcher.add_object_patch("diffusion_model.latent_frame_rate", latent_frame_rate)
patcher.add_object_patch("diffusion_model.vae_scale_factor", vae_scale_factor)
patcher.add_object_patch(
"model_sampling", LTXVSampling(wrapper.conditioning_mask)
"model_sampling", LTXVSampling(conditioning_mask, guiding_latent)
)
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))
vae.py
+92 -4
View File
@@ -1,3 +1,5 @@
from copy import copy
import comfy.latent_formats
import comfy.model_base
import comfy.model_management
@@ -13,11 +15,16 @@ from ltx_video.models.autoencoders.vae_encode import (
vae_encode,
)
from .nodes_registry import comfy_node
class LTXVVAE(comfy.sd.VAE):
def __init__(self):
def __init__(self, decode_timestep=0.05, decode_noise_scale=0.025, seed=42):
self.device = comfy.model_management.vae_device()
self.offload_device = comfy.model_management.vae_offload_device()
self.decode_timestep = decode_timestep
self.decode_noise_scale = decode_noise_scale
self.seed = seed
@classmethod
def from_pretrained(cls, vae_class, model_path, dtype=torch.bfloat16):
@@ -53,18 +60,48 @@ class LTXVVAE(comfy.sd.VAE):
# (batch, channels, frames, height, width)
def decode(self, samples_in):
is_video = samples_in.shape[2] > 1
decode_timestep = self.decode_timestep
if getattr(self.first_stage_model.decoder, "timestep_conditioning", False):
samples_in = self.add_noise(
decode_timestep, self.decode_noise_scale, self.seed, samples_in
)
else:
decode_timestep = None
result = vae_decode(
samples_in.to(self.device),
vae=self.first_stage_model,
is_video=is_video,
vae_per_channel_normalize=True,
timestep=decode_timestep,
)
result = self.image_processor.postprocess(
result, output_type="pt", do_denormalize=[True]
)
return (
result.squeeze(0).permute(1, 2, 3, 0).to(torch.float32)
) # .to(self.offload_device)
return result.squeeze(0).permute(1, 2, 3, 0).to(torch.float32)
@staticmethod
def add_noise(decode_timestep, decode_noise_scale, seed, latents):
generator = torch.Generator(device="cpu").manual_seed(seed)
noise = torch.randn(
latents.size(),
generator=generator,
device=latents.device,
dtype=latents.dtype,
)
if not isinstance(decode_timestep, list):
decode_timestep = [decode_timestep] * latents.shape[0]
if decode_noise_scale is None:
decode_noise_scale = decode_timestep
elif not isinstance(decode_noise_scale, list):
decode_noise_scale = [decode_noise_scale] * latents.shape[0]
decode_timestep = torch.tensor(decode_timestep).to(latents.device)
decode_noise_scale = torch.tensor(decode_noise_scale).to(latents.device)[
:, None, None, None, None
]
latents = latents * (1 - decode_noise_scale) + noise * decode_noise_scale
return latents
# Underlying VAE expects b, c, n, h, w dimensions order and dtype specific dtype.
# However in Comfy the convension is n, h, w, c.
@@ -77,3 +114,54 @@ class LTXVVAE(comfy.sd.VAE):
input, self.first_stage_model, vae_per_channel_normalize=True
).to(comfy.model_management.get_torch_device())
return latents
@comfy_node(name="Add VAE Decoder Noise")
class DecoderNoise:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"vae": ("VAE",),
"timestep": (
"FLOAT",
{
"default": 0.05,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"tooltip": "The timestep used for decoding the noise.",
},
),
"scale": (
"FLOAT",
{
"default": 0.025,
"min": 0.0,
"max": 1.0,
"step": 0.001,
"tooltip": "The scale of the noise added to the decoder.",
},
),
"seed": (
"INT",
{
"default": 42,
"min": 0,
"max": 0xFFFFFFFFFFFFFFFF,
"tooltip": "The random seed used for creating the noise.",
},
),
}
}
FUNCTION = "add_noise"
RETURN_TYPES = ("VAE",)
CATEGORY = "lightricks/LTXV"
def add_noise(self, vae, timestep, scale, seed):
result = copy(vae)
result.decode_timestep = timestep
result.decode_noise_scale = scale
result.seed = seed
return (result,)