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""" RETURN_TYPES = ("MODEL",) FUNCTION = "update" CATEGORY = "real-time/sampling" @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() inputs["required"].update({ "model": ("MODEL",), "cfg_type": (["self", "full", "initialize"], { "default": "self", "tooltip": "Type of CFG to use: full (standard), self (memory efficient), or initialize (memory efficient with initialization)" }), "residual_scale": ("FLOAT", { "default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Scale factor for residual (higher = more temporal consistency)" }), "delta": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1, "tooltip": "Delta parameter for self/initialize CFG types" }), "residual_decay": ("FLOAT", { "default": 0.6, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Decay factor for residual scale over time (1.0 = no decay)" }), "noise_decay": ("FLOAT", { "default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Decay factor for noise injection over time (1.0 = no decay)" }), "reset_interval": ("INT", { "default": 0, "min": 0, "max": 100, "step": 1, "tooltip": "Reset stock noise every N frames (0 = never reset)" }), }) return inputs def __init__(self): super().__init__() # Store the last model to detect when we need to reapply the hook self.last_model_hash = None self.post_cfg_function = None def update(self, model, always_execute=True, cfg_type="self", residual_scale=0.4, delta=1.0, residual_decay=0.95, noise_decay=0.8, reset_interval=0): print(f"[StreamCFG] Initializing with cfg_type={cfg_type}, residual_scale={residual_scale}, delta={delta}, residual_decay={residual_decay}, noise_decay={noise_decay}, reset_interval={reset_interval}") # Get state with defaults state = self.get_state({ "last_uncond": None, "initialized": False, "cfg_type": cfg_type, "residual_scale": residual_scale, "delta": delta, "workflow_count": 0, "current_sigmas": None, "seen_sigmas": set(), "is_last_step": False, # Add new state variables for proper scaling "alpha_prod_t": None, "beta_prod_t": None, "c_skip": None, "c_out": None, "residual_decay": residual_decay, "noise_decay": noise_decay, "reset_interval": reset_interval, }) def post_cfg_function(args): # Extract info denoised = args["denoised"] cond = args["cond"] uncond = args["uncond"] cond_denoised = args["cond_denoised"] uncond_denoised = args["uncond_denoised"] cond_scale = args["cond_scale"] # Debug prints for tensor stats print(f"\n[StreamCFG Debug] Step Info:") print(f"- Workflow count: {state['workflow_count']}") print(f"- CFG Type: {state['cfg_type']}") print(f"- Tensor Stats:") print(f" - denoised shape: {denoised.shape}, range: [{denoised.min():.3f}, {denoised.max():.3f}]") print(f" - uncond_denoised shape: {uncond_denoised.shape}, range: [{uncond_denoised.min():.3f}, {uncond_denoised.max():.3f}]") if state["last_uncond"] is not None: print(f" - last_uncond shape: {state['last_uncond'].shape}, range: [{state['last_uncond'].min():.3f}, {state['last_uncond'].max():.3f}]") # Handle both batched and single sigmas sigma = args["sigma"] if torch.is_tensor(sigma): sigma = sigma[0].item() if len(sigma.shape) > 0 else sigma.item() print(f"- Current sigma: {sigma:.6f}") # Get step info from model options model_options = args["model_options"] sample_sigmas = model_options["transformer_options"].get("sample_sigmas", None) # Update current sigmas if needed if sample_sigmas is not None and state["current_sigmas"] is None: sigmas = [s.item() for s in sample_sigmas] if sigmas[-1] == 0.0: sigmas = sigmas[:-1] state["current_sigmas"] = sigmas state["seen_sigmas"] = set() print(f"- New sigma sequence: {sigmas}") # Calculate paper's exact scaling factors state["alpha_prod_t"] = torch.tensor([1.0 / (1.0 + s**2) for s in sigmas], device=denoised.device, dtype=denoised.dtype) state["beta_prod_t"] = torch.tensor([s / (1.0 + s**2) for s in sigmas], device=denoised.device, dtype=denoised.dtype) # Calculate c_skip and c_out coefficients state["c_skip"] = torch.tensor([1.0 / (s**2 + 1.0) for s in sigmas], device=denoised.device, dtype=denoised.dtype) state["c_out"] = torch.tensor([-s / torch.sqrt(torch.tensor(s**2 + 1.0)) for s in sigmas], device=denoised.device, dtype=denoised.dtype) print(f"- Scaling factors for first step:") print(f" alpha: {state['alpha_prod_t'][0]:.6f}") print(f" beta: {state['beta_prod_t'][0]:.6f}") print(f" c_skip: {state['c_skip'][0]:.6f}") print(f" c_out: {state['c_out'][0]:.6f}") # Track this sigma state["seen_sigmas"].add(sigma) # Check if this is the last step state["is_last_step"] = False if state["current_sigmas"] is not None: # It's the last step if we've seen all sigmas is_last_step = len(state["seen_sigmas"]) >= len(state["current_sigmas"]) # Or if this is the smallest sigma in the sequence if not is_last_step and sigma == min(state["current_sigmas"]): is_last_step = True state["is_last_step"] = is_last_step print(f"- Is last step: {is_last_step}") print(f"- Seen sigmas: {sorted(state['seen_sigmas'])}") # First workflow case if state["last_uncond"] is None: if state["is_last_step"]: state["last_uncond"] = uncond_denoised.detach().clone() state["workflow_count"] += 1 state["current_sigmas"] = None if cfg_type == "initialize": state["initialized"] = True self.set_state(state) print("- First workflow complete, stored last_uncond") return denoised # Handle different CFG types if cfg_type == "full": result = denoised elif cfg_type == "initialize" and not state["initialized"]: result = denoised if state["is_last_step"]: state["initialized"] = True self.set_state(state) else: # self or initialized initialize current_idx = len(state["seen_sigmas"]) - 1 print(f"- Current step index: {current_idx}") # Apply decaying residual scale current_residual_scale = state["residual_scale"] * (state["residual_decay"] ** state["workflow_count"]) print(f"- Current residual scale: {current_residual_scale:.6f}") # Calculate temporal mixing weights current_weight = 1.0 # Base weight for current prediction temporal_weight = current_residual_scale * state["delta"] # Weight for temporal consistency # Normalize weights to ensure proper blending total_weight = current_weight + temporal_weight current_weight = current_weight / total_weight temporal_weight = temporal_weight / total_weight print(f"- Mixing weights - current: {current_weight:.3f}, temporal: {temporal_weight:.3f}") # Blend predictions with normalized weights noise_pred_uncond = current_weight * uncond_denoised + temporal_weight * state["last_uncond"] print(f"- Blended noise prediction range: [{noise_pred_uncond.min():.3f}, {noise_pred_uncond.max():.3f}]") # Apply CFG normally result = noise_pred_uncond + cond_scale * (cond_denoised - noise_pred_uncond) print(f"- Result range after CFG: [{result.min():.3f}, {result.max():.3f}]") # Store last prediction if this is the last step if state["is_last_step"]: # Calculate F_theta using paper's formulation F_theta = (uncond_denoised - state["beta_prod_t"][current_idx] * noise_pred_uncond) / state["alpha_prod_t"][current_idx] print(f"- F_theta range: [{F_theta.min():.3f}, {F_theta.max():.3f}]") delta_x = state["c_out"][current_idx] * F_theta + state["c_skip"][current_idx] * uncond_denoised print(f"- delta_x range: [{delta_x.min():.3f}, {delta_x.max():.3f}]") # Scale delta_x with next step's coefficients if current_idx < len(state["current_sigmas"]) - 1: next_alpha = state["alpha_prod_t"][current_idx + 1] next_beta = state["beta_prod_t"][current_idx + 1] else: next_alpha = torch.ones_like(state["alpha_prod_t"][0]) next_beta = torch.zeros_like(state["beta_prod_t"][0]) print(f"- Next step coefficients - alpha: {next_alpha:.6f}, beta: {next_beta:.6f}") # Update stored prediction with properly scaled residual if next_beta > 0: final_update = (next_alpha * delta_x) / next_beta # Add decaying noise when beta > 0 noise_scale = state["noise_decay"] ** state["workflow_count"] noise = torch.randn_like(delta_x) * (1 - next_alpha**2).sqrt() * noise_scale final_update = final_update + noise print(f"- Added noise (scale={noise_scale:.6f}) range: [{noise.min():.3f}, {noise.max():.3f}]") else: final_update = uncond_denoised print(f"- Final update range: [{final_update.min():.3f}, {final_update.max():.3f}]") # Reset stock noise if interval is reached if state["reset_interval"] > 0 and state["workflow_count"] % state["reset_interval"] == 0: print(f"- Resetting stock noise at workflow {state['workflow_count']}") state["last_uncond"] = uncond_denoised.detach().clone() else: state["last_uncond"] = final_update state["workflow_count"] += 1 state["current_sigmas"] = None self.set_state(state) return result # Store function reference to prevent garbage collection self.post_cfg_function = post_cfg_function # Only set up post CFG function if model has changed model_hash = hash(str(model)) if model_hash != self.last_model_hash: m = model.clone() m.model_options = m.model_options.copy() m.model_options["sampler_post_cfg_function"] = [self.post_cfg_function] self.last_model_hash = model_hash return (m,) # Make sure our function is still in the list if not any(f is self.post_cfg_function for f in model.model_options.get("sampler_post_cfg_function", [])): m = model.clone() m.model_options = m.model_options.copy() m.model_options["sampler_post_cfg_function"] = [self.post_cfg_function] return (m,) 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 @classmethod def INPUT_TYPES(s): return { "required": { "positive": ("CONDITIONING",), "negative": ("CONDITIONING",), "cfg_type": (["full", "self", "initialize"], { "default": "full", "tooltip": "Type of CFG to use: full (standard), self (memory efficient), or initialize (memory efficient with initialization)" }), "residual_scale": ("FLOAT", { "default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Scale factor for residual conditioning (higher = more temporal consistency)" }), "delta": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 5.0, "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, }), } } RETURN_TYPES = ("CONDITIONING", "CONDITIONING") RETURN_NAMES = ("positive", "negative") FUNCTION = "update" CATEGORY = "real-time/control/utility" def __init__(self): super().__init__() 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 "pos_context": [], # Store past positive conditionings "neg_context": [] # Store past negative conditionings }) # Extract conditioning tensors 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": state["stock_noise"] = current_neg_cond.detach().clone() elif cfg_type == "self": state["stock_noise"] = current_neg_cond.detach().clone() * delta self.set_state(state) return (positive, negative) # Handle different CFG types if cfg_type == "full": # 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 # Update state state["prev_positive"] = current_pos_cond.detach().clone() state["prev_negative"] = current_neg_cond.detach().clone() # Reconstruct conditioning format positive_out = [[blended_pos, positive[0][1]]] negative_out = [[blended_neg, negative[0][1]]] else: # self or initialize # Calculate residual for positive conditioning pos_residual = current_pos_cond - state["prev_positive"] blended_pos = current_pos_cond + residual_scale * pos_residual # Update stock noise based on current prediction if cfg_type == "initialize" and not state["initialized"]: # First prediction for initialize type state["stock_noise"] = current_neg_cond.detach().clone() state["initialized"] = True else: # Update stock noise with temporal consistency stock_residual = current_neg_cond - state["stock_noise"] state["stock_noise"] = current_neg_cond + residual_scale * stock_residual # Scale stock noise by delta scaled_stock = state["stock_noise"] * delta # Update state state["prev_positive"] = current_pos_cond.detach().clone() state["prev_negative"] = scaled_stock.detach().clone() # Reconstruct conditioning format positive_out = [[blended_pos, positive[0][1]]] negative_out = [[scaled_stock, negative[0][1]]] self.set_state(state) return (positive_out, negative_out) class StreamCrossAttention(ControlNodeBase): """Implements optimized cross attention with KV-cache for real-time generation Paper reference: StreamDiffusion Section 3.5 "Pre-computation" - Pre-computes and caches prompt embeddings - Stores Key-Value pairs for reuse with static prompts - Only recomputes KV pairs when prompt changes Additional optimizations beyond paper: - QK normalization for better numerical stability - Rotary position embeddings (RoPE) for improved temporal consistency - Configurable context window size for memory/quality tradeoff """ 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": "Additional optimization: Whether to apply layer normalization to query and key tensors" }), "use_rope": ("BOOLEAN", { "default": True, "tooltip": "Additional optimization: Whether to use rotary position embeddings for better temporal consistency" }), "context_size": ("INT", { "default": 4, "min": 1, "max": 32, "step": 1, "tooltip": "Additional optimization: Maximum number of past frames to keep in context. Higher values use more memory but may improve temporal consistency." }), "use_kv_cache": ("BOOLEAN", { "default": True, "tooltip": "Paper Section 3.5: Whether to cache key-value pairs for static prompts to avoid recomputation" }), }) 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, context_size=4, use_kv_cache=True): print(f"[StreamCrossAttention] Initializing with qk_norm={qk_norm}, use_rope={use_rope}, context_size={context_size}, use_kv_cache={use_kv_cache}") # Get state with defaults state = self.get_state({ "qk_norm": qk_norm, # Additional optimization "use_rope": use_rope, # Additional optimization "context_size": context_size, # Additional optimization "use_kv_cache": use_kv_cache, # From paper Section 3.5 "workflow_count": 0, "context_queue": [], # Additional: Store past context tensors for temporal consistency "kv_cache": {}, # From paper Section 3.5: Cache KV pairs for each prompt "last_prompt_embeds": None, # From paper Section 3.5: For cache validation }) 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 # Paper Section 3.5: KV Caching Logic cache_hit = False if state["use_kv_cache"] and state["last_prompt_embeds"] is not None: # Compare current context with cached prompt embeddings if torch.allclose(context, state["last_prompt_embeds"], rtol=1e-5, atol=1e-5): cache_hit = True k, v = state["kv_cache"].get(module, (None, None)) if k is not None and v is not None: print("[StreamCrossAttention] Using cached KV pairs") if not cache_hit: # Additional optimization: Temporal context management if len(state["context_queue"]) >= state["context_size"]: state["context_queue"].pop(0) state["context_queue"].append(context.detach().clone()) # Additional optimization: Use past context for temporal consistency full_context = torch.cat(state["context_queue"], dim=1) # Generate k/v for full context k = module.to_k(full_context) v = value if value is not None else module.to_v(full_context) # Paper Section 3.5: Cache KV pairs for static prompts if state["use_kv_cache"]: state["last_prompt_embeds"] = context.detach().clone() state["kv_cache"][module] = (k.detach().clone(), v.detach().clone()) # Additional optimization: QK normalization 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) # Additional optimization: Rotary position embeddings if state["use_rope"]: # Calculate position embeddings batch_size = q.shape[0] seq_len = q.shape[1] full_seq_len = k.shape[1] # Use full context length for k/v dim = q.shape[2] # Create position indices for q and k separately q_position = torch.arange(seq_len, device=q.device).unsqueeze(0).unsqueeze(-1) k_position = torch.arange(full_seq_len, device=k.device).unsqueeze(0).unsqueeze(-1) q_position = q_position.repeat(batch_size, 1, dim//2) k_position = k_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 q_theta = q_position * freq k_theta = k_position * freq # Apply rotations to q q_cos = torch.cos(q_theta) q_sin = torch.sin(q_theta) q_reshaped = q.view(*q.shape[:-1], -1, 2) q_out = torch.cat([ q_reshaped[..., 0] * q_cos - q_reshaped[..., 1] * q_sin, q_reshaped[..., 0] * q_sin + q_reshaped[..., 1] * q_cos ], dim=-1) # Apply rotations to k k_cos = torch.cos(k_theta) k_sin = torch.sin(k_theta) k_reshaped = k.view(*k.shape[:-1], -1, 2) k_out = torch.cat([ k_reshaped[..., 0] * k_cos - k_reshaped[..., 1] * k_sin, k_reshaped[..., 0] * k_sin + k_reshaped[..., 1] * k_cos ], dim=-1) q = q_out k = k_out # Standard attention computation with memory-efficient access pattern batch_size = q.shape[0] q_seq_len = q.shape[1] k_seq_len = k.shape[1] head_dim = q.shape[-1] // module.heads # Reshape for multi-head attention q = q.view(batch_size, q_seq_len, module.heads, head_dim) k = k.view(batch_size, k_seq_len, module.heads, head_dim) v = v.view(batch_size, k_seq_len, 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, q_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,)