properly limit input to 64 tokens

cog.yaml

@@ -1,12 +1,12 @@
 build:
-  cuda: "11.0"
+  cuda: "11.3"
   gpu: true
   python_version: "3.8"
   system_packages:
     - "libgl1-mesa-glx"
     - "libglib2.0-0"
   python_packages:
-    - "min-dalle==0.2.28"
+    - "min-dalle==0.2.29"
   run:
     - pip install torch==1.10.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html
 

min_dalle/min_dalle.py

@@ -172,6 +172,8 @@ class MinDalle:
         assert(log2_mid_count in range(5))
         if is_verbose: print("tokenizing text")
         tokens = self.tokenizer.tokenize(text, is_verbose=is_verbose)
+        if len(tokens) > self.text_token_count: 
+            tokens = tokens[:self.text_token_count]
         if is_verbose: print("text tokens", tokens)
         text_tokens = numpy.ones((2, 64), dtype=numpy.int32)
         text_tokens[0, :2] = [tokens[0], tokens[-1]]

replicate_predictor.py

@@ -11,7 +11,7 @@ class ReplicatePredictor(BasePredictor):
     def predict(
         self,
         text: str = Input(
-            description='Text',
+            description='For long prompts, only the first 64 tokens will be used to generate the image.',
             default='Dali painting of WALL·E'
         ),
         intermediate_outputs: bool = Input(

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.28',
+    version='0.2.29',
     author='Brett Kuprel',
     author_email='brkuprel@gmail.com',
     url='https://github.com/kuprel/min-dalle',