diff --git a/.gitignore b/.gitignore index 592d25b..391c06e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,8 @@ __pycache__ -node.zip \ No newline at end of file +node.zip +.vs/ +ComfyUI-EulerDiscreteScheduler/.vs +.idea +.iml +.vs +EulerDiscrete.iml diff --git a/README.md b/README.md index c87c7bc..e2cdb6f 100644 --- a/README.md +++ b/README.md @@ -53,7 +53,8 @@ More examples: image ## Changelog - +**1.0.4** +add start and end step by etupa, with some fixes **1.0.3** - node publish action diff --git a/__init__.py b/__init__.py index 04fcdde..64ddfe3 100644 --- a/__init__.py +++ b/__init__.py @@ -10,6 +10,7 @@ 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 @@ -22,7 +23,7 @@ except ImportError: from comfy.samplers import SchedulerHandler, SCHEDULER_HANDLERS, SCHEDULER_NAMES # Default config for registering in ComfyUI -default_config = { +default_config = { "base_image_seq_len": 256, "base_shift": math.log(3), "invert_sigmas": False, @@ -60,7 +61,19 @@ class FlowMatchEulerSchedulerNode: "default": 9, "min": 1, "max": 10000, - "tooltip": "Number of diffusion steps. Z-Image-Turbo uses 9 steps by default (8 DiT forwards). Higher = better quality but slower." + "tooltip": "Total number of diffusion steps to generate the full sigma schedule." + }), + "start_at_step": ("INT", { # <-- NEW INPUT + "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", { # <-- NEW INPUT + "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, @@ -129,11 +142,13 @@ class FlowMatchEulerSchedulerNode: RETURN_NAMES = ("sigmas",) FUNCTION = "create" CATEGORY = "sampling/schedulers" - DESCRIPTION = "FlowMatch Euler Discrete Scheduler with full parameter control. Outputs SIGMAS for use with SamplerCustom. Supports Karras sigmas, dynamic shifting, and stochastic sampling for advanced control over the diffusion process." + 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, # <-- New parameter + end_at_step, # <-- New parameter base_image_seq_len, base_shift, invert_sigmas, @@ -170,6 +185,8 @@ class FlowMatchEulerSchedulerNode: 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 @@ -183,20 +200,26 @@ class FlowMatchEulerSchedulerNode: # 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)) - return (sigmas,) - - -# Import Flash Attention node -# from .flash_attention_node import NODE_CLASS_MAPPINGS as FLASH_ATTN_MAPPINGS -# from .flash_attention_node import NODE_DISPLAY_NAME_MAPPINGS as FLASH_ATTN_DISPLAY_MAPPINGS + # 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,) NODE_CLASS_MAPPINGS = { "FlowMatchEulerDiscreteScheduler (Custom)": FlowMatchEulerSchedulerNode, - # **FLASH_ATTN_MAPPINGS } NODE_DISPLAY_NAME_MAPPINGS = { "FlowMatchEulerDiscreteScheduler (Custom)": "FlowMatch Euler Discrete Scheduler (Custom)", - # **FLASH_ATTN_DISPLAY_MAPPINGS -} +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index c2e6c6b..b7ed659 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "erosdiffusion-eulerflowmatchingdiscretescheduler" description = "Noise Free images with Euler Discrete Scheduler in ComfyUI with Z-Image or other models" -version = "1.0.3" +version = "1.0.4" license = {file = "LICENSE.TXT"} # classifiers = [ # # For OS-independent nodes (works on all operating systems)