add RCFG

src/batched_denoising/batched_denoising.py

@@ -1,6 +1,7 @@
 import torch
 import comfy.model_management
 import comfy.samplers
+import math
 
 class StreamBatchSampler:
     
@@ -8,27 +9,23 @@ class StreamBatchSampler:
     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."
-                }),
+                "cfg_type": (["none", "full", "self", "initialize"], {"default": "self", "tooltip": "'self' (RCFG) is fastest, 'full' is standard SD but slower, 'initialize' is a middle ground, 'none' disables guidance"}),
+                "delta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1, "tooltip": "Only used with 'self' CFG. Controls strength of self-guidance. Higher = stronger guidance but more artifacts. Default 1.0", "round": 0.01}),
             },
         }
     
     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."
+    DESCRIPTION = "Implements batched denoising with different CFG options (including RCFG) for faster inference and temporal consistency."
     
     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
+        self.delta = 1.0 
+        self.cfg_type = "self"
         
         # Initialize all buffers
         self.zeros_reference = None
@@ -37,118 +34,331 @@ class StreamBatchSampler:
         self.expanded_stock_noise = None
         self.working_buffer = None
         self.output_buffer = None
-    
+        
+        # Add RCFG state from reference
+        self.cached_sigmas = None
+        self.alpha_prod_t = None
+        self.beta_prod_t = None
+        self.alpha_prod_t_sqrt = None
+        self.beta_prod_t_sqrt = None
+        self.init_noise = None # Will be initialized in sample based on input shape
+        self.step_counter = 0
+
+    def compute_alpha_beta(self, sigmas):
+        """Pre-compute alpha and beta terms for the given sigmas (LCM-like using cosine schedule)"""
+        self.cached_sigmas = sigmas.clone()
+        print(f"\n[Alpha/Beta] Input sigmas: {sigmas}")
+
+        num_train_timesteps = 1000  # Standard for many diffusion models
+        timesteps = torch.linspace(
+            1, 0, num_train_timesteps, dtype=torch.float32, device=sigmas.device
+        )
+        alphas_cumprod = torch.cos((timesteps + 0.008) / 1.008 * math.pi / 2).pow(2)
+        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+
+        # Map sigmas to timesteps
+        scaled_sigmas = (
+            sigmas / sigmas.max() if sigmas.max() > 0 else torch.zeros_like(sigmas)
+        )
+        timesteps_for_sigma = (1 - scaled_sigmas) * (num_train_timesteps - 1)
+        indices = (
+            torch.round(timesteps_for_sigma).long().clamp(0, num_train_timesteps - 1)
+        )
+
+        alpha_prod_t = alphas_cumprod[indices].view(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        self.alpha_prod_t = alpha_prod_t
+        self.beta_prod_t = beta_prod_t
+        self.alpha_prod_t_sqrt = torch.sqrt(alpha_prod_t)
+        self.beta_prod_t_sqrt = torch.sqrt(beta_prod_t)
+
+        print(f"[Alpha/Beta] Final stats:")
+        print(
+            f"  alpha_prod_t: min={self.alpha_prod_t.min():.4f}, max={self.alpha_prod_t.max():.4f}, mean={self.alpha_prod_t.mean():.4f}"
+        )
+        print(
+            f"  beta_prod_t: min={self.beta_prod_t.min():.4f}, max={self.beta_prod_t.max():.4f}, mean={self.beta_prod_t.mean():.4f}"
+        )
+        print(
+            f"  alpha_sqrt: min={self.alpha_prod_t_sqrt.min():.4f}, max={self.alpha_prod_t_sqrt.max():.4f}, mean={self.alpha_prod_t_sqrt.mean():.4f}"
+        )
+        print(
+            f"  beta_sqrt: min={self.beta_prod_t_sqrt.min():.4f}, max={self.beta_prod_t_sqrt.max():.4f}, mean={self.beta_prod_t_sqrt.mean():.4f}\n"
+        )
+
+    def scheduler_step_batch(self, denoised_batch, scaled_noise):
+            """Compute scheduler step for the entire batch (LCM-like)"""
+            print("\n[Scheduler] Starting scheduler step (LCM-like):")
+            print(f"  Denoised shape: {denoised_batch.shape}")
+            print(f"  Scaled noise shape: {scaled_noise.shape}")
+            print(
+                f"  Denoised stats: min={denoised_batch.min():.4f}, max={denoised_batch.max():.4f}, mean={denoised_batch.mean():.4f}"
+            )
+            print(
+                f"  Scaled noise stats: min={scaled_noise.min():.4f}, max={scaled_noise.max():.4f}, mean={scaled_noise.mean():.4f}"
+            )
+
+            # Assuming self.alpha_prod_t and self.beta_prod_t are for the current timestep
+            alpha_t_sqrt = self.alpha_prod_t_sqrt[: denoised_batch.shape[0]]
+            beta_t_sqrt = self.beta_prod_t_sqrt[: denoised_batch.shape[0]]
+
+            # Calculate the predicted original sample (x_0) from the noisy sample and model output
+            # This is similar to the inverse of the forward diffusion step
+            pred_original_sample = (
+                denoised_batch - beta_t_sqrt.view(-1, 1, 1, 1) * scaled_noise
+            ) / (alpha_t_sqrt.view(-1, 1, 1, 1) + 1e-7)
+            print(
+                f"  Predicted original sample stats: min={pred_original_sample.min():.4f}, max={pred_original_sample.max():.4f}, mean={pred_original_sample.mean():.4f}"
+            )
+
+            # For the reverse step, we need the alpha and beta for the previous timestep
+            # Since we are operating on a batch of timesteps, we'll use the values corresponding to the first timestep in the batch for updating stock_noise
+            # In StreamDiffusion paper, the update happens based on the result of the last denoising step
+
+            # Get alpha and beta for the next step: (which would be the previous in the denoising sequence)
+            alpha_next_sqrt = torch.cat(
+                [
+                    self.alpha_prod_t_sqrt[1 : denoised_batch.shape[0]],
+                    torch.ones_like(self.alpha_prod_t_sqrt[0:1]),
+                ],
+                dim=0,
+            )
+
+            # Calculate the coefficient for the predicted original sample in the next step
+            sigma_next = (1 - alpha_next_sqrt.pow(2)) / (alpha_next_sqrt + 1e-7)
+            sigma_t = (1 - alpha_t_sqrt.pow(2)) / (alpha_t_sqrt + 1e-7)
+
+            # Guidance scale (set to 1 for unconditional prediction)
+            guidance_scale = 1
+
+            # Calculate delta_x (the change in the noisy sample)
+            delta_x = (sigma_next - sigma_t) * guidance_scale * pred_original_sample
+
+            print(
+                f"  Final delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}\n"
+            )
+
+            return delta_x
+
     def sample(self, model, noise, sigmas, extra_args=None, callback=None, disable=None):
-        """Sample with staggered batch denoising steps - Optimized version"""
+        """Sample with staggered batch denoising steps - Optimized version with multiple CFG types"""
         extra_args = {} if extra_args is None else extra_args
+        cond_scale = extra_args.get("model_options", {}).get("cfg_scale", 1.0) # Get cond_scale from model options
+        cfg_type = self.cfg_type
         
         # 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)
         
+        # Precompute alpha/beta if needed
+        self.compute_alpha_beta(sigmas) # Sigmas include the final 0.0, compute uses all
+        
         # Reuse zeros buffer for txt2img detection
-        if self.zeros_reference is None:
-            self.zeros_reference = torch.zeros(1, device=noise.device, dtype=noise.dtype)
+        if self.zeros_reference is None or self.zeros_reference.device != noise.device or self.zeros_reference.dtype != noise.dtype:
+             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
         
+        # Initialize init_noise sequence like reference if needed (once or if shape changes)
+        noise_shape_base = noise.shape[1:] # C, H, W
+        if self.init_noise is None or self.init_noise.shape[0] != num_sigmas or self.init_noise.shape[1:] != noise_shape_base:
+            self.init_noise = torch.randn(
+                (num_sigmas, *noise_shape_base),
+                device=noise.device,
+                dtype=noise.dtype,
+            )
+            # Initialize stock_noise to zeros when init_noise is created (like reference)
+            self.stock_noise = torch.zeros_like(self.init_noise[0]) 
+            self.step_counter = 0 # Reset step counter when noise shape changes
+
+        # Shift noise sequence
+        if self.step_counter > 0:
+            self.init_noise = torch.cat([self.init_noise[1:], self.init_noise[0:1]], dim=0)
+            # Update stock noise based on the *first* element of the *shifted* init_noise
+            self.stock_noise = self.init_noise[0].clone() 
+
+
         if self.is_txt2img_mode:
-            # If txt2img mode, reuse the noise tensor directly
+            # If txt2img mode, use the init_noise sequence directly as the base noise for the batch
+            # This ensures consistent noise across steps/frames for txt2img
             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_()
+            # Copy from the init_noise sequence for the current batch size
+            self.random_noise_buffer.copy_(self.init_noise[:batch_size])
             x = self.random_noise_buffer  # Use pre-allocated buffer
         else:
-            # If not txt2img, we'll still need to add noise later
-            x = noise
+             # If not txt2img, we'll still need to add noise later based on sigmas
+            x = noise # Start with the input noise (img2img)
         
         # 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
+            raise ValueError(f"Batch size ({batch_size}) must match number of sigmas ({num_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), 
+        # Ensure buffer matches the number of steps *in this batch* (num_sigmas)
+        if self.sigmas_view_buffer is None or self.sigmas_view_buffer.shape[0] != num_sigmas:
+            self.sigmas_view_buffer = torch.empty((num_sigmas, 1, 1, 1), 
                                                device=sigmas.device, 
                                                dtype=sigmas.dtype)
-        # In-place copy of sigmas view
+        # In-place copy of sigmas view (use sigmas[:-1] as these are the steps)
         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 not self.is_txt2img_mode: 
+            # For img2img, add noise based on sigmas and init_noise sequence
+            # Expand the relevant part of init_noise for the batch
+            current_batch_init_noise = self.init_noise[:batch_size].to(x.device) # Ensure device match
+            
+            # If x is the same object as noise, use a working buffer
+            if id(x) == id(noise):  
                 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)
-        
+                torch.add(noise, current_batch_init_noise * self.sigmas_view_buffer, out=x)
+            else: # Otherwise, add noise directly to x
+                torch.add(noise, current_batch_init_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
+        # Only buffer if more than one step
         if num_sigmas > 1:
+            if self.x_t_latent_buffer is None or self.x_t_latent_buffer.shape != x[0].shape or self.is_txt2img_mode:
+                # 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
             x[0].copy_(self.x_t_latent_buffer)
+        else: # Single step, no buffering needed
+             self.x_t_latent_buffer = None 
         
         with torch.no_grad():
             # Process all frames in parallel
-            sigma_batch = sigmas[:-1]
+            sigma_batch = sigmas[:-1] # Use sigmas for the steps, excluding the final 0.0
             
-            denoised_batch = model(x, sigma_batch, **extra_args)
+            # --- CFG Handling --- 
+            denoised_batch = None
+            stock_noise_update_value = None # Track value to update stock noise with (used in self-cfg)
             
+            # 1. No CFG
+            if cfg_type == "none" or cond_scale <= 1.0:
+                denoised_batch = model(x, sigma_batch, **extra_args)
+                stock_noise_update_value = denoised_batch[0].clone().detach()
+
+            # 2. Full CFG
+            elif cfg_type == "full":
+                x_double = torch.cat([x, x], dim=0)
+                sigma_double = torch.cat([sigma_batch, sigma_batch], dim=0)
+                model_output = model(x_double, sigma_double, **extra_args)
+                noise_pred_uncond, noise_pred_text = model_output.chunk(2)
+                denoised_batch = noise_pred_uncond + cond_scale * (noise_pred_text - noise_pred_uncond)
+                # Update stock noise with uncond prediction (using first frame)
+                stock_noise_update_value = noise_pred_uncond[0].clone().detach()
+
+            # 3. Initialize CFG
+            elif cfg_type == "initialize":
+                x_plus_uc = torch.cat([x[0:1], x], dim=0) # Add first frame again for uncond
+                sigma_plus_uc = torch.cat([sigma_batch[0:1], sigma_batch], dim=0)
+                model_output = model(x_plus_uc, sigma_plus_uc, **extra_args)
+                noise_pred_uncond = model_output[0:1] # Uncond is the first item
+                noise_pred_cond = model_output[1:]     # Cond are the rest
+                # Calculate final denoised batch using CFG scale
+                denoised_batch = noise_pred_uncond + cond_scale * (noise_pred_cond - noise_pred_uncond) 
+                # Update stock noise with uncond prediction 
+                stock_noise_update_value = noise_pred_uncond[0].clone().detach()
+            
+            # 4. Self CFG (RCFG)
+            elif cfg_type == "self":
+                eps_c = model(x, sigma_batch, **extra_args) # Conditional prediction
+                
+                # Initialize stock noise if needed (first run after init_noise creation)
+                if self.stock_noise is None: # Should have been initialized with init_noise
+                     # Fallback: initialize with first frame prediction if somehow missed
+                     self.stock_noise = eps_c[0].clone().detach()
+                
+                # RCFG equation
+                # Use current stock_noise (which should be based on init_noise[0])
+                noise_pred_uncond = self.stock_noise.to(eps_c.device) * self.delta
+                denoised_batch = noise_pred_uncond + cond_scale * (eps_c - noise_pred_uncond)
+                
+                # Defer stock noise update until after scheduler step below
+                # Store the conditional prediction for the scheduler step
+                stock_noise_update_value = eps_c.clone().detach() # Need the whole batch for scheduler
+
+
+            # --- CFG Handling End --- 
+
+            # --- RCFG Stock Noise Update ---
+            if cfg_type == "self" and stock_noise_update_value is not None:
+                 # Get the noise sequence for the current batch
+                 current_batch_init_noise = self.init_noise[:batch_size].to(denoised_batch.device)
+                 
+             
+                 beta_sqrt_batch = self.beta_prod_t_sqrt[:batch_size].to(denoised_batch.device)
+                 scaled_noise = beta_sqrt_batch * current_batch_init_noise # Use the init_noise for this step
+
+                 delta_x = self.scheduler_step_batch(denoised_batch, scaled_noise)
+                 num_steps_in_batch = denoised_batch.shape[0]
+                 if self.alpha_prod_t_sqrt.shape[0] >= num_steps_in_batch:
+                     one_tensor = torch.ones_like(self.alpha_prod_t_sqrt[0:1])
+                     zero_tensor = torch.zeros_like(self.beta_prod_t_sqrt[0:1])
+                     alpha_next_sqrt = torch.cat([self.alpha_prod_t_sqrt[1:num_steps_in_batch], one_tensor], dim=0).to(delta_x.device)
+                     beta_next_sqrt = torch.cat([self.beta_prod_t_sqrt[1:num_steps_in_batch], zero_tensor], dim=0).to(delta_x.device)
+                 else:
+                     raise ValueError("Precomputed alpha/beta shapes mismatch for final scaling.")
+                 
+                 # Apply the scaling
+                 scaled_delta_x = (alpha_next_sqrt / (beta_next_sqrt + 1e-7)) * delta_x
+                 
+                 # And the *scaled* delta_x calculated for the *first frame* of the batch
+                 self.stock_noise += scaled_delta_x[0].to(self.stock_noise.device) # Add scaled delta for first frame
+            elif stock_noise_update_value is not None:
+                 # For non-RCFG modes, update stock noise simply with the first frame's value
+                 if self.stock_noise is None or self.stock_noise.shape != stock_noise_update_value.shape:
+                     self.stock_noise = stock_noise_update_value.clone().detach()
+                 else:
+                     self.stock_noise.copy_(stock_noise_update_value.detach())
+
+
             # 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, 
+                output_shape = (1, *denoised_batch[-1].shape)
+                if self.output_buffer is None or self.output_buffer.shape != output_shape:
+                    self.output_buffer = torch.empty(output_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
+            else: # Single step
+                x_0_pred_out = denoised_batch.unsqueeze(0) # Ensure batch dim [1,C,H,W]
+                self.x_t_latent_buffer = None # No buffer needed
             
             # 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:]})
+                callback({'x': x_0_pred_out, 'i': self.step_counter, 'sigma': sigmas[0], 'sigma_hat': sigmas[0], 'denoised': denoised_batch[-1:]})
+        
+        # Increment step counter at the end of the sample method
+        self.step_counter += 1
         
         return x_0_pred_out
     
-    def update(self, num_steps=4):
+    def update(self, cfg_type="self", delta=1.0):
         """Create sampler with specified settings"""
-        self.num_steps = num_steps
+        self.cfg_type = cfg_type
+        self.delta = delta 
         sampler = comfy.samplers.KSAMPLER(self.sample)
         return (sampler,)
 
 
 class StreamScheduler:
-    
     @classmethod
     def INPUT_TYPES(cls):
         return {