diff --git a/__init__.py b/__init__.py index 040a383..88a4b50 100644 --- a/__init__.py +++ b/__init__.py @@ -3,8 +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 StreamCrossAttention -from .stream_sampler import StreamBatchSampler, StreamScheduler, StreamFrameBuffer +from .stream_sampler import StreamBatchSampler, StreamFrameBuffer, StreamScheduler from .stream_cfg import StreamCFG from .stream_conditioning import StreamConditioning import re @@ -23,9 +22,8 @@ NODE_CLASS_MAPPINGS = { "StreamCFG": StreamCFG, "StreamConditioning": StreamConditioning, "StreamBatchSampler": StreamBatchSampler, - "StreamScheduler": StreamScheduler, "StreamFrameBuffer": StreamFrameBuffer, - "StreamCrossAttention": StreamCrossAttention, + "StreamScheduler": StreamScheduler, "LazyCondition": LazyCondition, "MotionController": MotionController, "IntegerMotionController": IntegerMotionController, diff --git a/stream_sampler.py b/stream_sampler.py index c59275e..0915bc9 100644 --- a/stream_sampler.py +++ b/stream_sampler.py @@ -3,6 +3,7 @@ from .base.control_base import ControlNodeBase import comfy.model_management import comfy.samplers import random +import math class StreamBatchSampler(ControlNodeBase): @@ -23,7 +24,54 @@ class StreamBatchSampler(ControlNodeBase): "step": 1, "tooltip": "Number of denoising steps. Should match the frame buffer size." }), + "use_cfg": ("BOOLEAN", { + "default": False, + "tooltip": "Whether to use classifier-free guidance optimizations" + }), + "cfg_type": (["none", "full", "self", "initialize"], { + "default": "self", + "tooltip": "Type of CFG to use: none (no guidance), full (standard), self (memory efficient), or initialize (memory efficient with initialization)" + }), + "guidance_scale": ("FLOAT", { + "default": 1.2, + "min": 1.0, + "max": 20.0, + "step": 0.1, + "tooltip": "CFG guidance scale - 1.0 means no guidance" + }), + "delta": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 5.0, + "step": 0.1, + "tooltip": "Delta parameter for self/initialize CFG types" + }), + "use_similarity_filter": ("BOOLEAN", { + "default": False, + "tooltip": "Whether to use similarity filtering to skip similar frames" + }), + "similarity_threshold": ("FLOAT", { + "default": 0.98, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Threshold for similarity filtering - higher means more frames are skipped" + }), + "max_skip_frames": ("INT", { + "default": 10, + "min": 1, + "max": 100, + "step": 1, + "tooltip": "Maximum number of consecutive frames to skip" + }), }) + inputs["optional"] = { + "prompt": ("STRING", { + "multiline": True, + "default": "", + "tooltip": "Text prompt to use for cross-attention optimization. If provided, enables KV caching." + }), + } return inputs def __init__(self): @@ -32,12 +80,150 @@ class StreamBatchSampler(ControlNodeBase): self.frame_buffer = [] self.x_t_latent_buffer = None self.stock_noise = None + # Add CFG state + self.use_cfg = False + 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.use_similarity_filter = False + self.similarity_filter = None + def process_with_cfg(self, model, x, sigma_batch, extra_args): + """Apply classifier-free guidance based on current cfg_type""" + if not self.use_cfg or self.cfg_type == "none" or self.guidance_scale <= 1.0: + # No guidance - process batch normally + return model(x, sigma_batch, **extra_args) + + # Apply CFG based on type + if self.cfg_type == "full": + # 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) + elif self.cfg_type == "initialize": + # 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() + else: # self + # Use stock noise for uncond + model_output = model(x, sigma_batch, **extra_args) + noise_pred_text = model_output + noise_pred_uncond = self.stock_noise.unsqueeze(0).expand_as(noise_pred_text) * self.delta + + # Combine predictions with guidance + return noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + + 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 + + print(f"[StreamBatchSampler] Setting up cross attention with prompt: {prompt}") + + def optimized_attention_forward(module, x, context=None, mask=None): + """Optimized cross attention with KV caching""" + 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: + 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: + 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) + + # Reshape for attention + head_dim = q.shape[-1] // module.heads + q = q.view(batch_size, -1, module.heads, head_dim).transpose(1, 2) + k = k.view(-1, k.shape[1], module.heads, head_dim).transpose(1, 2) + v = v.view(-1, v.shape[1], module.heads, head_dim).transpose(1, 2) + + # Compute attention + scale = 1.0 / math.sqrt(head_dim) + scores = torch.matmul(q, k.transpose(-2, -1)) * scale + + if mask is not None: + scores = scores + mask + + attn = torch.softmax(scores, dim=-1) + out = torch.matmul(attn, v) + + # Reshape output + out = out.transpose(1, 2).reshape(batch_size, -1, module.heads * head_dim) + return module.to_out[0](out) + + 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 + + # Get the actual model from the sampler + if hasattr(model, "model"): + unet = model.model + else: + print("[StreamBatchSampler] Warning: Could not find UNet model in sampler") + return model + + # Only set up hooks if prompt has changed + model_hash = hash(str(unet)) + if model_hash != self.last_model_hash or self.last_prompt != prompt: + print("[StreamBatchSampler] Setting up new cross attention hooks") + # Remove old hooks if they exist + if self.cross_attention_hook is not None: + self.cross_attention_hook.remove() + + # 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_model_hash = model_hash + + return model + 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"[StreamBatchSampler] Starting sampling with {len(sigmas)-1} steps") + # Set up cross attention optimization if prompt provided + if self.prompt: + model = self.setup_cross_attention(model, self.prompt) + print("[StreamBatchSampler] Cross attention optimization enabled") + else: + print("[StreamBatchSampler] Cross attention optimization disabled (no prompt)") + + if self.use_cfg: + print(f"[StreamBatchSampler] CFG enabled - type: {self.cfg_type}, scale: {self.guidance_scale}, delta: {self.delta}") + # 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) @@ -85,7 +271,12 @@ class StreamBatchSampler(ControlNodeBase): sigma_batch = sigmas[:-1] print(f"[StreamBatchSampler] Using sigmas for denoising: {sigma_batch}") - denoised_batch = model(x, sigma_batch, **extra_args) + # Process with or without CFG + if self.use_cfg: + denoised_batch = self.process_with_cfg(model, x, sigma_batch, extra_args) + else: + 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}") @@ -101,15 +292,46 @@ class StreamBatchSampler(ControlNodeBase): x_0_pred_out = denoised_batch self.x_t_latent_buffer = None + # Apply similarity filter if enabled + if self.use_similarity_filter: + if self.similarity_filter is None: + from .controls.similar_image_filter import SimilarImageFilter + self.similarity_filter = SimilarImageFilter( + threshold=self.similarity_threshold, + max_skip_frame=self.max_skip_frames + ) + filtered_out = self.similarity_filter(x_0_pred_out) + if filtered_out is not None: + x_0_pred_out = filtered_out + print("[StreamBatchSampler] Frame passed similarity filter") + else: + print("[StreamBatchSampler] Frame skipped by similarity filter") + # 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:]}) return x_0_pred_out - def update(self, num_steps=4, always_execute=True): + def update(self, num_steps=4, use_cfg=False, cfg_type="self", guidance_scale=1.2, delta=1.0, prompt="", + use_similarity_filter=False, similarity_threshold=0.98, max_skip_frames=10, always_execute=True): """Create sampler with specified settings""" self.num_steps = num_steps + self.use_cfg = use_cfg + self.cfg_type = cfg_type + self.guidance_scale = guidance_scale + self.delta = delta + self.prompt = prompt + + # Update similarity filter settings + self.use_similarity_filter = use_similarity_filter + 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 without extra_args - we use self.prompt directly in sample() sampler = comfy.samplers.KSAMPLER(self.sample) return (sampler,)