276 lines
9.2 KiB
Python
276 lines
9.2 KiB
Python
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,)
|