fix a bug with bert encoding

lucidrains · Jul 15, 2022 · e4d9e9d · e4d9e9d
1 parent f692a12
commit e4d9e9d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'video-diffusion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.5.2',
+  version = '0.5.3',
   license='MIT',
   description = 'Video Diffusion - Pytorch',
   long_description_content_type = 'text/markdown',

diff --git a/video_diffusion_pytorch/video_diffusion_pytorch.py b/video_diffusion_pytorch/video_diffusion_pytorch.py
@@ -604,7 +604,12 @@ def __init__(
         register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
         register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod))
 
+        # text conditioning parameters
+
+        self.text_use_bert_cls = text_use_bert_cls
+
         # dynamic thresholding when sampling
+
         self.use_dynamic_thres = use_dynamic_thres
         self.dynamic_thres_percentile = dynamic_thres_percentile
 
@@ -715,7 +720,7 @@ def p_losses(self, x_start, t, cond = None, noise = None, **kwargs):
         x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
 
         if is_list_str(cond):
-            cond = bert_embed(tokenize(cond), return_cls_repr = text_use_bert_cls)
+            cond = bert_embed(tokenize(cond), return_cls_repr = self.text_use_bert_cls)
             cond = cond.to(device)
 
         x_recon = self.denoise_fn(x_noisy, t, cond = cond, **kwargs)