Update pipelines/pipeline_seesr.py

pipelines/pipeline_seesr.py

@@ -1270,6 +1270,25 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
                         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]
+
             # 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: