# ComfyUI custom node for FlowMatch Euler Scheduler # # This node creates a FlowMatchEulerDiscreteScheduler with configurable parameters # so it can be used with compatible sampler nodes. # # Also registers the scheduler in ComfyUI's scheduler list with default config. # # Place this file into: ComfyUI/custom_nodes/ # Then restart ComfyUI. It will show up as "FlowMatch Euler Discrete Scheduler (Custom)" import math import torch import numpy as np # <-- Required for robust slicing of PyTorch tensors try: from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler except ImportError as e: print("=" * 80) print("ERROR: Failed to import FlowMatchEulerDiscreteScheduler from diffusers") print("Please ensure dependencies are installed by running:") print(" pip install -r requirements.txt") print("=" * 80) raise ImportError( "FlowMatchEulerDiscreteScheduler not found. " "Please install dependencies from requirements.txt" ) from e try: from diffusers import VQDiffusionScheduler except ImportError: VQDiffusionScheduler = None print("[FlowMatch Scheduler] Warning: VQDiffusionScheduler not found in diffusers.") from comfy.samplers import SchedulerHandler, SCHEDULER_HANDLERS, SCHEDULER_NAMES # Import Nunchaku compatibility patches (auto-applies on import) try: from . import nunchaku_compat except Exception as e: print(f"[FlowMatch Scheduler] Warning: Could not load Nunchaku compatibility: {e}") # Default config for registering in ComfyUI default_config = { "base_image_seq_len": 256, "base_shift": math.log(3), "invert_sigmas": False, "max_image_seq_len": 8192, "max_shift": math.log(3), "num_train_timesteps": 1000, "shift": 1.0, "shift_terminal": None, "stochastic_sampling": False, "time_shift_type": "exponential", "use_beta_sigmas": False, "use_dynamic_shifting": True, "use_exponential_sigmas": False, "use_karras_sigmas": False, } def flow_match_euler_scheduler_handler(model_sampling, steps): scheduler = FlowMatchEulerDiscreteScheduler.from_config(default_config) scheduler.set_timesteps(steps, device=model_sampling.device if hasattr(model_sampling, 'device') else 'cpu', mu=0.0) sigmas = scheduler.sigmas return sigmas def vq_diffusion_scheduler_handler(model_sampling, steps): if VQDiffusionScheduler is None: raise ImportError("VQDiffusionScheduler is not available.") # VQDiffusionScheduler requires num_vec_classes. print("[FlowMatch Scheduler] WARNING: VQDiffusionScheduler is for discrete models (VQ-Diffusion).") print("It does not produce 'sigmas' for continuous diffusion.") print("Returning dummy linear sigmas to prevent crash, but sampling will likely fail with standard models.") # Dummy initialization # scheduler = VQDiffusionScheduler(num_vec_classes=4096, num_train_timesteps=1000) # Return dummy sigmas sigmas = torch.linspace(1.0, 0.0, steps + 1) if hasattr(model_sampling, 'device'): sigmas = sigmas.to(model_sampling.device) return sigmas # Register the schedulers in ComfyUI if "FlowMatchEulerDiscreteScheduler" not in SCHEDULER_HANDLERS: handler = SchedulerHandler(handler=flow_match_euler_scheduler_handler, use_ms=True) SCHEDULER_HANDLERS["FlowMatchEulerDiscreteScheduler"] = handler SCHEDULER_NAMES.append("FlowMatchEulerDiscreteScheduler") # Explicitly add to KSampler.SCHEDULERS to ensure compatibility with nodes # that might replace the list object (like RES4LYF) try: from comfy.samplers import KSampler if "FlowMatchEulerDiscreteScheduler" not in KSampler.SCHEDULERS: KSampler.SCHEDULERS.append("FlowMatchEulerDiscreteScheduler") except ImportError: pass # if "VQDiffusionScheduler" not in SCHEDULER_HANDLERS: # SCHEDULER_HANDLERS["VQDiffusionScheduler"] = SchedulerHandler(handler=vq_diffusion_scheduler_handler, use_ms=True) # SCHEDULER_NAMES.append("VQDiffusionScheduler") # try: # from comfy.samplers import KSampler # if "VQDiffusionScheduler" not in KSampler.SCHEDULERS: # KSampler.SCHEDULERS.append("VQDiffusionScheduler") # except ImportError: # pass class FlowMatchEulerSchedulerNode: @classmethod def INPUT_TYPES(cls): return { "required": { "steps": ("INT", { "default": 9, "min": 1, "max": 10000, "tooltip": "Total number of diffusion steps to generate the full sigma schedule." }), "start_at_step": ("INT", { "default": 0, "min": 0, "max": 10000, "tooltip": "The starting step (index) of the sigma schedule to use. Set to 0 to start at the beginning (first step)." }), "end_at_step": ("INT", { "default": 9999, "min": 0, "max": 10000, "tooltip": "The ending step (index) of the sigma schedule to use. Set higher than 'steps' to use all steps." }), "base_image_seq_len": ("INT", { "default": 256, "tooltip": "Base sequence length for dynamic shifting. Should match model's training resolution (e.g., 256 for 512x512 images)." }), "base_shift": ("FLOAT", { "default": 0.5, "step": 0.01, "tooltip": "Stabilizes generation. Higher values = more consistent/predictable outputs. Z-Image-Turbo uses default 0.5." }), "invert_sigmas": (["disable", "enable"], { "default": "disable", "tooltip": "Reverses the sigma schedule. Keep disabled unless experimenting with advanced techniques." }), "max_image_seq_len": ("INT", { "default": 8192, "tooltip": "Maximum sequence length for dynamic shifting. Affects how the scheduler adapts to large images." }), "max_shift": ("FLOAT", { "default": 1.15, "step": 0.01, "tooltip": "Maximum variation allowed. Higher = more exaggerated/stylized results. Z-Image-Turbo uses default 1.15." }), "num_train_timesteps": ("INT", { "default": 1000, "tooltip": "Timesteps the model was trained with. Should match your model's config (typically 1000)." }), "shift": ("FLOAT", { "default": 3.0, "step": 0.01, "tooltip": "Global timestep schedule shift. Z-Image-Turbo uses 3.0 for optimal performance with the Turbo model." }), "shift_terminal": ("FLOAT", { "default": 0.0, "step": 0.01, "tooltip": "End value for shifted schedule. Set to 0.0 to disable. Advanced parameter for timestep schedule control." }), "stochastic_sampling": (["disable", "enable"], { "default": "disable", "tooltip": "Adds controlled randomness to each step. Enable for more varied outputs (similar to ancestral samplers)." }), "time_shift_type": (["exponential", "linear"], { "default": "exponential", "tooltip": "Method for resolution-dependent shifting. Use 'exponential' for most cases, 'linear' for experiments." }), "use_beta_sigmas": (["disable", "enable"], { "default": "disable", "tooltip": "Uses beta distribution for sigmas. Experimental alternative noise schedule." }), "use_dynamic_shifting": (["disable", "enable"], { "default": "disable", "tooltip": "Auto-adjusts timesteps based on image resolution. Z-Image-Turbo disables this for consistent Turbo performance." }), "use_exponential_sigmas": (["disable", "enable"], { "default": "disable", "tooltip": "Uses exponential sigma spacing. Try enabling for different noise distribution characteristics." }), "use_karras_sigmas": (["disable", "enable"], { "default": "disable", "tooltip": "Uses Karras noise schedule for smoother results. Similar to DPM++ samplers, often improves quality." }), "device": (["auto", "cuda", "cpu"], { "default": "auto", "tooltip": "Device for sigma computation. 'auto' detects GPU if available, otherwise CPU. Using GPU avoids CPU->GPU transfers." }), } } RETURN_TYPES = ("SIGMAS",) RETURN_NAMES = ("sigmas",) FUNCTION = "create" CATEGORY = "sampling/schedulers" DESCRIPTION = "FlowMatch Euler Discrete Scheduler with full parameter control and ability to trim the schedule (start_at_step/end_at_step)." def create( self, steps, start_at_step, end_at_step, base_image_seq_len, base_shift, invert_sigmas, max_image_seq_len, max_shift, num_train_timesteps, shift, shift_terminal, stochastic_sampling, time_shift_type, use_beta_sigmas, use_dynamic_shifting, use_exponential_sigmas, use_karras_sigmas, device="auto", ): # Convert string combo values to boolean config = { "base_image_seq_len": base_image_seq_len, "base_shift": base_shift, "invert_sigmas": invert_sigmas == "enable", "max_image_seq_len": max_image_seq_len, "max_shift": max_shift, "num_train_timesteps": num_train_timesteps, "shift": shift, "shift_terminal": shift_terminal if shift_terminal != 0.0 else None, "stochastic_sampling": stochastic_sampling == "enable", "time_shift_type": time_shift_type, "use_beta_sigmas": use_beta_sigmas == "enable", "use_dynamic_shifting": use_dynamic_shifting == "enable", "use_exponential_sigmas": use_exponential_sigmas == "enable", "use_karras_sigmas": use_karras_sigmas == "enable", } scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) # 1. Generate the full sigma schedule # Determine device to use for sigma computation if device == "auto": # Auto-detect: use CUDA if available, otherwise CPU target_device = "cuda" if torch.cuda.is_available() else "cpu" print(f"[FlowMatch Scheduler] Auto-detected device: {target_device.upper()}") else: target_device = device print(f"[FlowMatch Scheduler] Using manually specified device: {target_device.upper()}") # Set timesteps and get sigmas for the specified number of steps # Using the model's device avoids unnecessary CPU->GPU transfers during sampling scheduler.set_timesteps(steps, device=target_device, mu=0.0) sigmas = scheduler.sigmas # 2. Apply start_at_step and end_at_step (Slicing the sigmas tensor) # Determine the exclusive end index for the slice # end_at_step is the step index (e.g., 5). We use 5+1=6 for the slice end index. end_index = min(end_at_step + 1, len(sigmas)) # Slice the tensor: [start:end] sigmas_sliced = sigmas[start_at_step:end_index] # Check for empty schedule resulting from slicing if sigmas_sliced.numel() == 0: print("Warning: start_at_step/end_at_step resulted in an empty sigma schedule. Using full schedule as fallback.") sigmas_sliced = sigmas return (sigmas_sliced,) class VQDiffusionSchedulerNode: @classmethod def INPUT_TYPES(cls): return { "required": { "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), "num_vec_classes": ("INT", {"default": 4096, "min": 1, "max": 65536, "tooltip": "Number of vector classes for VQ model."}), "num_train_timesteps": ("INT", {"default": 1000}), } } RETURN_TYPES = ("SIGMAS",) RETURN_NAMES = ("sigmas",) FUNCTION = "create" CATEGORY = "sampling/schedulers" DESCRIPTION = "VQ Diffusion Scheduler (Experimental). For VQ-Diffusion models. Returns dummy sigmas for compatibility." def create(self, steps, num_vec_classes, num_train_timesteps): if VQDiffusionScheduler is None: raise ImportError("VQDiffusionScheduler not found.") print("[FlowMatch Scheduler] Creating VQDiffusionScheduler (Experimental)") print("[FlowMatch Scheduler] WARNING: Returning dummy sigmas. This scheduler is for discrete latent models.") # We don't actually use the scheduler to generate sigmas because it can't. # We just return the dummy sigmas. sigmas = torch.linspace(1.0, 0.0, steps + 1) # Default to CPU, KSampler will move it if needed or we can try to detect # But here we don't have model context easily. return (sigmas,) NODE_CLASS_MAPPINGS = { "FlowMatchEulerDiscreteScheduler (Custom)": FlowMatchEulerSchedulerNode, # "VQDiffusionScheduler": VQDiffusionSchedulerNode, } NODE_DISPLAY_NAME_MAPPINGS = { "FlowMatchEulerDiscreteScheduler (Custom)": "FlowMatch Euler Discrete Scheduler (Custom)", # "VQDiffusionScheduler": "VQ Diffusion Scheduler (Experimental)", } from .extract_metadata_node import NODE_CLASS_MAPPINGS as METADATA_NODE_MAPPINGS from .extract_metadata_node import NODE_DISPLAY_NAME_MAPPINGS as METADATA_DISPLAY_MAPPINGS NODE_CLASS_MAPPINGS.update(METADATA_NODE_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(METADATA_DISPLAY_MAPPINGS) # Import Nunchaku nodes try: from .nunchaku_compat import NODE_CLASS_MAPPINGS as NUNCHAKU_NODES from .nunchaku_compat import NODE_DISPLAY_NAME_MAPPINGS as NUNCHAKU_NAMES NODE_CLASS_MAPPINGS.update(NUNCHAKU_NODES) NODE_DISPLAY_NAME_MAPPINGS.update(NUNCHAKU_NAMES) except Exception as e: print(f"[FlowMatch Scheduler] Could not load Nunchaku nodes: {e}")