diff --git a/__init__.py b/__init__.py index d7c660e..5671113 100644 --- a/__init__.py +++ b/__init__.py @@ -3,9 +3,11 @@ from .controls.sequence_controls import FloatSequence, IntSequence, StringSequen from .controls.utility_controls import FPSMonitor, SimilarityFilter, LazyCondition from .controls.motion_controls import MotionController, ROINode, IntegerMotionController from .misc_nodes import DTypeConverter, FastWebcamCapture, YOLOSimilarityCompare, TextRenderer, QuickShapeMask, MultilineText, LoadImageFromPath_ - +from .stream_diffusion_nodes import StreamConditioning, StreamCFG import re + + NODE_CLASS_MAPPINGS = { "FloatControl": FloatControl, "IntControl": IntControl, @@ -15,12 +17,16 @@ NODE_CLASS_MAPPINGS = { "StringSequence": StringSequence, "FPSMonitor": FPSMonitor, "SimilarityFilter": SimilarityFilter, + "StreamCFG": StreamCFG, + "StreamConditioning": StreamConditioning, "LazyCondition": LazyCondition, "MotionController": MotionController, "IntegerMotionController": IntegerMotionController, "YOLOSimilarityCompare": YOLOSimilarityCompare, "TextRenderer": TextRenderer, + "ROINode": ROINode, + #"IntervalControl": IntervalCo ntrol, #"DeltaControl": DeltaControl, "QuickShapeMask": QuickShapeMask, diff --git a/controls/utility_controls.py b/controls/utility_controls.py index 9b3f382..8f29533 100644 --- a/controls/utility_controls.py +++ b/controls/utility_controls.py @@ -269,5 +269,3 @@ class LazyCondition(ControlNodeBase): - - diff --git a/stream_diffusion_nodes.py b/stream_diffusion_nodes.py new file mode 100644 index 0000000..4ccb605 --- /dev/null +++ b/stream_diffusion_nodes.py @@ -0,0 +1,294 @@ +import torch +from .base.control_base import ControlNodeBase + +class StreamCFG(ControlNodeBase): + """Implements CFG approaches for temporal consistency between workflow runs""" + + RETURN_TYPES = ("MODEL",) + FUNCTION = "update" + CATEGORY = "real-time/sampling" + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "model": ("MODEL",), + "cfg_type": (["self", "full", "initialize"], { + "default": "self", + "tooltip": "Type of CFG to use: full (standard), self (memory efficient), or initialize (memory efficient with initialization)" + }), + "residual_scale": ("FLOAT", { + "default": 0.4, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Scale factor for residual (higher = more temporal consistency)" + }), + "delta": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 5.0, + "step": 0.1, + "tooltip": "Delta parameter for self/initialize CFG types" + }), + }) + return inputs + + def __init__(self): + super().__init__() + # Store the last model to detect when we need to reapply the hook + self.last_model_hash = None + self.post_cfg_function = None + + def update(self, model, always_execute=True, cfg_type="self", residual_scale=0.4, delta=1.0): + print(f"[StreamCFG] Initializing with cfg_type={cfg_type}, residual_scale={residual_scale}, delta={delta}") + + # Get state with defaults + state = self.get_state({ + "last_uncond": None, # Store last workflow's unconditioned prediction + "initialized": False, + "cfg_type": cfg_type, + "residual_scale": residual_scale, + "delta": delta, + "workflow_count": 0, # Track number of workflow runs + "current_sigmas": None, # Track sigmas for this workflow + "seen_sigmas": set(), # Track which sigmas we've seen this workflow + "is_last_step": False, # Track if we're on the last step + "alpha_prod_t": None, # Store alpha values for proper scaling + "beta_prod_t": None, # Store beta values for proper scaling + }) + + def post_cfg_function(args): + # Extract info + denoised = args["denoised"] + cond = args["cond"] + uncond = args["uncond"] + cond_denoised = args["cond_denoised"] + uncond_denoised = args["uncond_denoised"] + cond_scale = args["cond_scale"] + sigma = args["sigma"].item() if torch.is_tensor(args["sigma"]) else args["sigma"] + + # Get step info from model options + model_options = args["model_options"] + sample_sigmas = model_options["transformer_options"].get("sample_sigmas", None) + + # Update current sigmas if needed + if sample_sigmas is not None and state["current_sigmas"] is None: + # Filter out the trailing 0.0 if present + sigmas = [s.item() for s in sample_sigmas] + if sigmas[-1] == 0.0: + sigmas = sigmas[:-1] + state["current_sigmas"] = sigmas + state["seen_sigmas"] = set() + + # Calculate alpha and beta values for proper scaling + alphas = [1.0 / (1.0 + s**2) for s in sigmas] + state["alpha_prod_t"] = torch.tensor(alphas, device=denoised.device, dtype=denoised.dtype) + state["beta_prod_t"] = torch.sqrt(1 - state["alpha_prod_t"]) + + # Track this sigma + state["seen_sigmas"].add(sigma) + + # Check if this is the last step + state["is_last_step"] = False + if state["current_sigmas"] is not None: + # It's the last step if we've seen all sigmas + is_last_step = len(state["seen_sigmas"]) >= len(state["current_sigmas"]) + # Or if this is the smallest sigma in the sequence + if not is_last_step and sigma == min(state["current_sigmas"]): + is_last_step = True + state["is_last_step"] = is_last_step + + # First workflow case + if state["last_uncond"] is None: + if state["is_last_step"]: + state["last_uncond"] = uncond_denoised.detach().clone() + state["workflow_count"] += 1 + state["current_sigmas"] = None # Reset for next workflow + if cfg_type == "initialize": + state["initialized"] = True + self.set_state(state) + return denoised + + # Handle different CFG types for subsequent workflows + if cfg_type == "full": + result = denoised + + elif cfg_type == "initialize" and not state["initialized"]: + result = denoised + if state["is_last_step"]: + state["initialized"] = True + self.set_state(state) + + else: # self or initialized initialize + # Get current step index + current_idx = len(state["seen_sigmas"]) - 1 + + # Scale last prediction with proper alpha/beta values + noise_pred_uncond = state["last_uncond"] * delta + + # Apply CFG with scaled prediction + result = noise_pred_uncond + cond_scale * (cond_denoised - noise_pred_uncond) + + # Store last prediction if this is the last step + if state["is_last_step"]: + # Calculate properly scaled residual + scaled_noise = state["beta_prod_t"][current_idx] * state["last_uncond"] + delta_x = uncond_denoised - scaled_noise + + # Scale delta_x with next step's alpha/beta + if current_idx < len(state["current_sigmas"]) - 1: + alpha_next = state["alpha_prod_t"][current_idx + 1] + beta_next = state["beta_prod_t"][current_idx + 1] + else: + alpha_next = torch.ones_like(state["alpha_prod_t"][0]) + beta_next = torch.ones_like(state["beta_prod_t"][0]) + + delta_x = alpha_next * delta_x / beta_next + + # Update stored prediction with scaled residual + final_update = uncond_denoised + residual_scale * delta_x + state["last_uncond"] = final_update + state["workflow_count"] += 1 + state["current_sigmas"] = None # Reset for next workflow + self.set_state(state) + + return result + + # Store function reference to prevent garbage collection + self.post_cfg_function = post_cfg_function + + # Only set up post CFG function if model has changed + model_hash = hash(str(model)) + if model_hash != self.last_model_hash: + m = model.clone() + m.model_options = m.model_options.copy() + m.model_options["sampler_post_cfg_function"] = [self.post_cfg_function] + self.last_model_hash = model_hash + return (m,) + + # Make sure our function is still in the list + if not any(f is self.post_cfg_function for f in model.model_options.get("sampler_post_cfg_function", [])): + m = model.clone() + m.model_options = m.model_options.copy() + m.model_options["sampler_post_cfg_function"] = [self.post_cfg_function] + return (m,) + + return (model,) + + +class StreamConditioning(ControlNodeBase): + """Applies Residual CFG to conditioning for improved temporal consistency with different CFG types""" + #NOTE: experimental + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "cfg_type": (["full", "self", "initialize"], { + "default": "full", + "tooltip": "Type of CFG to use: full (standard), self (memory efficient), or initialize (memory efficient with initialization)" + }), + "residual_scale": ("FLOAT", { + "default": 0.4, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Scale factor for residual conditioning (higher = more temporal consistency)" + }), + "delta": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 5.0, + "step": 0.1, + "tooltip": "Delta parameter for self/initialize CFG types" + }), + "always_execute": ("BOOLEAN", { + "default": False, + }), + } + } + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING") + RETURN_NAMES = ("positive", "negative") + FUNCTION = "update" + CATEGORY = "real-time/control/utility" + + def __init__(self): + super().__init__() + + def update(self, positive, negative, cfg_type="full", residual_scale=0.4, delta=1.0, always_execute=False): + # Get state with defaults + state = self.get_state({ + "prev_positive": None, + "prev_negative": None, + "stock_noise": None, # For self/initialize CFG + "initialized": False # For initialize CFG + }) + + # Extract conditioning tensors + current_pos_cond = positive[0][0] # Assuming standard ComfyUI conditioning format + current_neg_cond = negative[0][0] + + # First frame case + if state["prev_positive"] is None: + state["prev_positive"] = current_pos_cond.detach().clone() + state["prev_negative"] = current_neg_cond.detach().clone() + if cfg_type == "initialize": + # For initialize, we use the first negative as our stock noise + state["stock_noise"] = current_neg_cond.detach().clone() + elif cfg_type == "self": + # For self, we start with a scaled version of the negative + state["stock_noise"] = current_neg_cond.detach().clone() * delta + self.set_state(state) + return (positive, negative) + + # Handle different CFG types + if cfg_type == "full": + # Standard R-CFG with full negative conditioning + pos_residual = current_pos_cond - state["prev_positive"] + neg_residual = current_neg_cond - state["prev_negative"] + + blended_pos = current_pos_cond + residual_scale * pos_residual + blended_neg = current_neg_cond + residual_scale * neg_residual + + # Update state + state["prev_positive"] = current_pos_cond.detach().clone() + state["prev_negative"] = current_neg_cond.detach().clone() + + # Reconstruct conditioning format + positive_out = [[blended_pos, positive[0][1]]] + negative_out = [[blended_neg, negative[0][1]]] + + else: # self or initialize + # Calculate residual for positive conditioning + pos_residual = current_pos_cond - state["prev_positive"] + blended_pos = current_pos_cond + residual_scale * pos_residual + + # Update stock noise based on current prediction + if cfg_type == "initialize" and not state["initialized"]: + # First prediction for initialize type + state["stock_noise"] = current_neg_cond.detach().clone() + state["initialized"] = True + else: + # Update stock noise with temporal consistency + stock_residual = current_neg_cond - state["stock_noise"] + state["stock_noise"] = current_neg_cond + residual_scale * stock_residual + + # Scale stock noise by delta + scaled_stock = state["stock_noise"] * delta + + # Update state + state["prev_positive"] = current_pos_cond.detach().clone() + state["prev_negative"] = scaled_stock.detach().clone() + + # Reconstruct conditioning format + positive_out = [[blended_pos, positive[0][1]]] + negative_out = [[scaled_stock, negative[0][1]]] + + self.set_state(state) + return (positive_out, negative_out) + + +