437 lines
14 KiB
Python
437 lines
14 KiB
Python
import math
|
|
import os
|
|
import sys
|
|
|
|
import comfy.model_management
|
|
import comfy.model_sampling
|
|
import comfy.nested_tensor
|
|
import comfy.sample
|
|
import comfy.samplers
|
|
import comfy.utils
|
|
import latent_preview
|
|
import torch
|
|
|
|
|
|
AITK_FLOWMATCH_CONFIG = {
|
|
"flux": {
|
|
"base_image_seq_len": 256,
|
|
"base_shift": 0.5,
|
|
"max_image_seq_len": 4096,
|
|
"max_shift": 1.15,
|
|
"shift": 3.0,
|
|
"use_dynamic_shifting": True,
|
|
},
|
|
"qwen": {
|
|
"base_image_seq_len": 256,
|
|
"base_shift": 0.5,
|
|
"max_image_seq_len": 8192,
|
|
"max_shift": 0.9,
|
|
"shift": 1.0,
|
|
"use_dynamic_shifting": True,
|
|
},
|
|
"z-image": {
|
|
"shift": 3.0,
|
|
"use_dynamic_shifting": False,
|
|
},
|
|
}
|
|
|
|
AITK_SCHEDULER_CONFIG = {
|
|
"flux": {
|
|
"base_image_seq_len": 256,
|
|
"base_shift": 0.5,
|
|
"max_image_seq_len": 4096,
|
|
"max_shift": 1.15,
|
|
"num_train_timesteps": 1000,
|
|
"shift": 3.0,
|
|
"use_dynamic_shifting": True,
|
|
},
|
|
"qwen": {
|
|
"base_image_seq_len": 256,
|
|
"base_shift": 0.5,
|
|
"invert_sigmas": False,
|
|
"max_image_seq_len": 8192,
|
|
"max_shift": 0.9,
|
|
"num_train_timesteps": 1000,
|
|
"shift": 1.0,
|
|
"shift_terminal": 0.02,
|
|
"time_shift_type": "exponential",
|
|
"use_beta_sigmas": False,
|
|
"use_dynamic_shifting": True,
|
|
"use_exponential_sigmas": False,
|
|
"use_karras_sigmas": False,
|
|
},
|
|
"z-image": {
|
|
"num_train_timesteps": 1000,
|
|
"shift": 3.0,
|
|
"use_dynamic_shifting": False,
|
|
},
|
|
}
|
|
|
|
_AITK_SCHEDULER_BACKEND = None
|
|
_AITK_SCHEDULER_BACKEND_FAILED = False
|
|
|
|
|
|
def _calculate_shift(image_seq_len, base_seq_len, max_seq_len, base_shift, max_shift):
|
|
m_val = (max_shift - base_shift) / (max_seq_len - base_seq_len)
|
|
b_val = base_shift - m_val * base_seq_len
|
|
return image_seq_len * m_val + b_val
|
|
|
|
|
|
def _dynamic_time_shift(mu, t):
|
|
t = torch.clamp(t, min=1e-6, max=1.0)
|
|
exp_mu = math.exp(float(mu))
|
|
return exp_mu / (exp_mu + (1.0 / t - 1.0))
|
|
|
|
|
|
def _static_time_shift(shift_value, t):
|
|
t = torch.clamp(t, min=1e-6, max=1.0)
|
|
return (shift_value * t) / (1.0 + (shift_value - 1.0) * t)
|
|
|
|
|
|
def _resolve_model_type(model, model_type):
|
|
if model_type != "auto":
|
|
return model_type
|
|
|
|
try:
|
|
image_model = str(model.model.model_config.unet_config.get("image_model", "")).lower()
|
|
except Exception:
|
|
image_model = ""
|
|
|
|
if image_model == "qwen_image":
|
|
return "qwen"
|
|
if image_model == "lumina2":
|
|
return "z-image"
|
|
if "flux" in image_model:
|
|
return "flux"
|
|
return "flux"
|
|
|
|
|
|
def _get_effective_shift(model_type, shift, resolution_aware, width, height):
|
|
config = AITK_FLOWMATCH_CONFIG[model_type]
|
|
if resolution_aware and config.get("use_dynamic_shifting", False):
|
|
seq_len = (width // 16) * (height // 16)
|
|
return _calculate_shift(
|
|
seq_len,
|
|
config["base_image_seq_len"],
|
|
config["max_image_seq_len"],
|
|
config["base_shift"],
|
|
config["max_shift"],
|
|
)
|
|
if shift > 0.0:
|
|
return shift
|
|
return config["shift"]
|
|
|
|
|
|
def _apply_dit_flow_patch(model, model_type, shift, resolution_aware, width, height):
|
|
resolved_model_type = _resolve_model_type(model, model_type)
|
|
effective_shift = _get_effective_shift(resolved_model_type, shift, resolution_aware, width, height)
|
|
m = model.clone()
|
|
|
|
if resolved_model_type in ("flux", "qwen") and resolution_aware:
|
|
sampling_base = comfy.model_sampling.ModelSamplingFlux
|
|
sampling_kwargs = {"shift": effective_shift}
|
|
else:
|
|
sampling_base = comfy.model_sampling.ModelSamplingDiscreteFlow
|
|
sampling_kwargs = {"shift": effective_shift, "multiplier": 1.0}
|
|
|
|
class ModelSamplingAdvanced(sampling_base, comfy.model_sampling.CONST):
|
|
pass
|
|
|
|
model_sampling = ModelSamplingAdvanced(model.model.model_config)
|
|
model_sampling.set_parameters(**sampling_kwargs)
|
|
m.add_object_patch("model_sampling", model_sampling)
|
|
return m
|
|
|
|
|
|
def _build_sigmas_formula(
|
|
model_type,
|
|
steps,
|
|
width,
|
|
height,
|
|
shift,
|
|
resolution_aware,
|
|
force_aitk_timesteps,
|
|
append_zero,
|
|
):
|
|
config = AITK_FLOWMATCH_CONFIG[model_type]
|
|
steps = max(int(steps), 1)
|
|
|
|
if force_aitk_timesteps:
|
|
t = torch.linspace(1.0, 1.0 / float(steps), steps, dtype=torch.float32)
|
|
else:
|
|
t = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float32)[:-1]
|
|
|
|
if resolution_aware and config.get("use_dynamic_shifting", False):
|
|
seq_len = (width // 16) * (height // 16)
|
|
mu = _calculate_shift(
|
|
seq_len,
|
|
config["base_image_seq_len"],
|
|
config["max_image_seq_len"],
|
|
config["base_shift"],
|
|
config["max_shift"],
|
|
)
|
|
sigmas = _dynamic_time_shift(mu, t)
|
|
else:
|
|
effective_shift = shift if shift > 0.0 else config["shift"]
|
|
sigmas = _static_time_shift(effective_shift, t)
|
|
|
|
sigmas = sigmas.to(dtype=torch.float32, device="cpu")
|
|
if append_zero:
|
|
sigmas = torch.cat([sigmas, sigmas.new_zeros((1,))])
|
|
return sigmas
|
|
|
|
|
|
def _load_aitk_scheduler_backend():
|
|
global _AITK_SCHEDULER_BACKEND, _AITK_SCHEDULER_BACKEND_FAILED
|
|
if _AITK_SCHEDULER_BACKEND is not None:
|
|
return _AITK_SCHEDULER_BACKEND
|
|
if _AITK_SCHEDULER_BACKEND_FAILED:
|
|
return None
|
|
|
|
try:
|
|
repo_dir = os.path.dirname(os.path.abspath(__file__))
|
|
candidates = (
|
|
os.path.join(repo_dir, "ai-toolkit"),
|
|
os.path.join(os.path.dirname(repo_dir), "ai-toolkit"),
|
|
)
|
|
for candidate in candidates:
|
|
if os.path.isdir(candidate) and candidate not in sys.path:
|
|
sys.path.insert(0, candidate)
|
|
|
|
from toolkit.samplers.custom_flowmatch_sampler import ( # pylint: disable=import-outside-toplevel
|
|
CustomFlowMatchEulerDiscreteScheduler,
|
|
calculate_shift as aitk_calculate_shift,
|
|
)
|
|
except Exception:
|
|
_AITK_SCHEDULER_BACKEND_FAILED = True
|
|
return None
|
|
|
|
_AITK_SCHEDULER_BACKEND = (CustomFlowMatchEulerDiscreteScheduler, aitk_calculate_shift)
|
|
return _AITK_SCHEDULER_BACKEND
|
|
|
|
|
|
def _build_sigmas_exact(
|
|
model_type,
|
|
steps,
|
|
width,
|
|
height,
|
|
shift,
|
|
resolution_aware,
|
|
force_aitk_timesteps,
|
|
append_zero,
|
|
):
|
|
backend = _load_aitk_scheduler_backend()
|
|
if backend is None:
|
|
return None
|
|
|
|
try:
|
|
import inspect # pylint: disable=import-outside-toplevel
|
|
import numpy as np # pylint: disable=import-outside-toplevel
|
|
|
|
scheduler_cls, aitk_calculate_shift = backend
|
|
scheduler = scheduler_cls(**AITK_SCHEDULER_CONFIG[model_type])
|
|
steps = max(int(steps), 1)
|
|
params = set(inspect.signature(scheduler.set_timesteps).parameters)
|
|
|
|
kwargs = {}
|
|
if "device" in params:
|
|
kwargs["device"] = "cpu"
|
|
if "num_inference_steps" in params:
|
|
kwargs["num_inference_steps"] = steps
|
|
if "sigmas" in params and force_aitk_timesteps:
|
|
kwargs["sigmas"] = np.linspace(1.0, 1.0 / float(steps), steps, dtype=np.float32)
|
|
if "mu" in params and resolution_aware and AITK_SCHEDULER_CONFIG[model_type].get("use_dynamic_shifting", False):
|
|
seq_len = (width // 16) * (height // 16)
|
|
kwargs["mu"] = aitk_calculate_shift(
|
|
seq_len,
|
|
AITK_SCHEDULER_CONFIG[model_type].get("base_image_seq_len", 256),
|
|
AITK_SCHEDULER_CONFIG[model_type].get("max_image_seq_len", 4096),
|
|
AITK_SCHEDULER_CONFIG[model_type].get("base_shift", 0.5),
|
|
AITK_SCHEDULER_CONFIG[model_type].get("max_shift", 1.15),
|
|
)
|
|
if "shift" in params and shift > 0.0 and not AITK_SCHEDULER_CONFIG[model_type].get("use_dynamic_shifting", False):
|
|
kwargs["shift"] = float(shift)
|
|
|
|
scheduler.set_timesteps(**kwargs)
|
|
sigmas = scheduler.sigmas.detach().to(dtype=torch.float32, device="cpu")
|
|
|
|
if append_zero:
|
|
if sigmas.numel() == 0 or not torch.isclose(sigmas[-1], sigmas.new_zeros(()), atol=1e-6):
|
|
sigmas = torch.cat([sigmas, sigmas.new_zeros((1,))])
|
|
else:
|
|
if sigmas.numel() > 0 and (
|
|
torch.isclose(sigmas[-1], sigmas.new_zeros(()), atol=1e-6)
|
|
or torch.isclose(sigmas[-1], sigmas.new_ones(()), atol=1e-6)
|
|
):
|
|
sigmas = sigmas[:-1]
|
|
return sigmas
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _build_flowmatch_sigmas(
|
|
model_type,
|
|
steps,
|
|
width,
|
|
height,
|
|
shift,
|
|
resolution_aware,
|
|
force_aitk_timesteps,
|
|
append_zero,
|
|
):
|
|
sigmas = _build_sigmas_exact(
|
|
model_type,
|
|
steps,
|
|
width,
|
|
height,
|
|
shift,
|
|
resolution_aware,
|
|
force_aitk_timesteps,
|
|
append_zero,
|
|
)
|
|
if sigmas is not None:
|
|
return sigmas
|
|
|
|
return _build_sigmas_formula(
|
|
model_type,
|
|
steps,
|
|
width,
|
|
height,
|
|
shift,
|
|
resolution_aware,
|
|
force_aitk_timesteps,
|
|
append_zero,
|
|
)
|
|
|
|
|
|
class AIToolkitFlowMatchSampler:
|
|
"""
|
|
All-in-one ai-toolkit style flowmatch sampler node.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent_image": ("LATENT",),
|
|
"noise_seed": ("INT", {"default": 42, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "control_after_generate": True}),
|
|
"add_noise": (["enable", "disable"], {"default": "enable"}),
|
|
"steps": ("INT", {"default": 8, "min": 1, "max": 10000}),
|
|
"guidance_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1}),
|
|
"sampler_name": (["euler", "res_multistep"], {"default": "euler"}),
|
|
"model_type": (["auto", "flux", "z-image", "qwen"], {"default": "z-image"}),
|
|
"width": ("INT", {"default": 768, "min": 64, "max": 8192, "step": 8}),
|
|
"height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8}),
|
|
"shift": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
|
"resolution_aware": ("BOOLEAN", {"default": False}),
|
|
"force_aitk_timesteps": ("BOOLEAN", {"default": True}),
|
|
"append_zero": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "LATENT", "SIGMAS", "MODEL")
|
|
RETURN_NAMES = ("output", "denoised_output", "sigmas", "patched_model")
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling/custom_sampling"
|
|
|
|
def sample(
|
|
self,
|
|
model,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
noise_seed,
|
|
add_noise,
|
|
steps,
|
|
guidance_scale,
|
|
sampler_name,
|
|
model_type,
|
|
width,
|
|
height,
|
|
shift,
|
|
resolution_aware,
|
|
force_aitk_timesteps,
|
|
append_zero,
|
|
):
|
|
patched_model = _apply_dit_flow_patch(model, model_type, shift, resolution_aware, width, height)
|
|
resolved_model_type = _resolve_model_type(patched_model, model_type)
|
|
sigmas = _build_flowmatch_sigmas(
|
|
resolved_model_type,
|
|
steps,
|
|
width,
|
|
height,
|
|
shift,
|
|
resolution_aware,
|
|
force_aitk_timesteps,
|
|
append_zero,
|
|
)
|
|
|
|
latent = latent_image.copy()
|
|
latent_samples = latent["samples"]
|
|
latent_samples = comfy.sample.fix_empty_latent_channels(
|
|
patched_model, latent_samples, latent.get("downscale_ratio_spacial", None)
|
|
)
|
|
latent["samples"] = latent_samples
|
|
|
|
if add_noise == "disable":
|
|
noise = torch.zeros(
|
|
latent_samples.size(),
|
|
dtype=latent_samples.dtype,
|
|
layout=latent_samples.layout,
|
|
device="cpu",
|
|
)
|
|
else:
|
|
batch_inds = latent["batch_index"] if "batch_index" in latent else None
|
|
noise = comfy.sample.prepare_noise(latent_samples, noise_seed, batch_inds)
|
|
|
|
noise_mask = latent.get("noise_mask", None)
|
|
|
|
guider = comfy.samplers.CFGGuider(patched_model)
|
|
guider.set_conds(positive, negative)
|
|
guider.set_cfg(guidance_scale)
|
|
sampler = comfy.samplers.ksampler(sampler_name)
|
|
|
|
x0_output = {}
|
|
callback = latent_preview.prepare_callback(patched_model, max(sigmas.shape[-1] - 1, 0), x0_output)
|
|
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
|
|
|
samples = guider.sample(
|
|
noise,
|
|
latent_samples,
|
|
sampler,
|
|
sigmas,
|
|
denoise_mask=noise_mask,
|
|
callback=callback,
|
|
disable_pbar=disable_pbar,
|
|
seed=noise_seed,
|
|
)
|
|
samples = samples.to(comfy.model_management.intermediate_device())
|
|
|
|
out = latent.copy()
|
|
out.pop("downscale_ratio_spacial", None)
|
|
out["samples"] = samples
|
|
|
|
if "x0" in x0_output:
|
|
x0_out = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu())
|
|
if samples.is_nested:
|
|
latent_shapes = [x.shape for x in samples.unbind()]
|
|
x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes))
|
|
out_denoised = latent.copy()
|
|
out_denoised["samples"] = x0_out
|
|
else:
|
|
out_denoised = out
|
|
|
|
return (out, out_denoised, sigmas, patched_model)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AIToolkitFlowMatchSampler": AIToolkitFlowMatchSampler,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AIToolkitFlowMatchSampler": "FlowMatch Sampler (ai-toolkit exact)",
|
|
}
|