592 lines
28 KiB
Python
592 lines
28 KiB
Python
import torch
|
|
from .base.control_base import ControlNodeBase
|
|
import comfy.model_management
|
|
import comfy.samplers
|
|
import random
|
|
import math
|
|
|
|
|
|
class StreamBatchSampler(ControlNodeBase):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = super().INPUT_TYPES()
|
|
inputs["required"].update({
|
|
"latent": ("LATENT",),
|
|
"num_timesteps": ("INT", {
|
|
"default": 2,
|
|
"min": 1,
|
|
"max": 10,
|
|
"step": 1,
|
|
"tooltip": "Number of denoising steps to use. More steps = better quality but MUCH slower. StreamDiffusion's speed comes from using just 2 steps for img2img (32,45) or 4 for txt2img (0,16,32,45)"
|
|
}),
|
|
"frame_buffer_size": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 10,
|
|
"step": 1,
|
|
"tooltip": "How many frames to process together for temporal consistency. Each frame is processed at each timestep (total batch = frame_buffer_size * num_timesteps). Higher values reduce flickering but use more VRAM and don't improve speed"
|
|
}),
|
|
"cfg_type": (["none", "full", "self", "initialize"], {
|
|
"default": "self",
|
|
"tooltip": "'self' is fastest and most memory efficient, 'full' is standard SD behavior but slower, 'initialize' is a middle ground, 'none' disables guidance"
|
|
}),
|
|
"delta": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 5.0,
|
|
"step": 0.1,
|
|
"tooltip": "Only used with 'self' CFG. Controls strength of self-guidance. Higher values = stronger guidance but more artifacts. Default 1.0 works well"
|
|
}),
|
|
"similarity_threshold": ("FLOAT", {
|
|
"default": 0.98,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"tooltip": "Skip frames that are too similar to save compute. 0 = process all frames. Try 0.98 if you want to skip static scenes. Higher = more skipping"
|
|
#TODO: see if this is backwards from original
|
|
}),
|
|
"max_skip_frames": ("INT", {
|
|
"default": 10,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1,
|
|
"tooltip": "Maximum frames to skip when using similarity filter. Prevents getting stuck on static scenes. Lower = more responsive to changes"
|
|
}),
|
|
})
|
|
inputs["optional"] = {
|
|
"prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "Optional text prompt. Enables KV caching for faster processing. Leave empty to disable. Same prompt across frames = faster"
|
|
}),
|
|
}
|
|
return inputs
|
|
|
|
RETURN_TYPES = ("SAMPLER", "LATENT",)
|
|
FUNCTION = "update"
|
|
CATEGORY = "real-time/sampling"
|
|
DESCRIPTION="Implements batched denoising for faster inference by processing multiple frames in parallel at different denoising steps (StreamDiffusion)"
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.frame_buffer_size = None
|
|
self.frame_buffer = None # Will be a tensor instead of list
|
|
self.x_t_latent_buffer = None
|
|
self.stock_noise = None
|
|
self.cached_noise = {} # Cache for different expanded versions
|
|
self.debug = False # Debug flag for print statements
|
|
self.last_batch_shape = None # Track last batch shape
|
|
# Add CFG state
|
|
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.similarity_threshold = 0.0
|
|
self.max_skip_frames = 10
|
|
self.similarity_filter = None
|
|
# Add buffer state
|
|
self.num_timesteps = None
|
|
# Add caching for alpha/beta and their square roots
|
|
self.cached_sigmas = None
|
|
self.alpha_prod_t = None
|
|
self.beta_prod_t = None
|
|
self.alpha_prod_t_sqrt = None
|
|
self.beta_prod_t_sqrt = None
|
|
# Add noise sequence and step counter
|
|
self.noise_sequence = None
|
|
self.step_counter = 0
|
|
|
|
def get_expanded_noise(self, batch_size, scale=1.0):
|
|
"""Get expanded noise with caching based on batch size and scale"""
|
|
cache_key = (batch_size, scale)
|
|
if self.stock_noise is None:
|
|
return None
|
|
|
|
if cache_key not in self.cached_noise:
|
|
self.cached_noise[cache_key] = self.stock_noise.unsqueeze(0).expand(batch_size, -1, -1, -1) * scale
|
|
if self.debug:
|
|
print(f"[StreamBatchSampler] Cached expanded noise for batch_size={batch_size}, scale={scale}")
|
|
|
|
return self.cached_noise[cache_key]
|
|
|
|
def invalidate_noise_cache(self):
|
|
"""Clear all cached noise when stock noise changes"""
|
|
self.cached_noise.clear()
|
|
|
|
|
|
def process_with_cfg(self, model, x, sigma_batch, extra_args):
|
|
"""Apply classifier-free guidance based on current cfg_type"""
|
|
cond_scale = extra_args["model_options"].get("cond_scale", 1.0)
|
|
|
|
if self.cfg_type == "none" or cond_scale <= 1.0:
|
|
if self.debug:
|
|
print("[CFG] No guidance - running model directly")
|
|
return model(x, sigma_batch, **extra_args)
|
|
|
|
if self.cfg_type == "full":
|
|
if self.debug:
|
|
print("[CFG] Using full CFG")
|
|
# 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":
|
|
if self.debug:
|
|
print("[CFG] Using initialize CFG")
|
|
# 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()
|
|
self.invalidate_noise_cache()
|
|
else: # self
|
|
if self.debug:
|
|
print("\n[CFG] Using RCFG (self)")
|
|
print(f"[CFG] Input shape: {x.shape}, Sigma shape: {sigma_batch.shape}")
|
|
|
|
# Get current model prediction
|
|
eps_c = model(x, sigma_batch, **extra_args)
|
|
if self.debug:
|
|
print(f"[CFG] Current prediction (eps_c) stats: min={eps_c.min():.3f}, max={eps_c.max():.3f}, mean={eps_c.mean():.3f}")
|
|
|
|
# Initialize stock noise if needed
|
|
if self.stock_noise is None:
|
|
if self.debug:
|
|
print("[CFG] Initializing stock noise with current prediction")
|
|
self.stock_noise = eps_c[0].clone().detach()
|
|
self.invalidate_noise_cache()
|
|
return eps_c
|
|
|
|
# Apply RCFG equation
|
|
noise_pred_uncond = self.stock_noise.to(eps_c.device) * self.delta
|
|
if self.debug:
|
|
print(f"[CFG] Using delta={self.delta}, cond_scale={cond_scale}")
|
|
print(f"[CFG] Stock noise stats: min={self.stock_noise.min():.3f}, max={self.stock_noise.max():.3f}, mean={self.stock_noise.mean():.3f}")
|
|
print(f"[CFG] Uncond prediction stats: min={noise_pred_uncond.min():.3f}, max={noise_pred_uncond.max():.3f}, mean={noise_pred_uncond.mean():.3f}")
|
|
|
|
# Combine predictions exactly as StreamDiffusion does
|
|
eps_cfg = noise_pred_uncond + cond_scale * (eps_c - noise_pred_uncond)
|
|
|
|
if self.debug:
|
|
print(f"[CFG] Final output (eps_cfg) stats: min={eps_cfg.min():.3f}, max={eps_cfg.max():.3f}, mean={eps_cfg.mean():.3f}")
|
|
print(f"[CFG] Difference (eps_cfg - eps_c) stats: min={(eps_cfg-eps_c).min():.3f}, max={(eps_cfg-eps_c).max():.3f}, mean={(eps_cfg-eps_c).mean():.3f}\n")
|
|
|
|
return eps_cfg
|
|
|
|
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
|
|
|
|
def optimized_attention_forward(module, x, context=None, mask=None):
|
|
"""Optimized cross attention with KV caching and ComfyUI's optimizations"""
|
|
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:
|
|
if self.debug:
|
|
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:
|
|
if self.debug:
|
|
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)
|
|
|
|
# Use ComfyUI's optimized attention
|
|
from comfy.ldm.modules.attention import optimized_attention
|
|
out = optimized_attention(q, k, v, module.heads, mask=mask)
|
|
return module.to_out[0](out)
|
|
|
|
# Get the actual model from the sampler
|
|
if hasattr(model, "model"):
|
|
unet = model.model
|
|
else:
|
|
return model
|
|
|
|
# Only set up hooks if prompt has changed
|
|
if self.last_prompt != prompt:
|
|
# Remove old hooks if they exist
|
|
if self.cross_attention_hook is not None:
|
|
self.cross_attention_hook.remove()
|
|
|
|
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
|
|
|
|
# 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_prompt = prompt
|
|
|
|
return model
|
|
|
|
def apply_similarity_filter(self, x_0_pred_out):
|
|
"""Apply similarity filter to skip similar frames if enabled"""
|
|
if not self.similarity_threshold > 0:
|
|
return x_0_pred_out
|
|
|
|
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:
|
|
return filtered_out
|
|
|
|
return x_0_pred_out
|
|
|
|
def compute_alpha_beta(self, sigmas):
|
|
"""Pre-compute alpha and beta terms for the given sigmas"""
|
|
self.cached_sigmas = sigmas.clone()
|
|
print(f"\n[Alpha/Beta] Input sigmas: {sigmas}")
|
|
|
|
# Compute alpha and beta terms using StreamDiffusion's approach
|
|
# In StreamDiffusion, alpha is based on alphas_cumprod which is 1/(1 + sigma^2)
|
|
self.alpha_prod_t = 1 / (1 + sigmas[:self.num_timesteps].pow(2)).view(-1, 1, 1, 1)
|
|
self.beta_prod_t = 1 - self.alpha_prod_t
|
|
|
|
# Also compute square root versions needed for noise scaling
|
|
self.alpha_prod_t_sqrt = torch.sqrt(self.alpha_prod_t)
|
|
self.beta_prod_t_sqrt = torch.sqrt(self.beta_prod_t)
|
|
|
|
print(f"[Alpha/Beta] Final stats:")
|
|
print(f" alpha_prod_t: min={self.alpha_prod_t.min():.4f}, max={self.alpha_prod_t.max():.4f}, mean={self.alpha_prod_t.mean():.4f}")
|
|
print(f" beta_prod_t: min={self.beta_prod_t.min():.4f}, max={self.beta_prod_t.max():.4f}, mean={self.beta_prod_t.mean():.4f}")
|
|
print(f" alpha_sqrt: min={self.alpha_prod_t_sqrt.min():.4f}, max={self.alpha_prod_t_sqrt.max():.4f}, mean={self.alpha_prod_t_sqrt.mean():.4f}")
|
|
print(f" beta_sqrt: min={self.beta_prod_t_sqrt.min():.4f}, max={self.beta_prod_t_sqrt.max():.4f}, mean={self.beta_prod_t_sqrt.mean():.4f}\n")
|
|
|
|
def scheduler_step_batch(self, denoised_batch, scaled_noise):
|
|
"""Compute scheduler step for the entire batch"""
|
|
print("\n[Scheduler] Starting scheduler step:")
|
|
print(f" Denoised shape: {denoised_batch.shape}")
|
|
print(f" Scaled noise shape: {scaled_noise.shape}")
|
|
print(f" Denoised stats: min={denoised_batch.min():.4f}, max={denoised_batch.max():.4f}, mean={denoised_batch.mean():.4f}")
|
|
print(f" Scaled noise stats: min={scaled_noise.min():.4f}, max={scaled_noise.max():.4f}, mean={scaled_noise.mean():.4f}")
|
|
|
|
# Compute F_theta using the scheduler equation
|
|
beta_scaled = self.beta_prod_t_sqrt * scaled_noise
|
|
print(f" Beta scaled noise stats: min={beta_scaled.min():.4f}, max={beta_scaled.max():.4f}, mean={beta_scaled.mean():.4f}")
|
|
|
|
diff = denoised_batch - beta_scaled
|
|
print(f" Difference stats: min={diff.min():.4f}, max={diff.max():.4f}, mean={diff.mean():.4f}")
|
|
|
|
F_theta = diff / (self.alpha_prod_t_sqrt + 1e-7)
|
|
print(f" F_theta stats: min={F_theta.min():.4f}, max={F_theta.max():.4f}, mean={F_theta.mean():.4f}")
|
|
|
|
# Get c_out and c_skip from scheduler like StreamDiffusion does
|
|
c_out = 1 / torch.sqrt(self.alpha_prod_t)
|
|
c_skip = torch.sqrt(1 - self.beta_prod_t)
|
|
|
|
# Apply scaling
|
|
delta_x = c_out * F_theta + c_skip * denoised_batch
|
|
print(f" Final delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}\n")
|
|
|
|
return delta_x
|
|
|
|
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"\n[Sample] Starting sample step {self.step_counter}:")
|
|
print(f" Input noise shape: {noise.shape}")
|
|
print(f" Sigmas shape: {sigmas.shape}")
|
|
print(f" Sigmas values: {sigmas[:self.num_timesteps]}")
|
|
|
|
# Set up cross attention optimization if prompt provided
|
|
if self.prompt:
|
|
model = self.setup_cross_attention(model, self.prompt)
|
|
|
|
# Get number of frames in batch and available sigmas
|
|
batch_size = noise.shape[0]
|
|
|
|
# Calculate total batch size based on frame_buffer_size and number of timesteps
|
|
total_batch_size = self.frame_buffer_size * self.num_timesteps
|
|
if batch_size != total_batch_size:
|
|
raise ValueError(f"Batch size ({batch_size}) must match frame_buffer_size * num_timesteps ({total_batch_size})")
|
|
|
|
# Pre-compute alpha and beta terms if not already computed
|
|
if self.cached_sigmas is None:
|
|
self.compute_alpha_beta(sigmas)
|
|
|
|
# Update noise sequence like StreamDiffusion
|
|
if self.step_counter == 0:
|
|
# First step - use initial noise
|
|
self.stock_noise = self.init_noise[0].clone()
|
|
else:
|
|
# Shift noise sequence
|
|
self.stock_noise = torch.cat([self.init_noise[1:], self.init_noise[0:1]], dim=0)[0]
|
|
|
|
print(f" Stock noise stats: min={self.stock_noise.min():.4f}, max={self.stock_noise.max():.4f}, mean={self.stock_noise.mean():.4f}")
|
|
|
|
# Scale noise for each frame based on its sigma - vectorized with caching
|
|
sigma_scaled_noise = self.get_expanded_noise(batch_size)
|
|
print(f" Expanded noise shape: {sigma_scaled_noise.shape}")
|
|
print(f" Expanded noise stats: min={sigma_scaled_noise.min():.4f}, max={sigma_scaled_noise.max():.4f}, mean={sigma_scaled_noise.mean():.4f}")
|
|
|
|
x = noise + sigma_scaled_noise * sigmas[:self.num_timesteps].view(-1, 1, 1, 1)
|
|
print(f" Initial x stats: min={x.min():.4f}, max={x.max():.4f}, mean={x.mean():.4f}")
|
|
|
|
# Initialize frame buffer if needed
|
|
if self.x_t_latent_buffer is None and self.num_timesteps > 1:
|
|
self.x_t_latent_buffer = x[0].clone()
|
|
print(" Initialized latent buffer with first frame")
|
|
|
|
# Use buffer for first frame to maintain temporal consistency
|
|
if self.num_timesteps > 1:
|
|
x = torch.cat([self.x_t_latent_buffer.unsqueeze(0), x[1:]], dim=0)
|
|
print(f" After buffer concat x stats: min={x.min():.4f}, max={x.max():.4f}, mean={x.mean():.4f}")
|
|
|
|
# Run model on entire batch at once
|
|
with torch.no_grad():
|
|
# Process all frames in parallel
|
|
sigma_batch = sigmas[:self.num_timesteps]
|
|
# Process with CFG handling inside process_with_cfg
|
|
denoised_batch = self.process_with_cfg(model, x, sigma_batch, extra_args)
|
|
print(f"\n[Sample] After denoising:")
|
|
print(f" Denoised shape: {denoised_batch.shape}")
|
|
print(f" Denoised stats: min={denoised_batch.min():.4f}, max={denoised_batch.max():.4f}, mean={denoised_batch.mean():.4f}")
|
|
|
|
# Update buffer with intermediate results
|
|
if self.num_timesteps > 1:
|
|
# Store result from first frame as buffer for next iteration
|
|
self.x_t_latent_buffer = denoised_batch[0].clone()
|
|
print(f" Updated buffer stats: min={self.x_t_latent_buffer.min():.4f}, max={self.x_t_latent_buffer.max():.4f}, mean={self.x_t_latent_buffer.mean():.4f}")
|
|
# 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
|
|
|
|
print(f" Output pred stats: min={x_0_pred_out.min():.4f}, max={x_0_pred_out.max():.4f}, mean={x_0_pred_out.mean():.4f}")
|
|
|
|
# Apply similarity filter if enabled
|
|
x_0_pred_out = self.apply_similarity_filter(x_0_pred_out)
|
|
|
|
# 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': x_0_pred_out})
|
|
|
|
# Update stock noise using scheduler steps
|
|
if self.cfg_type == "self":
|
|
print("\n[Sample] Updating stock noise:")
|
|
# Scale noise properly
|
|
scaled_noise = self.beta_prod_t_sqrt * self.stock_noise
|
|
print(f" Scaled noise stats: min={scaled_noise.min():.4f}, max={scaled_noise.max():.4f}, mean={scaled_noise.mean():.4f}")
|
|
|
|
# Get delta_x using scheduler step
|
|
delta_x = self.scheduler_step_batch(denoised_batch, scaled_noise)
|
|
print(f" Initial delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}")
|
|
|
|
# Prepare next step's alpha/beta terms
|
|
alpha_next = torch.cat([
|
|
self.alpha_prod_t_sqrt[1:],
|
|
torch.ones_like(self.alpha_prod_t_sqrt[0:1])
|
|
], dim=0)
|
|
beta_next = torch.cat([
|
|
self.beta_prod_t_sqrt[1:],
|
|
torch.ones_like(self.beta_prod_t_sqrt[0:1])
|
|
], dim=0)
|
|
|
|
print(f" Alpha next stats: min={alpha_next.min():.4f}, max={alpha_next.max():.4f}, mean={alpha_next.mean():.4f}")
|
|
print(f" Beta next stats: min={beta_next.min():.4f}, max={beta_next.max():.4f}, mean={beta_next.mean():.4f}")
|
|
|
|
# Update noise prediction
|
|
delta_x = (alpha_next * delta_x) / (beta_next + 1e-7)
|
|
print(f" Scaled delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}")
|
|
|
|
# Use noise sequence for init_noise
|
|
init_noise = torch.cat([
|
|
self.init_noise[1:],
|
|
self.init_noise[0:1]
|
|
], dim=0)
|
|
print(f" Init noise stats: min={init_noise.min():.4f}, max={init_noise.max():.4f}, mean={init_noise.mean():.4f}")
|
|
|
|
# Update stock noise with proper scaling
|
|
self.stock_noise = init_noise[0] + delta_x[0]
|
|
print(f" Final stock noise stats: min={self.stock_noise.min():.4f}, max={self.stock_noise.max():.4f}, mean={self.stock_noise.mean():.4f}\n")
|
|
|
|
# Update step counter
|
|
self.step_counter += 1
|
|
|
|
return x_0_pred_out
|
|
|
|
def buffer(self, latent, frame_buffer_size=1, num_timesteps=2, always_execute=True):
|
|
"""Add new frame to buffer and return batch when ready"""
|
|
self.frame_buffer_size = frame_buffer_size
|
|
self.num_timesteps = num_timesteps
|
|
|
|
# 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]
|
|
|
|
# Initialize or resize frame buffer if needed
|
|
if self.frame_buffer is None or self.frame_buffer.shape[0] != self.frame_buffer_size:
|
|
# Pre-allocate tensor buffer
|
|
self.frame_buffer = torch.zeros((self.frame_buffer_size,) + x.shape, dtype=x.dtype, device=x.device)
|
|
if self.debug:
|
|
print(f"[StreamFrameBuffer] Pre-allocated tensor buffer for {self.frame_buffer_size} frames")
|
|
|
|
# Rotate frames and add new frame using tensor operations
|
|
if self.frame_buffer_size > 1:
|
|
self.frame_buffer = torch.roll(self.frame_buffer, -1, dims=0)
|
|
self.frame_buffer[-1] = x
|
|
|
|
# Create batch by repeating frames for each timestep
|
|
batch = self.frame_buffer.repeat_interleave(self.num_timesteps, dim=0) # [buffer_size*num_timesteps,C,H,W]
|
|
|
|
# Return as latent dict
|
|
return {"samples": batch}
|
|
|
|
def update(self, latent, num_timesteps, frame_buffer_size=1, cfg_type="self", delta=1.0, prompt="",
|
|
similarity_threshold=0.0, max_skip_frames=10, always_execute=True):
|
|
"""Create sampler with specified settings and process frame"""
|
|
self.frame_buffer_size = frame_buffer_size
|
|
self.cfg_type = cfg_type
|
|
self.delta = delta
|
|
self.prompt = prompt
|
|
self.num_timesteps = num_timesteps
|
|
|
|
# Update similarity filter settings
|
|
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 with cusatom sample function
|
|
sampler = comfy.samplers.KSAMPLER(self.sample)
|
|
|
|
# Process frame
|
|
buffered = self.buffer(latent)
|
|
|
|
# Get latent dimensions from input
|
|
_, _, self.latent_height, self.latent_width = latent["samples"].shape
|
|
|
|
# Add noise sequence initialization
|
|
if self.noise_sequence is None or self.noise_sequence.shape[0] != num_timesteps:
|
|
# Initialize with random noise like StreamDiffusion
|
|
self.init_noise = torch.randn(
|
|
(num_timesteps, 4, self.latent_height, self.latent_width),
|
|
device=comfy.model_management.get_torch_device(),
|
|
dtype=latent["samples"].dtype
|
|
)
|
|
|
|
# Initialize stock noise as zeros like StreamDiffusion
|
|
self.stock_noise = torch.zeros_like(self.init_noise[0])
|
|
|
|
# Create noise sequence from init_noise
|
|
self.noise_sequence = self.init_noise.clone()
|
|
|
|
# No need for manual shifting as it's handled in sample() method
|
|
|
|
return (sampler, buffered)
|
|
|
|
|
|
class StreamScheduler(ControlNodeBase):
|
|
"""Implements StreamDiffusion's efficient timestep selection"""
|
|
|
|
RETURN_TYPES = ("SIGMAS",)
|
|
FUNCTION = "update"
|
|
CATEGORY = "real-time/sampling"
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.cached_indices = None
|
|
self.last_t_index_list = None
|
|
|
|
@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. For LCM: use '32,45' for img2img or '0,16,32,45' for txt2img. For SDXL Turbo: use '45,49' for 2 steps or '49' for 1 step (Turbo works best with steps near the end of denoising)"
|
|
}),
|
|
"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 parse_timesteps(self, t_index_list):
|
|
"""Parse and sort timesteps, with caching"""
|
|
if t_index_list == self.last_t_index_list and self.cached_indices is not None:
|
|
return self.cached_indices
|
|
|
|
try:
|
|
# Parse and sort in reverse order once
|
|
indices = sorted([int(t.strip()) for t in t_index_list.split(",")], reverse=True)
|
|
self.last_t_index_list = t_index_list
|
|
self.cached_indices = indices
|
|
return indices
|
|
except ValueError as e:
|
|
print(f"Error parsing timesteps: {e}. Using default [32,45]")
|
|
self.last_t_index_list = "32,45"
|
|
self.cached_indices = [45, 32] # Already sorted in reverse
|
|
return self.cached_indices
|
|
|
|
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")
|
|
|
|
# Get parsed and sorted timesteps
|
|
t_indices = self.parse_timesteps(t_index_list)
|
|
|
|
# Create full schedule using normal scheduler
|
|
full_sigmas = comfy.samplers.normal_scheduler(model_sampling, num_inference_steps)
|
|
|
|
# Select only the sigmas at our desired indices
|
|
# No need to sort t_indices here since they're already sorted
|
|
selected_sigmas = []
|
|
for t in t_indices:
|
|
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]))
|
|
|
|
# Add final sigma
|
|
selected_sigmas.append(0.0)
|
|
|
|
# Create tensor directly on target device
|
|
device = comfy.model_management.get_torch_device()
|
|
selected_sigmas = torch.tensor(selected_sigmas, dtype=torch.float32, device=device)
|
|
return (selected_sigmas,)
|
|
|
|
|