diff --git a/src/diffusers/modular_pipelines/anima/before_denoise.py b/src/diffusers/modular_pipelines/anima/before_denoise.py index 25f38cd0cb65..dbfe82d7f35d 100644 --- a/src/diffusers/modular_pipelines/anima/before_denoise.py +++ b/src/diffusers/modular_pipelines/anima/before_denoise.py @@ -61,6 +61,63 @@ def retrieve_timesteps( return timesteps, num_inference_steps +# Copied from diffusers.modular_pipelines.z_image.before_denoise.repeat_tensor_to_batch_size +def repeat_tensor_to_batch_size( + input_name: str, + input_tensor: torch.Tensor, + batch_size: int, + num_images_per_prompt: int = 1, +) -> torch.Tensor: + """Repeat tensor elements to match the final batch size. + + This function expands a tensor's batch dimension to match the final batch size (batch_size * num_images_per_prompt) + by repeating each element along dimension 0. + + The input tensor must have batch size 1 or batch_size. The function will: + - If batch size is 1: repeat each element (batch_size * num_images_per_prompt) times + - If batch size equals batch_size: repeat each element num_images_per_prompt times + + Args: + input_name (str): Name of the input tensor (used for error messages) + input_tensor (torch.Tensor): The tensor to repeat. Must have batch size 1 or batch_size. + batch_size (int): The base batch size (number of prompts) + num_images_per_prompt (int, optional): Number of images to generate per prompt. Defaults to 1. + + Returns: + torch.Tensor: The repeated tensor with final batch size (batch_size * num_images_per_prompt) + + Raises: + ValueError: If input_tensor is not a torch.Tensor or has invalid batch size + + Examples: + tensor = torch.tensor([[1, 2, 3]]) # shape: [1, 3] repeated = repeat_tensor_to_batch_size("image", tensor, + batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) - shape: + [4, 3] + + tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]) # shape: [2, 3] repeated = repeat_tensor_to_batch_size("image", + tensor, batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [4, 5, 6]]) + - shape: [4, 3] + """ + # make sure input is a tensor + if not isinstance(input_tensor, torch.Tensor): + raise ValueError(f"`{input_name}` must be a tensor") + + # make sure input tensor e.g. image_latents has batch size 1 or batch_size same as prompts + if input_tensor.shape[0] == 1: + repeat_by = batch_size * num_images_per_prompt + elif input_tensor.shape[0] == batch_size: + repeat_by = num_images_per_prompt + else: + raise ValueError( + f"`{input_name}` must have have batch size 1 or {batch_size}, but got {input_tensor.shape[0]}" + ) + + # expand the tensor to match the batch_size * num_images_per_prompt + input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0) + + return input_tensor + + class AnimaTextConditioningStep(ModularPipelineBlocks): model_name = "anima" @@ -270,6 +327,62 @@ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> Pi return components, state +class AnimaImageInputStep(ModularPipelineBlocks): + model_name = "anima" + + @property + def description(self) -> str: + return ( + "Input processing step that expands Anima image latents to the final denoising batch " + "and derives height/width from the latents when not provided." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image_latents"), + InputParam( + "batch_size", + required=True, + type_hint=int, + description="Number of input prompts before `num_images_per_prompt` expansion.", + ), + InputParam.template("num_images_per_prompt"), + InputParam.template("height"), + InputParam.template("width"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "image_latents", + type_hint=torch.Tensor, + description="Image latents expanded to the final denoising batch.", + ), + OutputParam("height", type_hint=int, description="Image height used for generation."), + OutputParam("width", type_hint=int, description="Image width used for generation."), + ] + + @torch.no_grad() + def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + latent_height, latent_width = block_state.image_latents.shape[-2:] + block_state.height = block_state.height or latent_height * components.vae_scale_factor + block_state.width = block_state.width or latent_width * components.vae_scale_factor + + block_state.image_latents = repeat_tensor_to_batch_size( + input_name="image_latents", + input_tensor=block_state.image_latents, + batch_size=block_state.batch_size, + num_images_per_prompt=block_state.num_images_per_prompt, + ) + + self.set_block_state(state, block_state) + return components, state + + class AnimaPrepareLatentsStep(ModularPipelineBlocks): model_name = "anima" @@ -370,6 +483,19 @@ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> Pi return components, state +# Copied from diffusers.modular_pipelines.qwenimage.before_denoise.get_timesteps +def get_timesteps(scheduler, num_inference_steps, strength): + # get the original timestep using init_timestep + init_timestep = min(num_inference_steps * strength, num_inference_steps) + + t_start = int(max(num_inference_steps - init_timestep, 0)) + timesteps = scheduler.timesteps[t_start * scheduler.order :] + if hasattr(scheduler, "set_begin_index"): + scheduler.set_begin_index(t_start * scheduler.order) + + return timesteps, num_inference_steps - t_start + + class AnimaSetTimestepsStep(ModularPipelineBlocks): model_name = "anima" @@ -414,3 +540,175 @@ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> Pi self.set_block_state(state, block_state) return components, state + + +class AnimaImg2ImgSetTimestepsStep(ModularPipelineBlocks): + """Set the scheduler timesteps for Anima image-to-image inference. + + This step computes the full timestep schedule, then slices it based on ``strength`` via ``get_timesteps()``, which + also sets the scheduler's begin index. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 0.9): + How much to transform the reference image. + + Outputs: + timesteps (`Tensor`): + Timestep schedule sliced by ``strength``. + num_inference_steps (`int`): + Number of denoising steps after strength-based slicing. + """ + + model_name = "anima" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def description(self) -> str: + return "Set the scheduler timesteps for Anima image-to-image inference, sliced by strength." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps"), + InputParam.template("sigmas"), + InputParam.template("strength"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "timesteps", + type_hint=torch.Tensor, + description="Timestep schedule sliced by strength.", + ), + OutputParam( + "num_inference_steps", + type_hint=int, + description="Number of denoising steps after strength-based slicing.", + ), + ] + + @torch.no_grad() + def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + sigmas = ( + np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps) + if block_state.sigmas is None + else block_state.sigmas + ) + block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps( + components.scheduler, + device=device, + sigmas=sigmas, + ) + block_state.timesteps, block_state.num_inference_steps = get_timesteps( + components.scheduler, block_state.num_inference_steps, block_state.strength + ) + + self.set_block_state(state, block_state) + return components, state + + +class AnimaImg2ImgPrepareLatentsStep(ModularPipelineBlocks): + """Prepares noisy latents for Anima image-to-image generation. + + Generates noise and mixes it with the image latents via ``scheduler.scale_noise()`` at the first sliced timestep. + The image latents are expected to already be expanded to the final batch size by ``AnimaImageInputStep``. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + image_latents (`Tensor`): + Encoded image latents, expanded to the final denoising batch. + timesteps (`Tensor`): + Timestep schedule sliced by ``strength`` from ``AnimaImg2ImgSetTimestepsStep``. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Pre-computed noise tensor. Generated randomly if ``None``. + dtype (`torch.dtype`): + Dtype used by the Anima denoiser. + height (`int`): + Image height. + width (`int`): + Image width. + + Outputs: + latents (`Tensor`): + Noisy image latents for the denoising loop. + padding_mask (`Tensor`): + Cosmos padding mask for the image latents. + """ + + model_name = "anima" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def description(self) -> str: + return ( + "Prepares noisy image-to-image latents for Anima by adding noise to the encoded " + "image latents via scheduler.scale_noise()." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image_latents"), + InputParam.template("timesteps", required=True), + InputParam.template("generator"), + InputParam.template("latents"), + InputParam("dtype", type_hint=torch.dtype, description="Dtype used by the Anima denoiser."), + InputParam.template("height"), + InputParam.template("width"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="Noisy latents for the denoising loop."), + OutputParam("padding_mask", type_hint=torch.Tensor, description="Cosmos padding mask for image latents."), + ] + + @torch.no_grad() + def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + image_latents = block_state.image_latents.to(device=device, dtype=torch.float32) + + if block_state.latents is None: + noise = randn_tensor( + image_latents.shape, + generator=block_state.generator, + device=device, + dtype=torch.float32, + ) + else: + noise = block_state.latents.to(device=device, dtype=torch.float32) + + latent_timestep = block_state.timesteps[:1].repeat(image_latents.shape[0]) + block_state.latents = components.scheduler.scale_noise(image_latents, latent_timestep, noise) + + block_state.padding_mask = block_state.latents.new_zeros( + 1, 1, block_state.height, block_state.width, dtype=block_state.dtype + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/anima/encoders.py b/src/diffusers/modular_pipelines/anima/encoders.py index bdeecd28737b..68950f97be83 100644 --- a/src/diffusers/modular_pipelines/anima/encoders.py +++ b/src/diffusers/modular_pipelines/anima/encoders.py @@ -17,6 +17,8 @@ from ...configuration_utils import FrozenDict from ...guiders import ClassifierFreeGuidance +from ...image_processor import VaeImageProcessor +from ...models import AutoencoderKLQwenImage from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import AnimaModularPipeline @@ -251,3 +253,152 @@ def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> Pi self.set_block_state(state, block_state) return components, state + + +# Copied from diffusers.modular_pipelines.qwenimage.encoders.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +# Copied from diffusers.modular_pipelines.qwenimage.encoders.encode_vae_image +def encode_vae_image( + image: torch.Tensor, + vae: AutoencoderKLQwenImage, + generator: torch.Generator, + device: torch.device, + dtype: torch.dtype, + latent_channels: int = 16, + sample_mode: str = "argmax", +): + if not isinstance(image, torch.Tensor): + raise ValueError(f"Expected image to be a tensor, got {type(image)}.") + + # preprocessed image should be a 4D tensor: batch_size, num_channels, height, width + if image.dim() == 4: + image = image.unsqueeze(2) + elif image.dim() != 5: + raise ValueError(f"Expected image dims 4 or 5, got {image.dim()}.") + + image = image.to(device=device, dtype=dtype) + + if isinstance(generator, list): + image_latents = [ + retrieve_latents(vae.encode(image[i : i + 1]), generator=generator[i], sample_mode=sample_mode) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents(vae.encode(image), generator=generator, sample_mode=sample_mode) + latents_mean = ( + torch.tensor(vae.config.latents_mean) + .view(1, latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + latents_std = ( + torch.tensor(vae.config.latents_std) + .view(1, latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + image_latents = (image_latents - latents_mean) / latents_std + + return image_latents + + +class AnimaImg2ImgVaeEncoderStep(ModularPipelineBlocks): + """VAE Encoder step for Anima image-to-image generation. + + Preprocesses the input image and encodes it with the VAE, producing ``image_latents``. Timestep slicing is handled + downstream by ``AnimaImg2ImgSetTimestepsStep`` and noise addition by ``AnimaImg2ImgPrepareLatentsStep``. + + Components: + vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + + Inputs: + image (`PIL.Image.Image`): + Input image to encode. + height (`int`, *optional*): + Height of the output image. Defaults to pipeline default. + width (`int`, *optional*): + Width of the output image. Defaults to pipeline default. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + image_latents (`Tensor`): + Encoded image latents. + height (`int`): + Output image height. + width (`int`): + Output image width. + """ + + model_name = "anima" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLQwenImage), + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 8}), + default_creation_method="from_config", + ), + ] + + @property + def description(self) -> str: + return ( + "VAE Encoder step for Anima image-to-image generation. Encodes the input image to produce image_latents." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image"), + InputParam.template("height"), + InputParam.template("width"), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("image_latents", type_hint=torch.Tensor, description="Encoded image latents."), + OutputParam("height", type_hint=int, description="Image height used for generation."), + OutputParam("width", type_hint=int, description="Image width used for generation."), + ] + + @torch.no_grad() + def __call__(self, components: AnimaModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + + block_state.height = block_state.height or components.default_height + block_state.width = block_state.width or components.default_width + + processed_image = components.image_processor.preprocess( + image=block_state.image, height=block_state.height, width=block_state.width + ) + + block_state.image_latents = encode_vae_image( + image=processed_image, + vae=components.vae, + generator=block_state.generator, + device=device, + dtype=components.vae.dtype, + latent_channels=components.num_channels_latents, + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py b/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py index fc71b87f62d8..ba48c7ecebb5 100644 --- a/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py +++ b/src/diffusers/modular_pipelines/anima/modular_blocks_anima.py @@ -12,9 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks from ..modular_pipeline_utils import OutputParam from .before_denoise import ( + AnimaImageInputStep, + AnimaImg2ImgPrepareLatentsStep, + AnimaImg2ImgSetTimestepsStep, AnimaPrepareLatentsStep, AnimaSetTimestepsStep, AnimaTextConditioningStep, @@ -22,7 +25,7 @@ ) from .decoders import AnimaProcessImagesOutputStep, AnimaVaeDecoderStep from .denoise import AnimaDenoiseStep -from .encoders import AnimaTextEncoderStep +from .encoders import AnimaImg2ImgVaeEncoderStep, AnimaTextEncoderStep # auto_docstring @@ -35,8 +38,6 @@ class AnimaCoreDenoiseStep(SequentialPipelineBlocks): (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) Inputs: - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. qwen_prompt_embeds (`Tensor`): Qwen prompt embeddings generated by the text encoder step. qwen_attention_mask (`Tensor`): @@ -53,6 +54,8 @@ class AnimaCoreDenoiseStep(SequentialPipelineBlocks): Negative T5 prompt token ids generated by the text encoder step. negative_t5_attention_mask (`Tensor`, *optional*): Negative T5 prompt attention mask generated by the text encoder step. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. height (`int`, *optional*): The height in pixels of the generated image. width (`int`, *optional*): @@ -122,18 +125,202 @@ def outputs(self): return [OutputParam.template("images")] +# auto_docstring +class AnimaImg2ImgCoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block for Anima image-to-image generation. Expects ``image_latents`` already in state (produced by + ``AnimaImg2ImgVaeEncoderStep`` upstream). + + Components: + text_conditioner (`AnimaTextConditioner`) transformer (`CosmosTransformer3DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + qwen_prompt_embeds (`Tensor`): + Qwen prompt embeddings generated by the text encoder step. + qwen_attention_mask (`Tensor`): + Qwen prompt attention mask generated by the text encoder step. + t5_input_ids (`Tensor`): + T5 prompt token ids generated by the text encoder step. + t5_attention_mask (`Tensor`): + T5 prompt attention mask generated by the text encoder step. + negative_qwen_prompt_embeds (`Tensor`, *optional*): + Negative Qwen prompt embeddings generated by the text encoder step. + negative_qwen_attention_mask (`Tensor`, *optional*): + Negative Qwen prompt attention mask generated by the text encoder step. + negative_t5_input_ids (`Tensor`, *optional*): + Negative T5 prompt token ids generated by the text encoder step. + negative_t5_attention_mask (`Tensor`, *optional*): + Negative T5 prompt attention mask generated by the text encoder step. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + image_latents (`Tensor`): + Encoded image latents from ``AnimaImg2ImgVaeEncoderStep``. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img transformation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + **denoiser_input_fields (`None`, *optional*): + The conditional model inputs for the Anima denoiser. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + block_classes = [ + AnimaTextConditioningStep, + AnimaTextInputStep, + AnimaImageInputStep, + AnimaImg2ImgSetTimestepsStep, + AnimaImg2ImgPrepareLatentsStep, + AnimaDenoiseStep, + ] + block_names = ["text_conditioning", "input", "image_input", "set_timesteps", "prepare_latents", "denoise"] + + @property + def description(self) -> str: + return ( + "Denoise block for Anima image-to-image generation. " + "Uses image_latents already in state from AnimaImg2ImgVaeEncoderStep." + ) + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# auto_docstring +class AnimaAutoCoreDenoiseStep(AutoPipelineBlocks): + """ + Denoise step that selects between text-to-image and image-to-image denoising based on whether ``image_latents`` is + present in state. - `AnimaCoreDenoiseStep` (text2image) is used when no ``image_latents`` are present. - + `AnimaImg2ImgCoreDenoiseStep` (img2img) is used when ``image_latents`` are present (set upstream by + ``AnimaAutoVaeImageEncoderStep``). + + Components: + text_conditioner (`AnimaTextConditioner`) transformer (`CosmosTransformer3DModel`) vae + (`AutoencoderKLQwenImage`) scheduler (`FlowMatchEulerDiscreteScheduler`) image_processor + (`VaeImageProcessor`) guider (`ClassifierFreeGuidance`) + + Inputs: + qwen_prompt_embeds (`Tensor`): + Qwen prompt embeddings generated by the text encoder step. + qwen_attention_mask (`Tensor`): + Qwen prompt attention mask generated by the text encoder step. + t5_input_ids (`Tensor`): + T5 prompt token ids generated by the text encoder step. + t5_attention_mask (`Tensor`): + T5 prompt attention mask generated by the text encoder step. + negative_qwen_prompt_embeds (`Tensor`, *optional*): + Negative Qwen prompt embeddings generated by the text encoder step. + negative_qwen_attention_mask (`Tensor`, *optional*): + Negative Qwen prompt attention mask generated by the text encoder step. + negative_t5_input_ids (`Tensor`, *optional*): + Negative T5 prompt token ids generated by the text encoder step. + negative_t5_attention_mask (`Tensor`, *optional*): + Negative T5 prompt attention mask generated by the text encoder step. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + image_latents (`Tensor`, *optional*): + Encoded image latents. When present in state, img2img denoising is selected. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img transformation. + **denoiser_input_fields (`None`, *optional*): + The conditional model inputs for the Anima denoiser. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + block_classes = [AnimaImg2ImgCoreDenoiseStep, AnimaCoreDenoiseStep] + block_names = ["img2img", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self) -> str: + return ( + "Denoise step that selects between text-to-image and image-to-image denoising based on whether " + "image_latents is present in state." + " - `AnimaCoreDenoiseStep` (text2image) is used when no image_latents are present." + " - `AnimaImg2ImgCoreDenoiseStep` (img2img) is used when image_latents are present." + ) + + +# auto_docstring +class AnimaAutoVaeImageEncoderStep(AutoPipelineBlocks): + """ + VAE Image Encoder step that encodes the input image to produce ``image_latents``. This step is skipped when no + image is provided (text-to-image workflow). + + Components: + vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + + Inputs: + image (`Image`, *optional*): + Input image for image-to-image generation. When provided, the image is encoded to ``image_latents``. When + not provided, this step is skipped. + height (`int`, *optional*): + Height of the output image. + width (`int`, *optional*): + Width of the output image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + image_latents (`Tensor`): + Encoded image latents used by ``AnimaAutoCoreDenoiseStep`` to trigger img2img denoising. + """ + + block_classes = [AnimaImg2ImgVaeEncoderStep] + block_names = ["vae_encoder"] + block_trigger_inputs = ["image"] + + @property + def description(self) -> str: + return ( + "VAE Image Encoder step that encodes the input image to produce image_latents. " + "Skipped when no image is provided (text-to-image workflow)." + ) + + # auto_docstring class AnimaAutoBlocks(SequentialPipelineBlocks): """ - Auto Modular pipeline for text-to-image generation using Anima. + Auto Modular pipeline for text-to-image and image-to-image generation using Anima. Supported workflows: - `text2image`: requires `prompt` + - `img2img`: requires `prompt`, `image` Components: - text_encoder (`Qwen3Model`) tokenizer (`Qwen2Tokenizer`) t5_tokenizer (`T5TokenizerFast`) text_conditioner - (`AnimaTextConditioner`) guider (`ClassifierFreeGuidance`) transformer (`CosmosTransformer3DModel`) scheduler - (`FlowMatchEulerDiscreteScheduler`) vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + text_encoder (`Qwen3Model`) tokenizer (`Qwen2Tokenizer`) t5_tokenizer (`T5Tokenizer`) guider + (`ClassifierFreeGuidance`) text_conditioner (`AnimaTextConditioner`) transformer (`CosmosTransformer3DModel`) + scheduler (`FlowMatchEulerDiscreteScheduler`) vae (`AutoencoderKLQwenImage`) image_processor + (`VaeImageProcessor`) Inputs: prompt (`str`): @@ -144,6 +331,8 @@ class AnimaAutoBlocks(SequentialPipelineBlocks): Maximum sequence length for prompt encoding. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. + image (`Image | list`, *optional*): + Reference image(s) for image-to-image generation. When provided, img2img workflow is used. height (`int`, *optional*): The height in pixels of the generated image. width (`int`, *optional*): @@ -156,6 +345,8 @@ class AnimaAutoBlocks(SequentialPipelineBlocks): The number of denoising steps. sigmas (`list`, *optional*): Custom sigmas for the denoising process. + strength (`float`, *optional*, defaults to 0.9): + How much to transform the reference image (img2img only). **denoiser_input_fields (`None`, *optional*): The conditional model inputs for the Anima denoiser. output_type (`str`, *optional*, defaults to pil): @@ -168,15 +359,19 @@ class AnimaAutoBlocks(SequentialPipelineBlocks): block_classes = [ AnimaTextEncoderStep, - AnimaCoreDenoiseStep, + AnimaAutoVaeImageEncoderStep, + AnimaAutoCoreDenoiseStep, AnimaDecodeStep, ] - block_names = ["text_encoder", "denoise", "decode"] - _workflow_map = {"text2image": {"prompt": True}} + block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] + _workflow_map = { + "text2image": {"prompt": True}, + "img2img": {"image": True, "prompt": True}, + } @property def description(self) -> str: - return "Auto Modular pipeline for text-to-image generation using Anima." + return "Auto Modular pipeline for text-to-image and image-to-image generation using Anima." @property def outputs(self): diff --git a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py index 9afe9ee6fab3..dc82483161db 100644 --- a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py +++ b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py @@ -15,6 +15,8 @@ import tempfile import unittest +import numpy as np +import PIL.Image import torch from transformers import Qwen2Tokenizer, Qwen3Config, Qwen3Model, T5TokenizerFast @@ -47,6 +49,21 @@ ], } +ANIMA_IMG2IMG_WORKFLOWS = { + "img2img": [ + ("text_encoder", "AnimaTextEncoderStep"), + ("vae_encoder", "AnimaImg2ImgVaeEncoderStep"), + ("denoise.text_conditioning", "AnimaTextConditioningStep"), + ("denoise.input", "AnimaTextInputStep"), + ("denoise.image_input", "AnimaImageInputStep"), + ("denoise.set_timesteps", "AnimaImg2ImgSetTimestepsStep"), + ("denoise.prepare_latents", "AnimaImg2ImgPrepareLatentsStep"), + ("denoise.denoise", "AnimaDenoiseStep"), + ("decode.decode", "AnimaVaeDecoderStep"), + ("decode.postprocess", "AnimaProcessImagesOutputStep"), + ], +} + def get_dummy_components(): torch.manual_seed(0) @@ -118,6 +135,11 @@ def get_dummy_components(): } +def get_dummy_image(height=32, width=32): + image_array = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8) + return PIL.Image.fromarray(image_array) + + class AnimaTextConditionerFastTests(unittest.TestCase): def test_conditioner_output_shape_and_padding(self): conditioner = AnimaTextConditioner( @@ -229,3 +251,72 @@ def test_load_lora_weights(self): assert "dummy" in pipe.transformer.peft_config assert "dummy" in pipe.text_conditioner.peft_config + + +class TestAnimaImg2ImgModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = AnimaModularPipeline + pipeline_blocks_class = AnimaAutoBlocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-anima-modular-pipe" + params = frozenset(["prompt", "image", "strength", "height", "width", "negative_prompt"]) + batch_params = frozenset(["prompt", "negative_prompt"]) + expected_workflow_blocks = ANIMA_IMG2IMG_WORKFLOWS + + def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): + pipe = self.pipeline_blocks_class().init_pipeline(components_manager=components_manager) + pipe.update_components(**get_dummy_components()) + pipe.to(dtype=torch_dtype) + pipe.set_progress_bar_config(disable=None) + return pipe + + def get_dummy_inputs(self, seed=0): + generator = torch.Generator(device="cpu").manual_seed(seed) + return { + "prompt": "dance monkey", + "negative_prompt": "bad quality", + "image": get_dummy_image(32, 32), + "strength": 0.8, + "generator": generator, + "num_inference_steps": 2, + "height": 32, + "width": 32, + "max_sequence_length": 16, + "output_type": "pt", + } + + def test_inference_basic(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + output = pipe(**inputs).images + + assert output.shape == (1, 3, 32, 32) + assert not torch.isnan(output).any() + + def test_inference_strength_low(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs["strength"] = 0.3 + output = pipe(**inputs).images + + assert output.shape == (1, 3, 32, 32) + assert not torch.isnan(output).any() + + def test_inference_strength_high(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs["strength"] = 0.95 + output = pipe(**inputs).images + + assert output.shape == (1, 3, 32, 32) + assert not torch.isnan(output).any() + + def test_inference_empty_negative_prompt(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs["negative_prompt"] = "" + output = pipe(**inputs).images + + assert output.shape == (1, 3, 32, 32) + assert not torch.isnan(output).any() + + def test_inference_batch_single_identical(self): + super().test_inference_batch_single_identical(expected_max_diff=5e-4)