347 lines
13 KiB
Python
347 lines
13 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"],
|
|
"model": ("MODEL",),
|
|
"latents": ("LATENT",),
|
|
"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 frames to buffer"
|
|
}),
|
|
"do_add_noise": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Whether to add noise to intermediate states"
|
|
}),
|
|
"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"
|
|
}),
|
|
}
|
|
}
|
|
|
|
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, 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, # 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((total_buffer_size, *samples.shape[1:]),
|
|
dtype=samples.dtype,
|
|
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
|
|
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]
|
|
|
|
# Update index for next frame
|
|
state["current_idx"] = (state["current_idx"] + 1) % frame_buffer_size
|
|
self.set_state(state)
|
|
|
|
# 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 [total_buffer_size, 4, H, W]
|
|
timesteps # Timesteps tensor [num_timesteps]
|
|
)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
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",),
|
|
"latents": ("LATENT",),
|
|
"timesteps": ("TIMESTEPS",),
|
|
"cfg_type": (["none", "full", "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"
|
|
}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
|
|
"scheduler": (comfy.samplers.SCHEDULER_NAMES,),
|
|
"steps": ("INT", {
|
|
"default": 20,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "update"
|
|
CATEGORY = "real-time/sampling/streamdiffusion"
|
|
DESCRIPTION = "Applies StreamDiffusion's Residual CFG optimizations to noise predictions"
|
|
|
|
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)
|
|
|
|
# Create sampler
|
|
sampler = comfy.samplers.sampler_object(sampler_name)
|
|
|
|
# Calculate sigmas
|
|
sigmas = comfy.samplers.calculate_sigmas(
|
|
model.model.model_sampling,
|
|
scheduler,
|
|
steps
|
|
).to(model.model.device)
|
|
|
|
# Get sigmas for selected timesteps
|
|
timestep_sigmas = sigmas[timesteps]
|
|
|
|
# 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):
|
|
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",),
|
|
"latents": ("LATENT",),
|
|
"timesteps": ("TIMESTEPS",),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
|
|
"cfg": ("FLOAT", {
|
|
"default": 8.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.1,
|
|
"tooltip": "Classifier-free guidance scale"
|
|
}),
|
|
"num_inference_steps": ("INT", {
|
|
"default": 50,
|
|
"min": 1,
|
|
"max": 100,
|
|
"step": 1
|
|
}),
|
|
"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, latents, timesteps, sampler_name, cfg, num_inference_steps, seed, always_execute=True):
|
|
state = self.get_state({
|
|
"generator": None
|
|
})
|
|
|
|
# Get input latents
|
|
x_t = latents["samples"] # Shape: [batch_size, 4, H, W]
|
|
|
|
# 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
|
|
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
|
|
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
|
|
)
|
|
|
|
self.set_state(state)
|
|
|
|
return ({"samples": samples},)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return super().IS_CHANGED(**kwargs) |