From 46c1aef082b2805f7e5945cd89fb162421ab7146 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 8 Feb 2025 14:39:16 -0500 Subject: [PATCH] before --- __init__.py | 4 +- stream_diffusion_nodes.py | 328 ++++++++++++++++++++++++++------------ 2 files changed, 227 insertions(+), 105 deletions(-) diff --git a/__init__.py b/__init__.py index 6281ca8..97dcb88 100644 --- a/__init__.py +++ b/__init__.py @@ -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 = " 🕒🅡🅣🅝" diff --git a/stream_diffusion_nodes.py b/stream_diffusion_nodes.py index a340b57..6c6e9bd 100644 --- a/stream_diffusion_nodes.py +++ b/stream_diffusion_nodes.py @@ -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):