From 1702d3c43928822a4ed6214c649174c374103a1e Mon Sep 17 00:00:00 2001
From: Brett Kuprel <brkuprel@gmail.com>
Date: Mon, 4 Jul 2022 17:27:23 -0400
Subject: [PATCH] log2_mid_count

---
 image_from_text.py     |  2 +-
 min_dalle/min_dalle.py | 10 +++++-----
 setup.py               |  2 +-
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/image_from_text.py b/image_from_text.py
index 582f24e..666aa7b 100644
--- a/image_from_text.py
+++ b/image_from_text.py
@@ -55,7 +55,7 @@ def generate_image(
         image_tokens = model.generate_image_tokens(
             text, 
             seed, 
-            grid_size ** 2, 
+            grid_size ** 2,
             row_count,
             is_verbose=True
         )
diff --git a/min_dalle/min_dalle.py b/min_dalle/min_dalle.py
index ab76f2d..a202501 100644
--- a/min_dalle/min_dalle.py
+++ b/min_dalle/min_dalle.py
@@ -165,7 +165,7 @@ class MinDalle:
         seed: int,
         grid_size: int,
         row_count: int,
-        mid_count: int = None,
+        log2_mid_count: int = 0,
         handle_intermediate_image: Callable[[int, Image.Image], None] = None,
         is_verbose: bool = False
     ) -> LongTensor:
@@ -206,8 +206,8 @@ class MinDalle:
                 attention_state,
                 image_tokens
             )
-            if mid_count is not None:
-                if ((row_index + 1) * mid_count) % row_count == 0:
+            if handle_intermediate_image is not None:
+                if ((row_index + 1) * (2 ** log2_mid_count)) % row_count == 0:
                     tokens = image_tokens[:, 1:]
                     image = self.image_from_tokens(grid_size, tokens, is_verbose)
                     handle_intermediate_image(row_index, image)
@@ -220,7 +220,7 @@ class MinDalle:
         text: str,
         seed: int = -1,
         grid_size: int = 1,
-        mid_count: int = None,
+        log2_mid_count: int = None,
         handle_intermediate_image: Callable[[Image.Image], None] = None,
         is_verbose: bool = False
     ) -> Image.Image:
@@ -229,7 +229,7 @@ class MinDalle:
             seed, 
             grid_size, 
             row_count = 16,
-            mid_count = mid_count,
+            log2_mid_count = log2_mid_count,
             handle_intermediate_image = handle_intermediate_image,
             is_verbose = is_verbose
         )
diff --git a/setup.py b/setup.py
index 7757bc8..2fc9f56 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ setuptools.setup(
     name='min-dalle',
     description = 'min(DALL·E)',
     long_description=(Path(__file__).parent / "README.rst").read_text(),
-    version='0.2.21',
+    version='0.2.22',
     author='Brett Kuprel',
     author_email='brkuprel@gmail.com',
     url='https://github.com/kuprel/min-dalle',