From 06ce29feb4bed92147b37516307bd4c08ff4fdbf Mon Sep 17 00:00:00 2001
From: Haoming <hmstudy02@gmail.com>
Date: Fri, 21 Nov 2025 14:55:39 +0800
Subject: [PATCH] xyz

---
 backend/args.py                         |  3 ++-
 backend/nn/lumina.py                    | 10 +++++++---
 backend/text_processing/gemma_engine.py |  6 +++---
 modules/processing.py                   |  1 -
 4 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/backend/args.py b/backend/args.py
index 3fd62de3..06a2f721 100644
--- a/backend/args.py
+++ b/backend/args.py
@@ -123,5 +123,6 @@ parser.add_argument("--sage-accum-dtype", type=Sage_pv_accum_dtype, default=Sage
 
 args, _ = parser.parse_known_args()
 
-dynamic_args = dict(embedding_dir="./embeddings", emphasis_name="original")
+# TODO: Stop using this to hack every problem...
+dynamic_args = dict(embedding_dir="./embeddings", emphasis_name="original", num_tokens={"c": None, "uc": None})
 """Some parameters that are changed by the Webui"""
diff --git a/backend/nn/lumina.py b/backend/nn/lumina.py
index 03a06b1b..6b85854c 100644
--- a/backend/nn/lumina.py
+++ b/backend/nn/lumina.py
@@ -403,6 +403,12 @@ class NextDiT(nn.Module):
 
         return padded_full_embed, mask, img_sizes, l_effective_cap_len, freqs_cis
 
+    @staticmethod
+    def load_tokens(transformer_options: dict) -> int:
+        opt: list[int] = transformer_options.get("cond_or_uncond", [0])
+        c_uc: str = "c" if opt[0] == 0 else "uc"
+        return dynamic_args["num_tokens"][c_uc]
+
     def forward(self, x, timesteps, context, num_tokens=None, attention_mask=None, **kwargs):
         t = 1.0 - timesteps
         cap_feats = context
@@ -416,9 +422,7 @@ class NextDiT(nn.Module):
         cap_feats = self.cap_embedder(cap_feats)
 
         transformer_options = kwargs.get("transformer_options", {})
-
-        _c = transformer_options.get("cond_or_uncond", [0])[0]
-        num_tokens = num_tokens or dynamic_args["num_tokens"][_c]
+        num_tokens = num_tokens or self.load_tokens(transformer_options)
 
         x_is_tensor = isinstance(x, torch.Tensor)
         x, mask, img_size, cap_size, freqs_cis = self.patchify_and_embed(x, cap_feats, cap_mask, t, num_tokens, transformer_options=transformer_options)
diff --git a/backend/text_processing/gemma_engine.py b/backend/text_processing/gemma_engine.py
index b9084ccd..98f306ab 100644
--- a/backend/text_processing/gemma_engine.py
+++ b/backend/text_processing/gemma_engine.py
@@ -114,7 +114,7 @@ class GemmaTextProcessingEngine:
                     #     tokens += [self.id_pad] * remaining_count
                     #     multipliers += [1.0] * remaining_count
 
-                    z = self.process_tokens([tokens], [multipliers])[0]
+                    z = self.process_tokens([tokens], [multipliers], texts.is_negative_prompt)[0]
                     line_z_values.append(z)
                 cache[line] = line_z_values
 
@@ -155,9 +155,9 @@ class GemmaTextProcessingEngine:
 
         return torch.cat(embeds_out), torch.tensor(attention_masks, device=device, dtype=torch.long), num_tokens
 
-    def process_tokens(self, batch_tokens, batch_multipliers):
+    def process_tokens(self, batch_tokens, batch_multipliers, negative: bool):
         embeds, mask, count = self.process_embeds(batch_tokens)
-        dynamic_args["num_tokens"] = [max(1, torch.sum(mask).item())] + dynamic_args.get("num_tokens", [])
+        dynamic_args["num_tokens"]["uc" if negative else "c"] = int(max(1, torch.sum(mask).item()))
         _, z = self.text_encoder(
             None,
             attention_mask=mask,
diff --git a/modules/processing.py b/modules/processing.py
index ea6facee..8f96a34c 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -1172,7 +1172,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         p.scripts.postprocess(p, res)
 
     dynamic_args.pop("ref_latents", None)
-    dynamic_args.pop("num_tokens", None)
     return res