From deaed13f0dabc981c042c3ca15b12438f57dab6e Mon Sep 17 00:00:00 2001 From: Sugam Date: Tue, 24 Feb 2026 14:57:03 +0545 Subject: [PATCH] Update nodes --- flowmatch_nodes.py | 234 +++++++++++++++++++++++++++++++++++---------- 1 file changed, 183 insertions(+), 51 deletions(-) diff --git a/flowmatch_nodes.py b/flowmatch_nodes.py index 7893026..edc7b6c 100644 --- a/flowmatch_nodes.py +++ b/flowmatch_nodes.py @@ -1,74 +1,206 @@ -import torch import math +import comfy.model_sampling +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, + }, +} + + +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"] + + class ModelSamplingDiTFlow: """ - Advanced Flow Matching Scaling for DiT models (Flux, Z-Image, Qwen). - Replicates the exact logic from lora-scripts for seamless ComfyUI integration. + ai-toolkit style FlowMatch model patch for DiT models. """ + @classmethod - def INPUT_TYPES(s): + def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), - "model_type": (["flux", "z-image", "qwen"], {"default": "flux"}), - "shift": ("FLOAT", {"default": 1.15, "min": 0.0, "max": 10.0, "step": 0.01}), + "model_type": (["auto", "flux", "z-image", "qwen"], {"default": "auto"}), + "shift": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.01}), "resolution_aware": ("BOOLEAN", {"default": True}), - }, - "optional": { "width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8}), "height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8}), - } + }, } RETURN_TYPES = ("MODEL",) FUNCTION = "apply_dit_flow" - CATEGORY = "sampling/custom_sampling" + CATEGORY = "advanced/model" - def apply_dit_flow(self, model, model_type, shift, resolution_aware, width=1024, height=1024): + def apply_dit_flow(self, 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() - def get_lin_function(x1, y1, x2, y2): - m_val = (y2 - y1) / (x2 - x1) - b_val = y1 - m_val * x1 - return lambda x: m_val * x + b_val + 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} - def time_shift(mu, t): - return torch.exp(torch.tensor(mu)) / (torch.exp(torch.tensor(mu)) + (1 / t - 1) ** 1.0) + class ModelSamplingAdvanced(sampling_base, comfy.model_sampling.CONST): + pass - def patch_sampling(filter_out, sigma_data): - t = sigma_data["timesteps"] - if t.max() > 10.0: - t = t / 1000.0 - - if resolution_aware: - # Calculate sequence length based on packed latents (typical for Flux/DiT) - # Flux uses 16x16 patch packing, so seq_len = (w/16) * (h/16) - # We use the provided width/height to calculate the mu shift - seq_len = (width // 16) * (height // 16) - - if model_type == "flux": - # Flux standard: linear interpolation between 0.5 and 1.15 - mu = get_lin_function(256, 0.5, 4096, 1.15)(seq_len) - elif model_type == "qwen": - # Qwen/Lumina often uses a slightly steeper shift for LLM-based embeddings - mu = get_lin_function(256, 0.5, 4096, shift)(seq_len) - else: # z-image - mu = shift - else: - mu = shift - - # Calculate new sigmas - sigmas = time_shift(mu, t) - return sigmas - - # Apply patches - m.add_object_patch("model_sampling", patch_sampling) - - # All these models are Flow-based and require v_prediction - m.add_object_patch("parameterization", "v_prediction") - + model_sampling = ModelSamplingAdvanced(model.model.model_config) + model_sampling.set_parameters(**sampling_kwargs) + m.add_object_patch("model_sampling", model_sampling) return (m,) -NODE_CLASS_MAPPINGS = { "ModelSamplingDiTFlow": ModelSamplingDiTFlow } -NODE_DISPLAY_NAME_MAPPINGS = { "ModelSamplingDiTFlow": "DiT FlowMatch (Flux/Z-Image/Qwen)" } + +class AIToolkitFlowMatchScheduler: + """ + ai-toolkit style SIGMAS output for SamplerCustomAdvanced. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_type": (["auto", "flux", "z-image", "qwen"], {"default": "auto"}), + "steps": ("INT", {"default": 25, "min": 1, "max": 10000}), + "width": ("INT", {"default": 1024, "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": True}), + "force_aitk_timesteps": ("BOOLEAN", {"default": True}), + "append_zero": ("BOOLEAN", {"default": True}), + }, + "optional": { + "model": ("MODEL",), + }, + } + + RETURN_TYPES = ("SIGMAS",) + FUNCTION = "get_sigmas" + CATEGORY = "sampling/custom_sampling/schedulers" + + def get_sigmas( + self, + model_type, + steps, + width, + height, + shift, + resolution_aware, + force_aitk_timesteps, + append_zero, + model=None, + ): + if model is not None: + resolved_model_type = _resolve_model_type(model, model_type) + else: + resolved_model_type = "flux" if model_type == "auto" else model_type + + config = AITK_FLOWMATCH_CONFIG[resolved_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) + if append_zero: + sigmas = torch.cat([sigmas, sigmas.new_zeros((1,))]) + + return (sigmas,) + + +NODE_CLASS_MAPPINGS = { + "ModelSamplingDiTFlow": ModelSamplingDiTFlow, + "AIToolkitFlowMatchScheduler": AIToolkitFlowMatchScheduler, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ModelSamplingDiTFlow": "DiT FlowMatch Model (ai-toolkit style)", + "AIToolkitFlowMatchScheduler": "FlowMatch Sigmas (ai-toolkit style)", +}