Update pipelines/pipeline_seesr.py

pipelines/pipeline_seesr.py

@@ -99,9 +99,7 @@ EXAMPLE_DOC_STRING = """
 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)
+    patches = unfold(x0, kernel_size=patch_size, stride=patch_size) # (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)
@@ -113,13 +111,15 @@ def kde_grad(x0: torch.Tensor, patch_size = 16, bandwidth = 0.1):
     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)
+         mshift / bandwidth**2,
+         output_size=(H, W),
+         kernel_size=patch_size,
+         stride=patch_size
+    )                                           # (N, C, H, W)
+
     return grad
 
 class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoaderMixin):