diff --git a/backend/diffusion_engine/anima.py b/backend/diffusion_engine/anima.py index 5eecd8d4..c361071c 100644 --- a/backend/diffusion_engine/anima.py +++ b/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): diff --git a/backend/nn/anima.py b/backend/nn/anima.py index eabcbe83..ec754171 100644 --- a/backend/nn/anima.py +++ b/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 diff --git a/backend/sampling/condition.py b/backend/sampling/condition.py index a1dc241c..524871e9 100644 --- a/backend/sampling/condition.py +++ b/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, diff --git a/backend/text_processing/anima_engine.py b/backend/text_processing/anima_engine.py index d1826346..b77d8ec5 100644 --- a/backend/text_processing/anima_engine.py +++ b/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()