ryanontheinside-ComfyUI_RealtimeNodes/stream_sampler.py

import torch
import comfy.model_management
import comfy.samplers

class StreamBatchSampler:
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "num_steps": ("INT", {
                    "default": 4,
                    "min": 1,
                    "max": 10,
                    "step": 1,
                    "tooltip": "Number of denoising steps. Should match the frame buffer size."
                }),
            },
        }
    
    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 = []
        self.x_t_latent_buffer = None
        self.stock_noise = None
        self.is_txt2img_mode = False
        
        # Initialize all buffers
        self.zeros_reference = None
        self.random_noise_buffer = None
        self.sigmas_view_buffer = None
        self.expanded_stock_noise = None
        self.working_buffer = None
        self.output_buffer = None
    
    def sample(self, model, noise, sigmas, extra_args=None, callback=None, disable=None):
        """Sample with staggered batch denoising steps - Optimized version"""
        extra_args = {} if extra_args is None else extra_args
        
        # Get number of frames in batch and available sigmas
        batch_size = noise.shape[0]
        num_sigmas = len(sigmas) - 1  # Subtract 1 because last sigma is the target (0.0)
        
        
        # Reuse zeros buffer for txt2img detection
        if self.zeros_reference is None:
            self.zeros_reference = torch.zeros(1, device=noise.device, dtype=noise.dtype)
        
        # Check if noise tensor is all zeros
        self.is_txt2img_mode = torch.abs(noise).sum() < 1e-5
        
        if self.is_txt2img_mode:
            # 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)
            
            # Generate random noise in-place
            self.random_noise_buffer.normal_()
            x = self.random_noise_buffer  # Use pre-allocated buffer
        else:
            # If not txt2img, we'll still need to add noise later
            x = noise
        
        # Verify batch size matches number of timesteps
        if batch_size != num_sigmas:
            raise ValueError(f"Batch size ({batch_size}) must match number of timesteps ({num_sigmas})")
        
        # Pre-compute alpha and beta terms
        alpha_prod_t = (sigmas[:-1] / sigmas[0]).view(-1, 1, 1, 1)  # [B,1,1,1]
        beta_prod_t = (1 - alpha_prod_t)
        
        # 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
        
        # 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, 
                                               dtype=sigmas.dtype)
        # In-place copy of sigmas view
        self.sigmas_view_buffer.copy_(sigmas[:-1].view(-1, 1, 1, 1))
        
        # 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)
        
        # Apply noise with pre-allocated buffers - no new memory allocation
        if not self.is_txt2img_mode:  # Already handled txt2img case above
            # If we need a working buffer separate from noise input:
            if id(x) == id(noise):  # They're the same object, need a separate buffer
                if self.working_buffer is None or self.working_buffer.shape != noise.shape:
                    self.working_buffer = torch.empty_like(noise)
                x = self.working_buffer
                # Add noise to input
                torch.add(noise, self.expanded_stock_noise * self.sigmas_view_buffer, out=x)
        
        # 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:
            # 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])
            self.x_t_latent_buffer.copy_(x[0])
        
        # Use buffer for first frame to maintain temporal consistency
        if num_sigmas > 1:
            x[0].copy_(self.x_t_latent_buffer)
        
        with torch.no_grad():
            # Process all frames in parallel
            sigma_batch = sigmas[:-1]
            
            denoised_batch = model(x, sigma_batch, **extra_args)
            
            # Update buffer with intermediate results
            if num_sigmas > 1:
                # Store result from first frame as buffer for next iteration
                self.x_t_latent_buffer.copy_(denoised_batch[0])  # In-place update
                
                # 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,
                                                  dtype=denoised_batch.dtype)
                # Copy the result directly to pre-allocated buffer
                self.output_buffer[0].copy_(denoised_batch[-1])
                x_0_pred_out = self.output_buffer
            else:
                x_0_pred_out = denoised_batch
                self.x_t_latent_buffer = None
            
            # Call callback if provided
            if callback is not None:
                callback({'x': x_0_pred_out, 'i': 0, 'sigma': sigmas[0], 'sigma_hat': sigmas[0], 'denoised': denoised_batch[-1:]})
        
        return x_0_pred_out
    
    def update(self, num_steps=4):
        """Create sampler with specified settings"""
        self.num_steps = num_steps
        sampler = comfy.samplers.KSAMPLER(self.sample)
        return (sampler,)


class StreamScheduler:
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "t_index_list": ("STRING", {
                    "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", {
                    "default": 50,
                    "min": 1,
                    "max": 1000,
                    "step": 1,
                    "tooltip": "Total number of timesteps in schedule. StreamDiffusion uses 50 by default. Only timesteps specified in t_index_list are actually used."
                }),
            },
        }

    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")
        
        # Parse timestep list
        try:
            t_index_list = [int(t.strip()) for t in t_index_list.split(",")]
        except ValueError as e:
            
            t_index_list = [32, 45]
            
        # Create full schedule using normal scheduler
        full_sigmas = comfy.samplers.normal_scheduler(model_sampling, num_inference_steps)
        
        # Select only the sigmas at our desired indices, but in reverse order
        # This ensures we go from high noise to low noise
        selected_sigmas = []
        for t in sorted(t_index_list, reverse=True):  # Sort in reverse to go from high noise to low
            if t < 0 or t >= num_inference_steps:
                
                continue
            selected_sigmas.append(float(full_sigmas[t]))
            
        # Add final sigma
        selected_sigmas.append(0.0)
        
        # Convert to tensor and move to appropriate device
        selected_sigmas = torch.FloatTensor(selected_sigmas).to(comfy.model_management.get_torch_device())
        return (selected_sigmas,)


class StreamFrameBuffer:
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "latent": ("LATENT",),
                "buffer_size": ("INT", {
                    "default": 4,
                    "min": 1,
                    "max": 10,
                    "step": 1,
                    "tooltip": "Number of frames to buffer before starting batch processing. Should match number of denoising steps."
                }),
            },
        }
    
    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
        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"""
        self.buffer_size = buffer_size
        
        # Extract latent tensor from input and remove batch dimension if present
        x = latent["samples"]
        
        # 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 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]
        
        # 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, *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[:] = init_frame.unsqueeze(0)  # Broadcast to [buffer_size, C, H, W]

            self.is_initialized = True
            self.buffer_pos = 0
        else:
            # 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
        
        # Optimization: frame_buffer is already a tensor batch, no need to stack
        batch = self.frame_buffer

        
        # Return as latent dict with preserved dimensions
        result = {"samples": batch}
        
        # Preserve height and width if present in input
        if "height" in latent:
            result["height"] = latent["height"]
        if "width" in latent:
            result["width"] = latent["width"]
            
        return (result,)