From cfc0a6a36b3657d0197af0faccdd6545157287ba Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Fri, 11 Oct 2024 11:08:34 +0300 Subject: [PATCH] only round the gamma once --- pyramid_dit/pyramid_dit_for_video_gen_pipeline.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/pyramid_dit/pyramid_dit_for_video_gen_pipeline.py b/pyramid_dit/pyramid_dit_for_video_gen_pipeline.py index 123239a..5d88121 100644 --- a/pyramid_dit/pyramid_dit_for_video_gen_pipeline.py +++ b/pyramid_dit/pyramid_dit_for_video_gen_pipeline.py @@ -56,6 +56,8 @@ class PyramidDiTForVideoGeneration: self.sample_ratios = sample_ratios self.corrupt_ratio = corrupt_ratio + + dit_path = os.path.join(model_path, model_variant) @@ -115,6 +117,9 @@ class PyramidDiTForVideoGeneration: shift=timestep_shift, stages=len(self.stages), stage_range=stage_range, gamma=scheduler_gamma, ) + #round the gamma as 1/3 seems to have issues on some systems + self.gamma = round(self.scheduler.config.gamma, 5) + print(f"The start sigmas and end sigmas of each stage is Start: {self.scheduler.start_sigmas}, End: {self.scheduler.end_sigmas}, Ori_start: {self.scheduler.ori_start_sigmas}") self.cfg_rate = 0.1 @@ -191,8 +196,7 @@ class PyramidDiTForVideoGeneration: return latents def sample_block_noise(self, bs, ch, temp, height, width): - gamma = round(self.scheduler.config.gamma, 5) - dist = torch.distributions.multivariate_normal.MultivariateNormal(torch.zeros(4), torch.eye(4) * (1 + gamma) - torch.ones(4, 4) * gamma) + dist = torch.distributions.multivariate_normal.MultivariateNormal(torch.zeros(4), torch.eye(4) * (1 + self.gamma) - torch.ones(4, 4) * self.gamma) block_number = bs * ch * temp * (height // 2) * (width // 2) noise = torch.stack([dist.sample() for _ in range(block_number)]) # [block number, 4] noise = rearrange(noise, '(b c t h w) (p q) -> b c t (h p) (w q)',b=bs,c=ch,t=temp,h=height//2,w=width//2,p=2,q=2) @@ -230,9 +234,8 @@ class PyramidDiTForVideoGeneration: latents = rearrange(latents, '(b t) c h w -> b c t h w', t=temp) # Fix the stage ori_sigma = 1 - self.scheduler.ori_start_sigmas[i_s] # the original coeff of signal - gamma = self.scheduler.config.gamma - alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) - beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) + alpha = 1 / (math.sqrt(1 + (1 / self.gamma)) * (1 - ori_sigma) + ori_sigma) + beta = alpha * (1 - ori_sigma) / math.sqrt(self.gamma) bs, ch, temp, height, width = latents.shape noise = self.sample_block_noise(bs, ch, temp, height, width)