227 lines
8.7 KiB
Python
227 lines
8.7 KiB
Python
from .base.control_base import ControlNodeBase
|
|
import torch
|
|
import comfy.model_management
|
|
import comfy.samplers
|
|
|
|
class StreamDiffusionLatentBuffer(ControlNodeBase):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
**super().INPUT_TYPES()["required"],
|
|
"latents": ("LATENT",),
|
|
"buffer_size": ("INT", {
|
|
"default": 4,
|
|
"min": 1,
|
|
"max": 16,
|
|
"step": 1,
|
|
"tooltip": "Number of latent frames to buffer"
|
|
}),
|
|
"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"
|
|
}),
|
|
"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")
|
|
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):
|
|
state = self.get_state({
|
|
"latent_buffer": None,
|
|
"current_idx": 0
|
|
})
|
|
|
|
# Get latent tensor from dict
|
|
samples = latents["samples"] # Shape: [1, 4, H, W]
|
|
|
|
# Initialize buffer if needed
|
|
if state["latent_buffer"] is None:
|
|
state["latent_buffer"] = torch.zeros((buffer_size, *samples.shape[1:]),
|
|
dtype=samples.dtype,
|
|
device=samples.device)
|
|
state["current_idx"] = 0
|
|
|
|
# Update buffer with current frame
|
|
state["latent_buffer"][state["current_idx"]] = samples[0].detach() # Remove batch dim
|
|
|
|
# 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
|
|
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 (
|
|
{"samples": samples}, # Current frame [1, 4, H, W]
|
|
{"samples": batched} # Batched buffer [buffer_size, 4, H, W]
|
|
)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return super().IS_CHANGED(**kwargs)
|
|
|
|
|
|
class StreamDiffusionRCFG(ControlNodeBase):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
**super().INPUT_TYPES()["required"],
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"rcfg_type": (["self", "initialize"],),
|
|
"guidance_scale": ("FLOAT", {
|
|
"default": 1.2,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.01,
|
|
"tooltip": "How strong the guidance should be"
|
|
}),
|
|
"delta": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"tooltip": "Scaling factor for residual noise"
|
|
})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
|
|
RETURN_NAMES = ("positive_cond", "negative_cond")
|
|
FUNCTION = "update"
|
|
CATEGORY = "real-time/conditioning/streamdiffusion"
|
|
DESCRIPTION = "Applies StreamDiffusion's Residual CFG optimizations to conditioning"
|
|
|
|
def __init__(self):
|
|
super().__init__() # Initialize ControlNodeBase
|
|
|
|
def update(self, positive, negative, rcfg_type, guidance_scale, delta, always_execute=True):
|
|
state = self.get_state({
|
|
"stock_noise": None,
|
|
"timestep_count": 0
|
|
})
|
|
|
|
# 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]
|
|
|
|
# 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
|
|
|
|
# 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
|
|
)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
# Always execute if enabled to maintain state between steps
|
|
return super().IS_CHANGED(**kwargs)
|
|
|
|
|
|
class StreamDiffusionSampler(ControlNodeBase):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
**super().INPUT_TYPES()["required"],
|
|
"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,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"tooltip": "Classifier-free guidance scale"
|
|
}),
|
|
"scheduler": (["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral", "lms"],),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "update"
|
|
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]
|
|
|
|
# Parse t_index_list from string
|
|
t_indices = [int(x.strip()) for x in t_index_list.split(",")]
|
|
|
|
# Create sampler object
|
|
sampler = comfy.samplers.sampler_object(scheduler)
|
|
|
|
# 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)
|
|
|
|
# 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)
|
|
|
|
# Return denoised latents - take only the first frame
|
|
return ({"samples": out[:1]},)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return super().IS_CHANGED(**kwargs) |