__init__.py

# 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}")