stop gap measure for supporting classifier free guidance condition sc…

…aling without kv cache
lucidrains · Dec 17, 2023 · cee61b4 · cee61b4
1 parent 919936d
commit cee61b4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/meshgpt_pytorch/meshgpt_pytorch.py b/meshgpt_pytorch/meshgpt_pytorch.py
@@ -1090,18 +1090,21 @@ def generate(
 
             can_eos = i != 0 and divisible_by(i, self.num_quantizers * 3)  # only allow for eos to be decoded at the end of each face, defined as 3 vertices with D residual VQ codes
 
-            logits, new_cache = self.forward_on_codes(
+            output = self.forward_on_codes(
                 codes,
                 cache = cache,
                 text_embeds = text_embeds,
                 return_loss = False,
-                return_cache = True,
+                return_cache = can_cache,
                 append_eos = False,
                 cond_scale = cond_scale
             )
 
             if can_cache:
+                logits, cache = output
                 cache = new_cache
+            else:
+                logits = output
 
             logits = logits[:, -1]
 

diff --git a/meshgpt_pytorch/version.py b/meshgpt_pytorch/version.py
@@ -1 +1 @@
-__version__ = '0.3.17'
+__version__ = '0.3.18'