diff --git a/.ipynb_checkpoints/__init__-checkpoint.py b/.ipynb_checkpoints/__init__-checkpoint.py new file mode 100644 index 0000000..fafac66 --- /dev/null +++ b/.ipynb_checkpoints/__init__-checkpoint.py @@ -0,0 +1,69 @@ +from .controls.value_controls import FloatControl, IntControl, StringControl +from .controls.sequence_controls import FloatSequence, IntSequence, StringSequence +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 StreamCrossAttention +from .stream_sampler import StreamBatchSampler, StreamScheduler, StreamFrameBuffer +from .stream_cfg import StreamCFG +from .stream_conditioning import StreamConditioning +import re +from .particle_nodes import TemporalParticleDepthNode + + +NODE_CLASS_MAPPINGS = { + "FloatControl": FloatControl, + "IntControl": IntControl, + "StringControl": StringControl, + "FloatSequence": FloatSequence, + "IntSequence": IntSequence, + "StringSequence": StringSequence, + "FPSMonitor": FPSMonitor, + "SimilarityFilter": SimilarityFilter, + "StreamCFG": StreamCFG, + "StreamConditioning": StreamConditioning, + "StreamBatchSampler": StreamBatchSampler, + "StreamScheduler": StreamScheduler, + "StreamFrameBuffer": StreamFrameBuffer, + "StreamCrossAttention": StreamCrossAttention, + "LazyCondition": LazyCondition, + "MotionController": MotionController, + "IntegerMotionController": IntegerMotionController, + "YOLOSimilarityCompare": YOLOSimilarityCompare, + "TextRenderer": TextRenderer, + + + "ROINode": ROINode, + + #"IntervalControl": IntervalCo ntrol, + #"DeltaControl": DeltaControl, + "QuickShapeMask": QuickShapeMask, + "DTypeConverter": DTypeConverter, + "FastWebcamCapture": FastWebcamCapture, + "MultilineText": MultilineText, + "LoadImageFromPath_": LoadImageFromPath_, + "TemporalParticleDepthNode": TemporalParticleDepthNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = {} + +suffix = " 🕒🅡🅣🅝" + +for node_name in NODE_CLASS_MAPPINGS.keys(): + # Convert camelCase or snake_case to Title Case + if node_name not in NODE_DISPLAY_NAME_MAPPINGS: + display_name = ' '.join(word.capitalize() for word in re.findall(r'[A-Z]?[a-z]+|[A-Z]+(?=[A-Z][a-z]|\d|\W|$)|\d+', node_name)) + else: + display_name = NODE_DISPLAY_NAME_MAPPINGS[node_name] + + # Add the suffix if it's not already present + if not display_name.endswith(suffix): + display_name += suffix + + # Assign the final display name to the mappings + NODE_DISPLAY_NAME_MAPPINGS[node_name] = display_name + + +WEB_DIRECTORY = "./web/js" + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/.ipynb_checkpoints/particle_nodes-checkpoint.py b/.ipynb_checkpoints/particle_nodes-checkpoint.py new file mode 100644 index 0000000..fc2a55c --- /dev/null +++ b/.ipynb_checkpoints/particle_nodes-checkpoint.py @@ -0,0 +1,350 @@ +import torch +import numpy as np +from .base.control_base import ControlNodeBase + +class TemporalParticleDepthNode(ControlNodeBase): + """ + A node that generates temporally consistent particle depth maps with a controllable origin point. + """ + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "origin_x": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "display": "slider", + "tooltip": "X coordinate of the origin point (0-1)" + }), + "origin_y": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "display": "slider", + "tooltip": "Y coordinate of the origin point (0-1)" + }), + "width": ("INT", { + "default": 512, + "min": 64, + "max": 2048, + "step": 8, + "tooltip": "Width of the output depth map" + }), + "height": ("INT", { + "default": 512, + "min": 64, + "max": 2048, + "step": 8, + "tooltip": "Height of the output depth map" + }), + "num_particles": ("INT", { + "default": 100, + "min": 10, + "max": 1000, + "step": 10, + "tooltip": "Number of particles to simulate" + }), + "particle_size": ("FLOAT", { + "default": 0.03, + "min": 0.001, + "max": 0.2, + "step": 0.001, + "tooltip": "Size of each particle (relative to image size)" + }), + "speed": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 3.0, + "step": 0.1, + "display": "slider", + "tooltip": "Overall movement speed of particles" + }), + "inertia_factor": ("FLOAT", { + "default": 0.1, + "min": 0.01, + "max": 0.5, + "step": 0.01, + "display": "slider", + "tooltip": "Controls how quickly particles respond to changes (lower = more inertia)" + }), + "distance_response": ("FLOAT", { + "default": 0.15, + "min": -0.3, + "max": 0.3, + "step": 0.01, + "display": "slider", + "tooltip": "How distance affects response speed (positive = closer particles respond faster, negative = distant particles respond faster, zero = uniform response)" + }), + }) + inputs["optional"] = { + "hand_data": ("HAND_DATA", { + "tooltip": "Optional hand tracking data to use for origin point" + }), + "hand_keypoint": (["palm", "thumb_tip", "index_tip", "middle_tip", "ring_tip", "pinky_tip"], { + "default": "index_tip", + "tooltip": "Which hand keypoint to use as origin (if hand_data is provided)" + }), + } + return inputs + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "update" + CATEGORY = "image/generators" + + def __init__(self): + super().__init__() + self.generator = None + self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + def update(self, origin_x, origin_y, width, height, num_particles, particle_size, speed, inertia_factor, distance_response, always_execute=True, hand_data=None, hand_keypoint="index_tip"): + # Initialize generator if needed or if parameters changed + if self.generator is None: + self.generator = TemporalParticleDepthGenerator( + batch_size=1, + height=height, + width=width, + num_particles=num_particles, + initial_origin_x=origin_x, + initial_origin_y=origin_y, + particle_size=particle_size, + speed=speed, + inertia_factor=inertia_factor, + distance_response=distance_response, + device=self.device + ) + elif (self.generator.shape[1] != height or + self.generator.shape[2] != width or + self.generator.num_particles != num_particles): + # Reinitialize if dimensions or particle count changed + self.generator = TemporalParticleDepthGenerator( + batch_size=1, + height=height, + width=width, + num_particles=num_particles, + initial_origin_x=origin_x, + initial_origin_y=origin_y, + particle_size=particle_size, + speed=speed, + inertia_factor=inertia_factor, + distance_response=distance_response, + device=self.device + ) + + # Use hand data if provided + if hand_data is not None and len(hand_data) > 0: + # Extract the first hand's data + hand_info = hand_data[0] + + # Check if hands are present + if hand_info["hands_present"]: + # Determine which hand to use (prefer right hand if available) + hand_landmarks = None + if hand_info["right_hand"] is not None: + hand_landmarks = hand_info["right_hand"] + elif hand_info["left_hand"] is not None: + hand_landmarks = hand_info["left_hand"] + + if hand_landmarks is not None: + # Map keypoint to index + keypoint_map = { + "palm": 0, # Center of palm + "thumb_tip": 4, # Thumb tip + "index_tip": 8, # Index finger tip + "middle_tip": 12, # Middle finger tip + "ring_tip": 16, # Ring finger tip + "pinky_tip": 20 # Pinky tip + } + + if hand_keypoint in keypoint_map: + keypoint_idx = keypoint_map[hand_keypoint] + # Get normalized coordinates (already in 0-1 range) + if keypoint_idx < len(hand_landmarks): + # X and Y coordinates are at indices 0 and 1 + origin_x = hand_landmarks[keypoint_idx][0] + origin_y = hand_landmarks[keypoint_idx][1] + + # Generate depth map with current origin + depth_map = self.generator.update( + new_origin_x=origin_x, + new_origin_y=origin_y, + particle_size=particle_size, + speed=speed, + inertia_factor=inertia_factor, + distance_response=distance_response + ) + + # ComfyUI already uses BHWC format, so no permute needed + return (depth_map,) + + +import torch + +class TemporalParticleDepthGenerator: + def __init__(self, + batch_size: int, + height: int, + width: int, + num_particles: int = 100, + initial_origin_x: float = 0.5, + initial_origin_y: float = 0.5, + particle_size: float = 0.03, + speed: float = 1.0, + inertia_factor: float = 0.1, + distance_response: float = 0.15, + device: str = 'cuda' if torch.cuda.is_available() else 'cpu'): + """ + Initialize a temporally consistent particle depth map generator with inertial origin movement. + """ + self.device = torch.device(device) + self.shape = (batch_size, height, width) + self.particle_size = particle_size + self.speed = speed + self.base_inertia_factor = inertia_factor + self.distance_response = distance_response + + # Coordinate grid (static) + y = torch.linspace(0, 1, height, device=device) + x = torch.linspace(0, 1, width, device=device) + self.y_grid, self.x_grid = torch.meshgrid(y, x, indexing='ij') + + # Persistent particle properties + self.num_particles = num_particles + self.angles = torch.rand(num_particles, device=device) * 2 * torch.pi + self.base_speeds = torch.rand(num_particles, device=device) * 0.02 + 0.005 + self.speeds = self.base_speeds * self.speed + self.distances = torch.zeros(num_particles, device=device) + + # Origin (current and target) + self.origin_x = torch.tensor(initial_origin_x, device=device) + self.origin_y = torch.tensor(initial_origin_y, device=device) + self.target_origin_x = self.origin_x.clone() + self.target_origin_y = self.origin_y.clone() + + # Particle-specific inertia factors (will be computed based on distance) + self.inertia_factors = torch.ones(num_particles, device=device) * inertia_factor + + # For respawning particles + self.active = torch.ones(num_particles, dtype=torch.bool, device=device) + + # Frame counters + self.frame_count = 0 + self.last_major_refresh = 0 + + def update(self, + new_origin_x: float = None, + new_origin_y: float = None, + particle_size: float = None, + speed: float = None, + inertia_factor: float = None, + distance_response: float = None) -> torch.Tensor: + """ + Update particle positions and generate next depth map frame with inertial origin movement. + + Args: + new_origin_x (float, optional): Target origin X (0-1) + new_origin_y (float, optional): Target origin Y (0-1) + particle_size (float, optional): Update particle size + speed (float, optional): Update overall movement speed + inertia_factor (float, optional): Update inertia factor + distance_response (float, optional): Update distance response factor + + Returns: + torch.Tensor: Depth map (B, H, W, 3) + """ + self.frame_count += 1 + + # Update target origin if provided + if new_origin_x is not None: + self.target_origin_x = torch.tensor(new_origin_x, device=self.device) + if new_origin_y is not None: + self.target_origin_y = torch.tensor(new_origin_y, device=self.device) + if particle_size is not None: + self.particle_size = particle_size + if speed is not None: + self.speed = speed + # Update actual speeds based on base speeds and speed multiplier + self.speeds = self.base_speeds * self.speed + if inertia_factor is not None: + self.base_inertia_factor = inertia_factor + if distance_response is not None: + self.distance_response = distance_response + + # Periodic randomization + if self.frame_count - self.last_major_refresh >= 200: + refresh_mask = torch.rand(self.num_particles, device=self.device) < 0.3 + if refresh_mask.any(): + num_refresh = refresh_mask.sum() + self.distances[refresh_mask] = torch.rand(num_refresh, device=self.device) * 0.3 + self.angles[refresh_mask] = torch.rand(num_refresh, device=self.device) * 2 * torch.pi + self.base_speeds[refresh_mask] = torch.rand(num_refresh, device=self.device) * 0.02 + 0.005 + self.speeds[refresh_mask] = self.base_speeds[refresh_mask] * self.speed + + speed_adjust = (torch.rand(self.num_particles, device=self.device) * 0.01) - 0.005 + self.base_speeds = torch.clamp(self.base_speeds + speed_adjust, 0.002, 0.03) + self.speeds = self.base_speeds * self.speed + self.last_major_refresh = self.frame_count + + # Update distances + self.distances = self.distances + self.speeds + + # Respawn off-screen particles + off_screen = self.distances > 1.414 + if off_screen.any(): + num_respawn = off_screen.sum() + self.distances[off_screen] = torch.rand(num_respawn, device=self.device) * 0.1 + self.angles[off_screen] = torch.rand(num_respawn, device=self.device) * 2 * torch.pi + self.speeds[off_screen] = torch.rand(num_respawn, device=self.device) * 0.02 + 0.005 + + # Calculate current particle positions (before origin update) + particle_x = self.origin_x + torch.cos(self.angles) * self.distances + particle_y = self.origin_y + torch.sin(self.angles) * self.distances + + # Compute distance-based inertia factors + particle_distances = torch.sqrt((particle_x - self.origin_x)**2 + (particle_y - self.origin_y)**2) + + # Calculate inertia based on distance and user parameters + # When distance_response is positive, closer particles respond faster + # When negative, distant particles respond faster + # When zero, all particles respond uniformly + base_response = self.base_inertia_factor * 2 + distance_effect = particle_distances * self.distance_response * self.base_inertia_factor + + if self.distance_response >= 0: + # Positive: closer particles respond faster (subtract distance effect) + self.inertia_factors = torch.clamp( + base_response - distance_effect, + 0.02 * self.base_inertia_factor, + 0.2 * self.base_inertia_factor + ) + else: + # Negative: distant particles respond faster (add distance effect) + self.inertia_factors = torch.clamp( + base_response + distance_effect, + 0.02 * self.base_inertia_factor, + 0.2 * self.base_inertia_factor + ) + + # Interpolate origin toward target with inertia + self.origin_x = self.origin_x + self.inertia_factors * (self.target_origin_x - self.origin_x) + self.origin_y = self.origin_y + self.inertia_factors * (self.target_origin_y - self.origin_y) + + # Recalculate particle positions with updated origin + particle_x = self.origin_x + torch.cos(self.angles) * self.distances + particle_y = self.origin_y + torch.sin(self.angles) * self.distances + + # Generate depth map + depth = torch.zeros(self.shape, device=self.device) + for i in range(self.num_particles): + if not self.active[i]: + continue + + dist = torch.sqrt((self.x_grid - particle_x[i])**2 + + (self.y_grid - particle_y[i])**2) + particle = torch.exp(-(dist**2) / (2 * self.particle_size**2)) + depth[0] = torch.maximum(depth[0], particle * self.distances[i]) + + depth = depth / (depth.max() + 1e-6) + return depth.unsqueeze(-1).expand(-1, -1, -1, 3).clamp(0, 1) \ No newline at end of file diff --git a/__init__.py b/__init__.py index 9e79adc..eb112a5 100644 --- a/__init__.py +++ b/__init__.py @@ -8,6 +8,7 @@ from .stream_sampler import StreamBatchSampler, StreamScheduler, StreamFrameBuf from .stream_cfg import StreamCFG from .stream_conditioning import StreamConditioning from .particle_nodes import DepthMapWarpNode +from .dimensions import ImageDimensions import re @@ -43,6 +44,7 @@ NODE_CLASS_MAPPINGS = { "FastWebcamCapture": FastWebcamCapture, "MultilineText": MultilineText, "LoadImageFromPath_": LoadImageFromPath_, + "ImageDimensions": ImageDimensions, } NODE_DISPLAY_NAME_MAPPINGS = {} diff --git a/dimensions.py b/dimensions.py new file mode 100644 index 0000000..db98845 --- /dev/null +++ b/dimensions.py @@ -0,0 +1,24 @@ +class ImageDimensions: + @classmethod + def INPUT_TYPES(s): + return {"required": {"image": ("IMAGE",)}} + + RETURN_TYPES = ("INT", "INT", "INT") + RETURN_NAMES = ("width", "height", "count") + FUNCTION = "get_dimensions" + CATEGORY = "image/info" + + def get_dimensions(self, image): + count = image.shape[0] # Batch size + height = image.shape[1] + width = image.shape[2] + print(f"Image dimensions: {width}x{height} (batch size: {count})") + return (width, height, count) + +NODE_CLASS_MAPPINGS = { + "ImageDimensions": ImageDimensions +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageDimensions": "Get Image Dimensions" +} \ No newline at end of file diff --git a/stream_sampler.py b/stream_sampler.py index e9cb1a2..6cb5f2e 100644 --- a/stream_sampler.py +++ b/stream_sampler.py @@ -1,17 +1,9 @@ import torch import comfy.model_management import comfy.samplers -import random -import time -import os class StreamBatchSampler: - - RETURN_TYPES = ("SAMPLER",) - FUNCTION = "update" - CATEGORY = "StreamPack/sampling" - DESCRIPTION = "Implements batched denoising for faster inference by processing multiple frames in parallel at different denoising steps. Also adds temportal consistency to the denoising process." @classmethod def INPUT_TYPES(cls): return { @@ -26,6 +18,11 @@ class StreamBatchSampler: }, } + RETURN_TYPES = ("SAMPLER",) + FUNCTION = "update" + CATEGORY = "StreamPack/sampling" + DESCRIPTION = "Implements batched denoising for faster inference by processing multiple frames in parallel at different denoising steps. Also adds temportal consistency to the denoising process." + def __init__(self): self.num_steps = None self.frame_buffer = [] @@ -33,7 +30,7 @@ class StreamBatchSampler: self.stock_noise = None self.is_txt2img_mode = False - # Initialize all optimization buffers as None + # Initialize all buffers self.zeros_reference = None self.random_noise_buffer = None self.sigmas_view_buffer = None @@ -49,18 +46,16 @@ class StreamBatchSampler: batch_size = noise.shape[0] num_sigmas = len(sigmas) - 1 # Subtract 1 because last sigma is the target (0.0) - # Optimization 1: Reuse zeros buffer for txt2img detection + + # Reuse zeros buffer for txt2img detection if self.zeros_reference is None: - # We only need a small reference tensor for comparison, not a full tensor self.zeros_reference = torch.zeros(1, device=noise.device, dtype=noise.dtype) - # Check if noise tensor is all zeros - functionally identical but more efficient + # Check if noise tensor is all zeros self.is_txt2img_mode = torch.abs(noise).sum() < 1e-5 - # Noise handling with memory optimization if self.is_txt2img_mode: - # Optimization 2: If txt2img mode, reuse the noise tensor directly - # instead of allocating new memory + # If txt2img mode, reuse the noise tensor directly if self.random_noise_buffer is None or self.random_noise_buffer.shape != noise.shape: self.random_noise_buffer = torch.empty_like(noise) @@ -69,7 +64,7 @@ class StreamBatchSampler: x = self.random_noise_buffer # Use pre-allocated buffer else: # If not txt2img, we'll still need to add noise later - x = noise # No need to copy, will add noise later + x = noise # Verify batch size matches number of timesteps if batch_size != num_sigmas: @@ -79,12 +74,12 @@ class StreamBatchSampler: alpha_prod_t = (sigmas[:-1] / sigmas[0]).view(-1, 1, 1, 1) # [B,1,1,1] beta_prod_t = (1 - alpha_prod_t) - # Optimization 3: Initialize stock noise with reuse + # Initialize stock noise with reuse if self.stock_noise is None or self.stock_noise.shape != noise[0].shape: self.stock_noise = torch.empty_like(noise[0]) self.stock_noise.normal_() # Generate random noise in-place - # Optimization 4: Pre-allocate and reuse view buffer for sigmas + # Pre-allocate and reuse view buffer for sigmas if self.sigmas_view_buffer is None or self.sigmas_view_buffer.shape[0] != len(sigmas)-1: self.sigmas_view_buffer = torch.empty((len(sigmas)-1, 1, 1, 1), device=sigmas.device, @@ -92,7 +87,7 @@ class StreamBatchSampler: # In-place copy of sigmas view self.sigmas_view_buffer.copy_(sigmas[:-1].view(-1, 1, 1, 1)) - # Optimization 5: Eliminate unsqueeze allocation by pre-expanding stock noise + # Eliminate unsqueeze allocation by pre-expanding stock noise if self.expanded_stock_noise is None or self.expanded_stock_noise.shape[0] != batch_size: self.expanded_stock_noise = self.stock_noise.expand(batch_size, *self.stock_noise.shape) @@ -108,18 +103,15 @@ class StreamBatchSampler: # Initialize and manage latent buffer with memory optimization if (self.x_t_latent_buffer is None or self.is_txt2img_mode) and num_sigmas > 1: - # Optimization 6: Pre-allocate or resize as needed + # Pre-allocate or resize as needed if self.x_t_latent_buffer is None or self.x_t_latent_buffer.shape != x[0].shape: self.x_t_latent_buffer = torch.empty_like(x[0]) - # In-place copy instead of clone self.x_t_latent_buffer.copy_(x[0]) # Use buffer for first frame to maintain temporal consistency if num_sigmas > 1: - # In-place update - no new allocation x[0].copy_(self.x_t_latent_buffer) - # Run model on entire batch at once with torch.no_grad(): # Process all frames in parallel sigma_batch = sigmas[:-1] @@ -131,7 +123,7 @@ class StreamBatchSampler: # Store result from first frame as buffer for next iteration self.x_t_latent_buffer.copy_(denoised_batch[0]) # In-place update - # Optimization 7: Pre-allocate output buffer + # Pre-allocate output buffer if self.output_buffer is None or self.output_buffer.shape != (1, *denoised_batch[-1].shape): self.output_buffer = torch.empty(1, *denoised_batch[-1].shape, device=denoised_batch.device, @@ -158,17 +150,13 @@ class StreamBatchSampler: class StreamScheduler: - RETURN_TYPES = ("SIGMAS",) - FUNCTION = "update" - CATEGORY = "StreamPack/sampling" - DESCRIPTION = "Implements StreamDiffusion's efficient timestep selection. Use in conjunction with StreamBatchSampler." @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), "t_index_list": ("STRING", { - "default": "32,45", + "default": "0,16,32,49", "tooltip": "Comma-separated list of timesteps to actually use for denoising. Examples: '32,45' for img2img or '0,16,32,45' for txt2img" }), "num_inference_steps": ("INT", { @@ -181,6 +169,11 @@ class StreamScheduler: }, } + RETURN_TYPES = ("SIGMAS",) + FUNCTION = "update" + CATEGORY = "StreamPack/sampling" + DESCRIPTION = "Implements StreamDiffusion's efficient timestep selection. Use in conjunction with StreamBatchSampler." + def update(self, model, t_index_list="32,45", num_inference_steps=50): # Get model's sampling parameters model_sampling = model.get_model_object("model_sampling") @@ -214,11 +207,6 @@ class StreamScheduler: class StreamFrameBuffer: - - RETURN_TYPES = ("LATENT",) - FUNCTION = "update" - CATEGORY = "StreamPack/sampling" - DESCRIPTION = "Accumulates frames to enable staggered batch denoising like StreamDiffusion. Use in conjunction with StreamBatchSampler" @classmethod def INPUT_TYPES(cls): return { @@ -234,12 +222,20 @@ class StreamFrameBuffer: }, } + RETURN_TYPES = ("LATENT",) + FUNCTION = "update" + CATEGORY = "StreamPack/sampling" + DESCRIPTION = "Accumulates frames to enable staggered batch denoising like StreamDiffusion. Use in conjunction with StreamBatchSampler" + + def __init__(self): self.frame_buffer = None # Tensor of shape [buffer_size, C, H, W] self.buffer_size = None - self.buffer_pos = 0 # Current position in ring buffer + self.buffer_pos = 0 # Current position self.is_initialized = False # Track buffer initialization self.is_txt2img_mode = False + self.last_valid_frame = None # Store the last valid frame as fallback + self.expected_shape = None # Store expected shape for validation def update(self, latent, buffer_size=4): """Add new frame to buffer and return batch when ready""" @@ -248,46 +244,48 @@ class StreamFrameBuffer: # Extract latent tensor from input and remove batch dimension if present x = latent["samples"] - # Check if this is an empty latent (for txt2img) - is_empty_latent = x.numel() == 0 or (x.dim() > 0 and x.shape[0] == 0) + # Check if this is a txt2img (zeros tensor) or img2img mode + # In ComfyUI, EmptyLatentImage returns a zeros tensor with shape [batch_size, 4, h//8, w//8] + # We consider it txt2img mode if the tensor contains all zeros + is_txt2img_mode = torch.sum(torch.abs(x)) < 1e-6 + self.is_txt2img_mode = is_txt2img_mode - if is_empty_latent: - self.is_txt2img_mode = True - - # Create empty latents with correct dimensions for txt2img - # Get dimensions from latent dict - height = latent.get("height", 512) - width = latent.get("width", 512) - - # Calculate latent dimensions (typically 1/8 of image dimensions for SD) - latent_height = height // 8 - latent_width = width // 8 - - # Create zero tensor with correct shape - x = torch.zeros((4, latent_height, latent_width), - device=comfy.model_management.get_torch_device()) - - elif x.dim() == 4: # [B,C,H,W] - self.is_txt2img_mode = False + # If it's a batch with size 1, remove the batch dimension for our buffer + if x.dim() == 4 and x.shape[0] == 1: # [1,C,H,W] x = x.squeeze(0) # Remove batch dimension -> [C,H,W] - # Optimization: Initialize or resize frame_buffer as a tensor - if not self.is_initialized or self.frame_buffer.shape[0] != self.buffer_size or \ - self.frame_buffer.shape[1:] != x.shape: + # Initialize buffer on first run or when buffer size changes + if not self.is_initialized or self.frame_buffer is None or self.frame_buffer.shape[0] != self.buffer_size: + # First initialization - set expected shape and store first frame + if not self.is_initialized: + self.expected_shape = x.shape + self.last_valid_frame = x.clone() + # Pre-allocate buffer with correct shape self.frame_buffer = torch.zeros( - (self.buffer_size, *x.shape), + (self.buffer_size, *self.expected_shape), device=x.device, dtype=x.dtype ) + if self.is_txt2img_mode or not self.is_initialized: + # Use the right-sized frame for initialization + init_frame = x if x.shape == self.expected_shape else self.last_valid_frame # Optimization: Use broadcasting to fill buffer with copies - self.frame_buffer[:] = x.unsqueeze(0) # Broadcast x to [buffer_size, C, H, W] + self.frame_buffer[:] = init_frame.unsqueeze(0) # Broadcast to [buffer_size, C, H, W] self.is_initialized = True self.buffer_pos = 0 else: - # Add new frame to buffer using ring buffer logic + # Check if incoming frame matches expected dimensions + if x.shape != self.expected_shape: + # Size mismatch - use last valid frame instead + x = self.last_valid_frame + else: + # Valid frame - update our reference + self.last_valid_frame = x.clone() + + # Add frame to buffer using ring buffer logic self.frame_buffer[self.buffer_pos] = x # In-place update self.buffer_pos = (self.buffer_pos + 1) % self.buffer_size # Circular increment