add annealed noising of token critic scores, from token critic paper

ZeroCool940711 · Jan 24, 2023 · 2cc82dc · 2cc82dc
1 parent 5570c74
commit 2cc82dc
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/muse_maskgit_pytorch/muse_maskgit_pytorch.py b/muse_maskgit_pytorch/muse_maskgit_pytorch.py
@@ -486,6 +486,7 @@ def generate(
         force_not_use_token_critic = False,
         timesteps = 18,  # ideal number of steps is 18 in maskgit paper
         cond_scale = 3,
+        critic_noise_scale = 1
     ):
         fmap_size = default(fmap_size, self.vae.get_encoded_fmap_size(self.image_size))
 
@@ -579,6 +580,9 @@ def generate(
                 )
 
                 scores = rearrange(scores, '... 1 -> ...')
+
+                scores = scores + (uniform(scores.shape, device = device) - 0.5) * critic_noise_scale * (steps_until_x0 / timesteps)
+
             else:
                 probs_without_temperature = logits.softmax(dim = -1)
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'muse-maskgit-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.24',
+  version = '0.0.25',
   license='MIT',
   description = 'MUSE - Text-to-Image Generation via Masked Generative Transformers, in Pytorch',
   author = 'Phil Wang',