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

coherent output

Browse Source
This commit is contained in:
RyanOnTheInside
2025-02-10 22:05:25 -05:00
parent e0cd08911e
commit 9ce8047fe8
3 changed files with 165 additions and 220 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+2 -2
View File
@@ -4,7 +4,7 @@ from .controls.utility_controls import FPSMonitor, SimilarityFilter, LazyConditi
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, StreamCrossAttention
from .stream_sampler import StreamBatchSampler, StreamScheduler
from .stream_sampler import StreamBatchSampler, StreamScheduler, StreamFrameBuffer
import re
@@ -22,6 +22,7 @@ NODE_CLASS_MAPPINGS = {
"StreamConditioning": StreamConditioning,
"StreamBatchSampler": StreamBatchSampler,
"StreamScheduler": StreamScheduler,
"StreamFrameBuffer": StreamFrameBuffer,
"StreamCrossAttention": StreamCrossAttention,
"LazyCondition": LazyCondition,
"MotionController": MotionController,
@@ -30,7 +31,6 @@ NODE_CLASS_MAPPINGS = {
"TextRenderer": TextRenderer,
"ROINode": ROINode,
#"IntervalControl": IntervalCo ntrol,
stream_diffusion_nodes.py
+32 -158
View File
@@ -186,6 +186,7 @@ class StreamCFG(ControlNodeBase):
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
@@ -213,6 +214,13 @@ class StreamConditioning(ControlNodeBase):
"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,
}),
@@ -227,37 +235,53 @@ class StreamConditioning(ControlNodeBase):
def __init__(self):
super().__init__()
def update(self, positive, negative, cfg_type="full", residual_scale=0.4, delta=1.0, always_execute=False):
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
"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] # Assuming standard ComfyUI conditioning format
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":
# For initialize, we use the first negative as our stock noise
state["stock_noise"] = current_neg_cond.detach().clone()
elif cfg_type == "self":
# For self, we start with a scaled version of the negative
state["stock_noise"] = current_neg_cond.detach().clone() * delta
self.set_state(state)
return (positive, negative)
# Handle different CFG types
if cfg_type == "full":
# Standard R-CFG with full negative conditioning
pos_residual = current_pos_cond - state["prev_positive"]
neg_residual = current_neg_cond - state["prev_negative"]
# 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
@@ -300,156 +324,6 @@ 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 with KV-cache for real-time generation
stream_sampler.py
+131 -60
View File
@@ -6,7 +6,7 @@ import random
class StreamBatchSampler(ControlNodeBase):
"""Implements batched denoising for faster inference by processing multiple steps in parallel"""
"""Implements batched denoising for faster inference by processing multiple frames in parallel at different denoising steps"""
RETURN_TYPES = ("SAMPLER",)
FUNCTION = "update"
@@ -16,79 +16,100 @@ class StreamBatchSampler(ControlNodeBase):
def INPUT_TYPES(cls):
inputs = super().INPUT_TYPES()
inputs["required"].update({
"batch_size": ("INT", {
"default": 2,
"num_steps": ("INT", {
"default": 4,
"min": 1,
"max": 10,
"step": 1,
"tooltip": "Number of steps to batch together. Higher values use more memory but are faster."
"tooltip": "Number of denoising steps. Should match the frame buffer size."
}),
})
return inputs
def __init__(self):
super().__init__()
self.batch_size = None
self.num_steps = None
self.frame_buffer = []
self.x_t_latent_buffer = None
self.stock_noise = None
def sample(self, model, noise, sigmas, extra_args=None, callback=None, disable=None):
"""Sample with batched denoising steps"""
"""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, batch_size={self.batch_size}")
print(f"[StreamBatchSampler] Input noise shape: {noise.shape}, device: {noise.device}")
print(f"[StreamBatchSampler] Sigmas: {sigmas.tolist()}")
print(f"[StreamBatchSampler] Starting sampling with {len(sigmas)-1} steps")
# Prepare batched sampling
num_sigmas = len(sigmas) - 1
num_batches = (num_sigmas + self.batch_size - 1) // self.batch_size
x = noise
# 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)
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()}")
print(f"[StreamBatchSampler] Input sigmas: {sigmas}")
print(f"[StreamBatchSampler] Input noise shape: {noise.shape}, min: {noise.min():.3f}, max: {noise.max():.3f}")
# Verify batch size matches number of timesteps
if batch_size != num_sigmas:
raise ValueError(f"Batch size ({batch_size}) must match number of timesteps ({num_sigmas})")
# Pre-compute alpha and beta terms
alpha_prod_t = (sigmas[:-1] / sigmas[0]).view(-1, 1, 1, 1) # [B,1,1,1]
beta_prod_t = (1 - alpha_prod_t)
print(f"[StreamBatchSampler] Alpha values: {alpha_prod_t.view(-1)}")
print(f"[StreamBatchSampler] Beta values: {beta_prod_t.view(-1)}")
# Initialize stock noise if needed
if self.stock_noise is None:
self.stock_noise = torch.randn_like(noise[0]) # Random noise instead of zeros
print(f"[StreamBatchSampler] Initialized random stock noise with shape: {self.stock_noise.shape}")
# Create batch of identical latents
batch_size = end_idx - start_idx
x_batch = x.repeat(batch_size, 1, 1, 1)
# Scale noise for each frame based on its sigma
scaled_noise = []
for i in range(batch_size):
frame_noise = noise[i] + self.stock_noise * sigmas[i] # Add scaled noise to input
scaled_noise.append(frame_noise)
x = torch.stack(scaled_noise, dim=0)
print(f"[StreamBatchSampler] Scaled noise shape: {x.shape}, min: {x.min():.3f}, max: {x.max():.3f}")
# Create batch of sigmas
sigma_batch = batch_sigmas[:-1] # All but last sigma
# Initialize frame buffer if needed
if self.x_t_latent_buffer is None and num_sigmas > 1:
self.x_t_latent_buffer = x[0].clone() # Initialize with noised first frame
print(f"[StreamBatchSampler] Initialized buffer with shape: {self.x_t_latent_buffer.shape}")
# 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}")
# Use buffer for first frame to maintain temporal consistency
if num_sigmas > 1:
x = torch.cat([self.x_t_latent_buffer.unsqueeze(0), x[1:]], dim=0)
print(f"[StreamBatchSampler] Combined with buffer, shape: {x.shape}")
# Run model on entire batch at once
with torch.no_grad():
# Process all frames in parallel
sigma_batch = sigmas[:-1]
print(f"[StreamBatchSampler] Using sigmas for denoising: {sigma_batch}")
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}")
# Update buffer with intermediate results
if num_sigmas > 1:
# Store result from first frame as buffer for next iteration
self.x_t_latent_buffer = denoised_batch[0].clone()
print(f"[StreamBatchSampler] Updated buffer with shape: {self.x_t_latent_buffer.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})
# Return result from last frame
x_0_pred_out = denoised_batch[-1].unsqueeze(0)
else:
x_0_pred_out = denoised_batch
self.x_t_latent_buffer = None
# 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:]})
print(f"\n[StreamBatchSampler] Sampling complete. Final x shape: {x.shape}")
return x
return x_0_pred_out
def update(self, batch_size=2, always_execute=True):
def update(self, num_steps=4, always_execute=True):
"""Create sampler with specified settings"""
self.batch_size = batch_size
self.num_steps = num_steps
sampler = comfy.samplers.KSAMPLER(self.sample)
return (sampler,)
@@ -122,7 +143,6 @@ class StreamScheduler(ControlNodeBase):
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:
@@ -130,11 +150,9 @@ class StreamScheduler(ControlNodeBase):
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
@@ -144,12 +162,65 @@ class StreamScheduler(ControlNodeBase):
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 StreamFrameBuffer(ControlNodeBase):
"""Accumulates frames to enable staggered batch denoising like StreamDiffusion"""
RETURN_TYPES = ("LATENT",)
FUNCTION = "update"
CATEGORY = "real-time/sampling"
@classmethod
def INPUT_TYPES(cls):
inputs = super().INPUT_TYPES()
inputs["required"].update({
"latent": ("LATENT",),
"buffer_size": ("INT", {
"default": 4,
"min": 1,
"max": 10,
"step": 1,
"tooltip": "Number of frames to buffer before starting batch processing. Should match number of denoising steps."
}),
})
return inputs
def __init__(self):
super().__init__()
self.frame_buffer = [] # List to store incoming frames
self.buffer_size = None
def update(self, latent, buffer_size=4, always_execute=True):
"""Add new frame to buffer and return batch when ready"""
self.buffer_size = buffer_size
# Extract latent tensor from input and remove batch dimension if present
x = latent["samples"]
if x.dim() == 4: # [B,C,H,W]
x = x.squeeze(0) # Remove batch dimension -> [C,H,W]
# Add new frame to buffer
if len(self.frame_buffer) == 0:
# First frame - initialize buffer with copies
self.frame_buffer = [x.clone() for _ in range(self.buffer_size)]
print(f"[StreamFrameBuffer] Initialized buffer with {self.buffer_size} copies of first frame")
else:
# Shift frames forward and add new frame
self.frame_buffer.pop(0) # Remove oldest frame
self.frame_buffer.append(x.clone()) # Add new frame
print(f"[StreamFrameBuffer] Added new frame to buffer")
# Stack frames into batch
batch = torch.stack(self.frame_buffer, dim=0) # [B,C,H,W]
print(f"[StreamFrameBuffer] Created batch with shape: {batch.shape}")
# Return as latent dict
return ({"samples": batch},)