From 135d012a0fad2acb8fdb3c1d47945305412de6ea Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Mon, 10 Feb 2025 16:53:42 -0500 Subject: [PATCH] streambatchsampler, streambatch scheduler, stream corss attention --- __init__.py | 7 +- stream_diffusion_nodes.py | 321 +++++++++++++++++++++++++++++++++++++- 2 files changed, 326 insertions(+), 2 deletions(-) diff --git a/__init__.py b/__init__.py index 5671113..2ee77c9 100644 --- a/__init__.py +++ b/__init__.py @@ -3,7 +3,7 @@ from .controls.sequence_controls import FloatSequence, IntSequence, StringSequen from .controls.utility_controls import FPSMonitor, SimilarityFilter, LazyCondition 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 +from .stream_diffusion_nodes import StreamConditioning, StreamCFG, StreamBatchSampler, StreamScheduler, StreamCrossAttention import re @@ -19,12 +19,17 @@ NODE_CLASS_MAPPINGS = { "SimilarityFilter": SimilarityFilter, "StreamCFG": StreamCFG, "StreamConditioning": StreamConditioning, + "StreamBatchSampler": StreamBatchSampler, + "StreamScheduler": StreamScheduler, + "StreamCrossAttention": StreamCrossAttention, "LazyCondition": LazyCondition, "MotionController": MotionController, "IntegerMotionController": IntegerMotionController, "YOLOSimilarityCompare": YOLOSimilarityCompare, "TextRenderer": TextRenderer, + + "ROINode": ROINode, #"IntervalControl": IntervalCo ntrol, diff --git a/stream_diffusion_nodes.py b/stream_diffusion_nodes.py index 4ccb605..bc13a65 100644 --- a/stream_diffusion_nodes.py +++ b/stream_diffusion_nodes.py @@ -1,5 +1,8 @@ import torch from .base.control_base import ControlNodeBase +import comfy.model_management +import comfy.samplers +import random class StreamCFG(ControlNodeBase): """Implements CFG approaches for temporal consistency between workflow runs""" @@ -66,7 +69,13 @@ class StreamCFG(ControlNodeBase): cond_denoised = args["cond_denoised"] uncond_denoised = args["uncond_denoised"] cond_scale = args["cond_scale"] - sigma = args["sigma"].item() if torch.is_tensor(args["sigma"]) else args["sigma"] + + # Handle both batched and single sigmas + sigma = args["sigma"] + if torch.is_tensor(sigma): + # For batched sampling, use first sigma in batch + # This is safe because we process in order and track seen sigmas + sigma = sigma[0].item() if len(sigma.shape) > 0 else sigma.item() # Get step info from model options model_options = args["model_options"] @@ -291,4 +300,314 @@ 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 for real-time generation""" + + RETURN_TYPES = ("MODEL",) + FUNCTION = "update" + CATEGORY = "real-time/sampling" + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "model": ("MODEL",), + "qk_norm": ("BOOLEAN", { + "default": True, + "tooltip": "Whether to apply layer normalization to query and key tensors" + }), + "use_rope": ("BOOLEAN", { + "default": True, + "tooltip": "Whether to use rotary position embeddings for better temporal consistency" + }), + }) + return inputs + + def __init__(self): + super().__init__() + self.last_model_hash = None + self.cross_attention_hook = None + + def update(self, model, always_execute=True, qk_norm=True, use_rope=True): + print(f"[StreamCrossAttention] Initializing with qk_norm={qk_norm}, use_rope={use_rope}") + + # Get state with defaults + state = self.get_state({ + "qk_norm": qk_norm, + "use_rope": use_rope, + "workflow_count": 0, + }) + + def cross_attention_forward(module, x, context=None, mask=None, value=None): + q = module.to_q(x) + context = x if context is None else context + k = module.to_k(context) + # Use provided value tensor if given, otherwise compute it + v = value if value is not None else module.to_v(context) + + # Apply QK normalization if enabled + if state["qk_norm"]: + q_norm = torch.nn.LayerNorm(q.shape[-1], device=q.device, dtype=q.dtype) + k_norm = torch.nn.LayerNorm(k.shape[-1], device=k.device, dtype=k.dtype) + q = q_norm(q) + k = k_norm(k) + + # Apply rotary embeddings if enabled + if state["use_rope"]: + # Calculate position embeddings + batch_size = q.shape[0] + seq_len = q.shape[1] + dim = q.shape[2] + + # Create position indices + position = torch.arange(seq_len, device=q.device).unsqueeze(0).unsqueeze(-1) + position = position.repeat(batch_size, 1, dim//2) + + # Calculate frequencies + freq = 10000.0 ** (-torch.arange(0, dim//2, 2, device=q.device) / dim) + freq = freq.repeat((dim + 1) // 2)[:dim//2] + + # Calculate rotation angles + theta = position * freq + + # Apply rotations + cos = torch.cos(theta) + sin = torch.sin(theta) + + # Reshape q and k for rotation + q_reshaped = q.view(*q.shape[:-1], -1, 2) + k_reshaped = k.view(*k.shape[:-1], -1, 2) + + # Apply rotations + q_out = torch.cat([ + q_reshaped[..., 0] * cos - q_reshaped[..., 1] * sin, + q_reshaped[..., 0] * sin + q_reshaped[..., 1] * cos + ], dim=-1) + + k_out = torch.cat([ + k_reshaped[..., 0] * cos - k_reshaped[..., 1] * sin, + k_reshaped[..., 0] * sin + k_reshaped[..., 1] * cos + ], dim=-1) + + q = q_out + k = k_out + + # Compute attention with optimized memory access pattern + batch_size, seq_len = q.shape[0], q.shape[1] + head_dim = q.shape[-1] // module.heads + + # Reshape for multi-head attention + q = q.view(batch_size, seq_len, module.heads, head_dim) + k = k.view(batch_size, -1, module.heads, head_dim) + v = v.view(batch_size, -1, module.heads, head_dim) + + # Transpose for attention computation + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + # Compute attention scores + scale = head_dim ** -0.5 + scores = torch.matmul(q, k.transpose(-2, -1)) * scale + + if mask is not None: + scores = scores + mask + + # Apply attention + attn = torch.softmax(scores, dim=-1) + out = torch.matmul(attn, v) + + # Reshape back + out = out.transpose(1, 2).contiguous() + out = out.view(batch_size, seq_len, -1) + + # Project back to original dimension + out = module.to_out[0](out) + + return out + + def hook_cross_attention(module, input, output): + if isinstance(module, torch.nn.Module) and hasattr(module, "to_q"): + # Store original forward + if not hasattr(module, "_original_forward"): + module._original_forward = module.forward + # Replace with our optimized version + module.forward = lambda *args, **kwargs: cross_attention_forward(module, *args, **kwargs) + return output + + # Only set up hooks if model has changed + model_hash = hash(str(model)) + if model_hash != self.last_model_hash: + m = model.clone() + + # Remove old hooks if they exist + if self.cross_attention_hook is not None: + self.cross_attention_hook.remove() + + # Register hook for cross attention modules + def register_hooks(module): + if isinstance(module, torch.nn.Module) and hasattr(module, "to_q"): + self.cross_attention_hook = module.register_forward_hook(hook_cross_attention) + + m.model.apply(register_hooks) + self.last_model_hash = model_hash + return (m,) + + return (model,) + + +