Automated PR - 2025-03-05
This commit is contained in:
sync-bot
@@ -1,16 +1,43 @@
|
||||
# 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).
|
||||
|
||||
ComfyUI-LTXVideo is a collection of custom nodes for ComfyUI, designed to provide useful tools for working with the LTXV model.
|
||||
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).
|
||||
|
||||
# 5.03.2025 :star: LTXVideo 0.9.5 Release :star:
|
||||
|
||||
### LTXVideo 0.9.5 introduces:
|
||||
|
||||
1. Improved quality with reduced artifacts.
|
||||
2. Support for higher resolution and longer sequences.
|
||||
3. Frame and sequence conditioning (beyond the first frame).
|
||||
4. Enhanced prompt understanding.
|
||||
5. Commercial license availability.
|
||||
|
||||
### Technical Updates
|
||||
|
||||
Since LTXVideo is now fully supported in the ComfyUI core, we have removed the custom model implementation. Instead, we provide updated workflows to showcase the new features:
|
||||
|
||||
1. **Frame Conditioning** – Enables interpolation between given frames.
|
||||
2. **Sequence Conditioning** – Allows motion interpolation from a given frame sequence, enabling video extension from the beginning, end, or middle of the original video.
|
||||
3. **Prompt Enhancer** – A new node that helps generate prompts optimized for the best model performance.
|
||||
See the [Example Workflows](#example-workflows) section for more details.
|
||||
|
||||
### LTXTricks Update
|
||||
|
||||
The LTXTricks code has been integrated into this repository (in the `/tricks` folder) and will be maintained here. The original [repo](https://github.com/logtd/ComfyUI-LTXTricks) is no longer maintained, but all existing workflows should continue to function as expected.
|
||||
|
||||
## 22.12.2024
|
||||
|
||||
Fixed a bug which caused the model to produce artifacts on short negative prompts when using a native CLIP Loader node.
|
||||
|
||||
## 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).
|
||||
4. Additional initial latent optional input to chain latents for high res generation.
|
||||
5. Image captioning in image to video [flow](assets/ltxvideo-i2v.json).
|
||||
|
||||
## Installation
|
||||
|
||||
@@ -21,10 +48,13 @@ Installation via [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager)
|
||||
1. Install ComfyUI
|
||||
2. Clone this repository to `custom-nodes` folder in your ComfyUI installation directory.
|
||||
3. Install the required packages:
|
||||
|
||||
```bash
|
||||
cd custom_nodes/ComfyUI-LTXVideo && pip install -r requirements.txt
|
||||
```
|
||||
|
||||
For portable ComfyUI installations, run
|
||||
|
||||
```
|
||||
.\python_embeded\python.exe -m pip install -r .\ComfyUI\custom_nodes\ComfyUI-LTXVideo\requirements.txt
|
||||
```
|
||||
@@ -47,3 +77,28 @@ Note that to run the example workflows, you need to have some additional custom
|
||||
|
||||
[Download workflow](assets/ltxvideo-t2v.json)
|
||||
|
||||
|
||||
### Frame Interpolation
|
||||
|
||||
[Download workflow](assets/ltxvideo-frame-interpolation.json)
|
||||
|
||||
|
||||
### First Sequence Conditioning
|
||||
|
||||
[Download workflow](assets/ltxvideo-first-sequence-conditioning.json)
|
||||
|
||||
|
||||
### Last Sequence Conditioning
|
||||
|
||||
[Download workflow](assets/ltxvideo-last-sequence-conditioning.json)
|
||||
|
||||
|
||||
### Flow Edit
|
||||
|
||||
[Download workflow](tricks/assets/ltxvideo-flow-edit.json)
|
||||
|
||||
|
||||
### RF Edit
|
||||
|
||||
[Download workflow](tricks/assets/ltxvideo-rf-edit.json)
|
||||
|
||||
|
||||
+10
-3
@@ -1,8 +1,15 @@
|
||||
from .loader_node import LTXVLoader # noqa: F401
|
||||
from .nodes_registry import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
from .prompt_enhancer_nodes import ( # noqa: F401
|
||||
LTXVPromptEnhancer,
|
||||
LTXVPromptEnhancerLoader,
|
||||
)
|
||||
from .stg import LTXVApplySTG, STGGuiderNode # noqa: F401
|
||||
from .t5_encoder import LTXVCLIPModelLoader # noqa: F401
|
||||
from .transformer import LTXVModelConfigurator, LTXVShiftSigmas # noqa: F401
|
||||
from .tricks import NODE_CLASS_MAPPINGS as TRICKS_NODE_CLASS_MAPPINGS
|
||||
from .tricks import NODE_DISPLAY_NAME_MAPPINGS as TRICKS_NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
# Merge the node mappings from tricks into the main mappings
|
||||
NODE_CLASS_MAPPINGS.update(TRICKS_NODE_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(TRICKS_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
|
||||
# Export so that ComfyUI can pick them up.
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 474 KiB |
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:5b8ac2148bdf1092f4342eb4f08961458a9e22743e8896f5dbcf749166f4f526
|
||||
size 1582922
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:35d14c91dc998cb39742998d4ebe4e79b82cb938d5145eb7f6719b7eb98bc961
|
||||
size 1007142
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:a66b6cee6f447cc24735bba839ce9bd99034220bcd15a1259795ba9ef1671508
|
||||
size 1243622
|
||||
+1086
-1227
File diff suppressed because it is too large
Load Diff
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:6ec16f1372ece0cd63395edf29d9997599580df3703d7fd30025fd0870ad5532
|
||||
size 1205996
|
||||
oid sha256:ffeaa0e0a24eb0de00f75cd96fb8ea46dd846c24fd221f723ab740169383f7af
|
||||
size 1253177
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:e49012d3d26c4f5e0e4320d9bd201c14adab322ebd9cd13aabce3580271c38e8
|
||||
size 1008949
|
||||
+808
-995
File diff suppressed because it is too large
Load Diff
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:1e508763410d009f524ed22c1c420e798904afe92098130a4d155d8318ef6b3a
|
||||
size 1112804
|
||||
oid sha256:eff8622aa7c3341934a69df8af5b44732c6e35284abb96609e1e5eb0bfd3dc82
|
||||
size 751672
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 349 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 570 KiB |
-99
@@ -1,99 +0,0 @@
|
||||
import io
|
||||
|
||||
import av
|
||||
import comfy.latent_formats
|
||||
import comfy.model_base
|
||||
import comfy.model_management
|
||||
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
|
||||
last_element = tensor.select(dim, -1).unsqueeze(dim)
|
||||
padding = last_element.repeat(1, 1, repeat_factor, 1, 1)
|
||||
return torch.cat([tensor, padding], dim=dim)
|
||||
|
||||
|
||||
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", ""
|
||||
).movedim(1, -1)
|
||||
encode_pixels = pixels[:, :, :, :3]
|
||||
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
|
||||
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(initial_latent, target_len)
|
||||
elif init_image_frame_number % 8 != 1:
|
||||
result = pad_tensor(initial_latent, target_len)
|
||||
else:
|
||||
result = initial_latent
|
||||
|
||||
return result
|
||||
-179
@@ -1,179 +0,0 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import comfy
|
||||
import comfy.model_management
|
||||
import comfy.model_patcher
|
||||
import folder_paths
|
||||
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 safetensors import safe_open
|
||||
|
||||
from .model import LTXVModel, LTXVModelConfig, LTXVTransformer3D
|
||||
from .nodes_registry import comfy_node
|
||||
from .vae import LTXVVAE
|
||||
|
||||
|
||||
@comfy_node(name="LTXVLoader")
|
||||
class LTXVLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (
|
||||
folder_paths.get_filename_list("checkpoints"),
|
||||
{"tooltip": "The name of the checkpoint (model) to load."},
|
||||
),
|
||||
"dtype": (["bfloat16", "float32"], {"default": "bfloat16"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "VAE")
|
||||
RETURN_NAMES = ("model", "vae")
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Loader"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def load(self, ckpt_name, dtype):
|
||||
dtype_map = {"bfloat16": torch.bfloat16, "float32": torch.float32}
|
||||
load_device = comfy.model_management.get_torch_device()
|
||||
offload_device = comfy.model_management.unet_offload_device()
|
||||
|
||||
ckpt_path = Path(folder_paths.get_full_path("checkpoints", ckpt_name))
|
||||
|
||||
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)
|
||||
|
||||
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(
|
||||
load_device,
|
||||
offload_device,
|
||||
weights,
|
||||
num_latent_channels,
|
||||
dtype=dtype_map[dtype],
|
||||
config=unet_config,
|
||||
)
|
||||
return (model, vae)
|
||||
|
||||
def _load_vae(self, weights, config=None):
|
||||
if not config:
|
||||
config = {
|
||||
"_class_name": "CausalVideoAutoencoder",
|
||||
"dims": 3,
|
||||
"in_channels": 3,
|
||||
"out_channels": 3,
|
||||
"latent_channels": 128,
|
||||
"blocks": [
|
||||
["res_x", 4],
|
||||
["compress_all", 1],
|
||||
["res_x_y", 1],
|
||||
["res_x", 3],
|
||||
["compress_all", 1],
|
||||
["res_x_y", 1],
|
||||
["res_x", 3],
|
||||
["compress_all", 1],
|
||||
["res_x", 3],
|
||||
["res_x", 4],
|
||||
],
|
||||
"scaling_factor": 1.0,
|
||||
"norm_layer": "pixel_norm",
|
||||
"patch_size": 4,
|
||||
"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()
|
||||
if key.startswith(vae_prefix)
|
||||
},
|
||||
)
|
||||
return vae
|
||||
|
||||
def _load_unet(
|
||||
self,
|
||||
load_device,
|
||||
offload_device,
|
||||
weights,
|
||||
num_latent_channels,
|
||||
dtype,
|
||||
config=None,
|
||||
):
|
||||
if not config:
|
||||
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."
|
||||
transformer.load_state_dict(
|
||||
{
|
||||
key.removeprefix(unet_prefix): value
|
||||
for key, value in weights.items()
|
||||
if key.startswith(unet_prefix)
|
||||
}
|
||||
)
|
||||
transformer.to(dtype).to(load_device).eval()
|
||||
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,
|
||||
device=comfy.model_management.get_torch_device(),
|
||||
)
|
||||
model.diffusion_model = diffusion_model
|
||||
|
||||
patcher = comfy.model_patcher.ModelPatcher(model, load_device, offload_device)
|
||||
|
||||
return patcher
|
||||
@@ -1,234 +0,0 @@
|
||||
import math
|
||||
from contextlib import nullcontext
|
||||
|
||||
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 torch
|
||||
import torch.nn as nn
|
||||
from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier
|
||||
from ltx_video.models.transformers.transformer3d import Transformer3DModel
|
||||
|
||||
|
||||
class LTXVModelConfig:
|
||||
def __init__(self, latent_channels, dtype):
|
||||
self.unet_config = {}
|
||||
self.unet_extra_config = {}
|
||||
self.latent_format = comfy.latent_formats.LatentFormat()
|
||||
self.latent_format.latent_channels = latent_channels
|
||||
self.manual_cast_dtype = dtype
|
||||
self.sampling_settings = {"multiplier": 1.0}
|
||||
self.memory_usage_factor = 2.7
|
||||
# denoiser is handled by extension
|
||||
self.unet_config["disable_unet_model_creation"] = True
|
||||
|
||||
|
||||
class LTXVSampling(torch.nn.Module, comfy.model_sampling.CONST):
|
||||
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(0, timesteps + 1, 1) / timesteps) * multiplier)
|
||||
self.register_buffer("sigmas", ts)
|
||||
|
||||
@property
|
||||
def sigma_min(self):
|
||||
return self.sigmas[0]
|
||||
|
||||
@property
|
||||
def sigma_max(self):
|
||||
return self.sigmas[-1]
|
||||
|
||||
def timestep(self, sigma):
|
||||
return sigma * self.multiplier
|
||||
|
||||
def sigma(self, timestep):
|
||||
return timestep
|
||||
|
||||
def percent_to_sigma(self, percent):
|
||||
if percent <= 0.0:
|
||||
return 1.0
|
||||
if percent >= 1.0:
|
||||
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
|
||||
result = latent_image * self.condition_mask + scaled * (1 - self.condition_mask)
|
||||
|
||||
return result
|
||||
|
||||
def calculate_denoised(self, sigma, model_output, model_input):
|
||||
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.
|
||||
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
|
||||
|
||||
|
||||
class LTXVModel(comfy.model_base.BaseModel):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.model_sampling = LTXVSampling(torch.zeros([1]))
|
||||
|
||||
def extra_conds(self, **kwargs):
|
||||
out = super().extra_conds(**kwargs)
|
||||
attention_mask = kwargs.get("attention_mask", None)
|
||||
if attention_mask is not None:
|
||||
out["attention_mask"] = comfy.conds.CONDRegular(attention_mask)
|
||||
return out
|
||||
|
||||
|
||||
class LTXVTransformer3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
transformer: Transformer3DModel,
|
||||
patchifier: SymmetricPatchifier,
|
||||
conditioning_mask,
|
||||
latent_frame_rate,
|
||||
vae_scale_factor,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
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,
|
||||
attention_mask,
|
||||
indices_grid,
|
||||
skip_layer_mask=None,
|
||||
skip_layer_strategy=None,
|
||||
img_hw=None,
|
||||
aspect_ratio=None,
|
||||
mixed_precision=True,
|
||||
**kwargs,
|
||||
):
|
||||
# infer mask from context padding, assumes padding vectors are all zero.
|
||||
latent = latent.to(self.transformer.dtype)
|
||||
latent_patchified = self.patchifier.patchify(latent)
|
||||
|
||||
if mixed_precision:
|
||||
context_manager = torch.autocast("cuda", dtype=torch.bfloat16)
|
||||
else:
|
||||
context_manager = nullcontext()
|
||||
with context_manager:
|
||||
noise_pred = self.transformer(
|
||||
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=attention_mask,
|
||||
timestep=timesteps,
|
||||
skip_layer_mask=skip_layer_mask,
|
||||
skip_layer_strategy=skip_layer_strategy,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
result = self.patchifier.unpatchify(
|
||||
latents=noise_pred,
|
||||
output_height=latent.shape[3],
|
||||
output_width=latent.shape[4],
|
||||
output_num_frames=latent.shape[2],
|
||||
out_channels=latent.shape[1] // math.prod(self.patchifier.patch_size),
|
||||
)
|
||||
return result
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
timesteps,
|
||||
context,
|
||||
attention_mask,
|
||||
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)
|
||||
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),
|
||||
)
|
||||
|
||||
result = self.wrapped_transformer(
|
||||
x,
|
||||
timesteps_masked,
|
||||
context,
|
||||
attention_mask,
|
||||
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
|
||||
@@ -0,0 +1,203 @@
|
||||
import os
|
||||
import shutil
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.model_patcher
|
||||
import folder_paths
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
from .prompt_enhancer_utils import generate_cinematic_prompt
|
||||
|
||||
LLM_NAME = ["unsloth/Llama-3.2-3B-Instruct"]
|
||||
|
||||
IMAGE_CAPTIONER = ["MiaoshouAI/Florence-2-large-PromptGen-v2.0"]
|
||||
|
||||
MODELS_PATH_KEY = "LLM"
|
||||
|
||||
|
||||
class PromptEnhancer(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
image_caption_processor: AutoProcessor,
|
||||
image_caption_model: AutoModelForCausalLM,
|
||||
llm_model: AutoModelForCausalLM,
|
||||
llm_tokenizer: AutoTokenizer,
|
||||
):
|
||||
super().__init__()
|
||||
self.image_caption_processor = image_caption_processor
|
||||
self.image_caption_model = image_caption_model
|
||||
self.llm_model = llm_model
|
||||
self.llm_tokenizer = llm_tokenizer
|
||||
self.device = image_caption_model.device
|
||||
# model parameters and buffer sizes plus some extra 1GB.
|
||||
self.model_size = (
|
||||
self.get_model_size(self.image_caption_model)
|
||||
+ self.get_model_size(self.llm_model)
|
||||
+ 1073741824
|
||||
)
|
||||
|
||||
def forward(self, prompt, image_conditioning, max_resulting_tokens):
|
||||
enhanced_prompt = generate_cinematic_prompt(
|
||||
self.image_caption_model,
|
||||
self.image_caption_processor,
|
||||
self.llm_model,
|
||||
self.llm_tokenizer,
|
||||
prompt,
|
||||
image_conditioning,
|
||||
max_new_tokens=max_resulting_tokens,
|
||||
)
|
||||
|
||||
return enhanced_prompt
|
||||
|
||||
@staticmethod
|
||||
def get_model_size(model):
|
||||
total_size = sum(p.numel() * p.element_size() for p in model.parameters())
|
||||
total_size += sum(b.numel() * b.element_size() for b in model.buffers())
|
||||
return total_size
|
||||
|
||||
def memory_required(self, input_shape):
|
||||
return self.model_size
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPromptEnhancerLoader")
|
||||
class LTXVPromptEnhancerLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"llm_name": (
|
||||
"STRING",
|
||||
{
|
||||
"default": LLM_NAME,
|
||||
"tooltip": "The hugging face name of the llm model to load.",
|
||||
},
|
||||
),
|
||||
"image_captioner_name": (
|
||||
"STRING",
|
||||
{
|
||||
"default": IMAGE_CAPTIONER,
|
||||
"tooltip": "The hugging face name of the image captioning model to load.",
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LTXV_PROMPT_ENHANCER",)
|
||||
RETURN_NAMES = ("prompt_enhancer",)
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Prompt Enhancer (Down)Loader"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def model_path_download_if_needed(self, model_name):
|
||||
model_directory = os.path.join(folder_paths.models_dir, MODELS_PATH_KEY)
|
||||
os.makedirs(model_directory, exist_ok=True)
|
||||
|
||||
model_name_ = model_name.rsplit("/", 1)[-1]
|
||||
model_path = os.path.join(model_directory, model_name_)
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
try:
|
||||
snapshot_download(
|
||||
repo_id=model_name,
|
||||
local_dir=model_path,
|
||||
local_dir_use_symlinks=False,
|
||||
)
|
||||
except Exception:
|
||||
shutil.rmtree(model_path, ignore_errors=True)
|
||||
raise
|
||||
return model_path
|
||||
|
||||
def down_load_llm_model(self, llm_name, load_device):
|
||||
model_path = self.model_path_download_if_needed(llm_name)
|
||||
llm_model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
llm_tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_path,
|
||||
)
|
||||
|
||||
return llm_model, llm_tokenizer
|
||||
|
||||
def down_load_image_captioner(self, image_captioner, load_device):
|
||||
model_path = self.model_path_download_if_needed(image_captioner)
|
||||
image_caption_model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path, trust_remote_code=True
|
||||
)
|
||||
|
||||
image_caption_processor = AutoProcessor.from_pretrained(
|
||||
model_path, trust_remote_code=True
|
||||
)
|
||||
|
||||
return image_caption_model, image_caption_processor
|
||||
|
||||
def load(self, llm_name, image_captioner_name):
|
||||
load_device = comfy.model_management.get_torch_device()
|
||||
offload_device = comfy.model_management.vae_offload_device()
|
||||
llm_model, llm_tokenizer = self.down_load_llm_model(llm_name, load_device)
|
||||
image_caption_model, image_caption_processor = self.down_load_image_captioner(
|
||||
image_captioner_name, load_device
|
||||
)
|
||||
|
||||
enhancer = PromptEnhancer(
|
||||
image_caption_processor, image_caption_model, llm_model, llm_tokenizer
|
||||
)
|
||||
patcher = comfy.model_patcher.ModelPatcher(
|
||||
enhancer,
|
||||
load_device,
|
||||
offload_device,
|
||||
)
|
||||
return (patcher,)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPromptEnhancer")
|
||||
class LTXVPromptEnhancer:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING",),
|
||||
"prompt_enhancer": ("LTXV_PROMPT_ENHANCER",),
|
||||
"max_resulting_tokens": (
|
||||
"INT",
|
||||
{"default": 256, "min": 32, "max": 512},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"image_prompt": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("str",)
|
||||
FUNCTION = "enhance"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Prompt Enhancer"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def enhance(
|
||||
self,
|
||||
prompt,
|
||||
prompt_enhancer: comfy.model_patcher.ModelPatcher,
|
||||
image_prompt: torch.Tensor = None,
|
||||
max_resulting_tokens=256,
|
||||
):
|
||||
comfy.model_management.free_memory(
|
||||
prompt_enhancer.memory_required([]),
|
||||
comfy.model_management.get_torch_device(),
|
||||
)
|
||||
comfy.model_management.load_model_gpu(prompt_enhancer)
|
||||
model = prompt_enhancer.model
|
||||
image_conditioning = None
|
||||
if image_prompt is not None:
|
||||
permuted_image = image_prompt.permute(3, 0, 1, 2)[None, :]
|
||||
image_conditioning = [(permuted_image, 0, 1.0)]
|
||||
|
||||
enhanced_prompt = model(prompt, image_conditioning, max_resulting_tokens)
|
||||
return (enhanced_prompt[0],)
|
||||
@@ -0,0 +1,258 @@
|
||||
import logging
|
||||
import random
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
T2V_CINEMATIC_PROMPT = """You are an expert cinematic director with many award winning movies, When writing prompts based on the user input, focus on detailed, chronological descriptions of actions and scenes.
|
||||
Include specific movements, appearances, camera angles, and environmental details - all in a single flowing paragraph.
|
||||
Start directly with the action, and keep descriptions literal and precise.
|
||||
Think like a cinematographer describing a shot list.
|
||||
Do not change the user input intent, just enhance it.
|
||||
Keep within 150 words.
|
||||
For best results, build your prompts using this structure:
|
||||
Start with main action in a single sentence
|
||||
Add specific details about movements and gestures
|
||||
Describe character/object appearances precisely
|
||||
Include background and environment details
|
||||
Specify camera angles and movements
|
||||
Describe lighting and colors
|
||||
Note any changes or sudden events
|
||||
Do not exceed the 150 word limit!
|
||||
Output the enhanced prompt only.
|
||||
|
||||
Examples:
|
||||
user prompt: A man drives a toyota car.
|
||||
enhanced prompt: A person is driving a car on a two-lane road, holding the steering wheel with both hands. The person's hands are light-skinned and they are wearing a black long-sleeved shirt. The steering wheel has a Toyota logo in the center and black leather around it. The car's dashboard is visible, showing a speedometer, tachometer, and navigation screen. The road ahead is straight and there are trees and fields visible on either side. The camera is positioned inside the car, providing a view from the driver's perspective. The lighting is natural and overcast, with a slightly cool tone.
|
||||
|
||||
user prompt: A young woman is sitting on a chair.
|
||||
enhanced prompt: A young woman with dark, curly hair and pale skin sits on a chair; she wears a dark, intricately patterned dress with a high collar and long, dark gloves that extend past her elbows; the scene is dimly lit, with light streaming in from a large window behind the characters.
|
||||
|
||||
user prompt: Aerial view of a city skyline.
|
||||
enhanced prompt: The camera pans across a cityscape of tall buildings with a circular building in the center. The camera moves from left to right, showing the tops of the buildings and the circular building in the center. The buildings are various shades of gray and white, and the circular building has a green roof. The camera angle is high, looking down at the city. The lighting is bright, with the sun shining from the upper left, casting shadows from the buildings.
|
||||
"""
|
||||
|
||||
I2V_CINEMATIC_PROMPT = """You are an expert cinematic director with many award winning movies, When writing prompts based on the user input, focus on detailed, chronological descriptions of actions and scenes.
|
||||
Include specific movements, appearances, camera angles, and environmental details - all in a single flowing paragraph.
|
||||
Start directly with the action, and keep descriptions literal and precise.
|
||||
Think like a cinematographer describing a shot list.
|
||||
Keep within 150 words.
|
||||
For best results, build your prompts using this structure:
|
||||
Describe the image first and then add the user input. Image description should be in first priority! Align to the image caption if it contradicts the user text input.
|
||||
Start with main action in a single sentence
|
||||
Add specific details about movements and gestures
|
||||
Describe character/object appearances precisely
|
||||
Include background and environment details
|
||||
Specify camera angles and movements
|
||||
Describe lighting and colors
|
||||
Note any changes or sudden events
|
||||
Align to the image caption if it contradicts the user text input.
|
||||
Do not exceed the 150 word limit!
|
||||
Output the enhanced prompt only.
|
||||
"""
|
||||
|
||||
|
||||
def tensor_to_pil(tensor):
|
||||
# Ensure tensor is in range [-1, 1]
|
||||
assert tensor.min() >= -1 and tensor.max() <= 1
|
||||
|
||||
# Convert from [-1, 1] to [0, 1]
|
||||
tensor = (tensor + 1) / 2
|
||||
|
||||
# Rearrange from [C, H, W] to [H, W, C]
|
||||
tensor = tensor.permute(1, 2, 0)
|
||||
|
||||
# Convert to numpy array and then to uint8 range [0, 255]
|
||||
numpy_image = (tensor.cpu().numpy() * 255).astype("uint8")
|
||||
|
||||
# Convert to PIL Image
|
||||
return Image.fromarray(numpy_image)
|
||||
|
||||
|
||||
def generate_cinematic_prompt(
|
||||
image_caption_model,
|
||||
image_caption_processor,
|
||||
prompt_enhancer_model,
|
||||
prompt_enhancer_tokenizer,
|
||||
prompt: Union[str, List[str]],
|
||||
conditioning_items: Optional[List[Tuple[torch.Tensor, int, float]]] = None,
|
||||
max_new_tokens: int = 256,
|
||||
) -> List[str]:
|
||||
prompts = [prompt] if isinstance(prompt, str) else prompt
|
||||
|
||||
if conditioning_items is None:
|
||||
prompts = _generate_t2v_prompt(
|
||||
prompt_enhancer_model,
|
||||
prompt_enhancer_tokenizer,
|
||||
prompts,
|
||||
max_new_tokens,
|
||||
T2V_CINEMATIC_PROMPT,
|
||||
)
|
||||
else:
|
||||
# if len(conditioning_items) > 1 or conditioning_items[0][1] != 0:
|
||||
# logger.warning(
|
||||
# "prompt enhancement does only support first frame of conditioning items, returning original prompts"
|
||||
# )
|
||||
# return prompts
|
||||
|
||||
first_frame_conditioning_item = conditioning_items[0]
|
||||
first_frames = _get_first_frames_from_conditioning_item(
|
||||
first_frame_conditioning_item
|
||||
)
|
||||
|
||||
assert len(first_frames) == len(
|
||||
prompts
|
||||
), "Number of conditioning frames must match number of prompts"
|
||||
|
||||
prompts = _generate_i2v_prompt(
|
||||
image_caption_model,
|
||||
image_caption_processor,
|
||||
prompt_enhancer_model,
|
||||
prompt_enhancer_tokenizer,
|
||||
prompts,
|
||||
first_frames,
|
||||
max_new_tokens,
|
||||
I2V_CINEMATIC_PROMPT,
|
||||
)
|
||||
|
||||
return prompts
|
||||
|
||||
|
||||
def _get_first_frames_from_conditioning_item(
|
||||
conditioning_item: Tuple[torch.Tensor, int, float]
|
||||
) -> List[Image.Image]:
|
||||
frames_tensor = conditioning_item[0]
|
||||
# tensor shape: [batch_size, 3, num_frames, height, width], take first frame from each sample
|
||||
return [
|
||||
tensor_to_pil(frames_tensor[i, :, 0, :, :])
|
||||
for i in range(frames_tensor.shape[0])
|
||||
]
|
||||
|
||||
|
||||
def _generate_t2v_prompt(
|
||||
prompt_enhancer_model,
|
||||
prompt_enhancer_tokenizer,
|
||||
prompts: List[str],
|
||||
max_new_tokens: int,
|
||||
system_prompt: str,
|
||||
) -> List[str]:
|
||||
messages = [
|
||||
[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": f"user_prompt: {p}"},
|
||||
]
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
texts = [
|
||||
prompt_enhancer_tokenizer.apply_chat_template(
|
||||
m, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
for m in messages
|
||||
]
|
||||
model_inputs = prompt_enhancer_tokenizer(texts, return_tensors="pt").to(
|
||||
prompt_enhancer_model.device
|
||||
)
|
||||
|
||||
return _generate_and_decode_prompts(
|
||||
prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens
|
||||
)
|
||||
|
||||
|
||||
def _generate_i2v_prompt(
|
||||
image_caption_model,
|
||||
image_caption_processor,
|
||||
prompt_enhancer_model,
|
||||
prompt_enhancer_tokenizer,
|
||||
prompts: List[str],
|
||||
first_frames: List[Image.Image],
|
||||
max_new_tokens: int,
|
||||
system_prompt: str,
|
||||
) -> List[str]:
|
||||
image_captions = _generate_image_captions(
|
||||
image_caption_model, image_caption_processor, first_frames
|
||||
)
|
||||
|
||||
messages = [
|
||||
[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": f"user_prompt: {p}\nimage_caption: {c}"},
|
||||
]
|
||||
for p, c in zip(prompts, image_captions)
|
||||
]
|
||||
|
||||
texts = [
|
||||
prompt_enhancer_tokenizer.apply_chat_template(
|
||||
m, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
for m in messages
|
||||
]
|
||||
model_inputs = prompt_enhancer_tokenizer(texts, return_tensors="pt").to(
|
||||
prompt_enhancer_model.device
|
||||
)
|
||||
|
||||
return _generate_and_decode_prompts(
|
||||
prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens
|
||||
)
|
||||
|
||||
|
||||
def _generate_image_captions(
|
||||
image_caption_model,
|
||||
image_caption_processor,
|
||||
images: List[Image.Image],
|
||||
system_prompt: str = "<DETAILED_CAPTION>",
|
||||
) -> List[str]:
|
||||
image_caption_prompts = [system_prompt] * len(images)
|
||||
inputs = image_caption_processor(
|
||||
image_caption_prompts, images, return_tensors="pt"
|
||||
).to(image_caption_model.device)
|
||||
|
||||
with torch.inference_mode():
|
||||
generated_ids = image_caption_model.generate(
|
||||
input_ids=inputs["input_ids"],
|
||||
pixel_values=inputs["pixel_values"],
|
||||
max_new_tokens=1024,
|
||||
do_sample=False,
|
||||
num_beams=3,
|
||||
)
|
||||
|
||||
return image_caption_processor.batch_decode(generated_ids, skip_special_tokens=True)
|
||||
|
||||
|
||||
def _get_random_scene_type():
|
||||
"""
|
||||
Randomly select a scene type to add to the prompt.
|
||||
"""
|
||||
types = [
|
||||
"The scene is captured in real-life footage.",
|
||||
"The scene is computer-generated imagery.",
|
||||
"The scene appears to be from a movie.",
|
||||
"The scene appears to be from a TV show.",
|
||||
"The scene is captured in a studio.",
|
||||
]
|
||||
return random.choice(types)
|
||||
|
||||
|
||||
def _generate_and_decode_prompts(
|
||||
prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens: int
|
||||
) -> List[str]:
|
||||
with torch.inference_mode():
|
||||
outputs = prompt_enhancer_model.generate(
|
||||
**model_inputs, max_new_tokens=max_new_tokens
|
||||
)
|
||||
generated_ids = [
|
||||
output_ids[len(input_ids) :]
|
||||
for input_ids, output_ids in zip(model_inputs.input_ids, outputs)
|
||||
]
|
||||
decoded_prompts = prompt_enhancer_tokenizer.batch_decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)
|
||||
|
||||
decoded_prompts = [p + f" {_get_random_scene_type()}." for p in decoded_prompts]
|
||||
|
||||
print(decoded_prompts)
|
||||
|
||||
return decoded_prompts
|
||||
+6
-2
@@ -9,11 +9,15 @@ requires-python = ">=3.10"
|
||||
readme = "README.md"
|
||||
license = { file = "LICENSE" }
|
||||
dependencies = [
|
||||
"ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.1",
|
||||
"av>=10.0.0"
|
||||
"huggingface_hub>=0.25.2",
|
||||
"transformers>=4.45.0"
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
internal = [
|
||||
"ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.1",
|
||||
"av>=10.0.0"
|
||||
]
|
||||
dev = [
|
||||
"pre-commit>=4.0.1",
|
||||
"pytest>=8.0.0",
|
||||
|
||||
+2
-2
@@ -1,2 +1,2 @@
|
||||
av>=10.0.0
|
||||
ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.1
|
||||
huggingface_hub>=0.25.2
|
||||
transformers>=4.45.0
|
||||
|
||||
@@ -1,8 +1,12 @@
|
||||
import contextlib
|
||||
from dataclasses import dataclass
|
||||
from typing import List
|
||||
|
||||
import comfy.ldm.modules.attention
|
||||
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
|
||||
|
||||
@@ -37,7 +41,6 @@ class LTXVApplySTG:
|
||||
"MODEL",
|
||||
{"tooltip": "The model to apply the STG to."},
|
||||
),
|
||||
"stg_mode": (["attention", "residual"],),
|
||||
"block_indices": (
|
||||
"STRING",
|
||||
{
|
||||
@@ -53,16 +56,10 @@ class LTXVApplySTG:
|
||||
RETURN_NAMES = ("model",)
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
def apply_stg(self, model: ModelPatcher, stg_mode: str, block_indices: str):
|
||||
def apply_stg(self, model: ModelPatcher, 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"]
|
||||
@@ -74,17 +71,105 @@ class LTXVApplySTG:
|
||||
return (new_model,)
|
||||
|
||||
|
||||
class STGGuider(comfy.samplers.CFGGuider):
|
||||
def set_conds(self, positive, negative):
|
||||
self.inner_set_conds(
|
||||
{"positive": positive, "negative": negative, "perturbed": positive}
|
||||
@dataclass
|
||||
class STGFlag:
|
||||
do_skip: bool = False
|
||||
skip_layers: List[int] = None
|
||||
|
||||
|
||||
# context manager that replaces the attention function in a transformer block
|
||||
class PatchAttention(contextlib.AbstractContextManager):
|
||||
def __init__(self, attn_idx=0):
|
||||
self.original_attention = comfy.ldm.modules.attention.optimized_attention
|
||||
self.original_attention_masked = (
|
||||
comfy.ldm.modules.attention.optimized_attention_masked
|
||||
)
|
||||
self.current_idx = -1
|
||||
self.attn_idx = attn_idx
|
||||
|
||||
def __enter__(self):
|
||||
comfy.ldm.modules.attention.optimized_attention = self.stg_attention
|
||||
comfy.ldm.modules.attention.optimized_attention_masked = (
|
||||
self.stg_attention_masked
|
||||
)
|
||||
|
||||
def set_cfg(self, cfg, stg_scale, rescale_scale: float = None):
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
comfy.ldm.modules.attention.optimized_attention = self.original_attention
|
||||
comfy.ldm.modules.attention.optimized_attention_masked = (
|
||||
self.original_attention_masked
|
||||
)
|
||||
|
||||
def stg_attention(self, q, k, v, heads, *args, **kwargs):
|
||||
self.current_idx += 1
|
||||
if self.current_idx == self.attn_idx:
|
||||
return v
|
||||
else:
|
||||
return self.original_attention(q, k, v, heads, *args, **kwargs)
|
||||
|
||||
def stg_attention_masked(self, q, k, v, heads, *args, **kwargs):
|
||||
self.current_idx += 1
|
||||
if self.current_idx == self.attn_idx:
|
||||
return v
|
||||
else:
|
||||
return self.original_attention_masked(q, k, v, heads, *args, **kwargs)
|
||||
|
||||
|
||||
class STGBlockWrapper:
|
||||
"""Wraps transformer blocks to be able to skip attention layers."""
|
||||
|
||||
def __init__(self, block, stg_flag: STGFlag, idx: int):
|
||||
self.flag = stg_flag
|
||||
self.idx = idx
|
||||
self.block = block
|
||||
|
||||
def __call__(self, args, extra_args):
|
||||
context_manager = contextlib.nullcontext()
|
||||
if self.flag.do_skip and self.idx in self.flag.skip_layers:
|
||||
context_manager = PatchAttention(0)
|
||||
|
||||
with context_manager:
|
||||
hidden_state = extra_args["original_block"](args)
|
||||
return hidden_state
|
||||
|
||||
|
||||
class STGGuider(comfy.samplers.CFGGuider):
|
||||
def __init__(
|
||||
self, model: ModelPatcher, cfg, stg_scale, rescale_scale: float = None
|
||||
):
|
||||
model = model.clone()
|
||||
super().__init__(model)
|
||||
|
||||
self.stg_flag = STGFlag(
|
||||
do_skip=False,
|
||||
skip_layers=model.model_options["transformer_options"]["skip_block_list"],
|
||||
)
|
||||
|
||||
self.patch_model(model, self.stg_flag)
|
||||
|
||||
self.cfg = cfg
|
||||
self.stg_scale = stg_scale
|
||||
self.rescale_scale = rescale_scale
|
||||
|
||||
@classmethod
|
||||
def patch_model(cls, model: ModelPatcher, stg_flag: STGFlag):
|
||||
transformer_blocks = cls.get_transformer_blocks(model)
|
||||
|
||||
for i, block in enumerate(transformer_blocks):
|
||||
model.set_model_patch_replace(
|
||||
STGBlockWrapper(block, stg_flag, i), "dit", "double_block", i
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_transformer_blocks(model: ModelPatcher):
|
||||
diffusion_model = model.get_model_object("diffusion_model")
|
||||
key = "diffusion_model.transformer_blocks"
|
||||
if diffusion_model.__class__.__name__ == "LTXVTransformer3D":
|
||||
key = "diffusion_model.transformer.transformer_blocks"
|
||||
return model.get_model_object(key)
|
||||
|
||||
def set_conds(self, positive, negative):
|
||||
self.inner_set_conds({"positive": positive, "negative": negative})
|
||||
|
||||
def predict_noise(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
@@ -95,32 +180,44 @@ class STGGuider(comfy.samplers.CFGGuider):
|
||||
# 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_pos = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[positive_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
|
||||
noise_pred_neg = 0
|
||||
# no similar optimization for stg=0, use CFG guider instead.
|
||||
noise_pred_perturbed = 0
|
||||
|
||||
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(
|
||||
noise_pred_neg = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[perturbed_cond, positive_cond],
|
||||
[negative_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)
|
||||
)[0]
|
||||
|
||||
if self.stg_scale > 0:
|
||||
try:
|
||||
model_options["transformer_options"]["ptb_index"] = 0
|
||||
self.stg_flag.do_skip = True
|
||||
noise_pred_perturbed = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[positive_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
finally:
|
||||
self.stg_flag.do_skip = False
|
||||
del model_options["transformer_options"]["ptb_index"]
|
||||
|
||||
stg_result = stg(
|
||||
noise_pred_pos,
|
||||
noise_pred_neg,
|
||||
@@ -190,7 +287,6 @@ class STGGuiderNode:
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
def get_guider(self, model, positive, negative, cfg, stg, rescale):
|
||||
guider = STGGuider(model)
|
||||
guider = STGGuider(model, cfg, stg, rescale)
|
||||
guider.set_conds(positive, negative)
|
||||
guider.set_cfg(cfg, stg, rescale)
|
||||
return (guider,)
|
||||
|
||||
-149
@@ -1,149 +0,0 @@
|
||||
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 folder_paths
|
||||
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):
|
||||
self.tokenizer = T5Tokenizer.from_pretrained(
|
||||
tokenizer_path, local_files_only=True
|
||||
)
|
||||
|
||||
def tokenize_with_weights(self, text: str, return_word_ids: bool = False):
|
||||
"""
|
||||
Takes a prompt and converts it to a list of (token, weight, word id) elements.
|
||||
Tokens can both be integer tokens and pre computed CLIP tensors.
|
||||
Word id values are unique per word and embedding, where the id 0 is reserved for non word tokens.
|
||||
Returned list has the dimensions NxM where M is the input size of CLIP
|
||||
"""
|
||||
text = text.lower().strip()
|
||||
text_inputs = self.tokenizer(
|
||||
text,
|
||||
padding="max_length", # do_not_pad, longest, max_length
|
||||
max_length=128,
|
||||
truncation=True,
|
||||
add_special_tokens=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
prompt_attention_mask = text_inputs.attention_mask
|
||||
|
||||
out = {
|
||||
"t5xxl": [
|
||||
(token, weight, i)
|
||||
for i, (token, weight) in enumerate(
|
||||
zip(text_input_ids[0], prompt_attention_mask[0])
|
||||
)
|
||||
]
|
||||
}
|
||||
|
||||
if not return_word_ids:
|
||||
out = {k: [(t, w) for t, w, _ in v] for k, v in out.items()}
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class LTXVTextEncoderModel(torch.nn.Module):
|
||||
def __init__(
|
||||
self, encoder_path, dtype_t5=None, device="cpu", dtype=None, model_options={}
|
||||
):
|
||||
super().__init__()
|
||||
dtype_t5 = comfy.model_management.pick_weight_dtype(dtype_t5, dtype, device)
|
||||
self.t5xxl = (
|
||||
T5EncoderModel.from_pretrained(encoder_path, local_files_only=True)
|
||||
.to(dtype_t5)
|
||||
.to(device)
|
||||
)
|
||||
self.dtypes = set([dtype, dtype_t5])
|
||||
|
||||
def set_clip_options(self, options):
|
||||
pass
|
||||
|
||||
def reset_clip_options(self):
|
||||
pass
|
||||
|
||||
def encode_token_weights(self, token_weight_pairs):
|
||||
token_weight_pairs_t5 = token_weight_pairs["t5xxl"]
|
||||
text_input_ids = torch.tensor(
|
||||
[[t[0] for t in token_weight_pairs_t5]],
|
||||
device=self.t5xxl.device,
|
||||
)
|
||||
prompt_attention_mask = torch.tensor(
|
||||
[[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
|
||||
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}
|
||||
|
||||
def load_sd(self, sd):
|
||||
return self.t5xxl.load_state_dict(sd, strict=False)
|
||||
|
||||
|
||||
def ltxv_clip(encoder_path, dtype_t5=None):
|
||||
class LTXVTextEncoderModel_(LTXVTextEncoderModel):
|
||||
def __init__(self, device="cpu", dtype=None, model_options={}):
|
||||
super().__init__(
|
||||
encoder_path=encoder_path,
|
||||
dtype_t5=dtype_t5,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
model_options=model_options,
|
||||
)
|
||||
|
||||
return LTXVTextEncoderModel_
|
||||
|
||||
|
||||
def ltxv_tokenizer(tokenizer_path):
|
||||
class LTXVTokenizer_(LTXVTokenizer):
|
||||
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
||||
super().__init__(tokenizer_path)
|
||||
|
||||
return LTXVTokenizer_
|
||||
|
||||
|
||||
@comfy_node(name="LTXVCLIPModelLoader", description="LTXV CLIP Model Loader")
|
||||
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_TYPES = ("CLIP",)
|
||||
RETURN_NAMES = ("clip",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Model Loader"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def load_model(self, clip_path):
|
||||
path = Path(folder_paths.get_full_path("text_encoders", clip_path))
|
||||
tokenizer_path = path.parents[1] / "tokenizer"
|
||||
encoder_path = path.parents[1] / "text_encoder"
|
||||
|
||||
clip_target = comfy.supported_models_base.ClipTarget(
|
||||
tokenizer=ltxv_tokenizer(tokenizer_path),
|
||||
clip=ltxv_clip(encoder_path, dtype_t5=torch.bfloat16),
|
||||
)
|
||||
|
||||
return (comfy.sd.CLIP(clip_target),)
|
||||
-268
@@ -1,268 +0,0 @@
|
||||
import math
|
||||
|
||||
import comfy.latent_formats
|
||||
import comfy.model_base
|
||||
import comfy.model_management
|
||||
import comfy.model_patcher
|
||||
import comfy.sd
|
||||
import comfy.supported_models_base
|
||||
import comfy.utils
|
||||
import torch
|
||||
from ltx_video.models.autoencoders.vae_encode import get_vae_size_scale_factor
|
||||
|
||||
from .img2vid import encode_media_conditioning
|
||||
from .model import LTXVSampling
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
def get_normal_shift(
|
||||
n_tokens: int,
|
||||
min_tokens: int = 1024,
|
||||
max_tokens: int = 4096,
|
||||
min_shift: float = 0.95,
|
||||
max_shift: float = 2.05,
|
||||
) -> float:
|
||||
m = (max_shift - min_shift) / (max_tokens - min_tokens)
|
||||
b = min_shift - m * min_tokens
|
||||
return m * n_tokens + b
|
||||
|
||||
|
||||
@comfy_node(name="LTXVModelConfigurator")
|
||||
class LTXVModelConfigurator:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
PRESETS = [
|
||||
"Custom",
|
||||
"1216x704 | 41",
|
||||
"1088x704 | 49",
|
||||
"1056x640 | 57",
|
||||
"992x608 | 65",
|
||||
"896x608 | 73",
|
||||
"896x544 | 81",
|
||||
"832x544 | 89",
|
||||
"800x512 | 97",
|
||||
"768x512 | 97",
|
||||
"800x480 | 105",
|
||||
"736x480 | 113",
|
||||
"704x480 | 121",
|
||||
"704x448 | 129",
|
||||
"672x448 | 137",
|
||||
"640x416 | 153",
|
||||
"672x384 | 161",
|
||||
"640x384 | 169",
|
||||
"608x384 | 177",
|
||||
"576x384 | 185",
|
||||
"608x352 | 193",
|
||||
"576x352 | 201",
|
||||
"544x352 | 209",
|
||||
"512x352 | 225",
|
||||
"512x352 | 233",
|
||||
"544x320 | 241",
|
||||
"512x320 | 249",
|
||||
"512x320 | 257",
|
||||
]
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"vae": ("VAE",),
|
||||
"preset": (
|
||||
PRESETS,
|
||||
{
|
||||
"default": "Custom",
|
||||
"tooltip": "Preset resolution and frame count. Custom allows manual input.",
|
||||
},
|
||||
),
|
||||
"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",
|
||||
},
|
||||
),
|
||||
"frame_rate": ("INT", {"default": 25, "min": 1, "max": 60}),
|
||||
"batch": ("INT", {"default": 1, "min": 1, "max": 60}),
|
||||
"mixed_precision": ("BOOLEAN", {"default": True}),
|
||||
"img_compression": (
|
||||
"INT",
|
||||
{
|
||||
"default": 29,
|
||||
"min": 0,
|
||||
"max": 100,
|
||||
"tooltip": "Amount of compression to apply on conditioning image.",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"conditioning": (
|
||||
"IMAGE",
|
||||
{"tooltip": "Optional conditioning image or video."},
|
||||
),
|
||||
"initial_latent": (
|
||||
"LATENT",
|
||||
{
|
||||
"tooltip": "initial latent that is combined with conditioning if given"
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "LATENT", "FLOAT")
|
||||
RETURN_NAMES = ("model", "latent", "sigma_shift")
|
||||
FUNCTION = "configure_sizes"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Model Configurator"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def latent_shape_and_frame_rate(
|
||||
self, vae, batch, height, width, frames_number, frame_rate
|
||||
):
|
||||
video_scale_factor, vae_scale_factor, _ = get_vae_size_scale_factor(
|
||||
vae.first_stage_model
|
||||
)
|
||||
video_scale_factor = video_scale_factor if frames_number > 1 else 1
|
||||
|
||||
latent_height = height // vae_scale_factor
|
||||
latent_width = width // vae_scale_factor
|
||||
latent_channels = vae.first_stage_model.config.latent_channels
|
||||
latent_num_frames = math.floor(frames_number / video_scale_factor) + 1
|
||||
latent_frame_rate = frame_rate / video_scale_factor
|
||||
|
||||
latent_shape = [
|
||||
batch,
|
||||
latent_channels,
|
||||
latent_num_frames,
|
||||
latent_height,
|
||||
latent_width,
|
||||
]
|
||||
return latent_shape, latent_frame_rate
|
||||
|
||||
def configure_sizes(
|
||||
self,
|
||||
model,
|
||||
vae,
|
||||
preset,
|
||||
width,
|
||||
height,
|
||||
frames_number,
|
||||
frame_rate,
|
||||
batch,
|
||||
mixed_precision,
|
||||
img_compression,
|
||||
conditioning=None,
|
||||
initial_latent=None,
|
||||
):
|
||||
load_device = comfy.model_management.get_torch_device()
|
||||
if preset != "Custom":
|
||||
preset = preset.split("|")
|
||||
width, height = map(int, preset[0].strip().split("x"))
|
||||
frames_number = int(preset[1].strip())
|
||||
latent_shape, latent_frame_rate = self.latent_shape_and_frame_rate(
|
||||
vae, batch, height, width, frames_number, frame_rate
|
||||
)
|
||||
mask_shape = [
|
||||
latent_shape[0],
|
||||
1,
|
||||
latent_shape[2],
|
||||
latent_shape[3],
|
||||
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,
|
||||
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, device=load_device)
|
||||
if initial_latent is not None:
|
||||
latent[:, :, : initial_latent.shape[2], ...] = initial_latent
|
||||
|
||||
_, vae_scale_factor, _ = get_vae_size_scale_factor(vae.first_stage_model)
|
||||
|
||||
patcher = model.clone()
|
||||
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(conditioning_mask, guiding_latent)
|
||||
)
|
||||
patcher.model_options.setdefault("transformer_options", {})[
|
||||
"mixed_precision"
|
||||
] = mixed_precision
|
||||
|
||||
num_latent_patches = latent_shape[2] * latent_shape[3] * latent_shape[4]
|
||||
return (patcher, {"samples": latent}, get_normal_shift(num_latent_patches))
|
||||
|
||||
|
||||
@comfy_node(name="LTXVShiftSigmas")
|
||||
class LTXVShiftSigmas:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"sigmas": ("SIGMAS",),
|
||||
"sigma_shift": ("FLOAT", {"default": 1.820833333}),
|
||||
"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.",
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SIGMAS",)
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
FUNCTION = "shift_sigmas"
|
||||
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
|
||||
sigmas = torch.where(
|
||||
sigmas != 0,
|
||||
math.exp(sigma_shift) / (math.exp(sigma_shift) + (1 / sigmas - 1) ** power),
|
||||
0,
|
||||
)
|
||||
|
||||
# Stretch sigmas so that its final value matches the given terminal value.
|
||||
if stretch:
|
||||
non_zero_mask = sigmas != 0
|
||||
non_zero_sigmas = sigmas[non_zero_mask]
|
||||
one_minus_z = 1.0 - non_zero_sigmas
|
||||
scale_factor = one_minus_z[-1] / (1.0 - terminal)
|
||||
stretched = 1.0 - (one_minus_z / scale_factor)
|
||||
sigmas[non_zero_mask] = stretched
|
||||
|
||||
return (sigmas,)
|
||||
@@ -0,0 +1,54 @@
|
||||
from .nodes.attn_bank_nodes import (
|
||||
LTXAttentionBankNode,
|
||||
LTXAttentioOverrideNode,
|
||||
LTXPrepareAttnInjectionsNode,
|
||||
)
|
||||
from .nodes.attn_override_node import LTXAttnOverrideNode
|
||||
from .nodes.latent_guide_node import AddLatentGuideNode
|
||||
from .nodes.ltx_feta_enhance_node import LTXFetaEnhanceNode
|
||||
from .nodes.ltx_flowedit_nodes import LTXFlowEditCFGGuiderNode, LTXFlowEditSamplerNode
|
||||
from .nodes.ltx_inverse_model_pred_nodes import (
|
||||
LTXForwardModelSamplingPredNode,
|
||||
LTXReverseModelSamplingPredNode,
|
||||
)
|
||||
from .nodes.ltx_pag_node import LTXPerturbedAttentionNode
|
||||
from .nodes.modify_ltx_model_node import ModifyLTXModelNode
|
||||
from .nodes.rectified_sampler_nodes import (
|
||||
LTXRFForwardODESamplerNode,
|
||||
LTXRFReverseODESamplerNode,
|
||||
)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ModifyLTXModel": ModifyLTXModelNode,
|
||||
"AddLatentGuide": AddLatentGuideNode,
|
||||
"LTXForwardModelSamplingPred": LTXForwardModelSamplingPredNode,
|
||||
"LTXReverseModelSamplingPred": LTXReverseModelSamplingPredNode,
|
||||
"LTXRFForwardODESampler": LTXRFForwardODESamplerNode,
|
||||
"LTXRFReverseODESampler": LTXRFReverseODESamplerNode,
|
||||
"LTXAttentionBank": LTXAttentionBankNode,
|
||||
"LTXPrepareAttnInjections": LTXPrepareAttnInjectionsNode,
|
||||
"LTXAttentioOverride": LTXAttentioOverrideNode,
|
||||
"LTXPerturbedAttention": LTXPerturbedAttentionNode,
|
||||
"LTXAttnOverride": LTXAttnOverrideNode,
|
||||
"LTXFlowEditCFGGuider": LTXFlowEditCFGGuiderNode,
|
||||
"LTXFlowEditSampler": LTXFlowEditSamplerNode,
|
||||
"LTXFetaEnhance": LTXFetaEnhanceNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ModifyLTXModel": "Modify LTX Model",
|
||||
"AddLatentGuide": "Add LTX Latent Guide",
|
||||
"LTXAddImageGuide": "Add LTX Image Guide",
|
||||
"LTXForwardModelSamplingPred": "LTX Forward Model Pred",
|
||||
"LTXReverseModelSamplingPred": "LTX Reverse Model Pred",
|
||||
"LTXRFForwardODESampler": "LTX Rf-Inv Forward Sampler",
|
||||
"LTXRFReverseODESampler": "LTX Rf-Inv Reverse Sampler",
|
||||
"LTXAttentionBank": "LTX Attention Bank",
|
||||
"LTXPrepareAttnInjections": "LTX Prepare Attn Injection",
|
||||
"LTXAttentioOverride": "LTX Attn Block Override",
|
||||
"LTXPerturbedAttention": "LTX Apply Perturbed Attention",
|
||||
"LTXAttnOverride": "LTX Attention Override",
|
||||
"LTXFlowEditCFGGuider": "LTX Flow Edit CFG Guider",
|
||||
"LTXFlowEditSampler": "LTX Flow Edit Sampler",
|
||||
"LTXFetaEnhance": "LTX Feta Enhance",
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:a8ea20c84d49818cf120b18c889cc8ead2e8298539b4db4c6bf769c45987ddfc
|
||||
size 2264904
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:9d2b002f12155d0b929366ad38561680a5bbaecea1925e5655658b46d281baf9
|
||||
size 2441562
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3726ed083a5c52873d29adab87e8fe0a65c7e390beaf8dfcf23e5c9091d6da73
|
||||
size 521574
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:01d44bb728e04879a5513a2c33bdd47d443bf901b33d3c1445918e75ba8b78b1
|
||||
size 678605
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:5432150d5477d19586b8d8049ca214038f886d3def9117830cf2af58879c48b5
|
||||
size 31584408
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 349 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 342 KiB |
@@ -0,0 +1,258 @@
|
||||
import math
|
||||
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.ldm.modules.attention
|
||||
import torch
|
||||
from comfy.ldm.lightricks.model import (
|
||||
BasicTransformerBlock,
|
||||
LTXVModel,
|
||||
apply_rotary_emb,
|
||||
precompute_freqs_cis,
|
||||
)
|
||||
from comfy.ldm.lightricks.symmetric_patchifier import latent_to_pixel_coords
|
||||
from torch import nn
|
||||
|
||||
from ..utils.feta_enhance_utils import get_feta_scores
|
||||
|
||||
|
||||
class LTXModifiedCrossAttention(nn.Module):
|
||||
def forward(self, x, context=None, mask=None, pe=None, transformer_options={}):
|
||||
context = x if context is None else context
|
||||
context_v = x if context is None else context
|
||||
|
||||
step = transformer_options.get("step", -1)
|
||||
total_steps = transformer_options.get("total_steps", 0)
|
||||
attn_bank = transformer_options.get("attn_bank", None)
|
||||
sample_mode = transformer_options.get("sample_mode", None)
|
||||
if attn_bank is not None and self.idx in attn_bank["block_map"]:
|
||||
len_conds = len(transformer_options["cond_or_uncond"])
|
||||
pred_order = transformer_options["pred_order"]
|
||||
if (
|
||||
sample_mode == "forward"
|
||||
and total_steps - step - 1 < attn_bank["save_steps"]
|
||||
):
|
||||
step_idx = f"{pred_order}_{total_steps-step-1}"
|
||||
attn_bank["block_map"][self.idx][step_idx] = x.cpu()
|
||||
elif sample_mode == "reverse" and step < attn_bank["inject_steps"]:
|
||||
step_idx = f"{pred_order}_{step}"
|
||||
inject_settings = attn_bank.get("inject_settings", {})
|
||||
if len(inject_settings) > 0:
|
||||
inj = (
|
||||
attn_bank["block_map"][self.idx][step_idx]
|
||||
.to(x.device)
|
||||
.repeat(len_conds, 1, 1)
|
||||
)
|
||||
if "q" in inject_settings:
|
||||
x = inj
|
||||
if "k" in inject_settings:
|
||||
context = inj
|
||||
if "v" in inject_settings:
|
||||
context_v = inj
|
||||
|
||||
q = self.to_q(x)
|
||||
k = self.to_k(context)
|
||||
v = self.to_v(context_v)
|
||||
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
|
||||
if pe is not None:
|
||||
q = apply_rotary_emb(q, pe)
|
||||
k = apply_rotary_emb(k, pe)
|
||||
|
||||
feta_score = None
|
||||
if (
|
||||
transformer_options.get("feta_weight", 0) > 0
|
||||
and self.idx in transformer_options["feta_layers"]["layers"]
|
||||
):
|
||||
feta_score = get_feta_scores(q, k, self.heads, transformer_options)
|
||||
|
||||
alt_attn_fn = (
|
||||
transformer_options.get("patches_replace", {})
|
||||
.get("layer", {})
|
||||
.get(("self_attn", self.idx), None)
|
||||
)
|
||||
if alt_attn_fn is not None:
|
||||
out = alt_attn_fn(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
self.heads,
|
||||
attn_precision=self.attn_precision,
|
||||
transformer_options=transformer_options,
|
||||
)
|
||||
elif mask is None:
|
||||
out = comfy.ldm.modules.attention.optimized_attention(
|
||||
q, k, v, self.heads, attn_precision=self.attn_precision
|
||||
)
|
||||
else:
|
||||
out = comfy.ldm.modules.attention.optimized_attention_masked(
|
||||
q, k, v, self.heads, mask, attn_precision=self.attn_precision
|
||||
)
|
||||
|
||||
if feta_score is not None:
|
||||
out *= feta_score
|
||||
|
||||
return self.to_out(out)
|
||||
|
||||
|
||||
class LTXModifiedBasicTransformerBlock(BasicTransformerBlock):
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
context=None,
|
||||
attention_mask=None,
|
||||
timestep=None,
|
||||
pe=None,
|
||||
transformer_options={},
|
||||
):
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
||||
self.scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
|
||||
+ timestep.reshape(
|
||||
x.shape[0], timestep.shape[1], self.scale_shift_table.shape[0], -1
|
||||
)
|
||||
).unbind(dim=2)
|
||||
x += (
|
||||
self.attn1(
|
||||
comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa,
|
||||
pe=pe,
|
||||
transformer_options=transformer_options,
|
||||
)
|
||||
* gate_msa
|
||||
)
|
||||
|
||||
x += self.attn2(x, context=context, mask=attention_mask)
|
||||
|
||||
y = comfy.ldm.common_dit.rms_norm(x) * (1 + scale_mlp) + shift_mlp
|
||||
x += self.ff(y) * gate_mlp
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class LTXVModelModified(LTXVModel):
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
timestep,
|
||||
context,
|
||||
attention_mask,
|
||||
frame_rate=25,
|
||||
transformer_options={},
|
||||
keyframe_idxs=None,
|
||||
**kwargs,
|
||||
):
|
||||
patches_replace = transformer_options.get("patches_replace", {})
|
||||
|
||||
orig_shape = list(x.shape)
|
||||
|
||||
x, latent_coords = self.patchifier.patchify(x)
|
||||
pixel_coords = latent_to_pixel_coords(
|
||||
latent_coords=latent_coords,
|
||||
scale_factors=self.vae_scale_factors,
|
||||
causal_fix=self.causal_temporal_positioning,
|
||||
)
|
||||
|
||||
if keyframe_idxs is not None:
|
||||
pixel_coords[:, :, -keyframe_idxs.shape[2] :] = keyframe_idxs
|
||||
|
||||
fractional_coords = pixel_coords.to(torch.float32)
|
||||
fractional_coords[:, 0] = fractional_coords[:, 0] * (1.0 / frame_rate)
|
||||
|
||||
x = self.patchify_proj(x)
|
||||
timestep = timestep * 1000.0
|
||||
|
||||
if attention_mask is not None and not torch.is_floating_point(attention_mask):
|
||||
attention_mask = (attention_mask - 1).to(x.dtype).reshape(
|
||||
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
|
||||
) * torch.finfo(x.dtype).max
|
||||
|
||||
pe = precompute_freqs_cis(
|
||||
fractional_coords, dim=self.inner_dim, out_dtype=x.dtype
|
||||
)
|
||||
|
||||
batch_size = x.shape[0]
|
||||
timestep, embedded_timestep = self.adaln_single(
|
||||
timestep.flatten(),
|
||||
{"resolution": None, "aspect_ratio": None},
|
||||
batch_size=batch_size,
|
||||
hidden_dtype=x.dtype,
|
||||
)
|
||||
# Second dimension is 1 or number of tokens (if timestep_per_token)
|
||||
timestep = timestep.view(batch_size, -1, timestep.shape[-1])
|
||||
embedded_timestep = embedded_timestep.view(
|
||||
batch_size, -1, embedded_timestep.shape[-1]
|
||||
)
|
||||
|
||||
# 2. Blocks
|
||||
if self.caption_projection is not None:
|
||||
batch_size = x.shape[0]
|
||||
context = self.caption_projection(context)
|
||||
context = context.view(batch_size, -1, x.shape[-1])
|
||||
|
||||
blocks_replace = patches_replace.get("dit", {})
|
||||
for i, block in enumerate(self.transformer_blocks):
|
||||
if ("double_block", i) in blocks_replace:
|
||||
|
||||
def block_wrap(args):
|
||||
out = {}
|
||||
out["img"] = block(
|
||||
args["img"],
|
||||
context=args["txt"],
|
||||
attention_mask=args["attention_mask"],
|
||||
timestep=args["vec"],
|
||||
pe=args["pe"],
|
||||
)
|
||||
return out
|
||||
|
||||
out = blocks_replace[("double_block", i)](
|
||||
{
|
||||
"img": x,
|
||||
"txt": context,
|
||||
"attention_mask": attention_mask,
|
||||
"vec": timestep,
|
||||
"pe": pe,
|
||||
},
|
||||
{"original_block": block_wrap},
|
||||
)
|
||||
x = out["img"]
|
||||
else:
|
||||
x = block(
|
||||
x,
|
||||
context=context,
|
||||
attention_mask=attention_mask,
|
||||
timestep=timestep,
|
||||
pe=pe,
|
||||
transformer_options=transformer_options,
|
||||
)
|
||||
|
||||
# 3. Output
|
||||
scale_shift_values = (
|
||||
self.scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
|
||||
+ embedded_timestep[:, :, None]
|
||||
)
|
||||
shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
|
||||
x = self.norm_out(x)
|
||||
# Modulation
|
||||
x = x * (1 + scale) + shift
|
||||
x = self.proj_out(x)
|
||||
|
||||
x = self.patchifier.unpatchify(
|
||||
latents=x,
|
||||
output_height=orig_shape[3],
|
||||
output_width=orig_shape[4],
|
||||
output_num_frames=orig_shape[2],
|
||||
out_channels=orig_shape[1] // math.prod(self.patchifier.patch_size),
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def inject_model(diffusion_model):
|
||||
diffusion_model.__class__ = LTXVModelModified
|
||||
for idx, transformer_block in enumerate(diffusion_model.transformer_blocks):
|
||||
transformer_block.__class__ = LTXModifiedBasicTransformerBlock
|
||||
transformer_block.idx = idx
|
||||
transformer_block.attn1.__class__ = LTXModifiedCrossAttention
|
||||
transformer_block.attn1.idx = idx
|
||||
return diffusion_model
|
||||
@@ -0,0 +1,90 @@
|
||||
from ..utils.attn_bank import AttentionBank
|
||||
|
||||
|
||||
class LTXAttentionBankNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"save_steps": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"blocks": ("STRING", {"multiline": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("ATTN_BANK",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(self, save_steps, blocks=""):
|
||||
block_map = {}
|
||||
block_list = blocks.split(",")
|
||||
for block in block_list:
|
||||
block_idx = int(block)
|
||||
block_map[block_idx] = {}
|
||||
|
||||
bank = AttentionBank(save_steps, block_map)
|
||||
return (bank,)
|
||||
|
||||
|
||||
class LTXPrepareAttnInjectionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latent": ("LATENT",),
|
||||
"attn_bank": ("ATTN_BANK",),
|
||||
"query": ("BOOLEAN", {"default": False}),
|
||||
"key": ("BOOLEAN", {"default": False}),
|
||||
"value": ("BOOLEAN", {"default": False}),
|
||||
"inject_steps": (
|
||||
"INT",
|
||||
{"default": 0, "min": 0, "max": 1000, "step": 1},
|
||||
),
|
||||
},
|
||||
"optional": {"blocks": ("LTX_BLOCKS",)},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT", "ATTN_INJ")
|
||||
FUNCTION = "prepare"
|
||||
|
||||
CATEGORY = "fluxtapoz"
|
||||
|
||||
def prepare(self, latent, attn_bank, query, key, value, inject_steps, blocks=None):
|
||||
if inject_steps > attn_bank["save_steps"]:
|
||||
raise ValueError("Can not inject more steps than were saved.")
|
||||
attn_bank = AttentionBank(
|
||||
attn_bank["save_steps"], attn_bank["block_map"], inject_steps
|
||||
)
|
||||
attn_bank["inject_settings"] = set([])
|
||||
if query:
|
||||
attn_bank["inject_settings"].add("q")
|
||||
if key:
|
||||
attn_bank["inject_settings"].add("k")
|
||||
if value:
|
||||
attn_bank["inject_settings"].add("v")
|
||||
|
||||
if blocks is not None:
|
||||
attn_bank["block_map"] = {**attn_bank["block_map"]}
|
||||
for key in list(attn_bank["block_map"].keys()):
|
||||
if key not in blocks:
|
||||
del attn_bank["block_map"][key]
|
||||
|
||||
# Hack to force order of operations in ComfyUI graph
|
||||
return (latent, attn_bank)
|
||||
|
||||
|
||||
class LTXAttentioOverrideNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"blocks": ("STRING", {"multiline": True})}}
|
||||
|
||||
RETURN_TYPES = ("LTX_BLOCKS",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(self, blocks=""):
|
||||
block_set = set(list(int(block) for block in blocks.split(",")))
|
||||
|
||||
return (block_set,)
|
||||
@@ -0,0 +1,30 @@
|
||||
def is_integer(string):
|
||||
try:
|
||||
int(string)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
class LTXAttnOverrideNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"layers": ("STRING", {"multiline": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("ATTN_OVERRIDE",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks/attn"
|
||||
|
||||
def build(self, layers):
|
||||
layers_map = set([])
|
||||
for block in layers.split(","):
|
||||
block = block.strip()
|
||||
if is_integer(block):
|
||||
layers_map.add(int(block))
|
||||
|
||||
return ({"layers": layers_map},)
|
||||
@@ -0,0 +1,39 @@
|
||||
import comfy_extras.nodes_lt as nodes_lt
|
||||
|
||||
|
||||
class AddLatentGuideNode(nodes_lt.LTXVAddGuide):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"latent": ("LATENT",),
|
||||
"image_latent": ("LATENT",),
|
||||
"index": ("INT", {"default": 0, "min": -1, "max": 9999, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "LATENT")
|
||||
RETURN_NAMES = ("model", "latent")
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
FUNCTION = "generate"
|
||||
|
||||
def generate(self, model, latent, image_latent, index):
|
||||
noise_mask = nodes_lt.get_noise_mask(latent)
|
||||
latent = latent["samples"]
|
||||
|
||||
image_latent = image_latent["samples"]
|
||||
|
||||
latent, noise_mask = self.replace_latent_frames(
|
||||
latent,
|
||||
noise_mask,
|
||||
image_latent,
|
||||
index,
|
||||
1.0,
|
||||
)
|
||||
|
||||
return (
|
||||
model,
|
||||
{"samples": latent, "noise_mask": noise_mask},
|
||||
)
|
||||
@@ -0,0 +1,36 @@
|
||||
DEFAULT_ATTN = {
|
||||
"layers": [i for i in range(0, 100, 1)],
|
||||
}
|
||||
|
||||
|
||||
class LTXFetaEnhanceNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"feta_weight": (
|
||||
"FLOAT",
|
||||
{"default": 4, "min": -100.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
},
|
||||
"optional": {"attn_override": ("ATTN_OVERRIDE",)},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
FUNCTION = "apply"
|
||||
|
||||
def apply(self, model, feta_weight, attn_override=DEFAULT_ATTN):
|
||||
model = model.clone()
|
||||
|
||||
model_options = model.model_options.copy()
|
||||
transformer_options = model_options["transformer_options"].copy()
|
||||
|
||||
transformer_options["feta_weight"] = feta_weight
|
||||
transformer_options["feta_layers"] = attn_override
|
||||
model_options["transformer_options"] = transformer_options
|
||||
|
||||
model.model_options = model_options
|
||||
return (model,)
|
||||
@@ -0,0 +1,180 @@
|
||||
import torch
|
||||
from comfy.samplers import KSAMPLER, CFGGuider, sampling_function
|
||||
from tqdm import trange
|
||||
|
||||
|
||||
class FlowEditGuider(CFGGuider):
|
||||
def __init__(self, model_patcher):
|
||||
super().__init__(model_patcher)
|
||||
self.cfgs = {}
|
||||
|
||||
def set_conds(self, **kwargs):
|
||||
self.inner_set_conds(kwargs)
|
||||
|
||||
def set_cfgs(self, **kwargs):
|
||||
self.cfgs = {**kwargs}
|
||||
|
||||
def predict_noise(self, x, timestep, model_options={}, seed=None):
|
||||
latent_type = model_options["transformer_options"]["latent_type"]
|
||||
positive = self.conds.get(f"{latent_type}_positive", None)
|
||||
negative = self.conds.get(f"{latent_type}_negative", None)
|
||||
cfg = self.cfgs.get(latent_type, self.cfg)
|
||||
return sampling_function(
|
||||
self.inner_model,
|
||||
x,
|
||||
timestep,
|
||||
negative,
|
||||
positive,
|
||||
cfg,
|
||||
model_options=model_options,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
|
||||
class LTXFlowEditCFGGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"source_pos": ("CONDITIONING",),
|
||||
"source_neg": ("CONDITIONING",),
|
||||
"target_pos": ("CONDITIONING",),
|
||||
"target_neg": ("CONDITIONING",),
|
||||
"source_cfg": (
|
||||
"FLOAT",
|
||||
{"default": 2, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "step": 0.01},
|
||||
),
|
||||
"target_cfg": (
|
||||
"FLOAT",
|
||||
{"default": 4.5, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "step": 0.01},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
|
||||
FUNCTION = "get_guider"
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def get_guider(
|
||||
self,
|
||||
model,
|
||||
source_pos,
|
||||
source_neg,
|
||||
target_pos,
|
||||
target_neg,
|
||||
source_cfg,
|
||||
target_cfg,
|
||||
):
|
||||
guider = FlowEditGuider(model)
|
||||
guider.set_conds(
|
||||
source_positive=source_pos,
|
||||
source_negative=source_neg,
|
||||
target_positive=target_pos,
|
||||
target_negative=target_neg,
|
||||
)
|
||||
guider.set_cfgs(source=source_cfg, target=target_cfg)
|
||||
return (guider,)
|
||||
|
||||
|
||||
def get_flowedit_sample(skip_steps, refine_steps, seed):
|
||||
generator = torch.manual_seed(seed)
|
||||
|
||||
@torch.no_grad()
|
||||
def flowedit_sample(
|
||||
model, x_init, sigmas, extra_args=None, callback=None, disable=None
|
||||
):
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
|
||||
model_options = extra_args.get("model_options", {})
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {**transformer_options}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
denoise_mask = extra_args.get("denoise_mask", 1)
|
||||
if denoise_mask is None:
|
||||
denoise_mask = 1
|
||||
else:
|
||||
extra_args["denoise_mask"] = torch.ones_like(denoise_mask)
|
||||
|
||||
source_extra_args = {
|
||||
**extra_args,
|
||||
"model_options": {
|
||||
"transformer_options": {**transformer_options, "latent_type": "source"}
|
||||
},
|
||||
}
|
||||
|
||||
sigmas = sigmas[skip_steps:]
|
||||
|
||||
x_tgt = x_init.clone()
|
||||
N = len(sigmas) - 1
|
||||
s_in = x_init.new_ones([x_init.shape[0]])
|
||||
|
||||
for i in trange(N, disable=disable):
|
||||
sigma = sigmas[i]
|
||||
noise = torch.randn(x_init.shape, generator=generator).to(x_init.device)
|
||||
|
||||
zt_src = (1 - sigma) * x_init + sigma * noise
|
||||
|
||||
if i < N - refine_steps:
|
||||
zt_tgt = x_tgt + (zt_src - x_init) * denoise_mask
|
||||
transformer_options["latent_type"] = "source"
|
||||
source_extra_args["model_options"]["transformer_options"][
|
||||
"latent_type"
|
||||
] = "source"
|
||||
vt_src = model(zt_src, sigma * s_in, **source_extra_args)
|
||||
else:
|
||||
if i == N - refine_steps:
|
||||
x_tgt = x_tgt + (zt_src - x_init) * denoise_mask
|
||||
zt_tgt = x_tgt
|
||||
vt_src = 0
|
||||
|
||||
transformer_options["latent_type"] = "target"
|
||||
vt_tgt = model(zt_tgt, sigma * s_in, **extra_args)
|
||||
|
||||
v_delta = (vt_tgt - vt_src) * denoise_mask
|
||||
x_tgt += (sigmas[i + 1] - sigmas[i]) * v_delta
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{
|
||||
"x": x_tgt,
|
||||
"denoised": x_tgt,
|
||||
"i": i + skip_steps,
|
||||
"sigma": sigmas[i],
|
||||
"sigma_hat": sigmas[i],
|
||||
}
|
||||
)
|
||||
|
||||
return x_tgt
|
||||
|
||||
return flowedit_sample
|
||||
|
||||
|
||||
class LTXFlowEditSamplerNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"skip_steps": (
|
||||
"INT",
|
||||
{"default": 4, "min": 0, "max": 0xFFFFFFFFFFFFFFFF},
|
||||
),
|
||||
"refine_steps": (
|
||||
"INT",
|
||||
{"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF},
|
||||
),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
||||
},
|
||||
"optional": {},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLER",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(self, skip_steps, refine_steps, seed):
|
||||
sampler = KSAMPLER(get_flowedit_sample(skip_steps, refine_steps, seed))
|
||||
return (sampler,)
|
||||
@@ -0,0 +1,91 @@
|
||||
import comfy.latent_formats
|
||||
import comfy.model_sampling
|
||||
import comfy.sd
|
||||
|
||||
|
||||
class InverseCONST:
|
||||
def calculate_input(self, sigma, noise):
|
||||
return noise
|
||||
|
||||
def calculate_denoised(self, sigma, model_output, model_input):
|
||||
sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
|
||||
return model_output
|
||||
|
||||
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
|
||||
return latent_image
|
||||
|
||||
def inverse_noise_scaling(self, sigma, latent):
|
||||
return latent
|
||||
|
||||
|
||||
class LTXForwardModelSamplingPredNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def patch(self, model):
|
||||
m = model.clone()
|
||||
|
||||
sampling_base = comfy.model_sampling.ModelSamplingFlux
|
||||
sampling_type = InverseCONST
|
||||
|
||||
class ModelSamplingAdvanced(sampling_base, sampling_type):
|
||||
pass
|
||||
|
||||
model_sampling = ModelSamplingAdvanced(model.model.model_config)
|
||||
model_sampling.set_parameters(shift=1.15)
|
||||
m.add_object_patch("model_sampling", model_sampling)
|
||||
return (m,)
|
||||
|
||||
|
||||
class ReverseCONST:
|
||||
def calculate_input(self, sigma, noise):
|
||||
return noise
|
||||
|
||||
def calculate_denoised(self, sigma, model_output, model_input):
|
||||
sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
|
||||
return model_output # model_input - model_output * sigma
|
||||
|
||||
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
|
||||
return latent_image
|
||||
|
||||
def inverse_noise_scaling(self, sigma, latent):
|
||||
return latent / (1.0 - sigma)
|
||||
|
||||
|
||||
class LTXReverseModelSamplingPredNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def patch(self, model):
|
||||
m = model.clone()
|
||||
|
||||
sampling_base = comfy.model_sampling.ModelSamplingFlux
|
||||
sampling_type = ReverseCONST
|
||||
|
||||
class ModelSamplingAdvanced(sampling_base, sampling_type):
|
||||
pass
|
||||
|
||||
model_sampling = ModelSamplingAdvanced(model.model.model_config)
|
||||
model_sampling.set_parameters(shift=1.15)
|
||||
m.add_object_patch("model_sampling", model_sampling)
|
||||
return (m,)
|
||||
@@ -0,0 +1,149 @@
|
||||
import math
|
||||
|
||||
import comfy.model_patcher
|
||||
import comfy.samplers
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from comfy.ldm.modules.attention import optimized_attention
|
||||
from einops import rearrange
|
||||
|
||||
DEFAULT_PAG_LTX = {"layers": set([14])}
|
||||
|
||||
|
||||
def gaussian_blur_2d(img, kernel_size, sigma):
|
||||
height = img.shape[-1]
|
||||
kernel_size = min(kernel_size, height - (height % 2 - 1))
|
||||
ksize_half = (kernel_size - 1) * 0.5
|
||||
|
||||
x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size)
|
||||
|
||||
pdf = torch.exp(-0.5 * (x / sigma).pow(2))
|
||||
|
||||
x_kernel = pdf / pdf.sum()
|
||||
x_kernel = x_kernel.to(device=img.device, dtype=img.dtype)
|
||||
|
||||
kernel2d = torch.mm(x_kernel[:, None], x_kernel[None, :])
|
||||
kernel2d = kernel2d.expand(img.shape[-3], 1, kernel2d.shape[0], kernel2d.shape[1])
|
||||
|
||||
padding = [kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2]
|
||||
|
||||
img = F.pad(img, padding, mode="reflect")
|
||||
img = F.conv2d(img, kernel2d, groups=img.shape[-3])
|
||||
|
||||
return img
|
||||
|
||||
|
||||
class LTXPerturbedAttentionNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"scale": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 2.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"rescale": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 100.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"cfg": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 3.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"attn_override": ("ATTN_OVERRIDE",),
|
||||
# "attn_type": (["PAG", "SEG"],),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch"
|
||||
|
||||
CATEGORY = "ltxtricks/attn"
|
||||
|
||||
def patch(
|
||||
self, model, scale, rescale, cfg, attn_override=DEFAULT_PAG_LTX, attn_type="PAG"
|
||||
):
|
||||
m = model.clone()
|
||||
|
||||
def pag_fn(q, k, v, heads, attn_precision=None, transformer_options=None):
|
||||
return v
|
||||
|
||||
def seg_fn(q, k, v, heads, attn_precision=None, transformer_options=None):
|
||||
_, sequence_length, _ = q.shape
|
||||
b, c, f, h, w = transformer_options["original_shape"]
|
||||
|
||||
q = rearrange(q, "b (f h w) d -> b (f d) w h", h=h, w=w)
|
||||
kernel_size = math.ceil(6 * scale) + 1 - math.ceil(6 * scale) % 2
|
||||
q = gaussian_blur_2d(q, kernel_size, scale)
|
||||
q = rearrange(q, "b (f d) w h -> b (f h w) d", f=f)
|
||||
return optimized_attention(q, k, v, heads, attn_precision=attn_precision)
|
||||
|
||||
def post_cfg_function(args):
|
||||
model = args["model"]
|
||||
|
||||
cond_pred = args["cond_denoised"]
|
||||
uncond_pred = args["uncond_denoised"]
|
||||
|
||||
len_conds = 1 if args.get("uncond", None) is None else 2
|
||||
|
||||
cond = args["cond"]
|
||||
sigma = args["sigma"]
|
||||
model_options = args["model_options"].copy()
|
||||
x = args["input"]
|
||||
|
||||
if scale == 0:
|
||||
if len_conds == 1:
|
||||
return cond_pred
|
||||
return uncond_pred + (cond_pred - uncond_pred)
|
||||
|
||||
attn_fn = pag_fn if attn_type == "PAG" else seg_fn
|
||||
for block_idx in attn_override["layers"]:
|
||||
model_options = comfy.model_patcher.set_model_options_patch_replace(
|
||||
model_options, attn_fn, "layer", "self_attn", int(block_idx)
|
||||
)
|
||||
|
||||
(perturbed,) = comfy.samplers.calc_cond_batch(
|
||||
model, [cond], x, sigma, model_options
|
||||
)
|
||||
|
||||
# if len_conds == 1:
|
||||
# output = cond_pred + scale * (cond_pred - pag)
|
||||
# else:
|
||||
# output = cond_pred + (scale-1.0) * (cond_pred - uncond_pred) + scale * (cond_pred - pag)
|
||||
|
||||
output = (
|
||||
uncond_pred
|
||||
+ cfg * (cond_pred - uncond_pred)
|
||||
+ scale * (cond_pred - perturbed)
|
||||
)
|
||||
if rescale > 0:
|
||||
factor = cond_pred.std() / output.std()
|
||||
factor = rescale * factor + (1 - rescale)
|
||||
output = output * factor
|
||||
|
||||
return output
|
||||
|
||||
m.set_model_sampler_post_cfg_function(post_cfg_function)
|
||||
|
||||
return (m,)
|
||||
@@ -0,0 +1,20 @@
|
||||
from ..modules.ltx_model import inject_model
|
||||
|
||||
|
||||
class ModifyLTXModelNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
FUNCTION = "modify"
|
||||
|
||||
def modify(self, model):
|
||||
model.model.diffusion_model = inject_model(model.model.diffusion_model)
|
||||
return (model,)
|
||||
@@ -0,0 +1,275 @@
|
||||
import torch
|
||||
from comfy.samplers import KSAMPLER
|
||||
from tqdm import trange
|
||||
|
||||
|
||||
def generate_trend_values(steps, start_time, end_time, eta, eta_trend):
|
||||
eta_values = [0] * steps
|
||||
|
||||
if eta_trend == "constant":
|
||||
for i in range(start_time, end_time):
|
||||
eta_values[i] = eta
|
||||
elif eta_trend == "linear_increase":
|
||||
for i in range(start_time, end_time):
|
||||
progress = (i - start_time) / (end_time - start_time - 1)
|
||||
eta_values[i] = eta * progress
|
||||
elif eta_trend == "linear_decrease":
|
||||
for i in range(start_time, end_time):
|
||||
progress = 1 - (i - start_time) / (end_time - start_time - 1)
|
||||
eta_values[i] = eta * progress
|
||||
|
||||
return eta_values
|
||||
|
||||
|
||||
def get_sample_forward(
|
||||
gamma, start_step, end_step, gamma_trend, seed, attn_bank=None, order="first"
|
||||
):
|
||||
# Controlled Forward ODE (Algorithm 1)
|
||||
generator = torch.Generator()
|
||||
generator.manual_seed(seed)
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_forward(model, y0, sigmas, extra_args=None, callback=None, disable=None):
|
||||
if attn_bank is not None:
|
||||
for block_idx in attn_bank["block_map"]:
|
||||
attn_bank["block_map"][block_idx].clear()
|
||||
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
model_options = extra_args.get("model_options", {})
|
||||
model_options = {**model_options}
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {
|
||||
**transformer_options,
|
||||
"total_steps": len(sigmas) - 1,
|
||||
"sample_mode": "forward",
|
||||
"attn_bank": attn_bank,
|
||||
}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
|
||||
Y = y0.clone()
|
||||
y1 = torch.randn(Y.shape, generator=generator).to(y0.device)
|
||||
N = len(sigmas) - 1
|
||||
s_in = y0.new_ones([y0.shape[0]])
|
||||
gamma_values = generate_trend_values(
|
||||
N, start_step, end_step, gamma, gamma_trend
|
||||
)
|
||||
for i in trange(N, disable=disable):
|
||||
transformer_options["step"] = i
|
||||
sigma = sigmas[i]
|
||||
sigma_next = sigmas[i + 1]
|
||||
t_i = model.inner_model.inner_model.model_sampling.timestep(sigmas[i])
|
||||
|
||||
conditional_vector_field = (y1 - Y) / (1 - t_i)
|
||||
|
||||
transformer_options["pred_order"] = "first"
|
||||
pred = model(
|
||||
Y, s_in * sigmas[i], **extra_args
|
||||
) # this implementation takes sigma instead of timestep
|
||||
|
||||
if order == "second":
|
||||
transformer_options["pred_order"] = "second"
|
||||
img_mid = Y + (sigma_next - sigma) / 2 * pred
|
||||
sigma_mid = sigma + (sigma_next - sigma) / 2
|
||||
pred_mid = model(img_mid, s_in * sigma_mid, **extra_args)
|
||||
|
||||
first_order = (pred_mid - pred) / ((sigma_next - sigma) / 2)
|
||||
pred = pred + gamma_values[i] * (conditional_vector_field - pred)
|
||||
# first_order = first_order + gamma_values[i] * (conditional_vector_field - first_order)
|
||||
Y = (
|
||||
Y
|
||||
+ (sigma_next - sigma) * pred
|
||||
+ 0.5 * (sigma_next - sigma) ** 2 * first_order
|
||||
)
|
||||
else:
|
||||
pred = pred + gamma_values[i] * (conditional_vector_field - pred)
|
||||
Y = Y + pred * (sigma_next - sigma)
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{"x": Y, "denoised": Y, "i": i, "sigma": sigma, "sigma_hat": sigma}
|
||||
)
|
||||
|
||||
return Y
|
||||
|
||||
return sample_forward
|
||||
|
||||
|
||||
def get_sample_reverse(
|
||||
latent_image, eta, start_time, end_time, eta_trend, attn_bank=None, order="first"
|
||||
):
|
||||
# Controlled Reverse ODE (Algorithm 2)
|
||||
@torch.no_grad()
|
||||
def sample_reverse(model, y1, sigmas, extra_args=None, callback=None, disable=None):
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
model_options = extra_args.get("model_options", {})
|
||||
model_options = {**model_options}
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {
|
||||
**transformer_options,
|
||||
"total_steps": len(sigmas) - 1,
|
||||
"sample_mode": "reverse",
|
||||
"attn_bank": attn_bank,
|
||||
}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
|
||||
X = y1.clone()
|
||||
N = len(sigmas) - 1
|
||||
y0 = latent_image.clone().to(y1.device)
|
||||
s_in = y0.new_ones([y0.shape[0]])
|
||||
eta_values = generate_trend_values(N, start_time, end_time, eta, eta_trend)
|
||||
for i in trange(N, disable=disable):
|
||||
transformer_options["step"] = i
|
||||
t_i = 1 - model.inner_model.inner_model.model_sampling.timestep(sigmas[i])
|
||||
sigma = sigmas[i]
|
||||
sigma_prev = sigmas[i + 1]
|
||||
|
||||
conditional_vector_field = (y0 - X) / (1 - t_i)
|
||||
|
||||
transformer_options["pred_order"] = "first"
|
||||
pred = model(
|
||||
X, sigma * s_in, **extra_args
|
||||
) # this implementation takes sigma instead of timestep
|
||||
|
||||
if order == "second":
|
||||
transformer_options["pred_order"] = "second"
|
||||
img_mid = X + (sigma_prev - sigma) / 2 * pred
|
||||
sigma_mid = sigma + (sigma_prev - sigma) / 2
|
||||
pred_mid = model(img_mid, s_in * sigma_mid, **extra_args)
|
||||
|
||||
first_order = (pred_mid - pred) / ((sigma_prev - sigma) / 2)
|
||||
pred = -pred + eta_values[i] * (conditional_vector_field + pred)
|
||||
|
||||
first_order = -first_order + eta_values[i] * (
|
||||
conditional_vector_field + first_order
|
||||
)
|
||||
X = (
|
||||
X
|
||||
+ (sigma - sigma_prev) * pred
|
||||
+ 0.5 * (sigma - sigma_prev) ** 2 * first_order
|
||||
)
|
||||
else:
|
||||
controlled_vector_field = -pred + eta_values[i] * (
|
||||
conditional_vector_field + pred
|
||||
)
|
||||
X = X + controlled_vector_field * (sigma - sigma_prev)
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{
|
||||
"x": X,
|
||||
"denoised": X,
|
||||
"i": i,
|
||||
"sigma": sigmas[i],
|
||||
"sigma_hat": sigmas[i],
|
||||
}
|
||||
)
|
||||
|
||||
return X
|
||||
|
||||
return sample_reverse
|
||||
|
||||
|
||||
class LTXRFForwardODESamplerNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"gamma": (
|
||||
"FLOAT",
|
||||
{"default": 0.5, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"start_step": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"end_step": ("INT", {"default": 5, "min": 0, "max": 1000, "step": 1}),
|
||||
"gamma_trend": (["linear_decrease", "linear_increase", "constant"],),
|
||||
},
|
||||
"optional": {
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
||||
"attn_bank": ("ATTN_BANK",),
|
||||
"order": (["first", "second"],),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLER",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(
|
||||
self,
|
||||
gamma,
|
||||
start_step,
|
||||
end_step,
|
||||
gamma_trend,
|
||||
seed=0,
|
||||
attn_bank=None,
|
||||
order="first",
|
||||
):
|
||||
sampler = KSAMPLER(
|
||||
get_sample_forward(
|
||||
gamma,
|
||||
start_step,
|
||||
end_step,
|
||||
gamma_trend,
|
||||
seed,
|
||||
attn_bank=attn_bank,
|
||||
order=order,
|
||||
)
|
||||
)
|
||||
|
||||
return (sampler,)
|
||||
|
||||
|
||||
class LTXRFReverseODESamplerNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"latent_image": ("LATENT",),
|
||||
"eta": (
|
||||
"FLOAT",
|
||||
{"default": 0.8, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"start_step": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"end_step": ("INT", {"default": 15, "min": 0, "max": 1000, "step": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"eta_trend": (["linear_decrease", "linear_increase", "constant"],),
|
||||
"attn_inj": ("ATTN_INJ",),
|
||||
"order": (["first", "second"],),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLER",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(
|
||||
self,
|
||||
model,
|
||||
latent_image,
|
||||
eta,
|
||||
start_step,
|
||||
end_step,
|
||||
eta_trend="constant",
|
||||
attn_inj=None,
|
||||
order="first",
|
||||
):
|
||||
process_latent_in = model.get_model_object("process_latent_in")
|
||||
latent_image = process_latent_in(latent_image["samples"])
|
||||
sampler = KSAMPLER(
|
||||
get_sample_reverse(
|
||||
latent_image,
|
||||
eta,
|
||||
start_step,
|
||||
end_step,
|
||||
eta_trend,
|
||||
attn_bank=attn_inj,
|
||||
order=order,
|
||||
)
|
||||
)
|
||||
|
||||
return (sampler,)
|
||||
@@ -0,0 +1,376 @@
|
||||
import torch
|
||||
from comfy.samplers import KSAMPLER
|
||||
from tqdm import trange
|
||||
|
||||
|
||||
def get_sample_forward(attn_bank, save_steps, single_layers, double_layers):
|
||||
@torch.no_grad()
|
||||
def sample_forward(model, x, sigmas, extra_args=None, callback=None, disable=None):
|
||||
attn_bank.clear()
|
||||
attn_bank["save_steps"] = save_steps
|
||||
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
|
||||
model_options = extra_args.get("model_options", {})
|
||||
model_options = {**model_options}
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {**transformer_options}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
|
||||
N = len(sigmas) - 1
|
||||
s_in = x.new_ones([x.shape[0]])
|
||||
for i in trange(N, disable=disable):
|
||||
sigma = sigmas[i]
|
||||
sigma_next = sigmas[i + 1]
|
||||
|
||||
if N - i - 1 < save_steps:
|
||||
attn_bank[N - i - 1] = {"first": {}, "mid": {}}
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": N - i - 1,
|
||||
"process": "forward" if N - i - 1 < save_steps else None,
|
||||
"pred": "first",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
pred = model(x, s_in * sigma, **extra_args)
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": N - i - 1,
|
||||
"process": "forward" if N - i - 1 < save_steps else None,
|
||||
"pred": "mid",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
img_mid = x + (sigma_next - sigma) / 2 * pred
|
||||
sigma_mid = sigma + (sigma_next - sigma) / 2
|
||||
pred_mid = model(img_mid, s_in * sigma_mid, **extra_args)
|
||||
|
||||
first_order = (pred_mid - pred) / ((sigma_next - sigma) / 2)
|
||||
x = (
|
||||
x
|
||||
+ (sigma_next - sigma) * pred
|
||||
+ 0.5 * (sigma_next - sigma) ** 2 * first_order
|
||||
)
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{
|
||||
"x": x,
|
||||
"denoised": x,
|
||||
"i": i,
|
||||
"sigma": sigmas[i],
|
||||
"sigma_hat": sigmas[i],
|
||||
}
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
return sample_forward
|
||||
|
||||
|
||||
def get_sample_reverse(attn_bank, inject_steps, single_layers, double_layers):
|
||||
@torch.no_grad()
|
||||
def sample_reverse(model, x, sigmas, extra_args=None, callback=None, disable=None):
|
||||
if inject_steps > attn_bank["save_steps"]:
|
||||
raise ValueError(
|
||||
f'You must save at least as many steps as you want to inject. save_steps: {attn_bank["save_steps"]}, inject_steps: {inject_steps}'
|
||||
)
|
||||
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
|
||||
model_options = extra_args.get("model_options", {})
|
||||
model_options = {**model_options}
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {**transformer_options}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
|
||||
N = len(sigmas) - 1
|
||||
s_in = x.new_ones([x.shape[0]])
|
||||
for i in trange(N, disable=disable):
|
||||
sigma = sigmas[i]
|
||||
sigma_prev = sigmas[i + 1]
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": i,
|
||||
"process": "reverse" if i < inject_steps else None,
|
||||
"pred": "first",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
pred = model(x, s_in * sigma, **extra_args)
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": i,
|
||||
"process": "reverse" if i < inject_steps else None,
|
||||
"pred": "mid",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
img_mid = x + (sigma_prev - sigma) / 2 * pred
|
||||
sigma_mid = sigma + (sigma_prev - sigma) / 2
|
||||
pred_mid = model(img_mid, s_in * sigma_mid, **extra_args)
|
||||
|
||||
first_order = (pred_mid - pred) / ((sigma_prev - sigma) / 2)
|
||||
x = (
|
||||
x
|
||||
+ (sigma_prev - sigma) * pred
|
||||
+ 0.5 * (sigma_prev - sigma) ** 2 * first_order
|
||||
)
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{
|
||||
"x": x,
|
||||
"denoised": x,
|
||||
"i": i,
|
||||
"sigma": sigmas[i],
|
||||
"sigma_hat": sigmas[i],
|
||||
}
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
return sample_reverse
|
||||
|
||||
|
||||
DEFAULT_SINGLE_LAYERS = {}
|
||||
for i in range(38):
|
||||
DEFAULT_SINGLE_LAYERS[f"{i}"] = i > 19
|
||||
|
||||
DEFAULT_DOUBLE_LAYERS = {}
|
||||
for i in range(19):
|
||||
DEFAULT_DOUBLE_LAYERS[f"{i}"] = False
|
||||
|
||||
|
||||
class FlowEditForwardSamplerNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"save_steps": (
|
||||
"INT",
|
||||
{"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"single_layers": ("SINGLE_LAYERS",),
|
||||
"double_layers": ("DOUBLE_LAYERS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLER", "ATTN_INJ")
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "fluxtapoz"
|
||||
|
||||
def build(
|
||||
self,
|
||||
save_steps,
|
||||
single_layers=DEFAULT_SINGLE_LAYERS,
|
||||
double_layers=DEFAULT_DOUBLE_LAYERS,
|
||||
):
|
||||
attn_bank = {}
|
||||
sampler = KSAMPLER(
|
||||
get_sample_forward(attn_bank, save_steps, single_layers, double_layers)
|
||||
)
|
||||
|
||||
return (sampler, attn_bank)
|
||||
|
||||
|
||||
# class FlowEditReverseSamplerNode:
|
||||
# @classmethod
|
||||
# def INPUT_TYPES(s):
|
||||
# return {
|
||||
# "required": {
|
||||
# "attn_inj": ("ATTN_INJ",),
|
||||
# "latent_image": ("LATENT",),
|
||||
# "eta": (
|
||||
# "FLOAT",
|
||||
# {"default": 0.8, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
# ),
|
||||
# "start_step": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
# "end_step": ("INT", {"default": 5, "min": 0, "max": 1000, "step": 1}),
|
||||
# },
|
||||
# "optional": {},
|
||||
# }
|
||||
#
|
||||
# RETURN_TYPES = ("SAMPLER",)
|
||||
# FUNCTION = "build"
|
||||
#
|
||||
# CATEGORY = "fluxtapoz"
|
||||
#
|
||||
# def build(self, latent_image, eta, start_step, end_step):
|
||||
# sampler = KSAMPLER(
|
||||
# get_sample_reverse(attn_inj, inject_steps, single_layers, double_layers)
|
||||
# )
|
||||
# return (sampler,)
|
||||
|
||||
|
||||
def get_sample_reverse2(attn_bank, inject_steps, single_layers, double_layers):
|
||||
@torch.no_grad()
|
||||
def sample_reverse(model, x, sigmas, extra_args=None, callback=None, disable=None):
|
||||
if inject_steps > attn_bank["save_steps"]:
|
||||
raise ValueError(
|
||||
f'You must save at least as many steps as you want to inject. save_steps: {attn_bank["save_steps"]}, inject_steps: {inject_steps}'
|
||||
)
|
||||
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
|
||||
model_options = extra_args.get("model_options", {})
|
||||
model_options = {**model_options}
|
||||
transformer_options = model_options.get("transformer_options", {})
|
||||
transformer_options = {**transformer_options}
|
||||
model_options["transformer_options"] = transformer_options
|
||||
extra_args["model_options"] = model_options
|
||||
|
||||
N = len(sigmas) - 1
|
||||
s_in = x.new_ones([x.shape[0]])
|
||||
for i in trange(N, disable=disable):
|
||||
sigma = sigmas[i]
|
||||
sigma_prev = sigmas[i + 1]
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": i,
|
||||
"process": "reverse" if i < inject_steps else None,
|
||||
"pred": "first",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
pred = model(x, s_in * sigma, **extra_args)
|
||||
|
||||
transformer_options["rfedit"] = {
|
||||
"step": i,
|
||||
"process": "reverse" if i < inject_steps else None,
|
||||
"pred": "mid",
|
||||
"bank": attn_bank,
|
||||
"single_layers": single_layers,
|
||||
"double_layers": double_layers,
|
||||
}
|
||||
|
||||
img_mid = x + (sigma_prev - sigma) / 2 * pred
|
||||
sigma_mid = sigma + (sigma_prev - sigma) / 2
|
||||
pred_mid = model(img_mid, s_in * sigma_mid, **extra_args)
|
||||
|
||||
first_order = (pred_mid - pred) / ((sigma_prev - sigma) / 2)
|
||||
x = (
|
||||
x
|
||||
+ (sigma_prev - sigma) * pred
|
||||
+ 0.5 * (sigma_prev - sigma) ** 2 * first_order
|
||||
)
|
||||
|
||||
if callback is not None:
|
||||
callback(
|
||||
{
|
||||
"x": x,
|
||||
"denoised": x,
|
||||
"i": i,
|
||||
"sigma": sigmas[i],
|
||||
"sigma_hat": sigmas[i],
|
||||
}
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
return sample_reverse
|
||||
|
||||
|
||||
class FlowEdit2ReverseSamplerNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"attn_inj": ("ATTN_INJ",),
|
||||
"inject_steps": (
|
||||
"INT",
|
||||
{"default": 0, "min": 0, "max": 1000, "step": 1},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"single_layers": ("SINGLE_LAYERS",),
|
||||
"double_layers": ("DOUBLE_LAYERS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLER",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(
|
||||
self,
|
||||
attn_inj,
|
||||
inject_steps,
|
||||
single_layers=DEFAULT_SINGLE_LAYERS,
|
||||
double_layers=DEFAULT_DOUBLE_LAYERS,
|
||||
):
|
||||
sampler = KSAMPLER(
|
||||
get_sample_reverse(attn_inj, inject_steps, single_layers, double_layers)
|
||||
)
|
||||
return (sampler,)
|
||||
|
||||
|
||||
class PrepareAttnBankNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latent": ("LATENT",),
|
||||
"attn_inj": ("ATTN_INJ",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT", "ATTN_INJ")
|
||||
FUNCTION = "prepare"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def prepare(self, latent, attn_inj):
|
||||
# Hack to force order of operations in ComfyUI graph
|
||||
return (latent, attn_inj)
|
||||
|
||||
|
||||
class RFSingleBlocksOverrideNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
layers = {}
|
||||
for i in range(38):
|
||||
layers[f"{i}"] = ("BOOLEAN", {"default": i > 19})
|
||||
return {"required": layers}
|
||||
|
||||
RETURN_TYPES = ("SINGLE_LAYERS",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(self, *args, **kwargs):
|
||||
return (kwargs,)
|
||||
|
||||
|
||||
class RFDoubleBlocksOverrideNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
layers = {}
|
||||
for i in range(19):
|
||||
layers[f"{i}"] = ("BOOLEAN", {"default": False})
|
||||
return {"required": layers}
|
||||
|
||||
RETURN_TYPES = ("DOUBLE_LAYERS",)
|
||||
FUNCTION = "build"
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
|
||||
def build(self, *args, **kwargs):
|
||||
return (kwargs,)
|
||||
@@ -0,0 +1,16 @@
|
||||
class AttentionBank:
|
||||
def __init__(self, save_steps, block_map, inject_steps=None):
|
||||
self._data = {
|
||||
"save_steps": save_steps,
|
||||
"block_map": block_map,
|
||||
"inject_steps": inject_steps,
|
||||
}
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self._data[key]
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
self._data[key] = value
|
||||
|
||||
def get(self, key, default=None):
|
||||
return self._data.get(key, default)
|
||||
@@ -0,0 +1,55 @@
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
def _feta_score(query_image, key_image, head_dim, num_frames, enhance_weight):
|
||||
scale = head_dim**-0.5
|
||||
query_image = query_image * scale
|
||||
attn_temp = query_image @ key_image.transpose(-2, -1) # translate attn to float32
|
||||
attn_temp = attn_temp.to(torch.float32)
|
||||
attn_temp = attn_temp.softmax(dim=-1)
|
||||
|
||||
# Reshape to [batch_size * num_tokens, num_frames, num_frames]
|
||||
attn_temp = attn_temp.reshape(-1, num_frames, num_frames)
|
||||
|
||||
# Create a mask for diagonal elements
|
||||
diag_mask = torch.eye(num_frames, device=attn_temp.device).bool()
|
||||
diag_mask = diag_mask.unsqueeze(0).expand(attn_temp.shape[0], -1, -1)
|
||||
|
||||
# Zero out diagonal elements
|
||||
attn_wo_diag = attn_temp.masked_fill(diag_mask, 0)
|
||||
|
||||
# Calculate mean for each token's attention matrix
|
||||
# Number of off-diagonal elements per matrix is n*n - n
|
||||
num_off_diag = num_frames * num_frames - num_frames
|
||||
mean_scores = attn_wo_diag.sum(dim=(1, 2)) / num_off_diag
|
||||
|
||||
enhance_scores = mean_scores.mean() * (num_frames + enhance_weight)
|
||||
enhance_scores = enhance_scores.clamp(min=1)
|
||||
return enhance_scores
|
||||
|
||||
|
||||
def get_feta_scores(img_q, img_k, num_heads, transformer_options):
|
||||
num_frames = transformer_options["original_shape"][2]
|
||||
_, ST, dim = img_q.shape
|
||||
head_dim = dim // num_heads
|
||||
spatial_dim = ST // num_frames
|
||||
|
||||
query_image = rearrange(
|
||||
img_q,
|
||||
"B (T S) (N C) -> (B S) N T C",
|
||||
T=num_frames,
|
||||
S=spatial_dim,
|
||||
N=num_heads,
|
||||
C=head_dim,
|
||||
)
|
||||
key_image = rearrange(
|
||||
img_k,
|
||||
"B (T S) (N C) -> (B S) N T C",
|
||||
T=num_frames,
|
||||
S=spatial_dim,
|
||||
N=num_heads,
|
||||
C=head_dim,
|
||||
)
|
||||
weight = transformer_options.get("feta_weight", 0)
|
||||
return _feta_score(query_image, key_image, head_dim, num_frames, weight)
|
||||
@@ -0,0 +1,8 @@
|
||||
import torch
|
||||
|
||||
|
||||
class LatentGuide(torch.nn.Module):
|
||||
def __init__(self, latent: torch.Tensor, index) -> None:
|
||||
super().__init__()
|
||||
self.index = index
|
||||
self.register_buffer("latent", latent)
|
||||
@@ -0,0 +1,6 @@
|
||||
def isinstance_str(x: object, cls_name: str):
|
||||
for _cls in x.__class__.__mro__:
|
||||
if _cls.__name__ == cls_name:
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -0,0 +1,23 @@
|
||||
def get_alphacumprod(sigma):
|
||||
return 1 / ((sigma * sigma) + 1)
|
||||
|
||||
|
||||
def add_noise(src_latent, noise, sigma):
|
||||
alphas_cumprod = get_alphacumprod(sigma)
|
||||
|
||||
sqrt_alpha_prod = alphas_cumprod**0.5
|
||||
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
|
||||
while len(sqrt_alpha_prod.shape) < len(src_latent.shape):
|
||||
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
|
||||
|
||||
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod) ** 0.5
|
||||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
|
||||
while len(sqrt_one_minus_alpha_prod.shape) < len(src_latent.shape):
|
||||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
|
||||
|
||||
noisy_samples = sqrt_alpha_prod * src_latent + sqrt_one_minus_alpha_prod * noise
|
||||
return noisy_samples
|
||||
|
||||
|
||||
def add_noise_flux(src_latent, noise, sigma):
|
||||
return sigma * noise + (1.0 - sigma) * src_latent
|
||||
@@ -1,167 +0,0 @@
|
||||
from copy import copy
|
||||
|
||||
import comfy.latent_formats
|
||||
import comfy.model_base
|
||||
import comfy.model_management
|
||||
import comfy.model_patcher
|
||||
import comfy.sd
|
||||
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 .nodes_registry import comfy_node
|
||||
|
||||
|
||||
class LTXVVAE(comfy.sd.VAE):
|
||||
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):
|
||||
instance = cls()
|
||||
model = vae_class.from_pretrained(
|
||||
pretrained_model_name_or_path=model_path,
|
||||
revision=None,
|
||||
torch_dtype=dtype,
|
||||
load_in_8bit=False,
|
||||
)
|
||||
instance._finalize_model(model)
|
||||
return instance
|
||||
|
||||
@classmethod
|
||||
def from_config_and_state_dict(
|
||||
cls, vae_class, config, state_dict, dtype=torch.bfloat16
|
||||
):
|
||||
instance = cls()
|
||||
model = vae_class.from_config(config)
|
||||
model.load_state_dict(state_dict)
|
||||
model.to(dtype)
|
||||
instance._finalize_model(model)
|
||||
return instance
|
||||
|
||||
def _finalize_model(self, model):
|
||||
self.video_scale_factor, self.vae_scale_factor, _ = get_vae_size_scale_factor(
|
||||
model
|
||||
)
|
||||
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
||||
self.first_stage_model = model.eval().to(self.device)
|
||||
|
||||
# Assumes that the input samples have dimensions in following order
|
||||
# (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)
|
||||
|
||||
@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.
|
||||
def encode(self, pixel_samples):
|
||||
preprocessed = self.image_processor.preprocess(
|
||||
pixel_samples.permute(3, 0, 1, 2)
|
||||
)
|
||||
input = preprocessed.unsqueeze(0).to(torch.bfloat16).to(self.device)
|
||||
latents = vae_encode(
|
||||
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,)
|
||||
Reference in New Issue
Block a user