Update pipelines/pipeline_seesr.py

pipelines/pipeline_seesr.py

@@ -1225,38 +1225,39 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
 
                 if use_KDS:
 
-                    # 2) Compute x₀ prediction
+                    # 2) Compute x₀ prediction for all particles
                     beta_t   = 1 - self.scheduler.alphas_cumprod[t]
                     alpha_t  = self.scheduler.alphas_cumprod[t].sqrt()
                     sigma_t  = beta_t.sqrt()
-                    x0_pred  = (latents - sigma_t * noise_pred) / alpha_t
-
-                    # 3) Apply KDE steering
-                    m_shift   = kde_grad(x0_pred, bandwidth=bandwidth)
-                    delta_t   = gamma_0 * (1 - i / (len(timesteps) - 1))
-                    x0_steer  = x0_pred + delta_t * m_shift
-                    # frac    = i / (len(timesteps) - 1)
-                    # delta_t = 0.0 if frac < 0.3 else 0.3
-                    # x0_steer = x0_pred + delta_t * gamma_0 * m_shift
+                    x0_pred  = (latents - sigma_t * noise_pred) / alpha_t  # shape [2N, C, H, W]
                     
+                    # — split into unconditional vs. conditional
+                    x0_uncond, x0_cond = x0_pred.chunk(2, dim=0)           # each [N, C, H, W]
                     
-                    # 4) Recompute “noise” for DDIM step
+                    # 3) Apply KDE steering *only* on the conditional batch
+                    m_shift_cond = kde_grad(x0_cond, bandwidth=bandwidth)  # [N, C, H, W]
+                    delta_t      = gamma_0 * (1 - i / (len(timesteps) - 1))
+                    x0_cond_steer = x0_cond + delta_t * m_shift_cond      # steered conditional
+                    
+                    # 4) Recombine the latents: leave uncond untouched, use steered cond
+                    x0_steer = torch.cat([x0_uncond, x0_cond_steer], dim=0)  # [2N, C, H, W]
+                    
+                    # 5) Recompute “noise” for DDIM step
                     noise_pred_kds = (latents - alpha_t * x0_steer) / sigma_t
                     
-                    # 5) Determine prev alphas
+                    # 6) Determine prev alphas and form next latent per DDIM
                     if i < len(timesteps) - 1:
-                      next_t = timesteps[i + 1]
-                      alpha_prev = self.scheduler.alphas_cumprod[next_t].sqrt()
+                        next_t = timesteps[i + 1]
+                        alpha_prev = self.scheduler.alphas_cumprod[next_t].sqrt()
                     else:
-                      alpha_prev = self.scheduler.final_alpha_cumprod.sqrt()
-                    
+                        alpha_prev = self.scheduler.final_alpha_cumprod.sqrt()
                     sigma_prev = (1 - alpha_prev**2).sqrt()
                     
-                    # 6) Form next latent per DDIM
                     latents = (
-                      alpha_prev * x0_steer
-                      + sigma_prev * noise_pred_kds
+                        alpha_prev * x0_steer
+                        + sigma_prev * noise_pred_kds
                     ).detach().requires_grad_(True)
+                    
                 else:
 
                     # compute the previous noisy sample x_t -> x_t-1