chore:remove sim filter

src/batched_denoising/batched_denoising.py

@@ -1,206 +1,761 @@
 import torch
+from .control_base import ControlNodeBase
 import comfy.model_management
 import comfy.samplers
+import math
 
-class StreamBatchSampler:
-    
+
+class StreamBatchSampler(ControlNodeBase):
     @classmethod
     def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "num_steps": ("INT", {
-                    "default": 4,
-                    "min": 1,
-                    "max": 10,
-                    "step": 1,
-                    "tooltip": "Number of denoising steps. Should match the frame buffer size."
-                }),
-            },
+        inputs = super().INPUT_TYPES()
+        inputs["required"].update(
+            {
+                "latent": ("LATENT",),
+                "num_timesteps": (
+                    "INT",
+                    {
+                        "default": 2,
+                        "min": 1,
+                        "max": 10,
+                        "step": 1,
+                        "tooltip": "Number of denoising steps to use. More steps = better quality but MUCH slower. StreamDiffusion's speed comes from using just 2 steps for img2img (32,45) or 4 for txt2img (0,16,32,45)",
+                    },
+                ),
+                "frame_buffer_size": (
+                    "INT",
+                    {
+                        "default": 1,
+                        "min": 1,
+                        "max": 10,
+                        "step": 1,
+                        "tooltip": "How many frames to process together for temporal consistency. Each frame is processed at each timestep (total batch = frame_buffer_size * num_timesteps). Higher values reduce flickering but use more VRAM and don't improve speed",
+                    },
+                ),
+                "cfg_type": (
+                    ["none", "full", "self", "initialize"],
+                    {
+                        "default": "self",
+                        "tooltip": "'self' is fastest and most memory efficient, 'full' is standard SD behavior but slower, 'initialize' is a middle ground, 'none' disables guidance",
+                    },
+                ),
+                "delta": (
+                    "FLOAT",
+                    {
+                        "default": 1.0,
+                        "min": 0.0,
+                        "max": 5.0,
+                        "step": 0.1,
+                        "tooltip": "Only used with 'self' CFG. Controls strength of self-guidance. Higher values = stronger guidance but more artifacts. Default 1.0 works well",
+                    },
+                ),
+                "similarity_threshold": (
+                    "FLOAT",
+                    {
+                        "default": 0.98,
+                        "min": 0.0,
+                        "max": 1.0,
+                        "step": 0.01,
+                        "tooltip": "Skip frames that are too similar to save compute. 0 = process all frames. Try 0.98 if you want to skip static scenes. Higher = more skipping",
+                        # TODO: see if this is backwards from original
+                    },
+                ),
+                "max_skip_frames": (
+                    "INT",
+                    {
+                        "default": 10,
+                        "min": 1,
+                        "max": 100,
+                        "step": 1,
+                        "tooltip": "Maximum frames to skip when using similarity filter. Prevents getting stuck on static scenes. Lower = more responsive to changes",
+                    },
+                ),
+            }
+        )
+        inputs["optional"] = {
+            "prompt": (
+                "STRING",
+                {
+                    "multiline": True,
+                    "default": "",
+                    "tooltip": "Optional text prompt. Enables KV caching for faster processing. Leave empty to disable. Same prompt across frames = faster",
+                },
+            ),
         }
-    
-    RETURN_TYPES = ("SAMPLER",)
+        return inputs
+
+    RETURN_TYPES = (
+        "SAMPLER",
+        "LATENT",
+    )
     FUNCTION = "update"
     CATEGORY = "StreamPack/sampling"
-    DESCRIPTION = "Implements batched denoising for faster inference by processing multiple frames in parallel at different denoising steps. Also adds temportal consistency to the denoising process."
-    
+    DESCRIPTION = "Implements batched denoising for faster inference by processing multiple frames in parallel at different denoising steps (StreamDiffusion)"
+
     def __init__(self):
-        self.num_steps = None
-        self.frame_buffer = []
+        super().__init__()
+        self.frame_buffer_size = None
+        self.frame_buffer = None  # Will be a tensor instead of list
         self.x_t_latent_buffer = None
         self.stock_noise = None
-        self.is_txt2img_mode = False
-        
-        # Initialize all buffers
-        self.zeros_reference = None
-        self.random_noise_buffer = None
-        self.sigmas_view_buffer = None
-        self.expanded_stock_noise = None
-        self.working_buffer = None
-        self.output_buffer = None
-    
-    def sample(self, model, noise, sigmas, extra_args=None, callback=None, disable=None):
-        """Sample with staggered batch denoising steps - Optimized version"""
+        self.cached_noise = {}  # Cache for different expanded versions
+        self.debug = False  # Debug flag for print statements
+        self.last_batch_shape = None  # Track last batch shape
+        # Add CFG state
+        self.cfg_type = "self"
+        self.guidance_scale = 1.2
+        self.delta = 1.0
+        # Add cross attention state
+        self.prompt = ""
+        self.last_prompt = None
+        self.kv_cache = None
+        self.cross_attention_hook = None
+        self.last_model_hash = None
+        # Add similarity filter state
+        self.similarity_threshold = 0.0
+        self.max_skip_frames = 10
+        self.similarity_filter = None
+        # Add buffer state
+        self.num_timesteps = None
+        # Add caching for alpha/beta and their square roots
+        self.cached_sigmas = None
+        self.alpha_prod_t = None
+        self.beta_prod_t = None
+        self.alpha_prod_t_sqrt = None
+        self.beta_prod_t_sqrt = None
+        # Add noise sequence and step counter
+        self.noise_sequence = None
+        self.step_counter = 0
+
+    def get_expanded_noise(self, batch_size, scale=1.0):
+        """Get expanded noise with caching based on batch size and scale"""
+        cache_key = (batch_size, scale)
+        if self.stock_noise is None:
+            return None
+
+        if cache_key not in self.cached_noise:
+            self.cached_noise[cache_key] = (
+                self.stock_noise.unsqueeze(0).expand(batch_size, -1, -1, -1) * scale
+            )
+            if self.debug:
+                print(
+                    f"[StreamBatchSampler] Cached expanded noise for batch_size={batch_size}, scale={scale}"
+                )
+
+        return self.cached_noise[cache_key]
+
+    def invalidate_noise_cache(self):
+        """Clear all cached noise when stock noise changes"""
+        self.cached_noise.clear()
+
+    def process_with_cfg(self, model, x, sigma_batch, extra_args):
+        """Apply classifier-free guidance based on current cfg_type"""
+        cond_scale = extra_args["model_options"].get("cond_scale", 1.0)
+
+        if self.cfg_type == "none" or cond_scale <= 1.0:
+            if self.debug:
+                print("[CFG] No guidance - running model directly")
+            return model(x, sigma_batch, **extra_args)
+
+        if self.cfg_type == "full":
+            if self.debug:
+                print("[CFG] Using full CFG")
+            # Double the batch for cond/uncond
+            x_double = torch.cat([x, x], dim=0)
+            sigma_double = torch.cat([sigma_batch, sigma_batch], dim=0)
+            model_output = model(x_double, sigma_double, **extra_args)
+            noise_pred_uncond, noise_pred_text = model_output.chunk(2)
+            # Combine predictions
+            eps_cfg = noise_pred_uncond + cond_scale * (
+                noise_pred_text - noise_pred_uncond
+            )
+            return eps_cfg
+
+        elif self.cfg_type == "initialize":
+            if self.debug:
+                print("[CFG] Using initialize CFG")
+            # Add single uncond at start
+            x_plus_uc = torch.cat([x[0:1], x], dim=0)
+            sigma_plus_uc = torch.cat([sigma_batch[0:1], sigma_batch], dim=0)
+            model_output = model(x_plus_uc, sigma_plus_uc, **extra_args)
+            noise_pred_uncond = model_output[0:1]
+            noise_pred_text = model_output[1:]
+            # Update stock noise with uncond prediction
+            self.stock_noise = noise_pred_uncond[0].clone()
+            self.invalidate_noise_cache()
+            return noise_pred_text  # Returning noise_pred_text as per original logic
+
+        else:  # self
+            if self.debug:
+                print("\n[CFG] Using RCFG (self)")
+                print(f"[CFG] Input shape: {x.shape}, Sigma shape: {sigma_batch.shape}")
+
+            # Get current model prediction
+            eps_c = model(x, sigma_batch, **extra_args)
+            if self.debug:
+                print(
+                    f"[CFG] Current prediction (eps_c) stats: min={eps_c.min():.3f}, max={eps_c.max():.3f}, mean={eps_c.mean():.3f}"
+                )
+
+            # Initialize stock noise if needed
+            if self.stock_noise is None:
+                if self.debug:
+                    print("[CFG] Initializing stock noise with current prediction")
+                self.stock_noise = eps_c[0].clone().detach()
+                self.invalidate_noise_cache()
+                return eps_c
+
+            # Apply RCFG equation
+            noise_pred_uncond = self.stock_noise.to(eps_c.device) * self.delta
+            if self.debug:
+                print(f"[CFG] Using delta={self.delta}, cond_scale={cond_scale}")
+                print(
+                    f"[CFG] Stock noise stats: min={self.stock_noise.min():.3f}, max={self.stock_noise.max():.3f}, mean={self.stock_noise.mean():.3f}"
+                )
+                print(
+                    f"[CFG] Uncond prediction stats: min={noise_pred_uncond.min():.3f}, max={noise_pred_uncond.max():.3f}, mean={noise_pred_uncond.mean():.3f}"
+                )
+
+            # Combine predictions exactly as StreamDiffusion does
+            eps_cfg = noise_pred_uncond + cond_scale * (eps_c - noise_pred_uncond)
+
+            if self.debug:
+                print(
+                    f"[CFG] Final output (eps_cfg) stats: min={eps_cfg.min():.3f}, max={eps_cfg.max():.3f}, mean={eps_cfg.mean():.3f}"
+                )
+                print(
+                    f"[CFG] Difference (eps_cfg - eps_c) stats: min={(eps_cfg-eps_c).min():.3f}, max={(eps_cfg-eps_c).max():.3f}, mean={(eps_cfg-eps_c).mean():.3f}\n"
+                )
+
+            return eps_cfg
+
+    def setup_cross_attention(self, model, prompt):
+        """Set up cross attention optimization for the model"""
+        if not prompt:  # Don't use optimization if no prompt provided
+            return model
+
+        def optimized_attention_forward(module, x, context=None, mask=None):
+            """Optimized cross attention with KV caching and ComfyUI's optimizations"""
+            batch_size = x.shape[0]
+
+            # Generate query from input
+            q = module.to_q(x)
+
+            # Check if we need to recompute KV cache
+            if self.kv_cache is None or self.last_prompt != prompt:
+                if self.debug:
+                    print("[StreamBatchSampler] Computing new KV cache")
+                k = module.to_k(context)
+                v = module.to_v(context)
+                self.kv_cache = (k, v)
+                self.last_prompt = prompt
+            else:
+                if self.debug:
+                    print("[StreamBatchSampler] Using cached KV pairs")
+                k, v = self.kv_cache
+
+            # Handle batch size expansion if needed
+            if k.shape[0] == 1 and batch_size > 1:
+                k = k.expand(batch_size, -1, -1)
+                v = v.expand(batch_size, -1, -1)
+
+            # Use ComfyUI's optimized attention
+            from comfy.ldm.modules.attention import optimized_attention
+
+            out = optimized_attention(q, k, v, module.heads, mask=mask)
+            return module.to_out[0](out)
+
+        # Get the actual model from the sampler
+        if hasattr(model, "model"):
+            unet = model.model
+        else:
+            return model
+
+        # Only set up hooks if prompt has changed
+        if self.last_prompt != prompt:
+            # Remove old hooks if they exist
+            if self.cross_attention_hook is not None:
+                self.cross_attention_hook.remove()
+
+            def create_cross_attn_patch(module, input, output):
+                """Hook to replace cross attention modules with optimized version"""
+                if isinstance(module, torch.nn.Module) and hasattr(module, "to_q"):
+                    if not hasattr(module, "_original_forward"):
+                        module._original_forward = module.forward
+                    module.forward = (
+                        lambda *args, **kwargs: optimized_attention_forward(
+                            module, *args, **kwargs
+                        )
+                    )
+                return output
+
+            # Register hooks on cross attention blocks
+            def register_hooks(module):
+                if isinstance(module, torch.nn.Module) and hasattr(module, "to_q"):
+                    self.cross_attention_hook = module.register_forward_hook(
+                        create_cross_attn_patch
+                    )
+
+            unet.apply(register_hooks)
+            self.last_prompt = prompt
+
+        return model
+
+    def compute_alpha_beta(self, sigmas):
+        """Pre-compute alpha and beta terms for the given sigmas (LCM-like using cosine schedule)"""
+        self.cached_sigmas = sigmas.clone()
+        print(f"\n[Alpha/Beta] Input sigmas: {sigmas}")
+
+        num_train_timesteps = 1000  # Standard for many diffusion models
+        timesteps = torch.linspace(
+            1, 0, num_train_timesteps, dtype=torch.float32, device=sigmas.device
+        )
+        alphas_cumprod = torch.cos((timesteps + 0.008) / 1.008 * math.pi / 2).pow(2)
+        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+
+        # Map sigmas to timesteps
+        scaled_sigmas = (
+            sigmas / sigmas.max() if sigmas.max() > 0 else torch.zeros_like(sigmas)
+        )
+        timesteps_for_sigma = (1 - scaled_sigmas) * (num_train_timesteps - 1)
+        indices = (
+            torch.round(timesteps_for_sigma).long().clamp(0, num_train_timesteps - 1)
+        )
+
+        alpha_prod_t = alphas_cumprod[indices].view(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        self.alpha_prod_t = alpha_prod_t
+        self.beta_prod_t = beta_prod_t
+        self.alpha_prod_t_sqrt = torch.sqrt(alpha_prod_t)
+        self.beta_prod_t_sqrt = torch.sqrt(beta_prod_t)
+
+        print(f"[Alpha/Beta] Final stats:")
+        print(
+            f"  alpha_prod_t: min={self.alpha_prod_t.min():.4f}, max={self.alpha_prod_t.max():.4f}, mean={self.alpha_prod_t.mean():.4f}"
+        )
+        print(
+            f"  beta_prod_t: min={self.beta_prod_t.min():.4f}, max={self.beta_prod_t.max():.4f}, mean={self.beta_prod_t.mean():.4f}"
+        )
+        print(
+            f"  alpha_sqrt: min={self.alpha_prod_t_sqrt.min():.4f}, max={self.alpha_prod_t_sqrt.max():.4f}, mean={self.alpha_prod_t_sqrt.mean():.4f}"
+        )
+        print(
+            f"  beta_sqrt: min={self.beta_prod_t_sqrt.min():.4f}, max={self.beta_prod_t_sqrt.max():.4f}, mean={self.beta_prod_t_sqrt.mean():.4f}\n"
+        )
+
+    def scheduler_step_batch(self, denoised_batch, scaled_noise):
+        """Compute scheduler step for the entire batch (LCM-like)"""
+        print("\n[Scheduler] Starting scheduler step (LCM-like):")
+        print(f"  Denoised shape: {denoised_batch.shape}")
+        print(f"  Scaled noise shape: {scaled_noise.shape}")
+        print(
+            f"  Denoised stats: min={denoised_batch.min():.4f}, max={denoised_batch.max():.4f}, mean={denoised_batch.mean():.4f}"
+        )
+        print(
+            f"  Scaled noise stats: min={scaled_noise.min():.4f}, max={scaled_noise.max():.4f}, mean={scaled_noise.mean():.4f}"
+        )
+
+        # Assuming self.alpha_prod_t and self.beta_prod_t are for the current timestep
+        alpha_t_sqrt = self.alpha_prod_t_sqrt[: denoised_batch.shape[0]]
+        beta_t_sqrt = self.beta_prod_t_sqrt[: denoised_batch.shape[0]]
+
+        # Calculate the predicted original sample (x_0) from the noisy sample and model output
+        # This is similar to the inverse of the forward diffusion step
+        pred_original_sample = (
+            denoised_batch - beta_t_sqrt.view(-1, 1, 1, 1) * scaled_noise
+        ) / (alpha_t_sqrt.view(-1, 1, 1, 1) + 1e-7)
+        print(
+            f"  Predicted original sample stats: min={pred_original_sample.min():.4f}, max={pred_original_sample.max():.4f}, mean={pred_original_sample.mean():.4f}"
+        )
+
+        # For the reverse step, we need the alpha and beta for the previous timestep
+        # Since we are operating on a batch of timesteps, we'll use the values corresponding to the first timestep in the batch for updating stock_noise
+        # In StreamDiffusion paper, the update happens based on the result of the last denoising step
+
+        # Get alpha and beta for the next step: (which would be the previous in the denoising sequence)
+        alpha_next_sqrt = torch.cat(
+            [
+                self.alpha_prod_t_sqrt[1 : denoised_batch.shape[0]],
+                torch.ones_like(self.alpha_prod_t_sqrt[0:1]),
+            ],
+            dim=0,
+        )
+
+        # Calculate the coefficient for the predicted original sample in the next step
+        sigma_next = (1 - alpha_next_sqrt.pow(2)) / (alpha_next_sqrt + 1e-7)
+        sigma_t = (1 - alpha_t_sqrt.pow(2)) / (alpha_t_sqrt + 1e-7)
+
+        # Guidance scale (set to 1 for unconditional prediction)
+        guidance_scale = 1
+
+        # Calculate delta_x (the change in the noisy sample)
+        delta_x = (sigma_next - sigma_t) * guidance_scale * pred_original_sample
+
+        print(
+            f"  Final delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}\n"
+        )
+
+        return delta_x
+
+    def sample(
+        self, model, noise, sigmas, extra_args=None, callback=None, disable=None
+    ):
+        """Sample with staggered batch denoising steps"""
         extra_args = {} if extra_args is None else extra_args
-        
+
+        print(f"\n[Sample] Starting sample step {self.step_counter}:")
+        print(f"  Input noise shape: {noise.shape}")
+        print(f"  Sigmas shape: {sigmas.shape}")
+        print(f"  Sigmas values: {sigmas[:self.num_timesteps]}")
+
+        # Set up cross attention optimization if prompt provided
+        if self.prompt:
+            model = self.setup_cross_attention(model, self.prompt)
+
         # Get number of frames in batch and available sigmas
         batch_size = noise.shape[0]
-        num_sigmas = len(sigmas) - 1  # Subtract 1 because last sigma is the target (0.0)
-        
-        # Reuse zeros buffer for txt2img detection
-        if self.zeros_reference is None:
-            self.zeros_reference = torch.zeros(1, device=noise.device, dtype=noise.dtype)
-        
-        # Check if noise tensor is all zeros
-        self.is_txt2img_mode = torch.abs(noise).sum() < 1e-5
-        
-        if self.is_txt2img_mode:
-            # If txt2img mode, reuse the noise tensor directly
-            if self.random_noise_buffer is None or self.random_noise_buffer.shape != noise.shape:
-                self.random_noise_buffer = torch.empty_like(noise)
-            
-            # Generate random noise in-place
-            self.random_noise_buffer.normal_()
-            x = self.random_noise_buffer  # Use pre-allocated buffer
+
+        # Calculate total batch size based on frame_buffer_size and number of timesteps
+        total_batch_size = self.frame_buffer_size * self.num_timesteps
+        if batch_size != total_batch_size:
+            raise ValueError(
+                f"Batch size ({batch_size}) must match frame_buffer_size * num_timesteps ({total_batch_size})"
+            )
+
+        # Pre-compute alpha and beta terms if not already computed or if sigmas changed
+        if self.cached_sigmas is None or not torch.equal(self.cached_sigmas, sigmas):
+            self.compute_alpha_beta(sigmas)
+
+        # Update noise sequence like StreamDiffusion
+        if self.step_counter == 0:
+            # First step - use initial noise
+            self.stock_noise = self.init_noise[0].clone()
         else:
-            # If not txt2img, we'll still need to add noise later
-            x = noise
-        
-        # Verify batch size matches number of timesteps
-        if batch_size != num_sigmas:
-            raise ValueError(f"Batch size ({batch_size}) must match number of timesteps ({num_sigmas})")
-        
-        # Pre-compute alpha and beta terms
-        alpha_prod_t = (sigmas[:-1] / sigmas[0]).view(-1, 1, 1, 1)  # [B,1,1,1]
-        beta_prod_t = (1 - alpha_prod_t)
-        
-        # Initialize stock noise with reuse
-        if self.stock_noise is None or self.stock_noise.shape != noise[0].shape:
-            self.stock_noise = torch.empty_like(noise[0])
-            self.stock_noise.normal_()  # Generate random noise in-place
-        
-        # Pre-allocate and reuse view buffer for sigmas
-        if self.sigmas_view_buffer is None or self.sigmas_view_buffer.shape[0] != len(sigmas)-1:
-            self.sigmas_view_buffer = torch.empty((len(sigmas)-1, 1, 1, 1), 
-                                               device=sigmas.device, 
-                                               dtype=sigmas.dtype)
-        # In-place copy of sigmas view
-        self.sigmas_view_buffer.copy_(sigmas[:-1].view(-1, 1, 1, 1))
-        
-        # Eliminate unsqueeze allocation by pre-expanding stock noise
-        if self.expanded_stock_noise is None or self.expanded_stock_noise.shape[0] != batch_size:
-            self.expanded_stock_noise = self.stock_noise.expand(batch_size, *self.stock_noise.shape)
-        
-        # Apply noise with pre-allocated buffers - no new memory allocation
-        if not self.is_txt2img_mode:  # Already handled txt2img case above
-            # If we need a working buffer separate from noise input:
-            if id(x) == id(noise):  # They're the same object, need a separate buffer
-                if self.working_buffer is None or self.working_buffer.shape != noise.shape:
-                    self.working_buffer = torch.empty_like(noise)
-                x = self.working_buffer
-                # Add noise to input
-                torch.add(noise, self.expanded_stock_noise * self.sigmas_view_buffer, out=x)
-        
-        # Initialize and manage latent buffer with memory optimization
-        if (self.x_t_latent_buffer is None or self.is_txt2img_mode) and num_sigmas > 1:
-            # Pre-allocate or resize as needed
-            if self.x_t_latent_buffer is None or self.x_t_latent_buffer.shape != x[0].shape:
-                self.x_t_latent_buffer = torch.empty_like(x[0])
-            self.x_t_latent_buffer.copy_(x[0])
-        
+            # Shift noise sequence
+            self.stock_noise = torch.cat(
+                [self.init_noise[1:], self.init_noise[0:1]], dim=0
+            )[0]
+
+        print(
+            f"  Stock noise stats: min={self.stock_noise.min():.4f}, max={self.stock_noise.max():.4f}, mean={self.stock_noise.mean():.4f}"
+        )
+
+        # Scale noise for each frame based on its sigma - vectorized with caching
+        sigma_scaled_noise = self.get_expanded_noise(batch_size)
+        print(f"  Expanded noise shape: {sigma_scaled_noise.shape}")
+        print(
+            f"  Expanded noise stats: min={sigma_scaled_noise.min():.4f}, max={sigma_scaled_noise.max():.4f}, mean={sigma_scaled_noise.mean():.4f}"
+        )
+
+        x = noise + sigma_scaled_noise * sigmas[: self.num_timesteps].view(-1, 1, 1, 1)
+        print(
+            f"  Initial x stats: min={x.min():.4f}, max={x.max():.4f}, mean={x.mean():.4f}"
+        )
+
+        # Initialize frame buffer if needed
+        if self.x_t_latent_buffer is None and self.num_timesteps > 1:
+            self.x_t_latent_buffer = x[0].clone()
+            print("  Initialized latent buffer with first frame")
+
         # Use buffer for first frame to maintain temporal consistency
-        if num_sigmas > 1:
-            x[0].copy_(self.x_t_latent_buffer)
-        
+        if self.num_timesteps > 1:
+            x = torch.cat([self.x_t_latent_buffer.unsqueeze(0), x[1:]], dim=0)
+            print(
+                f"  After buffer concat x stats: min={x.min():.4f}, max={x.max():.4f}, mean={x.mean():.4f}"
+            )
+
+        # Run model on entire batch at once
         with torch.no_grad():
             # Process all frames in parallel
-            sigma_batch = sigmas[:-1]
-            
-            denoised_batch = model(x, sigma_batch, **extra_args)
-            
+            sigma_batch = sigmas[: self.num_timesteps]
+            # Process with CFG handling inside process_with_cfg
+            denoised_batch = self.process_with_cfg(model, x, sigma_batch, extra_args)
+            print(f"\n[Sample] After denoising:")
+            print(f"  Denoised shape: {denoised_batch.shape}")
+            print(
+                f"  Denoised stats: min={denoised_batch.min():.4f}, max={denoised_batch.max():.4f}, mean={denoised_batch.mean():.4f}"
+            )
+
             # Update buffer with intermediate results
-            if num_sigmas > 1:
+            if self.num_timesteps > 1:
                 # Store result from first frame as buffer for next iteration
-                self.x_t_latent_buffer.copy_(denoised_batch[0])  # In-place update
-                
-                # Pre-allocate output buffer
-                if self.output_buffer is None or self.output_buffer.shape != (1, *denoised_batch[-1].shape):
-                    self.output_buffer = torch.empty(1, *denoised_batch[-1].shape, 
-                                                  device=denoised_batch.device,
-                                                  dtype=denoised_batch.dtype)
-                # Copy the result directly to pre-allocated buffer
-                self.output_buffer[0].copy_(denoised_batch[-1])
-                x_0_pred_out = self.output_buffer
+                self.x_t_latent_buffer = denoised_batch[0].clone()
+                print(
+                    f"  Updated buffer stats: min={self.x_t_latent_buffer.min():.4f}, max={self.x_t_latent_buffer.max():.4f}, mean={self.x_t_latent_buffer.mean():.4f}"
+                )
+                # Return result from last frame
+                x_0_pred_out = denoised_batch[-1].unsqueeze(0)
             else:
                 x_0_pred_out = denoised_batch
                 self.x_t_latent_buffer = None
-            
+
+            print(
+                f"  Output pred stats: min={x_0_pred_out.min():.4f}, max={x_0_pred_out.max():.4f}, mean={x_0_pred_out.mean():.4f}"
+            )
+
             # Call callback if provided
             if callback is not None:
-                callback({'x': x_0_pred_out, 'i': 0, 'sigma': sigmas[0], 'sigma_hat': sigmas[0], 'denoised': denoised_batch[-1:]})
-        
+                callback(
+                    {
+                        "x": x_0_pred_out,
+                        "i": 0,
+                        "sigma": sigmas[0],
+                        "sigma_hat": sigmas[0],
+                        "denoised": x_0_pred_out,
+                    }
+                )
+
+            # Update stock noise using scheduler steps
+            if self.cfg_type == "self":
+                print("\n[Sample] Updating stock noise:")
+                # Scale noise properly
+                scaled_noise = (
+                    self.beta_prod_t_sqrt[: denoised_batch.shape[0]].view(-1, 1, 1, 1)
+                    * self.stock_noise
+                )
+                print(
+                    f"  Scaled noise stats: min={scaled_noise.min():.4f}, max={scaled_noise.max():.4f}, mean={scaled_noise.mean():.4f}"
+                )
+
+                # Get delta_x using scheduler step
+                delta_x = self.scheduler_step_batch(denoised_batch, scaled_noise)
+                print(
+                    f"  Initial delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}"
+                )
+
+                # Prepare next step's alpha/beta terms
+                alpha_next = torch.cat(
+                    [
+                        self.alpha_prod_t_sqrt[1 : denoised_batch.shape[0]],
+                        torch.ones_like(self.alpha_prod_t_sqrt[0:1]),
+                    ],
+                    dim=0,
+                )
+                beta_next = torch.cat(
+                    [
+                        self.beta_prod_t_sqrt[1 : denoised_batch.shape[0]],
+                        torch.zeros_like(self.beta_prod_t_sqrt[0:1]),
+                    ],
+                    dim=0,
+                )
+
+                print(
+                    f"  Alpha next stats: min={alpha_next.min():.4f}, max={alpha_next.max():.4f}, mean={alpha_next.mean():.4f}"
+                )
+                print(
+                    f"  Beta next stats: min={beta_next.min():.4f}, max={beta_next.max():.4f}, mean={beta_next.mean():.4f}"
+                )
+
+                # Update noise prediction
+                delta_x = (alpha_next.view(-1, 1, 1, 1) * delta_x) / (
+                    beta_next.view(-1, 1, 1, 1) + 1e-7
+                )
+                print(
+                    f"  Scaled delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}"
+                )
+
+                # Use noise sequence for init_noise
+                init_noise = torch.cat(
+                    [self.init_noise[1:], self.init_noise[0:1]], dim=0
+                )
+                print(
+                    f"  Init noise stats: min={init_noise.min():.4f}, max={init_noise.max():.4f}, mean={init_noise.mean():.4f}"
+                )
+
+                # Update stock noise with proper scaling
+                self.stock_noise = init_noise[0] + delta_x[0]
+                print(
+                    f"  Final stock noise stats: min={self.stock_noise.min():.4f}, max={self.stock_noise.max():.4f}, mean={self.stock_noise.mean():.4f}\n"
+                )
+
+            # Update step counter
+            self.step_counter += 1
+
         return x_0_pred_out
-    
-    def update(self, num_steps=4):
-        """Create sampler with specified settings"""
-        self.num_steps = num_steps
+
+    def buffer(self, latent, frame_buffer_size=1, num_timesteps=2, always_execute=True):
+        """Add new frame to buffer and return batch when ready"""
+        self.frame_buffer_size = frame_buffer_size
+        self.num_timesteps = num_timesteps
+
+        # Extract latent tensor from input and remove batch dimension if present
+        x = latent["samples"]
+        if x.dim() == 4:  # [B,C,H,W]
+            x = x.squeeze(0)  # Remove batch dimension -> [C,H,W]
+
+        # Initialize or resize frame buffer if needed
+        if (
+            self.frame_buffer is None
+            or self.frame_buffer.shape[0] != self.frame_buffer_size
+        ):
+            # Pre-allocate tensor buffer
+            self.frame_buffer = torch.zeros(
+                (self.frame_buffer_size,) + x.shape, dtype=x.dtype, device=x.device
+            )
+            if self.debug:
+                print(
+                    f"[StreamFrameBuffer] Pre-allocated tensor buffer for {self.frame_buffer_size} frames"
+                )
+
+        # Rotate frames and add new frame using tensor operations
+        if self.frame_buffer_size > 1:
+            self.frame_buffer = torch.roll(self.frame_buffer, -1, dims=0)
+        self.frame_buffer[-1] = x
+
+        # Create batch by repeating frames for each timestep
+        batch = self.frame_buffer.repeat_interleave(
+            self.num_timesteps, dim=0
+        )  # [buffer_size*num_timesteps,C,H,W]
+
+        # Return as latent dict
+        return {"samples": batch}
+
+    def update(
+        self,
+        latent,
+        num_timesteps,
+        frame_buffer_size=1,
+        cfg_type="self",
+        delta=1.0,
+        prompt="",
+        similarity_threshold=0.0,
+        max_skip_frames=10,
+        always_execute=True,
+    ):
+        """Create sampler with specified settings and process frame"""
+        self.frame_buffer_size = frame_buffer_size
+        self.cfg_type = cfg_type
+        self.delta = delta
+        self.prompt = prompt
+        self.num_timesteps = num_timesteps
+
+        # Update similarity filter settings
+        self.similarity_threshold = similarity_threshold
+        self.max_skip_frames = max_skip_frames
+        if self.similarity_filter is not None:
+            self.similarity_filter.set_threshold(similarity_threshold)
+            self.similarity_filter.set_max_skip_frame(max_skip_frames)
+
+        # Create sampler with custom sample function
         sampler = comfy.samplers.KSAMPLER(self.sample)
-        return (sampler,)
+
+        # Process frame
+        buffered = self.buffer(latent)
+
+        # Get latent dimensions from input
+        _, _, self.latent_height, self.latent_width = latent["samples"].shape
+
+        # Add noise sequence initialization
+        if self.noise_sequence is None or self.noise_sequence.shape[0] != num_timesteps:
+            # Initialize with random noise like StreamDiffusion
+            self.init_noise = torch.randn(
+                (num_timesteps, 4, self.latent_height, self.latent_width),
+                device=comfy.model_management.get_torch_device(),
+                dtype=latent["samples"].dtype,
+            )
+
+            # Initialize stock noise as zeros like StreamDiffusion
+            self.stock_noise = torch.zeros_like(self.init_noise[0])
+
+            # Create noise sequence from init_noise
+            self.noise_sequence = self.init_noise.clone()
+
+            # No need for manual shifting as it's handled in sample() method
+
+        return (sampler, buffered)
 
 
-class StreamScheduler:
-    
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "model": ("MODEL",),
-                "t_index_list": ("STRING", {
-                    "default": "0,16,32,49",
-                    "tooltip": "Comma-separated list of timesteps to actually use for denoising. Examples: '32,45' for img2img or '0,16,32,45' for txt2img"
-                }),
-                "num_inference_steps": ("INT", {
-                    "default": 50,
-                    "min": 1,
-                    "max": 1000,
-                    "step": 1,
-                    "tooltip": "Total number of timesteps in schedule. StreamDiffusion uses 50 by default. Only timesteps specified in t_index_list are actually used."
-                }),
-            },
-        }
+class StreamScheduler(ControlNodeBase):
+    """Implements StreamDiffusion's efficient timestep selection"""
 
     RETURN_TYPES = ("SIGMAS",)
     FUNCTION = "update"
-    CATEGORY = "StreamPack/sampling"
-    DESCRIPTION = "Implements StreamDiffusion's efficient timestep selection. Use in conjunction with StreamBatchSampler."
-    
-    def update(self, model, t_index_list="32,45", num_inference_steps=50):
+    CATEGORY = "real-time/sampling"
+
+    def __init__(self):
+        super().__init__()
+        self.cached_indices = None
+        self.last_t_index_list = None
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        inputs = super().INPUT_TYPES()
+        inputs["required"].update(
+            {
+                "model": ("MODEL",),
+                "t_index_list": (
+                    "STRING",
+                    {
+                        "default": "32,45",
+                        "tooltip": "Comma-separated list of timesteps to actually use for denoising. For LCM: use '32,45' for img2img or '0,16,32,45' for txt2img. For SDXL Turbo: use '45,49' for 2 steps or '49' for 1 step (Turbo works best with steps near the end of denoising)",
+                    },
+                ),
+                "num_inference_steps": (
+                    "INT",
+                    {
+                        "default": 50,
+                        "min": 1,
+                        "max": 1000,
+                        "step": 1,
+                        "tooltip": "Total number of timesteps in schedule. StreamDiffusion uses 50 by default. Only timesteps specified in t_index_list are actually used.",
+                    },
+                ),
+            }
+        )
+        return inputs
+
+    def parse_timesteps(self, t_index_list):
+        """Parse and sort timesteps, with caching"""
+        if t_index_list == self.last_t_index_list and self.cached_indices is not None:
+            return self.cached_indices
+
+        try:
+            # Parse and sort in reverse order once
+            indices = sorted(
+                [int(t.strip()) for t in t_index_list.split(",")], reverse=True
+            )
+            self.last_t_index_list = t_index_list
+            self.cached_indices = indices
+            return indices
+        except ValueError as e:
+            print(f"Error parsing timesteps: {e}. Using default [32,45]")
+            self.last_t_index_list = "32,45"
+            self.cached_indices = [45, 32]  # Already sorted in reverse
+            return self.cached_indices
+
+    def update(
+        self, model, t_index_list="32,45", num_inference_steps=50, always_execute=True
+    ):
         # Get model's sampling parameters
         model_sampling = model.get_model_object("model_sampling")
-        
-        # Parse timestep list
-        try:
-            t_index_list = [int(t.strip()) for t in t_index_list.split(",")]
-        except ValueError as e:
-            
-            t_index_list = [32, 45]
-            
+
+        # Get parsed and sorted timesteps
+        t_indices = self.parse_timesteps(t_index_list)
+
         # Create full schedule using normal scheduler
-        full_sigmas = comfy.samplers.normal_scheduler(model_sampling, num_inference_steps)
-        
-        # Select only the sigmas at our desired indices, but in reverse order
-        # This ensures we go from high noise to low noise
+        full_sigmas = comfy.samplers.normal_scheduler(
+            model_sampling, num_inference_steps
+        )
+
+        # Select only the sigmas at our desired indices
         selected_sigmas = []
-        for t in sorted(t_index_list, reverse=True):  # Sort in reverse to go from high noise to low
+        for t in t_indices:
             if t < 0 or t >= num_inference_steps:
-                
+                print(
+                    f"Warning: timestep {t} out of range [0,{num_inference_steps}), skipping"
+                )
                 continue
             selected_sigmas.append(float(full_sigmas[t]))
-            
+
         # Add final sigma
         selected_sigmas.append(0.0)
-        
-        # Convert to tensor and move to appropriate device
-        selected_sigmas = torch.FloatTensor(selected_sigmas).to(comfy.model_management.get_torch_device())
+
+        # Create tensor directly on target device
+        device = comfy.model_management.get_torch_device()
+        selected_sigmas = torch.tensor(
+            selected_sigmas, dtype=torch.float32, device=device
+        )
         return (selected_sigmas,)
 
 
@@ -284,4 +839,4 @@ class StreamFrameBuffer:
         if "width" in latent:
             result["width"] = latent["width"]
             
-        return (result,)
+        return (result,)

src/batched_denoising/control_base.py

@@ -0,0 +1,92 @@
+from abc import ABC, abstractmethod
+import time
+
+class Pattern(ABC):
+    """Base class for all patterns (movement, timing, etc)"""
+    @abstractmethod
+    def calculate(self, phase: float, min_val: float, max_val: float) -> float:
+        """Calculate pattern value at given phase (0-1)"""
+        pass
+
+    @abstractmethod
+    def get_name(self) -> str:
+        """Pattern name for UI"""
+        pass
+
+class StateManager:
+    """Handles persistent state between executions"""
+    def __init__(self):
+            self._states = {}
+    
+    def get_state(self, node_id: str, default=None):
+        return self._states.get(node_id, default)
+    
+    def set_state(self, node_id: str, state):
+        self._states[node_id] = state
+        
+    def clear_state(self, node_id: str) -> bool:
+        """Clear state for a specific node. Returns True if state was found and cleared."""
+        if node_id in self._states:
+            del self._states[node_id]
+            return True
+        return False
+        
+    def clear_all_states(self):
+        """Clear all states"""
+        self._states.clear()
+
+    def get_all_node_ids(self):
+        """Get list of all node IDs with state"""
+        return list(self._states.keys())
+        
+    def cleanup_orphaned_states(self, active_node_ids: set[str]):
+        """Remove states for nodes that no longer exist"""
+        orphaned = [nid for nid in self._states.keys() if nid not in active_node_ids]
+        for nid in orphaned:
+            del self._states[nid]
+        return len(orphaned)
+
+class ControlNodeBase(ABC):
+    """Base class for all control nodes"""
+    
+    # Shared state manager
+    state_manager = StateManager()
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        """Base input types common to all control nodes"""
+        return {
+            "required": {
+                "always_execute": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "When enabled, the node updates every execution"
+                }),
+            },
+        }
+    
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        if kwargs.get('always_execute', True):
+            return float(time.time())
+        return False
+    
+    def __init__(self):
+        self.node_id = str(id(self))
+        
+    def get_state(self, default=None):
+        return self.state_manager.get_state(self.node_id, default)
+        
+    def set_state(self, state):
+        self.state_manager.set_state(self.node_id, state)
+
+    def cleanup(self):
+        """Clean up node state when node is destroyed"""
+        self.state_manager.clear_state(self.node_id)
+
+    def __del__(self):
+        self.cleanup()
+
+    @abstractmethod
+    def update(self, *args, **kwargs):
+        """Main update method to be implemented by control nodes"""
+        pass