Update pipelines/pipeline_seesr.py

pipelines/pipeline_seesr.py

@@ -1258,6 +1258,26 @@ 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: