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

before

Browse Source
This commit is contained in:
RyanOnTheInside
2025-02-08 14:39:16 -05:00
parent 44759ab6d8
commit 46c1aef082
2 changed files with 227 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+3 -1
View File
@@ -21,7 +21,7 @@ NODE_CLASS_MAPPINGS = {
"IntegerMotionController": IntegerMotionController,
"YOLOSimilarityCompare": YOLOSimilarityCompare,
"TextRenderer": TextRenderer,
"ROINode": ROINode,
"ROINode": ROINode,
"QuickShapeMask": QuickShapeMask,
"DTypeConverter": DTypeConverter,
"FastWebcamCapture": FastWebcamCapture,
@@ -33,6 +33,8 @@ NODE_CLASS_MAPPINGS = {
}
NODE_DISPLAY_NAME_MAPPINGS = {}
suffix = " 🕒🅡🅣🅝"
stream_diffusion_nodes.py
+224 -104
View File
@@ -9,80 +9,137 @@ class StreamDiffusionLatentBuffer(ControlNodeBase):
return {
"required": {
**super().INPUT_TYPES()["required"],
"model": ("MODEL",),
"latents": ("LATENT",),
"buffer_size": ("INT", {
"default": 4,
"t_index_list": ("STRING", {
"default": "0,16,32,45",
"tooltip": "Comma-separated list of timestep indices"
}),
"frame_buffer_size": ("INT", {
"default": 1,
"min": 1,
"max": 16,
"step": 1,
"tooltip": "Number of latent frames to buffer"
"tooltip": "Number of frames to buffer"
}),
"add_noise": ("BOOLEAN", {
"do_add_noise": ("BOOLEAN", {
"default": True,
"tooltip": "Whether to add noise to intermediate states"
}),
"noise_alpha": ("FLOAT", {
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"tooltip": "Alpha for noise scaling"
"scheduler": (comfy.samplers.SCHEDULER_NAMES, {
"default": "lcm",
"tooltip": "Which scheduler to use for noise scheduling"
}),
"num_inference_steps": ("INT", {
"default": 50,
"min": 1,
"max": 100,
"step": 1,
"tooltip": "Total number of denoising steps"
}),
"noise_beta": ("FLOAT", {
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"tooltip": "Beta for noise scaling"
})
}
}
RETURN_TYPES = ("LATENT", "LATENT")
RETURN_NAMES = ("current_latents", "batched_latents")
RETURN_TYPES = ("LATENT", "LATENT", "TIMESTEPS")
RETURN_NAMES = ("current_latents", "batched_latents", "timesteps")
FUNCTION = "update"
CATEGORY = "real-time/latents/streamdiffusion"
DESCRIPTION = "Buffers and processes latent frames for StreamDiffusion's batch denoising"
def update(self, latents, buffer_size, add_noise, noise_alpha, noise_beta, always_execute=True):
def update(self, model, latents, t_index_list, frame_buffer_size, do_add_noise, scheduler, num_inference_steps, always_execute=True):
state = self.get_state({
"latent_buffer": None,
"current_idx": 0
"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
})
# Get latent tensor from dict
samples = latents["samples"] # Shape: [1, 4, H, W]
device = samples.device
# Initialize scheduler and timesteps if needed
if state["sigmas"] is None:
# Calculate sigmas using ComfyUI's scheduler system
sigmas = comfy.samplers.calculate_sigmas(
model.model.model_sampling,
scheduler,
num_inference_steps
).to(device)
state["sigmas"] = sigmas
# Parse t_index_list and get timesteps
state["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 = []
# 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)
# Calculate total buffer size based on timesteps and frame buffer
total_buffer_size = len(state["t_indices"]) * frame_buffer_size
# Initialize buffer if needed
if state["latent_buffer"] is None:
state["latent_buffer"] = torch.zeros((buffer_size, *samples.shape[1:]),
state["latent_buffer"] = torch.zeros((total_buffer_size, *samples.shape[1:]),
dtype=samples.dtype,
device=samples.device)
device=device)
# Initialize noise with proper seed for reproducibility
state["init_noise"] = torch.randn_like(state["latent_buffer"])
state["current_idx"] = 0
# Update buffer with current frame
state["latent_buffer"][state["current_idx"]] = samples[0].detach() # Remove batch dim
for i in range(len(state["t_indices"])):
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"][buffer_idx] # Add scaled noise
)
else:
state["latent_buffer"][buffer_idx] = samples[0]
# Create batched latents for processing
if add_noise:
# Add scaled noise to previous frames like StreamDiffusion
noise = torch.randn_like(state["latent_buffer"])
batched = noise_alpha * state["latent_buffer"] + noise_beta * noise
else:
batched = state["latent_buffer"].clone()
# Update index for next frame
state["current_idx"] = (state["current_idx"] + 1) % buffer_size
state["current_idx"] = (state["current_idx"] + 1) % frame_buffer_size
self.set_state(state)
# Return current frame and batched buffer
# Both in ComfyUI's LATENT format with batch dimension
# Reshape batched to [batch_size, channels, H, W] where batch_size = buffer_size
batched = batched.view(-1, *samples.shape[1:]) # Flatten to [buffer_size, 4, H, W]
# Return current frame, batched buffer and timesteps
batched = state["latent_buffer"].view(-1, *samples.shape[1:]) # Reshape to [total_buffer_size, 4, H, W]
# Create timesteps tensor
timesteps = torch.tensor(state["t_indices"], dtype=torch.long, device=device)
return (
{"samples": samples}, # Current frame [1, 4, H, W]
{"samples": batched} # Batched buffer [buffer_size, 4, H, W]
{"samples": batched}, # Batched buffer [total_buffer_size, 4, H, W]
timesteps # Timesteps tensor [num_timesteps]
)
@classmethod
@@ -90,15 +147,56 @@ class StreamDiffusionLatentBuffer(ControlNodeBase):
return super().IS_CHANGED(**kwargs)
class StreamDiffusionRCFGGuider:
def __init__(self, model, cfg_type="self", guidance_scale=1.0, delta=1.0):
self.model = model
self.cfg_type = cfg_type
self.guidance_scale = guidance_scale
self.delta = delta
self.stock_noise = None
self.init_noise = None
def __call__(self, x, timestep, cond, uncond, cfg_scale, model_options={}):
# Initialize stock noise if needed
if self.stock_noise is None:
self.stock_noise = torch.zeros_like(x)
self.init_noise = torch.randn_like(x)
with torch.no_grad():
# Get conditional prediction
noise_pred_text = self.model.apply_model(x, timestep, cond)
# Handle different CFG types
if self.cfg_type == "none" or cfg_scale <= 1.0:
return noise_pred_text
if self.cfg_type == "full":
# Get unconditional prediction
noise_pred_uncond = self.model.apply_model(x, timestep, uncond)
elif self.cfg_type == "self" or self.cfg_type == "initialize":
# Use stock noise for unconditional
noise_pred_uncond = self.stock_noise * self.delta
# Update stock noise for next iteration
if self.cfg_type == "self":
self.stock_noise = noise_pred_text.detach()
elif self.cfg_type == "initialize" and timestep == timestep[0]: # First step
self.stock_noise = noise_pred_text.detach()
# Apply CFG
return noise_pred_uncond + cfg_scale * (noise_pred_text - noise_pred_uncond)
class StreamDiffusionRCFG(ControlNodeBase):
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
**super().INPUT_TYPES()["required"],
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"rcfg_type": (["self", "initialize"],),
"latents": ("LATENT",),
"timesteps": ("TIMESTEPS",),
"cfg_type": (["none", "full", "self", "initialize"],),
"guidance_scale": ("FLOAT", {
"default": 1.2,
"min": 0.0,
@@ -112,55 +210,57 @@ class StreamDiffusionRCFG(ControlNodeBase):
"max": 1.0,
"step": 0.01,
"tooltip": "Scaling factor for residual noise"
})
}),
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
"scheduler": (comfy.samplers.SCHEDULER_NAMES,),
"steps": ("INT", {
"default": 20,
"min": 1,
"max": 100,
"step": 1
}),
}
}
RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
RETURN_NAMES = ("positive_cond", "negative_cond")
RETURN_TYPES = ("LATENT",)
FUNCTION = "update"
CATEGORY = "real-time/conditioning/streamdiffusion"
DESCRIPTION = "Applies StreamDiffusion's Residual CFG optimizations to conditioning"
CATEGORY = "real-time/sampling/streamdiffusion"
DESCRIPTION = "Applies StreamDiffusion's Residual CFG optimizations to noise predictions"
def __init__(self):
super().__init__() # Initialize ControlNodeBase
def update(self, model, positive, negative, latents, timesteps, cfg_type, guidance_scale, delta, sampler_name, scheduler, steps, always_execute=True):
# Create RCFG guider
rcfg_guider = StreamDiffusionRCFGGuider(model.model, cfg_type, guidance_scale, delta)
def update(self, positive, negative, rcfg_type, guidance_scale, delta, always_execute=True):
state = self.get_state({
"stock_noise": None,
"timestep_count": 0
})
# Create sampler
sampler = comfy.samplers.sampler_object(sampler_name)
# Get first conditioning tuple from input lists
pos_cond, pos_meta = positive[0] # Each conditioning is a list of tuples
neg_cond, neg_meta = negative[0]
# Calculate sigmas
sigmas = comfy.samplers.calculate_sigmas(
model.model.model_sampling,
scheduler,
steps
).to(model.model.device)
# Initialize stock_noise if needed
if state["stock_noise"] is None:
if rcfg_type == "self":
state["stock_noise"] = torch.zeros_like(neg_cond)
elif rcfg_type == "initialize" and state["timestep_count"] == 0:
state["stock_noise"] = neg_cond.clone() # Store initial negative cond
# Get sigmas for selected timesteps
timestep_sigmas = sigmas[timesteps]
# Calculate residual negative conditioning
residual_neg_cond = state["stock_noise"] * delta
# Update stock noise for next iteration (only in self mode)
if rcfg_type == "self":
state["stock_noise"] = pos_cond.detach() # Store current positive cond
state["timestep_count"] += 1
self.set_state(state)
# Return lists of tuples in ComfyUI's format
return (
[(pos_cond, pos_meta)], # Positive conditioning
[(residual_neg_cond, neg_meta)] # Negative conditioning with residual
# Sample with RCFG
samples = comfy.samplers.sample(
model,
latents["samples"], # Use input as noise
positive,
negative,
guidance_scale,
model.model.device,
sampler,
timestep_sigmas,
model_options={"custom_cfg": rcfg_guider} # Use our custom CFG
)
return ({"samples": samples},)
@classmethod
def IS_CHANGED(cls, **kwargs):
# Always execute if enabled to maintain state between steps
return super().IS_CHANGED(**kwargs)
@@ -173,19 +273,22 @@ class StreamDiffusionSampler(ControlNodeBase):
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"batched_latents": ("LATENT",), # Buffered latents from StreamDiffusionLatentBuffer
"t_index_list": ("STRING", {
"default": "0,16,32,45",
"tooltip": "Comma-separated list of timestep indices to use"
}),
"cfg_scale": ("FLOAT", {
"default": 1.2,
"latents": ("LATENT",),
"timesteps": ("TIMESTEPS",),
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
"cfg": ("FLOAT", {
"default": 8.0,
"min": 0.0,
"max": 10.0,
"max": 100.0,
"step": 0.1,
"tooltip": "Classifier-free guidance scale"
}),
"scheduler": (["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral", "lms"],),
"num_inference_steps": ("INT", {
"default": 50,
"min": 1,
"max": 100,
"step": 1
}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff})
}
}
@@ -195,32 +298,49 @@ class StreamDiffusionSampler(ControlNodeBase):
CATEGORY = "real-time/sampling"
DESCRIPTION = "StreamDiffusion's optimized sampler for real-time generation"
def update(self, model, positive, negative, batched_latents, t_index_list, cfg_scale, scheduler, seed, always_execute=True):
device = comfy.model_management.get_torch_device()
buffered_latents = batched_latents["samples"].to(device) # Shape: [buffer_size, 4, H, W]
def update(self, model, positive, negative, latents, timesteps, sampler_name, cfg, num_inference_steps, seed, always_execute=True):
state = self.get_state({
"generator": None
})
# Parse t_index_list from string
t_indices = [int(x.strip()) for x in t_index_list.split(",")]
# Get input latents
x_t = latents["samples"] # Shape: [batch_size, 4, H, W]
# Create sampler object
sampler = comfy.samplers.sampler_object(scheduler)
# Initialize generator if needed
if state["generator"] is None:
state["generator"] = torch.Generator(device=model.model.device)
state["generator"].manual_seed(seed)
# Get sampler object
sampler = comfy.samplers.sampler_object(sampler_name)
# Calculate sigmas for timesteps
timesteps = torch.linspace(999, 0, max(t_indices) + 2, device=device)[:-1] # +2 for safety
sigmas = []
for t_idx in t_indices:
sigma = model.model.model_sampling.sigma(torch.tensor([timesteps[t_idx]], device=device))
sigmas.append(sigma[0])
sigmas = torch.stack(sigmas)
sigmas = comfy.samplers.calculate_sigmas(
model.model.model_sampling,
"exponential", # Use exponential scheduler by default
num_inference_steps
).to(model.model.device)
# Get sigmas for selected timesteps
timestep_sigmas = sigmas[timesteps]
# Sample using ComfyUI's infrastructure
torch.manual_seed(seed)
noise = torch.randn_like(buffered_latents, device=device)
out = comfy.samplers.sample(model, noise, positive, negative, cfg_scale,
device, sampler, sigmas, latent_image=buffered_latents)
samples = comfy.samplers.sample(
model,
x_t, # Use input latents as noise
positive,
negative,
cfg,
model.model.device,
sampler,
timestep_sigmas,
denoise_mask=None, # No mask support yet
seed=seed
)
# Return denoised latents - take only the first frame
return ({"samples": out[:1]},)
self.set_state(state)
return ({"samples": samples},)
@classmethod
def IS_CHANGED(cls, **kwargs):