Merge pull request #18 from argmaxinc/long_context_fix

Truncating long prompts

Author: atiorh

python/src/diffusionkit/mlx/__init__.py

@@ -40,6 +40,11 @@
     "FLUX.1-schnell": "argmaxinc/mlx-FLUX.1-schnell",
 }
 
+T5_MAX_LENGTH = {
+    "stable-diffusion-3-medium": 512,
+    "FLUX.1-schnell": 256,
+}
+
 
 class DiffusionKitInferenceContext(AppleSiliconContextMixin, InferenceContextSpec):
     def code_spec(self):
@@ -138,7 +143,9 @@ def set_up_t5(self):
                 low_memory_mode=self.low_memory_mode,
             )
         if not hasattr(self, "t5_tokenizer") or self.t5_tokenizer is None:
-            self.t5_tokenizer = load_t5_tokenizer()
+            self.t5_tokenizer = load_t5_tokenizer(
+                max_context_length=T5_MAX_LENGTH[self.model_version]
+            )
         self.use_t5 = True
 
     def unload_t5(self):

python/src/diffusionkit/mlx/model_io.py

@@ -871,6 +871,6 @@ def load_tokenizer(
     return Tokenizer(bpe_ranks, vocab, pad_with_eos)
 
 
-def load_t5_tokenizer():
+def load_t5_tokenizer(max_context_length: int = 256):
     config = T5Config.from_pretrained("google/t5-v1_1-xxl")
-    return T5Tokenizer(config)
+    return T5Tokenizer(config, max_context_length)

python/src/diffusionkit/mlx/tokenizer.py

@@ -5,8 +5,11 @@
 import mlx.core as mx
 import numpy as np
 import regex
+from argmaxtools.utils import get_logger
 from transformers import AutoTokenizer, T5Config
 
+logger = get_logger(__name__)
+
 
 class Tokenizer:
     """A simple port of CLIPTokenizer from https://github.com/huggingface/transformers/ ."""
@@ -101,6 +104,14 @@ def tokenize(self, text, prepend_bos=True, append_eos=True):
 
         # Map to token ids and return
         tokens = [self.vocab[t] for t in bpe_tokens]
+
+        # Truncate
+        max_length = self.max_length - int(prepend_bos) - int(append_eos)
+        if len(tokens) > max_length:
+            tokens = tokens[:max_length]
+            logger.warning(
+                f"Length of tokens exceeds {self.max_length}. Truncating to {self.max_length}."
+            )
         if prepend_bos:
             tokens = [self.bos_token] + tokens
         if append_eos:
@@ -110,16 +121,16 @@ def tokenize(self, text, prepend_bos=True, append_eos=True):
 
 
 class T5Tokenizer:
-    def __init__(self, config: T5Config):
+    def __init__(self, config: T5Config, max_context_length: int):
+        self.max_length = max_context_length
         self._decoder_start_id = config.decoder_start_token_id
         self._tokenizer = AutoTokenizer.from_pretrained(
             "google/t5-v1_1-xxl",
             legacy=False,
-            model_max_length=getattr(config, "n_positions", 512),
+            model_max_length=self.max_length,
         )
 
         self.pad_to_max_length = True
-        self.max_length = 77
         self.pad_with_eos = False
 
     @property
@@ -136,6 +147,8 @@ def encode(self, s: str) -> mx.array:
                 s,
                 return_tensors="np",
                 return_attention_mask=False,
+                max_length=self.max_length,
+                truncation=True,
             )["input_ids"]
         )