From 2938fda9b872d0844b4f62d13feffbc968d564cc Mon Sep 17 00:00:00 2001 From: lucasruan1618 Date: Mon, 5 Oct 2026 05:27:13 +0000 Subject: [PATCH] Fix Anima text encoder exports and standalone core denoising --- .../modular_pipelines/anima/encoders.py | 8 +++ .../anima/modular_pipeline.py | 2 +- .../anima/test_modular_pipeline_anima.py | 51 +++++++++++++++++++ 3 files changed, 60 insertions(+), 1 deletion(-) diff --git a/src/diffusers/modular_pipelines/anima/encoders.py b/src/diffusers/modular_pipelines/anima/encoders.py index 726a68280d82..2f86334fd507 100644 --- a/src/diffusers/modular_pipelines/anima/encoders.py +++ b/src/diffusers/modular_pipelines/anima/encoders.py @@ -58,41 +58,49 @@ def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "qwen_prompt_embeds", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Qwen prompt embeddings to be consumed by the Anima text conditioner.", ), OutputParam( "qwen_attention_mask", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Qwen prompt attention mask to be consumed by the Anima text conditioner.", ), OutputParam( "t5_input_ids", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="T5 prompt token ids to be consumed by the Anima text conditioner.", ), OutputParam( "t5_attention_mask", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="T5 prompt attention mask to be consumed by the Anima text conditioner.", ), OutputParam( "negative_qwen_prompt_embeds", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Negative Qwen prompt embeddings to be consumed by the Anima text conditioner.", ), OutputParam( "negative_qwen_attention_mask", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Negative Qwen prompt attention mask to be consumed by the Anima text conditioner.", ), OutputParam( "negative_t5_input_ids", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Negative T5 prompt token ids to be consumed by the Anima text conditioner.", ), OutputParam( "negative_t5_attention_mask", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Negative T5 prompt attention mask to be consumed by the Anima text conditioner.", ), diff --git a/src/diffusers/modular_pipelines/anima/modular_pipeline.py b/src/diffusers/modular_pipelines/anima/modular_pipeline.py index 1430ac555851..a7136bb97ca7 100644 --- a/src/diffusers/modular_pipelines/anima/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/anima/modular_pipeline.py @@ -39,7 +39,7 @@ def default_sample_size(self): @property def vae_scale_factor(self): vae_scale_factor = 8 - if self.vae is not None: + if getattr(self, "vae", None) is not None: vae_scale_factor = 2 ** len(self.vae.temperal_downsample) return vae_scale_factor diff --git a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py index c3ae0d9f2e0d..ff8a52839df1 100644 --- a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py +++ b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py @@ -17,6 +17,7 @@ import numpy as np import PIL.Image +import pytest import torch from diffusers import ( @@ -25,6 +26,7 @@ AnimaTextConditioner, CosmosTransformer3DModel, ) +from diffusers.modular_pipelines.anima.modular_blocks_anima import AnimaCoreDenoiseStep from ...testing_utils import enable_full_determinism, is_lora, require_peft_backend from ..testing_utils import ( @@ -127,6 +129,55 @@ def get_dummy_inputs(self, seed=0): class TestAnimaModularPipelineFast(AnimaModularPipelineTesterConfig, ModularPipelineTesterMixin): + @pytest.mark.parametrize("negative_prompt", [None, "", "bad quality"]) + def test_split_text_encoder(self, negative_prompt): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs["negative_prompt"] = negative_prompt + expected = pipe(**inputs).images + + blocks = self.pipeline_blocks_class() + text_block = blocks.sub_blocks.pop("text_encoder") + text_pipe = text_block.init_pipeline(self.pretrained_model_name_or_path) + text_pipe.load_components(dtype=torch.float32) + denoise_pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) + denoise_pipe.load_components(dtype=torch.float32) + + inputs = self.get_dummy_inputs() + inputs["negative_prompt"] = negative_prompt + text_inputs = {name: inputs.pop(name) for name in ("prompt", "negative_prompt", "max_sequence_length")} + embeddings = text_pipe(**text_inputs).get_by_kwargs("denoiser_input_fields") + output = denoise_pipe(**embeddings, **inputs).images + + torch.testing.assert_close(output, expected, rtol=1e-5, atol=1e-5) + + def test_standalone_core_denoise(self): + text_block = self.pipeline_blocks_class().sub_blocks["text_encoder"] + text_pipe = text_block.init_pipeline(self.pretrained_model_name_or_path) + text_pipe.load_components(dtype=torch.float32) + inputs = self.get_dummy_inputs() + text_inputs = {name: inputs.pop(name) for name in ("prompt", "negative_prompt", "max_sequence_length")} + embeddings = text_pipe(**text_inputs).get( + [ + "qwen_prompt_embeds", + "qwen_attention_mask", + "t5_input_ids", + "t5_attention_mask", + "negative_qwen_prompt_embeds", + "negative_qwen_attention_mask", + "negative_t5_input_ids", + "negative_t5_attention_mask", + ] + ) + + denoise_pipe = AnimaCoreDenoiseStep().init_pipeline(self.pretrained_model_name_or_path) + denoise_pipe.load_components(dtype=torch.float32) + inputs.pop("output_type") + output = denoise_pipe(**embeddings, **inputs).get("latents") + + assert output.shape == (1, 4, 1, 4, 4) + assert torch.isfinite(output).all() + def test_inference_empty_negative_prompt(self): pipe = self.get_pipeline()