nodes/ryanontheinside-ComfyUI_RealtimeNodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

streambatchsampler, streambatch scheduler, stream corss attention

Browse Source
This commit is contained in:
RyanOnTheInside
2025-02-10 16:53:42 -05:00
parent ac67be542a
commit 135d012a0f
2 changed files with 326 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+6 -1
View File
@@ -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,
stream_diffusion_nodes.py
+320 -1
View File
@@ -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,)