cleaner way to organize code around RiverHasWings application of clas…

…sifier free guidance to transformers

dalle_pytorch/dalle_pytorch.py

@@ -525,18 +525,7 @@ def generate_images(
 
             text, image = out[:, :text_seq_len], out[:, text_seq_len:]
 
-            if cond_scale != 1 and use_cache:
-                # copy the cache state to infer from the same place twice
-                prev_cache = cache.copy()
-
-            logits = self(text, image, cache = cache)
-
-            if cond_scale != 1:
-                # discovery by Katherine Crowson
-                # https://twitter.com/RiversHaveWings/status/1478093658716966912
-                null_cond_logits = self(text, image, null_cond_prob = 1., cache = prev_cache)
-                logits = null_cond_logits + (logits - null_cond_logits) * cond_scale
-
+            logits = self.forward_with_cond_scale(text, image, cond_scale = cond_scale, cache = cache)
             logits = logits[:, -1, :]
 
             filtered_logits = top_k(logits, thres = filter_thres)
@@ -556,6 +545,18 @@ def generate_images(
 
         return images
 
+    def forward_with_cond_scale(self, *args, cond_scale = 1, cache = None, **kwargs):
+        if cond_scale == 1:
+            return self(*args, **kwargs)
+
+        prev_cache = cache.copy() if exists(cache) else None
+        logits = self(*args, cache = cache, **kwargs)
+
+        # discovery by Katherine Crowson
+        # https://twitter.com/RiversHaveWings/status/1478093658716966912
+        null_cond_logits = self(*args, null_cond_prob = 1., cache = prev_cache, **kwargs)
+        return null_cond_logits + (logits - null_cond_logits) * cond_scale
+
     def forward(
         self,
         text,

setup.py

@@ -4,7 +4,7 @@
   name = 'dalle-pytorch',
   packages = find_packages(),
   include_package_data = True,
-  version = '1.5.1',
+  version = '1.5.2',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',