batch

backend/diffusion_engine/anima.py

@@ -30,6 +30,7 @@ class Anima(ForgeDiffusionEngine):
             text_encoder=clip.cond_stage_model.qwen3_06b,
             qwen_tokenizer=clip.tokenizer.qwen3_06b,
             t5_tokenizer=clip.tokenizer.t5xxl,
+            unet=unet,
         )
 
         self.forge_objects = ForgeObjects(unet=unet, clip=clip, vae=vae, clipvision=None)
@@ -46,7 +47,7 @@ class Anima(ForgeDiffusionEngine):
     @torch.inference_mode()
     def get_prompt_lengths_on_ui(self, prompt):
         token_count = len(self.text_processing_engine_anima.tokenize([prompt])[0][0])
-        return token_count, max(999, token_count)
+        return token_count, max(512, token_count)
 
     @torch.inference_mode()
     def encode_first_stage(self, x: torch.Tensor):

backend/nn/anima.py

@@ -797,6 +797,7 @@ class Anima(MiniTrainDIT):
 
     def preprocess_text_embeds(self, text_embeds, text_ids):
         if text_ids is not None:
-            return self.llm_adapter(text_embeds, text_ids)
+            device = self.llm_adapter.embed.weight.device
+            return self.llm_adapter(text_embeds.to(device), text_ids.to(device))
         else:
             return text_embeds

backend/sampling/condition.py

@@ -91,27 +91,10 @@ class ConditionConstant(Condition):
         return self.cond
 
 
-def anima_preprocess(cross_attn: torch.Tensor, t5xxl_ids: torch.Tensor, t5xxl_weights: torch.Tensor) -> torch.Tensor:
-    device: torch.device = cross_attn.device
-    dtype: torch.dtype = shared.sd_model.forge_objects.unet.model.computation_dtype
-
-    cross_attn = shared.sd_model.forge_objects.unet.model.diffusion_model.preprocess_text_embeds(cross_attn.to(dtype=dtype), t5xxl_ids.to(device=device))
-    if t5xxl_weights is not None:
-        cross_attn *= t5xxl_weights.unsqueeze(-1).to(cross_attn)
-
-    if cross_attn.shape[1] < 512:
-        cross_attn = torch.nn.functional.pad(cross_attn, (0, 0, 0, 512 - cross_attn.shape[1]))
-
-    return cross_attn
-
-
 def compile_conditions(cond):
     if cond is None:
         return None
 
-    if isinstance(cond, dict) and "qwen_cond" in cond and "t5_ids" in cond:
-        cond = anima_preprocess(cond["qwen_cond"], cond["t5_ids"], cond["t5_weights"])
-
     if isinstance(cond, torch.Tensor):
         result = dict(
             cross_attn=cond,

backend/text_processing/anima_engine.py

@@ -1,3 +1,10 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from backend.patcher.unet import UnetPatcher
+
+import weakref
+
 import torch
 
 from backend import memory_management
@@ -14,16 +21,22 @@ class PromptChunk:
 
 
 class AnimaTextProcessingEngine:
-    def __init__(self, text_encoder, qwen_tokenizer, t5_tokenizer):
+    def __init__(self, text_encoder, qwen_tokenizer, t5_tokenizer, unet):
         super().__init__()
 
         self.text_encoder = text_encoder
         self.qwen_tokenizer = qwen_tokenizer
         self.t5_tokenizer = t5_tokenizer
 
+        self._unet = weakref.ref(unet)
+
         self.id_pad = 151643
         self.id_end = 1
 
+    @property
+    def unet(self) -> " UnetPatcher":
+        return self._unet()
+
     def tokenize(self, texts):
         return (
             self.qwen_tokenizer(texts, truncation=False, add_special_tokens=False)["input_ids"],
@@ -72,7 +85,7 @@ class AnimaTextProcessingEngine:
         return chunks
 
     def __call__(self, texts):
-        zs, ti, tw = [], [], []
+        zs = []
         cache = {}
 
         self.emphasis = emphasis.get_current_option(opts.emphasis)()
@@ -92,17 +105,30 @@ class AnimaTextProcessingEngine:
 
                 cache[line] = z
 
-            zs.append(z)
-            ti.append(torch.tensor(chunk.t5_tokens, dtype=torch.int))
-            tw.append(torch.tensor(chunk.t5_multipliers))
+            zs.append(
+                self.anima_preprocess(
+                    z,
+                    torch.tensor(chunk.t5_tokens, dtype=torch.int),
+                    torch.tensor(chunk.t5_multipliers),
+                )
+            )
 
-        z = {
-            "qwen_cond": zs,
-            "t5_ids": ti,
-            "t5_weights": tw,
-        }
+        return zs
 
-        return z
+    def anima_preprocess(self, cross_attn: torch.Tensor, t5xxl_ids: torch.Tensor, t5xxl_weights: torch.Tensor) -> torch.Tensor:
+        dtype: torch.dtype = self.unet.model.computation_dtype
+
+        cross_attn = cross_attn.unsqueeze(0).to(dtype=dtype)
+        t5xxl_ids = t5xxl_ids.unsqueeze(0)
+
+        cross_attn = self.unet.model.diffusion_model.preprocess_text_embeds(cross_attn, t5xxl_ids)
+        if t5xxl_weights is not None:
+            cross_attn *= t5xxl_weights.unsqueeze(0).unsqueeze(-1).to(cross_attn)
+
+        if cross_attn.shape[1] < 512:
+            cross_attn = torch.nn.functional.pad(cross_attn, (0, 0, 0, 512 - cross_attn.shape[1]))
+
+        return cross_attn
 
     def process_embeds(self, batch_tokens):
         device = memory_management.text_encoder_device()