nodes/ryanontheinside-ComfyUI_RealtimeNodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update

Browse Source
This commit is contained in:
RyanOnTheInside
2025-03-14 00:38:15 +00:00
parent d9e05164fa
commit fb8c84f844
5 changed files with 504 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.ipynb_checkpoints/__init__-checkpoint.py
+69
View File
@@ -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']
.ipynb_checkpoints/particle_nodes-checkpoint.py
+350
View File
@@ -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)
__init__.py
+2
View File
@@ -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 = {}
dimensions.py
+24
View File
@@ -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"
}
stream_sampler.py
+59 -61
View File
@@ -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