From 04733873cf1da8060fce01316ec45daedbd4a918 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Mon, 19 Oct 2026 10:15:45 -0600
Subject: [PATCH] Improve performance of encoding Moss prompts with a batch
 size > 1

---
 .../diffusion_models/moss_image/moss_image.py | 55 ++++++++++++++-----
 1 file changed, 41 insertions(+), 14 deletions(-)

diff --git a/extensions_built_in/diffusion_models/moss_image/moss_image.py b/extensions_built_in/diffusion_models/moss_image/moss_image.py
index 112b975..c5b67d7 100644
--- a/extensions_built_in/diffusion_models/moss_image/moss_image.py
+++ b/extensions_built_in/diffusion_models/moss_image/moss_image.py
@@ -117,6 +117,30 @@ def _hub_file_folder(repo_id: str, filename: str) -> str:
     return os.path.dirname(path)
 
 
+@torch.no_grad()
+def encode_text_batch(text_encoder, tokenizer, texts, max_length, device) -> List[torch.Tensor]:
+    """Chat-templated prompts -> per-item (L_i, D) final states from ONE forward.
+    Right-padded, no attention mask: for the causal model the pads sit after
+    each prompt's real tokens, so neither causal attention nor the gated-delta
+    recurrence can reach them -- each item's first L_i states equal its
+    unpadded forward (the HF mask path only zeroes padded states, which is
+    left-padding hygiene, and would also drop attention off its fast path).
+    Default arange positions are right for right padding."""
+    enc = tokenizer(
+        texts,
+        return_tensors="pt",
+        padding=True,
+        padding_side="right",
+        truncation=True,
+        max_length=max_length,
+    )
+    lengths = enc["attention_mask"].sum(1).tolist()
+    states = text_encoder(
+        input_ids=enc["input_ids"].to(device), use_cache=False
+    ).last_hidden_state
+    return [states[i, :n] for i, n in enumerate(lengths)]
+
+
 def _drop_repeats(images):
     """ctrl_img_1 defaults to ctrl_img, so a single reference can arrive twice."""
     kept = []
@@ -559,21 +583,24 @@ class MossImageModel(BaseModel):
         else:
             control_images = [None] * len(prompt)
 
-        embeds = []
-        # one unpadded forward per prompt, like the vendor (padding would leak
-        # into the gated-delta recurrence)
-        for text, images in zip(prompt, control_images):
+        # text-only prompts go through ONE right-padded forward (exact, see
+        # encode_text_batch); prompts with references stay one forward each
+        # (their image-pad runs differ per prompt)
+        embeds: List[Optional[torch.Tensor]] = [None] * len(prompt)
+        text_only = [i for i, images in enumerate(control_images) if not images]
+        if text_only:
+            states = encode_text_batch(
+                self.text_encoder,
+                self.tokenizer,
+                [self._chat(prompt[i]) for i in text_only],
+                self.max_text_length,
+                self.device_torch,
+            )
+            for i, st in zip(text_only, states):
+                embeds[i] = st.to(self.torch_dtype)
+        for i, (text, images) in enumerate(zip(prompt, control_images)):
             if images:
-                states = self._encode_with_images(text, images)
-            else:
-                input_ids = self.tokenizer(
-                    self._chat(text),
-                    return_tensors="pt",
-                    truncation=True,
-                    max_length=self.max_text_length,
-                )["input_ids"].to(self.device_torch)
-                states = self.text_encoder(input_ids=input_ids, use_cache=False).last_hidden_state[0]
-            embeds.append(states.to(self.torch_dtype))
+                embeds[i] = self._encode_with_images(text, images).to(self.torch_dtype)
         return AdvancedPromptEmbeds(text_embeds=embeds)
 
     def get_loss_target(self, *args, **kwargs):