From 5e2e11dd0a717ce04758d85ebe7cda0385781a0d Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 6 Oct 2024 12:58:07 +0300
Subject: [PATCH] Update pipeline_cogvideox_inpaint.py

---
 cogvideox_fun/pipeline_cogvideox_inpaint.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/cogvideox_fun/pipeline_cogvideox_inpaint.py b/cogvideox_fun/pipeline_cogvideox_inpaint.py
index e9aff57..4c9d505 100644
--- a/cogvideox_fun/pipeline_cogvideox_inpaint.py
+++ b/cogvideox_fun/pipeline_cogvideox_inpaint.py
@@ -432,7 +432,17 @@ class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline):
         if accepts_generator:
             extra_step_kwargs["generator"] = generator
         return extra_step_kwargs
+    
+    def _gaussian_weights(self, t_tile_length, t_batch_size):
+        from numpy import pi, exp, sqrt
 
+        var = 0.01
+        midpoint = (t_tile_length - 1) / 2  # -1 because index goes from 0 to latent_width - 1
+        t_probs = [exp(-(t-midpoint)*(t-midpoint)/(t_tile_length*t_tile_length)/(2*var)) / sqrt(2*pi*var) for t in range(t_tile_length)]
+        weights = torch.tensor(t_probs)
+        weights = weights.unsqueeze(0).unsqueeze(2).unsqueeze(3).unsqueeze(4).repeat(1, t_batch_size,1, 1, 1)
+        return weights
+    
     # Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs
     def check_inputs(
         self,
@@ -869,7 +879,6 @@ class CogVideoX_Fun_Pipeline_Inpaint(VideoSysPipeline):
             use_context_schedule = True
             from .context import get_context_scheduler
             context = get_context_scheduler(context_schedule)
-
         else:
             use_temporal_tiling = False
             use_context_schedule = False