diff --git a/__init__.py b/__init__.py index 4b98126..80bc9af 100644 --- a/__init__.py +++ b/__init__.py @@ -4,7 +4,7 @@ from .controls.utility_controls import FPSMonitor, SimilarityFilter, LazyConditi from .controls.motion_controls import MotionController, ROINode, IntegerMotionController from .misc_nodes import DTypeConverter, FastWebcamCapture, YOLOSimilarityCompare, TextRenderer, QuickShapeMask, MultilineText, LoadImageFromPath_ from .stream_diffusion_nodes import StreamConditioning, StreamCFG, StreamCrossAttention -from .stream_sampler import StreamBatchSampler, StreamScheduler +from .stream_sampler import StreamBatchSampler, StreamScheduler, StreamFrameBuffer import re @@ -22,6 +22,7 @@ NODE_CLASS_MAPPINGS = { "StreamConditioning": StreamConditioning, "StreamBatchSampler": StreamBatchSampler, "StreamScheduler": StreamScheduler, + "StreamFrameBuffer": StreamFrameBuffer, "StreamCrossAttention": StreamCrossAttention, "LazyCondition": LazyCondition, "MotionController": MotionController, @@ -30,7 +31,6 @@ NODE_CLASS_MAPPINGS = { "TextRenderer": TextRenderer, - "ROINode": ROINode, #"IntervalControl": IntervalCo ntrol, diff --git a/stream_diffusion_nodes.py b/stream_diffusion_nodes.py index 2c45253..b47677a 100644 --- a/stream_diffusion_nodes.py +++ b/stream_diffusion_nodes.py @@ -186,6 +186,7 @@ class StreamCFG(ControlNodeBase): return (model,) +#NOTE: totally and utterly experimental. No theoretical backing whatsoever. class StreamConditioning(ControlNodeBase): """Applies Residual CFG to conditioning for improved temporal consistency with different CFG types""" #NOTE: experimental @@ -213,6 +214,13 @@ class StreamConditioning(ControlNodeBase): "step": 0.1, "tooltip": "Delta parameter for self/initialize CFG types" }), + "context_size": ("INT", { + "default": 4, + "min": 1, + "max": 32, + "step": 1, + "tooltip": "Number of past conditionings to keep in context. Higher values = smoother transitions but more memory usage." + }), "always_execute": ("BOOLEAN", { "default": False, }), @@ -227,37 +235,53 @@ class StreamConditioning(ControlNodeBase): def __init__(self): super().__init__() - def update(self, positive, negative, cfg_type="full", residual_scale=0.4, delta=1.0, always_execute=False): + def update(self, positive, negative, cfg_type="full", residual_scale=0.4, delta=1.0, context_size=4, always_execute=False): # Get state with defaults state = self.get_state({ "prev_positive": None, "prev_negative": None, "stock_noise": None, # For self/initialize CFG - "initialized": False # For initialize CFG + "initialized": False, # For initialize CFG + "pos_context": [], # Store past positive conditionings + "neg_context": [] # Store past negative conditionings }) # Extract conditioning tensors - current_pos_cond = positive[0][0] # Assuming standard ComfyUI conditioning format + current_pos_cond = positive[0][0] current_neg_cond = negative[0][0] + # Update context queues + if len(state["pos_context"]) >= context_size: + state["pos_context"].pop(0) + state["neg_context"].pop(0) + state["pos_context"].append(current_pos_cond.detach().clone()) + state["neg_context"].append(current_neg_cond.detach().clone()) + # First frame case if state["prev_positive"] is None: state["prev_positive"] = current_pos_cond.detach().clone() state["prev_negative"] = current_neg_cond.detach().clone() if cfg_type == "initialize": - # For initialize, we use the first negative as our stock noise state["stock_noise"] = current_neg_cond.detach().clone() elif cfg_type == "self": - # For self, we start with a scaled version of the negative state["stock_noise"] = current_neg_cond.detach().clone() * delta self.set_state(state) return (positive, negative) # Handle different CFG types if cfg_type == "full": - # Standard R-CFG with full negative conditioning - pos_residual = current_pos_cond - state["prev_positive"] - neg_residual = current_neg_cond - state["prev_negative"] + # Use entire context for smoother transitions + pos_context = torch.stack(state["pos_context"], dim=0) + neg_context = torch.stack(state["neg_context"], dim=0) + + # Calculate weighted residuals across context + weights = torch.linspace(0.5, 1.0, len(state["pos_context"]), device=current_pos_cond.device) + pos_residual = (current_pos_cond - pos_context) * weights.view(-1, 1, 1) + neg_residual = (current_neg_cond - neg_context) * weights.view(-1, 1, 1) + + # Average residuals + pos_residual = pos_residual.mean(dim=0) + neg_residual = neg_residual.mean(dim=0) blended_pos = current_pos_cond + residual_scale * pos_residual blended_neg = current_neg_cond + residual_scale * neg_residual @@ -300,156 +324,6 @@ class StreamConditioning(ControlNodeBase): return (positive_out, negative_out) -class StreamBatchSampler(ControlNodeBase): - """Implements batched denoising for faster inference by processing multiple steps in parallel""" - - RETURN_TYPES = ("SAMPLER",) - FUNCTION = "update" - CATEGORY = "real-time/sampling" - - @classmethod - def INPUT_TYPES(cls): - inputs = super().INPUT_TYPES() - inputs["required"].update({ - "batch_size": ("INT", { - "default": 2, - "min": 1, - "max": 10, - "step": 1, - "tooltip": "Number of steps to batch together. Higher values use more memory but are faster." - }), - }) - return inputs - - def __init__(self): - super().__init__() - self.batch_size = None - - def sample(self, model, noise, sigmas, extra_args=None, callback=None, disable=None): - """Sample with batched denoising steps""" - extra_args = {} if extra_args is None else extra_args - print(f"[StreamBatchSampler] Starting sampling with {len(sigmas)-1} steps, batch_size={self.batch_size}") - print(f"[StreamBatchSampler] Input noise shape: {noise.shape}, device: {noise.device}") - print(f"[StreamBatchSampler] Sigmas: {sigmas.tolist()}") - - # Prepare batched sampling - num_sigmas = len(sigmas) - 1 - num_batches = (num_sigmas + self.batch_size - 1) // self.batch_size - x = noise - - for batch_idx in range(num_batches): - # Get sigmas for this batch - start_idx = batch_idx * self.batch_size - end_idx = min(start_idx + self.batch_size, num_sigmas) - batch_sigmas = sigmas[start_idx:end_idx+1] - print(f"\n[StreamBatchSampler] Batch {batch_idx+1}/{num_batches}") - print(f"[StreamBatchSampler] Processing steps {start_idx}-{end_idx}") - print(f"[StreamBatchSampler] Batch sigmas: {batch_sigmas.tolist()}") - - # Create batch of identical latents - batch_size = end_idx - start_idx - x_batch = x.repeat(batch_size, 1, 1, 1) - - # Create batch of sigmas - sigma_batch = batch_sigmas[:-1] # All but last sigma - - # Run model on entire batch at once - with torch.no_grad(): - # Process all steps in parallel - denoised_batch = model(x_batch, sigma_batch, **extra_args) - print(f"[StreamBatchSampler] Denoised batch shape: {denoised_batch.shape}") - - # Process results one at a time to maintain callback - for i in range(batch_size): - sigma = sigma_batch[i] - sigma_next = batch_sigmas[i + 1] - denoised = denoised_batch[i:i+1] - - # Calculate step size (now always positive as we go from high to low sigma) - dt = sigma - sigma_next - - # Update x using Euler method - # The (denoised - x) term gives us the direction to move - # dt/sigma scales how far we move based on current noise level - x = x + (denoised - x) * (dt / sigma) - print(f"[StreamBatchSampler] Step {start_idx+i}: sigma={sigma:.4f}, next_sigma={sigma_next:.4f}, dt={dt:.4f}") - - # Call callback if provided - if callback is not None: - callback({'x': x, 'i': start_idx + i, 'sigma': sigma, 'sigma_hat': sigma, 'denoised': denoised}) - - print(f"\n[StreamBatchSampler] Sampling complete. Final x shape: {x.shape}") - return x - - def update(self, batch_size=2, always_execute=True): - """Create sampler with specified settings""" - self.batch_size = batch_size - sampler = comfy.samplers.KSAMPLER(self.sample) - return (sampler,) - - -class StreamScheduler(ControlNodeBase): - """Implements StreamDiffusion's efficient timestep selection""" - - RETURN_TYPES = ("SIGMAS",) - FUNCTION = "update" - CATEGORY = "real-time/sampling" - - @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. 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." - }), - }) - return inputs - - 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") - print(f"[StreamScheduler] Model sampling max sigma: {model_sampling.sigma_max}, min sigma: {model_sampling.sigma_min}") - - # Parse timestep list - try: - t_index_list = [int(t.strip()) for t in t_index_list.split(",")] - except ValueError as e: - print(f"Error parsing timesteps: {e}. Using default [32,45]") - t_index_list = [32, 45] - print(f"[StreamScheduler] Using timesteps: {t_index_list}") - - # Create full schedule using normal scheduler - full_sigmas = comfy.samplers.normal_scheduler(model_sampling, num_inference_steps) - print(f"[StreamScheduler] Full sigma schedule: {full_sigmas.tolist()}") - - # Select only the sigmas at our desired indices, but in reverse order - # This ensures we go from high noise to low noise - selected_sigmas = [] - for t in sorted(t_index_list, reverse=True): # Sort in reverse to go from high noise to low - 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])) - print(f"[StreamScheduler] Selected sigmas: {selected_sigmas}") - - # Add final sigma - selected_sigmas.append(0.0) - print(f"[StreamScheduler] Final sigma schedule: {selected_sigmas}") - - # Convert to tensor and move to appropriate device - selected_sigmas = torch.FloatTensor(selected_sigmas).to(comfy.model_management.get_torch_device()) - return (selected_sigmas,) - - class StreamCrossAttention(ControlNodeBase): """Implements optimized cross attention with KV-cache for real-time generation diff --git a/stream_sampler.py b/stream_sampler.py index 8269ed1..c59275e 100644 --- a/stream_sampler.py +++ b/stream_sampler.py @@ -6,7 +6,7 @@ import random class StreamBatchSampler(ControlNodeBase): - """Implements batched denoising for faster inference by processing multiple steps in parallel""" + """Implements batched denoising for faster inference by processing multiple frames in parallel at different denoising steps""" RETURN_TYPES = ("SAMPLER",) FUNCTION = "update" @@ -16,79 +16,100 @@ class StreamBatchSampler(ControlNodeBase): def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() inputs["required"].update({ - "batch_size": ("INT", { - "default": 2, + "num_steps": ("INT", { + "default": 4, "min": 1, "max": 10, "step": 1, - "tooltip": "Number of steps to batch together. Higher values use more memory but are faster." + "tooltip": "Number of denoising steps. Should match the frame buffer size." }), }) return inputs def __init__(self): super().__init__() - self.batch_size = None + self.num_steps = None + self.frame_buffer = [] + self.x_t_latent_buffer = None + self.stock_noise = None def sample(self, model, noise, sigmas, extra_args=None, callback=None, disable=None): - """Sample with batched denoising steps""" + """Sample with staggered batch denoising steps""" extra_args = {} if extra_args is None else extra_args - print(f"[StreamBatchSampler] Starting sampling with {len(sigmas)-1} steps, batch_size={self.batch_size}") - print(f"[StreamBatchSampler] Input noise shape: {noise.shape}, device: {noise.device}") - print(f"[StreamBatchSampler] Sigmas: {sigmas.tolist()}") + print(f"[StreamBatchSampler] Starting sampling with {len(sigmas)-1} steps") - # Prepare batched sampling - num_sigmas = len(sigmas) - 1 - num_batches = (num_sigmas + self.batch_size - 1) // self.batch_size - x = noise + # 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) - for batch_idx in range(num_batches): - # Get sigmas for this batch - start_idx = batch_idx * self.batch_size - end_idx = min(start_idx + self.batch_size, num_sigmas) - batch_sigmas = sigmas[start_idx:end_idx+1] - print(f"\n[StreamBatchSampler] Batch {batch_idx+1}/{num_batches}") - print(f"[StreamBatchSampler] Processing steps {start_idx}-{end_idx}") - print(f"[StreamBatchSampler] Batch sigmas: {batch_sigmas.tolist()}") + print(f"[StreamBatchSampler] Input sigmas: {sigmas}") + print(f"[StreamBatchSampler] Input noise shape: {noise.shape}, min: {noise.min():.3f}, max: {noise.max():.3f}") + + # 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) + + print(f"[StreamBatchSampler] Alpha values: {alpha_prod_t.view(-1)}") + print(f"[StreamBatchSampler] Beta values: {beta_prod_t.view(-1)}") + + # Initialize stock noise if needed + if self.stock_noise is None: + self.stock_noise = torch.randn_like(noise[0]) # Random noise instead of zeros + print(f"[StreamBatchSampler] Initialized random stock noise with shape: {self.stock_noise.shape}") - # Create batch of identical latents - batch_size = end_idx - start_idx - x_batch = x.repeat(batch_size, 1, 1, 1) + # Scale noise for each frame based on its sigma + scaled_noise = [] + for i in range(batch_size): + frame_noise = noise[i] + self.stock_noise * sigmas[i] # Add scaled noise to input + scaled_noise.append(frame_noise) + x = torch.stack(scaled_noise, dim=0) + print(f"[StreamBatchSampler] Scaled noise shape: {x.shape}, min: {x.min():.3f}, max: {x.max():.3f}") - # Create batch of sigmas - sigma_batch = batch_sigmas[:-1] # All but last sigma + # Initialize frame buffer if needed + if self.x_t_latent_buffer is None and num_sigmas > 1: + self.x_t_latent_buffer = x[0].clone() # Initialize with noised first frame + print(f"[StreamBatchSampler] Initialized buffer with shape: {self.x_t_latent_buffer.shape}") - # Run model on entire batch at once - with torch.no_grad(): - # Process all steps in parallel - denoised_batch = model(x_batch, sigma_batch, **extra_args) - print(f"[StreamBatchSampler] Denoised batch shape: {denoised_batch.shape}") + # Use buffer for first frame to maintain temporal consistency + if num_sigmas > 1: + x = torch.cat([self.x_t_latent_buffer.unsqueeze(0), x[1:]], dim=0) + print(f"[StreamBatchSampler] Combined with buffer, shape: {x.shape}") + + # Run model on entire batch at once + with torch.no_grad(): + # Process all frames in parallel + sigma_batch = sigmas[:-1] + print(f"[StreamBatchSampler] Using sigmas for denoising: {sigma_batch}") + + denoised_batch = model(x, sigma_batch, **extra_args) + print(f"[StreamBatchSampler] Denoised batch shape: {denoised_batch.shape}") + print(f"[StreamBatchSampler] Denoised stats - min: {denoised_batch.min():.3f}, max: {denoised_batch.max():.3f}") + + # Update buffer with intermediate results + if num_sigmas > 1: + # Store result from first frame as buffer for next iteration + self.x_t_latent_buffer = denoised_batch[0].clone() + print(f"[StreamBatchSampler] Updated buffer with shape: {self.x_t_latent_buffer.shape}") - # Process results one at a time to maintain callback - for i in range(batch_size): - sigma = sigma_batch[i] - sigma_next = batch_sigmas[i + 1] - denoised = denoised_batch[i:i+1] - - # Calculate step size (now always positive as we go from high to low sigma) - dt = sigma - sigma_next - - # Update x using Euler method - # The (denoised - x) term gives us the direction to move - # dt/sigma scales how far we move based on current noise level - x = x + (denoised - x) * (dt / sigma) - print(f"[StreamBatchSampler] Step {start_idx+i}: sigma={sigma:.4f}, next_sigma={sigma_next:.4f}, dt={dt:.4f}") - - # Call callback if provided - if callback is not None: - callback({'x': x, 'i': start_idx + i, 'sigma': sigma, 'sigma_hat': sigma, 'denoised': denoised}) + # 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 + + # 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:]}) - print(f"\n[StreamBatchSampler] Sampling complete. Final x shape: {x.shape}") - return x + return x_0_pred_out - def update(self, batch_size=2, always_execute=True): + def update(self, num_steps=4, always_execute=True): """Create sampler with specified settings""" - self.batch_size = batch_size + self.num_steps = num_steps sampler = comfy.samplers.KSAMPLER(self.sample) return (sampler,) @@ -122,7 +143,6 @@ class StreamScheduler(ControlNodeBase): 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") - print(f"[StreamScheduler] Model sampling max sigma: {model_sampling.sigma_max}, min sigma: {model_sampling.sigma_min}") # Parse timestep list try: @@ -130,11 +150,9 @@ class StreamScheduler(ControlNodeBase): except ValueError as e: print(f"Error parsing timesteps: {e}. Using default [32,45]") t_index_list = [32, 45] - print(f"[StreamScheduler] Using timesteps: {t_index_list}") - + # Create full schedule using normal scheduler full_sigmas = comfy.samplers.normal_scheduler(model_sampling, num_inference_steps) - print(f"[StreamScheduler] Full sigma schedule: {full_sigmas.tolist()}") # Select only the sigmas at our desired indices, but in reverse order # This ensures we go from high noise to low noise @@ -144,12 +162,65 @@ class StreamScheduler(ControlNodeBase): print(f"Warning: timestep {t} out of range [0,{num_inference_steps}), skipping") continue selected_sigmas.append(float(full_sigmas[t])) - print(f"[StreamScheduler] Selected sigmas: {selected_sigmas}") - + # Add final sigma selected_sigmas.append(0.0) - print(f"[StreamScheduler] Final sigma schedule: {selected_sigmas}") # Convert to tensor and move to appropriate device selected_sigmas = torch.FloatTensor(selected_sigmas).to(comfy.model_management.get_torch_device()) return (selected_sigmas,) + + +class StreamFrameBuffer(ControlNodeBase): + """Accumulates frames to enable staggered batch denoising like StreamDiffusion""" + + RETURN_TYPES = ("LATENT",) + FUNCTION = "update" + CATEGORY = "real-time/sampling" + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "latent": ("LATENT",), + "buffer_size": ("INT", { + "default": 4, + "min": 1, + "max": 10, + "step": 1, + "tooltip": "Number of frames to buffer before starting batch processing. Should match number of denoising steps." + }), + }) + return inputs + + def __init__(self): + super().__init__() + self.frame_buffer = [] # List to store incoming frames + self.buffer_size = None + + def update(self, latent, buffer_size=4, always_execute=True): + """Add new frame to buffer and return batch when ready""" + self.buffer_size = buffer_size + + # 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] + + # Add new frame to buffer + if len(self.frame_buffer) == 0: + # First frame - initialize buffer with copies + self.frame_buffer = [x.clone() for _ in range(self.buffer_size)] + print(f"[StreamFrameBuffer] Initialized buffer with {self.buffer_size} copies of first frame") + else: + # Shift frames forward and add new frame + self.frame_buffer.pop(0) # Remove oldest frame + self.frame_buffer.append(x.clone()) # Add new frame + print(f"[StreamFrameBuffer] Added new frame to buffer") + + # Stack frames into batch + batch = torch.stack(self.frame_buffer, dim=0) # [B,C,H,W] + print(f"[StreamFrameBuffer] Created batch with shape: {batch.shape}") + + # Return as latent dict + return ({"samples": batch},)