Update pipelines/pipeline_seesr.py

pipelines/pipeline_seesr.py

@@ -95,7 +95,31 @@ EXAMPLE_DOC_STRING = """
         ... ).images[0]
         ```
 """
-
+def kde_grad(x0: torch.Tensor, patch_size = 16, bandwidth = 0.1):
+    # x0: (N, C, H, W) in float32
+    N, C, H, W = x0.shape
+    patches = unfold(
+        x0, kernel_size=patch_size, stride=patch_size//2
+    )                           # (N, C*ps*ps, M)
+    P, M = patches.shape[1], patches.shape[2]
+    p_i = patches.unsqueeze(1)  # (N,1,P,M)
+    p_j = patches.unsqueeze(0)  # (1,N,P,M)
+    diff = p_j - p_i            # (N,N,P,M)
+    # Gaussian weights
+    w = torch.exp((-0.5 / bandwidth**2) *
+                  (diff.square().sum(dim=2)))  # (N,N,M)
+    # mean-shift numerator & normalizer
+    num = (w.unsqueeze(2) * diff).sum(dim=1)          # (N,P,M)
+    denom = w.sum(dim=1, keepdim=True) + 1e-8         # (N,1,M)
+    mshift = num / denom                             # (N,P,M)
+    # fold back
+    grad = fold(
+        mshift / bandwidth**2,
+        output_size=(H, W),
+        kernel_size=patch_size,
+        stride=patch_size//2
+    )                                                # (N, C, H, W)
+    return grad
 
 class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoaderMixin):
     r"""
@@ -807,6 +831,8 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
         ram_encoder_hidden_states=None,
         latent_tiled_size=320,
         latent_tiled_overlap=4,
+        num_particles: Optional[int] = 4,
+        gamma_0: Optional[float] = 0.1,  # base steering strength
         use_KDS=True,
         args=None,
     ):