793 lines
28 KiB
Python
793 lines
28 KiB
Python
import contextlib
|
|
import json
|
|
import math
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import List, Optional, Union
|
|
|
|
import comfy.ldm.modules.attention
|
|
import comfy.samplers
|
|
import torch
|
|
from comfy.model_patcher import ModelPatcher
|
|
|
|
from .nodes_registry import comfy_node
|
|
|
|
|
|
def stg(
|
|
noise_pred_pos,
|
|
noise_pred_neg,
|
|
noise_pred_pertubed,
|
|
cfg_scale,
|
|
stg_scale,
|
|
rescale_scale,
|
|
):
|
|
noise_pred = (
|
|
noise_pred_pos
|
|
+ (cfg_scale - 1) * (noise_pred_pos - noise_pred_neg)
|
|
+ stg_scale * (noise_pred_pos - noise_pred_pertubed)
|
|
)
|
|
if rescale_scale != 0:
|
|
factor = noise_pred_pos.std() / noise_pred.std()
|
|
factor = rescale_scale * factor + (1 - rescale_scale)
|
|
noise_pred = noise_pred * factor
|
|
return noise_pred
|
|
|
|
|
|
def project(v0: torch.Tensor, v1: torch.Tensor):
|
|
dtype = v0.dtype
|
|
v0, v1 = v0.double(), v1.double()
|
|
v1 = torch.nn.functional.normalize(v1, dim=[-1, -2, -3])
|
|
v0_parallel = (v0 * v1).sum(dim=[-1, -2, -3], keepdim=True) * v1
|
|
v0_orthogonal = v0 - v0_parallel
|
|
return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
|
|
|
|
|
|
def apg(
|
|
noise_pred_pos: torch.Tensor,
|
|
noise_pred_neg: torch.Tensor,
|
|
cfg_scale: float,
|
|
eta: float = 1.0,
|
|
norm_threshold: float = 0.0,
|
|
):
|
|
diff = noise_pred_pos - noise_pred_neg
|
|
if norm_threshold > 0:
|
|
ones = torch.ones_like(diff)
|
|
diff_norm = diff.norm(p=2, dim=[-1, -2, -3], keepdim=True)
|
|
scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
|
|
diff = diff * scale_factor
|
|
diff_parallel, diff_orthogonal = project(diff, noise_pred_pos)
|
|
normalized_update = diff_orthogonal + eta * diff_parallel
|
|
noise_pred = noise_pred_pos + (cfg_scale - 1) * normalized_update
|
|
return noise_pred
|
|
|
|
|
|
@comfy_node(name="LTXVApplySTG")
|
|
class LTXVApplySTG:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": (
|
|
"MODEL",
|
|
{"tooltip": "The model to apply the STG to."},
|
|
),
|
|
"block_indices": (
|
|
"STRING",
|
|
{
|
|
"default": "14, 19",
|
|
"tooltip": "Comma-separated indices of the blocks to apply the STG to.",
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
FUNCTION = "apply_stg"
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
CATEGORY = "lightricks/LTXV"
|
|
|
|
DESCRIPTION = "Defines the blocks to apply the STG to."
|
|
|
|
def apply_stg(self, model: ModelPatcher, block_indices: str):
|
|
skip_block_list = [int(i.strip()) for i in block_indices.split(",")]
|
|
new_model = model.clone()
|
|
|
|
if "skip_block_list" in new_model.model_options["transformer_options"]:
|
|
skip_block_list.extend(
|
|
new_model.model_options["transformer_options"]["skip_block_list"]
|
|
)
|
|
new_model.model_options["transformer_options"][
|
|
"skip_block_list"
|
|
] = skip_block_list
|
|
|
|
return (new_model,)
|
|
|
|
|
|
@dataclass
|
|
class STGFlag:
|
|
do_skip: bool = False
|
|
skip_layers: List[int] = None
|
|
|
|
|
|
def _unwrap_attention_tensor(v):
|
|
# ComfyUI wraps q/k/v in AttentionTensorContainer before calling attention.
|
|
# The skip path returns v as the attention output, so it must be a plain tensor.
|
|
if isinstance(v, comfy.ldm.modules.attention.AttentionTensorContainer):
|
|
return v.take()
|
|
return v
|
|
|
|
|
|
# context manager that replaces the attention function in a transformer block
|
|
class PatchAttention(contextlib.AbstractContextManager):
|
|
def __init__(self, attn_idx: Optional[Union[int, List[int]]] = None):
|
|
self.current_idx = -1
|
|
|
|
if isinstance(attn_idx, int):
|
|
self.attn_idx = [attn_idx]
|
|
elif attn_idx is None:
|
|
self.attn_idx = [0]
|
|
else:
|
|
self.attn_idx = list(attn_idx)
|
|
|
|
def __enter__(self):
|
|
self.original_attention = comfy.ldm.modules.attention.optimized_attention
|
|
self.original_attention_masked = (
|
|
comfy.ldm.modules.attention.optimized_attention_masked
|
|
)
|
|
|
|
comfy.ldm.modules.attention.optimized_attention = self.stg_attention
|
|
comfy.ldm.modules.attention.optimized_attention_masked = (
|
|
self.stg_attention_masked
|
|
)
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
comfy.ldm.modules.attention.optimized_attention = self.original_attention
|
|
comfy.ldm.modules.attention.optimized_attention_masked = (
|
|
self.original_attention_masked
|
|
)
|
|
|
|
self.original_attention = None
|
|
self.original_attention_masked = None
|
|
|
|
def stg_attention(self, q, k, v, heads, *args, **kwargs):
|
|
self.current_idx += 1
|
|
if self.current_idx in self.attn_idx:
|
|
return _unwrap_attention_tensor(v)
|
|
else:
|
|
return self.original_attention(q, k, v, heads, *args, **kwargs)
|
|
|
|
def stg_attention_masked(self, q, k, v, heads, *args, **kwargs):
|
|
self.current_idx += 1
|
|
if self.current_idx in self.attn_idx:
|
|
return _unwrap_attention_tensor(v)
|
|
else:
|
|
return self.original_attention_masked(q, k, v, heads, *args, **kwargs)
|
|
|
|
|
|
class STGBlockWrapper:
|
|
"""Wraps transformer blocks to be able to skip attention layers."""
|
|
|
|
def __init__(self, block, stg_flag: STGFlag, idx: int):
|
|
self.flag = stg_flag
|
|
self.idx = idx
|
|
self.block = block
|
|
|
|
def __call__(self, args, extra_args):
|
|
context_manager = contextlib.nullcontext()
|
|
|
|
stg_indexes = args["transformer_options"].get("stg_indexes", [0])
|
|
if self.flag.do_skip and self.idx in self.flag.skip_layers:
|
|
context_manager = PatchAttention(stg_indexes)
|
|
|
|
with context_manager:
|
|
hidden_state = extra_args["original_block"](args)
|
|
return hidden_state
|
|
|
|
|
|
class STGGuider(comfy.samplers.CFGGuider):
|
|
def __init__(
|
|
self, model: ModelPatcher, cfg, stg_scale, rescale_scale: float = None
|
|
):
|
|
model = model.clone()
|
|
super().__init__(model)
|
|
|
|
self.stg_flag = STGFlag(
|
|
do_skip=False,
|
|
skip_layers=model.model_options["transformer_options"]["skip_block_list"],
|
|
)
|
|
|
|
self.patch_model(model, self.stg_flag)
|
|
|
|
self.cfg = cfg
|
|
self.stg_scale = stg_scale
|
|
self.rescale_scale = rescale_scale
|
|
|
|
@classmethod
|
|
def patch_model(cls, model: ModelPatcher, stg_flag: STGFlag):
|
|
transformer_blocks = cls.get_transformer_blocks(model)
|
|
|
|
for i, block in enumerate(transformer_blocks):
|
|
model.set_model_patch_replace(
|
|
STGBlockWrapper(block, stg_flag, i), "dit", "double_block", i
|
|
)
|
|
|
|
@staticmethod
|
|
def get_transformer_blocks(model: ModelPatcher):
|
|
diffusion_model = model.get_model_object("diffusion_model")
|
|
key = "diffusion_model.transformer_blocks"
|
|
if diffusion_model.__class__.__name__ == "LTXVTransformer3D":
|
|
key = "diffusion_model.transformer.transformer_blocks"
|
|
return model.get_model_object(key)
|
|
|
|
def set_conds(self, positive, negative):
|
|
self.inner_set_conds({"positive": positive, "negative": negative})
|
|
|
|
def predict_noise(
|
|
self,
|
|
x: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
model_options: dict = {},
|
|
seed=None,
|
|
):
|
|
# in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
|
|
# but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly
|
|
|
|
positive_cond = self.conds.get("positive", None)
|
|
negative_cond = self.conds.get("negative", None)
|
|
|
|
if model_options.get("sigma_to_params_mapping", None) is not None:
|
|
cfg_value, stg_scale, stg_layer_skip_layer_indices, stg_rescale = (
|
|
model_options["sigma_to_params_mapping"](timestep)
|
|
)
|
|
self.stg_flag.skip_layers = stg_layer_skip_layer_indices
|
|
self.patch_model(self.model_patcher, self.stg_flag)
|
|
|
|
else:
|
|
cfg_value = self.cfg
|
|
stg_scale = self.stg_scale
|
|
stg_rescale = self.rescale_scale
|
|
|
|
noise_pred_pos = comfy.samplers.calc_cond_batch(
|
|
self.inner_model,
|
|
[positive_cond],
|
|
x,
|
|
timestep,
|
|
model_options,
|
|
)[0]
|
|
|
|
noise_pred_neg = 0
|
|
noise_pred_perturbed = 0
|
|
|
|
if not math.isclose(cfg_value, 1.0):
|
|
noise_pred_neg = comfy.samplers.calc_cond_batch(
|
|
self.inner_model,
|
|
[negative_cond],
|
|
x,
|
|
timestep,
|
|
model_options,
|
|
)[0]
|
|
|
|
if not math.isclose(stg_scale, 0.0):
|
|
try:
|
|
model_options["transformer_options"]["ptb_index"] = 0
|
|
self.stg_flag.do_skip = True
|
|
noise_pred_perturbed = comfy.samplers.calc_cond_batch(
|
|
self.inner_model,
|
|
[positive_cond],
|
|
x,
|
|
timestep,
|
|
model_options,
|
|
)[0]
|
|
finally:
|
|
self.stg_flag.do_skip = False
|
|
del model_options["transformer_options"]["ptb_index"]
|
|
|
|
stg_result = stg(
|
|
noise_pred_pos,
|
|
noise_pred_neg,
|
|
noise_pred_perturbed,
|
|
cfg_value,
|
|
stg_scale,
|
|
stg_rescale,
|
|
)
|
|
|
|
# normally this would be done in cfg_function, but we skipped
|
|
# that for efficiency: we can compute the noise predictions in
|
|
# a single call to calc_cond_batch() (rather than two)
|
|
# so we replicate the hook here
|
|
for fn in model_options.get("sampler_post_cfg_function", []):
|
|
args = {
|
|
"denoised": stg_result,
|
|
"cond": positive_cond,
|
|
"uncond": negative_cond,
|
|
"model": self.inner_model,
|
|
"uncond_denoised": noise_pred_neg,
|
|
"cond_denoised": noise_pred_pos,
|
|
"sigma": timestep,
|
|
"model_options": model_options,
|
|
"input": x,
|
|
# not in the original call in samplers.py:cfg_function, but made available for future hooks
|
|
"perturbed_cond": positive_cond,
|
|
"perturbed_cond_denoised": noise_pred_perturbed,
|
|
}
|
|
stg_result = fn(args)
|
|
|
|
return stg_result
|
|
|
|
|
|
class STGGuiderAdvanced(comfy.samplers.CFGGuider):
|
|
def __init__(
|
|
self,
|
|
model: ModelPatcher,
|
|
sigma_list,
|
|
cfg_list,
|
|
stg_scale_list,
|
|
stg_rescale_list,
|
|
stg_layers_indices_list,
|
|
skip_steps_sigma_threshold,
|
|
cfg_star_rescale,
|
|
apply_apg,
|
|
apg_cfg_scale,
|
|
eta,
|
|
norm_threshold,
|
|
):
|
|
model = model.clone()
|
|
super().__init__(model)
|
|
|
|
self.stg_flag = STGFlag(
|
|
do_skip=False,
|
|
skip_layers=model.model_options["transformer_options"].get(
|
|
"skip_block_list"
|
|
),
|
|
)
|
|
|
|
self.sigma_list = sigma_list
|
|
self.cfg_list = cfg_list
|
|
self.stg_scale_list = stg_scale_list
|
|
self.stg_rescale_list = stg_rescale_list
|
|
self.stg_layers_indices_list = stg_layers_indices_list
|
|
self.skip_steps_sigma_threshold = skip_steps_sigma_threshold
|
|
self.cfg_star_rescale = cfg_star_rescale
|
|
self.apply_apg = apply_apg
|
|
self.apg_cfg_scale = apg_cfg_scale
|
|
self.eta = eta
|
|
self.norm_threshold = norm_threshold
|
|
STGGuider.patch_model(model, self.stg_flag)
|
|
|
|
def sigma_to_params_mapping(self, sigma):
|
|
# Find the closest higher sigma value and return corresponding cfg
|
|
higher_sigmas = [s for s in self.sigma_list if s >= sigma]
|
|
if not higher_sigmas:
|
|
closest_idx = -1 # Return last cfg if no higher sigma exists
|
|
else:
|
|
closest_higher = min(higher_sigmas)
|
|
closest_idx = self.sigma_list.index(closest_higher)
|
|
return (
|
|
self.cfg_list[closest_idx],
|
|
self.stg_scale_list[closest_idx],
|
|
self.stg_rescale_list[closest_idx],
|
|
self.stg_layers_indices_list[closest_idx],
|
|
)
|
|
|
|
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.
|
|
if timestep > self.skip_steps_sigma_threshold:
|
|
return torch.zeros_like(x)
|
|
|
|
# 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)
|
|
|
|
cfg_value, stg_scale, stg_rescale, stg_layer_skip_layer_indices = (
|
|
self.sigma_to_params_mapping(timestep)
|
|
)
|
|
|
|
if stg_layer_skip_layer_indices is not None:
|
|
self.stg_flag.skip_layers = stg_layer_skip_layer_indices
|
|
STGGuider.patch_model(self.model_patcher, self.stg_flag)
|
|
|
|
noise_pred_pos = comfy.samplers.calc_cond_batch(
|
|
self.inner_model,
|
|
[positive_cond],
|
|
x,
|
|
timestep,
|
|
model_options,
|
|
)[0]
|
|
|
|
noise_pred_neg = 0
|
|
noise_pred_perturbed = 0
|
|
|
|
if not math.isclose(cfg_value, 1.0) or (
|
|
self.apply_apg and not math.isclose(self.apg_cfg_scale, 1.0)
|
|
):
|
|
noise_pred_neg = comfy.samplers.calc_cond_batch(
|
|
self.inner_model,
|
|
[negative_cond],
|
|
x,
|
|
timestep,
|
|
model_options,
|
|
)[0]
|
|
|
|
if self.cfg_star_rescale:
|
|
batch_size = noise_pred_pos.shape[0]
|
|
|
|
positive_flat = noise_pred_pos.view(batch_size, -1)
|
|
negative_flat = noise_pred_neg.view(batch_size, -1)
|
|
dot_product = torch.sum(
|
|
positive_flat * negative_flat, dim=1, keepdim=True
|
|
)
|
|
squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
|
|
alpha = dot_product / squared_norm
|
|
noise_pred_neg = alpha * noise_pred_neg
|
|
|
|
if not math.isclose(stg_scale, 0.0):
|
|
try:
|
|
model_options["transformer_options"]["ptb_index"] = 0
|
|
self.stg_flag.do_skip = True
|
|
noise_pred_perturbed = comfy.samplers.calc_cond_batch(
|
|
self.inner_model,
|
|
[positive_cond],
|
|
x,
|
|
timestep,
|
|
model_options,
|
|
)[0]
|
|
finally:
|
|
self.stg_flag.do_skip = False
|
|
del model_options["transformer_options"]["ptb_index"]
|
|
|
|
stg_result = stg(
|
|
noise_pred_pos,
|
|
noise_pred_neg,
|
|
noise_pred_perturbed,
|
|
cfg_value,
|
|
stg_scale,
|
|
stg_rescale,
|
|
)
|
|
|
|
if self.apply_apg:
|
|
stg_result = apg(
|
|
stg_result,
|
|
noise_pred_neg,
|
|
cfg_scale=self.apg_cfg_scale,
|
|
momentum_buffer=None,
|
|
eta=self.eta,
|
|
norm_threshold=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": stg_result,
|
|
"cond": positive_cond,
|
|
"uncond": negative_cond,
|
|
"model": self.inner_model,
|
|
"uncond_denoised": noise_pred_neg,
|
|
"cond_denoised": noise_pred_pos,
|
|
"sigma": timestep,
|
|
"model_options": model_options,
|
|
"input": x,
|
|
# not in the original call in samplers.py:cfg_function, but made available for future hooks
|
|
"perturbed_cond": positive_cond,
|
|
"perturbed_cond_denoised": noise_pred_perturbed,
|
|
}
|
|
stg_result = fn(args)
|
|
|
|
return stg_result
|
|
|
|
|
|
@comfy_node(name="STGGuider")
|
|
class STGGuiderNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"cfg": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.1,
|
|
"round": 0.01,
|
|
},
|
|
),
|
|
"stg": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
|
),
|
|
"rescale": (
|
|
"FLOAT",
|
|
{"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("GUIDER",)
|
|
|
|
FUNCTION = "get_guider"
|
|
CATEGORY = "lightricks/LTXV"
|
|
|
|
DESCRIPTION = (
|
|
"Implements Spatiotemporal Skip Guidance (STG), a training-free method enhancing transformer-based "
|
|
"video diffusion models by selectively skipping layers during sampling. This approach improves video "
|
|
"quality without sacrificing diversity or motion fidelity."
|
|
"Reference: https://arxiv.org/abs/2411.18664."
|
|
)
|
|
|
|
def get_guider(self, model, positive, negative, cfg, stg, rescale):
|
|
guider = STGGuider(model, cfg, stg, rescale)
|
|
guider.set_conds(positive, negative)
|
|
return (guider,)
|
|
|
|
|
|
def load_stg_presets():
|
|
preset_file_path = os.path.join(
|
|
os.path.dirname(__file__), "presets", "stg_advanced_presets.json"
|
|
)
|
|
if os.path.exists(preset_file_path):
|
|
presets = json.load(open(preset_file_path))
|
|
preset_names = [preset["name"] for preset in presets]
|
|
else:
|
|
presets = []
|
|
preset_names = ["Custom"]
|
|
return presets, preset_names
|
|
|
|
|
|
STG_ADVANCED_PRESETS, STG_ADVANCED_PRESET_NAMES = load_stg_presets()
|
|
|
|
|
|
@comfy_node(name="STGGuiderAdvanced")
|
|
class STGGuiderAdvancedNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"skip_steps_sigma_threshold": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.998,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.001,
|
|
"tooltip": "Steps with sigma greater than this values will be skipped.",
|
|
},
|
|
),
|
|
"cfg_star_rescale": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "If true, applies the cfg star rescale, scales the negative prediction according to dot product between positive and negative.",
|
|
},
|
|
),
|
|
"sigmas": (
|
|
"STRING",
|
|
{
|
|
"default": "1.0, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180",
|
|
"tooltip": "Comma-separated list sigmas, the actual parameters will be selected according to the closest sigma from this list to current timestep sigma.",
|
|
},
|
|
),
|
|
"cfg_values": (
|
|
"STRING",
|
|
{
|
|
"default": "8, 6, 6, 4, 3, 1",
|
|
"tooltip": "Comma-separated list of cfg values. Should be same length as sigmas list.",
|
|
},
|
|
),
|
|
"stg_scale_values": (
|
|
"STRING",
|
|
{
|
|
"default": "4, 4, 3, 2, 1, 0",
|
|
"tooltip": "Comma-separated list of stg scale values. Should be same length as sigmas list.",
|
|
},
|
|
),
|
|
"stg_rescale_values": (
|
|
"STRING",
|
|
{
|
|
"default": "1, 1, 1, 1, 1, 1",
|
|
"tooltip": "Comma-separated list of stg rescale values. Should be same length as sigmas list.",
|
|
},
|
|
),
|
|
"stg_layers_indices": (
|
|
"STRING",
|
|
{
|
|
"default": "[29], [29], [29], [29], [29], [29]",
|
|
"tooltip": "Comma-separated list of list of layer indices. Should be same length as sigmas list.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"preset": (
|
|
"STG_ADVANCED_PRESET",
|
|
{
|
|
"tooltip": "Preset resolution and frame count. Custom allows manual input.",
|
|
},
|
|
),
|
|
"apply_apg": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": False,
|
|
"tooltip": "If true, applies the APG (Adaptive Projections Guidance) to the STG.",
|
|
},
|
|
),
|
|
"apg_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},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("GUIDER",)
|
|
|
|
FUNCTION = "get_guider"
|
|
CATEGORY = "lightricks/LTXV"
|
|
DESCRIPTION = """
|
|
The Advanced STG Guider implements sophisticated techniques for controlling the denoising process:
|
|
|
|
It creates a dynamic mapping from scheduler-defined sigma values to CFG and STG (Spatio-Temporal Skip Guidance [1]) parameters.
|
|
This approach establishes distinct sigma value ranges that operate independently of step numbers, allowing precise control over:
|
|
• CFG scale
|
|
• STG scale and rescale factors
|
|
• STG attention layer skipping patterns
|
|
|
|
The guider also supports:
|
|
• CFG-Zero* [2] rescaling, which dynamically adjusts negative predictions based on the dot product between positive and negative signals
|
|
• Threshold-based noise prediction zeroing for steps with sigma values exceeding a specified threshold
|
|
|
|
For example if the sigma ranges are defined as [1.0, 0.9, 0.85, 0.6] and the CFG values are defined as [4, 3, 2, 1] and STG scale values
|
|
are defined as [2, 2, 2, 1] and STG rescale values are defined as [1, 1, 1, 1] and STG layers indices are defined as [[14, 17], [14, 16], [14], [14]], then the guider will:
|
|
- use CFG=4, STG scale=2, STG rescale=1 and STG layers indices = [14, 17] for sigma in the range (0.9, 1.0]
|
|
- use CFG=3, STG scale=2, STG rescale=1 and STG layers indices = [14, 16] for sigma in the range (0.85, 0.9]
|
|
- use CFG=2, STG scale=2, STG rescale=1 and STG layers indices = [14] for sigma in the range (0.8, 0.85]
|
|
- use CFG=1, STG scale=1, STG rescale=1 and STG layers indices = [14] for sigma in the range (0.6, 0.8]
|
|
|
|
The guider will use the same parameters for the same sigma values, regardless of the step number.
|
|
|
|
References:
|
|
[1] https://arxiv.org/abs/2411.18664
|
|
[2] https://arxiv.org/abs/2503.18886
|
|
|
|
"""
|
|
|
|
@classmethod
|
|
def parse_stg_layers_indices(cls, stg_layers_indices: str) -> List[List[int]]:
|
|
# First split by "], " to separate the lists, but preserve the closing bracket
|
|
lists = [s + "]" for s in stg_layers_indices.split("],")[:-1]]
|
|
# Add the last list which already has its closing bracket
|
|
if stg_layers_indices.strip():
|
|
lists.append(stg_layers_indices.split("],")[-1])
|
|
|
|
result = []
|
|
for s in lists:
|
|
s = s.strip()
|
|
if s == "[]": # Empty list case
|
|
result.append([])
|
|
else:
|
|
# Remove brackets and whitespace
|
|
s = s.strip("[]").strip()
|
|
if not s: # Handle case like "[]" after split
|
|
result.append([])
|
|
else:
|
|
# Split by comma and convert to integers
|
|
numbers = [int(n.strip()) for n in s.split(",") if n.strip()]
|
|
result.append(numbers)
|
|
|
|
return result
|
|
|
|
def get_guider(
|
|
self,
|
|
model,
|
|
positive,
|
|
negative,
|
|
skip_steps_sigma_threshold,
|
|
cfg_star_rescale,
|
|
sigmas,
|
|
cfg_values,
|
|
stg_scale_values,
|
|
stg_rescale_values,
|
|
stg_layers_indices,
|
|
preset=None,
|
|
apply_apg=False,
|
|
apg_cfg_scale=1.0,
|
|
eta=1.0,
|
|
norm_threshold=0.0,
|
|
):
|
|
if preset and preset != "Custom":
|
|
preset_data = next(
|
|
(item for item in STG_ADVANCED_PRESETS if item["name"] == preset), None
|
|
)
|
|
if preset_data:
|
|
skip_steps_sigma_threshold = preset_data["skip_steps_sigma_threshold"]
|
|
cfg_star_rescale = preset_data["cfg_star_rescale"]
|
|
sigma_list = preset_data["sigmas"]
|
|
cfg_list = preset_data["cfg_values"]
|
|
stg_scale_list = preset_data["stg_scale_values"]
|
|
stg_rescale_list = preset_data["stg_rescale_values"]
|
|
stg_layers_indices_list = preset_data["stg_layers_indices"]
|
|
else:
|
|
raise ValueError(f"Preset {preset} not found in the presets list.")
|
|
|
|
else:
|
|
sigma_list = [float(s.strip()) for s in sigmas.split(",")]
|
|
cfg_list = [float(c.strip()) for c in cfg_values.split(",")]
|
|
stg_scale_list = [float(s.strip()) for s in stg_scale_values.split(",")]
|
|
stg_rescale_list = [float(s.strip()) for s in stg_rescale_values.split(",")]
|
|
stg_layers_indices_list = self.parse_stg_layers_indices(stg_layers_indices)
|
|
print("Using preset: ", preset)
|
|
print("Skip steps sigma threshold: ", skip_steps_sigma_threshold)
|
|
print("Cfg star rescale: ", cfg_star_rescale)
|
|
print("Sigma list: ", sigma_list)
|
|
print("Cfg list: ", cfg_list)
|
|
print("Stg scale list: ", stg_scale_list)
|
|
print("Stg rescale list: ", stg_rescale_list)
|
|
print("Stg layers indices list: ", stg_layers_indices_list)
|
|
|
|
guider = STGGuiderAdvanced(
|
|
model,
|
|
sigma_list,
|
|
cfg_list,
|
|
stg_scale_list,
|
|
stg_rescale_list,
|
|
stg_layers_indices_list,
|
|
skip_steps_sigma_threshold,
|
|
cfg_star_rescale,
|
|
apply_apg,
|
|
apg_cfg_scale=apg_cfg_scale,
|
|
eta=eta,
|
|
norm_threshold=norm_threshold,
|
|
)
|
|
guider.set_conds(positive, negative)
|
|
guider.raw_conds = (positive, negative)
|
|
return (guider,)
|
|
|
|
|
|
@comfy_node(name="STGAdvancedPresets")
|
|
class STGAdvancedPresetsNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"preset": (
|
|
STG_ADVANCED_PRESET_NAMES,
|
|
{
|
|
"default": "13b Balanced",
|
|
"tooltip": "Preset resolution and frame count. Custom allows manual input.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STG_ADVANCED_PRESET",)
|
|
|
|
FUNCTION = "get_preset"
|
|
CATEGORY = "lightricks/LTXV"
|
|
|
|
def get_preset(self, preset=None):
|
|
return (preset,)
|