Merge pull request #1 from Etupa/master

Add start and end step

.gitignore

@@ -1,2 +1,8 @@
 __pycache__
-node.zip
+node.zip
+.vs/
+ComfyUI-EulerDiscreteScheduler/.vs
+.idea
+.iml
+.vs
+EulerDiscrete.iml

README.md

@@ -53,7 +53,8 @@ More examples:
 <img width="1536" height="1088" alt="image" src="https://github.com/user-attachments/assets/1931af7e-1b3e-47c9-ac20-27add5135a71" />
 
 ## Changelog
-
+**1.0.4**
+add start and end step by etupa, with some fixes
 **1.0.3**
 
 - node publish action

__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
-}
+}

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)