Merge pull request #577 from Lightricks/pr-2026-10-13
Update ComfyUI-LTXVideo - 2026-10-13
This commit is contained in:
DoronGenzelHass
+3
-27
@@ -1,6 +1,5 @@
|
||||
from .audio_only import LTXVAudioOnlyEmptyVideoLatent, LTXVAudioOnlyModel
|
||||
from .conditioning_loader import LTXVLoadConditioning
|
||||
from .conditioning_saver import LTXVSaveConditioning
|
||||
from .decoder_noise import DecoderNoise
|
||||
from .dynamic_conditioning import DynamicConditioning
|
||||
from .easy_samplers import (
|
||||
@@ -12,7 +11,7 @@ from .easy_samplers import (
|
||||
)
|
||||
from .gemma_api_conditioning import GemmaAPITextEncode
|
||||
from .gemma_encoder import LTXVGemmaCLIPModelLoader, LTXVGemmaEnhancePrompt
|
||||
from .guide import LTXVAddGuideAdvanced, LTXVAddGuideAdvancedAttention
|
||||
from .guide import LTXVImgToVideoAdvanced
|
||||
from .guiders import GuiderParametersNode, MultimodalGuiderNode
|
||||
from .hdr_nodes import (
|
||||
LTXVHDRDecodePostprocess,
|
||||
@@ -27,12 +26,7 @@ from .iclora import (
|
||||
LTXICLoRALoaderModelOnly,
|
||||
LTXVSetAudioRefTokens,
|
||||
)
|
||||
from .latent_norm import (
|
||||
LTXVAdainLatent,
|
||||
LTXVPerStepAdainPatcher,
|
||||
LTXVPerStepStatNormPatcher,
|
||||
LTXVStatNormLatent,
|
||||
)
|
||||
from .latent_norm import LTXVAdainLatent
|
||||
from .latents import (
|
||||
LTXVImgToVideoConditionOnly,
|
||||
LTXVSelectLatents,
|
||||
@@ -44,13 +38,11 @@ from .low_vram_loaders import (
|
||||
LowVRAMCheckpointLoader,
|
||||
LowVRAMLatentUpscaleModelLoader,
|
||||
)
|
||||
from .masks import LTXVPreprocessMasks
|
||||
from .nodes_registry import NODE_CLASS_MAPPINGS as RUNTIME_NODE_CLASS_MAPPINGS
|
||||
from .nodes_registry import (
|
||||
NODE_DISPLAY_NAME_MAPPINGS as RUNTIME_NODE_DISPLAY_NAME_MAPPINGS,
|
||||
)
|
||||
from .nodes_registry import NODES_DISPLAY_NAME_PREFIX, camel_case_to_spaces
|
||||
from .prompt_enhancer_nodes import LTXVPromptEnhancer, LTXVPromptEnhancerLoader
|
||||
from .pyramid_blending import LTXVLaplacianPyramidBlend
|
||||
from .q8_nodes import LTXVQ8LoraModelLoader, LTXVQ8Patch
|
||||
from .sparse_tracks import LTXVDrawTracks, LTXVSparseTrackEditor
|
||||
@@ -64,30 +56,21 @@ from .tiled_fusion_sampler import LTXVTiledFusionSampler
|
||||
from .tiled_fusion_sizes import LTXVGetTilingSizes
|
||||
from .tiled_sampler import LTXVTiledSampler
|
||||
from .tiled_vae_decode import LTXVTiledVAEDecode
|
||||
from .tricks import NODE_CLASS_MAPPINGS as TRICKS_NODE_CLASS_MAPPINGS
|
||||
from .tricks import NODE_DISPLAY_NAME_MAPPINGS as TRICKS_NODE_DISPLAY_NAME_MAPPINGS
|
||||
from .utiltily_nodes import FloatToInt, ImageToCPU
|
||||
from .vae_patcher import LTXVPatcherVAE
|
||||
from .vanish_nodes import LTXVDilateVideoMask, LTXVInpaintPreprocess
|
||||
|
||||
# Static node mappings, required for ComfyUI-Manager mapping to work
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Set VAE Decoder Noise": DecoderNoise,
|
||||
"LTXVLinearOverlapLatentTransition": LinearOverlapLatentTransition,
|
||||
"LTXVAddGuideAdvanced": LTXVAddGuideAdvanced,
|
||||
"LTXVAddGuideAdvancedAttention": LTXVAddGuideAdvancedAttention,
|
||||
"LTXVImgToVideoAdvanced": LTXVImgToVideoAdvanced,
|
||||
"LTXVAdainLatent": LTXVAdainLatent,
|
||||
"LTXVImgToVideoConditionOnly": LTXVImgToVideoConditionOnly,
|
||||
"LTXVPerStepAdainPatcher": LTXVPerStepAdainPatcher,
|
||||
"LTXVApplySTG": LTXVApplySTG,
|
||||
"LTXVBaseSampler": LTXVBaseSampler,
|
||||
"LTXVInContextSampler": LTXVInContextSampler,
|
||||
"LTXVExtendSampler": LTXVExtendSampler,
|
||||
"LTXVNormalizingSampler": LTXVNormalizingSampler,
|
||||
"LTXVPreprocessMasks": LTXVPreprocessMasks,
|
||||
"LTXVPatcherVAE": LTXVPatcherVAE,
|
||||
"LTXVPromptEnhancer": LTXVPromptEnhancer,
|
||||
"LTXVPromptEnhancerLoader": LTXVPromptEnhancerLoader,
|
||||
"LTXQ8Patch": LTXVQ8Patch,
|
||||
"LTXVQ8LoraModelLoader": LTXVQ8LoraModelLoader,
|
||||
"LTXVSelectLatents": LTXVSelectLatents,
|
||||
@@ -105,8 +88,6 @@ NODE_CLASS_MAPPINGS = {
|
||||
"LTXVMultiPromptProvider": MultiPromptProvider,
|
||||
"ImageToCPU": ImageToCPU,
|
||||
"LTXFloatToInt": FloatToInt,
|
||||
"LTXVStatNormLatent": LTXVStatNormLatent,
|
||||
"LTXVPerStepStatNormPatcher": LTXVPerStepStatNormPatcher,
|
||||
"LTXVGemmaCLIPModelLoader": LTXVGemmaCLIPModelLoader,
|
||||
"LTXVGemmaEnhancePrompt": LTXVGemmaEnhancePrompt,
|
||||
"GemmaAPITextEncode": GemmaAPITextEncode,
|
||||
@@ -119,7 +100,6 @@ NODE_CLASS_MAPPINGS = {
|
||||
"LTXICLoRALoaderModelOnly": LTXICLoRALoaderModelOnly,
|
||||
"LTXVSetAudioRefTokens": LTXVSetAudioRefTokens,
|
||||
"LTXVLoadConditioning": LTXVLoadConditioning,
|
||||
"LTXVSaveConditioning": LTXVSaveConditioning,
|
||||
"LTXVDrawTracks": LTXVDrawTracks,
|
||||
"LTXVSparseTrackEditor": LTXVSparseTrackEditor,
|
||||
"LTXVDilateVideoMask": LTXVDilateVideoMask,
|
||||
@@ -141,10 +121,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
for name in NODE_CLASS_MAPPINGS.keys()
|
||||
}
|
||||
|
||||
# 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)
|
||||
|
||||
# Update with runtime mappings (these will override static mappings if there are any differences)
|
||||
NODE_CLASS_MAPPINGS.update(RUNTIME_NODE_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(RUNTIME_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
|
||||
@@ -1,69 +0,0 @@
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
import torch
|
||||
from comfy_api.latest import io, ui
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVSaveConditioning")
|
||||
class LTXVSaveConditioning(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
node_id="LTXVSaveConditioning",
|
||||
display_name="🅛🅣🅧 LTXV Save Conditioning",
|
||||
category="lightricks/LTXV",
|
||||
inputs=[
|
||||
io.Conditioning.Input("conditioning"),
|
||||
io.String.Input("filename", default="conditioning"),
|
||||
io.Combo.Input("dtype", options=["bfloat16", "float16"]),
|
||||
],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, conditioning: list, filename: str, dtype: str) -> io.NodeOutput:
|
||||
if not conditioning or len(conditioning) == 0:
|
||||
raise ValueError("Conditioning is empty")
|
||||
|
||||
embeddings_folder = Path(folder_paths.get_folder_paths("embeddings")[0])
|
||||
embeddings_folder.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
sanitized_filename = "".join(
|
||||
c for c in filename if c.isalnum() or c in ("_", "-", ".")
|
||||
)
|
||||
if not sanitized_filename:
|
||||
sanitized_filename = "conditioning"
|
||||
|
||||
output_path = embeddings_folder / f"{sanitized_filename}.safetensors"
|
||||
|
||||
target_dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float16
|
||||
|
||||
tensors_to_save: dict[str, torch.Tensor] = {}
|
||||
|
||||
for idx, (cond_tensor, cond_options) in enumerate(conditioning):
|
||||
tensor_converted = cond_tensor.to(dtype=target_dtype).contiguous()
|
||||
tensors_to_save[f"conditioning_data_{idx}"] = tensor_converted
|
||||
|
||||
if "attention_mask" in cond_options:
|
||||
mask = cond_options["attention_mask"].contiguous()
|
||||
tensors_to_save[f"attention_mask_{idx}"] = mask
|
||||
|
||||
metadata = {
|
||||
"num_conditionings": str(len(conditioning)),
|
||||
"dtype": dtype,
|
||||
"created_at": str(datetime.now()),
|
||||
}
|
||||
|
||||
comfy.utils.save_torch_file(
|
||||
tensors_to_save, str(output_path), metadata=metadata
|
||||
)
|
||||
|
||||
file_size_mb = output_path.stat().st_size / (1024 * 1024)
|
||||
return io.NodeOutput(
|
||||
ui=ui.PreviewText(f"Saved: {output_path.name} ({file_size_mb:.2f} MB)")
|
||||
)
|
||||
+1
-1
@@ -1570,7 +1570,7 @@
|
||||
"ver": "0.3.56"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"widgets_values_named": {
|
||||
|
||||
@@ -92,7 +92,7 @@
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
|
||||
@@ -263,7 +263,7 @@
|
||||
"ver": "0.3.56"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -502,7 +502,7 @@
|
||||
"ver": "0.3.56"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -823,7 +823,7 @@
|
||||
"ver": "0.3.56"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
|
||||
@@ -837,7 +837,7 @@
|
||||
"ver": "0.3.56"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
|
||||
@@ -536,7 +536,7 @@
|
||||
"ver": "0.3.56"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"example.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#232",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -17,126 +17,6 @@ def blur_internal(image, blur_radius):
|
||||
return image
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAddGuideAdvanced")
|
||||
class LTXVAddGuideAdvanced:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"vae": ("VAE",),
|
||||
"latent": ("LATENT",),
|
||||
"image": ("IMAGE",),
|
||||
"frame_idx": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": -9999,
|
||||
"max": 9999,
|
||||
"tooltip": "Frame index to start the conditioning at. For single-frame images or "
|
||||
"videos with 1-8 frames, any frame_idx value is acceptable. For videos with 9+ "
|
||||
"frames, frame_idx must be divisible by 8, otherwise it will be rounded down to "
|
||||
"the nearest multiple of 8. Negative values are counted from the end of the video.",
|
||||
},
|
||||
),
|
||||
"strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"tooltip": "Strength of the conditioning. Higher values will make the conditioning more exact.",
|
||||
},
|
||||
),
|
||||
"crf": (
|
||||
"INT",
|
||||
{
|
||||
"default": 29,
|
||||
"min": 0,
|
||||
"max": 51,
|
||||
"step": 1,
|
||||
"tooltip": "CRF value for the video. Higher values mean more motion, lower values mean higher quality.",
|
||||
},
|
||||
),
|
||||
"blur_radius": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 7,
|
||||
"step": 1,
|
||||
"tooltip": "Blur kernel radius size. Higher values mean more motion, lower values mean higher quality.",
|
||||
},
|
||||
),
|
||||
"interpolation": (
|
||||
[
|
||||
"lanczos",
|
||||
"bislerp",
|
||||
"nearest",
|
||||
"bilinear",
|
||||
"bicubic",
|
||||
"area",
|
||||
"nearest-exact",
|
||||
],
|
||||
{"default": "lanczos"},
|
||||
),
|
||||
"crop": (["center", "disabled"], {"default": "disabled"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
||||
RETURN_NAMES = ("positive", "negative", "latent")
|
||||
|
||||
CATEGORY = "conditioning/video_models"
|
||||
FUNCTION = "generate"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Adds a conditioning frame or a video at a specific frame index. "
|
||||
"This node is used to add a keyframe or a video segment which should appear in the "
|
||||
"generated video at a specified index. It resizes the image to the correct size and "
|
||||
"applies preprocessing to it."
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent,
|
||||
image,
|
||||
frame_idx,
|
||||
strength,
|
||||
crf,
|
||||
blur_radius,
|
||||
interpolation,
|
||||
crop,
|
||||
):
|
||||
_, width_scale_factor, height_scale_factor = vae.downscale_index_formula
|
||||
width, height = (
|
||||
latent["samples"].shape[4] * width_scale_factor,
|
||||
latent["samples"].shape[3] * height_scale_factor,
|
||||
)
|
||||
image = (
|
||||
comfy.utils.common_upscale(
|
||||
image.movedim(-1, 1), width, height, interpolation, crop=crop
|
||||
)
|
||||
.movedim(1, -1)
|
||||
.clamp(0, 1)
|
||||
)
|
||||
image = nodes_lt.LTXVPreprocess().execute(image, crf)[0]
|
||||
image = blur_internal(image, blur_radius)
|
||||
return nodes_lt.LTXVAddGuide().execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
vae=vae,
|
||||
latent=latent,
|
||||
image=image,
|
||||
frame_idx=frame_idx,
|
||||
strength=strength,
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVImgToVideoAdvanced")
|
||||
class LTXVImgToVideoAdvanced:
|
||||
@classmethod
|
||||
@@ -252,205 +132,3 @@ class LTXVImgToVideoAdvanced:
|
||||
batch_size=batch_size,
|
||||
strength=strength,
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAddGuideAdvancedAttention")
|
||||
class LTXVAddGuideAdvancedAttention:
|
||||
"""Extended keyframe guide node with per-guide attention strength control.
|
||||
|
||||
Same preprocessing as LTXVAddGuideAdvanced (CRF, blur, interpolation, crop),
|
||||
plus attention_strength and attention_mask inputs to control how strongly
|
||||
this guide's conditioning influences generation via self-attention.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"vae": ("VAE",),
|
||||
"latent": ("LATENT",),
|
||||
"image": ("IMAGE",),
|
||||
"frame_idx": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": -9999,
|
||||
"max": 9999,
|
||||
"tooltip": (
|
||||
"Frame index to start the conditioning at. "
|
||||
"Negative values are counted from the end of the video."
|
||||
),
|
||||
},
|
||||
),
|
||||
"strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"tooltip": "Strength of the conditioning. Higher values make it more exact.",
|
||||
},
|
||||
),
|
||||
"crf": (
|
||||
"INT",
|
||||
{
|
||||
"default": 29,
|
||||
"min": 0,
|
||||
"max": 51,
|
||||
"step": 1,
|
||||
"tooltip": "CRF value. Higher = more motion, lower = higher quality.",
|
||||
},
|
||||
),
|
||||
"blur_radius": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 7,
|
||||
"step": 1,
|
||||
"tooltip": "Blur kernel radius. Higher = more motion.",
|
||||
},
|
||||
),
|
||||
"interpolation": (
|
||||
[
|
||||
"lanczos",
|
||||
"bislerp",
|
||||
"nearest",
|
||||
"bilinear",
|
||||
"bicubic",
|
||||
"area",
|
||||
"nearest-exact",
|
||||
],
|
||||
{"default": "lanczos"},
|
||||
),
|
||||
"crop": (["center", "disabled"], {"default": "disabled"}),
|
||||
"attention_strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": (
|
||||
"Controls how strongly this guide influences generation via "
|
||||
"self-attention. 1.0 = full conditioning, 0.0 = ignore."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"attention_mask": (
|
||||
"MASK",
|
||||
{
|
||||
"tooltip": (
|
||||
"Optional pixel-space spatial mask. Shape (F, H, W) or (H, W). "
|
||||
"Values in [0, 1]. Controls per-region conditioning influence. "
|
||||
"Multiplied by attention_strength."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
||||
RETURN_NAMES = ("positive", "negative", "latent")
|
||||
|
||||
CATEGORY = "conditioning/video_models"
|
||||
FUNCTION = "generate"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Adds a conditioning frame/video at a specific frame index with per-guide "
|
||||
"attention strength control. Same preprocessing as LTXVAddGuideAdvanced, "
|
||||
"plus attention_strength and optional spatial attention_mask."
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent,
|
||||
image,
|
||||
frame_idx,
|
||||
strength,
|
||||
crf,
|
||||
blur_radius,
|
||||
interpolation,
|
||||
crop,
|
||||
attention_strength=1.0,
|
||||
attention_mask=None,
|
||||
):
|
||||
from .iclora_attention import append_guide_attention_entry, normalize_mask
|
||||
|
||||
# Preprocessing: resize, CRF, blur (same as LTXVAddGuideAdvanced)
|
||||
scale_factors = vae.downscale_index_formula
|
||||
_, width_scale_factor, height_scale_factor = scale_factors
|
||||
latent_image = latent["samples"]
|
||||
noise_mask = nodes_lt.get_noise_mask(latent)
|
||||
_, _, latent_length, latent_height, latent_width = latent_image.shape
|
||||
|
||||
width = latent_width * width_scale_factor
|
||||
height = latent_height * height_scale_factor
|
||||
image = (
|
||||
comfy.utils.common_upscale(
|
||||
image.movedim(-1, 1), width, height, interpolation, crop=crop
|
||||
)
|
||||
.movedim(1, -1)
|
||||
.clamp(0, 1)
|
||||
)
|
||||
image = nodes_lt.LTXVPreprocess().execute(image, crf)[0]
|
||||
image = blur_internal(image, blur_radius)
|
||||
|
||||
# Encode
|
||||
_, t = nodes_lt.LTXVAddGuide.encode(
|
||||
vae, latent_width, latent_height, image, scale_factors
|
||||
)
|
||||
|
||||
# Compute latent index
|
||||
frame_idx, latent_idx = nodes_lt.LTXVAddGuide.get_latent_index(
|
||||
positive, latent_length, len(image), frame_idx, scale_factors
|
||||
)
|
||||
assert (
|
||||
latent_idx + t.shape[2] <= latent_length
|
||||
), "Conditioning frames exceed the length of the latent sequence."
|
||||
|
||||
# Append keyframe
|
||||
positive, negative, latent_image, noise_mask = (
|
||||
nodes_lt.LTXVAddGuide.append_keyframe(
|
||||
positive,
|
||||
negative,
|
||||
frame_idx,
|
||||
latent_image,
|
||||
noise_mask,
|
||||
t,
|
||||
strength,
|
||||
scale_factors,
|
||||
)
|
||||
)
|
||||
|
||||
# Track with custom attention strength/mask
|
||||
pre_filter_count = t.shape[2] * t.shape[3] * t.shape[4]
|
||||
guide_latent_shape = list(t.shape[2:])
|
||||
norm_mask = normalize_mask(attention_mask)
|
||||
positive = append_guide_attention_entry(
|
||||
positive,
|
||||
pre_filter_count,
|
||||
guide_latent_shape,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=norm_mask,
|
||||
)
|
||||
negative = append_guide_attention_entry(
|
||||
negative,
|
||||
pre_filter_count,
|
||||
guide_latent_shape,
|
||||
attention_strength=attention_strength,
|
||||
attention_mask=norm_mask,
|
||||
)
|
||||
|
||||
return (
|
||||
positive,
|
||||
negative,
|
||||
{"samples": latent_image, "noise_mask": noise_mask},
|
||||
)
|
||||
|
||||
-259
@@ -1,8 +1,6 @@
|
||||
import copy
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
@@ -72,260 +70,3 @@ class LTXVAdainLatent:
|
||||
|
||||
latents_copy["samples"] = torch.lerp(latents["samples"], t, factor)
|
||||
return (latents_copy,)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVStatNormLatent")
|
||||
class LTXVStatNormLatent:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latents": ("LATENT",),
|
||||
"target_mean": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"target_std": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.01,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"percentile": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 95.0,
|
||||
"min": 50.0,
|
||||
"max": 100.0,
|
||||
"step": 0.1,
|
||||
"round": 0.1,
|
||||
"tooltip": "Percentile of distribution to use for statistics calculation",
|
||||
},
|
||||
),
|
||||
"factor": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"clip_outliers": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "statistical_normalize"
|
||||
|
||||
CATEGORY = "Lightricks/latents"
|
||||
|
||||
def statistical_normalize(
|
||||
self, latents, target_mean, target_std, percentile, factor, clip_outliers
|
||||
):
|
||||
latents_copy = copy.deepcopy(latents)
|
||||
t = latents_copy["samples"] # B x C x F x H x W
|
||||
|
||||
# For 95% of distribution, we want to exclude 2.5% from each tail
|
||||
lower_percentile = (100 - percentile) / 2
|
||||
upper_percentile = 100 - lower_percentile
|
||||
|
||||
for i in range(t.size(0)): # batch
|
||||
for c in range(t.size(1)): # channel
|
||||
channel_data = t[i, c]
|
||||
original_shape = channel_data.shape
|
||||
channel_flat = channel_data.flatten()
|
||||
|
||||
# Calculate percentiles
|
||||
lower_bound = torch.quantile(channel_flat, lower_percentile / 100)
|
||||
upper_bound = torch.quantile(channel_flat, upper_percentile / 100)
|
||||
|
||||
# Create mask for values within the percentile range
|
||||
mask = (channel_flat >= lower_bound) & (channel_flat <= upper_bound)
|
||||
|
||||
# Calculate mean and std only on the masked values
|
||||
if mask.sum() > 0:
|
||||
filtered_data = channel_flat[mask]
|
||||
current_mean = filtered_data.mean()
|
||||
current_std = filtered_data.std()
|
||||
|
||||
# Avoid division by zero
|
||||
if current_std > 1e-8:
|
||||
# Normalize all values
|
||||
normalized_flat = (
|
||||
(channel_flat - current_mean) / current_std
|
||||
) * target_std + target_mean
|
||||
|
||||
if clip_outliers:
|
||||
# Calculate the normalized bounds
|
||||
normalized_lower = (
|
||||
(lower_bound - current_mean) / current_std
|
||||
) * target_std + target_mean
|
||||
normalized_upper = (
|
||||
(upper_bound - current_mean) / current_std
|
||||
) * target_std + target_mean
|
||||
|
||||
# Clip outliers to the normalized bounds
|
||||
normalized_flat = torch.where(
|
||||
channel_flat < lower_bound,
|
||||
normalized_lower,
|
||||
normalized_flat,
|
||||
)
|
||||
normalized_flat = torch.where(
|
||||
channel_flat > upper_bound,
|
||||
normalized_upper,
|
||||
normalized_flat,
|
||||
)
|
||||
|
||||
# Reshape back to original shape
|
||||
t[i, c] = normalized_flat.reshape(original_shape)
|
||||
else:
|
||||
# If std is too small, just shift by mean
|
||||
t[i, c] = channel_data - current_mean + target_mean
|
||||
|
||||
latents_copy["samples"] = torch.lerp(latents["samples"], t, factor)
|
||||
return (latents_copy,)
|
||||
|
||||
|
||||
class PerStepNormPatcher:
|
||||
"""
|
||||
Base class for per-step normalization nodes.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def required(s):
|
||||
return {
|
||||
"model": ("MODEL",),
|
||||
"factors": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "0.9, 0.75, 0.0",
|
||||
"tooltip": "Comma-separated list of factors, each factor will be used for one step.",
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def optional(s):
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def patch_smooth_norm(
|
||||
model: ModelPatcher, factors: str, normalize_fn: Callable
|
||||
) -> ModelPatcher:
|
||||
model = model.clone()
|
||||
factors = [float(x) for x in factors.split(",")]
|
||||
step = 0
|
||||
|
||||
def norm_fn(args):
|
||||
nonlocal step
|
||||
latent = args["denoised"]
|
||||
factor = factors[min(step, len(factors) - 1)]
|
||||
step += 1
|
||||
out_latent = normalize_fn(latent={"samples": latent}, factor=factor)
|
||||
return out_latent[0]["samples"]
|
||||
|
||||
model.set_model_sampler_post_cfg_function(norm_fn)
|
||||
return (model,)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(s, *args, **kwargs):
|
||||
return float("NaN")
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPerStepAdainPatcher")
|
||||
class LTXVPerStepAdainPatcher(PerStepNormPatcher):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": super().required()
|
||||
| {
|
||||
"reference": ("LATENT",),
|
||||
},
|
||||
"optional": super().optional()
|
||||
| {
|
||||
"per_frame": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch_model"
|
||||
|
||||
CATEGORY = "Lightricks/latents"
|
||||
|
||||
def patch_model(self, model, factors, reference, per_frame=False):
|
||||
def cfg_adain(latent, factor):
|
||||
return LTXVAdainLatent().batch_normalize(
|
||||
latent, reference, factor, per_frame
|
||||
)
|
||||
|
||||
return self.patch_smooth_norm(model, factors, cfg_adain)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPerStepStatNormPatcher")
|
||||
class LTXVPerStepStatNormPatcher(PerStepNormPatcher):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": super().required()
|
||||
| {
|
||||
"target_mean": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"target_std": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.01,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"percentile": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 95.0,
|
||||
"min": 50.0,
|
||||
"max": 100.0,
|
||||
"step": 0.1,
|
||||
"round": 0.1,
|
||||
"tooltip": "Percentile of distribution to use for statistics calculation",
|
||||
},
|
||||
),
|
||||
"clip_outliers": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch_model"
|
||||
|
||||
CATEGORY = "Lightricks/latents"
|
||||
|
||||
def patch_model(
|
||||
self, model, factors, target_mean, target_std, percentile, clip_outliers
|
||||
):
|
||||
def cfg_stat_norm(latent, factor):
|
||||
return LTXVStatNormLatent().statistical_normalize(
|
||||
latent, target_mean, target_std, percentile, factor, clip_outliers
|
||||
)
|
||||
|
||||
return self.patch_smooth_norm(model, factors, cfg_stat_norm)
|
||||
|
||||
-117
@@ -1,5 +1,3 @@
|
||||
from typing import Optional
|
||||
|
||||
import comfy.utils
|
||||
import comfy_extras.nodes_lt as nodes_lt
|
||||
import numpy as np
|
||||
@@ -91,121 +89,6 @@ class LTXVSelectLatents:
|
||||
raise
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAddLatents")
|
||||
class LTXVAddLatents:
|
||||
"""
|
||||
Concatenates two video latents along the frames dimension.
|
||||
|
||||
Features:
|
||||
- Validates dimension compatibility
|
||||
- Handles device placement
|
||||
- Preserves noise masks with proper handling
|
||||
- Supports batch processing
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latents1": ("LATENT",),
|
||||
"latents2": ("LATENT",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "add_latents"
|
||||
CATEGORY = "latent/video"
|
||||
DESCRIPTION = (
|
||||
"Concatenates two video latents along the frames dimension. "
|
||||
"latents1 and latents2 must have the same dimensions except for the frames dimension."
|
||||
)
|
||||
|
||||
def add_latents(
|
||||
self, latents1: torch.Tensor, latents2: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Concatenates two video latents along the frames dimension.
|
||||
|
||||
Args:
|
||||
latents1 (dict): First video latent dictionary
|
||||
latents2 (dict): Second video latent dictionary
|
||||
|
||||
Returns:
|
||||
tuple: Contains concatenated latent dictionary
|
||||
|
||||
Raises:
|
||||
ValueError: If latent dimensions don't match
|
||||
RuntimeError: If tensor operations fail
|
||||
"""
|
||||
try:
|
||||
s = latents1.copy()
|
||||
video_latent1 = latents1["samples"]
|
||||
video_latent2 = latents2["samples"]
|
||||
|
||||
# Ensure tensors are on the same device
|
||||
target_device = video_latent1.device
|
||||
video_latent2 = video_latent2.to(target_device)
|
||||
|
||||
# Validate dimensions
|
||||
self._validate_dimensions(video_latent1, video_latent2)
|
||||
|
||||
# Concatenate along frames dimension
|
||||
s["samples"] = torch.cat([video_latent1, video_latent2], dim=2)
|
||||
|
||||
# Handle noise masks
|
||||
s["noise_mask"] = self._merge_noise_masks(
|
||||
latents1, latents2, video_latent1.shape[2], video_latent2.shape[2]
|
||||
)
|
||||
|
||||
return (s,)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[LTXVAddLatents] Error: {str(e)}")
|
||||
raise
|
||||
|
||||
def _validate_dimensions(self, latent1: torch.Tensor, latent2: torch.Tensor):
|
||||
"""Validates that latent dimensions match except for frames."""
|
||||
# video latents
|
||||
if latent1.ndim == 5 and latent2.ndim == 5:
|
||||
b1, c1, f1, h1, w1 = latent1.shape
|
||||
b2, c2, f2, h2, w2 = latent2.shape
|
||||
|
||||
if not (b1 == b2 and c1 == c2 and h1 == h2 and w1 == w2):
|
||||
raise ValueError(
|
||||
f"Latent dimensions must match (except frames dimension).\n"
|
||||
f"Got shapes {latent1.shape} and {latent2.shape}"
|
||||
)
|
||||
# audio latents
|
||||
elif latent1.ndim == 4 and latent2.ndim == 4:
|
||||
b1, c1, f1, s1 = latent1.shape
|
||||
b2, c2, f2, s2 = latent2.shape
|
||||
|
||||
if not (b1 == b2 and c1 == c2 and s1 == s2):
|
||||
raise ValueError(
|
||||
f"Latent dimensions must match (except frames dimension).\n"
|
||||
f"Got shapes {latent1.shape} and {latent2.shape}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Latent dimensions must be 4 (audio) or 5 (video) for both inputs.\n"
|
||||
f"Got shapes {latent1.shape} and {latent2.shape}"
|
||||
)
|
||||
|
||||
def _merge_noise_masks(
|
||||
self, latents1: torch.Tensor, latents2: torch.Tensor, frames1: int, frames2: int
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""Merges noise masks from both latents with proper handling."""
|
||||
if "noise_mask" in latents1 and "noise_mask" in latents2:
|
||||
return torch.cat([latents1["noise_mask"], latents2["noise_mask"]], dim=2)
|
||||
elif "noise_mask" in latents1 and latents1["noise_mask"] is not None:
|
||||
zeros = torch.zeros_like(latents1["noise_mask"][:, :, :frames2, :, :])
|
||||
return torch.cat([latents1["noise_mask"], zeros], dim=2)
|
||||
elif "noise_mask" in latents2 and latents2["noise_mask"] is not None:
|
||||
zeros = torch.zeros_like(latents2["noise_mask"][:, :, :frames1, :, :])
|
||||
return torch.cat([zeros, latents2["noise_mask"]], dim=2)
|
||||
return None
|
||||
|
||||
|
||||
@comfy_node(name="LTXVSetVideoLatentNoiseMasks")
|
||||
class LTXVSetVideoLatentNoiseMasks:
|
||||
"""
|
||||
|
||||
@@ -1,212 +0,0 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from comfy.model_management import is_device_cpu
|
||||
from comfy_extras.nodes_mask import GrowMask
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPreprocessMasks")
|
||||
class LTXVPreprocessMasks:
|
||||
"""
|
||||
Preprocesses masks for use with the LTXVideo model's latent masking.
|
||||
|
||||
This node handles temporal mask processing by:
|
||||
1. Validating mask dimensions against VAE downscaling factors
|
||||
2. Optionally inverting masks
|
||||
3. Handling the first frame mask separately (can be ignored)
|
||||
4. Combining temporal masks using max pooling
|
||||
5. Growing/shrinking masks with morphological operations
|
||||
6. Clamping mask values to ensure proper opacity
|
||||
|
||||
The output is a set of masks ready for latent-space masking operations.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls: "LTXVPreprocessMasks") -> dict[str, Any]: # noqa: N802
|
||||
return {
|
||||
"required": {
|
||||
"masks": ("MASK",),
|
||||
"vae": ("VAE",),
|
||||
"invert_input_masks": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "Invert the input masks before processing."
|
||||
"Useful for masking vs unmasking operations.",
|
||||
},
|
||||
),
|
||||
"ignore_first_mask": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Zero out the first mask, typically used when"
|
||||
"it corresponds to the conditioning frame.",
|
||||
},
|
||||
),
|
||||
"pooling_method": (
|
||||
["max", "mean", "min"],
|
||||
{
|
||||
"default": "max",
|
||||
"tooltip": "Method to combine temporal masks. Max preserves strongest values, "
|
||||
"mean averages them, min takes weakest values.",
|
||||
},
|
||||
),
|
||||
"grow_mask": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": -16384,
|
||||
"max": 16384,
|
||||
"step": 1,
|
||||
"tooltip": "Pixels to grow (positive) or shrink (negative) the mask."
|
||||
"Uses morphological operations.",
|
||||
},
|
||||
),
|
||||
"tapered_corners": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Use tapered corners in morphological operations for smoother mask edges.",
|
||||
},
|
||||
),
|
||||
"clamp_min": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Minimum value to clamp masks to, ensuring minimum opacity.",
|
||||
},
|
||||
),
|
||||
"clamp_max": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Maximum value to clamp masks to.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "preprocess_masks"
|
||||
CATEGORY = "Lightricks/mask_operations"
|
||||
DESCRIPTION = (
|
||||
"Preprocess masks to be used for masking latents in the LTXVideo model."
|
||||
)
|
||||
|
||||
def preprocess_masks(
|
||||
self,
|
||||
masks: torch.Tensor,
|
||||
vae, # noqa: ANN001
|
||||
pooling_method: str,
|
||||
grow_mask: int,
|
||||
tapered_corners: bool,
|
||||
clamp_min: float,
|
||||
clamp_max: float,
|
||||
ignore_first_mask: bool,
|
||||
invert_input_masks: bool,
|
||||
) -> tuple[torch.Tensor]:
|
||||
"""
|
||||
Preprocess masks for temporal latent masking.
|
||||
|
||||
Args:
|
||||
masks: Input masks tensor of shape (batch_size, H, W) where batch_size = 1 + (num_frames - 1)
|
||||
The first mask is for the conditioning frame, remaining masks are for generated frames
|
||||
vae: VAE model object containing downscale_index_formula attribute
|
||||
pooling_method: Method to combine temporal masks ('max', 'mean', 'min')
|
||||
grow_mask: Pixels to grow (positive) or shrink (negative) the mask
|
||||
tapered_corners: Whether to use tapered corners in morphological operations
|
||||
clamp_min: Minimum value to clamp masks to (ensures minimum opacity)
|
||||
clamp_max: Maximum value to clamp masks to
|
||||
ignore_first_mask: Whether to zero out the first mask (often used for conditioning)
|
||||
invert_input_masks: Whether to invert all input masks (1 - mask)
|
||||
|
||||
Returns:
|
||||
Tuple containing the processed masks tensor
|
||||
|
||||
Raises:
|
||||
ValueError: If mask dimensions are not compatible with VAE downscaling factors
|
||||
"""
|
||||
|
||||
if masks.ndim != 3:
|
||||
raise ValueError("Masks must be of shape (batch_size, H, W).")
|
||||
|
||||
batch_size, H, W = masks.shape
|
||||
|
||||
# Get VAE downscaling factors for temporal and spatial dimensions
|
||||
time_scale_factor, width_scale_factor, height_scale_factor = (
|
||||
vae.downscale_index_formula
|
||||
)
|
||||
|
||||
# Validate that mask dimensions are compatible with VAE downscaling
|
||||
# The batch size minus 1 (for the first frame) must be divisible by time_scale_factor
|
||||
if (batch_size - 1) % time_scale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Masks batch size must have a multiple of {time_scale_factor} masks + 1."
|
||||
)
|
||||
|
||||
if H % height_scale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Masks height must have a multiple of {height_scale_factor}."
|
||||
)
|
||||
|
||||
if W % width_scale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Masks width must have a multiple of {width_scale_factor}."
|
||||
)
|
||||
|
||||
# Optionally invert all masks (useful for masking vs unmasking operations)
|
||||
if invert_input_masks:
|
||||
masks = 1.0 - masks
|
||||
|
||||
# Handle first mask separately - often used as conditioning frame
|
||||
if ignore_first_mask:
|
||||
first_mask = masks[0:1, :, :] * 0 # Zero out the first mask
|
||||
else:
|
||||
first_mask = masks[0:1, :, :]
|
||||
|
||||
# Process remaining masks by grouping them according to time_scale_factor
|
||||
# and taking the maximum value across each temporal group
|
||||
remaining = masks[1:].view(-1, time_scale_factor, H, W)
|
||||
|
||||
# Apply selected pooling method
|
||||
pooling_methods = {
|
||||
"max": lambda x: x.max(dim=1)[0],
|
||||
"mean": lambda x: x.mean(dim=1),
|
||||
"min": lambda x: x.min(dim=1)[0],
|
||||
}
|
||||
|
||||
combined = pooling_methods[pooling_method](remaining)
|
||||
|
||||
# Apply morphological operations to grow or shrink masks
|
||||
if grow_mask != 0:
|
||||
grow_mask_node = GrowMask()
|
||||
|
||||
# GrowMask node is not supported on GPU, so we need to move the combined tensor to CPU
|
||||
# if it is not already on CPU
|
||||
input_device = combined.device
|
||||
if not is_device_cpu(input_device):
|
||||
combined = combined.cpu()
|
||||
|
||||
(combined,) = grow_mask_node.expand_mask(
|
||||
mask=combined, expand=grow_mask, tapered_corners=tapered_corners
|
||||
)
|
||||
|
||||
# Move the combined tensor back to the original device
|
||||
if not is_device_cpu(input_device):
|
||||
combined = combined.to(input_device)
|
||||
|
||||
# Combine first mask with processed remaining masks and clamp values
|
||||
# Clamping ensures masks have proper opacity values
|
||||
processed_masks = torch.cat(
|
||||
[first_mask, torch.clamp(combined, clamp_min, clamp_max)], dim=0
|
||||
)
|
||||
|
||||
return (processed_masks,)
|
||||
@@ -1,208 +0,0 @@
|
||||
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
|
||||
DESCRIPTION = "Downloads and initializes LLM and image captioning models from Hugging Face to enhance text prompts for image generation."
|
||||
|
||||
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
|
||||
DESCRIPTION = (
|
||||
"Enhances text prompts for image generation using LLMs. "
|
||||
"Optionally incorporates reference images to create more contextually relevant descriptions."
|
||||
)
|
||||
|
||||
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],)
|
||||
@@ -1,258 +0,0 @@
|
||||
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
|
||||
+1
-1
@@ -7,4 +7,4 @@ kornia
|
||||
ninja~=1.11.1.4
|
||||
# Required for HDR EXR read/write (same as ltx-pipelines); no OpenCV EXR fallback.
|
||||
openimageio
|
||||
transformers[timm]>=4.50.0
|
||||
transformers>=4.50.0
|
||||
|
||||
@@ -33,16 +33,6 @@ def stg(
|
||||
return noise_pred
|
||||
|
||||
|
||||
class MomentumBuffer:
|
||||
def __init__(self, momentum: float):
|
||||
self.momentum = momentum
|
||||
self.running_average = 0
|
||||
|
||||
def update(self, update_value: torch.Tensor):
|
||||
new_average = self.momentum * self.running_average
|
||||
self.running_average = update_value + new_average
|
||||
|
||||
|
||||
def project(v0: torch.Tensor, v1: torch.Tensor):
|
||||
dtype = v0.dtype
|
||||
v0, v1 = v0.double(), v1.double()
|
||||
@@ -792,153 +782,3 @@ class STGAdvancedPresetsNode:
|
||||
|
||||
def get_preset(self, preset=None):
|
||||
return (preset,)
|
||||
|
||||
|
||||
class APGGuider(comfy.samplers.CFGGuider):
|
||||
def __init__(
|
||||
self,
|
||||
model: ModelPatcher,
|
||||
cfg_scale,
|
||||
eta,
|
||||
norm_threshold,
|
||||
momentum_coefficient: float = -0.9,
|
||||
):
|
||||
self.model = model.clone()
|
||||
super().__init__(self.model)
|
||||
self.momentum_coefficient = momentum_coefficient
|
||||
self.momentum_buffer = MomentumBuffer(self.momentum_coefficient)
|
||||
self.eta = eta
|
||||
self.norm_threshold = norm_threshold
|
||||
self.cfg_scale = cfg_scale
|
||||
self.previous_timestep = None
|
||||
|
||||
def set_conds(self, positive, negative):
|
||||
self.inner_set_conds({"positive": positive, "negative": negative})
|
||||
|
||||
def predict_noise(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
model_options: dict = {},
|
||||
seed=None,
|
||||
):
|
||||
# CFG zero init - skipping steps with timestep bigger than given threshold.
|
||||
|
||||
# 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
|
||||
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]
|
||||
|
||||
if (
|
||||
self.previous_timestep is not None
|
||||
and timestep.item() > self.previous_timestep
|
||||
):
|
||||
print("Resetting momentum buffer")
|
||||
self.momentum_buffer = MomentumBuffer(self.momentum_coefficient)
|
||||
|
||||
noise_pred_neg = 0
|
||||
if not math.isclose(self.cfg_scale, 1.0):
|
||||
noise_pred_neg = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[negative_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
|
||||
apg_result = apg(
|
||||
noise_pred_pos,
|
||||
noise_pred_neg,
|
||||
self.cfg_scale,
|
||||
self.eta,
|
||||
self.norm_threshold,
|
||||
)
|
||||
|
||||
# normally this would be done in cfg_function, but we skipped
|
||||
# that for efficiency: we can compute the noise predictions in
|
||||
# a single call to calc_cond_batch() (rather than two)
|
||||
# so we replicate the hook here
|
||||
for fn in model_options.get("sampler_post_cfg_function", []):
|
||||
args = {
|
||||
"denoised": apg_result,
|
||||
"cond": positive_cond,
|
||||
"uncond": negative_cond,
|
||||
"model": self.inner_model,
|
||||
"uncond_denoised": noise_pred_neg,
|
||||
"cond_denoised": noise_pred_pos,
|
||||
"sigma": timestep,
|
||||
"model_options": model_options,
|
||||
"input": x,
|
||||
}
|
||||
apg_result = fn(args)
|
||||
|
||||
self.previous_timestep = timestep.item()
|
||||
return apg_result
|
||||
|
||||
|
||||
@comfy_node(name="APGGuider")
|
||||
class APGGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"cfg_scale": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"eta": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"norm_threshold": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"momentum_coefficient": (
|
||||
"FLOAT",
|
||||
{"default": -0.9, "min": -3.0, "max": 3.0, "step": 0.01},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
|
||||
FUNCTION = "get_guider"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
DESCRIPTION = """
|
||||
The APG Guider implements Adaptive Projected Guidance (APG).
|
||||
Reference: https://arxiv.org/abs/2410.02416.
|
||||
"""
|
||||
|
||||
def get_guider(
|
||||
self,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
cfg_scale,
|
||||
eta,
|
||||
norm_threshold,
|
||||
momentum_coefficient,
|
||||
):
|
||||
guider = APGGuider(
|
||||
model,
|
||||
cfg_scale,
|
||||
eta,
|
||||
norm_threshold,
|
||||
momentum_coefficient,
|
||||
)
|
||||
guider.set_conds(positive, negative)
|
||||
guider.raw_conds = (positive, negative)
|
||||
return (guider,)
|
||||
|
||||
@@ -1,54 +0,0 @@
|
||||
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",
|
||||
}
|
||||
@@ -1,257 +0,0 @@
|
||||
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,
|
||||
)
|
||||
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 = self._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
|
||||
@@ -1,90 +0,0 @@
|
||||
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,)
|
||||
@@ -1,30 +0,0 @@
|
||||
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},)
|
||||
@@ -1,40 +0,0 @@
|
||||
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}),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "LATENT")
|
||||
RETURN_NAMES = ("model", "latent")
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
FUNCTION = "generate"
|
||||
|
||||
def generate(self, model, latent, image_latent, index, strength):
|
||||
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,
|
||||
strength,
|
||||
)
|
||||
|
||||
return (
|
||||
model,
|
||||
{"samples": latent, "noise_mask": noise_mask},
|
||||
)
|
||||
@@ -1,36 +0,0 @@
|
||||
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,)
|
||||
@@ -1,180 +0,0 @@
|
||||
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,)
|
||||
@@ -1,91 +0,0 @@
|
||||
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,)
|
||||
@@ -1,149 +0,0 @@
|
||||
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,)
|
||||
@@ -1,20 +0,0 @@
|
||||
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,)
|
||||
@@ -1,275 +0,0 @@
|
||||
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,)
|
||||
@@ -1,376 +0,0 @@
|
||||
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,)
|
||||
@@ -1,16 +0,0 @@
|
||||
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)
|
||||
@@ -1,55 +0,0 @@
|
||||
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)
|
||||
@@ -1,8 +0,0 @@
|
||||
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)
|
||||
@@ -1,6 +0,0 @@
|
||||
def isinstance_str(x: object, cls_name: str):
|
||||
for _cls in x.__class__.__mro__:
|
||||
if _cls.__name__ == cls_name:
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -1,23 +0,0 @@
|
||||
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,38 +0,0 @@
|
||||
from comfy.model_management import dtype_size
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPatcherVAE")
|
||||
class LTXVPatcherVAE:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VAE",)
|
||||
FUNCTION = "patch"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV VAE Patcher"
|
||||
|
||||
def patch(self, vae):
|
||||
from q8_kernels.integration.patch_vae import patch_vae
|
||||
|
||||
vae_model = vae.first_stage_model
|
||||
vae.memory_used_decode = (
|
||||
lambda shape, dtype: shape[1]
|
||||
* shape[2]
|
||||
* shape[3]
|
||||
* shape[4]
|
||||
* 8
|
||||
* 8
|
||||
* 8
|
||||
* 3
|
||||
* dtype_size(dtype)
|
||||
)
|
||||
|
||||
patch_vae(vae_model, patch_block=4)
|
||||
return (vae,)
|
||||
Reference in New Issue
Block a user