Update pipelines/pipeline_seesr.py

pipelines/pipeline_seesr.py

@@ -1258,39 +1258,37 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
                         alpha_prev * x0_steer
                         + sigma_prev * noise_pred_kds
                     ).detach().requires_grad_(True)
-
-                    uncond_latents, cond_latents = latents.chunk(2, dim=0)  # each is [N, C, H, W]
-
-                    # 1) Compute ensemble mean of the conditional latents
-                    mean_cond = cond_latents.mean(dim=0, keepdim=True)      # shape [1, C, H, W]
-                
-                    # 2) Compute squared distances to the mean for each particle
-                    #    Flatten each latent to [N, C*H*W], then sum-of-squares
-                    dists = ((cond_latents - mean_cond).view(cond_latents.size(0), -1) ** 2).sum(dim=1)  # [N]
-                
-                    # 3) Find the index of the particle closest to the mean
-                    best_idx = dists.argmin().item()
-                
-                    # 4) Select that one latent
-                    best_latent = cond_latents[best_idx : best_idx + 1]      # shape [1, C, H, W]
-                
-                    # (Optional) If you need to keep classifier-free guidance structure,
-                    # you can reconstruct a 2-sample batch with its uncond pair:
-                    best_uncond = uncond_latents[best_idx : best_idx + 1]
-                    latents = torch.cat([best_uncond, best_latent], dim=0)   # shape [2, C, H, W]
-                    
                 else:
 
                     # compute the previous noisy sample x_t -> x_t-1
                     latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
 
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-                    if callback is not None and i % callback_steps == 0:
-                        callback(i, t, latents)
-
         with torch.no_grad():
+            if use_KDS:
+                # Final-latent selection (once!)
+                # latents shape: [2*N, C, H, W]
+                uncond_latents, cond_latents = latents.chunk(2, dim=0)  # each [N, C, H, W]
+                # 1) ensemble mean
+                mean_cond = cond_latents.mean(dim=0, keepdim=True)      # [1, C, H, W]
+                # 2) distances
+                dists = ((cond_latents - mean_cond)
+                          .view(cond_latents.size(0), -1)
+                          .pow(2)
+                          .sum(dim=1))                             # [N]
+                # 3) best index
+                best_idx = dists.argmin().item()
+                # 4) select that latent (and its uncond pair)
+                best_uncond = uncond_latents[best_idx:best_idx+1]
+                best_cond   = cond_latents  [best_idx:best_idx+1]
+                latents     = torch.cat([best_uncond, best_cond], dim=0)  # [2, C, H, W]
+
+            # call the callback, if provided
+            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                progress_bar.update()
+                if callback is not None and i % callback_steps == 0:
+                    callback(i, t, latents)
+
+            
             # If we do sequential model offloading, let's offload unet and controlnet
             # manually for max memory savings
             if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: