diff --git a/stream_attention.py b/stream_attention.py index 76dc81a..e0c80c9 100644 --- a/stream_attention.py +++ b/stream_attention.py @@ -3,6 +3,8 @@ import comfy.model_management from comfy.ldm.modules.attention import optimized_attention class StreamCrossAttention: + original_attention = None + @classmethod def get_attention_scores(cls, q, k, v, extra_options): """ @@ -11,55 +13,39 @@ class StreamCrossAttention: - KV cache preservation - Sliced attention fallback """ - print(f"[DEBUG] Attention input q stats - min: {q.min():.4f}, max: {q.max():.4f}, mean: {q.mean():.4f}") - print(f"[DEBUG] Attention input k stats - min: {k.min():.4f}, max: {k.max():.4f}, mean: {k.mean():.4f}") - print(f"[DEBUG] Attention input v stats - min: {v.min():.4f}, max: {v.max():.4f}, mean: {v.mean():.4f}") - - if not hasattr(cls, "attention_impl"): - # Initialize optimized implementation once - cls.attention_impl = cls.initialize_attention() + # Get transformer options + transformer_options = extra_options.get("transformer_options", {}) # Get cached key/values if available - cache = extra_options.get("_stream_kv_cache", {}) + kv_cache = transformer_options.get("kv_cache", {}) layer_id = extra_options["layer_id"] - if layer_id not in cache: - cache[layer_id] = {"k": None, "v": None} - - # Use cached KV if available - if cache[layer_id]["k"] is not None: - k = cache[layer_id]["k"] - v = cache[layer_id]["v"] - print(f"[DEBUG] Using cached k stats - min: {k.min():.4f}, max: {k.max():.4f}, mean: {k.mean():.4f}") - print(f"[DEBUG] Using cached v stats - min: {v.min():.4f}, max: {v.max():.4f}, mean: {v.mean():.4f}") - else: # First run or cache disabled - cache[layer_id]["k"] = k - cache[layer_id]["v"] = v + # Check if we should use cached KV + if kv_cache is not None and layer_id in kv_cache: + cached = kv_cache[layer_id] + k = cached["k"] + v = cached["v"] + else: + # Store KV in cache if enabled + if kv_cache is not None: + kv_cache[layer_id] = {"k": k, "v": v} - # Use ComfyUI's memory-efficient attention as base - output = optimized_attention( + # Use ComfyUI's memory-efficient attention + return optimized_attention( q, k, v, heads=extra_options["n_heads"], - mask=extra_options["mask"], - attn_precision=None, # Let ComfyUI handle precision - skip_reshape=False, # Default reshape behavior - skip_output_reshape=False # Default output reshape behavior + mask=extra_options.get("mask"), ) - print(f"[DEBUG] Attention output stats - min: {output.min():.4f}, max: {output.max():.4f}, mean: {output.mean():.4f}") - return output @classmethod def initialize_attention(cls): """Patch ComfyUI's attention modules with our optimized version""" from comfy.ldm.modules.attention import CrossAttention - original_attention = CrossAttention.forward + + if cls.original_attention is None: + cls.original_attention = CrossAttention.forward def stream_attention_forward(self, x, context=None, mask=None, **kwargs): - # Convert to BHWC format if needed - if x.shape[1] == self.dim_head * self.heads: - x = x.view(x.shape[0], -1, self.dim_head).permute(0, 2, 1, 3) - - # Use our optimized attention implementation q = self.to_q(x) context = context if context is not None else x k = self.to_k(context) @@ -69,8 +55,9 @@ class StreamCrossAttention: q, k, v, extra_options={ "n_heads": self.heads, - "layer_id": id(self), # Unique ID for KV caching - "mask": mask + "layer_id": id(self), + "mask": mask, + "transformer_options": kwargs.get("transformer_options", {}) } ) return self.to_out(out) @@ -90,7 +77,7 @@ class UNetBatchedStream: self.inner_model = model self.kv_cache = {} self.noise_base = None - self.current_channels = None + self.current_batch_size = None def __call__(self, apply_model, args): """ @@ -106,46 +93,31 @@ class UNetBatchedStream: timestep = args["timestep"] c = args["c"] - print(f"[DEBUG] Initial x stats - min: {x.min():.4f}, max: {x.max():.4f}, mean: {x.mean():.4f}") + # Get current batch size + batch_size = x.shape[0] - # Get current channel count from input - current_channels = x.shape[1] - - # Reinitialize noise base if channels change - if self.current_channels != current_channels: + # Reinitialize noise base if batch size changes + if self.current_batch_size != batch_size: self.noise_base = None - self.current_channels = current_channels + self.current_batch_size = batch_size - # Generate shared noise pattern matching current channels + # Generate shared noise pattern for batch if self.noise_base is None: self.noise_base = torch.randn_like(x) - print(f"[DEBUG] Generated noise_base - min: {self.noise_base.min():.4f}, max: {self.noise_base.max():.4f}, mean: {self.noise_base.mean():.4f}") - - # Apply noise blending using ComfyUI's native sigma scheduling - x = self.blend_noise(x, timestep) - print(f"[DEBUG] After blend_noise - min: {x.min():.4f}, max: {x.max():.4f}, mean: {x.mean():.4f}") - # Pass through model + # Get transformer options + transformer_options = c.get("transformer_options", {}) + + # Add KV cache to transformer options + transformer_options["kv_cache"] = self.kv_cache + + # Pass through model with updated options return apply_model( x, timestep, c_concat=c.get("c_concat"), c_crossattn=c.get("c_crossattn"), control=c.get("control"), - transformer_options=c.get("transformer_options", {}), - cond_or_uncond=args.get("cond_or_uncond", None), - kv_cache=self.kv_cache - ) - def blend_noise(self, x, timesteps): - """ComfyUI-native noise blending""" - # Convert timesteps to sigmas (this matches ComfyUI's internal scheduling) - sigma = timesteps - - # Simple sigma-based noise mixing - noise_level = 1.0 / (1.0 + sigma) - noise_level = noise_level.reshape(-1, 1, 1, 1) - - # Mix clean latent with noise base - x = x * (1.0 - noise_level) + self.noise_base * noise_level - - return x + transformer_options=transformer_options, + cond_or_uncond=args.get("cond_or_uncond", None) + ) diff --git a/stream_diffusion_nodes.py b/stream_diffusion_nodes.py index e3bfb19..9d766fc 100644 --- a/stream_diffusion_nodes.py +++ b/stream_diffusion_nodes.py @@ -71,11 +71,9 @@ class StreamDiffusionPipeline(ControlNodeBase): "latent_buffer": None, # Shape: [num_timesteps * frame_buffer_size, 4, H, W] "current_idx": 0, "t_indices": None, - "alpha_prod_t_sqrt": None, - "beta_prod_t_sqrt": None, - "init_noise": None, "sigmas": None, "stock_noise": None, # For RCFG + "timesteps": None, }) # Get latent tensor from dict @@ -93,71 +91,50 @@ class StreamDiffusionPipeline(ControlNodeBase): state["sigmas"] = sigmas # Parse t_index_list and get timesteps - state["t_indices"] = [int(x.strip()) for x in t_index_list.split(",")] + t_indices = [int(x.strip()) for x in t_index_list.split(",")] # Get actual timesteps from scheduler timesteps = [] - alpha_prod_t_sqrt_list = [] - beta_prod_t_sqrt_list = [] + for t_idx in t_indices: + timesteps.append(len(sigmas) - 1 - t_idx) - # Convert sigmas to timesteps and calculate alpha/beta - for t_idx in state["t_indices"]: - sigma = sigmas[t_idx] - # Convert sigma to timestep - timestep = len(sigmas) - 1 - t_idx - timesteps.append(timestep) - - # Calculate alpha/beta from sigma - alpha = 1.0 / (1.0 + sigma * sigma) - beta = sigma * sigma / (1.0 + sigma * sigma) - - alpha_prod_t_sqrt_list.append(alpha.sqrt()) - beta_prod_t_sqrt_list.append(beta.sqrt()) - - # Update t_indices to use actual timesteps - state["t_indices"] = timesteps - - # Stack and reshape for broadcasting - state["alpha_prod_t_sqrt"] = torch.stack(alpha_prod_t_sqrt_list).view(-1, 1, 1, 1).to(device) - state["beta_prod_t_sqrt"] = torch.stack(beta_prod_t_sqrt_list).view(-1, 1, 1, 1).to(device) + # Store timesteps for later use + state["timesteps"] = torch.tensor(timesteps, dtype=torch.long, device=device) + state["t_indices"] = torch.tensor(t_indices, dtype=torch.long, device=device) # Calculate total buffer size based on timesteps and frame buffer - total_buffer_size = len(state["t_indices"]) * frame_buffer_size + total_buffer_size = len(state["timesteps"]) * frame_buffer_size # Initialize buffers if needed if state["latent_buffer"] is None: state["latent_buffer"] = torch.zeros((total_buffer_size, *samples.shape[1:]), dtype=samples.dtype, device=device) - # Generate fixed noise pattern - state["init_noise"] = torch.randn(samples.shape, dtype=samples.dtype, device=device) - state["stock_noise"] = torch.zeros_like(samples) # For RCFG + # Initialize stock noise for RCFG + state["stock_noise"] = torch.zeros_like(samples) state["current_idx"] = 0 # Update buffer with current frame - for i in range(len(state["t_indices"])): + for i in range(len(state["timesteps"])): buffer_idx = i * frame_buffer_size + (state["current_idx"] % frame_buffer_size) if do_add_noise: - # Add scaled noise using proper scheduler parameters - alpha = state["alpha_prod_t_sqrt"][i] - beta = state["beta_prod_t_sqrt"][i] - state["latent_buffer"][buffer_idx] = ( - alpha * samples[0] + # Scale clean latent - beta * state["init_noise"][0] # Use fixed noise pattern - ) + # Get sigma for current timestep + sigma = state["sigmas"][state["t_indices"][i]] + # Add noise using sigma scheduling + noise = torch.randn_like(samples) + state["latent_buffer"][buffer_idx] = samples[0] + noise[0] * sigma else: state["latent_buffer"][buffer_idx] = samples[0] # Create batched input for denoising - batched = state["latent_buffer"].view(-1, *samples.shape[1:]) # [total_buffer_size, 4, H, W] - timesteps = torch.tensor(state["t_indices"], dtype=torch.long, device=device) + batched = state["latent_buffer"] # Already in right shape [total_buffer_size, 4, H, W] + + # Get sigmas for selected timesteps + timestep_sigmas = state["sigmas"][state["t_indices"]] # Create sampler sampler = comfy.samplers.sampler_object(sampler_name) - # Get sigmas for selected timesteps - timestep_sigmas = state["sigmas"][timesteps] - # Create custom CFG function for the sampler def custom_cfg(x, t, cond, uncond, cfg_scale, model_options={}): with torch.no_grad(): @@ -172,13 +149,14 @@ class StreamDiffusionPipeline(ControlNodeBase): # Get unconditional prediction noise_pred_uncond = model.model.apply_model(x, t, uncond) elif cfg_type == "self" or cfg_type == "initialize": - # Use stock noise for unconditional - noise_pred_uncond = state["stock_noise"] * delta + # Use stock noise for unconditional with delta scaling + noise_pred_uncond = state["stock_noise"].repeat(x.shape[0], 1, 1, 1) * delta # Update stock noise for next iteration if cfg_type == "self": - state["stock_noise"] = noise_pred_text.detach() - elif cfg_type == "initialize" and t == timesteps[0]: # First step - state["stock_noise"] = noise_pred_text.detach() + # Update per batch element + state["stock_noise"] = noise_pred_text.reshape(-1, frame_buffer_size, *noise_pred_text.shape[1:])[-1] + elif cfg_type == "initialize" and t == timestep_sigmas[0]: # First step + state["stock_noise"] = noise_pred_text[-1:] # Use last prediction # Apply CFG return noise_pred_uncond + cfg_scale * (noise_pred_text - noise_pred_uncond) @@ -192,14 +170,11 @@ class StreamDiffusionPipeline(ControlNodeBase): if not hasattr(model, "_stream_wrapped"): model._stream_wrapped = UNetBatchedStream(model) model.set_model_unet_function_wrapper(model._stream_wrapped) - - # Convert latent to BHWC format - samples = samples.permute(0, 2, 3, 1) # BCHW -> BHWC # Modified denoising call: denoised = comfy.samplers.sample( model, - state["init_noise"].repeat(len(batched), 1, 1, 1), + torch.randn_like(batched), # Fresh noise per batch positive, negative, guidance_scale,