diff --git a/diffusers/modular_pipelines/flux2/decoders.py b/diffusers/modular_pipelines/flux2/decoders.py new file mode 100644 index 0000000000000000000000000000000000000000..81f5ca00dc33dfb91c9f783f88262719fc30e7b3 --- /dev/null +++ b/diffusers/modular_pipelines/flux2/decoders.py @@ -0,0 +1,185 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import Any, Union + +import numpy as np +import PIL +import torch + +from ...configuration_utils import FrozenDict +from ...models import AutoencoderKLFlux2 +from ...pipelines.flux2.image_processor import Flux2ImageProcessor +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class Flux2UnpackLatentsStep(ModularPipelineBlocks): + model_name = "flux2" + + @property + def description(self) -> str: + return "Step that unpacks the latents from the denoising step" + + @property + def inputs(self) -> list[tuple[str, Any]]: + return [ + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The denoised latents from the denoising step", + ), + InputParam( + "latent_ids", + required=True, + type_hint=torch.Tensor, + description="Position IDs for the latents, used for unpacking", + ), + ] + + @property + def intermediate_outputs(self) -> list[str]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="The denoise latents from denoising step, unpacked with position IDs.", + ) + ] + + @staticmethod + def _unpack_latents_with_ids(x: torch.Tensor, x_ids: torch.Tensor) -> torch.Tensor: + """ + Unpack latents using position IDs to scatter tokens into place. + + Args: + x: Packed latents tensor of shape (B, seq_len, C) + x_ids: Position IDs tensor of shape (B, seq_len, 4) with (T, H, W, L) coordinates + + Returns: + Unpacked latents tensor of shape (B, C, H, W) + """ + x_list = [] + for data, pos in zip(x, x_ids): + _, ch = data.shape # noqa: F841 + h_ids = pos[:, 1].to(torch.int64) + w_ids = pos[:, 2].to(torch.int64) + + h = torch.max(h_ids) + 1 + w = torch.max(w_ids) + 1 + + flat_ids = h_ids * w + w_ids + + out = torch.zeros((h * w, ch), device=data.device, dtype=data.dtype) + out.scatter_(0, flat_ids.unsqueeze(1).expand(-1, ch), data) + + out = out.view(h, w, ch).permute(2, 0, 1) + x_list.append(out) + + return torch.stack(x_list, dim=0) + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + latents = block_state.latents + latent_ids = block_state.latent_ids + + latents = self._unpack_latents_with_ids(latents, latent_ids) + + block_state.latents = latents + + self.set_block_state(state, block_state) + return components, state + + +class Flux2DecodeStep(ModularPipelineBlocks): + model_name = "flux2" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLFlux2), + ComponentSpec( + "image_processor", + Flux2ImageProcessor, + config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 32}), + default_creation_method="from_config", + ), + ] + + @property + def description(self) -> str: + return "Step that decodes the denoised latents into images using Flux2 VAE with batch norm denormalization" + + @property + def inputs(self) -> list[tuple[str, Any]]: + return [ + InputParam("output_type", default="pil"), + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The denoised latents from the denoising step", + ), + ] + + @property + def intermediate_outputs(self) -> list[str]: + return [ + OutputParam( + "images", + type_hint=Union[list[PIL.Image.Image], torch.Tensor, np.ndarray], + description="The generated images, can be a list of PIL.Image.Image, torch.Tensor or a numpy array", + ) + ] + + @staticmethod + def _unpatchify_latents(latents): + """Convert patchified latents back to regular format.""" + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width) + latents = latents.permute(0, 1, 4, 2, 5, 3) + latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2) + return latents + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + vae = components.vae + + latents = block_state.latents + + latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to( + latents.device, latents.dtype + ) + latents = latents * latents_bn_std + latents_bn_mean + + latents = self._unpatchify_latents(latents) + + block_state.images = vae.decode(latents, return_dict=False)[0] + block_state.images = components.image_processor.postprocess( + block_state.images, output_type=block_state.output_type + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/flux2/denoise.py b/diffusers/modular_pipelines/flux2/denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..675f14b03c63a0d777dfc7ca2638641d91460403 --- /dev/null +++ b/diffusers/modular_pipelines/flux2/denoise.py @@ -0,0 +1,501 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any + +import torch + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import Flux2Transformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + PipelineState, +) +from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam +from .modular_pipeline import Flux2KleinModularPipeline, Flux2ModularPipeline + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class Flux2LoopDenoiser(ModularPipelineBlocks): + model_name = "flux2" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Flux2Transformer2DModel)] + + @property + def description(self) -> str: + return ( + "Step within the denoising loop that denoises the latents for Flux2. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `Flux2DenoiseLoopWrapper`)" + ) + + @property + def inputs(self) -> list[tuple[str, Any]]: + return [ + InputParam("joint_attention_kwargs"), + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The latents to denoise. Shape: (B, seq_len, C)", + ), + InputParam( + "image_latents", + type_hint=torch.Tensor, + description="Packed image latents for conditioning. Shape: (B, img_seq_len, C)", + ), + InputParam( + "image_latent_ids", + type_hint=torch.Tensor, + description="Position IDs for image latents. Shape: (B, img_seq_len, 4)", + ), + InputParam( + "guidance", + required=True, + type_hint=torch.Tensor, + description="Guidance scale as a tensor", + ), + InputParam( + "prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Text embeddings from Mistral3", + ), + InputParam( + "txt_ids", + required=True, + type_hint=torch.Tensor, + description="4D position IDs for text tokens (T, H, W, L)", + ), + InputParam( + "latent_ids", + required=True, + type_hint=torch.Tensor, + description="4D position IDs for latent tokens (T, H, W, L)", + ), + ] + + @torch.no_grad() + def __call__( + self, components: Flux2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> PipelineState: + latents = block_state.latents + latent_model_input = latents.to(components.transformer.dtype) + img_ids = block_state.latent_ids + + image_latents = getattr(block_state, "image_latents", None) + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1).to(components.transformer.dtype) + image_latent_ids = block_state.image_latent_ids + img_ids = torch.cat([img_ids, image_latent_ids], dim=1) + + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + noise_pred = components.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=block_state.guidance, + encoder_hidden_states=block_state.prompt_embeds, + txt_ids=block_state.txt_ids, + img_ids=img_ids, + joint_attention_kwargs=block_state.joint_attention_kwargs, + return_dict=False, + )[0] + + noise_pred = noise_pred[:, : latents.size(1)] + block_state.noise_pred = noise_pred + + return components, block_state + + +# same as Flux2LoopDenoiser but guidance=None +class Flux2KleinLoopDenoiser(ModularPipelineBlocks): + model_name = "flux2-klein" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Flux2Transformer2DModel)] + + @property + def description(self) -> str: + return ( + "Step within the denoising loop that denoises the latents for Flux2. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `Flux2DenoiseLoopWrapper`)" + ) + + @property + def inputs(self) -> list[tuple[str, Any]]: + return [ + InputParam("joint_attention_kwargs"), + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The latents to denoise. Shape: (B, seq_len, C)", + ), + InputParam( + "image_latents", + type_hint=torch.Tensor, + description="Packed image latents for conditioning. Shape: (B, img_seq_len, C)", + ), + InputParam( + "image_latent_ids", + type_hint=torch.Tensor, + description="Position IDs for image latents. Shape: (B, img_seq_len, 4)", + ), + InputParam( + "prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Text embeddings from Qwen3", + ), + InputParam( + "txt_ids", + required=True, + type_hint=torch.Tensor, + description="4D position IDs for text tokens (T, H, W, L)", + ), + InputParam( + "latent_ids", + required=True, + type_hint=torch.Tensor, + description="4D position IDs for latent tokens (T, H, W, L)", + ), + ] + + @torch.no_grad() + def __call__( + self, components: Flux2KleinModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> PipelineState: + latents = block_state.latents + latent_model_input = latents.to(components.transformer.dtype) + img_ids = block_state.latent_ids + + image_latents = getattr(block_state, "image_latents", None) + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1).to(components.transformer.dtype) + image_latent_ids = block_state.image_latent_ids + img_ids = torch.cat([img_ids, image_latent_ids], dim=1) + + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + noise_pred = components.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=None, + encoder_hidden_states=block_state.prompt_embeds, + txt_ids=block_state.txt_ids, + img_ids=img_ids, + joint_attention_kwargs=block_state.joint_attention_kwargs, + return_dict=False, + )[0] + + noise_pred = noise_pred[:, : latents.size(1)] + block_state.noise_pred = noise_pred + + return components, block_state + + +# support CFG for Flux2-Klein base model +class Flux2KleinBaseLoopDenoiser(ModularPipelineBlocks): + model_name = "flux2-klein" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", Flux2Transformer2DModel), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 4.0}), + default_creation_method="from_config", + ), + ] + + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ + ConfigSpec(name="is_distilled", default=False), + ] + + @property + def description(self) -> str: + return ( + "Step within the denoising loop that denoises the latents for Flux2. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `Flux2DenoiseLoopWrapper`)" + ) + + @property + def inputs(self) -> list[tuple[str, Any]]: + return [ + InputParam("joint_attention_kwargs"), + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The latents to denoise. Shape: (B, seq_len, C)", + ), + InputParam( + "image_latents", + type_hint=torch.Tensor, + description="Packed image latents for conditioning. Shape: (B, img_seq_len, C)", + ), + InputParam( + "image_latent_ids", + type_hint=torch.Tensor, + description="Position IDs for image latents. Shape: (B, img_seq_len, 4)", + ), + InputParam( + "prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Text embeddings from Qwen3", + ), + InputParam( + "negative_prompt_embeds", + required=False, + type_hint=torch.Tensor, + description="Negative text embeddings from Qwen3", + ), + InputParam( + "txt_ids", + required=True, + type_hint=torch.Tensor, + description="4D position IDs for text tokens (T, H, W, L)", + ), + InputParam( + "negative_txt_ids", + required=False, + type_hint=torch.Tensor, + description="4D position IDs for negative text tokens (T, H, W, L)", + ), + InputParam( + "latent_ids", + required=True, + type_hint=torch.Tensor, + description="4D position IDs for latent tokens (T, H, W, L)", + ), + ] + + @torch.no_grad() + def __call__( + self, components: Flux2KleinModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> PipelineState: + latents = block_state.latents + latent_model_input = latents.to(components.transformer.dtype) + img_ids = block_state.latent_ids + + image_latents = getattr(block_state, "image_latents", None) + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1).to(components.transformer.dtype) + image_latent_ids = block_state.image_latent_ids + img_ids = torch.cat([img_ids, image_latent_ids], dim=1) + + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + guider_inputs = { + "encoder_hidden_states": ( + getattr(block_state, "prompt_embeds", None), + getattr(block_state, "negative_prompt_embeds", None), + ), + "txt_ids": ( + getattr(block_state, "txt_ids", None), + getattr(block_state, "negative_txt_ids", None), + ), + } + + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()} + + noise_pred = components.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=None, + img_ids=img_ids, + joint_attention_kwargs=block_state.joint_attention_kwargs, + return_dict=False, + **cond_kwargs, + )[0] + guider_state_batch.noise_pred = noise_pred[:, : latents.size(1)] + components.guider.cleanup_models(components.transformer) + + # perform guidance + block_state.noise_pred = components.guider(guider_state)[0] + + return components, block_state + + +class Flux2LoopAfterDenoiser(ModularPipelineBlocks): + model_name = "flux2" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def description(self) -> str: + return ( + "Step within the denoising loop that updates the latents after denoising. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `Flux2DenoiseLoopWrapper`)" + ) + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")] + + @torch.no_grad() + def __call__(self, components: Flux2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + latents_dtype = block_state.latents.dtype + block_state.latents = components.scheduler.step( + block_state.noise_pred, + t, + block_state.latents, + return_dict=False, + )[0] + + if block_state.latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + block_state.latents = block_state.latents.to(latents_dtype) + + return components, block_state + + +class Flux2DenoiseLoopWrapper(LoopSequentialPipelineBlocks): + model_name = "flux2" + + @property + def description(self) -> str: + return ( + "Pipeline block that iteratively denoises the latents over `timesteps`. " + "The specific steps within each iteration can be customized with `sub_blocks` attribute" + ) + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ComponentSpec("transformer", Flux2Transformer2DModel), + ] + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam( + "timesteps", + required=True, + type_hint=torch.Tensor, + description="The timesteps to use for the denoising process.", + ), + InputParam( + "num_inference_steps", + required=True, + type_hint=int, + description="The number of inference steps to use for the denoising process.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + block_state.num_warmup_steps = max( + len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0 + ) + + with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, block_state = self.loop_step(components, block_state, i=i, t=t) + + if i == len(block_state.timesteps) - 1 or ( + (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0 + ): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + self.set_block_state(state, block_state) + return components, state + + +class Flux2DenoiseStep(Flux2DenoiseLoopWrapper): + block_classes = [Flux2LoopDenoiser, Flux2LoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoises the latents for Flux2. \n" + "Its loop logic is defined in `Flux2DenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `Flux2LoopDenoiser`\n" + " - `Flux2LoopAfterDenoiser`\n" + "This block supports both text-to-image and image-conditioned generation." + ) + + +class Flux2KleinDenoiseStep(Flux2DenoiseLoopWrapper): + block_classes = [Flux2KleinLoopDenoiser, Flux2LoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoises the latents for Flux2. \n" + "Its loop logic is defined in `Flux2DenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `Flux2KleinLoopDenoiser`\n" + " - `Flux2LoopAfterDenoiser`\n" + "This block supports both text-to-image and image-conditioned generation." + ) + + +class Flux2KleinBaseDenoiseStep(Flux2DenoiseLoopWrapper): + block_classes = [Flux2KleinBaseLoopDenoiser, Flux2LoopAfterDenoiser] + block_names = ["denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoises the latents for Flux2. \n" + "Its loop logic is defined in `Flux2DenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `Flux2KleinBaseLoopDenoiser`\n" + " - `Flux2LoopAfterDenoiser`\n" + "This block supports both text-to-image and image-conditioned generation." + ) diff --git a/diffusers/modular_pipelines/flux2/encoders.py b/diffusers/modular_pipelines/flux2/encoders.py new file mode 100644 index 0000000000000000000000000000000000000000..09615c4becb6865ce5e76fd7576c0ee8048bd329 --- /dev/null +++ b/diffusers/modular_pipelines/flux2/encoders.py @@ -0,0 +1,608 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +from transformers import AutoProcessor, Mistral3ForConditionalGeneration, Qwen2TokenizerFast, Qwen3ForCausalLM + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import AutoencoderKLFlux2 +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam +from .modular_pipeline import Flux2KleinModularPipeline, Flux2ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def format_text_input(prompts: list[str], system_message: str = None): + """Format prompts for Mistral3 chat template.""" + cleaned_txt = [prompt.replace("[IMG]", "") for prompt in prompts] + + return [ + [ + { + "role": "system", + "content": [{"type": "text", "text": system_message}], + }, + {"role": "user", "content": [{"type": "text", "text": prompt}]}, + ] + for prompt in cleaned_txt + ] + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.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") + + +class Flux2TextEncoderStep(ModularPipelineBlocks): + model_name = "flux2" + + # fmt: off + DEFAULT_SYSTEM_MESSAGE = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation." + # fmt: on + + @property + def description(self) -> str: + return "Text Encoder step that generates text embeddings using Mistral3 to guide the image generation" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Mistral3ForConditionalGeneration), + ComponentSpec("tokenizer", AutoProcessor), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("prompt"), + InputParam("max_sequence_length", type_hint=int, default=512, required=False), + InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(10, 20, 30), required=False), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "prompt_embeds", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Text embeddings from Mistral3 used to guide the image generation", + ), + ] + + @staticmethod + def check_inputs(block_state): + prompt = block_state.prompt + if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + @staticmethod + def _get_mistral_3_prompt_embeds( + text_encoder: Mistral3ForConditionalGeneration, + tokenizer: AutoProcessor, + prompt: str | list[str], + dtype: torch.dtype | None = None, + device: torch.device | None = None, + max_sequence_length: int = 512, + # fmt: off + system_message: str = "You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object attribution and actions without speculation.", + # fmt: on + hidden_states_layers: tuple[int] = (10, 20, 30), + ): + dtype = text_encoder.dtype if dtype is None else dtype + device = text_encoder.device if device is None else device + + prompt = [prompt] if isinstance(prompt, str) else prompt + + messages_batch = format_text_input(prompts=prompt, system_message=system_message) + + inputs = tokenizer.apply_chat_template( + messages_batch, + add_generation_prompt=False, + tokenize=True, + return_dict=True, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=max_sequence_length, + ) + + input_ids = inputs["input_ids"].to(device) + attention_mask = inputs["attention_mask"].to(device) + + output = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False, + ) + + out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1) + out = out.to(dtype=dtype, device=device) + + batch_size, num_channels, seq_len, hidden_dim = out.shape + prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim) + + return prompt_embeds + + @torch.no_grad() + def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + block_state.device = components._execution_device + + prompt = block_state.prompt + if prompt is None: + prompt = "" + prompt = [prompt] if isinstance(prompt, str) else prompt + + block_state.prompt_embeds = self._get_mistral_3_prompt_embeds( + text_encoder=components.text_encoder, + tokenizer=components.tokenizer, + prompt=prompt, + device=block_state.device, + max_sequence_length=block_state.max_sequence_length, + system_message=self.DEFAULT_SYSTEM_MESSAGE, + hidden_states_layers=block_state.text_encoder_out_layers, + ) + + self.set_block_state(state, block_state) + return components, state + + +class Flux2RemoteTextEncoderStep(ModularPipelineBlocks): + model_name = "flux2" + + REMOTE_URL = "https://remote-text-encoder-flux-2.huggingface.co/predict" + + @property + def description(self) -> str: + return "Text Encoder step that generates text embeddings using a remote API endpoint" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("prompt"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "prompt_embeds", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Text embeddings from remote API used to guide the image generation", + ), + ] + + @staticmethod + def check_inputs(block_state): + prompt = block_state.prompt + if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(block_state.prompt)}") + + @torch.no_grad() + def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: + import io + + import requests + from huggingface_hub import get_token + + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + block_state.device = components._execution_device + + prompt = block_state.prompt + if prompt is None: + prompt = "" + prompt = [prompt] if isinstance(prompt, str) else prompt + + response = requests.post( + self.REMOTE_URL, + json={"prompt": prompt}, + headers={ + "Authorization": f"Bearer {get_token()}", + "Content-Type": "application/json", + }, + ) + response.raise_for_status() + + block_state.prompt_embeds = torch.load(io.BytesIO(response.content), weights_only=True) + block_state.prompt_embeds = block_state.prompt_embeds.to(block_state.device) + + self.set_block_state(state, block_state) + return components, state + + +class Flux2KleinTextEncoderStep(ModularPipelineBlocks): + model_name = "flux2-klein" + + @property + def description(self) -> str: + return "Text Encoder step that generates text embeddings using Qwen3 to guide the image generation" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3ForCausalLM), + ComponentSpec("tokenizer", Qwen2TokenizerFast), + ] + + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ + ConfigSpec(name="is_distilled", default=True), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("prompt"), + InputParam("max_sequence_length", type_hint=int, default=512, required=False), + InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(9, 18, 27), required=False), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "prompt_embeds", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Text embeddings from qwen3 used to guide the image generation", + ), + ] + + @staticmethod + def check_inputs(block_state): + prompt = block_state.prompt + + if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2_klein.Flux2KleinPipeline._get_qwen3_prompt_embeds + def _get_qwen3_prompt_embeds( + text_encoder: Qwen3ForCausalLM, + tokenizer: Qwen2TokenizerFast, + prompt: str | list[str], + dtype: torch.dtype | None = None, + device: torch.device | None = None, + max_sequence_length: int = 512, + hidden_states_layers: list[int] = (9, 18, 27), + ): + dtype = text_encoder.dtype if dtype is None else dtype + device = text_encoder.device if device is None else device + + prompt = [prompt] if isinstance(prompt, str) else prompt + + all_input_ids = [] + all_attention_masks = [] + + for single_prompt in prompt: + messages = [{"role": "user", "content": single_prompt}] + text = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + inputs = tokenizer( + text, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=max_sequence_length, + ) + + all_input_ids.append(inputs["input_ids"]) + all_attention_masks.append(inputs["attention_mask"]) + + input_ids = torch.cat(all_input_ids, dim=0).to(device) + attention_mask = torch.cat(all_attention_masks, dim=0).to(device) + + # Forward pass through the model + output = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False, + ) + + # Only use outputs from intermediate layers and stack them + out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1) + out = out.to(dtype=dtype, device=device) + + batch_size, num_channels, seq_len, hidden_dim = out.shape + prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim) + + return prompt_embeds + + @torch.no_grad() + def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + device = components._execution_device + + prompt = block_state.prompt + if prompt is None: + prompt = "" + prompt = [prompt] if isinstance(prompt, str) else prompt + + block_state.prompt_embeds = self._get_qwen3_prompt_embeds( + text_encoder=components.text_encoder, + tokenizer=components.tokenizer, + prompt=prompt, + device=device, + max_sequence_length=block_state.max_sequence_length, + hidden_states_layers=block_state.text_encoder_out_layers, + ) + + self.set_block_state(state, block_state) + return components, state + + +class Flux2KleinBaseTextEncoderStep(ModularPipelineBlocks): + model_name = "flux2-klein" + + @property + def description(self) -> str: + return "Text Encoder step that generates text embeddings using Qwen3 to guide the image generation" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3ForCausalLM), + ComponentSpec("tokenizer", Qwen2TokenizerFast), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 4.0}), + default_creation_method="from_config", + ), + ] + + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ + ConfigSpec(name="is_distilled", default=False), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("prompt"), + InputParam("max_sequence_length", type_hint=int, default=512, required=False), + InputParam("text_encoder_out_layers", type_hint=tuple[int], default=(9, 18, 27), required=False), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "prompt_embeds", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Text embeddings from qwen3 used to guide the image generation", + ), + OutputParam( + "negative_prompt_embeds", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Negative text embeddings from qwen3 used to guide the image generation", + ), + ] + + @staticmethod + def check_inputs(block_state): + prompt = block_state.prompt + + if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2_klein.Flux2KleinPipeline._get_qwen3_prompt_embeds + def _get_qwen3_prompt_embeds( + text_encoder: Qwen3ForCausalLM, + tokenizer: Qwen2TokenizerFast, + prompt: str | list[str], + dtype: torch.dtype | None = None, + device: torch.device | None = None, + max_sequence_length: int = 512, + hidden_states_layers: list[int] = (9, 18, 27), + ): + dtype = text_encoder.dtype if dtype is None else dtype + device = text_encoder.device if device is None else device + + prompt = [prompt] if isinstance(prompt, str) else prompt + + all_input_ids = [] + all_attention_masks = [] + + for single_prompt in prompt: + messages = [{"role": "user", "content": single_prompt}] + text = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + inputs = tokenizer( + text, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=max_sequence_length, + ) + + all_input_ids.append(inputs["input_ids"]) + all_attention_masks.append(inputs["attention_mask"]) + + input_ids = torch.cat(all_input_ids, dim=0).to(device) + attention_mask = torch.cat(all_attention_masks, dim=0).to(device) + + # Forward pass through the model + output = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False, + ) + + # Only use outputs from intermediate layers and stack them + out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1) + out = out.to(dtype=dtype, device=device) + + batch_size, num_channels, seq_len, hidden_dim = out.shape + prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim) + + return prompt_embeds + + @torch.no_grad() + def __call__(self, components: Flux2KleinModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + device = components._execution_device + + prompt = block_state.prompt + if prompt is None: + prompt = "" + prompt = [prompt] if isinstance(prompt, str) else prompt + + block_state.prompt_embeds = self._get_qwen3_prompt_embeds( + text_encoder=components.text_encoder, + tokenizer=components.tokenizer, + prompt=prompt, + device=device, + max_sequence_length=block_state.max_sequence_length, + hidden_states_layers=block_state.text_encoder_out_layers, + ) + + if components.requires_unconditional_embeds: + negative_prompt = [""] * len(prompt) + block_state.negative_prompt_embeds = self._get_qwen3_prompt_embeds( + text_encoder=components.text_encoder, + tokenizer=components.tokenizer, + prompt=negative_prompt, + device=device, + max_sequence_length=block_state.max_sequence_length, + hidden_states_layers=block_state.text_encoder_out_layers, + ) + else: + block_state.negative_prompt_embeds = None + + self.set_block_state(state, block_state) + return components, state + + +class Flux2VaeEncoderStep(ModularPipelineBlocks): + model_name = "flux2" + + @property + def description(self) -> str: + return "VAE Encoder step that encodes preprocessed images into latent representations for Flux2." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("vae", AutoencoderKLFlux2)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("condition_images", type_hint=list[torch.Tensor]), + InputParam("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "image_latents", + type_hint=list[torch.Tensor], + description="List of latent representations for each reference image", + ), + ] + + @staticmethod + def _patchify_latents(latents): + """Convert latents to patchified format for Flux2.""" + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 3, 5, 2, 4) + latents = latents.reshape(batch_size, num_channels_latents * 4, height // 2, width // 2) + return latents + + def _encode_vae_image(self, vae: AutoencoderKLFlux2, image: torch.Tensor, generator: torch.Generator): + """Encode a single image using Flux2 VAE with batch norm normalization.""" + if image.ndim != 4: + raise ValueError(f"Expected image dims 4, got {image.ndim}.") + + image_latents = retrieve_latents(vae.encode(image), generator=generator, sample_mode="argmax") + image_latents = self._patchify_latents(image_latents) + + latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype) + latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps) + latents_bn_std = latents_bn_std.to(image_latents.device, image_latents.dtype) + image_latents = (image_latents - latents_bn_mean) / latents_bn_std + + return image_latents + + @torch.no_grad() + def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + condition_images = block_state.condition_images + + if condition_images is None: + return components, state + + device = components._execution_device + dtype = components.vae.dtype + + image_latents = [] + for image in condition_images: + image = image.to(device=device, dtype=dtype) + latent = self._encode_vae_image( + vae=components.vae, + image=image, + generator=block_state.generator, + ) + image_latents.append(latent) + + block_state.image_latents = image_latents + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/flux2/inputs.py b/diffusers/modular_pipelines/flux2/inputs.py new file mode 100644 index 0000000000000000000000000000000000000000..6bfe6aec97fd00621ebc695d8eaa896740c29d34 --- /dev/null +++ b/diffusers/modular_pipelines/flux2/inputs.py @@ -0,0 +1,242 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...configuration_utils import FrozenDict +from ...pipelines.flux2.image_processor import Flux2ImageProcessor +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Flux2ModularPipeline + + +logger = logging.get_logger(__name__) + + +class Flux2TextInputStep(ModularPipelineBlocks): + model_name = "flux2" + + @property + def description(self) -> str: + return ( + "This step:\n" + " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n" + " 2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt)" + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("num_images_per_prompt", default=1), + InputParam( + "prompt_embeds", + required=True, + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Pre-generated text embeddings. Can be generated from text_encoder step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[str]: + return [ + OutputParam( + "batch_size", + type_hint=int, + description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt", + ), + OutputParam( + "dtype", + type_hint=torch.dtype, + description="Data type of model tensor inputs (determined by `prompt_embeds`)", + ), + OutputParam( + "prompt_embeds", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Text embeddings used to guide the image generation", + ), + ] + + @torch.no_grad() + def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + block_state.batch_size = block_state.prompt_embeds.shape[0] + block_state.dtype = block_state.prompt_embeds.dtype + + _, seq_len, _ = block_state.prompt_embeds.shape + block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1) + block_state.prompt_embeds = block_state.prompt_embeds.view( + block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1 + ) + + self.set_block_state(state, block_state) + return components, state + + +class Flux2KleinBaseTextInputStep(ModularPipelineBlocks): + model_name = "flux2-klein" + + @property + def description(self) -> str: + return ( + "This step:\n" + " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n" + " 2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt)" + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("num_images_per_prompt", default=1), + InputParam( + "prompt_embeds", + required=True, + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Pre-generated text embeddings. Can be generated from text_encoder step.", + ), + InputParam( + "negative_prompt_embeds", + required=False, + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="Pre-generated negative text embeddings. Can be generated from text_encoder step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[str]: + return [ + OutputParam( + "batch_size", + type_hint=int, + description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt", + ), + OutputParam( + "dtype", + type_hint=torch.dtype, + description="Data type of model tensor inputs (determined by `prompt_embeds`)", + ), + OutputParam( + "prompt_embeds", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Text embeddings used to guide the image generation", + ), + OutputParam( + "negative_prompt_embeds", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Negative text embeddings used to guide the image generation", + ), + ] + + @torch.no_grad() + def __call__(self, components: Flux2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + block_state.batch_size = block_state.prompt_embeds.shape[0] + block_state.dtype = block_state.prompt_embeds.dtype + + _, seq_len, _ = block_state.prompt_embeds.shape + block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1) + block_state.prompt_embeds = block_state.prompt_embeds.view( + block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1 + ) + + if block_state.negative_prompt_embeds is not None: + _, seq_len, _ = block_state.negative_prompt_embeds.shape + block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat( + 1, block_state.num_images_per_prompt, 1 + ) + block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view( + block_state.batch_size * block_state.num_images_per_prompt, seq_len, -1 + ) + + self.set_block_state(state, block_state) + return components, state + + +class Flux2ProcessImagesInputStep(ModularPipelineBlocks): + model_name = "flux2" + + @property + def description(self) -> str: + return "Image preprocess step for Flux2. Validates and preprocesses reference images." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + Flux2ImageProcessor, + config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 32}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("image"), + InputParam("height"), + InputParam("width"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam(name="condition_images", type_hint=list[torch.Tensor])] + + @torch.no_grad() + def __call__(self, components: Flux2ModularPipeline, state: PipelineState): + block_state = self.get_block_state(state) + images = block_state.image + + if images is None: + block_state.condition_images = None + self.set_block_state(state, block_state) + return components, state + + if not isinstance(images, list): + images = [images] + + condition_images = [] + for img in images: + components.image_processor.check_image_input(img) + + image_width, image_height = img.size + if image_width * image_height > 1024 * 1024: + img = components.image_processor._resize_to_target_area(img, 1024 * 1024) + image_width, image_height = img.size + + multiple_of = components.vae_scale_factor * 2 + image_width = (image_width // multiple_of) * multiple_of + image_height = (image_height // multiple_of) * multiple_of + condition_img = components.image_processor.preprocess( + img, height=image_height, width=image_width, resize_mode="crop" + ) + condition_images.append(condition_img) + + if block_state.height is None: + block_state.height = image_height + if block_state.width is None: + block_state.width = image_width + + block_state.condition_images = condition_images + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py b/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py new file mode 100644 index 0000000000000000000000000000000000000000..2bbb7975a9834ff84f09063cb587adf985cda8e9 --- /dev/null +++ b/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py @@ -0,0 +1,356 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + Flux2PrepareGuidanceStep, + Flux2PrepareImageLatentsStep, + Flux2PrepareLatentsStep, + Flux2RoPEInputsStep, + Flux2SetTimestepsStep, +) +from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep +from .denoise import Flux2DenoiseStep +from .encoders import ( + Flux2TextEncoderStep, + Flux2VaeEncoderStep, +) +from .inputs import ( + Flux2ProcessImagesInputStep, + Flux2TextInputStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# auto_docstring +class Flux2VaeEncoderSequentialStep(SequentialPipelineBlocks): + """ + VAE encoder step that preprocesses, encodes, and prepares image latents for Flux2 conditioning. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + model_name = "flux2" + + block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "VAE encoder step that preprocesses, encodes, and prepares image latents for Flux2 conditioning." + + +# auto_docstring +class Flux2AutoVaeEncoderStep(AutoPipelineBlocks): + """ + VAE encoder step that encodes the image inputs into their latent representations. + This is an auto pipeline block that works for image conditioning tasks. + - `Flux2VaeEncoderSequentialStep` is used when `image` is provided. + - If `image` is not provided, step will be skipped. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + block_classes = [Flux2VaeEncoderSequentialStep] + block_names = ["img_conditioning"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "VAE encoder step that encodes the image inputs into their latent representations.\n" + "This is an auto pipeline block that works for image conditioning tasks.\n" + " - `Flux2VaeEncoderSequentialStep` is used when `image` is provided.\n" + " - If `image` is not provided, step will be skipped." + ) + + +Flux2CoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2TextInputStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_guidance", Flux2PrepareGuidanceStep()), + ("prepare_rope_inputs", Flux2RoPEInputsStep()), + ("denoise", Flux2DenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2CoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-dev. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 4.0): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latents (`Tensor`, *optional*): + Packed image latents for conditioning. Shape: (B, img_seq_len, C) + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2" + + block_classes = Flux2CoreDenoiseBlocks.values() + block_names = Flux2CoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-dev." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +Flux2ImageConditionedCoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2TextInputStep()), + ("prepare_image_latents", Flux2PrepareImageLatentsStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_guidance", Flux2PrepareGuidanceStep()), + ("prepare_rope_inputs", Flux2RoPEInputsStep()), + ("denoise", Flux2DenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2ImageConditionedCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-dev with image conditioning. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + image_latents (`list`, *optional*): + TODO: Add description. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 4.0): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2" + + block_classes = Flux2ImageConditionedCoreDenoiseBlocks.values() + block_names = Flux2ImageConditionedCoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-dev with image conditioning." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +class Flux2AutoCoreDenoiseStep(AutoPipelineBlocks): + model_name = "flux2" + + block_classes = [Flux2ImageConditionedCoreDenoiseStep, Flux2CoreDenoiseStep] + block_names = ["image_conditioned", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Auto core denoise step that performs the denoising process for Flux2-dev." + "This is an auto pipeline block that works for text-to-image and image-conditioned generation." + " - `Flux2CoreDenoiseStep` is used for text-to-image generation.\n" + " - `Flux2ImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n" + ) + + +AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", Flux2TextEncoderStep()), + ("vae_encoder", Flux2AutoVaeEncoderStep()), + ("denoise", Flux2AutoCoreDenoiseStep()), + ("decode", Flux2DecodeStep()), + ] +) + + +# auto_docstring +class Flux2AutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-image and image-conditioned generation using Flux2. + + Supported workflows: + - `text2image`: requires `prompt` + - `image_conditioned`: requires `image`, `prompt` + + Components: + text_encoder (`Mistral3ForConditionalGeneration`) tokenizer (`AutoProcessor`) image_processor + (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer + (`Flux2Transformer2DModel`) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`int`, *optional*, defaults to 512): + TODO: Add description. + text_encoder_out_layers (`tuple`, *optional*, defaults to (10, 20, 30)): + TODO: Add description. + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + image_latents (`list`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + timesteps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + guidance_scale (`None`, *optional*, defaults to 4.0): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + output_type (`None`, *optional*, defaults to pil): + TODO: Add description. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "flux2" + + block_classes = AUTO_BLOCKS.values() + block_names = AUTO_BLOCKS.keys() + _workflow_map = { + "text2image": {"prompt": True}, + "image_conditioned": {"image": True, "prompt": True}, + } + + @property + def description(self): + return "Auto Modular pipeline for text-to-image and image-conditioned generation using Flux2." + + @property + def outputs(self): + return [ + OutputParam.template("images"), + ] diff --git a/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py b/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py new file mode 100644 index 0000000000000000000000000000000000000000..689cf808c4ba93ac373226e375070bc6ffa5dd64 --- /dev/null +++ b/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py @@ -0,0 +1,399 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + Flux2PrepareImageLatentsStep, + Flux2PrepareLatentsStep, + Flux2RoPEInputsStep, + Flux2SetTimestepsStep, +) +from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep +from .denoise import Flux2KleinDenoiseStep +from .encoders import ( + Flux2KleinTextEncoderStep, + Flux2VaeEncoderStep, +) +from .inputs import ( + Flux2ProcessImagesInputStep, + Flux2TextInputStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +################ +# VAE encoder +################ + + +# auto_docstring +class Flux2KleinVaeEncoderSequentialStep(SequentialPipelineBlocks): + """ + VAE encoder step that preprocesses and encodes the image inputs into their latent representations. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + model_name = "flux2-klein" + + block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "VAE encoder step that preprocesses and encodes the image inputs into their latent representations." + + +# auto_docstring +class Flux2KleinAutoVaeEncoderStep(AutoPipelineBlocks): + """ + VAE encoder step that encodes the image inputs into their latent representations. + This is an auto pipeline block that works for image conditioning tasks. + - `Flux2KleinVaeEncoderSequentialStep` is used when `image` is provided. + - If `image` is not provided, step will be skipped. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + model_name = "flux2-klein" + + block_classes = [Flux2KleinVaeEncoderSequentialStep] + block_names = ["img_conditioning"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "VAE encoder step that encodes the image inputs into their latent representations.\n" + "This is an auto pipeline block that works for image conditioning tasks.\n" + " - `Flux2KleinVaeEncoderSequentialStep` is used when `image` is provided.\n" + " - If `image` is not provided, step will be skipped." + ) + + +### +### Core denoise +### + +Flux2KleinCoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2TextInputStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_rope_inputs", Flux2RoPEInputsStep()), + ("denoise", Flux2KleinDenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2KleinCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-Klein (distilled model), for text-to-image + generation. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latents (`Tensor`, *optional*): + Packed image latents for conditioning. Shape: (B, img_seq_len, C) + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + + block_classes = Flux2KleinCoreDenoiseBlocks.values() + block_names = Flux2KleinCoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-Klein (distilled model), for text-to-image generation." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +Flux2KleinImageConditionedCoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2TextInputStep()), + ("prepare_image_latents", Flux2PrepareImageLatentsStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_rope_inputs", Flux2RoPEInputsStep()), + ("denoise", Flux2KleinDenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2KleinImageConditionedCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-Klein (distilled model) with image conditioning. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + image_latents (`list`, *optional*): + TODO: Add description. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + + block_classes = Flux2KleinImageConditionedCoreDenoiseBlocks.values() + block_names = Flux2KleinImageConditionedCoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-Klein (distilled model) with image conditioning." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# auto_docstring +class Flux2KleinAutoCoreDenoiseStep(AutoPipelineBlocks): + """ + Auto core denoise step that performs the denoising process for Flux2-Klein. + This is an auto pipeline block that works for text-to-image and image-conditioned generation. + - `Flux2KleinCoreDenoiseStep` is used for text-to-image generation. + - `Flux2KleinImageConditionedCoreDenoiseStep` is used for image-conditioned generation. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + image_latents (`list`, *optional*): + TODO: Add description. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + timesteps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + block_classes = [Flux2KleinImageConditionedCoreDenoiseStep, Flux2KleinCoreDenoiseStep] + block_names = ["image_conditioned", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Auto core denoise step that performs the denoising process for Flux2-Klein.\n" + "This is an auto pipeline block that works for text-to-image and image-conditioned generation.\n" + " - `Flux2KleinCoreDenoiseStep` is used for text-to-image generation.\n" + " - `Flux2KleinImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n" + ) + + +### +### Auto blocks +### + + +# auto_docstring +class Flux2KleinAutoBlocks(SequentialPipelineBlocks): + """ + Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein. + + Supported workflows: + - `text2image`: requires `prompt` + - `image_conditioned`: requires `image`, `prompt` + + Components: + text_encoder (`Qwen3ForCausalLM`) tokenizer (`Qwen2Tokenizer`) image_processor (`Flux2ImageProcessor`) vae + (`AutoencoderKLFlux2`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Configs: + is_distilled (default: True) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`int`, *optional*, defaults to 512): + TODO: Add description. + text_encoder_out_layers (`tuple`, *optional*, defaults to (9, 18, 27)): + TODO: Add description. + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + image_latents (`list`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + timesteps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + output_type (`None`, *optional*, defaults to pil): + TODO: Add description. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "flux2-klein" + block_classes = [ + Flux2KleinTextEncoderStep(), + Flux2KleinAutoVaeEncoderStep(), + Flux2KleinAutoCoreDenoiseStep(), + Flux2DecodeStep(), + ] + block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] + _workflow_map = { + "text2image": {"prompt": True}, + "image_conditioned": {"image": True, "prompt": True}, + } + + @property + def description(self): + return "Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein." + + @property + def outputs(self): + return [ + OutputParam.template("images"), + ] diff --git a/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py b/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py new file mode 100644 index 0000000000000000000000000000000000000000..f3108bdadeacdb20c3fdbf22ea3ae215018c9c8c --- /dev/null +++ b/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py @@ -0,0 +1,413 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + Flux2KleinBaseRoPEInputsStep, + Flux2PrepareImageLatentsStep, + Flux2PrepareLatentsStep, + Flux2SetTimestepsStep, +) +from .decoders import Flux2DecodeStep, Flux2UnpackLatentsStep +from .denoise import Flux2KleinBaseDenoiseStep +from .encoders import ( + Flux2KleinBaseTextEncoderStep, + Flux2VaeEncoderStep, +) +from .inputs import ( + Flux2KleinBaseTextInputStep, + Flux2ProcessImagesInputStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +################ +# VAE encoder +################ + + +# auto_docstring +class Flux2KleinBaseVaeEncoderSequentialStep(SequentialPipelineBlocks): + """ + VAE encoder step that preprocesses and encodes the image inputs into their latent representations. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + model_name = "flux2" + + block_classes = [Flux2ProcessImagesInputStep(), Flux2VaeEncoderStep()] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "VAE encoder step that preprocesses and encodes the image inputs into their latent representations." + + +# auto_docstring +class Flux2KleinBaseAutoVaeEncoderStep(AutoPipelineBlocks): + """ + VAE encoder step that encodes the image inputs into their latent representations. + This is an auto pipeline block that works for image conditioning tasks. + - `Flux2KleinBaseVaeEncoderSequentialStep` is used when `image` is provided. + - If `image` is not provided, step will be skipped. + + Components: + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + + Outputs: + condition_images (`list`): + TODO: Add description. + image_latents (`list`): + List of latent representations for each reference image + """ + + block_classes = [Flux2KleinBaseVaeEncoderSequentialStep] + block_names = ["img_conditioning"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "VAE encoder step that encodes the image inputs into their latent representations.\n" + "This is an auto pipeline block that works for image conditioning tasks.\n" + " - `Flux2KleinBaseVaeEncoderSequentialStep` is used when `image` is provided.\n" + " - If `image` is not provided, step will be skipped." + ) + + +### +### Core denoise +### + +Flux2KleinBaseCoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2KleinBaseTextInputStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_rope_inputs", Flux2KleinBaseRoPEInputsStep()), + ("denoise", Flux2KleinBaseDenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2KleinBaseCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-Klein (base model), for text-to-image generation. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider + (`ClassifierFreeGuidance`) + + Configs: + is_distilled (default: False) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latents (`Tensor`, *optional*): + Packed image latents for conditioning. Shape: (B, img_seq_len, C) + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + block_classes = Flux2KleinBaseCoreDenoiseBlocks.values() + block_names = Flux2KleinBaseCoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-Klein (base model), for text-to-image generation." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +Flux2KleinBaseImageConditionedCoreDenoiseBlocks = InsertableDict( + [ + ("input", Flux2KleinBaseTextInputStep()), + ("prepare_latents", Flux2PrepareLatentsStep()), + ("prepare_image_latents", Flux2PrepareImageLatentsStep()), + ("set_timesteps", Flux2SetTimestepsStep()), + ("prepare_rope_inputs", Flux2KleinBaseRoPEInputsStep()), + ("denoise", Flux2KleinBaseDenoiseStep()), + ("after_denoise", Flux2UnpackLatentsStep()), + ] +) + + +# auto_docstring +class Flux2KleinBaseImageConditionedCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that performs the denoising process for Flux2-Klein (base model) with image conditioning. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider + (`ClassifierFreeGuidance`) + + Configs: + is_distilled (default: False) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + image_latents (`list`, *optional*): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + block_classes = Flux2KleinBaseImageConditionedCoreDenoiseBlocks.values() + block_names = Flux2KleinBaseImageConditionedCoreDenoiseBlocks.keys() + + @property + def description(self): + return "Core denoise step that performs the denoising process for Flux2-Klein (base model) with image conditioning." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +# auto_docstring +class Flux2KleinBaseAutoCoreDenoiseStep(AutoPipelineBlocks): + """ + Auto core denoise step that performs the denoising process for Flux2-Klein (base model). + This is an auto pipeline block that works for text-to-image and image-conditioned generation. + - `Flux2KleinBaseCoreDenoiseStep` is used for text-to-image generation. + - `Flux2KleinBaseImageConditionedCoreDenoiseStep` is used for image-conditioned generation. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) guider + (`ClassifierFreeGuidance`) + + Configs: + is_distilled (default: False) + + Inputs: + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + image_latents (`list`, *optional*): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + timesteps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "flux2-klein" + block_classes = [Flux2KleinBaseImageConditionedCoreDenoiseStep, Flux2KleinBaseCoreDenoiseStep] + block_names = ["image_conditioned", "text2image"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Auto core denoise step that performs the denoising process for Flux2-Klein (base model).\n" + "This is an auto pipeline block that works for text-to-image and image-conditioned generation.\n" + " - `Flux2KleinBaseCoreDenoiseStep` is used for text-to-image generation.\n" + " - `Flux2KleinBaseImageConditionedCoreDenoiseStep` is used for image-conditioned generation.\n" + ) + + +### +### Auto blocks +### + + +# auto_docstring +class Flux2KleinBaseAutoBlocks(SequentialPipelineBlocks): + """ + Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein (base model). + + Supported workflows: + - `text2image`: requires `prompt` + - `image_conditioned`: requires `image`, `prompt` + + Components: + text_encoder (`Qwen3ForCausalLM`) tokenizer (`Qwen2Tokenizer`) guider (`ClassifierFreeGuidance`) + image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) scheduler + (`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`) + + Configs: + is_distilled (default: False) + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`int`, *optional*, defaults to 512): + TODO: Add description. + text_encoder_out_layers (`tuple`, *optional*, defaults to (9, 18, 27)): + TODO: Add description. + image (`None`, *optional*): + TODO: Add description. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_images_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + latents (`Tensor | NoneType`): + TODO: Add description. + image_latents (`list`, *optional*): + TODO: Add description. + num_inference_steps (`None`): + TODO: Add description. + timesteps (`None`): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + joint_attention_kwargs (`None`, *optional*): + TODO: Add description. + image_latent_ids (`Tensor`, *optional*): + Position IDs for image latents. Shape: (B, img_seq_len, 4) + output_type (`None`, *optional*, defaults to pil): + TODO: Add description. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "flux2-klein" + block_classes = [ + Flux2KleinBaseTextEncoderStep(), + Flux2KleinBaseAutoVaeEncoderStep(), + Flux2KleinBaseAutoCoreDenoiseStep(), + Flux2DecodeStep(), + ] + block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] + _workflow_map = { + "text2image": {"prompt": True}, + "image_conditioned": {"image": True, "prompt": True}, + } + + @property + def description(self): + return "Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein (base model)." + + @property + def outputs(self): + return [ + OutputParam.template("images"), + ] diff --git a/diffusers/modular_pipelines/flux2/modular_pipeline.py b/diffusers/modular_pipelines/flux2/modular_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..ed070206c319625dba440dea9baf46037cba91c7 --- /dev/null +++ b/diffusers/modular_pipelines/flux2/modular_pipeline.py @@ -0,0 +1,99 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...loaders import Flux2LoraLoaderMixin +from ...utils import logging +from ..modular_pipeline import ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class Flux2ModularPipeline(ModularPipeline, Flux2LoraLoaderMixin): + """ + A ModularPipeline for Flux2. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "Flux2AutoBlocks" + + @property + def default_height(self): + return self.default_sample_size * self.vae_scale_factor + + @property + def default_width(self): + return self.default_sample_size * self.vae_scale_factor + + @property + def default_sample_size(self): + return 128 + + @property + def vae_scale_factor(self): + vae_scale_factor = 8 + if getattr(self, "vae", None) is not None: + vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + return vae_scale_factor + + @property + def num_channels_latents(self): + num_channels_latents = 32 + if getattr(self, "transformer", None): + num_channels_latents = self.transformer.config.in_channels // 4 + return num_channels_latents + + +class Flux2KleinModularPipeline(Flux2ModularPipeline): + """ + A ModularPipeline for Flux2-Klein (distilled model). + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "Flux2KleinAutoBlocks" + + @property + def requires_unconditional_embeds(self): + if hasattr(self.config, "is_distilled") and self.config.is_distilled: + return False + + requires_unconditional_embeds = False + if hasattr(self, "guider") and self.guider is not None: + requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1 + + return requires_unconditional_embeds + + +class Flux2KleinBaseModularPipeline(Flux2ModularPipeline): + """ + A ModularPipeline for Flux2-Klein (base model). + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "Flux2KleinBaseAutoBlocks" + + @property + def requires_unconditional_embeds(self): + if hasattr(self.config, "is_distilled") and self.config.is_distilled: + return False + + requires_unconditional_embeds = False + if hasattr(self, "guider") and self.guider is not None: + requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1 + + return requires_unconditional_embeds diff --git a/diffusers/modular_pipelines/helios/__init__.py b/diffusers/modular_pipelines/helios/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..26551399a3e81881ac30ae9d995fe352d3afd3da --- /dev/null +++ b/diffusers/modular_pipelines/helios/__init__.py @@ -0,0 +1,59 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_helios"] = ["HeliosAutoBlocks"] + _import_structure["modular_blocks_helios_pyramid"] = ["HeliosPyramidAutoBlocks"] + _import_structure["modular_blocks_helios_pyramid_distilled"] = ["HeliosPyramidDistilledAutoBlocks"] + _import_structure["modular_pipeline"] = [ + "HeliosModularPipeline", + "HeliosPyramidDistilledModularPipeline", + "HeliosPyramidModularPipeline", + ] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_helios import HeliosAutoBlocks + from .modular_blocks_helios_pyramid import HeliosPyramidAutoBlocks + from .modular_blocks_helios_pyramid_distilled import HeliosPyramidDistilledAutoBlocks + from .modular_pipeline import ( + HeliosModularPipeline, + HeliosPyramidDistilledModularPipeline, + HeliosPyramidModularPipeline, + ) +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/diffusers/modular_pipelines/helios/before_denoise.py b/diffusers/modular_pipelines/helios/before_denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..64407db63cca0e0b0e7326e17c32c8613bc886e8 --- /dev/null +++ b/diffusers/modular_pipelines/helios/before_denoise.py @@ -0,0 +1,836 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import torch + +from ...models import HeliosTransformer3DModel +from ...schedulers import HeliosScheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import HeliosModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +class HeliosTextInputStep(ModularPipelineBlocks): + model_name = "helios" + + @property + def description(self) -> str: + return ( + "Input processing step that:\n" + " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n" + " 2. Adjusts input tensor shapes based on `batch_size` (number of prompts) and `num_videos_per_prompt`\n\n" + "All input tensors are expected to have either batch_size=1 or match the batch_size\n" + "of prompt_embeds. The tensors will be duplicated across the batch dimension to\n" + "have a final batch_size of batch_size * num_videos_per_prompt." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "num_videos_per_prompt", + default=1, + type_hint=int, + description="Number of videos to generate per prompt.", + ), + InputParam.template("prompt_embeds"), + InputParam.template("negative_prompt_embeds"), + ] + + @property + def intermediate_outputs(self) -> list[str]: + return [ + OutputParam( + "batch_size", + type_hint=int, + description="Number of prompts, the final batch size of model inputs should be batch_size * num_videos_per_prompt", + ), + OutputParam( + "dtype", + type_hint=torch.dtype, + description="Data type of model tensor inputs (determined by `prompt_embeds.dtype`)", + ), + ] + + def check_inputs(self, components, block_state): + if block_state.prompt_embeds is not None and block_state.negative_prompt_embeds is not None: + if block_state.prompt_embeds.shape != block_state.negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {block_state.prompt_embeds.shape} != `negative_prompt_embeds`" + f" {block_state.negative_prompt_embeds.shape}." + ) + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(components, block_state) + + block_state.batch_size = block_state.prompt_embeds.shape[0] + block_state.dtype = block_state.prompt_embeds.dtype + + _, seq_len, _ = block_state.prompt_embeds.shape + block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_videos_per_prompt, 1) + block_state.prompt_embeds = block_state.prompt_embeds.view( + block_state.batch_size * block_state.num_videos_per_prompt, seq_len, -1 + ) + + if block_state.negative_prompt_embeds is not None: + _, seq_len, _ = block_state.negative_prompt_embeds.shape + block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat( + 1, block_state.num_videos_per_prompt, 1 + ) + block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view( + block_state.batch_size * block_state.num_videos_per_prompt, seq_len, -1 + ) + + self.set_block_state(state, block_state) + + return components, state + + +# Copied from diffusers.modular_pipelines.wan.before_denoise.repeat_tensor_to_batch_size +def repeat_tensor_to_batch_size( + input_name: str, + input_tensor: torch.Tensor, + batch_size: int, + num_videos_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_videos_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_videos_per_prompt) times + - If batch size equals batch_size: repeat each element num_videos_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_videos_per_prompt (int, optional): Number of videos to generate per prompt. Defaults to 1. + + Returns: + torch.Tensor: The repeated tensor with final batch size (batch_size * num_videos_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_videos_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_videos_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_videos_per_prompt + elif input_tensor.shape[0] == batch_size: + repeat_by = num_videos_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_videos_per_prompt + input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0) + + return input_tensor + + +# Copied from diffusers.modular_pipelines.wan.before_denoise.calculate_dimension_from_latents +def calculate_dimension_from_latents( + latents: torch.Tensor, vae_scale_factor_temporal: int, vae_scale_factor_spatial: int +) -> tuple[int, int]: + """Calculate image dimensions from latent tensor dimensions. + + This function converts latent temporal and spatial dimensions to image temporal and spatial dimensions by + multiplying the latent num_frames/height/width by the VAE scale factor. + + Args: + latents (torch.Tensor): The latent tensor. Must have 4 or 5 dimensions. + Expected shapes: [batch, channels, height, width] or [batch, channels, frames, height, width] + vae_scale_factor_temporal (int): The scale factor used by the VAE to compress temporal dimension. + Typically 4 for most VAEs (video is 4x larger than latents in temporal dimension) + vae_scale_factor_spatial (int): The scale factor used by the VAE to compress spatial dimension. + Typically 8 for most VAEs (image is 8x larger than latents in each dimension) + + Returns: + tuple[int, int]: The calculated image dimensions as (height, width) + + Raises: + ValueError: If latents tensor doesn't have 4 or 5 dimensions + + """ + if latents.ndim != 5: + raise ValueError(f"latents must have 5 dimensions, but got {latents.ndim}") + + _, _, num_latent_frames, latent_height, latent_width = latents.shape + + num_frames = (num_latent_frames - 1) * vae_scale_factor_temporal + 1 + height = latent_height * vae_scale_factor_spatial + width = latent_width * vae_scale_factor_spatial + + return num_frames, height, width + + +class HeliosAdditionalInputsStep(ModularPipelineBlocks): + """Configurable step that standardizes inputs for the denoising step. + + This step handles: + 1. For encoded image latents: Computes height/width from latents and expands batch size + 2. For additional_batch_inputs: Expands batch dimensions to match final batch size + """ + + model_name = "helios" + + def __init__( + self, + image_latent_inputs: list[InputParam] | None = None, + additional_batch_inputs: list[InputParam] | None = None, + ): + if image_latent_inputs is None: + image_latent_inputs = [InputParam.template("image_latents")] + if additional_batch_inputs is None: + additional_batch_inputs = [] + + if not isinstance(image_latent_inputs, list): + raise ValueError(f"image_latent_inputs must be a list, but got {type(image_latent_inputs)}") + else: + for input_param in image_latent_inputs: + if not isinstance(input_param, InputParam): + raise ValueError(f"image_latent_inputs must be a list of InputParam, but got {type(input_param)}") + + if not isinstance(additional_batch_inputs, list): + raise ValueError(f"additional_batch_inputs must be a list, but got {type(additional_batch_inputs)}") + else: + for input_param in additional_batch_inputs: + if not isinstance(input_param, InputParam): + raise ValueError( + f"additional_batch_inputs must be a list of InputParam, but got {type(input_param)}" + ) + + self._image_latent_inputs = image_latent_inputs + self._additional_batch_inputs = additional_batch_inputs + super().__init__() + + @property + def description(self) -> str: + summary_section = ( + "Input processing step that:\n" + " 1. For image latent inputs: Computes height/width from latents and expands batch size\n" + " 2. For additional batch inputs: Expands batch dimensions to match final batch size" + ) + + inputs_info = "" + if self._image_latent_inputs or self._additional_batch_inputs: + inputs_info = "\n\nConfigured inputs:" + if self._image_latent_inputs: + inputs_info += f"\n - Image latent inputs: {[p.name for p in self._image_latent_inputs]}" + if self._additional_batch_inputs: + inputs_info += f"\n - Additional batch inputs: {[p.name for p in self._additional_batch_inputs]}" + + placement_section = "\n\nThis block should be placed after the encoder steps and the text input step." + + return summary_section + inputs_info + placement_section + + @property + def inputs(self) -> list[InputParam]: + inputs = [ + InputParam(name="num_videos_per_prompt", default=1), + InputParam(name="batch_size", required=True), + ] + inputs += self._image_latent_inputs + self._additional_batch_inputs + + return inputs + + @property + def intermediate_outputs(self) -> list[OutputParam]: + outputs = [ + OutputParam("height", type_hint=int), + OutputParam("width", type_hint=int), + ] + + for input_param in self._image_latent_inputs: + outputs.append(OutputParam(input_param.name, type_hint=torch.Tensor)) + + for input_param in self._additional_batch_inputs: + outputs.append(OutputParam(input_param.name, type_hint=torch.Tensor)) + + return outputs + + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + for input_param in self._image_latent_inputs: + image_latent_tensor = getattr(block_state, input_param.name) + if image_latent_tensor is None: + continue + + # Calculate height/width from latents + _, height, width = calculate_dimension_from_latents( + image_latent_tensor, components.vae_scale_factor_temporal, components.vae_scale_factor_spatial + ) + block_state.height = height + block_state.width = width + + # Expand batch size + image_latent_tensor = repeat_tensor_to_batch_size( + input_name=input_param.name, + input_tensor=image_latent_tensor, + num_videos_per_prompt=block_state.num_videos_per_prompt, + batch_size=block_state.batch_size, + ) + + setattr(block_state, input_param.name, image_latent_tensor) + + for input_param in self._additional_batch_inputs: + input_tensor = getattr(block_state, input_param.name) + if input_tensor is None: + continue + + input_tensor = repeat_tensor_to_batch_size( + input_name=input_param.name, + input_tensor=input_tensor, + num_videos_per_prompt=block_state.num_videos_per_prompt, + batch_size=block_state.batch_size, + ) + + setattr(block_state, input_param.name, input_tensor) + + self.set_block_state(state, block_state) + return components, state + + +class HeliosAddNoiseToImageLatentsStep(ModularPipelineBlocks): + """Adds noise to image_latents and fake_image_latents for I2V conditioning. + + Applies single-sigma noise to image_latents (using image_noise_sigma range) and single-sigma noise to + fake_image_latents (using video_noise_sigma range). + """ + + model_name = "helios" + + @property + def description(self) -> str: + return ( + "Adds noise to image_latents and fake_image_latents for I2V conditioning. " + "Uses random sigma from configured ranges for each." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image_latents"), + InputParam( + "fake_image_latents", + required=True, + type_hint=torch.Tensor, + description="Fake image latents used as history seed for I2V generation.", + ), + InputParam( + "image_noise_sigma_min", + default=0.111, + type_hint=float, + description="Minimum sigma for image latent noise.", + ), + InputParam( + "image_noise_sigma_max", + default=0.135, + type_hint=float, + description="Maximum sigma for image latent noise.", + ), + InputParam( + "video_noise_sigma_min", + default=0.111, + type_hint=float, + description="Minimum sigma for video/fake-image latent noise.", + ), + InputParam( + "video_noise_sigma_max", + default=0.135, + type_hint=float, + description="Maximum sigma for video/fake-image latent noise.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("image_latents"), + OutputParam("fake_image_latents", type_hint=torch.Tensor, description="Noisy fake image latents"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + image_latents = block_state.image_latents + fake_image_latents = block_state.fake_image_latents + + # Add noise to image_latents + image_noise_sigma = ( + torch.rand(1, device=device, generator=block_state.generator) + * (block_state.image_noise_sigma_max - block_state.image_noise_sigma_min) + + block_state.image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=block_state.generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + + # Add noise to fake_image_latents + fake_image_noise_sigma = ( + torch.rand(1, device=device, generator=block_state.generator) + * (block_state.video_noise_sigma_max - block_state.video_noise_sigma_min) + + block_state.video_noise_sigma_min + ) + fake_image_latents = ( + fake_image_noise_sigma + * randn_tensor(fake_image_latents.shape, generator=block_state.generator, device=device) + + (1 - fake_image_noise_sigma) * fake_image_latents + ) + + block_state.image_latents = image_latents.to(device=device, dtype=torch.float32) + block_state.fake_image_latents = fake_image_latents.to(device=device, dtype=torch.float32) + + self.set_block_state(state, block_state) + return components, state + + +class HeliosAddNoiseToVideoLatentsStep(ModularPipelineBlocks): + """Adds noise to image_latents and video_latents for V2V conditioning. + + Applies single-sigma noise to image_latents (using image_noise_sigma range) and per-frame noise to video_latents in + chunks (using video_noise_sigma range). + """ + + model_name = "helios" + + @property + def description(self) -> str: + return ( + "Adds noise to image_latents and video_latents for V2V conditioning. " + "Uses single-sigma noise for image_latents and per-frame noise for video chunks." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image_latents"), + InputParam( + "video_latents", + required=True, + type_hint=torch.Tensor, + description="Encoded video latents for V2V generation.", + ), + InputParam( + "num_latent_frames_per_chunk", + default=9, + type_hint=int, + description="Number of latent frames per temporal chunk.", + ), + InputParam( + "image_noise_sigma_min", + default=0.111, + type_hint=float, + description="Minimum sigma for image latent noise.", + ), + InputParam( + "image_noise_sigma_max", + default=0.135, + type_hint=float, + description="Maximum sigma for image latent noise.", + ), + InputParam( + "video_noise_sigma_min", + default=0.111, + type_hint=float, + description="Minimum sigma for video latent noise.", + ), + InputParam( + "video_noise_sigma_max", + default=0.135, + type_hint=float, + description="Maximum sigma for video latent noise.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("image_latents"), + OutputParam("video_latents", type_hint=torch.Tensor, description="Noisy video latents"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + image_latents = block_state.image_latents + video_latents = block_state.video_latents + num_latent_frames_per_chunk = block_state.num_latent_frames_per_chunk + + # Add noise to first frame (single sigma) + image_noise_sigma = ( + torch.rand(1, device=device, generator=block_state.generator) + * (block_state.image_noise_sigma_max - block_state.image_noise_sigma_min) + + block_state.image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=block_state.generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + + # Add per-frame noise to video chunks + noisy_latents_chunks = [] + num_latent_chunks = video_latents.shape[2] // num_latent_frames_per_chunk + for i in range(num_latent_chunks): + chunk_start = i * num_latent_frames_per_chunk + chunk_end = chunk_start + num_latent_frames_per_chunk + latent_chunk = video_latents[:, :, chunk_start:chunk_end, :, :] + + chunk_frames = latent_chunk.shape[2] + frame_sigmas = ( + torch.rand(chunk_frames, device=device, generator=block_state.generator) + * (block_state.video_noise_sigma_max - block_state.video_noise_sigma_min) + + block_state.video_noise_sigma_min + ) + frame_sigmas = frame_sigmas.view(1, 1, chunk_frames, 1, 1) + + noisy_chunk = ( + frame_sigmas * randn_tensor(latent_chunk.shape, generator=block_state.generator, device=device) + + (1 - frame_sigmas) * latent_chunk + ) + noisy_latents_chunks.append(noisy_chunk) + video_latents = torch.cat(noisy_latents_chunks, dim=2) + + block_state.image_latents = image_latents.to(device=device, dtype=torch.float32) + block_state.video_latents = video_latents.to(device=device, dtype=torch.float32) + + self.set_block_state(state, block_state) + return components, state + + +class HeliosPrepareHistoryStep(ModularPipelineBlocks): + """Prepares chunk/history indices and initializes history state for the chunk loop.""" + + model_name = "helios" + + @property + def description(self) -> str: + return ( + "Prepares the chunk loop by computing latent dimensions, number of chunks, " + "history indices, and initializing history state (history_latents, image_latents, latent_chunks)." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", HeliosTransformer3DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("height", default=384), + InputParam.template("width", default=640), + InputParam( + "num_frames", default=132, type_hint=int, description="Total number of video frames to generate." + ), + InputParam("batch_size", required=True, type_hint=int), + InputParam( + "num_latent_frames_per_chunk", + default=9, + type_hint=int, + description="Number of latent frames per temporal chunk.", + ), + InputParam( + "history_sizes", + default=[16, 2, 1], + type_hint=list, + description="Sizes of long/mid/short history buffers for temporal context.", + ), + InputParam( + "keep_first_frame", + default=True, + type_hint=bool, + description="Whether to keep the first frame as a prefix in history.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("num_latent_chunk", type_hint=int, description="Number of temporal chunks"), + OutputParam("latent_shape", type_hint=tuple, description="Shape of latent tensor per chunk"), + OutputParam("history_sizes", type_hint=list, description="Adjusted history sizes (sorted, descending)"), + OutputParam("indices_hidden_states", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields"), + OutputParam("indices_latents_history_short", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields"), + OutputParam("indices_latents_history_mid", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields"), + OutputParam("indices_latents_history_long", type_hint=torch.Tensor, kwargs_type="denoiser_input_fields"), + OutputParam("history_latents", type_hint=torch.Tensor, description="Initialized zero history latents"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + batch_size = block_state.batch_size + device = components._execution_device + + block_state.num_frames = max(block_state.num_frames, 1) + history_sizes = sorted(block_state.history_sizes, reverse=True) + + num_channels_latents = components.num_channels_latents + h_latent = block_state.height // components.vae_scale_factor_spatial + w_latent = block_state.width // components.vae_scale_factor_spatial + + # Compute number of chunks + block_state.window_num_frames = ( + block_state.num_latent_frames_per_chunk - 1 + ) * components.vae_scale_factor_temporal + 1 + block_state.num_latent_chunk = max( + 1, (block_state.num_frames + block_state.window_num_frames - 1) // block_state.window_num_frames + ) + + # Modify history_sizes for non-keep_first_frame (matching pipeline behavior) + if not block_state.keep_first_frame: + history_sizes = history_sizes.copy() + history_sizes[-1] = history_sizes[-1] + 1 + + # Compute indices ONCE (same structure for all chunks) + if block_state.keep_first_frame: + indices = torch.arange(0, sum([1, *history_sizes, block_state.num_latent_frames_per_chunk])) + ( + indices_prefix, + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_1x, + indices_hidden_states, + ) = indices.split([1, *history_sizes, block_state.num_latent_frames_per_chunk], dim=0) + indices_latents_history_short = torch.cat([indices_prefix, indices_latents_history_1x], dim=0) + else: + indices = torch.arange(0, sum([*history_sizes, block_state.num_latent_frames_per_chunk])) + ( + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_short, + indices_hidden_states, + ) = indices.split([*history_sizes, block_state.num_latent_frames_per_chunk], dim=0) + + # Latent shape per chunk + block_state.latent_shape = ( + batch_size, + num_channels_latents, + block_state.num_latent_frames_per_chunk, + h_latent, + w_latent, + ) + + # Set outputs + block_state.history_sizes = history_sizes + block_state.indices_hidden_states = indices_hidden_states.unsqueeze(0) + block_state.indices_latents_history_short = indices_latents_history_short.unsqueeze(0) + block_state.indices_latents_history_mid = indices_latents_history_mid.unsqueeze(0) + block_state.indices_latents_history_long = indices_latents_history_long.unsqueeze(0) + block_state.history_latents = torch.zeros( + batch_size, + num_channels_latents, + sum(history_sizes), + h_latent, + w_latent, + device=device, + dtype=torch.float32, + ) + + self.set_block_state(state, block_state) + + return components, state + + +class HeliosI2VSeedHistoryStep(ModularPipelineBlocks): + """Seeds history_latents with fake_image_latents for I2V pipelines. + + This small additive step runs after HeliosPrepareHistoryStep and appends fake_image_latents to the initialized + history_latents tensor. + """ + + model_name = "helios" + + @property + def description(self) -> str: + return "I2V history seeding: appends fake_image_latents to history_latents." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("history_latents", required=True, type_hint=torch.Tensor), + InputParam("fake_image_latents", required=True, type_hint=torch.Tensor), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "history_latents", type_hint=torch.Tensor, description="History latents seeded with fake_image_latents" + ), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + block_state.history_latents = torch.cat([block_state.history_latents, block_state.fake_image_latents], dim=2) + + self.set_block_state(state, block_state) + return components, state + + +class HeliosV2VSeedHistoryStep(ModularPipelineBlocks): + """Seeds history_latents with video_latents for V2V pipelines. + + This step runs after HeliosPrepareHistoryStep and replaces the tail of history_latents with video_latents. If the + video has fewer frames than the history, the beginning of history is preserved. + """ + + model_name = "helios" + + @property + def description(self) -> str: + return "V2V history seeding: replaces the tail of history_latents with video_latents." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("history_latents", required=True, type_hint=torch.Tensor), + InputParam("video_latents", required=True, type_hint=torch.Tensor), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "history_latents", type_hint=torch.Tensor, description="History latents seeded with video_latents" + ), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + history_latents = block_state.history_latents + video_latents = block_state.video_latents + + history_frames = history_latents.shape[2] + video_frames = video_latents.shape[2] + if video_frames < history_frames: + keep_frames = history_frames - video_frames + history_latents = torch.cat([history_latents[:, :, :keep_frames, :, :], video_latents], dim=2) + else: + history_latents = video_latents + + block_state.history_latents = history_latents + + self.set_block_state(state, block_state) + return components, state + + +class HeliosSetTimestepsStep(ModularPipelineBlocks): + """Computes scheduler parameters (mu, sigmas) for the chunk loop.""" + + model_name = "helios" + + @property + def description(self) -> str: + return "Computes scheduler shift parameter (mu) and default sigmas for the Helios chunk loop." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", HeliosTransformer3DModel), + ComponentSpec("scheduler", HeliosScheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("latent_shape", required=True, type_hint=tuple), + InputParam.template("num_inference_steps"), + InputParam.template("sigmas"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("mu", type_hint=float, description="Scheduler shift parameter"), + OutputParam("sigmas", type_hint=list, description="Sigma schedule for diffusion"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + patch_size = components.transformer.config.patch_size + latent_shape = block_state.latent_shape + image_seq_len = (latent_shape[-1] * latent_shape[-2] * latent_shape[-3]) // ( + patch_size[0] * patch_size[1] * patch_size[2] + ) + + if block_state.sigmas is None: + block_state.sigmas = np.linspace(0.999, 0.0, block_state.num_inference_steps + 1)[:-1] + + block_state.mu = calculate_shift( + image_seq_len, + components.scheduler.config.get("base_image_seq_len", 256), + components.scheduler.config.get("max_image_seq_len", 4096), + components.scheduler.config.get("base_shift", 0.5), + components.scheduler.config.get("max_shift", 1.15), + ) + + self.set_block_state(state, block_state) + + return components, state diff --git a/diffusers/modular_pipelines/helios/decoders.py b/diffusers/modular_pipelines/helios/decoders.py new file mode 100644 index 0000000000000000000000000000000000000000..c448d36136e6da7ab2a5a03fc9e7c4867e43759c --- /dev/null +++ b/diffusers/modular_pipelines/helios/decoders.py @@ -0,0 +1,112 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import PIL +import torch + +from ...configuration_utils import FrozenDict +from ...models import AutoencoderKLWan +from ...utils import logging +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class HeliosDecodeStep(ModularPipelineBlocks): + """Decode all chunk latents with VAE, trim frames, and postprocess into final video output.""" + + model_name = "helios" + + @property + def description(self) -> str: + return ( + "Decodes all chunk latents with the VAE, concatenates them, " + "trims to the target frame count, and postprocesses into the final video output." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 8}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "latent_chunks", required=True, type_hint=list, description="List of per-chunk denoised latent tensors" + ), + InputParam("num_frames", required=True, type_hint=int, description="The target number of output frames"), + InputParam.template("output_type", default="np"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "videos", + type_hint=list[list[PIL.Image.Image]] | list[torch.Tensor] | list[np.ndarray], + description="The generated videos, can be a PIL.Image.Image, torch.Tensor or a numpy array", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + vae = components.vae + device = components._execution_device + decode_dtype = vae.dtype + + latents_mean = ( + torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(device, decode_dtype) + ) + latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( + device, decode_dtype + ) + + history_video = None + for chunk_latents in block_state.latent_chunks: + current_latents = chunk_latents.to(device=device, dtype=decode_dtype) / latents_std + latents_mean + current_video = vae.decode(current_latents, return_dict=False)[0] + + if history_video is None: + history_video = current_video + else: + history_video = torch.cat([history_video, current_video], dim=2) + + # Trim to proper frame count + generated_frames = history_video.size(2) + generated_frames = ( + generated_frames - 1 + ) // components.vae_scale_factor_temporal * components.vae_scale_factor_temporal + 1 + history_video = history_video[:, :, :generated_frames] + + block_state.videos = components.video_processor.postprocess_video( + history_video, output_type=block_state.output_type + ) + + self.set_block_state(state, block_state) + + return components, state diff --git a/diffusers/modular_pipelines/helios/denoise.py b/diffusers/modular_pipelines/helios/denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..5fcf01a73ffc189f2399fd43f8ba3ea163cfd448 --- /dev/null +++ b/diffusers/modular_pipelines/helios/denoise.py @@ -0,0 +1,1069 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import math + +import torch +import torch.nn.functional as F +from tqdm.auto import tqdm + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance, ClassifierFreeZeroStarGuidance +from ...models import HeliosTransformer3DModel +from ...schedulers import HeliosScheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + PipelineState, +) +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .before_denoise import calculate_shift +from .modular_pipeline import HeliosModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def sample_block_noise( + batch_size, + channel, + num_frames, + height, + width, + gamma, + patch_size=(1, 2, 2), + device=None, + generator=None, +): + """Generate spatially-correlated block noise for pyramid upsampling correction. + + Uses a multivariate normal distribution with covariance based on `gamma` to produce noise with block structure, + matching the upsampling artifacts that need correction. + """ + # NOTE: A generator must be provided to ensure correct and reproducible results. + # Creating a default generator here is a fallback only — without a fixed seed, + # the output will be non-deterministic and may produce incorrect results in CP context. + if generator is None: + generator = torch.Generator(device=device) + elif isinstance(generator, list): + generator = generator[0] + + _, ph, pw = patch_size + block_size = ph * pw + + cov = ( + torch.eye(block_size, device=device) * (1 + gamma) - torch.ones(block_size, block_size, device=device) * gamma + ) + cov += torch.eye(block_size, device=device) * 1e-8 + cov = cov.float() # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16. + + L = torch.linalg.cholesky(cov) + block_number = batch_size * channel * num_frames * (height // ph) * (width // pw) + z = torch.randn(block_number, block_size, device=generator.device, generator=generator).to(device) + noise = z @ L.T + + noise = noise.view(batch_size, channel, num_frames, height // ph, width // pw, ph, pw) + noise = noise.permute(0, 1, 2, 3, 5, 4, 6).reshape(batch_size, channel, num_frames, height, width) + return noise + + +# ======================================== +# Chunk Loop Leaf Blocks +# ======================================== + + +class HeliosChunkHistorySliceStep(ModularPipelineBlocks): + """Slices history latents into short/mid/long for a T2V chunk. + + At k==0 with no image_latents, creates a zero prefix. Otherwise uses image_latents (either provided or captured + from first chunk by HeliosChunkUpdateStep). + """ + + model_name = "helios" + + @property + def description(self) -> str: + return ( + "T2V history slice: splits history into long/mid/short. At k==0 with no image_latents, " + "creates a zero prefix; otherwise uses image_latents as prefix for short history." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "keep_first_frame", + default=True, + type_hint=bool, + description="Whether to keep the first frame as a prefix in history.", + ), + InputParam( + "history_sizes", + required=True, + type_hint=list, + description="Sizes of long/mid/short history buffers for temporal context.", + ), + InputParam( + "history_latents", + required=True, + type_hint=torch.Tensor, + description="Accumulated history latents from previous chunks.", + ), + InputParam("latent_shape", required=True, type_hint=tuple), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int): + keep_first_frame = block_state.keep_first_frame + history_sizes = block_state.history_sizes + image_latents = block_state.image_latents + device = components._execution_device + + batch_size, num_channels_latents, _, h_latent, w_latent = block_state.latent_shape + + if keep_first_frame: + latents_history_long, latents_history_mid, latents_history_1x = block_state.history_latents[ + :, :, -sum(history_sizes) : + ].split(history_sizes, dim=2) + if image_latents is None and k == 0: + latents_prefix = torch.zeros( + batch_size, + num_channels_latents, + 1, + h_latent, + w_latent, + device=device, + dtype=torch.float32, + ) + else: + latents_prefix = image_latents + latents_history_short = torch.cat([latents_prefix, latents_history_1x], dim=2) + else: + latents_history_long, latents_history_mid, latents_history_short = block_state.history_latents[ + :, :, -sum(history_sizes) : + ].split(history_sizes, dim=2) + + block_state.latents_history_short = latents_history_short + block_state.latents_history_mid = latents_history_mid + block_state.latents_history_long = latents_history_long + + return components, block_state + + +class HeliosI2VChunkHistorySliceStep(ModularPipelineBlocks): + """Slices history latents into short/mid/long for an I2V chunk. + + Always uses image_latents as prefix (assumes history pre-seeded with fake_image_latents). + """ + + model_name = "helios" + + @property + def description(self) -> str: + return ( + "I2V history slice: splits pre-seeded history into long/mid/short, " + "always using image_latents as prefix for short history." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "keep_first_frame", + default=True, + type_hint=bool, + description="Whether to keep the first frame as a prefix in history.", + ), + InputParam( + "history_sizes", + required=True, + type_hint=list, + description="Sizes of long/mid/short history buffers for temporal context.", + ), + InputParam( + "history_latents", + required=True, + type_hint=torch.Tensor, + description="Accumulated history latents from previous chunks.", + ), + InputParam( + "image_latents", + required=True, + type_hint=torch.Tensor, + description="First-frame latents used as prefix for short history.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int): + keep_first_frame = block_state.keep_first_frame + history_sizes = block_state.history_sizes + image_latents = block_state.image_latents + + if keep_first_frame: + latents_history_long, latents_history_mid, latents_history_1x = block_state.history_latents[ + :, :, -sum(history_sizes) : + ].split(history_sizes, dim=2) + latents_history_short = torch.cat([image_latents, latents_history_1x], dim=2) + else: + latents_history_long, latents_history_mid, latents_history_short = block_state.history_latents[ + :, :, -sum(history_sizes) : + ].split(history_sizes, dim=2) + + block_state.latents_history_short = latents_history_short + block_state.latents_history_mid = latents_history_mid + block_state.latents_history_long = latents_history_long + + return components, block_state + + +class HeliosChunkNoiseGenStep(ModularPipelineBlocks): + """Generates noise latents for a chunk using randn_tensor.""" + + model_name = "helios" + + @property + def description(self) -> str: + return "Generates random noise latents at full resolution for a single chunk." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("latent_shape", required=True, type_hint=tuple), + InputParam.template("generator"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int): + device = components._execution_device + block_state.latents = randn_tensor( + block_state.latent_shape, generator=block_state.generator, device=device, dtype=torch.float32 + ) + return components, block_state + + +class HeliosPyramidChunkNoiseGenStep(ModularPipelineBlocks): + """Generates noise latents and downsamples to smallest pyramid level.""" + + model_name = "helios-pyramid" + + @property + def description(self) -> str: + return ( + "Generates random noise at full resolution, then downsamples to the smallest " + "pyramid level via bilinear interpolation." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("latent_shape", required=True, type_hint=tuple), + InputParam( + "pyramid_num_inference_steps_list", + default=[10, 10, 10], + type_hint=list, + description="Number of denoising steps per pyramid stage.", + ), + InputParam.template("generator"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int): + device = components._execution_device + batch_size, num_channels_latents, num_latent_frames, h_latent, w_latent = block_state.latent_shape + + latents = randn_tensor( + block_state.latent_shape, generator=block_state.generator, device=device, dtype=torch.float32 + ) + + # Downsample to smallest pyramid level + h, w = h_latent, w_latent + latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_latent_frames, num_channels_latents, h, w) + for _ in range(len(block_state.pyramid_num_inference_steps_list) - 1): + h //= 2 + w //= 2 + latents = F.interpolate(latents, size=(h, w), mode="bilinear") * 2 + block_state.latents = latents.reshape(batch_size, num_latent_frames, num_channels_latents, h, w).permute( + 0, 2, 1, 3, 4 + ) + + return components, block_state + + +class HeliosChunkSchedulerResetStep(ModularPipelineBlocks): + """Resets the scheduler with timesteps for a single chunk.""" + + model_name = "helios" + + @property + def description(self) -> str: + return "Resets the scheduler with the correct timesteps and shift parameter (mu) for this chunk." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", HeliosScheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("mu", required=True, type_hint=float), + InputParam.template("sigmas", required=True), + InputParam.template("num_inference_steps"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int): + device = components._execution_device + components.scheduler.set_timesteps( + block_state.num_inference_steps, device=device, sigmas=block_state.sigmas, mu=block_state.mu + ) + block_state.timesteps = components.scheduler.timesteps + + return components, block_state + + +# ======================================== +# Inner Denoising Blocks +# ======================================== + + +class HeliosChunkDenoiseInner(ModularPipelineBlocks): + """Inner timestep loop for denoising a single chunk, using guider for guidance.""" + + model_name = "helios" + + @property + def description(self) -> str: + return ( + "Inner denoising loop that iterates over timesteps for a single chunk. " + "Uses the guider to manage conditional/unconditional forward passes with cache_context, " + "applies guidance, and runs scheduler step." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", HeliosTransformer3DModel), + ComponentSpec("scheduler", HeliosScheduler), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 5.0}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents"), + InputParam.template("timesteps"), + InputParam("prompt_embeds", type_hint=torch.Tensor), + InputParam("negative_prompt_embeds", type_hint=torch.Tensor), + InputParam.template("denoiser_input_fields"), + InputParam.template("num_inference_steps"), + InputParam.template("attention_kwargs"), + InputParam.template("generator"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int): + latents = block_state.latents + timesteps = block_state.timesteps + num_inference_steps = block_state.num_inference_steps + + transformer_dtype = components.transformer.dtype + num_warmup_steps = len(timesteps) - num_inference_steps * components.scheduler.order + + # Guider inputs: only encoder_hidden_states differs between cond/uncond + guider_inputs = { + "encoder_hidden_states": (block_state.prompt_embeds, block_state.negative_prompt_embeds), + } + + # Build shared kwargs from denoiser_input_fields (excludes guider-managed ones) + transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys()) + shared_kwargs = {} + for field_name, field_value in block_state.denoiser_input_fields.items(): + if field_name in transformer_args and field_name not in guider_inputs: + shared_kwargs[field_name] = field_value + + # Add loop-internal history latents with dtype casting + shared_kwargs["latents_history_short"] = block_state.latents_history_short.to(transformer_dtype) + shared_kwargs["latents_history_mid"] = block_state.latents_history_mid.to(transformer_dtype) + shared_kwargs["latents_history_long"] = block_state.latents_history_long.to(transformer_dtype) + shared_kwargs["attention_kwargs"] = block_state.attention_kwargs + + with tqdm(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + timestep = t.expand(latents.shape[0]).to(torch.int64) + latent_model_input = latents.to(transformer_dtype) + + components.guider.set_state(step=i, num_inference_steps=num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = {k: getattr(guider_state_batch, k) for k in guider_inputs.keys()} + + context_name = getattr(guider_state_batch, components.guider._identifier_key) + with components.transformer.cache_context(context_name): + guider_state_batch.noise_pred = components.transformer( + hidden_states=latent_model_input, + timestep=timestep, + return_dict=False, + **cond_kwargs, + **shared_kwargs, + )[0] + components.guider.cleanup_models(components.transformer) + + noise_pred = components.guider(guider_state)[0] + + # Scheduler step + latents = components.scheduler.step( + noise_pred, + t, + latents, + generator=block_state.generator, + return_dict=False, + )[0] + + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0 + ): + progress_bar.update() + + block_state.latents = latents + return components, block_state + + +class HeliosPyramidChunkDenoiseInner(ModularPipelineBlocks): + """Nested pyramid stage loop with inner timestep denoising. + + For each pyramid stage (small -> full resolution): + 1. Upsample latents + block noise correction (stages > 0) + 2. Compute mu from current resolution, set scheduler timesteps + 3. Run timestep denoising loop (same logic as HeliosChunkDenoiseInner) + """ + + model_name = "helios-pyramid" + + @property + def description(self) -> str: + return ( + "Pyramid denoising inner block: loops over pyramid stages from smallest to full resolution. " + "Each stage upsamples latents (with block noise correction), recomputes scheduler parameters, " + "and runs the timestep denoising loop." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", HeliosTransformer3DModel), + ComponentSpec("scheduler", HeliosScheduler), + ComponentSpec( + "guider", + ClassifierFreeZeroStarGuidance, + config=FrozenDict({"guidance_scale": 5.0, "zero_init_steps": 2}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents"), + InputParam("prompt_embeds", type_hint=torch.Tensor), + InputParam("negative_prompt_embeds", type_hint=torch.Tensor), + InputParam.template("denoiser_input_fields"), + InputParam( + "pyramid_num_inference_steps_list", + default=[10, 10, 10], + type_hint=list, + description="Number of denoising steps per pyramid stage.", + ), + InputParam.template("attention_kwargs"), + InputParam.template("generator"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int): + device = components._execution_device + transformer_dtype = components.transformer.dtype + latents = block_state.latents + pyramid_num_stages = len(block_state.pyramid_num_inference_steps_list) + + # Guider inputs: only encoder_hidden_states differs between cond/uncond + guider_inputs = { + "encoder_hidden_states": (block_state.prompt_embeds, block_state.negative_prompt_embeds), + } + + # Build shared kwargs from denoiser_input_fields (excludes guider-managed ones) + transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys()) + shared_kwargs = {} + for field_name, field_value in block_state.denoiser_input_fields.items(): + if field_name in transformer_args and field_name not in guider_inputs: + shared_kwargs[field_name] = field_value + + # Add loop-internal history latents with dtype casting + shared_kwargs["latents_history_short"] = block_state.latents_history_short.to(transformer_dtype) + shared_kwargs["latents_history_mid"] = block_state.latents_history_mid.to(transformer_dtype) + shared_kwargs["latents_history_long"] = block_state.latents_history_long.to(transformer_dtype) + shared_kwargs["attention_kwargs"] = block_state.attention_kwargs + + # Save original zero_init_steps if the guider supports it (e.g. ClassifierFreeZeroStarGuidance). + # Helios only applies zero init in pyramid stage 0 (lowest resolution), so we disable it + # for subsequent stages by temporarily setting zero_init_steps=0. + orig_zero_init_steps = getattr(components.guider, "zero_init_steps", None) + + for i_s in range(pyramid_num_stages): + # --- Stage setup --- + + # Disable zero init for stages > 0 (only stage 0 should have zero init) + if orig_zero_init_steps is not None and i_s > 0: + components.guider.zero_init_steps = 0 + + # a. Compute mu from current resolution (before upsample, matching standard pipeline) + patch_size = components.transformer.config.patch_size + image_seq_len = (latents.shape[-1] * latents.shape[-2] * latents.shape[-3]) // ( + patch_size[0] * patch_size[1] * patch_size[2] + ) + mu = calculate_shift( + image_seq_len, + components.scheduler.config.get("base_image_seq_len", 256), + components.scheduler.config.get("max_image_seq_len", 4096), + components.scheduler.config.get("base_shift", 0.5), + components.scheduler.config.get("max_shift", 1.15), + ) + + # b. Set scheduler timesteps for this stage + num_inference_steps = block_state.pyramid_num_inference_steps_list[i_s] + components.scheduler.set_timesteps( + num_inference_steps, + i_s, + device=device, + mu=mu, + ) + timesteps = components.scheduler.timesteps + + # c. Upsample + block noise correction for stages > 0 + if i_s > 0: + batch_size, num_channels_latents, num_frames, current_h, current_w = latents.shape + new_h = current_h * 2 + new_w = current_w * 2 + + latents = latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_frames, num_channels_latents, current_h, current_w + ) + latents = F.interpolate(latents, size=(new_h, new_w), mode="nearest") + latents = latents.reshape(batch_size, num_frames, num_channels_latents, new_h, new_w).permute( + 0, 2, 1, 3, 4 + ) + + # Block noise correction + ori_sigma = 1 - components.scheduler.ori_start_sigmas[i_s] + gamma = components.scheduler.config.gamma + alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) + beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) + + batch_size, num_channels_latents, num_frames, h, w = latents.shape + noise = sample_block_noise( + batch_size, + num_channels_latents, + num_frames, + h, + w, + gamma, + patch_size, + device=device, + generator=block_state.generator, + ) + noise = noise.to(dtype=transformer_dtype) + latents = alpha * latents + beta * noise + + # --- Timestep denoising loop --- + num_warmup_steps = len(timesteps) - num_inference_steps * components.scheduler.order + + with tqdm(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + timestep = t.expand(latents.shape[0]).to(torch.int64) + latent_model_input = latents.to(transformer_dtype) + + components.guider.set_state(step=i, num_inference_steps=num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = {kk: getattr(guider_state_batch, kk) for kk in guider_inputs.keys()} + + context_name = getattr(guider_state_batch, components.guider._identifier_key) + with components.transformer.cache_context(context_name): + guider_state_batch.noise_pred = components.transformer( + hidden_states=latent_model_input, + timestep=timestep, + return_dict=False, + **cond_kwargs, + **shared_kwargs, + )[0] + components.guider.cleanup_models(components.transformer) + + noise_pred = components.guider(guider_state)[0] + + # Scheduler step + latents = components.scheduler.step( + noise_pred, + t, + latents, + generator=block_state.generator, + return_dict=False, + )[0] + + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0 + ): + progress_bar.update() + + # Restore original zero_init_steps + if orig_zero_init_steps is not None: + components.guider.zero_init_steps = orig_zero_init_steps + + block_state.latents = latents + return components, block_state + + +# ======================================== +# Post-Denoise Update +# ======================================== + + +class HeliosChunkUpdateStep(ModularPipelineBlocks): + """Updates chunk collection and history after denoising a single chunk.""" + + model_name = "helios" + + @property + def description(self) -> str: + return ( + "Post-denoising update step: appends the denoised latents to the chunk list, " + "captures image_latents from the first chunk if needed, and extends history_latents." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("latents", type_hint=torch.Tensor), + InputParam("history_latents", type_hint=torch.Tensor), + InputParam("keep_first_frame", default=True, type_hint=bool), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int): + # e. Collect denoised latents for this chunk + block_state.latent_chunks.append(block_state.latents) + + # f. Update history + if block_state.keep_first_frame and k == 0 and block_state.image_latents is None: + block_state.image_latents = block_state.latents[:, :, 0:1, :, :] + + block_state.history_latents = torch.cat([block_state.history_latents, block_state.latents], dim=2) + + return components, block_state + + +# ======================================== +# Chunk Loop Wrapper +# ======================================== + + +class HeliosChunkLoopWrapper(LoopSequentialPipelineBlocks): + """Outer chunk loop that iterates over temporal chunks. + + History indices, scheduler params, and history state are prepared by HeliosPrepareHistoryStep and + HeliosSetTimestepsStep before this block runs. Sub-blocks handle per-chunk preparation, denoising, and history + updates. + """ + + model_name = "helios" + + @property + def description(self) -> str: + return ( + "Pipeline block that iterates over temporal chunks for progressive video generation. " + "At each chunk iteration, it runs sub-blocks for preparation, denoising, and history updates." + ) + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam("num_latent_chunk", required=True, type_hint=int), + ] + + @property + def loop_intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.latent_chunks = [] + + if not hasattr(block_state, "image_latents"): + block_state.image_latents = None + + for k in range(block_state.num_latent_chunk): + components, block_state = self.loop_step(components, block_state, k=k) + + self.set_block_state(state, block_state) + + return components, state + + +# ======================================== +# Composed Chunk Denoise Steps +# ======================================== + + +class HeliosChunkDenoiseStep(HeliosChunkLoopWrapper): + """T2V chunk-based denoising: history slice -> noise gen -> scheduler reset -> denoise -> update.""" + + block_classes = [ + HeliosChunkHistorySliceStep, + HeliosChunkNoiseGenStep, + HeliosChunkSchedulerResetStep, + HeliosChunkDenoiseInner, + HeliosChunkUpdateStep, + ] + block_names = ["history_slice", "noise_gen", "scheduler_reset", "denoise_inner", "update_chunk"] + + @property + def description(self) -> str: + return ( + "T2V chunk denoise step that iterates over temporal chunks.\n" + "At each chunk: history_slice -> noise_gen -> scheduler_reset -> denoise_inner -> update_chunk." + ) + + +class HeliosI2VChunkDenoiseStep(HeliosChunkLoopWrapper): + """I2V chunk-based denoising: I2V history slice -> noise gen -> scheduler reset -> denoise -> update.""" + + block_classes = [ + HeliosI2VChunkHistorySliceStep, + HeliosChunkNoiseGenStep, + HeliosChunkSchedulerResetStep, + HeliosChunkDenoiseInner, + HeliosChunkUpdateStep, + ] + block_names = ["history_slice", "noise_gen", "scheduler_reset", "denoise_inner", "update_chunk"] + + @property + def description(self) -> str: + return ( + "I2V chunk denoise step that iterates over temporal chunks.\n" + "At each chunk: history_slice (I2V) -> noise_gen -> scheduler_reset -> denoise_inner -> update_chunk." + ) + + +class HeliosPyramidDistilledChunkDenoiseInner(ModularPipelineBlocks): + """Nested pyramid stage loop with DMD denoising for distilled checkpoints. + + Same progressive multi-resolution strategy as HeliosPyramidChunkDenoiseInner, but: + - Guidance is disabled (guidance_scale=1.0, no unconditional pass) + - Supports is_amplify_first_chunk (doubles first chunk's timesteps via scheduler) + - Tracks start_point_list and passes DMD-specific args to scheduler.step() + """ + + model_name = "helios-pyramid" + + @property + def description(self) -> str: + return ( + "Distilled pyramid denoising inner block for DMD checkpoints. Loops over pyramid stages " + "from smallest to full resolution with guidance disabled and DMD scheduler support." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", HeliosTransformer3DModel), + ComponentSpec("scheduler", HeliosScheduler), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 1.0}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents"), + InputParam("prompt_embeds", type_hint=torch.Tensor), + InputParam("negative_prompt_embeds", type_hint=torch.Tensor), + InputParam.template("denoiser_input_fields"), + InputParam( + "pyramid_num_inference_steps_list", + default=[2, 2, 2], + type_hint=list, + description="Number of denoising steps per pyramid stage.", + ), + InputParam( + "is_amplify_first_chunk", + default=True, + type_hint=bool, + description="Whether to double the first chunk's timesteps via the scheduler for amplified generation.", + ), + InputParam.template("attention_kwargs"), + InputParam.template("generator"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, block_state: BlockState, k: int): + device = components._execution_device + transformer_dtype = components.transformer.dtype + latents = block_state.latents + pyramid_num_stages = len(block_state.pyramid_num_inference_steps_list) + is_first_chunk = k == 0 + + # Track start points for DMD scheduler + start_point_list = [latents] + + # Guider inputs: only encoder_hidden_states differs between cond/uncond + guider_inputs = { + "encoder_hidden_states": (block_state.prompt_embeds, block_state.negative_prompt_embeds), + } + + # Build shared kwargs from denoiser_input_fields (excludes guider-managed ones) + transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys()) + shared_kwargs = {} + for field_name, field_value in block_state.denoiser_input_fields.items(): + if field_name in transformer_args and field_name not in guider_inputs: + shared_kwargs[field_name] = field_value + + # Add loop-internal history latents with dtype casting + shared_kwargs["latents_history_short"] = block_state.latents_history_short.to(transformer_dtype) + shared_kwargs["latents_history_mid"] = block_state.latents_history_mid.to(transformer_dtype) + shared_kwargs["latents_history_long"] = block_state.latents_history_long.to(transformer_dtype) + shared_kwargs["attention_kwargs"] = block_state.attention_kwargs + + for i_s in range(pyramid_num_stages): + # --- Stage setup --- + patch_size = components.transformer.config.patch_size + + # a. Compute mu from current resolution (before upsample, matching standard pipeline) + image_seq_len = (latents.shape[-1] * latents.shape[-2] * latents.shape[-3]) // ( + patch_size[0] * patch_size[1] * patch_size[2] + ) + mu = calculate_shift( + image_seq_len, + components.scheduler.config.get("base_image_seq_len", 256), + components.scheduler.config.get("max_image_seq_len", 4096), + components.scheduler.config.get("base_shift", 0.5), + components.scheduler.config.get("max_shift", 1.15), + ) + + # b. Set scheduler timesteps for this stage (with DMD amplification) + num_inference_steps = block_state.pyramid_num_inference_steps_list[i_s] + components.scheduler.set_timesteps( + num_inference_steps, + i_s, + device=device, + mu=mu, + is_amplify_first_chunk=block_state.is_amplify_first_chunk and is_first_chunk, + ) + timesteps = components.scheduler.timesteps + + # c. Upsample + block noise correction for stages > 0 + if i_s > 0: + batch_size, num_channels_latents, num_frames, current_h, current_w = latents.shape + new_h = current_h * 2 + new_w = current_w * 2 + + latents = latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_frames, num_channels_latents, current_h, current_w + ) + latents = F.interpolate(latents, size=(new_h, new_w), mode="nearest") + latents = latents.reshape(batch_size, num_frames, num_channels_latents, new_h, new_w).permute( + 0, 2, 1, 3, 4 + ) + + # Block noise correction + ori_sigma = 1 - components.scheduler.ori_start_sigmas[i_s] + gamma = components.scheduler.config.gamma + alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) + beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) + + batch_size, num_channels_latents, num_frames, h, w = latents.shape + noise = sample_block_noise( + batch_size, + num_channels_latents, + num_frames, + h, + w, + gamma, + patch_size, + device=device, + generator=block_state.generator, + ) + noise = noise.to(dtype=transformer_dtype) + latents = alpha * latents + beta * noise + + start_point_list.append(latents) + + # --- Timestep denoising loop --- + num_warmup_steps = len(timesteps) - num_inference_steps * components.scheduler.order + + with tqdm(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + timestep = t.expand(latents.shape[0]).to(torch.int64) + latent_model_input = latents.to(transformer_dtype) + + components.guider.set_state(step=i, num_inference_steps=num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = {k: getattr(guider_state_batch, k) for k in guider_inputs.keys()} + + context_name = getattr(guider_state_batch, components.guider._identifier_key) + with components.transformer.cache_context(context_name): + guider_state_batch.noise_pred = components.transformer( + hidden_states=latent_model_input, + timestep=timestep, + return_dict=False, + **cond_kwargs, + **shared_kwargs, + )[0] + components.guider.cleanup_models(components.transformer) + + noise_pred = components.guider(guider_state)[0] + + # Scheduler step with DMD args + latents = components.scheduler.step( + noise_pred, + t, + latents, + generator=block_state.generator, + return_dict=False, + cur_sampling_step=i, + dmd_noisy_tensor=start_point_list[i_s], + dmd_sigmas=components.scheduler.sigmas, + dmd_timesteps=components.scheduler.timesteps, + all_timesteps=timesteps, + )[0] + + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % components.scheduler.order == 0 + ): + progress_bar.update() + + block_state.latents = latents + return components, block_state + + +class HeliosPyramidChunkDenoiseStep(HeliosChunkLoopWrapper): + """T2V pyramid chunk denoising: history slice -> pyramid noise gen -> pyramid denoise inner -> update.""" + + block_classes = [ + HeliosChunkHistorySliceStep, + HeliosPyramidChunkNoiseGenStep, + HeliosPyramidChunkDenoiseInner, + HeliosChunkUpdateStep, + ] + block_names = ["history_slice", "noise_gen", "denoise_inner", "update_chunk"] + + @property + def description(self) -> str: + return ( + "T2V pyramid chunk denoise step that iterates over temporal chunks.\n" + "At each chunk: history_slice -> noise_gen (pyramid) -> denoise_inner (pyramid stages) -> update_chunk.\n" + "Denoising starts at the smallest resolution and progressively upsamples." + ) + + +class HeliosPyramidI2VChunkDenoiseStep(HeliosChunkLoopWrapper): + """I2V pyramid chunk denoising: I2V history slice -> pyramid noise gen -> pyramid denoise inner -> update.""" + + block_classes = [ + HeliosI2VChunkHistorySliceStep, + HeliosPyramidChunkNoiseGenStep, + HeliosPyramidChunkDenoiseInner, + HeliosChunkUpdateStep, + ] + block_names = ["history_slice", "noise_gen", "denoise_inner", "update_chunk"] + + @property + def description(self) -> str: + return ( + "I2V pyramid chunk denoise step that iterates over temporal chunks.\n" + "At each chunk: history_slice (I2V) -> noise_gen (pyramid) -> denoise_inner (pyramid stages) -> update_chunk.\n" + "Denoising starts at the smallest resolution and progressively upsamples." + ) + + +class HeliosPyramidDistilledChunkDenoiseStep(HeliosChunkLoopWrapper): + """T2V distilled pyramid chunk denoising with DMD scheduler and no CFG.""" + + block_classes = [ + HeliosChunkHistorySliceStep, + HeliosPyramidChunkNoiseGenStep, + HeliosPyramidDistilledChunkDenoiseInner, + HeliosChunkUpdateStep, + ] + block_names = ["history_slice", "noise_gen", "denoise_inner", "update_chunk"] + + @property + def description(self) -> str: + return ( + "T2V distilled pyramid chunk denoise step with DMD scheduler.\n" + "At each chunk: history_slice -> noise_gen (pyramid) -> denoise_inner (distilled/DMD) -> update_chunk." + ) + + +class HeliosPyramidDistilledI2VChunkDenoiseStep(HeliosChunkLoopWrapper): + """I2V distilled pyramid chunk denoising with DMD scheduler and no CFG.""" + + block_classes = [ + HeliosI2VChunkHistorySliceStep, + HeliosPyramidChunkNoiseGenStep, + HeliosPyramidDistilledChunkDenoiseInner, + HeliosChunkUpdateStep, + ] + block_names = ["history_slice", "noise_gen", "denoise_inner", "update_chunk"] + + @property + def description(self) -> str: + return ( + "I2V distilled pyramid chunk denoise step with DMD scheduler.\n" + "At each chunk: history_slice (I2V) -> noise_gen (pyramid) -> denoise_inner (distilled/DMD) -> update_chunk." + ) diff --git a/diffusers/modular_pipelines/helios/encoders.py b/diffusers/modular_pipelines/helios/encoders.py new file mode 100644 index 0000000000000000000000000000000000000000..ce11f1b5876297bff1f35b36f1791e299041971a --- /dev/null +++ b/diffusers/modular_pipelines/helios/encoders.py @@ -0,0 +1,392 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import html + +import regex as re +import torch +from transformers import AutoTokenizer, UMT5EncoderModel + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import AutoencoderKLWan +from ...utils import is_ftfy_available, logging +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import HeliosModularPipeline + + +if is_ftfy_available(): + import ftfy + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r"\s+", " ", text) + text = text.strip() + return text + + +def prompt_clean(text): + text = whitespace_clean(basic_clean(text)) + return text + + +def get_t5_prompt_embeds( + text_encoder: UMT5EncoderModel, + tokenizer: AutoTokenizer, + prompt: str | list[str], + max_sequence_length: int, + device: torch.device, + dtype: torch.dtype | None = None, +): + """Encode text prompts into T5 embeddings for Helios. + + Args: + text_encoder: The T5 text encoder model. + tokenizer: The tokenizer for the text encoder. + prompt: The prompt or prompts to encode. + max_sequence_length: Maximum sequence length for tokenization. + device: Device to place tensors on. + dtype: Optional dtype override. Defaults to `text_encoder.dtype`. + + Returns: + A tuple of `(prompt_embeds, attention_mask)` where `prompt_embeds` is the encoded text embeddings and + `attention_mask` is a boolean mask. + """ + dtype = dtype or text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt = [prompt_clean(u) for u in prompt] + + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask + seq_lens = mask.gt(0).sum(dim=1).long() + + prompt_embeds = text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 + ) + + return prompt_embeds, text_inputs.attention_mask.bool() + + +class HeliosTextEncoderStep(ModularPipelineBlocks): + model_name = "helios" + + @property + def description(self) -> str: + return "Text Encoder step that generates text embeddings to guide the video generation" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", UMT5EncoderModel), + ComponentSpec("tokenizer", AutoTokenizer), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 5.0}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt"), + InputParam.template("negative_prompt"), + InputParam.template("max_sequence_length"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("negative_prompt_embeds"), + ] + + @staticmethod + def check_inputs(prompt, negative_prompt): + if prompt is not None and not isinstance(prompt, (str, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and not isinstance(negative_prompt, (str, list)): + raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") + + if prompt is not None and negative_prompt is not None: + prompt_list = [prompt] if isinstance(prompt, str) else prompt + neg_list = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + if type(prompt_list) is not type(neg_list): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + if len(prompt_list) != len(neg_list): + raise ValueError( + f"`negative_prompt` has batch size {len(neg_list)}, but `prompt` has batch size" + f" {len(prompt_list)}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + prompt = block_state.prompt + negative_prompt = block_state.negative_prompt + max_sequence_length = block_state.max_sequence_length + device = components._execution_device + + self.check_inputs(prompt, negative_prompt) + + # Encode prompt + block_state.prompt_embeds, _ = get_t5_prompt_embeds( + text_encoder=components.text_encoder, + tokenizer=components.tokenizer, + prompt=prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + + # Encode negative prompt + block_state.negative_prompt_embeds = None + if components.requires_unconditional_embeds: + negative_prompt = negative_prompt or "" + if isinstance(prompt, list) and isinstance(negative_prompt, str): + negative_prompt = len(prompt) * [negative_prompt] + + block_state.negative_prompt_embeds, _ = get_t5_prompt_embeds( + text_encoder=components.text_encoder, + tokenizer=components.tokenizer, + prompt=negative_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + + self.set_block_state(state, block_state) + return components, state + + +class HeliosImageVaeEncoderStep(ModularPipelineBlocks): + """Encodes an input image into VAE latent space for image-to-video generation.""" + + model_name = "helios" + + @property + def description(self) -> str: + return ( + "Image Encoder step that encodes an input image into VAE latent space, " + "producing image_latents (first frame prefix) and fake_image_latents (history seed) " + "for image-to-video generation." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 8}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image"), + InputParam.template("height", default=384), + InputParam.template("width", default=640), + InputParam( + "num_latent_frames_per_chunk", + default=9, + type_hint=int, + description="Number of latent frames per temporal chunk.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("image_latents"), + OutputParam( + "fake_image_latents", type_hint=torch.Tensor, description="Fake image latents for history seeding" + ), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + vae = components.vae + device = components._execution_device + + latents_mean = ( + torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(vae.device, vae.dtype) + ) + latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( + vae.device, vae.dtype + ) + + # Preprocess image to 4D tensor (B, C, H, W) + image = components.video_processor.preprocess( + block_state.image, height=block_state.height, width=block_state.width + ) + image_5d = image.unsqueeze(2).to(device=device, dtype=vae.dtype) # (B, C, 1, H, W) + + # Encode image to get image_latents + image_latents = vae.encode(image_5d).latent_dist.sample(generator=block_state.generator) + image_latents = (image_latents - latents_mean) * latents_std + + # Encode fake video to get fake_image_latents + min_frames = (block_state.num_latent_frames_per_chunk - 1) * components.vae_scale_factor_temporal + 1 + fake_video = image_5d.repeat(1, 1, min_frames, 1, 1) # (B, C, min_frames, H, W) + fake_latents_full = vae.encode(fake_video).latent_dist.sample(generator=block_state.generator) + fake_latents_full = (fake_latents_full - latents_mean) * latents_std + fake_image_latents = fake_latents_full[:, :, -1:, :, :] + + block_state.image_latents = image_latents.to(device=device, dtype=torch.float32) + block_state.fake_image_latents = fake_image_latents.to(device=device, dtype=torch.float32) + + self.set_block_state(state, block_state) + return components, state + + +class HeliosVideoVaeEncoderStep(ModularPipelineBlocks): + """Encodes an input video into VAE latent space for video-to-video generation. + + Produces `image_latents` (first frame) and `video_latents` (remaining frames encoded in chunks). + """ + + model_name = "helios" + + @property + def description(self) -> str: + return ( + "Video Encoder step that encodes an input video into VAE latent space, " + "producing image_latents (first frame) and video_latents (chunked video frames) " + "for video-to-video generation." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 8}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("video", required=True, description="Input video for video-to-video generation"), + InputParam.template("height", default=384), + InputParam.template("width", default=640), + InputParam( + "num_latent_frames_per_chunk", + default=9, + type_hint=int, + description="Number of latent frames per temporal chunk.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("image_latents"), + OutputParam("video_latents", type_hint=torch.Tensor, description="Encoded video latents (chunked)"), + ] + + @torch.no_grad() + def __call__(self, components: HeliosModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + vae = components.vae + device = components._execution_device + num_latent_frames_per_chunk = block_state.num_latent_frames_per_chunk + + latents_mean = ( + torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(vae.device, vae.dtype) + ) + latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( + vae.device, vae.dtype + ) + + # Preprocess video + video = components.video_processor.preprocess_video( + block_state.video, height=block_state.height, width=block_state.width + ) + video = video.to(device=device, dtype=vae.dtype) + + # Encode video into latents + num_frames = video.shape[2] + min_frames = (num_latent_frames_per_chunk - 1) * 4 + 1 + num_chunks = num_frames // min_frames + if num_chunks == 0: + raise ValueError( + f"Video must have at least {min_frames} frames " + f"(got {num_frames} frames). " + f"Required: (num_latent_frames_per_chunk - 1) * 4 + 1 = ({num_latent_frames_per_chunk} - 1) * 4 + 1 = {min_frames}" + ) + total_valid_frames = num_chunks * min_frames + start_frame = num_frames - total_valid_frames + + # Encode first frame + first_frame = video[:, :, 0:1, :, :] + image_latents = vae.encode(first_frame).latent_dist.sample(generator=block_state.generator) + image_latents = (image_latents - latents_mean) * latents_std + + # Encode remaining frames in chunks + latents_chunks = [] + for i in range(num_chunks): + chunk_start = start_frame + i * min_frames + chunk_end = chunk_start + min_frames + video_chunk = video[:, :, chunk_start:chunk_end, :, :] + chunk_latents = vae.encode(video_chunk).latent_dist.sample(generator=block_state.generator) + chunk_latents = (chunk_latents - latents_mean) * latents_std + latents_chunks.append(chunk_latents) + video_latents = torch.cat(latents_chunks, dim=2) + + block_state.image_latents = image_latents.to(device=device, dtype=torch.float32) + block_state.video_latents = video_latents.to(device=device, dtype=torch.float32) + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/helios/modular_blocks_helios.py b/diffusers/modular_pipelines/helios/modular_blocks_helios.py new file mode 100644 index 0000000000000000000000000000000000000000..c3d5cb4efc774a8fb911e57882bf021181146ae9 --- /dev/null +++ b/diffusers/modular_pipelines/helios/modular_blocks_helios.py @@ -0,0 +1,542 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam +from .before_denoise import ( + HeliosAdditionalInputsStep, + HeliosAddNoiseToImageLatentsStep, + HeliosAddNoiseToVideoLatentsStep, + HeliosI2VSeedHistoryStep, + HeliosPrepareHistoryStep, + HeliosSetTimestepsStep, + HeliosTextInputStep, + HeliosV2VSeedHistoryStep, +) +from .decoders import HeliosDecodeStep +from .denoise import HeliosChunkDenoiseStep, HeliosI2VChunkDenoiseStep +from .encoders import HeliosImageVaeEncoderStep, HeliosTextEncoderStep, HeliosVideoVaeEncoderStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# ==================== +# 1. Vae Encoder +# ==================== + + +# auto_docstring +class HeliosAutoVaeEncoderStep(AutoPipelineBlocks): + """ + Encoder step that encodes video or image inputs. This is an auto pipeline block. + - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided. + - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided. + - If neither is provided, step will be skipped. + + Components: + vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Inputs: + video (`None`, *optional*): + Input video for video-to-video generation + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + + Outputs: + image_latents (`Tensor`): + The latent representation of the input image. + video_latents (`Tensor`): + Encoded video latents (chunked) + fake_image_latents (`Tensor`): + Fake image latents for history seeding + """ + + block_classes = [HeliosVideoVaeEncoderStep, HeliosImageVaeEncoderStep] + block_names = ["video_encoder", "image_encoder"] + block_trigger_inputs = ["video", "image"] + + @property + def description(self): + return ( + "Encoder step that encodes video or image inputs. This is an auto pipeline block.\n" + " - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided.\n" + " - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided.\n" + " - If neither is provided, step will be skipped." + ) + + +# ==================== +# 2. DENOISE +# ==================== + + +# DENOISE (T2V) +# auto_docstring +class HeliosCoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block that takes encoded conditions and runs the chunk-based denoising process. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + model_name = "helios" + block_classes = [ + HeliosTextInputStep, + HeliosPrepareHistoryStep, + HeliosSetTimestepsStep, + HeliosChunkDenoiseStep, + ] + block_names = ["input", "prepare_history", "set_timesteps", "chunk_denoise"] + + @property + def description(self): + return "Denoise block that takes encoded conditions and runs the chunk-based denoising process." + + @property + def outputs(self): + return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")] + + +# DENOISE (I2V) +# auto_docstring +class HeliosI2VCoreDenoiseStep(SequentialPipelineBlocks): + """ + I2V denoise block that seeds history with image latents and uses I2V-aware chunk preparation. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + fake_image_latents (`Tensor`, *optional*): + Fake image latents used as history seed for I2V generation. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video/fake-image latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video/fake-image latent noise. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + model_name = "helios" + block_classes = [ + HeliosTextInputStep, + HeliosAdditionalInputsStep( + image_latent_inputs=[InputParam.template("image_latents")], + additional_batch_inputs=[ + InputParam( + "fake_image_latents", + type_hint=torch.Tensor, + description="Fake image latents used as history seed for I2V generation.", + ), + ], + ), + HeliosAddNoiseToImageLatentsStep, + HeliosPrepareHistoryStep, + HeliosI2VSeedHistoryStep, + HeliosSetTimestepsStep, + HeliosI2VChunkDenoiseStep, + ] + block_names = [ + "input", + "additional_inputs", + "add_noise_image", + "prepare_history", + "seed_history", + "set_timesteps", + "chunk_denoise", + ] + + @property + def description(self): + return "I2V denoise block that seeds history with image latents and uses I2V-aware chunk preparation." + + @property + def outputs(self): + return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")] + + +# DENOISE (V2V) +# auto_docstring +class HeliosV2VCoreDenoiseStep(SequentialPipelineBlocks): + """ + V2V denoise block that seeds history with video latents and uses I2V-aware chunk preparation. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + video_latents (`Tensor`, *optional*): + Encoded video latents for V2V generation. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video latent noise. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + model_name = "helios" + block_classes = [ + HeliosTextInputStep, + HeliosAdditionalInputsStep( + image_latent_inputs=[InputParam.template("image_latents")], + additional_batch_inputs=[ + InputParam( + "video_latents", type_hint=torch.Tensor, description="Encoded video latents for V2V generation." + ), + ], + ), + HeliosAddNoiseToVideoLatentsStep, + HeliosPrepareHistoryStep, + HeliosV2VSeedHistoryStep, + HeliosSetTimestepsStep, + HeliosI2VChunkDenoiseStep, + ] + block_names = [ + "input", + "additional_inputs", + "add_noise_video", + "prepare_history", + "seed_history", + "set_timesteps", + "chunk_denoise", + ] + + @property + def description(self): + return "V2V denoise block that seeds history with video latents and uses I2V-aware chunk preparation." + + @property + def outputs(self): + return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")] + + +# AUTO DENOISE +# auto_docstring +class HeliosAutoCoreDenoiseStep(ConditionalPipelineBlocks): + """ + Core denoise step that selects the appropriate denoising block. + - `HeliosV2VCoreDenoiseStep` (video2video) for video-to-video tasks. + - `HeliosI2VCoreDenoiseStep` (image2video) for image-to-video tasks. + - `HeliosCoreDenoiseStep` (text2video) for text-to-video tasks. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + video_latents (`Tensor`, *optional*): + Encoded video latents for V2V generation. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video latent noise. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + history_sizes (`list`): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`): + Custom sigmas for the denoising process. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + fake_image_latents (`Tensor`, *optional*): + Fake image latents used as history seed for I2V generation. + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + block_classes = [HeliosV2VCoreDenoiseStep, HeliosI2VCoreDenoiseStep, HeliosCoreDenoiseStep] + block_names = ["video2video", "image2video", "text2video"] + block_trigger_inputs = ["video_latents", "fake_image_latents"] + default_block_name = "text2video" + + def select_block(self, video_latents=None, fake_image_latents=None): + if video_latents is not None: + return "video2video" + elif fake_image_latents is not None: + return "image2video" + return None + + @property + def description(self): + return ( + "Core denoise step that selects the appropriate denoising block.\n" + " - `HeliosV2VCoreDenoiseStep` (video2video) for video-to-video tasks.\n" + " - `HeliosI2VCoreDenoiseStep` (image2video) for image-to-video tasks.\n" + " - `HeliosCoreDenoiseStep` (text2video) for text-to-video tasks." + ) + + +AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", HeliosTextEncoderStep()), + ("vae_encoder", HeliosAutoVaeEncoderStep()), + ("denoise", HeliosAutoCoreDenoiseStep()), + ("decode", HeliosDecodeStep()), + ] +) + +# ==================== +# 3. Auto Blocks +# ==================== + + +# auto_docstring +class HeliosAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-video, image-to-video, and video-to-video tasks using Helios. + + Supported workflows: + - `text2video`: requires `prompt` + - `image2video`: requires `prompt`, `image` + - `video2video`: requires `prompt`, `video` + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) vae + (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`HeliosTransformer3DModel`) scheduler + (`HeliosScheduler`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for prompt encoding. + video (`None`, *optional*): + Input video for video-to-video generation + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + video_latents (`Tensor`, *optional*): + Encoded video latents for V2V generation. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video latent noise. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + history_sizes (`list`): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`): + Custom sigmas for the denoising process. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + fake_image_latents (`Tensor`, *optional*): + Fake image latents used as history seed for I2V generation. + output_type (`str`, *optional*, defaults to np): + Output format: 'pil', 'np', 'pt'. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "helios" + + block_classes = AUTO_BLOCKS.values() + block_names = AUTO_BLOCKS.keys() + + _workflow_map = { + "text2video": {"prompt": True}, + "image2video": {"prompt": True, "image": True}, + "video2video": {"prompt": True, "video": True}, + } + + @property + def description(self): + return "Auto Modular pipeline for text-to-video, image-to-video, and video-to-video tasks using Helios." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid.py b/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid.py new file mode 100644 index 0000000000000000000000000000000000000000..fea11786de21a5242f145c3eac47a9063cb8c333 --- /dev/null +++ b/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid.py @@ -0,0 +1,520 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam +from .before_denoise import ( + HeliosAdditionalInputsStep, + HeliosAddNoiseToImageLatentsStep, + HeliosAddNoiseToVideoLatentsStep, + HeliosI2VSeedHistoryStep, + HeliosPrepareHistoryStep, + HeliosTextInputStep, + HeliosV2VSeedHistoryStep, +) +from .decoders import HeliosDecodeStep +from .denoise import HeliosPyramidChunkDenoiseStep, HeliosPyramidI2VChunkDenoiseStep +from .encoders import HeliosImageVaeEncoderStep, HeliosTextEncoderStep, HeliosVideoVaeEncoderStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# ==================== +# 1. Vae Encoder +# ==================== + + +# auto_docstring +class HeliosPyramidAutoVaeEncoderStep(AutoPipelineBlocks): + """ + Encoder step that encodes video or image inputs. This is an auto pipeline block. + - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided. + - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided. + - If neither is provided, step will be skipped. + + Components: + vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Inputs: + video (`None`, *optional*): + Input video for video-to-video generation + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + + Outputs: + image_latents (`Tensor`): + The latent representation of the input image. + video_latents (`Tensor`): + Encoded video latents (chunked) + fake_image_latents (`Tensor`): + Fake image latents for history seeding + """ + + block_classes = [HeliosVideoVaeEncoderStep, HeliosImageVaeEncoderStep] + block_names = ["video_encoder", "image_encoder"] + block_trigger_inputs = ["video", "image"] + + @property + def description(self): + return ( + "Encoder step that encodes video or image inputs. This is an auto pipeline block.\n" + " - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided.\n" + " - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided.\n" + " - If neither is provided, step will be skipped." + ) + + +# ==================== +# 2. DENOISE +# ==================== + + +# DENOISE (T2V) +# auto_docstring +class HeliosPyramidCoreDenoiseStep(SequentialPipelineBlocks): + """ + T2V pyramid denoise block with progressive multi-resolution denoising. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider + (`ClassifierFreeZeroStarGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + Number of denoising steps per pyramid stage. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + model_name = "helios-pyramid" + block_classes = [ + HeliosTextInputStep, + HeliosPrepareHistoryStep, + HeliosPyramidChunkDenoiseStep, + ] + block_names = ["input", "prepare_history", "pyramid_chunk_denoise"] + + @property + def description(self): + return "T2V pyramid denoise block with progressive multi-resolution denoising." + + @property + def outputs(self): + return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")] + + +# DENOISE (I2V) +# auto_docstring +class HeliosPyramidI2VCoreDenoiseStep(SequentialPipelineBlocks): + """ + I2V pyramid denoise block with progressive multi-resolution denoising. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider + (`ClassifierFreeZeroStarGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + fake_image_latents (`Tensor`, *optional*): + Fake image latents used as history seed for I2V generation. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video/fake-image latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video/fake-image latent noise. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + Number of denoising steps per pyramid stage. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + model_name = "helios-pyramid" + block_classes = [ + HeliosTextInputStep, + HeliosAdditionalInputsStep( + image_latent_inputs=[InputParam.template("image_latents")], + additional_batch_inputs=[ + InputParam( + "fake_image_latents", + type_hint=torch.Tensor, + description="Fake image latents used as history seed for I2V generation.", + ), + ], + ), + HeliosAddNoiseToImageLatentsStep, + HeliosPrepareHistoryStep, + HeliosI2VSeedHistoryStep, + HeliosPyramidI2VChunkDenoiseStep, + ] + block_names = [ + "input", + "additional_inputs", + "add_noise_image", + "prepare_history", + "seed_history", + "pyramid_chunk_denoise", + ] + + @property + def description(self): + return "I2V pyramid denoise block with progressive multi-resolution denoising." + + @property + def outputs(self): + return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")] + + +# DENOISE (V2V) +# auto_docstring +class HeliosPyramidV2VCoreDenoiseStep(SequentialPipelineBlocks): + """ + V2V pyramid denoise block with progressive multi-resolution denoising. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider + (`ClassifierFreeZeroStarGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + video_latents (`Tensor`, *optional*): + Encoded video latents for V2V generation. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video latent noise. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + Number of denoising steps per pyramid stage. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + model_name = "helios-pyramid" + block_classes = [ + HeliosTextInputStep, + HeliosAdditionalInputsStep( + image_latent_inputs=[InputParam.template("image_latents")], + additional_batch_inputs=[ + InputParam( + "video_latents", type_hint=torch.Tensor, description="Encoded video latents for V2V generation." + ), + ], + ), + HeliosAddNoiseToVideoLatentsStep, + HeliosPrepareHistoryStep, + HeliosV2VSeedHistoryStep, + HeliosPyramidI2VChunkDenoiseStep, + ] + block_names = [ + "input", + "additional_inputs", + "add_noise_video", + "prepare_history", + "seed_history", + "pyramid_chunk_denoise", + ] + + @property + def description(self): + return "V2V pyramid denoise block with progressive multi-resolution denoising." + + @property + def outputs(self): + return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")] + + +# AUTO DENOISE +# auto_docstring +class HeliosPyramidAutoCoreDenoiseStep(ConditionalPipelineBlocks): + """ + Pyramid core denoise step that selects the appropriate denoising block. + - `HeliosPyramidV2VCoreDenoiseStep` (video2video) for video-to-video tasks. + - `HeliosPyramidI2VCoreDenoiseStep` (image2video) for image-to-video tasks. + - `HeliosPyramidCoreDenoiseStep` (text2video) for text-to-video tasks. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider + (`ClassifierFreeZeroStarGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + video_latents (`Tensor`, *optional*): + Encoded video latents for V2V generation. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video latent noise. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + history_sizes (`list`): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + Number of denoising steps per pyramid stage. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + fake_image_latents (`Tensor`, *optional*): + Fake image latents used as history seed for I2V generation. + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + block_classes = [HeliosPyramidV2VCoreDenoiseStep, HeliosPyramidI2VCoreDenoiseStep, HeliosPyramidCoreDenoiseStep] + block_names = ["video2video", "image2video", "text2video"] + block_trigger_inputs = ["video_latents", "fake_image_latents"] + default_block_name = "text2video" + + def select_block(self, video_latents=None, fake_image_latents=None): + if video_latents is not None: + return "video2video" + elif fake_image_latents is not None: + return "image2video" + return None + + @property + def description(self): + return ( + "Pyramid core denoise step that selects the appropriate denoising block.\n" + " - `HeliosPyramidV2VCoreDenoiseStep` (video2video) for video-to-video tasks.\n" + " - `HeliosPyramidI2VCoreDenoiseStep` (image2video) for image-to-video tasks.\n" + " - `HeliosPyramidCoreDenoiseStep` (text2video) for text-to-video tasks." + ) + + +# ==================== +# 3. Auto Blocks +# ==================== + +PYRAMID_AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", HeliosTextEncoderStep()), + ("vae_encoder", HeliosPyramidAutoVaeEncoderStep()), + ("denoise", HeliosPyramidAutoCoreDenoiseStep()), + ("decode", HeliosDecodeStep()), + ] +) + + +# auto_docstring +class HeliosPyramidAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for pyramid progressive generation (T2V/I2V/V2V) using Helios. + + Supported workflows: + - `text2video`: requires `prompt` + - `image2video`: requires `prompt`, `image` + - `video2video`: requires `prompt`, `video` + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) vae + (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`HeliosTransformer3DModel`) scheduler + (`HeliosScheduler`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for prompt encoding. + video (`None`, *optional*): + Input video for video-to-video generation + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + video_latents (`Tensor`, *optional*): + Encoded video latents for V2V generation. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video latent noise. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + history_sizes (`list`): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + Number of denoising steps per pyramid stage. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + fake_image_latents (`Tensor`, *optional*): + Fake image latents used as history seed for I2V generation. + output_type (`str`, *optional*, defaults to np): + Output format: 'pil', 'np', 'pt'. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "helios-pyramid" + + block_classes = PYRAMID_AUTO_BLOCKS.values() + block_names = PYRAMID_AUTO_BLOCKS.keys() + + _workflow_map = { + "text2video": {"prompt": True}, + "image2video": {"prompt": True, "image": True}, + "video2video": {"prompt": True, "video": True}, + } + + @property + def description(self): + return "Auto Modular pipeline for pyramid progressive generation (T2V/I2V/V2V) using Helios." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid_distilled.py b/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid_distilled.py new file mode 100644 index 0000000000000000000000000000000000000000..3e0b32f0df7e23f740208f17fe29eda95534aa6e --- /dev/null +++ b/diffusers/modular_pipelines/helios/modular_blocks_helios_pyramid_distilled.py @@ -0,0 +1,530 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam +from .before_denoise import ( + HeliosAdditionalInputsStep, + HeliosAddNoiseToImageLatentsStep, + HeliosAddNoiseToVideoLatentsStep, + HeliosI2VSeedHistoryStep, + HeliosPrepareHistoryStep, + HeliosTextInputStep, + HeliosV2VSeedHistoryStep, +) +from .decoders import HeliosDecodeStep +from .denoise import HeliosPyramidDistilledChunkDenoiseStep, HeliosPyramidDistilledI2VChunkDenoiseStep +from .encoders import HeliosImageVaeEncoderStep, HeliosTextEncoderStep, HeliosVideoVaeEncoderStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# ==================== +# 1. Vae Encoder +# ==================== + + +# auto_docstring +class HeliosPyramidDistilledAutoVaeEncoderStep(AutoPipelineBlocks): + """ + Encoder step for distilled pyramid pipeline. + - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided. + - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided. + - If neither is provided, step will be skipped. + + Components: + vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Inputs: + video (`None`, *optional*): + Input video for video-to-video generation + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + + Outputs: + image_latents (`Tensor`): + The latent representation of the input image. + video_latents (`Tensor`): + Encoded video latents (chunked) + fake_image_latents (`Tensor`): + Fake image latents for history seeding + """ + + block_classes = [HeliosVideoVaeEncoderStep, HeliosImageVaeEncoderStep] + block_names = ["video_encoder", "image_encoder"] + block_trigger_inputs = ["video", "image"] + + @property + def description(self): + return ( + "Encoder step for distilled pyramid pipeline.\n" + " - `HeliosVideoVaeEncoderStep` (video_encoder) is used when `video` is provided.\n" + " - `HeliosImageVaeEncoderStep` (image_encoder) is used when `image` is provided.\n" + " - If neither is provided, step will be skipped." + ) + + +# ==================== +# 2. DENOISE +# ==================== + + +# DENOISE (T2V) +# auto_docstring +class HeliosPyramidDistilledCoreDenoiseStep(SequentialPipelineBlocks): + """ + T2V distilled pyramid denoise block with DMD scheduler and no CFG. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + Number of denoising steps per pyramid stage. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + is_amplify_first_chunk (`bool`, *optional*, defaults to True): + Whether to double the first chunk's timesteps via the scheduler for amplified generation. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + model_name = "helios-pyramid" + block_classes = [ + HeliosTextInputStep, + HeliosPrepareHistoryStep, + HeliosPyramidDistilledChunkDenoiseStep, + ] + block_names = ["input", "prepare_history", "pyramid_chunk_denoise"] + + @property + def description(self): + return "T2V distilled pyramid denoise block with DMD scheduler and no CFG." + + @property + def outputs(self): + return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")] + + +# DENOISE (I2V) +# auto_docstring +class HeliosPyramidDistilledI2VCoreDenoiseStep(SequentialPipelineBlocks): + """ + I2V distilled pyramid denoise block with DMD scheduler and no CFG. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + fake_image_latents (`Tensor`, *optional*): + Fake image latents used as history seed for I2V generation. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video/fake-image latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video/fake-image latent noise. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + Number of denoising steps per pyramid stage. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + is_amplify_first_chunk (`bool`, *optional*, defaults to True): + Whether to double the first chunk's timesteps via the scheduler for amplified generation. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + model_name = "helios-pyramid" + block_classes = [ + HeliosTextInputStep, + HeliosAdditionalInputsStep( + image_latent_inputs=[InputParam.template("image_latents")], + additional_batch_inputs=[ + InputParam( + "fake_image_latents", + type_hint=torch.Tensor, + description="Fake image latents used as history seed for I2V generation.", + ), + ], + ), + HeliosAddNoiseToImageLatentsStep, + HeliosPrepareHistoryStep, + HeliosI2VSeedHistoryStep, + HeliosPyramidDistilledI2VChunkDenoiseStep, + ] + block_names = [ + "input", + "additional_inputs", + "add_noise_image", + "prepare_history", + "seed_history", + "pyramid_chunk_denoise", + ] + + @property + def description(self): + return "I2V distilled pyramid denoise block with DMD scheduler and no CFG." + + @property + def outputs(self): + return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")] + + +# DENOISE (V2V) +# auto_docstring +class HeliosPyramidDistilledV2VCoreDenoiseStep(SequentialPipelineBlocks): + """ + V2V distilled pyramid denoise block with DMD scheduler and no CFG. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + video_latents (`Tensor`, *optional*): + Encoded video latents for V2V generation. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video latent noise. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + history_sizes (`list`, *optional*, defaults to [16, 2, 1]): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + Number of denoising steps per pyramid stage. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + is_amplify_first_chunk (`bool`, *optional*, defaults to True): + Whether to double the first chunk's timesteps via the scheduler for amplified generation. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + model_name = "helios-pyramid" + block_classes = [ + HeliosTextInputStep, + HeliosAdditionalInputsStep( + image_latent_inputs=[InputParam.template("image_latents")], + additional_batch_inputs=[ + InputParam( + "video_latents", type_hint=torch.Tensor, description="Encoded video latents for V2V generation." + ), + ], + ), + HeliosAddNoiseToVideoLatentsStep, + HeliosPrepareHistoryStep, + HeliosV2VSeedHistoryStep, + HeliosPyramidDistilledI2VChunkDenoiseStep, + ] + block_names = [ + "input", + "additional_inputs", + "add_noise_video", + "prepare_history", + "seed_history", + "pyramid_chunk_denoise", + ] + + @property + def description(self): + return "V2V distilled pyramid denoise block with DMD scheduler and no CFG." + + @property + def outputs(self): + return [OutputParam("latent_chunks", type_hint=list, description="List of per-chunk denoised latent tensors")] + + +# AUTO DENOISE +# auto_docstring +class HeliosPyramidDistilledAutoCoreDenoiseStep(ConditionalPipelineBlocks): + """ + Distilled pyramid core denoise step that selects the appropriate denoising block. + - `HeliosPyramidDistilledV2VCoreDenoiseStep` (video2video) for video-to-video tasks. + - `HeliosPyramidDistilledI2VCoreDenoiseStep` (image2video) for image-to-video tasks. + - `HeliosPyramidDistilledCoreDenoiseStep` (text2video) for text-to-video tasks. + + Components: + transformer (`HeliosTransformer3DModel`) scheduler (`HeliosScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + video_latents (`Tensor`, *optional*): + Encoded video latents for V2V generation. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video latent noise. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + history_sizes (`list`): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + Number of denoising steps per pyramid stage. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + is_amplify_first_chunk (`bool`, *optional*, defaults to True): + Whether to double the first chunk's timesteps via the scheduler for amplified generation. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + fake_image_latents (`Tensor`, *optional*): + Fake image latents used as history seed for I2V generation. + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + + Outputs: + latent_chunks (`list`): + List of per-chunk denoised latent tensors + """ + + block_classes = [ + HeliosPyramidDistilledV2VCoreDenoiseStep, + HeliosPyramidDistilledI2VCoreDenoiseStep, + HeliosPyramidDistilledCoreDenoiseStep, + ] + block_names = ["video2video", "image2video", "text2video"] + block_trigger_inputs = ["video_latents", "fake_image_latents"] + default_block_name = "text2video" + + def select_block(self, video_latents=None, fake_image_latents=None): + if video_latents is not None: + return "video2video" + elif fake_image_latents is not None: + return "image2video" + return None + + @property + def description(self): + return ( + "Distilled pyramid core denoise step that selects the appropriate denoising block.\n" + " - `HeliosPyramidDistilledV2VCoreDenoiseStep` (video2video) for video-to-video tasks.\n" + " - `HeliosPyramidDistilledI2VCoreDenoiseStep` (image2video) for image-to-video tasks.\n" + " - `HeliosPyramidDistilledCoreDenoiseStep` (text2video) for text-to-video tasks." + ) + + +# ==================== +# 3. Auto Blocks +# ==================== + +DISTILLED_PYRAMID_AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", HeliosTextEncoderStep()), + ("vae_encoder", HeliosPyramidDistilledAutoVaeEncoderStep()), + ("denoise", HeliosPyramidDistilledAutoCoreDenoiseStep()), + ("decode", HeliosDecodeStep()), + ] +) + + +# auto_docstring +class HeliosPyramidDistilledAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for distilled pyramid progressive generation (T2V/I2V/V2V) using Helios. + + Supported workflows: + - `text2video`: requires `prompt` + - `image2video`: requires `prompt`, `image` + - `video2video`: requires `prompt`, `video` + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) vae + (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`HeliosTransformer3DModel`) scheduler + (`HeliosScheduler`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for prompt encoding. + video (`None`, *optional*): + Input video for video-to-video generation + height (`int`, *optional*, defaults to 384): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 640): + The width in pixels of the generated image. + num_latent_frames_per_chunk (`int`, *optional*, defaults to 9): + Number of latent frames per temporal chunk. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos to generate per prompt. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + video_latents (`Tensor`, *optional*): + Encoded video latents for V2V generation. + image_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for image latent noise. + image_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for image latent noise. + video_noise_sigma_min (`float`, *optional*, defaults to 0.111): + Minimum sigma for video latent noise. + video_noise_sigma_max (`float`, *optional*, defaults to 0.135): + Maximum sigma for video latent noise. + num_frames (`int`, *optional*, defaults to 132): + Total number of video frames to generate. + history_sizes (`list`): + Sizes of long/mid/short history buffers for temporal context. + keep_first_frame (`bool`, *optional*, defaults to True): + Whether to keep the first frame as a prefix in history. + pyramid_num_inference_steps_list (`list`, *optional*, defaults to [10, 10, 10]): + Number of denoising steps per pyramid stage. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + is_amplify_first_chunk (`bool`, *optional*, defaults to True): + Whether to double the first chunk's timesteps via the scheduler for amplified generation. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + fake_image_latents (`Tensor`, *optional*): + Fake image latents used as history seed for I2V generation. + output_type (`str`, *optional*, defaults to np): + Output format: 'pil', 'np', 'pt'. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "helios-pyramid" + + block_classes = DISTILLED_PYRAMID_AUTO_BLOCKS.values() + block_names = DISTILLED_PYRAMID_AUTO_BLOCKS.keys() + + _workflow_map = { + "text2video": {"prompt": True}, + "image2video": {"prompt": True, "image": True}, + "video2video": {"prompt": True, "video": True}, + } + + @property + def description(self): + return "Auto Modular pipeline for distilled pyramid progressive generation (T2V/I2V/V2V) using Helios." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/diffusers/modular_pipelines/helios/modular_pipeline.py b/diffusers/modular_pipelines/helios/modular_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..1fc338e67f05c6714bdb0634c516b732bb16256d --- /dev/null +++ b/diffusers/modular_pipelines/helios/modular_pipeline.py @@ -0,0 +1,87 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...loaders import HeliosLoraLoaderMixin +from ...utils import logging +from ..modular_pipeline import ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class HeliosModularPipeline( + ModularPipeline, + HeliosLoraLoaderMixin, +): + """ + A ModularPipeline for Helios text-to-video generation. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "HeliosAutoBlocks" + + @property + def vae_scale_factor_spatial(self): + vae_scale_factor = 8 + if hasattr(self, "vae") and self.vae is not None: + vae_scale_factor = self.vae.config.scale_factor_spatial + return vae_scale_factor + + @property + def vae_scale_factor_temporal(self): + vae_scale_factor = 4 + if hasattr(self, "vae") and self.vae is not None: + vae_scale_factor = self.vae.config.scale_factor_temporal + return vae_scale_factor + + @property + def num_channels_latents(self): + # YiYi TODO: find out default value + num_channels_latents = 16 + if hasattr(self, "transformer") and self.transformer is not None: + num_channels_latents = self.transformer.config.in_channels + return num_channels_latents + + @property + def requires_unconditional_embeds(self): + requires_unconditional_embeds = False + + if hasattr(self, "guider") and self.guider is not None: + requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1 + + return requires_unconditional_embeds + + +class HeliosPyramidModularPipeline(HeliosModularPipeline): + """ + A ModularPipeline for Helios pyramid (progressive resolution) video generation. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "HeliosPyramidAutoBlocks" + + +class HeliosPyramidDistilledModularPipeline(HeliosModularPipeline): + """ + A ModularPipeline for Helios distilled pyramid video generation using DMD scheduler. + + Uses guidance_scale=1.0 (no CFG) and supports is_amplify_first_chunk for the DMD scheduler. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "HeliosPyramidDistilledAutoBlocks" diff --git a/diffusers/modular_pipelines/hunyuan_video1_5/__init__.py b/diffusers/modular_pipelines/hunyuan_video1_5/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9c12e4a78ce2a4d5bc743176cc4199b01a7e035 --- /dev/null +++ b/diffusers/modular_pipelines/hunyuan_video1_5/__init__.py @@ -0,0 +1,49 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_hunyuan_video1_5"] = [ + "HunyuanVideo15AutoBlocks", + ] + _import_structure["modular_pipeline"] = ["HunyuanVideo15ModularPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_hunyuan_video1_5 import HunyuanVideo15AutoBlocks + from .modular_pipeline import HunyuanVideo15ModularPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/diffusers/modular_pipelines/hunyuan_video1_5/before_denoise.py b/diffusers/modular_pipelines/hunyuan_video1_5/before_denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..4c02eb9dd0846976db2d731a340a4c0197f36e92 --- /dev/null +++ b/diffusers/modular_pipelines/hunyuan_video1_5/before_denoise.py @@ -0,0 +1,324 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect + +import numpy as np +import torch + +from ...configuration_utils import FrozenDict +from ...models import HunyuanVideo15Transformer3DModel +from ...pipelines.hunyuan_video1_5.image_processor import HunyuanVideo15ImageProcessor +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import HunyuanVideo15ModularPipeline + + +logger = logging.get_logger(__name__) + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class HunyuanVideo15TextInputStep(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + @property + def description(self) -> str: + return "Input processing step that determines batch_size" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt_embeds"), + InputParam.template("batch_size", default=None), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("batch_size", type_hint=int), + ] + + @torch.no_grad() + def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.batch_size = getattr(block_state, "batch_size", None) or block_state.prompt_embeds.shape[0] + self.set_block_state(state, block_state) + return components, state + + +class HunyuanVideo15SetTimestepsStep(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def description(self) -> str: + return "Step that sets the scheduler's timesteps for inference" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps"), + InputParam.template("sigmas"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("timesteps", type_hint=torch.Tensor), + OutputParam("num_inference_steps", type_hint=int), + ] + + @torch.no_grad() + def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + sigmas = block_state.sigmas + if sigmas is None: + sigmas = np.linspace(1.0, 0.0, block_state.num_inference_steps + 1)[:-1] + + block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps( + components.scheduler, block_state.num_inference_steps, device, sigmas=sigmas + ) + + self.set_block_state(state, block_state) + return components, state + + +class HunyuanVideo15PrepareLatentsStep(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + @property + def description(self) -> str: + return "Prepare latents, conditioning latents, mask, and image_embeds for T2V" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", HunyuanVideo15Transformer3DModel), + ComponentSpec( + "video_processor", + HunyuanVideo15ImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("height"), + InputParam.template("width"), + InputParam("num_frames", type_hint=int, default=121, description="Number of video frames to generate."), + InputParam.template("latents"), + InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), + InputParam.template("generator"), + InputParam.template("batch_size", required=True, default=None), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="Pure noise latents"), + OutputParam("cond_latents_concat", type_hint=torch.Tensor), + OutputParam("mask_concat", type_hint=torch.Tensor), + OutputParam("image_embeds", type_hint=torch.Tensor), + ] + + @torch.no_grad() + def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + + height = block_state.height + width = block_state.width + if height is None and width is None: + height, width = components.video_processor.calculate_default_height_width( + components.default_aspect_ratio[1], components.default_aspect_ratio[0], components.target_size + ) + + batch_size = block_state.batch_size * block_state.num_videos_per_prompt + num_frames = block_state.num_frames + + latents = block_state.latents + if latents is not None: + latents = latents.to(device=device, dtype=dtype) + else: + shape = ( + batch_size, + components.num_channels_latents, + (num_frames - 1) // components.vae_scale_factor_temporal + 1, + int(height) // components.vae_scale_factor_spatial, + int(width) // components.vae_scale_factor_spatial, + ) + if isinstance(block_state.generator, list) and len(block_state.generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(block_state.generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=dtype) + + block_state.latents = latents + + b, c, f, h, w = latents.shape + block_state.cond_latents_concat = torch.zeros(b, c, f, h, w, dtype=dtype, device=device) + block_state.mask_concat = torch.zeros(b, 1, f, h, w, dtype=dtype, device=device) + + block_state.image_embeds = torch.zeros( + block_state.batch_size, + components.vision_num_semantic_tokens, + components.vision_states_dim, + dtype=dtype, + device=device, + ) + + self.set_block_state(state, block_state) + return components, state + + +class HunyuanVideo15Image2VideoPrepareLatentsStep(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + @property + def description(self) -> str: + return ( + "Prepare I2V conditioning from image_latents and image_embeds. " + "Expects pure noise `latents` from HunyuanVideo15PrepareLatentsStep. " + "Builds cond_latents_concat and mask_concat for the denoiser." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", HunyuanVideo15Transformer3DModel)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "image_latents", + type_hint=torch.Tensor, + required=True, + description="Pre-encoded image latents from the VAE encoder step, used as conditioning for I2V.", + ), + InputParam( + "image_embeds", + type_hint=torch.Tensor, + required=True, + description="Siglip image embeddings from the image encoder step, used as extra conditioning for I2V.", + ), + InputParam.template("latents", required=True), + InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), + InputParam.template("batch_size", required=True, default=None), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("cond_latents_concat", type_hint=torch.Tensor), + OutputParam("mask_concat", type_hint=torch.Tensor), + OutputParam("image_embeds", type_hint=torch.Tensor), + ] + + @torch.no_grad() + def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + + batch_size = block_state.batch_size * block_state.num_videos_per_prompt + + b, c, f, h, w = block_state.latents.shape + + latent_condition = block_state.image_latents.to(device=device, dtype=dtype) + latent_condition = latent_condition.repeat(batch_size, 1, f, 1, 1) + latent_condition[:, :, 1:, :, :] = 0 + block_state.cond_latents_concat = latent_condition + + latent_mask = torch.zeros(b, 1, f, h, w, dtype=dtype, device=device) + latent_mask[:, :, 0, :, :] = 1.0 + block_state.mask_concat = latent_mask + + image_embeds = block_state.image_embeds.to(device=device, dtype=dtype) + if image_embeds.shape[0] == 1 and batch_size > 1: + image_embeds = image_embeds.repeat(batch_size, 1, 1) + block_state.image_embeds = image_embeds + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/hunyuan_video1_5/decoders.py b/diffusers/modular_pipelines/hunyuan_video1_5/decoders.py new file mode 100644 index 0000000000000000000000000000000000000000..630af85c1b10f48f37ac43d1c98e5b9fc4264ceb --- /dev/null +++ b/diffusers/modular_pipelines/hunyuan_video1_5/decoders.py @@ -0,0 +1,70 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from ...configuration_utils import FrozenDict +from ...models import AutoencoderKLHunyuanVideo15 +from ...pipelines.hunyuan_video1_5.image_processor import HunyuanVideo15ImageProcessor +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +logger = logging.get_logger(__name__) + + +class HunyuanVideo15VaeDecoderStep(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLHunyuanVideo15), + ComponentSpec( + "video_processor", + HunyuanVideo15ImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def description(self) -> str: + return "Step that decodes the denoised latents into videos" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True), + InputParam.template("output_type", default="np"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("videos"), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + latents = block_state.latents.to(components.vae.dtype) / components.vae.config.scaling_factor + video = components.vae.decode(latents, return_dict=False)[0] + block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type) + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/hunyuan_video1_5/denoise.py b/diffusers/modular_pipelines/hunyuan_video1_5/denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..293fad57c93f41f469d09dab81b565b457a3a9d8 --- /dev/null +++ b/diffusers/modular_pipelines/hunyuan_video1_5/denoise.py @@ -0,0 +1,401 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import HunyuanVideo15Transformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + PipelineState, +) +from ..modular_pipeline_utils import ComponentSpec, InputParam +from .modular_pipeline import HunyuanVideo15ModularPipeline + + +logger = logging.get_logger(__name__) + + +class HunyuanVideo15LoopBeforeDenoiser(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + @property + def description(self) -> str: + return "Step within the denoising loop that prepares the latent input" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True), + InputParam("cond_latents_concat", required=True, type_hint=torch.Tensor), + InputParam("mask_concat", required=True, type_hint=torch.Tensor), + ] + + @torch.no_grad() + def __call__(self, components: HunyuanVideo15ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + block_state.latent_model_input = torch.cat( + [block_state.latents, block_state.cond_latents_concat, block_state.mask_concat], dim=1 + ) + return components, block_state + + +class HunyuanVideo15LoopDenoiser(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + def __init__(self, guider_input_fields=None): + if guider_input_fields is None: + guider_input_fields = { + "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"), + "encoder_attention_mask": ("prompt_embeds_mask", "negative_prompt_embeds_mask"), + "encoder_hidden_states_2": ("prompt_embeds_2", "negative_prompt_embeds_2"), + "encoder_attention_mask_2": ("prompt_embeds_mask_2", "negative_prompt_embeds_mask_2"), + } + if not isinstance(guider_input_fields, dict): + raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}") + self._guider_input_fields = guider_input_fields + super().__init__() + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 7.5}), + default_creation_method="from_config", + ), + ComponentSpec("transformer", HunyuanVideo15Transformer3DModel), + ] + + @property + def description(self) -> str: + return "Step within the denoising loop that denoises the latents with guidance" + + @property + def inputs(self) -> list[InputParam]: + inputs = [ + InputParam.template("attention_kwargs"), + InputParam.template("num_inference_steps", required=True, default=None), + InputParam( + "image_embeds", + type_hint=torch.Tensor, + description="Siglip image embeddings used as extra conditioning for I2V. Zero-filled for T2V.", + ), + ] + for value in self._guider_input_fields.values(): + if isinstance(value, tuple): + inputs.append( + InputParam( + name=value[0], + required=True, + type_hint=torch.Tensor, + description=f"Positive branch of the {value[0]!r} field fed into the guider.", + ) + ) + for neg_name in value[1:]: + inputs.append( + InputParam( + name=neg_name, + type_hint=torch.Tensor, + description=f"Negative branch of the {neg_name!r} field fed into the guider.", + ) + ) + else: + inputs.append( + InputParam( + name=value, + required=True, + type_hint=torch.Tensor, + description=f"{value!r} field fed into the guider.", + ) + ) + return inputs + + @torch.no_grad() + def __call__( + self, components: HunyuanVideo15ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> PipelineState: + timestep = t.expand(block_state.latent_model_input.shape[0]).to(block_state.latent_model_input.dtype) + + # Step 1: Collect model inputs + guider_inputs = { + input_name: tuple(getattr(block_state, v) for v in value) + if isinstance(value, tuple) + else getattr(block_state, value) + for input_name, value in self._guider_input_fields.items() + } + + # Step 2: Update guider state + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + + # Step 3: Prepare batched inputs + guider_state = components.guider.prepare_inputs(guider_inputs) + + # Step 4: Run denoiser for each batch + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + + cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()} + + context_name = getattr(guider_state_batch, components.guider._identifier_key) + with components.transformer.cache_context(context_name): + guider_state_batch.noise_pred = components.transformer( + hidden_states=block_state.latent_model_input, + image_embeds=block_state.image_embeds, + timestep=timestep, + attention_kwargs=block_state.attention_kwargs, + return_dict=False, + **cond_kwargs, + )[0] + + components.guider.cleanup_models(components.transformer) + + # Step 5: Combine predictions + block_state.noise_pred = components.guider(guider_state)[0] + + return components, block_state + + +class HunyuanVideo15LoopAfterDenoiser(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def description(self) -> str: + return "Step within the denoising loop that updates the latents" + + @torch.no_grad() + def __call__(self, components: HunyuanVideo15ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + latents_dtype = block_state.latents.dtype + block_state.latents = components.scheduler.step( + block_state.noise_pred, t, block_state.latents, return_dict=False + )[0] + + if block_state.latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + block_state.latents = block_state.latents.to(latents_dtype) + + return components, block_state + + +class HunyuanVideo15DenoiseLoopWrapper(LoopSequentialPipelineBlocks): + model_name = "hunyuan-video-1.5" + + @property + def description(self) -> str: + return "Pipeline block that iteratively denoises the latents over timesteps" + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ComponentSpec("transformer", HunyuanVideo15Transformer3DModel), + ] + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam.template("timesteps", required=True), + InputParam.template("num_inference_steps", required=True, default=None), + ] + + @torch.no_grad() + def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + block_state.num_warmup_steps = max( + len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0 + ) + + with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, block_state = self.loop_step(components, block_state, i=i, t=t) + if i == len(block_state.timesteps) - 1 or ( + (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0 + ): + progress_bar.update() + + self.set_block_state(state, block_state) + return components, state + + +class HunyuanVideo15DenoiseStep(HunyuanVideo15DenoiseLoopWrapper): + block_classes = [ + HunyuanVideo15LoopBeforeDenoiser, + HunyuanVideo15LoopDenoiser(), + HunyuanVideo15LoopAfterDenoiser, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoises the latents.\n" + "At each iteration:\n" + " - `HunyuanVideo15LoopBeforeDenoiser`\n" + " - `HunyuanVideo15LoopDenoiser`\n" + " - `HunyuanVideo15LoopAfterDenoiser`\n" + "This block supports text-to-video tasks." + ) + + +class HunyuanVideo15Image2VideoLoopDenoiser(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + def __init__(self, guider_input_fields=None): + if guider_input_fields is None: + guider_input_fields = { + "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"), + "encoder_attention_mask": ("prompt_embeds_mask", "negative_prompt_embeds_mask"), + "encoder_hidden_states_2": ("prompt_embeds_2", "negative_prompt_embeds_2"), + "encoder_attention_mask_2": ("prompt_embeds_mask_2", "negative_prompt_embeds_mask_2"), + } + if not isinstance(guider_input_fields, dict): + raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}") + self._guider_input_fields = guider_input_fields + super().__init__() + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 7.5}), + default_creation_method="from_config", + ), + ComponentSpec("transformer", HunyuanVideo15Transformer3DModel), + ] + + @property + def description(self) -> str: + return "I2V denoiser with MeanFlow timestep_r support" + + @property + def inputs(self) -> list[InputParam]: + inputs = [ + InputParam.template("attention_kwargs"), + InputParam.template("num_inference_steps", required=True, default=None), + InputParam( + "image_embeds", + type_hint=torch.Tensor, + description="Siglip image embeddings used as extra conditioning for I2V. Zero-filled for T2V.", + ), + InputParam.template("timesteps", required=True), + ] + for value in self._guider_input_fields.values(): + if isinstance(value, tuple): + inputs.append( + InputParam( + name=value[0], + required=True, + type_hint=torch.Tensor, + description=f"Positive branch of the {value[0]!r} field fed into the guider.", + ) + ) + for neg_name in value[1:]: + inputs.append( + InputParam( + name=neg_name, + type_hint=torch.Tensor, + description=f"Negative branch of the {neg_name!r} field fed into the guider.", + ) + ) + else: + inputs.append( + InputParam( + name=value, + required=True, + type_hint=torch.Tensor, + description=f"{value!r} field fed into the guider.", + ) + ) + return inputs + + @torch.no_grad() + def __call__( + self, components: HunyuanVideo15ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> PipelineState: + timestep = t.expand(block_state.latent_model_input.shape[0]).to(block_state.latent_model_input.dtype) + + # MeanFlow timestep_r (lines 855-862) + if components.transformer.config.use_meanflow: + if i == len(block_state.timesteps) - 1: + timestep_r = torch.tensor([0.0], device=timestep.device) + else: + timestep_r = block_state.timesteps[i + 1] + timestep_r = timestep_r.expand(block_state.latents.shape[0]).to(block_state.latents.dtype) + else: + timestep_r = None + + guider_inputs = { + input_name: tuple(getattr(block_state, v) for v in value) + if isinstance(value, tuple) + else getattr(block_state, value) + for input_name, value in self._guider_input_fields.items() + } + + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + + cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()} + + context_name = getattr(guider_state_batch, components.guider._identifier_key) + with components.transformer.cache_context(context_name): + guider_state_batch.noise_pred = components.transformer( + hidden_states=block_state.latent_model_input, + image_embeds=block_state.image_embeds, + timestep=timestep, + timestep_r=timestep_r, + attention_kwargs=block_state.attention_kwargs, + return_dict=False, + **cond_kwargs, + )[0] + + components.guider.cleanup_models(components.transformer) + + block_state.noise_pred = components.guider(guider_state)[0] + + return components, block_state + + +class HunyuanVideo15Image2VideoDenoiseStep(HunyuanVideo15DenoiseLoopWrapper): + block_classes = [ + HunyuanVideo15LoopBeforeDenoiser, + HunyuanVideo15Image2VideoLoopDenoiser(), + HunyuanVideo15LoopAfterDenoiser, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step for image-to-video with MeanFlow support.\n" + "At each iteration:\n" + " - `HunyuanVideo15LoopBeforeDenoiser`\n" + " - `HunyuanVideo15Image2VideoLoopDenoiser`\n" + " - `HunyuanVideo15LoopAfterDenoiser`" + ) diff --git a/diffusers/modular_pipelines/hunyuan_video1_5/encoders.py b/diffusers/modular_pipelines/hunyuan_video1_5/encoders.py new file mode 100644 index 0000000000000000000000000000000000000000..9d340cc88194c0322c8696d22ae04b1e4b47856a --- /dev/null +++ b/diffusers/modular_pipelines/hunyuan_video1_5/encoders.py @@ -0,0 +1,441 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re + +import torch +from transformers import ( + ByT5Tokenizer, + Qwen2_5_VLTextModel, + Qwen2TokenizerFast, + SiglipImageProcessor, + SiglipVisionModel, + T5EncoderModel, +) + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import AutoencoderKLHunyuanVideo15 +from ...pipelines.hunyuan_video1_5.image_processor import HunyuanVideo15ImageProcessor +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import HunyuanVideo15ModularPipeline + + +logger = logging.get_logger(__name__) + + +def format_text_input(prompt, system_message): + return [ + [{"role": "system", "content": system_message}, {"role": "user", "content": p if p else " "}] for p in prompt + ] + + +def extract_glyph_texts(prompt): + pattern = r"\"(.*?)\"|\"(.*?)\"" + matches = re.findall(pattern, prompt) + result = [match[0] or match[1] for match in matches] + result = list(dict.fromkeys(result)) if len(result) > 1 else result + if result: + formatted_result = ". ".join([f'Text "{text}"' for text in result]) + ". " + else: + formatted_result = None + return formatted_result + + +def _get_mllm_prompt_embeds( + text_encoder, + tokenizer, + prompt, + device, + tokenizer_max_length=1000, + num_hidden_layers_to_skip=2, + # fmt: off + system_message="You are a helpful assistant. Describe the video by detailing the following aspects: \ + 1. The main content and theme of the video. \ + 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \ + 3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \ + 4. background environment, light, style and atmosphere. \ + 5. camera angles, movements, and transitions used in the video.", + # fmt: on + crop_start=108, +): + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt = format_text_input(prompt, system_message) + + text_inputs = tokenizer.apply_chat_template( + prompt, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + padding="max_length", + max_length=tokenizer_max_length + crop_start, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids.to(device=device) + prompt_attention_mask = text_inputs.attention_mask.to(device=device) + + prompt_embeds = text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_attention_mask, + output_hidden_states=True, + ).hidden_states[-(num_hidden_layers_to_skip + 1)] + + if crop_start is not None and crop_start > 0: + prompt_embeds = prompt_embeds[:, crop_start:] + prompt_attention_mask = prompt_attention_mask[:, crop_start:] + + return prompt_embeds, prompt_attention_mask + + +def _get_byt5_prompt_embeds(tokenizer, text_encoder, prompt, device, tokenizer_max_length=256): + prompt = [prompt] if isinstance(prompt, str) else prompt + glyph_texts = [extract_glyph_texts(p) for p in prompt] + + prompt_embeds_list = [] + prompt_embeds_mask_list = [] + + for glyph_text in glyph_texts: + if glyph_text is None: + glyph_text_embeds = torch.zeros( + (1, tokenizer_max_length, text_encoder.config.d_model), device=device, dtype=text_encoder.dtype + ) + glyph_text_embeds_mask = torch.zeros((1, tokenizer_max_length), device=device, dtype=torch.int64) + else: + txt_tokens = tokenizer( + glyph_text, + padding="max_length", + max_length=tokenizer_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ).to(device) + + glyph_text_embeds = text_encoder( + input_ids=txt_tokens.input_ids, + attention_mask=txt_tokens.attention_mask.float(), + )[0] + glyph_text_embeds = glyph_text_embeds.to(device=device) + glyph_text_embeds_mask = txt_tokens.attention_mask.to(device=device) + + prompt_embeds_list.append(glyph_text_embeds) + prompt_embeds_mask_list.append(glyph_text_embeds_mask) + + return torch.cat(prompt_embeds_list, dim=0), torch.cat(prompt_embeds_mask_list, dim=0) + + +class HunyuanVideo15TextEncoderStep(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + @property + def description(self) -> str: + return "Dual text encoder step using Qwen2.5-VL (MLLM) and ByT5 (glyph text)" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen2_5_VLTextModel), + ComponentSpec("tokenizer", Qwen2TokenizerFast), + ComponentSpec("text_encoder_2", T5EncoderModel), + ComponentSpec("tokenizer_2", ByT5Tokenizer), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 7.5}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", required=False), + InputParam.template("negative_prompt"), + InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("prompt_embeds_mask"), + OutputParam.template("negative_prompt_embeds"), + OutputParam.template("negative_prompt_embeds_mask"), + OutputParam( + "prompt_embeds_2", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="ByT5 glyph-text embeddings used as a second conditioning stream for the transformer.", + ), + OutputParam( + "prompt_embeds_mask_2", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Attention mask for the ByT5 glyph-text embeddings.", + ), + OutputParam( + "negative_prompt_embeds_2", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="ByT5 glyph-text negative embeddings for classifier-free guidance.", + ), + OutputParam( + "negative_prompt_embeds_mask_2", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Attention mask for the ByT5 glyph-text negative embeddings.", + ), + ] + + @staticmethod + def encode_prompt( + components, + prompt, + device=None, + dtype=None, + batch_size=1, + num_videos_per_prompt=1, + ): + device = device or components._execution_device + dtype = dtype or components.text_encoder.dtype + + if prompt is None: + prompt = [""] * batch_size + prompt = [prompt] if isinstance(prompt, str) else prompt + + prompt_embeds, prompt_embeds_mask = _get_mllm_prompt_embeds( + tokenizer=components.tokenizer, + text_encoder=components.text_encoder, + prompt=prompt, + device=device, + tokenizer_max_length=components.tokenizer_max_length, + system_message=components.system_message, + crop_start=components.prompt_template_encode_start_idx, + ) + + prompt_embeds_2, prompt_embeds_mask_2 = _get_byt5_prompt_embeds( + tokenizer=components.tokenizer_2, + text_encoder=components.text_encoder_2, + prompt=prompt, + device=device, + tokenizer_max_length=components.tokenizer_2_max_length, + ) + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1).view( + batch_size * num_videos_per_prompt, seq_len, -1 + ) + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_videos_per_prompt, 1).view( + batch_size * num_videos_per_prompt, seq_len + ) + + _, seq_len_2, _ = prompt_embeds_2.shape + prompt_embeds_2 = prompt_embeds_2.repeat(1, num_videos_per_prompt, 1).view( + batch_size * num_videos_per_prompt, seq_len_2, -1 + ) + prompt_embeds_mask_2 = prompt_embeds_mask_2.repeat(1, num_videos_per_prompt, 1).view( + batch_size * num_videos_per_prompt, seq_len_2 + ) + + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + prompt_embeds_mask = prompt_embeds_mask.to(dtype=dtype, device=device) + prompt_embeds_2 = prompt_embeds_2.to(dtype=dtype, device=device) + prompt_embeds_mask_2 = prompt_embeds_mask_2.to(dtype=dtype, device=device) + + return prompt_embeds, prompt_embeds_mask, prompt_embeds_2, prompt_embeds_mask_2 + + @torch.no_grad() + def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + + prompt = block_state.prompt + negative_prompt = block_state.negative_prompt + num_videos_per_prompt = block_state.num_videos_per_prompt + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = 1 + + ( + block_state.prompt_embeds, + block_state.prompt_embeds_mask, + block_state.prompt_embeds_2, + block_state.prompt_embeds_mask_2, + ) = self.encode_prompt( + components, + prompt=prompt, + device=device, + dtype=dtype, + batch_size=batch_size, + num_videos_per_prompt=num_videos_per_prompt, + ) + + if components.requires_unconditional_embeds: + ( + block_state.negative_prompt_embeds, + block_state.negative_prompt_embeds_mask, + block_state.negative_prompt_embeds_2, + block_state.negative_prompt_embeds_mask_2, + ) = self.encode_prompt( + components, + prompt=negative_prompt, + device=device, + dtype=dtype, + batch_size=batch_size, + num_videos_per_prompt=num_videos_per_prompt, + ) + + state.set("batch_size", batch_size) + + self.set_block_state(state, block_state) + return components, state + + +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") + + +class HunyuanVideo15VaeEncoderStep(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + @property + def description(self) -> str: + return "VAE Encoder step that encodes an input image into latent space for image-to-video generation" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLHunyuanVideo15), + ComponentSpec( + "video_processor", + HunyuanVideo15ImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image", required=True), + InputParam.template("height"), + InputParam.template("width"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "image_latents", + type_hint=torch.Tensor, + description="Encoded image latents from the VAE encoder", + ), + OutputParam("height", type_hint=int, description="Target height resolved from image"), + OutputParam("width", type_hint=int, description="Target width resolved from image"), + ] + + @torch.no_grad() + def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + image = block_state.image + height = block_state.height + width = block_state.width + if height is None or width is None: + height, width = components.video_processor.calculate_default_height_width( + height=image.size[1], width=image.size[0], target_size=components.target_size + ) + image = components.video_processor.resize(image, height=height, width=width, resize_mode="crop") + + vae_dtype = components.vae.dtype + image_tensor = components.video_processor.preprocess(image, height=height, width=width).to( + device=device, dtype=vae_dtype + ) + image_tensor = image_tensor.unsqueeze(2) + image_latents = retrieve_latents(components.vae.encode(image_tensor), sample_mode="argmax") + image_latents = image_latents * components.vae.config.scaling_factor + + block_state.image_latents = image_latents + block_state.height = height + block_state.width = width + state.set("image", image) + + self.set_block_state(state, block_state) + return components, state + + +class HunyuanVideo15ImageEncoderStep(ModularPipelineBlocks): + model_name = "hunyuan-video-1.5" + + @property + def description(self) -> str: + return "Siglip image encoder step that produces image_embeds for image-to-video generation" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("image_encoder", SiglipVisionModel), + ComponentSpec("feature_extractor", SiglipImageProcessor), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image", required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "image_embeds", + type_hint=torch.Tensor, + description="Image embeddings from the Siglip vision encoder", + ), + ] + + @torch.no_grad() + def __call__(self, components: HunyuanVideo15ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + image_encoder_dtype = next(components.image_encoder.parameters()).dtype + image_inputs = components.feature_extractor.preprocess( + images=block_state.image, do_resize=True, return_tensors="pt", do_convert_rgb=True + ) + image_inputs = image_inputs.to(device=device, dtype=image_encoder_dtype) + image_embeds = components.image_encoder(**image_inputs).last_hidden_state + + block_state.image_embeds = image_embeds + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/hunyuan_video1_5/modular_blocks_hunyuan_video1_5.py b/diffusers/modular_pipelines/hunyuan_video1_5/modular_blocks_hunyuan_video1_5.py new file mode 100644 index 0000000000000000000000000000000000000000..bdbdba1ecdd914e7301f73ca5a8a812f89c21342 --- /dev/null +++ b/diffusers/modular_pipelines/hunyuan_video1_5/modular_blocks_hunyuan_video1_5.py @@ -0,0 +1,535 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import ( + HunyuanVideo15Image2VideoPrepareLatentsStep, + HunyuanVideo15PrepareLatentsStep, + HunyuanVideo15SetTimestepsStep, + HunyuanVideo15TextInputStep, +) +from .decoders import HunyuanVideo15VaeDecoderStep +from .denoise import HunyuanVideo15DenoiseStep, HunyuanVideo15Image2VideoDenoiseStep +from .encoders import ( + HunyuanVideo15ImageEncoderStep, + HunyuanVideo15TextEncoderStep, + HunyuanVideo15VaeEncoderStep, +) + + +logger = logging.get_logger(__name__) + + +# auto_docstring +class HunyuanVideo15CoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block that takes encoded conditions and runs the denoising process. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`HunyuanVideo15Transformer3DModel`) + video_processor (`HunyuanVideo15ImageProcessor`) guider (`ClassifierFreeGuidance`) + + Inputs: + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + batch_size (`int`, *optional*): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + num_frames (`int`, *optional*, defaults to 121): + Number of video frames to generate. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + negative_prompt_embeds (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds' field fed into the guider. + prompt_embeds_mask (`Tensor`): + Positive branch of the 'prompt_embeds_mask' field fed into the guider. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds_mask' field fed into the guider. + prompt_embeds_2 (`Tensor`): + Positive branch of the 'prompt_embeds_2' field fed into the guider. + negative_prompt_embeds_2 (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds_2' field fed into the guider. + prompt_embeds_mask_2 (`Tensor`): + Positive branch of the 'prompt_embeds_mask_2' field fed into the guider. + negative_prompt_embeds_mask_2 (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds_mask_2' field fed into the guider. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "hunyuan-video-1.5" + block_classes = [ + HunyuanVideo15TextInputStep, + HunyuanVideo15SetTimestepsStep, + HunyuanVideo15PrepareLatentsStep, + HunyuanVideo15DenoiseStep, + ] + block_names = ["input", "set_timesteps", "prepare_latents", "denoise"] + + @property + def description(self): + return "Denoise block that takes encoded conditions and runs the denoising process." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# auto_docstring +class HunyuanVideo15Blocks(SequentialPipelineBlocks): + """ + Modular pipeline blocks for HunyuanVideo 1.5 text-to-video. + + Components: + text_encoder (`Qwen2_5_VLTextModel`) tokenizer (`Qwen2Tokenizer`) text_encoder_2 (`T5EncoderModel`) + tokenizer_2 (`ByT5Tokenizer`) guider (`ClassifierFreeGuidance`) scheduler (`FlowMatchEulerDiscreteScheduler`) + transformer (`HunyuanVideo15Transformer3DModel`) video_processor (`HunyuanVideo15ImageProcessor`) vae + (`AutoencoderKLHunyuanVideo15`) + + Inputs: + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`, *optional*): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + num_frames (`int`, *optional*, defaults to 121): + Number of video frames to generate. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to np): + Output format: 'pil', 'np', 'pt'. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "hunyuan-video-1.5" + block_classes = [ + HunyuanVideo15TextEncoderStep, + HunyuanVideo15CoreDenoiseStep, + HunyuanVideo15VaeDecoderStep, + ] + block_names = ["text_encoder", "denoise", "decode"] + + @property + def description(self): + return "Modular pipeline blocks for HunyuanVideo 1.5 text-to-video." + + @property + def outputs(self): + return [OutputParam.template("videos")] + + +# auto_docstring +class HunyuanVideo15Image2VideoCoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block for image-to-video that takes encoded conditions and runs the denoising process. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`HunyuanVideo15Transformer3DModel`) + video_processor (`HunyuanVideo15ImageProcessor`) guider (`ClassifierFreeGuidance`) + + Inputs: + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + batch_size (`int`, *optional*): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + num_frames (`int`, *optional*, defaults to 121): + Number of video frames to generate. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image_latents (`Tensor`): + Pre-encoded image latents from the VAE encoder step, used as conditioning for I2V. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + negative_prompt_embeds (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds' field fed into the guider. + prompt_embeds_mask (`Tensor`): + Positive branch of the 'prompt_embeds_mask' field fed into the guider. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds_mask' field fed into the guider. + prompt_embeds_2 (`Tensor`): + Positive branch of the 'prompt_embeds_2' field fed into the guider. + negative_prompt_embeds_2 (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds_2' field fed into the guider. + prompt_embeds_mask_2 (`Tensor`): + Positive branch of the 'prompt_embeds_mask_2' field fed into the guider. + negative_prompt_embeds_mask_2 (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds_mask_2' field fed into the guider. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "hunyuan-video-1.5" + block_classes = [ + HunyuanVideo15TextInputStep, + HunyuanVideo15SetTimestepsStep, + HunyuanVideo15PrepareLatentsStep, + HunyuanVideo15Image2VideoPrepareLatentsStep, + HunyuanVideo15Image2VideoDenoiseStep, + ] + block_names = ["input", "set_timesteps", "prepare_latents", "prepare_i2v_latents", "denoise"] + + @property + def description(self): + return "Denoise block for image-to-video that takes encoded conditions and runs the denoising process." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# auto_docstring +class HunyuanVideo15AutoVaeEncoderStep(AutoPipelineBlocks): + """ + VAE encoder step that encodes the image input into its latent representation. + This is an auto pipeline block that works for image-to-video tasks. + - `HunyuanVideo15VaeEncoderStep` is used when `image` is provided. + - If `image` is not provided, step will be skipped. + + Components: + vae (`AutoencoderKLHunyuanVideo15`) video_processor (`HunyuanVideo15ImageProcessor`) + + Inputs: + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + + Outputs: + image_latents (`Tensor`): + Encoded image latents from the VAE encoder + height (`int`): + Target height resolved from image + width (`int`): + Target width resolved from image + """ + + model_name = "hunyuan-video-1.5" + block_classes = [HunyuanVideo15VaeEncoderStep] + block_names = ["vae_encoder"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "VAE encoder step that encodes the image input into its latent representation.\n" + "This is an auto pipeline block that works for image-to-video tasks.\n" + " - `HunyuanVideo15VaeEncoderStep` is used when `image` is provided.\n" + " - If `image` is not provided, step will be skipped." + ) + + +# auto_docstring +class HunyuanVideo15AutoImageEncoderStep(AutoPipelineBlocks): + """ + Siglip image encoder step that produces image_embeds. + This is an auto pipeline block that works for image-to-video tasks. + - `HunyuanVideo15ImageEncoderStep` is used when `image` is provided. + - If `image` is not provided, step will be skipped. + + Components: + image_encoder (`SiglipVisionModel`) feature_extractor (`SiglipImageProcessor`) + + Inputs: + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + + Outputs: + image_embeds (`Tensor`): + Image embeddings from the Siglip vision encoder + """ + + model_name = "hunyuan-video-1.5" + block_classes = [HunyuanVideo15ImageEncoderStep] + block_names = ["image_encoder"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "Siglip image encoder step that produces image_embeds.\n" + "This is an auto pipeline block that works for image-to-video tasks.\n" + " - `HunyuanVideo15ImageEncoderStep` is used when `image` is provided.\n" + " - If `image` is not provided, step will be skipped." + ) + + +# auto_docstring +class HunyuanVideo15AutoCoreDenoiseStep(AutoPipelineBlocks): + """ + Auto denoise block that selects the appropriate denoise pipeline based on inputs. + - `HunyuanVideo15Image2VideoCoreDenoiseStep` is used when `image_latents` is provided. + - `HunyuanVideo15CoreDenoiseStep` is used otherwise (text-to-video). + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`HunyuanVideo15Transformer3DModel`) + video_processor (`HunyuanVideo15ImageProcessor`) guider (`ClassifierFreeGuidance`) + + Inputs: + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + num_inference_steps (`int`): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + num_frames (`int`, *optional*, defaults to 121): + Number of video frames to generate. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image_latents (`Tensor`, *optional*): + Pre-encoded image latents from the VAE encoder step, used as conditioning for I2V. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + negative_prompt_embeds (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds' field fed into the guider. + prompt_embeds_mask (`Tensor`): + Positive branch of the 'prompt_embeds_mask' field fed into the guider. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds_mask' field fed into the guider. + prompt_embeds_2 (`Tensor`): + Positive branch of the 'prompt_embeds_2' field fed into the guider. + negative_prompt_embeds_2 (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds_2' field fed into the guider. + prompt_embeds_mask_2 (`Tensor`): + Positive branch of the 'prompt_embeds_mask_2' field fed into the guider. + negative_prompt_embeds_mask_2 (`Tensor`, *optional*): + Negative branch of the 'negative_prompt_embeds_mask_2' field fed into the guider. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "hunyuan-video-1.5" + block_classes = [HunyuanVideo15Image2VideoCoreDenoiseStep, HunyuanVideo15CoreDenoiseStep] + block_names = ["image2video", "text2video"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Auto denoise block that selects the appropriate denoise pipeline based on inputs.\n" + " - `HunyuanVideo15Image2VideoCoreDenoiseStep` is used when `image_latents` is provided.\n" + " - `HunyuanVideo15CoreDenoiseStep` is used otherwise (text-to-video)." + ) + + +# auto_docstring +class HunyuanVideo15AutoBlocks(SequentialPipelineBlocks): + """ + Auto blocks for HunyuanVideo 1.5 that support both text-to-video and image-to-video workflows. + + Supported workflows: + - `text2video`: requires `prompt` + - `image2video`: requires `image`, `prompt` + + Components: + text_encoder (`Qwen2_5_VLTextModel`) tokenizer (`Qwen2Tokenizer`) text_encoder_2 (`T5EncoderModel`) + tokenizer_2 (`ByT5Tokenizer`) guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLHunyuanVideo15`) + video_processor (`HunyuanVideo15ImageProcessor`) image_encoder (`SiglipVisionModel`) feature_extractor + (`SiglipImageProcessor`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer + (`HunyuanVideo15Transformer3DModel`) + + Inputs: + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + batch_size (`int`): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + num_inference_steps (`int`): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + num_frames (`int`, *optional*, defaults to 121): + Number of video frames to generate. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image_latents (`Tensor`, *optional*): + Pre-encoded image latents from the VAE encoder step, used as conditioning for I2V. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to np): + Output format: 'pil', 'np', 'pt'. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "hunyuan-video-1.5" + block_classes = [ + HunyuanVideo15TextEncoderStep, + HunyuanVideo15AutoVaeEncoderStep, + HunyuanVideo15AutoImageEncoderStep, + HunyuanVideo15AutoCoreDenoiseStep, + HunyuanVideo15VaeDecoderStep, + ] + block_names = ["text_encoder", "vae_encoder", "image_encoder", "denoise", "decode"] + _workflow_map = { + "text2video": {"prompt": True}, + "image2video": {"image": True, "prompt": True}, + } + + @property + def description(self): + return "Auto blocks for HunyuanVideo 1.5 that support both text-to-video and image-to-video workflows." + + @property + def outputs(self): + return [OutputParam.template("videos")] + + +# auto_docstring +class HunyuanVideo15Image2VideoBlocks(SequentialPipelineBlocks): + """ + Modular pipeline blocks for HunyuanVideo 1.5 image-to-video. + + Components: + text_encoder (`Qwen2_5_VLTextModel`) tokenizer (`Qwen2Tokenizer`) text_encoder_2 (`T5EncoderModel`) + tokenizer_2 (`ByT5Tokenizer`) guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLHunyuanVideo15`) + video_processor (`HunyuanVideo15ImageProcessor`) image_encoder (`SiglipVisionModel`) feature_extractor + (`SiglipImageProcessor`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer + (`HunyuanVideo15Transformer3DModel`) + + Inputs: + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + batch_size (`int`, *optional*): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + num_frames (`int`, *optional*, defaults to 121): + Number of video frames to generate. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image_latents (`Tensor`): + Pre-encoded image latents from the VAE encoder step, used as conditioning for I2V. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to np): + Output format: 'pil', 'np', 'pt'. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "hunyuan-video-1.5" + block_classes = [ + HunyuanVideo15TextEncoderStep, + HunyuanVideo15AutoVaeEncoderStep, + HunyuanVideo15AutoImageEncoderStep, + HunyuanVideo15Image2VideoCoreDenoiseStep, + HunyuanVideo15VaeDecoderStep, + ] + block_names = ["text_encoder", "vae_encoder", "image_encoder", "denoise", "decode"] + + @property + def description(self): + return "Modular pipeline blocks for HunyuanVideo 1.5 image-to-video." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/diffusers/modular_pipelines/hunyuan_video1_5/modular_pipeline.py b/diffusers/modular_pipelines/hunyuan_video1_5/modular_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..e83aa33f201dc418f390d99a0427aa5fcb6f2ce6 --- /dev/null +++ b/diffusers/modular_pipelines/hunyuan_video1_5/modular_pipeline.py @@ -0,0 +1,90 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...loaders import HunyuanVideoLoraLoaderMixin +from ...utils import logging +from ..modular_pipeline import ModularPipeline + + +logger = logging.get_logger(__name__) + + +class HunyuanVideo15ModularPipeline( + ModularPipeline, + HunyuanVideoLoraLoaderMixin, +): + """ + A ModularPipeline for HunyuanVideo 1.5. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "HunyuanVideo15AutoBlocks" + + @property + def vae_scale_factor_spatial(self): + return self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 16 + + @property + def vae_scale_factor_temporal(self): + return self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 4 + + @property + def num_channels_latents(self): + return self.vae.config.latent_channels if getattr(self, "vae", None) else 32 + + @property + def target_size(self): + return self.transformer.config.target_size if getattr(self, "transformer", None) else 640 + + @property + def default_aspect_ratio(self): + return (16, 9) + + @property + def vision_num_semantic_tokens(self): + return 729 + + @property + def vision_states_dim(self): + return self.transformer.config.image_embed_dim if getattr(self, "transformer", None) else 1152 + + @property + def tokenizer_max_length(self): + return 1000 + + @property + def tokenizer_2_max_length(self): + return 256 + + # fmt: off + @property + def system_message(self): + return "You are a helpful assistant. Describe the video by detailing the following aspects: \ + 1. The main content and theme of the video. \ + 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. \ + 3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. \ + 4. background environment, light, style and atmosphere. \ + 5. camera angles, movements, and transitions used in the video." + # fmt: on + + @property + def prompt_template_encode_start_idx(self): + return 108 + + @property + def requires_unconditional_embeds(self): + if hasattr(self, "guider") and self.guider is not None: + return self.guider._enabled and self.guider.num_conditions > 1 + return False diff --git a/diffusers/modular_pipelines/ideogram4/__init__.py b/diffusers/modular_pipelines/ideogram4/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c7c733dda1418a238edeb0408a28b85b46a9dfa0 --- /dev/null +++ b/diffusers/modular_pipelines/ideogram4/__init__.py @@ -0,0 +1,47 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_ideogram4"] = ["Ideogram4AutoBlocks"] + _import_structure["modular_pipeline"] = ["Ideogram4ModularPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_ideogram4 import Ideogram4AutoBlocks + from .modular_pipeline import Ideogram4ModularPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/diffusers/modular_pipelines/ideogram4/before_denoise.py b/diffusers/modular_pipelines/ideogram4/before_denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..98be3b141aecf156722a4f2d767f8d11b0ff62b2 --- /dev/null +++ b/diffusers/modular_pipelines/ideogram4/before_denoise.py @@ -0,0 +1,558 @@ +# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import math + +import torch + +from ...models.transformers.transformer_ideogram4 import ( + IMAGE_POSITION_OFFSET, + LLM_TOKEN_INDICATOR, + OUTPUT_IMAGE_INDICATOR, + SEQUENCE_PADDING_INDICATOR, + Ideogram4Transformer2DModel, +) +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Ideogram4ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +# Default per-step guidance schedule (length must equal `num_inference_steps`): 7.0 for the main steps, +# dropping to 3.0 for the final 3 "polish" steps. +DEFAULT_GUIDANCE_SCHEDULE = (7.0,) * 45 + (3.0,) * 3 + + +# Copied from diffusers.pipelines.ideogram4.pipeline_ideogram4._logit_normal_sigmas +def _logit_normal_sigmas( + num_inference_steps: int, + mu: float, + std: float = 1.0, + logsnr_min: float = -15.0, + logsnr_max: float = 18.0, + device: torch.device | None = None, +) -> torch.Tensor: + r""" + Build a length-`num_inference_steps` sigma schedule using the Ideogram4 logit-normal flow-matching schedule. + + Sigmas are returned in `[0, 1]` in decreasing order (sigma close to 1 corresponds to pure noise, sigma close to 0 + to clean data), matching diffusers conventions. + + The Ideogram4 schedule applies `sigma(s) = 1 - logit_normal_cdf_inverse(1 - s)` to `s = linspace(0, 1, N + 1)` and + keeps the first `N` entries; a terminal zero is appended downstream by the scheduler. + """ + intervals = torch.linspace(0.0, 1.0, num_inference_steps + 1, dtype=torch.float64) + # Apply the inverse CDF of a normal then push through the logistic to obtain a logit-normal CDF inverse. + z = torch.special.ndtri(intervals) + y = mu + std * z + t = 1.0 - torch.special.expit(y) + t_min = 1.0 / (1.0 + math.exp(0.5 * logsnr_max)) + t_max = 1.0 / (1.0 + math.exp(0.5 * logsnr_min)) + t = t.clamp(t_min, t_max) + # Convert from model time (0 = noise, 1 = data) to diffusers sigma (1 = noise, 0 = data) and reverse. + sigmas = (1.0 - t).flip(0) + # Drop the trailing 0; FlowMatchEulerDiscreteScheduler.set_timesteps appends one back internally. + sigmas = sigmas[:-1].to(dtype=torch.float32, device=device) + return sigmas + + +# Copied from diffusers.pipelines.ideogram4.pipeline_ideogram4._resolution_aware_mu +def _resolution_aware_mu( + height: int, + width: int, + base_mu: float, + base_resolution: tuple[int, int] = (512, 512), +) -> float: + """Shift the schedule mean as a function of image resolution.""" + num_pixels = height * width + base_pixels = base_resolution[0] * base_resolution[1] + return base_mu + 0.5 * math.log(num_pixels / base_pixels) + + +# Copied from diffusers.pipelines.ideogram4.pipeline_ideogram4._expand_tensor_to_effective_batch +def _expand_tensor_to_effective_batch( + tensor: torch.Tensor, + batch_size: int, + num_per_prompt: int, + tensor_name: str | None = None, +) -> torch.Tensor: + """Replicate `tensor` along dim 0 from `batch_size` (or 1) to `batch_size * num_per_prompt`.""" + target_batch_size = batch_size * num_per_prompt + + if tensor.shape[0] == target_batch_size: + return tensor + + if tensor.shape[0] == 1: + repeat_by = target_batch_size + elif tensor.shape[0] == batch_size: + repeat_by = num_per_prompt + else: + tensor_name = f"`{tensor_name}`" if tensor_name is not None else "Tensor" + raise ValueError( + f"{tensor_name} batch size must be 1, `batch_size` ({batch_size}), or " + f"`batch_size * num_*_per_prompt` ({target_batch_size}), but got {tensor.shape[0]}." + ) + + return torch.repeat_interleave(tensor, repeats=repeat_by, dim=0, output_size=tensor.shape[0] * repeat_by) + + +# auto_docstring +class Ideogram4TextInputsStep(ModularPipelineBlocks): + """ + Input step that determines `batch_size`/`dtype` from the per-prompt `text_features` and replicates the text outputs + to `batch_size * num_images_per_prompt`. Place after the text encoder. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + text_features (`Tensor`): + Per-prompt text features from the encoder. + text_lengths (`list`): + Per-prompt text-token counts from the encoder. + + Outputs: + batch_size (`int`): + Effective batch size (num prompts * num_images_per_prompt). + dtype (`dtype`): + The dtype of the text features. + text_features (`Tensor`): + Text features, batch-expanded. + text_lengths (`list`): + Text-token counts, batch-expanded. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return ( + "Input step that determines `batch_size`/`dtype` from the per-prompt `text_features` and replicates the " + "text outputs to `batch_size * num_images_per_prompt`. Place after the text encoder." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_images_per_prompt", default=1), + InputParam( + name="text_features", + required=True, + type_hint=torch.Tensor, + description="Per-prompt text features from the encoder.", + ), + InputParam( + name="text_lengths", + required=True, + type_hint=list, + description="Per-prompt text-token counts from the encoder.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="batch_size", + type_hint=int, + description="Effective batch size (num prompts * num_images_per_prompt).", + ), + OutputParam(name="dtype", type_hint=torch.dtype, description="The dtype of the text features."), + OutputParam(name="text_features", type_hint=torch.Tensor, description="Text features, batch-expanded."), + OutputParam(name="text_lengths", type_hint=list, description="Text-token counts, batch-expanded."), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + prompt_batch = block_state.text_features.shape[0] + num_per_prompt = block_state.num_images_per_prompt + + block_state.dtype = block_state.text_features.dtype + block_state.text_features = _expand_tensor_to_effective_batch( + block_state.text_features, prompt_batch, num_per_prompt, "text_features" + ) + block_state.text_lengths = [n for n in block_state.text_lengths for _ in range(num_per_prompt)] + block_state.batch_size = prompt_batch * num_per_prompt + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Ideogram4PrepareLatentsStep(ModularPipelineBlocks): + """ + Step that prepares the packed image latents (B, num_image_tokens, latent_dim) for the denoising loop. + + Components: + transformer (`Ideogram4Transformer2DModel`) + + Inputs: + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + batch_size (`int`): + Effective batch size. + + Outputs: + latents (`Tensor`): + The initial packed image latents (B, num_image_tokens, latent_dim). + num_image_tokens (`int`): + Number of image tokens (grid_h * grid_w). + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Step that prepares the packed image latents (B, num_image_tokens, latent_dim) for the denoising loop." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Ideogram4Transformer2DModel)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents"), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam.template("generator"), + InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="latents", + type_hint=torch.Tensor, + description="The initial packed image latents (B, num_image_tokens, latent_dim).", + ), + OutputParam( + name="num_image_tokens", type_hint=int, description="Number of image tokens (grid_h * grid_w)." + ), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + patch = components.patch_size + grid_h = block_state.height // (components.vae_scale_factor * patch) + grid_w = block_state.width // (components.vae_scale_factor * patch) + num_image_tokens = grid_h * grid_w + latent_dim = components.transformer.config.in_channels + + shape = (block_state.batch_size, num_image_tokens, latent_dim) + if block_state.latents is None: + block_state.latents = randn_tensor( + shape, generator=block_state.generator, device=device, dtype=torch.float32 + ) + else: + block_state.latents = block_state.latents.to(device=device, dtype=torch.float32) + + block_state.num_image_tokens = num_image_tokens + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Ideogram4SetTimestepsStep(ModularPipelineBlocks): + """ + Step that sets the resolution-aware logit-normal sigma schedule on the scheduler and resolves the per-step guidance + weights. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_inference_steps (`int`, *optional*, defaults to 48): + The number of denoising steps. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + mu (`float`, *optional*, defaults to 0.0): + Base mean of the logit-normal schedule. + std (`float`, *optional*, defaults to 1.5): + Std of the logit-normal schedule. + guidance_schedule (`list`, *optional*, defaults to (7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 3.0, 3.0, 3.0)): + Per-step guidance scale schedule (length num_inference_steps). + + Outputs: + timesteps (`Tensor`): + The denoising timesteps. + gw (`Tensor`): + Per-step guidance weights (num_inference_steps,). + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return ( + "Step that sets the resolution-aware logit-normal sigma schedule on the scheduler and resolves the " + "per-step guidance weights." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps", default=48), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam(name="mu", default=0.0, type_hint=float, description="Base mean of the logit-normal schedule."), + InputParam(name="std", default=1.5, type_hint=float, description="Std of the logit-normal schedule."), + InputParam( + name="guidance_schedule", + default=DEFAULT_GUIDANCE_SCHEDULE, + type_hint=list, + description="Per-step guidance scale schedule (length num_inference_steps).", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam(name="timesteps", type_hint=torch.Tensor, description="The denoising timesteps."), + OutputParam( + name="gw", type_hint=torch.Tensor, description="Per-step guidance weights (num_inference_steps,)." + ), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + if len(block_state.guidance_schedule) != block_state.num_inference_steps: + raise ValueError( + f"`guidance_schedule` must have length `num_inference_steps` ({block_state.num_inference_steps}), " + f"got {len(block_state.guidance_schedule)}." + ) + + schedule_mu = _resolution_aware_mu(height=block_state.height, width=block_state.width, base_mu=block_state.mu) + sigmas = _logit_normal_sigmas(block_state.num_inference_steps, schedule_mu, std=block_state.std, device=device) + components.scheduler.set_timesteps(sigmas=sigmas.tolist(), device=device) + + block_state.timesteps = components.scheduler.timesteps + block_state.gw = torch.as_tensor(block_state.guidance_schedule, dtype=torch.float32, device=device) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Ideogram4PrepareAdditionalInputsStep(ModularPipelineBlocks): + """ + Step that prepares the additional denoiser inputs from the packed-sequence layout: the conditional + encoder_hidden_states (text features packed with image padding) and the position_ids/segment_ids/indicator, plus + the unconditional (image-only) counterparts. Place after prepare_latents. + + Inputs: + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + text_features (`Tensor`): + Batch-expanded text features. + text_lengths (`list`): + Batch-expanded text-token counts. + batch_size (`int`): + Effective batch size. + + Outputs: + prompt_embeds (`Tensor`): + Packed conditional encoder_hidden_states (B, total_seq, dim). + position_ids (`Tensor`): + Conditional 3-axis MRoPE position ids. + segment_ids (`Tensor`): + Conditional block-diagonal segment ids. + indicator (`Tensor`): + Conditional per-token text/image/pad role. + negative_prompt_embeds (`Tensor`): + Unconditional (zeroed) text features (B, num_image_tokens, dim). + negative_position_ids (`Tensor`): + Unconditional position ids (image region). + negative_segment_ids (`Tensor`): + Unconditional segment ids (image region). + negative_indicator (`Tensor`): + Unconditional indicator (image region). + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return ( + "Step that prepares the additional denoiser inputs from the packed-sequence layout: the conditional " + "encoder_hidden_states (text features packed with image padding) and the position_ids/segment_ids/" + "indicator, plus the unconditional (image-only) counterparts. Place after prepare_latents." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam( + name="text_features", + required=True, + type_hint=torch.Tensor, + description="Batch-expanded text features.", + ), + InputParam( + name="text_lengths", required=True, type_hint=list, description="Batch-expanded text-token counts." + ), + InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="prompt_embeds", + type_hint=torch.Tensor, + description="Packed conditional encoder_hidden_states (B, total_seq, dim).", + ), + OutputParam( + name="position_ids", type_hint=torch.Tensor, description="Conditional 3-axis MRoPE position ids." + ), + OutputParam( + name="segment_ids", type_hint=torch.Tensor, description="Conditional block-diagonal segment ids." + ), + OutputParam( + name="indicator", type_hint=torch.Tensor, description="Conditional per-token text/image/pad role." + ), + OutputParam( + name="negative_prompt_embeds", + type_hint=torch.Tensor, + description="Unconditional (zeroed) text features (B, num_image_tokens, dim).", + ), + OutputParam( + name="negative_position_ids", + type_hint=torch.Tensor, + description="Unconditional position ids (image region).", + ), + OutputParam( + name="negative_segment_ids", + type_hint=torch.Tensor, + description="Unconditional segment ids (image region).", + ), + OutputParam( + name="negative_indicator", + type_hint=torch.Tensor, + description="Unconditional indicator (image region).", + ), + ] + + @staticmethod + # Copied from diffusers.pipelines.ideogram4.pipeline_ideogram4.Ideogram4Pipeline._prepare_ids + def _prepare_ids( + text_lengths: list[int], + grid_h: int, + grid_w: int, + max_text_tokens: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the packed `[left-pad][text][image]` layout from the per-prompt text lengths and the image grid. + + Returns `position_ids` (3-axis MRoPE), `segment_ids` (block-diagonal attention) and `indicator` (per-token + text/image/pad role). + """ + batch_size = len(text_lengths) + num_image_tokens = grid_h * grid_w + total_seq_len = max_text_tokens + num_image_tokens + + # Image position ids (t=0, h, w); offset keeps them disjoint from text positions. + h_idx = torch.arange(grid_h).view(-1, 1).expand(grid_h, grid_w).reshape(-1) + w_idx = torch.arange(grid_w).view(1, -1).expand(grid_h, grid_w).reshape(-1) + t_idx = torch.zeros_like(h_idx) + image_pos = torch.stack([t_idx, h_idx, w_idx], dim=1) + IMAGE_POSITION_OFFSET + + position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long) + segment_ids = torch.full((batch_size, total_seq_len), SEQUENCE_PADDING_INDICATOR, dtype=torch.long) + indicator = torch.zeros(batch_size, total_seq_len, dtype=torch.long) + + for b, num_text in enumerate(text_lengths): + offset = max_text_tokens - num_text + + text_pos = torch.arange(num_text) + text_pos_3d = torch.stack([text_pos, text_pos, text_pos], dim=1) + position_ids[b, offset : offset + num_text] = text_pos_3d + position_ids[b, offset + num_text :] = image_pos + + indicator[b, offset : offset + num_text] = LLM_TOKEN_INDICATOR + indicator[b, offset + num_text :] = OUTPUT_IMAGE_INDICATOR + + segment_ids[b, offset : offset + num_text + num_image_tokens] = 1 + + return position_ids.to(device), segment_ids.to(device), indicator.to(device) + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + patch = components.patch_size + grid_h = block_state.height // (components.vae_scale_factor * patch) + grid_w = block_state.width // (components.vae_scale_factor * patch) + num_image_tokens = grid_h * grid_w + + text_features = block_state.text_features + max_text_tokens = text_features.shape[1] + feature_dim = text_features.shape[-1] + + position_ids, segment_ids, indicator = self._prepare_ids( + block_state.text_lengths, grid_h, grid_w, max_text_tokens, device + ) + + # Pack the text features into the full sequence; image positions carry no text features. + image_feature_padding = torch.zeros( + block_state.batch_size, num_image_tokens, feature_dim, dtype=text_features.dtype, device=device + ) + block_state.prompt_embeds = torch.cat([text_features, image_feature_padding], dim=1) + + # Unconditional (image-only) branch, derived from the conditioning. + block_state.negative_prompt_embeds = torch.zeros( + block_state.batch_size, num_image_tokens, feature_dim, dtype=text_features.dtype, device=device + ) + block_state.position_ids = position_ids + block_state.segment_ids = segment_ids + block_state.indicator = indicator + block_state.negative_position_ids = position_ids[:, max_text_tokens:] + block_state.negative_segment_ids = segment_ids[:, max_text_tokens:] + block_state.negative_indicator = indicator[:, max_text_tokens:] + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/ideogram4/decoders.py b/diffusers/modular_pipelines/ideogram4/decoders.py new file mode 100644 index 0000000000000000000000000000000000000000..bf5d69270b7c15a8cfa573970b3e3d867384eded --- /dev/null +++ b/diffusers/modular_pipelines/ideogram4/decoders.py @@ -0,0 +1,112 @@ +# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from ...configuration_utils import FrozenDict +from ...image_processor import VaeImageProcessor +from ...models import AutoencoderKLFlux2 +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Ideogram4ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# auto_docstring +class Ideogram4DecodeStep(ModularPipelineBlocks): + """ + Step that decodes the unpatchified (B, ae_channels, H, W) latents into images: de-normalizes with the VAE + batch-norm statistics and decodes through the VAE. + + Components: + vae (`AutoencoderKLFlux2`) image_processor (`VaeImageProcessor`) + + Inputs: + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + latents (`Tensor`): + The unpatchified (B, ae_channels, H, W) latents to decode, from the after-denoise step. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return ( + "Step that decodes the unpatchified (B, ae_channels, H, W) latents into images: de-normalizes with the " + "VAE batch-norm statistics and decodes through the VAE." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLFlux2), + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("output_type", default="pil"), + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The unpatchified (B, ae_channels, H, W) latents to decode, from the after-denoise step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("images")] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + z = block_state.latents + patch = components.patch_size + ae_channels = z.shape[1] + grid_h, grid_w = z.shape[2] // patch, z.shape[3] // patch + + # VAE bn stores per-channel statistics over the packed channels, laid out as (patch_row, patch_col, + # ae_channel). Reshape them into an (ae_channels, patch, patch) tile and repeat across the grid so the + # denormalization on the unpatchified latents matches the packed-space statistics. + bn_mean = components.vae.bn.running_mean.view(patch, patch, ae_channels).permute(2, 0, 1) + bn_std = torch.sqrt(components.vae.bn.running_var + components.vae.config.batch_norm_eps) + bn_std = bn_std.view(patch, patch, ae_channels).permute(2, 0, 1) + bn_mean = bn_mean.repeat(1, grid_h, grid_w).to(device=z.device, dtype=z.dtype) + bn_std = bn_std.repeat(1, grid_h, grid_w).to(device=z.device, dtype=z.dtype) + z = z * bn_std + bn_mean + + decoded = components.vae.decode(z.to(components.vae.dtype), return_dict=False)[0] + block_state.images = components.image_processor.postprocess( + decoded.float(), output_type=block_state.output_type + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/ideogram4/denoise.py b/diffusers/modular_pipelines/ideogram4/denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..871db69d344c3383fa613709c7ef01e7643bbc3c --- /dev/null +++ b/diffusers/modular_pipelines/ideogram4/denoise.py @@ -0,0 +1,363 @@ +# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from ...models.transformers.transformer_ideogram4 import Ideogram4Transformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + PipelineState, +) +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Ideogram4ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class Ideogram4LoopBeforeDenoiser(ModularPipelineBlocks): + model_name = "ideogram4" + + @property + def description(self) -> str: + return ( + "Within the denoising loop: build the conditional packed input `[text-padding][image latents]` and the " + "model timestep. Compose into the `sub_blocks` of `Ideogram4DenoiseLoopWrapper`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."), + InputParam( + name="position_ids", required=True, type_hint=torch.Tensor, description="Conditional position ids." + ), + InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + # Conditional packed sequence is [text-padding][image latents]; text region length = total - image tokens. + max_text_tokens = block_state.position_ids.shape[1] - block_state.latents.shape[1] + text_z_padding = torch.zeros( + block_state.latents.shape[0], + max_text_tokens, + block_state.latents.shape[-1], + dtype=block_state.latents.dtype, + device=block_state.latents.device, + ) + block_state.pos_z = torch.cat([text_z_padding, block_state.latents], dim=1) + block_state.max_text_tokens = max_text_tokens + + # Map sigma-domain timestep to model time t in [0, 1] (0 = noise, 1 = clean data). + num_train_timesteps = components.scheduler.config.num_train_timesteps + t_model = 1.0 - (t.float() / num_train_timesteps) + block_state.t_model = t_model.expand(block_state.batch_size) + return components, block_state + + +class Ideogram4LoopDenoiser(ModularPipelineBlocks): + model_name = "ideogram4" + + @property + def description(self) -> str: + return ( + "Within the denoising loop: run the conditional `transformer` on the full packed sequence and the " + "`unconditional_transformer` on the image-only sequence, then blend with the per-step guidance weight " + "(asymmetric CFG, no guider). Compose into `Ideogram4DenoiseLoopWrapper`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", Ideogram4Transformer2DModel), + ComponentSpec("unconditional_transformer", Ideogram4Transformer2DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Packed conditional encoder_hidden_states.", + ), + InputParam( + name="position_ids", + required=True, + type_hint=torch.Tensor, + description="Conditional 3-axis MRoPE position ids.", + ), + InputParam( + name="segment_ids", + required=True, + type_hint=torch.Tensor, + description="Conditional block-diagonal segment ids.", + ), + InputParam( + name="indicator", + required=True, + type_hint=torch.Tensor, + description="Conditional per-token text/image/pad role.", + ), + InputParam( + name="negative_prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Unconditional (zeroed) text features.", + ), + InputParam( + name="negative_position_ids", + required=True, + type_hint=torch.Tensor, + description="Unconditional position ids (image region).", + ), + InputParam( + name="negative_segment_ids", + required=True, + type_hint=torch.Tensor, + description="Unconditional segment ids (image region).", + ), + InputParam( + name="negative_indicator", + required=True, + type_hint=torch.Tensor, + description="Unconditional indicator (image region).", + ), + InputParam(name="gw", required=True, type_hint=torch.Tensor, description="Per-step guidance weights."), + InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + transformer = components.transformer + unconditional_transformer = components.unconditional_transformer + + # Conditional pass operates on the full packed sequence; the velocity is the image-token region. + pos_out = transformer( + hidden_states=block_state.pos_z.to(transformer.dtype), + timestep=block_state.t_model.to(transformer.dtype), + encoder_hidden_states=block_state.prompt_embeds.to(transformer.dtype), + position_ids=block_state.position_ids, + segment_ids=block_state.segment_ids, + indicator=block_state.indicator, + return_dict=False, + )[0] + pos_v = pos_out[:, block_state.max_text_tokens :].to(torch.float32) + + # Unconditional pass uses the image-only positions with zeroed text features. + neg_v = unconditional_transformer( + hidden_states=block_state.latents.to(unconditional_transformer.dtype), + timestep=block_state.t_model.to(unconditional_transformer.dtype), + encoder_hidden_states=block_state.negative_prompt_embeds.to(unconditional_transformer.dtype), + position_ids=block_state.negative_position_ids, + segment_ids=block_state.negative_segment_ids, + indicator=block_state.negative_indicator, + return_dict=False, + )[0].to(torch.float32) + + gw_i = block_state.gw[i] + v = gw_i * pos_v + (1.0 - gw_i) * neg_v + # The scheduler integrates `-v` (Ideogram predicts velocity v = x0 - noise). + block_state.noise_pred = -v + return components, block_state + + +class Ideogram4LoopAfterDenoiser(ModularPipelineBlocks): + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Within the denoising loop: scheduler step. Compose into `Ideogram4DenoiseLoopWrapper`." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam(name="latents", type_hint=torch.Tensor, description="The denoised latents.")] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + block_state.latents = components.scheduler.step( + block_state.noise_pred, t, block_state.latents, return_dict=False + )[0] + return components, block_state + + +# auto_docstring +class Ideogram4DenoiseStep(LoopSequentialPipelineBlocks): + """ + Denoising loop that iteratively denoises the packed image latents over `timesteps`, running both the conditional + and unconditional transformers and blending with the per-step guidance schedule. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Ideogram4Transformer2DModel`) + unconditional_transformer (`Ideogram4Transformer2DModel`) + + Inputs: + timesteps (`Tensor`): + Denoising timesteps from set_timesteps. + num_inference_steps (`int`, *optional*, defaults to 48): + The number of denoising steps. + latents (`Tensor`): + Packed image latents. + position_ids (`Tensor`): + Conditional position ids. + batch_size (`int`): + Effective batch size. + prompt_embeds (`Tensor`): + Packed conditional encoder_hidden_states. + position_ids (`Tensor`): + Conditional 3-axis MRoPE position ids. + segment_ids (`Tensor`): + Conditional block-diagonal segment ids. + indicator (`Tensor`): + Conditional per-token text/image/pad role. + negative_prompt_embeds (`Tensor`): + Unconditional (zeroed) text features. + negative_position_ids (`Tensor`): + Unconditional position ids (image region). + negative_segment_ids (`Tensor`): + Unconditional segment ids (image region). + negative_indicator (`Tensor`): + Unconditional indicator (image region). + gw (`Tensor`): + Per-step guidance weights. + + Outputs: + latents (`Tensor`): + The denoised latents. + """ + + model_name = "ideogram4" + block_classes = [Ideogram4LoopBeforeDenoiser, Ideogram4LoopDenoiser, Ideogram4LoopAfterDenoiser] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoising loop that iteratively denoises the packed image latents over `timesteps`, running both the " + "conditional and unconditional transformers and blending with the per-step guidance schedule." + ) + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam( + name="timesteps", + required=True, + type_hint=torch.Tensor, + description="Denoising timesteps from set_timesteps.", + ), + InputParam.template("num_inference_steps", default=48), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, block_state = self.loop_step(components, block_state, i=i, t=t) + progress_bar.update() + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Ideogram4AfterDenoiseStep(ModularPipelineBlocks): + """ + Step that runs after the denoising loop: unpatchifies the packed image latents (B, num_image_tokens, ae_channels * + patch ** 2) into a (B, ae_channels, H, W) latent for the decoder. + + Inputs: + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + latents (`Tensor`): + The denoised packed image latents (B, num_image_tokens, latent_dim). + + Outputs: + latents (`Tensor`): + Unpatchified latents (B, ae_channels, H, W) ready for the VAE decoder. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return ( + "Step that runs after the denoising loop: unpatchifies the packed image latents " + "(B, num_image_tokens, ae_channels * patch ** 2) into a (B, ae_channels, H, W) latent for the decoder." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The denoised packed image latents (B, num_image_tokens, latent_dim).", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="latents", + type_hint=torch.Tensor, + description="Unpatchified latents (B, ae_channels, H, W) ready for the VAE decoder.", + ) + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + z = block_state.latents + patch = components.patch_size + grid_h = block_state.height // (components.vae_scale_factor * patch) + grid_w = block_state.width // (components.vae_scale_factor * patch) + + ae_channels = z.shape[-1] // (patch * patch) + z = z.view(z.shape[0], grid_h, grid_w, patch, patch, ae_channels) + z = z.permute(0, 5, 1, 3, 2, 4).contiguous() + z = z.view(z.shape[0], ae_channels, grid_h * patch, grid_w * patch) + + block_state.latents = z + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/ideogram4/encoders.py b/diffusers/modular_pipelines/ideogram4/encoders.py new file mode 100644 index 0000000000000000000000000000000000000000..6e149fa8392e2d21ad1948154751bfd57ce8b6a7 --- /dev/null +++ b/diffusers/modular_pipelines/ideogram4/encoders.py @@ -0,0 +1,327 @@ +# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +from transformers import Qwen2Tokenizer, Qwen3VLModel +from transformers.masking_utils import create_causal_mask + +from ...pipelines.ideogram4.prompt_enhancer import ( + PROMPT_UPSAMPLE_TEMPERATURE, + Ideogram4PromptEnhancerHead, + build_caption_logits_processor, + build_prompt_enhancer, + generate_captions, +) +from ...utils import is_outlines_available, logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Ideogram4ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Hidden states of these Qwen3-VL decoder layers are concatenated to form the per-token +# text conditioning consumed by the Ideogram4 transformer. +QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35) + + +# auto_docstring +class Ideogram4PromptUpsampleStep(ModularPipelineBlocks): + """ + Optional step that rewrites the prompt(s) into Ideogram4's native structured JSON caption when + `prompt_upsampling=True` (the format the model is trained on). Requires a generative `text_encoder` (a + `Qwen3VLForConditionalGeneration`); install `outlines` for schema-constrained captions. + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`Qwen2Tokenizer`): The tokenizer paired + with the text encoder. prompt_enhancer_head (`Ideogram4PromptEnhancerHead`): LM head grafted onto the text + encoder for prompt upsampling. + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + prompt_upsampling (`bool`, *optional*, defaults to False): + If True, rewrite the prompt into Ideogram4's native JSON caption before encoding. + prompt_upsampling_temperature (`float`, *optional*, defaults to 1.0): + Sampling temperature for prompt upsampling. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + max_sequence_length (`int`, *optional*, defaults to 2048): + Maximum sequence length for prompt encoding. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + prompt (`list`): + The (possibly upsampled) prompt forwarded to the text encoder. + """ + + model_name = "ideogram4" + + def __init__(self): + # Built lazily on first upsample: the head-less encoder body + `prompt_enhancer_head`, combined. + self._prompt_enhancer = None + # Outlines logits processor for schema-constrained captions; built lazily on first upsample. + self._caption_logits_processor = None + super().__init__() + + @property + def description(self) -> str: + return ( + "Optional step that rewrites the prompt(s) into Ideogram4's native structured JSON caption when " + "`prompt_upsampling=True` (the format the model is trained on). Requires a generative `text_encoder` " + "(a `Qwen3VLForConditionalGeneration`); install `outlines` for schema-constrained captions." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3VLModel, description="The Qwen3-VL text encoder."), + ComponentSpec("tokenizer", Qwen2Tokenizer, description="The tokenizer paired with the text encoder."), + ComponentSpec( + "prompt_enhancer_head", + Ideogram4PromptEnhancerHead, + description="LM head grafted onto the text encoder for prompt upsampling.", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", required=True), + InputParam( + name="prompt_upsampling", + type_hint=bool, + default=False, + description="If True, rewrite the prompt into Ideogram4's native JSON caption before encoding.", + ), + InputParam( + name="prompt_upsampling_temperature", + type_hint=float, + default=PROMPT_UPSAMPLE_TEMPERATURE, + description="Sampling temperature for prompt upsampling.", + ), + InputParam.template("height"), + InputParam.template("width"), + InputParam.template("max_sequence_length", default=2048), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="prompt", + type_hint=list, + description="The (possibly upsampled) prompt forwarded to the text encoder.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + if block_state.prompt_upsampling: + if components.prompt_enhancer_head is None: + raise ValueError( + "Prompt upsampling requires the `prompt_enhancer_head` component, which is not loaded. Load an " + "`Ideogram4PromptEnhancerHead` and add it to the pipeline." + ) + if self._prompt_enhancer is None: + self._prompt_enhancer = build_prompt_enhancer(components.text_encoder, components.prompt_enhancer_head) + if self._caption_logits_processor is None and is_outlines_available(): + self._caption_logits_processor = build_caption_logits_processor( + self._prompt_enhancer, components.tokenizer + ) + if self._caption_logits_processor is None: + logger.warning_once( + "`outlines` is not installed; prompt upsampling runs unconstrained and may not return " + "schema-valid JSON. Install with `pip install outlines` for structured captions." + ) + height = block_state.height or components.default_height + width = block_state.width or components.default_width + block_state.prompt = generate_captions( + self._prompt_enhancer, + components.tokenizer, + self._caption_logits_processor, + block_state.prompt, + height, + width, + temperature=block_state.prompt_upsampling_temperature, + max_new_tokens=block_state.max_sequence_length, + generator=block_state.generator, + device=components._execution_device, + ) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Ideogram4TextEncoderStep(ModularPipelineBlocks): + """ + Text encoder step that tokenizes the prompt(s) and runs the Qwen3-VL text encoder, returning the per-token text + features (concatenated from a fixed set of activation layers). Only the text tokens are encoded; the packed image + tokens are appended later (the encoder is causal with image after text, so they never affect the text features). + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`Qwen2Tokenizer`): The tokenizer paired + with the text encoder. + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + max_sequence_length (`int`, *optional*, defaults to 2048): + Maximum sequence length for prompt encoding. + + Outputs: + text_features (`Tensor`): + Per-prompt text features (B, max_sequence_length, llm_features_dim), padding zeroed. + text_lengths (`list`): + Per-prompt real text-token counts, used to lay out the packed sequence. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return ( + "Text encoder step that tokenizes the prompt(s) and runs the Qwen3-VL text encoder, returning the " + "per-token text features (concatenated from a fixed set of activation layers). Only the text tokens are " + "encoded; the packed image tokens are appended later (the encoder is causal with image after text, so " + "they never affect the text features)." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3VLModel, description="The Qwen3-VL text encoder."), + ComponentSpec("tokenizer", Qwen2Tokenizer, description="The tokenizer paired with the text encoder."), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", required=True), + InputParam.template("max_sequence_length", default=2048), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="text_features", + type_hint=torch.Tensor, + description="Per-prompt text features (B, max_sequence_length, llm_features_dim), padding zeroed.", + ), + OutputParam( + name="text_lengths", + type_hint=list, + description="Per-prompt real text-token counts, used to lay out the packed sequence.", + ), + ] + + @staticmethod + # Copied from diffusers.pipelines.ideogram4.pipeline_ideogram4.Ideogram4Pipeline._get_text_encoder_hidden_states + def _get_text_encoder_hidden_states( + text_encoder, + token_ids: torch.Tensor, + attention_mask: torch.Tensor, + pos_2d: torch.Tensor, + ) -> list[torch.Tensor]: + """Run the text encoder's decoder layers, returning the hidden states tapped at each activation layer.""" + + language_model = text_encoder.language_model + + inputs_embeds = language_model.embed_tokens(token_ids) + + position_ids_4d = pos_2d[None, ...].expand(4, pos_2d.shape[0], -1) + text_position_ids = position_ids_4d[0] + mrope_position_ids = position_ids_4d[1:] + + causal_mask = create_causal_mask( + config=language_model.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=None, + position_ids=text_position_ids, + ) + position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids) + + tap_set = set(QWEN3_VL_ACTIVATION_LAYERS) + captured: dict[int, torch.Tensor] = {} + hidden_states = inputs_embeds + for layer_idx, decoder_layer in enumerate(language_model.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=text_position_ids, + past_key_values=None, + position_embeddings=position_embeddings, + ) + if layer_idx in tap_set: + captured[layer_idx] = hidden_states + + return [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + tokenizer = components.tokenizer + max_text_tokens = block_state.max_sequence_length + + prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else list(block_state.prompt) + batch_size = len(prompts) + + # Tokenize each chat-formatted prompt and left-pad to `max_sequence_length`. + token_ids = torch.zeros(batch_size, max_text_tokens, dtype=torch.long) + attention_mask = torch.zeros(batch_size, max_text_tokens, dtype=torch.long) + text_position_ids = torch.zeros(batch_size, max_text_tokens, dtype=torch.long) + text_lengths = [] + for b, text_prompt in enumerate(prompts): + messages = [{"role": "user", "content": [{"type": "text", "text": text_prompt}]}] + text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + toks = tokenizer(text, return_tensors="pt", add_special_tokens=False)["input_ids"][0] + n = int(toks.shape[0]) + if n > max_text_tokens: + raise ValueError(f"prompt has {n} tokens, exceeds max_sequence_length={max_text_tokens}") + text_lengths.append(n) + offset = max_text_tokens - n + token_ids[b, offset:] = toks + attention_mask[b, offset:] = 1 + text_position_ids[b, offset:] = torch.arange(n) + + token_ids = token_ids.to(device) + attention_mask = attention_mask.to(device) + text_position_ids = text_position_ids.to(device) + + # Run the text encoder, tapping the activation-layer hidden states, then concatenate them into per-token + # text features (padding zeroed). + selected = self._get_text_encoder_hidden_states( + components.text_encoder, token_ids, attention_mask, text_position_ids + ) + text_features = torch.stack(selected, dim=0).permute(1, 2, 3, 0).reshape(batch_size, max_text_tokens, -1) + text_features = (text_features * attention_mask.to(text_features.dtype).unsqueeze(-1)).to(torch.float32) + + block_state.text_features = text_features + block_state.text_lengths = text_lengths + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py b/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py new file mode 100644 index 0000000000000000000000000000000000000000..0b788fe236be669d79e865101dd74b392db55f71 --- /dev/null +++ b/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py @@ -0,0 +1,185 @@ +# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + Ideogram4PrepareAdditionalInputsStep, + Ideogram4PrepareLatentsStep, + Ideogram4SetTimestepsStep, + Ideogram4TextInputsStep, +) +from .decoders import Ideogram4DecodeStep +from .denoise import Ideogram4AfterDenoiseStep, Ideogram4DenoiseStep +from .encoders import Ideogram4PromptUpsampleStep, Ideogram4TextEncoderStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Core denoise: consumes the per-prompt text features and produces the unpatchified latents +# (batch/latents/timesteps/ids inputs -> denoising loop -> unpatchify). +CORE_DENOISE_BLOCKS = InsertableDict( + [ + ("input", Ideogram4TextInputsStep()), + ("prepare_latents", Ideogram4PrepareLatentsStep()), + ("set_timesteps", Ideogram4SetTimestepsStep()), + ("prepare_additional_inputs", Ideogram4PrepareAdditionalInputsStep()), + ("denoise", Ideogram4DenoiseStep()), + ("after_denoise", Ideogram4AfterDenoiseStep()), + ] +) + + +# auto_docstring +class Ideogram4CoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoising workflow for Ideogram4 text-to-image: prepares the batch/latents/timesteps and the packed denoiser + inputs, runs the asymmetric-CFG denoising loop over the conditional and unconditional transformers, and + unpatchifies the result for the decoder. + + Components: + transformer (`Ideogram4Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + unconditional_transformer (`Ideogram4Transformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + text_features (`Tensor`): + Per-prompt text features from the encoder. + text_lengths (`list`): + Per-prompt text-token counts from the encoder. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 48): + The number of denoising steps. + mu (`float`, *optional*, defaults to 0.0): + Base mean of the logit-normal schedule. + std (`float`, *optional*, defaults to 1.5): + Std of the logit-normal schedule. + guidance_schedule (`list`, *optional*, defaults to (7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 3.0, 3.0, 3.0)): + Per-step guidance scale schedule (length num_inference_steps). + + Outputs: + latents (`Tensor`): + Unpatchified (B, ae_channels, H, W) latents. + """ + + model_name = "ideogram4" + block_classes = list(CORE_DENOISE_BLOCKS.values()) + block_names = list(CORE_DENOISE_BLOCKS.keys()) + + @property + def description(self) -> str: + return ( + "Core denoising workflow for Ideogram4 text-to-image: prepares the batch/latents/timesteps and the packed " + "denoiser inputs, runs the asymmetric-CFG denoising loop over the conditional and unconditional " + "transformers, and unpatchifies the result for the decoder." + ) + + @property + def outputs(self) -> list[OutputParam]: + # The only meaningful product of the core step is the unpatchified latents; the batch/timesteps/packed-sequence + # inputs prepared along the way are consumed within the loop and are not updated by it. + return [OutputParam.template("latents", description="Unpatchified (B, ae_channels, H, W) latents.")] + + +# auto_docstring +class Ideogram4AutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-image generation using Ideogram4: (optional) prompt upsampling -> encode text -> + core denoise (asymmetric CFG over two transformers) -> decode. + + Supported workflows: + - `text2image`: requires `prompt` + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`Qwen2Tokenizer`): The tokenizer paired + with the text encoder. prompt_enhancer_head (`Ideogram4PromptEnhancerHead`): LM head grafted onto the text + encoder for prompt upsampling. transformer (`Ideogram4Transformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) unconditional_transformer (`Ideogram4Transformer2DModel`) vae + (`AutoencoderKLFlux2`) image_processor (`VaeImageProcessor`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + prompt_upsampling (`bool`, *optional*, defaults to False): + If True, rewrite the prompt into Ideogram4's native JSON caption before encoding. + prompt_upsampling_temperature (`float`, *optional*, defaults to 1.0): + Sampling temperature for prompt upsampling. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + max_sequence_length (`int`, *optional*, defaults to 2048): + Maximum sequence length for prompt encoding. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + num_inference_steps (`int`, *optional*, defaults to 48): + The number of denoising steps. + mu (`float`, *optional*, defaults to 0.0): + Base mean of the logit-normal schedule. + std (`float`, *optional*, defaults to 1.5): + Std of the logit-normal schedule. + guidance_schedule (`list`, *optional*, defaults to (7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 3.0, 3.0, 3.0)): + Per-step guidance scale schedule (length num_inference_steps). + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "ideogram4" + block_classes = [ + Ideogram4PromptUpsampleStep(), + Ideogram4TextEncoderStep(), + Ideogram4CoreDenoiseStep(), + Ideogram4DecodeStep(), + ] + block_names = ["prompt_upsample", "text_encoder", "denoise", "decode"] + + # Workflow map declaring the trigger conditions for each supported workflow. + # `True` means the workflow triggers when the input is not None. + _workflow_map = { + "text2image": {"prompt": True}, + } + + @property + def description(self) -> str: + return ( + "Auto Modular pipeline for text-to-image generation using Ideogram4: (optional) prompt upsampling -> " + "encode text -> core denoise (asymmetric CFG over two transformers) -> decode." + ) + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("images")] diff --git a/diffusers/modular_pipelines/ideogram4/modular_pipeline.py b/diffusers/modular_pipelines/ideogram4/modular_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..9c0ff00b880ae97089127a04ebe83a4b34b772c8 --- /dev/null +++ b/diffusers/modular_pipelines/ideogram4/modular_pipeline.py @@ -0,0 +1,46 @@ +# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...loaders import Ideogram4LoraLoaderMixin +from ..modular_pipeline import ModularPipeline + + +class Ideogram4ModularPipeline(ModularPipeline, Ideogram4LoraLoaderMixin): + """ + A ModularPipeline for Ideogram4. + + > [!WARNING] > This is an experimental feature! + """ + + default_blocks_name = "Ideogram4AutoBlocks" + + # Ideogram4 patchifies the VAE output by a factor of 2 before feeding the transformer. + @property + def patch_size(self): + return 2 + + @property + def default_height(self): + return 2048 + + @property + def default_width(self): + return 2048 + + @property + def vae_scale_factor(self): + vae_scale_factor = 8 + if getattr(self, "vae", None) is not None: + vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + return vae_scale_factor diff --git a/diffusers/modular_pipelines/krea2/__init__.py b/diffusers/modular_pipelines/krea2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..12e51c7c3018b0e9e147183d6599935024b80b9d --- /dev/null +++ b/diffusers/modular_pipelines/krea2/__init__.py @@ -0,0 +1,49 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_krea2"] = ["Krea2AutoBlocks"] + _import_structure["modular_blocks_krea2_turbo"] = ["Krea2TurboAutoBlocks"] + _import_structure["modular_pipeline"] = ["Krea2ModularPipeline", "Krea2TurboModularPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_krea2 import Krea2AutoBlocks + from .modular_blocks_krea2_turbo import Krea2TurboAutoBlocks + from .modular_pipeline import Krea2ModularPipeline, Krea2TurboModularPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/diffusers/modular_pipelines/krea2/before_denoise.py b/diffusers/modular_pipelines/krea2/before_denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..63810d30a9035fa58b564dccdafaf550d92a2b27 --- /dev/null +++ b/diffusers/modular_pipelines/krea2/before_denoise.py @@ -0,0 +1,590 @@ +# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import numpy as np +import torch + +from ...models.transformers.transformer_krea2 import Krea2Transformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Krea2ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Copied from diffusers.pipelines.krea2.pipeline_krea2.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# auto_docstring +class Krea2TextInputsStep(ModularPipelineBlocks): + """ + Input step that determines `batch_size`/`dtype` from the per-prompt `prompt_embeds` and replicates the text + conditioning (and the optional negative branch) to `batch_size * num_images_per_prompt`. Place after the text + encoder. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + negative_prompt_embeds (`Tensor`, *optional*): + Per-prompt negative text features. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Per-prompt negative text mask. + + Outputs: + batch_size (`int`): + Effective batch size (num prompts * num_images_per_prompt). + dtype (`dtype`): + The dtype of the text features. + prompt_embeds (`Tensor`): + Text features, batch-expanded. + prompt_embeds_mask (`Tensor`): + Text mask, batch-expanded. + negative_prompt_embeds (`Tensor`): + Negative text features, batch-expanded. + negative_prompt_embeds_mask (`Tensor`): + Negative text mask, batch-expanded. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Input step that determines `batch_size`/`dtype` from the per-prompt `prompt_embeds` and replicates the " + "text conditioning (and the optional negative branch) to `batch_size * num_images_per_prompt`. Place after " + "the text encoder." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_images_per_prompt", default=1), + InputParam( + name="prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).", + ), + InputParam( + name="prompt_embeds_mask", + required=True, + type_hint=torch.Tensor, + description="Per-prompt boolean text mask (B, text_seq_len).", + ), + InputParam( + name="negative_prompt_embeds", + type_hint=torch.Tensor, + description="Per-prompt negative text features.", + ), + InputParam( + name="negative_prompt_embeds_mask", + type_hint=torch.Tensor, + description="Per-prompt negative text mask.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="batch_size", + type_hint=int, + description="Effective batch size (num prompts * num_images_per_prompt).", + ), + OutputParam(name="dtype", type_hint=torch.dtype, description="The dtype of the text features."), + OutputParam(name="prompt_embeds", type_hint=torch.Tensor, description="Text features, batch-expanded."), + OutputParam(name="prompt_embeds_mask", type_hint=torch.Tensor, description="Text mask, batch-expanded."), + OutputParam( + name="negative_prompt_embeds", + type_hint=torch.Tensor, + description="Negative text features, batch-expanded.", + ), + OutputParam( + name="negative_prompt_embeds_mask", + type_hint=torch.Tensor, + description="Negative text mask, batch-expanded.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + prompt_batch, seq_len, num_layers, dim = block_state.prompt_embeds.shape + n = block_state.num_images_per_prompt + + block_state.dtype = block_state.prompt_embeds.dtype + block_state.batch_size = prompt_batch * n + + block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, n, 1, 1).view( + prompt_batch * n, seq_len, num_layers, dim + ) + block_state.prompt_embeds_mask = block_state.prompt_embeds_mask.repeat(1, n).view(prompt_batch * n, seq_len) + + if block_state.negative_prompt_embeds is not None: + block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(1, n, 1, 1).view( + prompt_batch * n, seq_len, num_layers, dim + ) + block_state.negative_prompt_embeds_mask = block_state.negative_prompt_embeds_mask.repeat(1, n).view( + prompt_batch * n, seq_len + ) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2TurboTextInputsStep(ModularPipelineBlocks): + """ + Input step for the distilled Krea 2 turbo checkpoint that determines `batch_size`/`dtype` from the per-prompt + `prompt_embeds` and replicates the text conditioning to `batch_size * num_images_per_prompt`. The distilled + checkpoint runs without classifier-free guidance, so there is no negative branch. Place after the text encoder. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + + Outputs: + batch_size (`int`): + Effective batch size (num prompts * num_images_per_prompt). + dtype (`dtype`): + The dtype of the text features. + prompt_embeds (`Tensor`): + Text features, batch-expanded. + prompt_embeds_mask (`Tensor`): + Text mask, batch-expanded. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Input step for the distilled Krea 2 turbo checkpoint that determines `batch_size`/`dtype` from the " + "per-prompt `prompt_embeds` and replicates the text conditioning to `batch_size * num_images_per_prompt`. " + "The distilled checkpoint runs without classifier-free guidance, so there is no negative branch. Place " + "after the text encoder." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_images_per_prompt", default=1), + InputParam( + name="prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).", + ), + InputParam( + name="prompt_embeds_mask", + required=True, + type_hint=torch.Tensor, + description="Per-prompt boolean text mask (B, text_seq_len).", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="batch_size", + type_hint=int, + description="Effective batch size (num prompts * num_images_per_prompt).", + ), + OutputParam(name="dtype", type_hint=torch.dtype, description="The dtype of the text features."), + OutputParam(name="prompt_embeds", type_hint=torch.Tensor, description="Text features, batch-expanded."), + OutputParam(name="prompt_embeds_mask", type_hint=torch.Tensor, description="Text mask, batch-expanded."), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + prompt_batch, seq_len, num_layers, dim = block_state.prompt_embeds.shape + n = block_state.num_images_per_prompt + + block_state.dtype = block_state.prompt_embeds.dtype + block_state.batch_size = prompt_batch * n + + block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, n, 1, 1).view( + prompt_batch * n, seq_len, num_layers, dim + ) + block_state.prompt_embeds_mask = block_state.prompt_embeds_mask.repeat(1, n).view(prompt_batch * n, seq_len) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2PrepareLatentsStep(ModularPipelineBlocks): + """ + Step that samples the spatial image latents and patch-packs them into (B, image_seq_len, in_channels) for the + denoising loop. + + Components: + transformer (`Krea2Transformer2DModel`) + + Inputs: + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + batch_size (`int`): + Effective batch size. + dtype (`dtype`): + The working dtype. + + Outputs: + latents (`Tensor`): + The initial packed image latents (B, image_seq_len, in_channels). + image_seq_len (`int`): + Number of image tokens (grid_h * grid_w). + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Step that samples the spatial image latents and patch-packs them into (B, image_seq_len, in_channels) " + "for the denoising loop." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Krea2Transformer2DModel)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents"), + InputParam.template("height", default=1024), + InputParam.template("width", default=1024), + InputParam.template("generator"), + InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."), + InputParam(name="dtype", required=True, type_hint=torch.dtype, description="The working dtype."), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="latents", + type_hint=torch.Tensor, + description="The initial packed image latents (B, image_seq_len, in_channels).", + ), + OutputParam(name="image_seq_len", type_hint=int, description="Number of image tokens (grid_h * grid_w)."), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + p = components.patch_size + num_channels_latents = components.transformer.config.in_channels // (p**2) + + multiple = components.vae_scale_factor * components.patch_size + if block_state.height % multiple != 0 or block_state.width % multiple != 0: + rounded_height = ((block_state.height + multiple - 1) // multiple) * multiple + rounded_width = ((block_state.width + multiple - 1) // multiple) * multiple + logger.warning( + f"`height` and `width` must be multiples of {multiple}; rounding up from {block_state.height}x{block_state.width} to" + f" {rounded_height}x{rounded_width}." + ) + block_state.height, block_state.width = rounded_height, rounded_width + + latent_height = block_state.height // components.vae_scale_factor + latent_width = block_state.width // components.vae_scale_factor + + if block_state.latents is not None: + block_state.latents = block_state.latents.to(device=device, dtype=block_state.dtype) + else: + latents = randn_tensor( + (block_state.batch_size, num_channels_latents, latent_height, latent_width), + generator=block_state.generator, + device=device, + dtype=block_state.dtype, + ) + latents = latents.view( + block_state.batch_size, num_channels_latents, latent_height // p, p, latent_width // p, p + ) + latents = latents.permute(0, 2, 4, 1, 3, 5) + block_state.latents = latents.reshape( + block_state.batch_size, (latent_height // p) * (latent_width // p), num_channels_latents * p * p + ) + + block_state.image_seq_len = block_state.latents.shape[1] + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2SetTimestepsStep(ModularPipelineBlocks): + """ + Step that sets the Krea 2 flow-matching schedule on the scheduler: a linear sigma schedule with a resolution-aware + dynamic time shift `mu`. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_inference_steps (`int`, *optional*, defaults to 28): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + image_seq_len (`int`): + Number of image tokens, used to compute the resolution-aware shift. + + Outputs: + timesteps (`Tensor`): + The denoising timesteps. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Step that sets the Krea 2 flow-matching schedule on the scheduler: a linear sigma schedule with a " + "resolution-aware dynamic time shift `mu`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps", default=28), + InputParam( + name="sigmas", type_hint=list, description="Custom sigma schedule (defaults to a linear ramp)." + ), + InputParam( + name="image_seq_len", + required=True, + type_hint=int, + description="Number of image tokens, used to compute the resolution-aware shift.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam(name="timesteps", type_hint=torch.Tensor, description="The denoising timesteps.")] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + num_inference_steps = block_state.num_inference_steps + + sigmas = block_state.sigmas + if sigmas is None: + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + else: + block_state.num_inference_steps = len(sigmas) + + config = components.scheduler.config + mu = calculate_shift( + block_state.image_seq_len, + config.get("base_image_seq_len", 256), + config.get("max_image_seq_len", 6400), + config.get("base_shift", 0.5), + config.get("max_shift", 1.15), + ) + + components.scheduler.set_timesteps(sigmas=sigmas, mu=mu, device=device) + components.scheduler.set_begin_index(0) + block_state.timesteps = components.scheduler.timesteps + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2TurboSetTimestepsStep(ModularPipelineBlocks): + """ + Step that sets the flow-matching schedule for the distilled Krea 2 turbo checkpoint on the scheduler: a linear + sigma schedule with the fixed time shift `mu=1.15` the checkpoint was distilled with. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_inference_steps (`int`, *optional*, defaults to 8): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + + Outputs: + timesteps (`Tensor`): + The denoising timesteps. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Step that sets the flow-matching schedule for the distilled Krea 2 turbo checkpoint on the scheduler: a " + "linear sigma schedule with the fixed time shift `mu=1.15` the checkpoint was distilled with." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps", default=8), + InputParam( + name="sigmas", type_hint=list, description="Custom sigma schedule (defaults to a linear ramp)." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam(name="timesteps", type_hint=torch.Tensor, description="The denoising timesteps.")] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + num_inference_steps = block_state.num_inference_steps + + sigmas = block_state.sigmas + if sigmas is None: + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + else: + block_state.num_inference_steps = len(sigmas) + + components.scheduler.set_timesteps(sigmas=sigmas, mu=1.15, device=device) + components.scheduler.set_begin_index(0) + block_state.timesteps = components.scheduler.timesteps + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2PreparePositionIdsStep(ModularPipelineBlocks): + """ + Step that builds the shared rotary position ids for the combined [text | image] sequence: text at the origin, image + tokens at their (0, h, w) latent-grid coordinates. Place after prepare_latents. + + Inputs: + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + prompt_embeds (`Tensor`): + Batch-expanded text features (only text_seq_len is used). + + Outputs: + position_ids (`Tensor`): + Shared rotary coordinates (text_seq_len + grid_h * grid_w, 3). + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Step that builds the shared rotary position ids for the combined [text | image] sequence: text at the " + "origin, image tokens at their (0, h, w) latent-grid coordinates. Place after prepare_latents." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("height", default=1024), + InputParam.template("width", default=1024), + InputParam( + name="prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Batch-expanded text features (only text_seq_len is used).", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="position_ids", + type_hint=torch.Tensor, + description="Shared rotary coordinates (text_seq_len + grid_h * grid_w, 3).", + ) + ] + + @staticmethod + # Copied from diffusers.pipelines.krea2.pipeline_krea2.Krea2Pipeline.prepare_position_ids + def prepare_position_ids(text_seq_len: int, grid_height: int, grid_width: int, device: torch.device): + """Build the `(text_seq_len + grid_height * grid_width, 3)` rotary coordinates for the combined sequence: + text tokens sit at the origin, image tokens carry their `(0, h, w)` latent-grid coordinates.""" + text_ids = torch.zeros(text_seq_len, 3, device=device) + image_ids = torch.zeros(grid_height, grid_width, 3, device=device) + image_ids[..., 1] = torch.arange(grid_height, device=device)[:, None] + image_ids[..., 2] = torch.arange(grid_width, device=device)[None, :] + image_ids = image_ids.reshape(grid_height * grid_width, 3) + return torch.cat([text_ids, image_ids], dim=0) + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + p = components.patch_size + grid_h = block_state.height // (components.vae_scale_factor * p) + grid_w = block_state.width // (components.vae_scale_factor * p) + text_seq_len = block_state.prompt_embeds.shape[1] + + block_state.position_ids = self.prepare_position_ids(text_seq_len, grid_h, grid_w, device) + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/krea2/decoders.py b/diffusers/modular_pipelines/krea2/decoders.py new file mode 100644 index 0000000000000000000000000000000000000000..fd308b5ef64844470de7ebde464cce2e46f77fec --- /dev/null +++ b/diffusers/modular_pipelines/krea2/decoders.py @@ -0,0 +1,121 @@ +# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from ...configuration_utils import FrozenDict +from ...image_processor import VaeImageProcessor +from ...models import AutoencoderKLQwenImage +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Krea2ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# auto_docstring +class Krea2DecodeStep(ModularPipelineBlocks): + """ + Step that unpacks the denoised packed latents back to the spatial grid, de-normalizes them with the VAE's + per-channel statistics, and decodes them through the Qwen-Image VAE into images. + + Components: + vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + + Inputs: + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + latents (`Tensor`): + The denoised packed latents (B, image_seq_len, in_channels) from the denoising loop. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Step that unpacks the denoised packed latents back to the spatial grid, de-normalizes them with the " + "VAE's per-channel statistics, and decodes them through the Qwen-Image VAE into images." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLQwenImage), + ComponentSpec( + "image_processor", + VaeImageProcessor, + # Effective pixel-to-token downsampling factor: vae_scale_factor (8) * patch_size (2). + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("output_type", default="pil"), + InputParam.template("height", default=1024), + InputParam.template("width", default=1024), + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The denoised packed latents (B, image_seq_len, in_channels) from the denoising loop.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("images")] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + vae = components.vae + p = components.patch_size + latents = block_state.latents + + batch_size, _, channels = latents.shape + height = p * (int(block_state.height) // (components.vae_scale_factor * p)) + width = p * (int(block_state.width) // (components.vae_scale_factor * p)) + latents = latents.view(batch_size, height // p, width // p, channels // (p * p), p, p) + latents = latents.permute(0, 3, 1, 4, 2, 5) + latents = latents.reshape(batch_size, channels // (p * p), 1, height, width) + + latents = latents.to(vae.dtype) + latents_mean = ( + torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_std + latents_mean + image = vae.decode(latents, return_dict=False)[0][:, :, 0] + block_state.images = components.image_processor.postprocess(image, output_type=block_state.output_type) + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/krea2/denoise.py b/diffusers/modular_pipelines/krea2/denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..88c6cdca7aba093440d66b2cf65da4d278c01eec --- /dev/null +++ b/diffusers/modular_pipelines/krea2/denoise.py @@ -0,0 +1,369 @@ +# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models.transformers.transformer_krea2 import Krea2Transformer2DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + PipelineState, +) +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Krea2ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class Krea2LoopBeforeDenoiser(ModularPipelineBlocks): + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Within the denoising loop: normalize the scheduler timestep into the model's flow time and broadcast it " + "across the batch. Compose into the `sub_blocks` of a `Krea2DenoiseLoopWrapper`-based step." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."), + InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + num_train_timesteps = components.scheduler.config.num_train_timesteps + block_state.timestep = (t / num_train_timesteps).expand(block_state.batch_size) + return components, block_state + + +class Krea2LoopDenoiser(ModularPipelineBlocks): + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Within the denoising loop: run the `transformer` on the conditional (and, when the guider enables CFG, " + "the negative) text features and combine them through the `guider`. Compose into `Krea2DenoiseStep`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "guider", + ClassifierFreeGuidance, + # Krea 2 uses cond-anchored CFG (`cond + scale * (cond - uncond)`), which is the + # `use_original_formulation` branch of ClassifierFreeGuidance; scale 0 disables it (distilled TDM). + config=FrozenDict({"guidance_scale": 4.5, "use_original_formulation": True}), + default_creation_method="from_config", + ), + ComponentSpec("transformer", Krea2Transformer2DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."), + InputParam.template("num_inference_steps", required=True), + InputParam( + name="prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Conditional stacked text features.", + ), + InputParam( + name="prompt_embeds_mask", required=True, type_hint=torch.Tensor, description="Conditional text mask." + ), + InputParam( + name="position_ids", + required=True, + type_hint=torch.Tensor, + description="Shared rotary coordinates for the [text | image] sequence.", + ), + InputParam( + name="negative_prompt_embeds", type_hint=torch.Tensor, description="Negative stacked text features." + ), + InputParam(name="negative_prompt_embeds_mask", type_hint=torch.Tensor, description="Negative text mask."), + InputParam.template("attention_kwargs"), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + transformer = components.transformer + + latents = block_state.latents.to(transformer.dtype) + timestep = block_state.timestep.to(transformer.dtype) + + guider_inputs = { + "encoder_hidden_states": ( + block_state.prompt_embeds.to(transformer.dtype), + block_state.negative_prompt_embeds.to(transformer.dtype) + if block_state.negative_prompt_embeds is not None + else None, + ), + "encoder_attention_mask": ( + block_state.prompt_embeds_mask, + block_state.negative_prompt_embeds_mask, + ), + } + + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = {name: getattr(guider_state_batch, name) for name in guider_inputs} + guider_state_batch.noise_pred = transformer( + hidden_states=latents, + timestep=timestep, + position_ids=block_state.position_ids, + attention_kwargs=block_state.attention_kwargs, + return_dict=False, + **cond_kwargs, + )[0] + components.guider.cleanup_models(components.transformer) + + block_state.noise_pred = components.guider(guider_state).pred + return components, block_state + + +class Krea2TurboLoopDenoiser(ModularPipelineBlocks): + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Within the denoising loop: run the `transformer` on the conditional text features. The distilled Krea 2 " + "turbo checkpoint runs without classifier-free guidance, so there is no negative branch or guider. Compose " + "into the `sub_blocks` of `Krea2TurboDenoiseStep`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Krea2Transformer2DModel)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."), + InputParam( + name="prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Conditional stacked text features.", + ), + InputParam( + name="prompt_embeds_mask", required=True, type_hint=torch.Tensor, description="Conditional text mask." + ), + InputParam( + name="position_ids", + required=True, + type_hint=torch.Tensor, + description="Shared rotary coordinates for the [text | image] sequence.", + ), + InputParam.template("attention_kwargs"), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + transformer = components.transformer + + latents = block_state.latents.to(transformer.dtype) + timestep = block_state.timestep.to(transformer.dtype) + + block_state.noise_pred = transformer( + hidden_states=latents, + timestep=timestep, + position_ids=block_state.position_ids, + attention_kwargs=block_state.attention_kwargs, + encoder_hidden_states=block_state.prompt_embeds.to(transformer.dtype), + encoder_attention_mask=block_state.prompt_embeds_mask, + return_dict=False, + )[0] + return components, block_state + + +class Krea2LoopAfterDenoiser(ModularPipelineBlocks): + model_name = "krea2" + + @property + def description(self) -> str: + return "Within the denoising loop: scheduler step. Compose into a `Krea2DenoiseLoopWrapper`-based step." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam(name="latents", type_hint=torch.Tensor, description="The denoised latents.")] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + latents_dtype = block_state.latents.dtype + block_state.latents = components.scheduler.step( + block_state.noise_pred, t, block_state.latents, return_dict=False + )[0] + block_state.latents = block_state.latents.to(latents_dtype) + return components, block_state + + +class Krea2DenoiseLoopWrapper(LoopSequentialPipelineBlocks): + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Pipeline block that iteratively denoises the packed image latents over `timesteps`. " + "The specific steps within each iteration can be customized with the `sub_blocks` attribute." + ) + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam( + name="timesteps", + required=True, + type_hint=torch.Tensor, + description="Denoising timesteps from set_timesteps.", + ), + InputParam.template("num_inference_steps", required=True), + InputParam.template("attention_kwargs"), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, block_state = self.loop_step(components, block_state, i=i, t=t) + progress_bar.update() + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2DenoiseStep(Krea2DenoiseLoopWrapper): + """ + Denoising loop that iteratively denoises the packed image latents over `timesteps`, running the transformer on the + conditional (and, when the guider enables CFG, the negative) text features and combining them through the `guider`. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer + (`Krea2Transformer2DModel`) + + Inputs: + timesteps (`Tensor`): + Denoising timesteps from set_timesteps. + num_inference_steps (`int`): + The number of denoising steps. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + latents (`Tensor`): + Packed image latents. + batch_size (`int`): + Effective batch size. + prompt_embeds (`Tensor`): + Conditional stacked text features. + prompt_embeds_mask (`Tensor`): + Conditional text mask. + position_ids (`Tensor`): + Shared rotary coordinates for the [text | image] sequence. + negative_prompt_embeds (`Tensor`, *optional*): + Negative stacked text features. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Negative text mask. + + Outputs: + latents (`Tensor`): + The denoised latents. + """ + + model_name = "krea2" + block_classes = [Krea2LoopBeforeDenoiser, Krea2LoopDenoiser, Krea2LoopAfterDenoiser] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoising loop that iteratively denoises the packed image latents over `timesteps`, running the " + "transformer on the conditional (and, when the guider enables CFG, the negative) text features and " + "combining them through the `guider`." + ) + + +# auto_docstring +class Krea2TurboDenoiseStep(Krea2DenoiseLoopWrapper): + """ + Denoising loop for the distilled Krea 2 turbo checkpoint that iteratively denoises the packed image latents over + `timesteps`, running the transformer on the conditional text features. The distilled checkpoint runs without + classifier-free guidance. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Krea2Transformer2DModel`) + + Inputs: + timesteps (`Tensor`): + Denoising timesteps from set_timesteps. + num_inference_steps (`int`): + The number of denoising steps. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + latents (`Tensor`): + Packed image latents. + batch_size (`int`): + Effective batch size. + prompt_embeds (`Tensor`): + Conditional stacked text features. + prompt_embeds_mask (`Tensor`): + Conditional text mask. + position_ids (`Tensor`): + Shared rotary coordinates for the [text | image] sequence. + + Outputs: + latents (`Tensor`): + The denoised latents. + """ + + model_name = "krea2" + block_classes = [Krea2LoopBeforeDenoiser, Krea2TurboLoopDenoiser, Krea2LoopAfterDenoiser] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoising loop for the distilled Krea 2 turbo checkpoint that iteratively denoises the packed image " + "latents over `timesteps`, running the transformer on the conditional text features. The distilled " + "checkpoint runs without classifier-free guidance." + ) diff --git a/diffusers/modular_pipelines/krea2/encoders.py b/diffusers/modular_pipelines/krea2/encoders.py new file mode 100644 index 0000000000000000000000000000000000000000..7640222e9ad2bd69588447b6c8a8bee8dcbf63b4 --- /dev/null +++ b/diffusers/modular_pipelines/krea2/encoders.py @@ -0,0 +1,276 @@ +# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +from transformers import AutoTokenizer, Qwen3VLModel + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Krea2ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Indices into the Qwen3-VL `hidden_states` tuple (0 is the embedding output) whose states are stacked per token as the +# transformer's text conditioning. Must have `transformer.config.num_text_layers` entries. +KREA2_TEXT_ENCODER_SELECT_LAYERS = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35) + +# Krea 2 wraps the prompt in this Qwen-Image chat template before encoding. The prompt is padded to a fixed length +# first and the assistant suffix is appended *after* the padding (matching how the model was sampled at training time); +# the first `_PROMPT_TEMPLATE_ENCODE_START_IDX` (system prefix) tokens are dropped from the encoder outputs. +_PROMPT_TEMPLATE_ENCODE_PREFIX = ( + "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, " + "spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n" +) +_PROMPT_TEMPLATE_ENCODE_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n" +_PROMPT_TEMPLATE_ENCODE_START_IDX = 34 +_PROMPT_TEMPLATE_ENCODE_NUM_SUFFIX_TOKENS = 5 + + +# auto_docstring +class Krea2TextEncoderStep(ModularPipelineBlocks): + """ + Text encoder step that tokenizes the prompt(s) with the Krea 2 chat template, runs the Qwen3-VL text encoder, and + stacks a fixed set of decoder-layer hidden states per token as the transformer's text conditioning. The negative + prompt is encoded the same way when the guider enables CFG. + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`AutoTokenizer`): The tokenizer paired + with the text encoder. guider (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The negative prompt(s) for CFG. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for prompt encoding. + + Outputs: + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + negative_prompt_embeds (`Tensor`): + Per-prompt negative text features (only when guidance is enabled). + negative_prompt_embeds_mask (`Tensor`): + Per-prompt negative text mask (only when guidance is enabled). + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Text encoder step that tokenizes the prompt(s) with the Krea 2 chat template, runs the Qwen3-VL text " + "encoder, and stacks a fixed set of decoder-layer hidden states per token as the transformer's text " + "conditioning. The negative prompt is encoded the same way when the guider enables CFG." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3VLModel, description="The Qwen3-VL text encoder."), + ComponentSpec("tokenizer", AutoTokenizer, description="The tokenizer paired with the text encoder."), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 4.5, "use_original_formulation": True}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", required=True), + InputParam(name="negative_prompt", type_hint=str, description="The negative prompt(s) for CFG."), + InputParam.template("max_sequence_length", default=512), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="prompt_embeds", + type_hint=torch.Tensor, + description="Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).", + ), + OutputParam( + name="prompt_embeds_mask", + type_hint=torch.Tensor, + description="Per-prompt boolean text mask (B, text_seq_len).", + ), + OutputParam( + name="negative_prompt_embeds", + type_hint=torch.Tensor, + description="Per-prompt negative text features (only when guidance is enabled).", + ), + OutputParam( + name="negative_prompt_embeds_mask", + type_hint=torch.Tensor, + description="Per-prompt negative text mask (only when guidance is enabled).", + ), + ] + + def _encode_prompt(self, components, prompt, max_sequence_length, device): + """Tokenize `prompt` into the fixed-length Krea 2 layout and tap the selected encoder hidden states. + + Mirrors `Krea2Pipeline.get_text_hidden_states`. Returns a `(hidden_states, attention_mask)` tuple of shapes + `(batch_size, text_seq_len, num_text_layers, text_hidden_dim)` and `(batch_size, text_seq_len)` (bool). + """ + tokenizer = components.tokenizer + prompt = [prompt] if isinstance(prompt, str) else prompt + prefix_idx = _PROMPT_TEMPLATE_ENCODE_START_IDX + text = [_PROMPT_TEMPLATE_ENCODE_PREFIX + e for e in prompt] + text_tokens = tokenizer( + text, + truncation=True, + padding="max_length", + max_length=max_sequence_length + prefix_idx - _PROMPT_TEMPLATE_ENCODE_NUM_SUFFIX_TOKENS, + return_tensors="pt", + ).to(device) + suffix_tokens = tokenizer([_PROMPT_TEMPLATE_ENCODE_SUFFIX] * len(text), return_tensors="pt").to(device) + + input_ids = torch.cat([text_tokens.input_ids, suffix_tokens.input_ids], dim=1) + attention_mask = torch.cat([text_tokens.attention_mask, suffix_tokens.attention_mask], dim=1).bool() + + # Krea 2 pads in the middle of the template (`[prefix | prompt | PAD | suffix]`), so the suffix tokens sit + # downstream of the padding. The text features must use positions that count only real tokens (padding does + # not consume a position) to match how the model was trained; otherwise the suffix gets a shifted mRoPE phase. + position_ids = (attention_mask.long().cumsum(dim=-1) - 1).clamp(min=0) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + outputs = components.text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_hidden_states=True, + ) + hidden_states = torch.stack([outputs.hidden_states[i] for i in KREA2_TEXT_ENCODER_SELECT_LAYERS], dim=2) + + hidden_states = hidden_states[:, prefix_idx:] + attention_mask = attention_mask[:, prefix_idx:] + return hidden_states, attention_mask + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else list(block_state.prompt) + + block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt( + components, prompts, block_state.max_sequence_length, device + ) + + block_state.negative_prompt_embeds = None + block_state.negative_prompt_embeds_mask = None + if components.requires_unconditional_embeds: + negative_prompt = block_state.negative_prompt + if negative_prompt is None: + negative_prompt = "" + if isinstance(negative_prompt, str): + negative_prompt = [negative_prompt] * len(prompts) + block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = self._encode_prompt( + components, negative_prompt, block_state.max_sequence_length, device + ) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2TurboTextEncoderStep(Krea2TextEncoderStep): + """ + Text encoder step for the distilled Krea 2 turbo checkpoint that tokenizes the prompt(s) with the Krea 2 chat + template, runs the Qwen3-VL text encoder, and stacks a fixed set of decoder-layer hidden states per token as the + transformer's text conditioning. The distilled checkpoint runs without classifier-free guidance, so it takes no + negative prompt and has no guider. + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`AutoTokenizer`): The tokenizer paired + with the text encoder. + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for prompt encoding. + + Outputs: + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Text encoder step for the distilled Krea 2 turbo checkpoint that tokenizes the prompt(s) with the Krea 2 " + "chat template, runs the Qwen3-VL text encoder, and stacks a fixed set of decoder-layer hidden states per " + "token as the transformer's text conditioning. The distilled checkpoint runs without classifier-free " + "guidance, so it takes no negative prompt and has no guider." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3VLModel, description="The Qwen3-VL text encoder."), + ComponentSpec("tokenizer", AutoTokenizer, description="The tokenizer paired with the text encoder."), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", required=True), + InputParam.template("max_sequence_length", default=512), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="prompt_embeds", + type_hint=torch.Tensor, + description="Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim).", + ), + OutputParam( + name="prompt_embeds_mask", + type_hint=torch.Tensor, + description="Per-prompt boolean text mask (B, text_seq_len).", + ), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else list(block_state.prompt) + + block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt( + components, prompts, block_state.max_sequence_length, device + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py b/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py new file mode 100644 index 0000000000000000000000000000000000000000..ae3b2ac4fb52b1f35c29adb907c13e6939d142a7 --- /dev/null +++ b/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py @@ -0,0 +1,170 @@ +# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + Krea2PrepareLatentsStep, + Krea2PreparePositionIdsStep, + Krea2SetTimestepsStep, + Krea2TextInputsStep, +) +from .decoders import Krea2DecodeStep +from .denoise import Krea2DenoiseStep +from .encoders import Krea2TextEncoderStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +CORE_DENOISE_BLOCKS = InsertableDict( + [ + ("input", Krea2TextInputsStep()), + ("prepare_latents", Krea2PrepareLatentsStep()), + ("set_timesteps", Krea2SetTimestepsStep()), + ("prepare_position_ids", Krea2PreparePositionIdsStep()), + ("denoise", Krea2DenoiseStep()), + ] +) + + +# auto_docstring +class Krea2CoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoising workflow for Krea 2 text-to-image: prepares the batch/latents/timesteps and the shared position ids, + then runs the symmetric-CFG denoising loop, producing the denoised packed latents for the decoder. + + Components: + transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + negative_prompt_embeds (`Tensor`, *optional*): + Per-prompt negative text features. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Per-prompt negative text mask. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 28): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latents (`Tensor`): + The denoised packed latents (B, image_seq_len, in_channels). + """ + + model_name = "krea2" + block_classes = list(CORE_DENOISE_BLOCKS.values()) + block_names = list(CORE_DENOISE_BLOCKS.keys()) + + @property + def description(self) -> str: + return ( + "Core denoising workflow for Krea 2 text-to-image: prepares the batch/latents/timesteps and the shared " + "position ids, then runs the symmetric-CFG denoising loop, producing the denoised packed latents for the " + "decoder." + ) + + @property + def outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("latents", description="The denoised packed latents (B, image_seq_len, in_channels).") + ] + + +# auto_docstring +class Krea2AutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-image generation using Krea 2: encode text -> core denoise (symmetric CFG) -> + decode. + + Supported workflows: + - `text2image`: requires `prompt` + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`AutoTokenizer`): The tokenizer paired + with the text encoder. guider (`ClassifierFreeGuidance`) transformer (`Krea2Transformer2DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The negative prompt(s) for CFG. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for prompt encoding. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 28): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "krea2" + block_classes = [ + Krea2TextEncoderStep, + Krea2CoreDenoiseStep, + Krea2DecodeStep, + ] + block_names = ["text_encoder", "denoise", "decode"] + + _workflow_map = { + "text2image": {"prompt": True}, + } + + @property + def description(self) -> str: + return ( + "Auto Modular pipeline for text-to-image generation using Krea 2: encode text -> core denoise " + "(symmetric CFG) -> decode." + ) + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("images")] diff --git a/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py b/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py new file mode 100644 index 0000000000000000000000000000000000000000..79fa5406c4e5af819a7a96f7b8ea5fafff14a8b7 --- /dev/null +++ b/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py @@ -0,0 +1,164 @@ +# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + Krea2PrepareLatentsStep, + Krea2PreparePositionIdsStep, + Krea2TurboSetTimestepsStep, + Krea2TurboTextInputsStep, +) +from .decoders import Krea2DecodeStep +from .denoise import Krea2TurboDenoiseStep +from .encoders import Krea2TurboTextEncoderStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +CORE_DENOISE_BLOCKS = InsertableDict( + [ + ("input", Krea2TurboTextInputsStep()), + ("prepare_latents", Krea2PrepareLatentsStep()), + ("set_timesteps", Krea2TurboSetTimestepsStep()), + ("prepare_position_ids", Krea2PreparePositionIdsStep()), + ("denoise", Krea2TurboDenoiseStep()), + ] +) + + +# auto_docstring +class Krea2TurboCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoising workflow for the distilled Krea 2 turbo text-to-image checkpoint: prepares the + batch/latents/timesteps and the shared position ids, then runs the guidance-free denoising loop, producing the + denoised packed latents for the decoder. + + Components: + transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 8): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latents (`Tensor`): + The denoised packed latents (B, image_seq_len, in_channels). + """ + + model_name = "krea2" + block_classes = list(CORE_DENOISE_BLOCKS.values()) + block_names = list(CORE_DENOISE_BLOCKS.keys()) + + @property + def description(self) -> str: + return ( + "Core denoising workflow for the distilled Krea 2 turbo text-to-image checkpoint: prepares the " + "batch/latents/timesteps and the shared position ids, then runs the guidance-free denoising loop, " + "producing the denoised packed latents for the decoder." + ) + + @property + def outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("latents", description="The denoised packed latents (B, image_seq_len, in_channels).") + ] + + +# auto_docstring +class Krea2TurboAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-image generation using the distilled Krea 2 turbo checkpoint: encode text -> core + denoise (guidance-free) -> decode. + + Supported workflows: + - `text2image`: requires `prompt` + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`AutoTokenizer`): The tokenizer paired + with the text encoder. transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for prompt encoding. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 8): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "krea2" + block_classes = [ + Krea2TurboTextEncoderStep, + Krea2TurboCoreDenoiseStep, + Krea2DecodeStep, + ] + block_names = ["text_encoder", "denoise", "decode"] + + _workflow_map = { + "text2image": {"prompt": True}, + } + + @property + def description(self) -> str: + return ( + "Auto Modular pipeline for text-to-image generation using the distilled Krea 2 turbo checkpoint: encode " + "text -> core denoise (guidance-free) -> decode." + ) + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("images")] diff --git a/diffusers/modular_pipelines/krea2/modular_pipeline.py b/diffusers/modular_pipelines/krea2/modular_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..70d709573eecd5d8903e4d17b94be5ede54cea02 --- /dev/null +++ b/diffusers/modular_pipelines/krea2/modular_pipeline.py @@ -0,0 +1,67 @@ +# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...loaders import Krea2LoraLoaderMixin +from ..modular_pipeline import ModularPipeline + + +class Krea2ModularPipeline(ModularPipeline, Krea2LoraLoaderMixin): + """ + A ModularPipeline for Krea 2. + + > [!WARNING] > This is an experimental feature! + """ + + default_blocks_name = "Krea2AutoBlocks" + + @property + def patch_size(self): + return 2 + + @property + def default_height(self): + return 1024 + + @property + def default_width(self): + return 1024 + + @property + def vae_scale_factor(self): + vae_scale_factor = 8 + if getattr(self, "vae", None) is not None: + vae_scale_factor = 2 ** len(self.vae.temperal_downsample) + return vae_scale_factor + + @property + def requires_unconditional_embeds(self): + requires_unconditional_embeds = False + if hasattr(self, "guider") and self.guider is not None: + requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1 + return requires_unconditional_embeds + + +class Krea2TurboModularPipeline(Krea2ModularPipeline): + """ + A ModularPipeline for the distilled Krea 2 turbo (TDM) checkpoint. It runs without classifier-free guidance, so it + takes no negative prompt and has no guider. + + > [!WARNING] > This is an experimental feature! + """ + + default_blocks_name = "Krea2TurboAutoBlocks" + + @property + def requires_unconditional_embeds(self): + return False diff --git a/diffusers/modular_pipelines/ltx/__init__.py b/diffusers/modular_pipelines/ltx/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..531d9d3e4b20c786245dce67ca43e77066fc76ff --- /dev/null +++ b/diffusers/modular_pipelines/ltx/__init__.py @@ -0,0 +1,47 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_ltx"] = ["LTXAutoBlocks", "LTXBlocks", "LTXImage2VideoBlocks"] + _import_structure["modular_pipeline"] = ["LTXModularPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_ltx import LTXAutoBlocks, LTXBlocks, LTXImage2VideoBlocks + from .modular_pipeline import LTXModularPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/diffusers/modular_pipelines/ltx/before_denoise.py b/diffusers/modular_pipelines/ltx/before_denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..cd8b3ea82b821dccb2edff0e9deb9062fd7e9e62 --- /dev/null +++ b/diffusers/modular_pipelines/ltx/before_denoise.py @@ -0,0 +1,392 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect + +import numpy as np +import torch + +from ...configuration_utils import FrozenDict +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import LTXModularPipeline, LTXVideoPachifier + + +logger = logging.get_logger(__name__) + + +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class LTXTextInputStep(ModularPipelineBlocks): + model_name = "ltx" + + @property + def description(self) -> str: + return ( + "Input processing step that:\n" + " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n" + " 2. Adjusts input tensor shapes based on `batch_size` and `num_videos_per_prompt`" + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), + InputParam.template("prompt_embeds", required=True), + InputParam.template("prompt_embeds_mask", name="prompt_attention_mask"), + InputParam.template("negative_prompt_embeds"), + InputParam.template("negative_prompt_embeds_mask", name="negative_prompt_attention_mask"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("batch_size", type_hint=int), + OutputParam("dtype", type_hint=torch.dtype), + ] + + @torch.no_grad() + def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + block_state.batch_size = block_state.prompt_embeds.shape[0] + block_state.dtype = block_state.prompt_embeds.dtype + num_videos = block_state.num_videos_per_prompt + + # Repeat prompt_embeds for num_videos_per_prompt + _, seq_len, _ = block_state.prompt_embeds.shape + block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, num_videos, 1) + block_state.prompt_embeds = block_state.prompt_embeds.view(block_state.batch_size * num_videos, seq_len, -1) + + if block_state.prompt_attention_mask is not None: + block_state.prompt_attention_mask = block_state.prompt_attention_mask.repeat(num_videos, 1) + + if block_state.negative_prompt_embeds is not None: + _, seq_len, _ = block_state.negative_prompt_embeds.shape + block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(1, num_videos, 1) + block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view( + block_state.batch_size * num_videos, seq_len, -1 + ) + + if block_state.negative_prompt_attention_mask is not None: + block_state.negative_prompt_attention_mask = block_state.negative_prompt_attention_mask.repeat( + num_videos, 1 + ) + + self.set_block_state(state, block_state) + return components, state + + +class LTXSetTimestepsStep(ModularPipelineBlocks): + model_name = "ltx" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ] + + @property + def description(self) -> str: + return "Step that sets the scheduler's timesteps for inference" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps"), + InputParam.template("timesteps"), + InputParam.template("sigmas"), + InputParam.template("height", default=512), + InputParam.template("width", default=704), + InputParam("num_frames", type_hint=int, default=161), + InputParam("frame_rate", type_hint=int, default=25), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("timesteps", type_hint=torch.Tensor), + OutputParam("num_inference_steps", type_hint=int), + OutputParam("rope_interpolation_scale", type_hint=tuple), + ] + + @torch.no_grad() + def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + height = block_state.height + width = block_state.width + num_frames = block_state.num_frames + frame_rate = block_state.frame_rate + + latent_num_frames = (num_frames - 1) // components.vae_temporal_compression_ratio + 1 + latent_height = height // components.vae_spatial_compression_ratio + latent_width = width // components.vae_spatial_compression_ratio + video_sequence_length = latent_num_frames * latent_height * latent_width + + custom_timesteps = block_state.timesteps + sigmas = block_state.sigmas + + if custom_timesteps is not None: + # User provided custom timesteps, don't compute sigmas + block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps( + components.scheduler, + block_state.num_inference_steps, + device, + custom_timesteps, + ) + else: + if sigmas is None: + sigmas = np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps) + + mu = calculate_shift( + video_sequence_length, + components.scheduler.config.get("base_image_seq_len", 256), + components.scheduler.config.get("max_image_seq_len", 4096), + components.scheduler.config.get("base_shift", 0.5), + components.scheduler.config.get("max_shift", 1.15), + ) + + block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps( + components.scheduler, + block_state.num_inference_steps, + device, + sigmas=sigmas, + mu=mu, + ) + + block_state.rope_interpolation_scale = ( + components.vae_temporal_compression_ratio / frame_rate, + components.vae_spatial_compression_ratio, + components.vae_spatial_compression_ratio, + ) + + self.set_block_state(state, block_state) + return components, state + + +class LTXPrepareLatentsStep(ModularPipelineBlocks): + model_name = "ltx" + + @property + def description(self) -> str: + return "Prepare latents step that prepares the latents for the text-to-video generation process" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "pachifier", + LTXVideoPachifier, + config=FrozenDict({"patch_size": 1, "patch_size_t": 1}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("height", default=512), + InputParam.template("width", default=704), + InputParam("num_frames", type_hint=int, default=161), + InputParam.template("latents"), + InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), + InputParam.template("generator"), + InputParam.template("batch_size", required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor), + ] + + @torch.no_grad() + def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + batch_size = block_state.batch_size * block_state.num_videos_per_prompt + num_channels_latents = components.transformer.config.in_channels + + if block_state.latents is not None: + block_state.latents = block_state.latents.to(device=device, dtype=torch.float32) + else: + height = block_state.height // components.vae_spatial_compression_ratio + width = block_state.width // components.vae_spatial_compression_ratio + num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 + + shape = (batch_size, num_channels_latents, num_frames, height, width) + block_state.latents = randn_tensor( + shape, generator=block_state.generator, device=device, dtype=torch.float32 + ) + block_state.latents = components.pachifier.pack_latents(block_state.latents) + + self.set_block_state(state, block_state) + return components, state + + +class LTXImage2VideoPrepareLatentsStep(ModularPipelineBlocks): + model_name = "ltx" + + @property + def description(self) -> str: + return ( + "Prepare image-to-video latents: adds noise to pre-encoded image latents and creates a conditioning mask. " + "Expects pure noise `latents` from LTXPrepareLatentsStep." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "pachifier", + LTXVideoPachifier, + config=FrozenDict({"patch_size": 1, "patch_size_t": 1}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("image_latents", type_hint=torch.Tensor, required=True), + InputParam.template("latents", required=True), + InputParam.template("height", default=512), + InputParam.template("width", default=704), + InputParam("num_frames", type_hint=int, default=161), + InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), + InputParam.template("batch_size", required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor), + OutputParam("conditioning_mask", type_hint=torch.Tensor), + ] + + @torch.no_grad() + def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + batch_size = block_state.batch_size * block_state.num_videos_per_prompt + + height = block_state.height // components.vae_spatial_compression_ratio + width = block_state.width // components.vae_spatial_compression_ratio + num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 + + init_latents = block_state.image_latents.to(device=device, dtype=torch.float32) + if init_latents.shape[0] < batch_size: + init_latents = init_latents.repeat_interleave(batch_size // init_latents.shape[0], dim=0) + init_latents = init_latents.repeat(1, 1, num_frames, 1, 1) + + conditioning_mask = torch.zeros( + init_latents.shape[0], + 1, + init_latents.shape[2], + init_latents.shape[3], + init_latents.shape[4], + device=device, + dtype=torch.float32, + ) + conditioning_mask[:, :, 0] = 1.0 + + noise = components.pachifier.unpack_latents(block_state.latents, num_frames, height, width) + latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask) + + conditioning_mask = components.pachifier.pack_latents(conditioning_mask).squeeze(-1) + latents = components.pachifier.pack_latents(latents) + + block_state.latents = latents + block_state.conditioning_mask = conditioning_mask + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/ltx/decoders.py b/diffusers/modular_pipelines/ltx/decoders.py new file mode 100644 index 0000000000000000000000000000000000000000..8664dee25bfe73d841336d78c49193dfdc8c133e --- /dev/null +++ b/diffusers/modular_pipelines/ltx/decoders.py @@ -0,0 +1,132 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any + +import torch + +from ...configuration_utils import FrozenDict +from ...models import AutoencoderKLLTXVideo +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import LTXVideoPachifier + + +logger = logging.get_logger(__name__) + + +def _denormalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 +) -> torch.Tensor: + # Denormalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents * latents_std / scaling_factor + latents_mean + return latents + + +class LTXVaeDecoderStep(ModularPipelineBlocks): + model_name = "ltx" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLLTXVideo), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 32}), + default_creation_method="from_config", + ), + ComponentSpec( + "pachifier", + LTXVideoPachifier, + config=FrozenDict({"patch_size": 1, "patch_size_t": 1}), + default_creation_method="from_config", + ), + ] + + @property + def description(self) -> str: + return "Step that decodes the denoised latents into videos" + + @property + def inputs(self) -> list[tuple[str, Any]]: + return [ + InputParam.template("latents", required=True), + InputParam.template("output_type", default="np"), + InputParam.template("height", default=512), + InputParam.template("width", default=704), + InputParam("num_frames", type_hint=int, default=161), + InputParam("decode_timestep", default=0.0), + InputParam("decode_noise_scale", default=None), + InputParam.template("generator"), + InputParam.template("batch_size"), + InputParam.template("dtype", required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("videos")] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + vae = components.vae + + latents = block_state.latents + + height = block_state.height + width = block_state.width + num_frames = block_state.num_frames + + latent_num_frames = (num_frames - 1) // components.vae_temporal_compression_ratio + 1 + latent_height = height // components.vae_spatial_compression_ratio + latent_width = width // components.vae_spatial_compression_ratio + + latents = components.pachifier.unpack_latents(latents, latent_num_frames, latent_height, latent_width) + latents = _denormalize_latents(latents, vae.latents_mean, vae.latents_std, vae.config.scaling_factor) + latents = latents.to(block_state.dtype) + + if not vae.config.timestep_conditioning: + timestep = None + else: + device = latents.device + batch_size = block_state.batch_size + decode_timestep = block_state.decode_timestep + decode_noise_scale = block_state.decode_noise_scale + + noise = randn_tensor(latents.shape, generator=block_state.generator, device=device, dtype=latents.dtype) + if not isinstance(decode_timestep, list): + decode_timestep = [decode_timestep] * batch_size + if decode_noise_scale is None: + decode_noise_scale = decode_timestep + elif not isinstance(decode_noise_scale, list): + decode_noise_scale = [decode_noise_scale] * batch_size + + timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype) + decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ + :, None, None, None, None + ] + latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise + + latents = latents.to(vae.dtype) + video = vae.decode(latents, timestep, return_dict=False)[0] + block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type) + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/ltx/denoise.py b/diffusers/modular_pipelines/ltx/denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..b3ed86b5167934c665178987d58e5623c7a3ca90 --- /dev/null +++ b/diffusers/modular_pipelines/ltx/denoise.py @@ -0,0 +1,458 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any + +import torch + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import LTXVideoTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + PipelineState, +) +from ..modular_pipeline_utils import ComponentSpec, InputParam +from .modular_pipeline import LTXModularPipeline, LTXVideoPachifier + + +class LTXLoopBeforeDenoiser(ModularPipelineBlocks): + model_name = "ltx" + + @property + def description(self) -> str: + return ( + "Step within the denoising loop that prepares the latent input for the denoiser. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `LTXDenoiseLoopWrapper`)" + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True), + InputParam.template("dtype", required=True), + ] + + @torch.no_grad() + def __call__(self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + block_state.latent_model_input = block_state.latents.to(block_state.dtype) + return components, block_state + + +class LTXLoopDenoiser(ModularPipelineBlocks): + model_name = "ltx" + + def __init__( + self, + guider_input_fields: dict[str, Any] | None = None, + ): + if guider_input_fields is None: + guider_input_fields = { + "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"), + "encoder_attention_mask": ("prompt_attention_mask", "negative_prompt_attention_mask"), + } + if not isinstance(guider_input_fields, dict): + raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}") + self._guider_input_fields = guider_input_fields + super().__init__() + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 3.0}), + default_creation_method="from_config", + ), + ComponentSpec("transformer", LTXVideoTransformer3DModel), + ] + + @property + def description(self) -> str: + return ( + "Step within the denoising loop that denoises the latents with guidance. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `LTXDenoiseLoopWrapper`)" + ) + + @property + def inputs(self) -> list[tuple[str, Any]]: + inputs = [ + InputParam.template("attention_kwargs"), + InputParam.template("num_inference_steps", required=True), + InputParam("rope_interpolation_scale", type_hint=tuple), + InputParam.template("height"), + InputParam.template("width"), + InputParam("num_frames", type_hint=int), + ] + guider_input_names = [] + for value in self._guider_input_fields.values(): + if isinstance(value, tuple): + guider_input_names.extend(value) + else: + guider_input_names.append(value) + + for name in guider_input_names: + inputs.append(InputParam(name=name, required=True, type_hint=torch.Tensor)) + return inputs + + @torch.no_grad() + def __call__( + self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> PipelineState: + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + + latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 + latent_height = block_state.height // components.vae_spatial_compression_ratio + latent_width = block_state.width // components.vae_spatial_compression_ratio + + guider_state = components.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = guider_state_batch.as_dict() + cond_kwargs = { + k: v.to(block_state.dtype) if isinstance(v, torch.Tensor) else v + for k, v in cond_kwargs.items() + if k in self._guider_input_fields.keys() + } + + context_name = getattr(guider_state_batch, components.guider._identifier_key, None) + with components.transformer.cache_context(context_name): + guider_state_batch.noise_pred = components.transformer( + hidden_states=block_state.latent_model_input, + timestep=t.expand(block_state.latent_model_input.shape[0]).to(block_state.dtype), + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + rope_interpolation_scale=block_state.rope_interpolation_scale, + attention_kwargs=block_state.attention_kwargs, + return_dict=False, + **cond_kwargs, + )[0] + components.guider.cleanup_models(components.transformer) + + block_state.noise_pred = components.guider(guider_state)[0] + + return components, block_state + + +class LTXLoopAfterDenoiser(ModularPipelineBlocks): + model_name = "ltx" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ] + + @property + def description(self) -> str: + return ( + "Step within the denoising loop that updates the latents. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `LTXDenoiseLoopWrapper`)" + ) + + @torch.no_grad() + def __call__(self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + latents_dtype = block_state.latents.dtype + block_state.latents = components.scheduler.step( + block_state.noise_pred, + t, + block_state.latents, + return_dict=False, + )[0] + + if block_state.latents.dtype != latents_dtype: + block_state.latents = block_state.latents.to(latents_dtype) + + return components, block_state + + +class LTXDenoiseLoopWrapper(LoopSequentialPipelineBlocks): + model_name = "ltx" + + @property + def description(self) -> str: + return ( + "Pipeline block that iteratively denoises the latents over `timesteps`. " + "The specific steps within each iteration can be customized with `sub_blocks` attributes" + ) + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ComponentSpec("transformer", LTXVideoTransformer3DModel), + ] + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam.template("timesteps", required=True), + InputParam.template("num_inference_steps", required=True), + ] + + @torch.no_grad() + def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + block_state.num_warmup_steps = max( + len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0 + ) + + with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, block_state = self.loop_step(components, block_state, i=i, t=t) + if i == len(block_state.timesteps) - 1 or ( + (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0 + ): + progress_bar.update() + + self.set_block_state(state, block_state) + return components, state + + +class LTXDenoiseStep(LTXDenoiseLoopWrapper): + block_classes = [ + LTXLoopBeforeDenoiser, + LTXLoopDenoiser( + guider_input_fields={ + "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"), + "encoder_attention_mask": ("prompt_attention_mask", "negative_prompt_attention_mask"), + } + ), + LTXLoopAfterDenoiser, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoises the latents.\n" + "Its loop logic is defined in `LTXDenoiseLoopWrapper.__call__` method.\n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `LTXLoopBeforeDenoiser`\n" + " - `LTXLoopDenoiser`\n" + " - `LTXLoopAfterDenoiser`\n" + "This block supports text-to-video tasks." + ) + + +class LTXImage2VideoLoopBeforeDenoiser(ModularPipelineBlocks): + model_name = "ltx" + + @property + def description(self) -> str: + return ( + "Step within the i2v denoising loop that prepares the latent input and modulates " + "the timestep with the conditioning mask." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True), + InputParam("conditioning_mask", required=True, type_hint=torch.Tensor), + InputParam.template("dtype", required=True), + ] + + @torch.no_grad() + def __call__(self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + block_state.latent_model_input = block_state.latents.to(block_state.dtype) + block_state.timestep_adjusted = t.expand(block_state.latent_model_input.shape[0]).unsqueeze(-1) * ( + 1 - block_state.conditioning_mask + ) + return components, block_state + + +class LTXImage2VideoLoopDenoiser(ModularPipelineBlocks): + model_name = "ltx" + + def __init__( + self, + guider_input_fields: dict[str, Any] | None = None, + ): + if guider_input_fields is None: + guider_input_fields = { + "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"), + "encoder_attention_mask": ("prompt_attention_mask", "negative_prompt_attention_mask"), + } + if not isinstance(guider_input_fields, dict): + raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}") + self._guider_input_fields = guider_input_fields + super().__init__() + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 3.0}), + default_creation_method="from_config", + ), + ComponentSpec("transformer", LTXVideoTransformer3DModel), + ] + + @property + def description(self) -> str: + return ( + "Step within the i2v denoising loop that denoises the latents with guidance " + "using timestep modulated by the conditioning mask." + ) + + @property + def inputs(self) -> list[tuple[str, Any]]: + inputs = [ + InputParam.template("attention_kwargs"), + InputParam.template("num_inference_steps", required=True), + InputParam("rope_interpolation_scale", type_hint=tuple), + InputParam.template("height"), + InputParam.template("width"), + InputParam("num_frames", type_hint=int), + ] + guider_input_names = [] + for value in self._guider_input_fields.values(): + if isinstance(value, tuple): + guider_input_names.extend(value) + else: + guider_input_names.append(value) + for name in guider_input_names: + inputs.append(InputParam(name=name, required=True, type_hint=torch.Tensor)) + return inputs + + @torch.no_grad() + def __call__( + self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> PipelineState: + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + + latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 + latent_height = block_state.height // components.vae_spatial_compression_ratio + latent_width = block_state.width // components.vae_spatial_compression_ratio + + guider_state = components.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = guider_state_batch.as_dict() + cond_kwargs = { + k: v.to(block_state.dtype) if isinstance(v, torch.Tensor) else v + for k, v in cond_kwargs.items() + if k in self._guider_input_fields.keys() + } + + context_name = getattr(guider_state_batch, components.guider._identifier_key, None) + with components.transformer.cache_context(context_name): + guider_state_batch.noise_pred = components.transformer( + hidden_states=block_state.latent_model_input, + timestep=block_state.timestep_adjusted, + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + rope_interpolation_scale=block_state.rope_interpolation_scale, + attention_kwargs=block_state.attention_kwargs, + return_dict=False, + **cond_kwargs, + )[0] + components.guider.cleanup_models(components.transformer) + + block_state.noise_pred = components.guider(guider_state)[0] + + return components, block_state + + +class LTXImage2VideoLoopAfterDenoiser(ModularPipelineBlocks): + model_name = "ltx" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ComponentSpec( + "pachifier", + LTXVideoPachifier, + config=FrozenDict({"patch_size": 1, "patch_size_t": 1}), + default_creation_method="from_config", + ), + ] + + @property + def description(self) -> str: + return ( + "Step within the i2v denoising loop that updates the latents, " + "applying the scheduler step only to frames after the first (conditioned) frame." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("height"), + InputParam.template("width"), + InputParam("num_frames", type_hint=int), + ] + + @torch.no_grad() + def __call__(self, components: LTXModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 + latent_height = block_state.height // components.vae_spatial_compression_ratio + latent_width = block_state.width // components.vae_spatial_compression_ratio + + noise_pred = components.pachifier.unpack_latents( + block_state.noise_pred, latent_num_frames, latent_height, latent_width + ) + latents = components.pachifier.unpack_latents( + block_state.latents, latent_num_frames, latent_height, latent_width + ) + + noise_pred = noise_pred[:, :, 1:] + noise_latents = latents[:, :, 1:] + pred_latents = components.scheduler.step(noise_pred, t, noise_latents, return_dict=False)[0] + + latents = torch.cat([latents[:, :, :1], pred_latents], dim=2) + block_state.latents = components.pachifier.pack_latents(latents) + + return components, block_state + + +class LTXImage2VideoDenoiseStep(LTXDenoiseLoopWrapper): + block_classes = [ + LTXImage2VideoLoopBeforeDenoiser, + LTXImage2VideoLoopDenoiser( + guider_input_fields={ + "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"), + "encoder_attention_mask": ("prompt_attention_mask", "negative_prompt_attention_mask"), + } + ), + LTXImage2VideoLoopAfterDenoiser, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step for image-to-video that iteratively denoises the latents.\n" + "The first frame is kept fixed via a conditioning mask.\n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `LTXImage2VideoLoopBeforeDenoiser`\n" + " - `LTXImage2VideoLoopDenoiser`\n" + " - `LTXImage2VideoLoopAfterDenoiser`" + ) diff --git a/diffusers/modular_pipelines/ltx/encoders.py b/diffusers/modular_pipelines/ltx/encoders.py new file mode 100644 index 0000000000000000000000000000000000000000..55405ad0aefefd43ff58239c992fe3ec9c9d329b --- /dev/null +++ b/diffusers/modular_pipelines/ltx/encoders.py @@ -0,0 +1,273 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from transformers import T5EncoderModel, T5TokenizerFast + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import AutoencoderKLLTXVideo +from ...utils import logging +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import LTXModularPipeline + + +logger = logging.get_logger(__name__) + + +def _get_t5_prompt_embeds( + components, + prompt: str | list[str], + max_sequence_length: int, + device: torch.device, + dtype: torch.dtype, +): + prompt = [prompt] if isinstance(prompt, str) else prompt + + text_inputs = components.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + prompt_attention_mask = text_inputs.attention_mask + prompt_attention_mask = prompt_attention_mask.bool().to(device) + + prompt_embeds = components.text_encoder(text_input_ids.to(device))[0] + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + return prompt_embeds, prompt_attention_mask + + +class LTXTextEncoderStep(ModularPipelineBlocks): + model_name = "ltx" + + @property + def description(self) -> str: + return "Text Encoder step that generates text embeddings to guide the video generation" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", T5EncoderModel), + ComponentSpec("tokenizer", T5TokenizerFast), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 3.0}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt"), + InputParam.template("negative_prompt"), + InputParam.template("max_sequence_length", default=128), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("prompt_embeds_mask", name="prompt_attention_mask"), + OutputParam.template("negative_prompt_embeds"), + OutputParam.template("negative_prompt_embeds_mask", name="negative_prompt_attention_mask"), + ] + + @staticmethod + def check_inputs(block_state): + if block_state.prompt is not None and ( + not isinstance(block_state.prompt, str) and not isinstance(block_state.prompt, list) + ): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(block_state.prompt)}") + + @staticmethod + def encode_prompt( + components, + prompt: str, + device: torch.device | None = None, + prepare_unconditional_embeds: bool = True, + negative_prompt: str | None = None, + max_sequence_length: int = 128, + ): + device = device or components._execution_device + dtype = components.text_encoder.dtype + + if not isinstance(prompt, list): + prompt = [prompt] + batch_size = len(prompt) + + prompt_embeds, prompt_attention_mask = _get_t5_prompt_embeds( + components=components, + prompt=prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + negative_prompt_embeds = None + negative_prompt_attention_mask = None + + if prepare_unconditional_embeds: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds, negative_prompt_attention_mask = _get_t5_prompt_embeds( + components=components, + prompt=negative_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + @torch.no_grad() + def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + block_state.device = components._execution_device + + ( + block_state.prompt_embeds, + block_state.prompt_attention_mask, + block_state.negative_prompt_embeds, + block_state.negative_prompt_attention_mask, + ) = self.encode_prompt( + components=components, + prompt=block_state.prompt, + device=block_state.device, + prepare_unconditional_embeds=components.requires_unconditional_embeds, + negative_prompt=block_state.negative_prompt, + max_sequence_length=block_state.max_sequence_length, + ) + + self.set_block_state(state, block_state) + return components, state + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.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") + + +def _normalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 +) -> torch.Tensor: + # Normalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = (latents - latents_mean) * scaling_factor / latents_std + return latents + + +class LTXVaeEncoderStep(ModularPipelineBlocks): + model_name = "ltx" + + @property + def description(self) -> str: + return "VAE Encoder step that encodes an input image into latent space for image-to-video generation" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLLTXVideo), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 32}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image", required=True), + InputParam.template("height", default=512), + InputParam.template("width", default=704), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "image_latents", + type_hint=torch.Tensor, + description="Encoded image latents from the VAE encoder", + ), + ] + + @torch.no_grad() + def __call__(self, components: LTXModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + image = block_state.image + if not isinstance(image, torch.Tensor): + image = components.video_processor.preprocess(image, height=block_state.height, width=block_state.width) + image = image.to(device=device, dtype=torch.float32) + + vae_dtype = components.vae.dtype + + num_images = image.shape[0] + if isinstance(block_state.generator, list): + init_latents = [ + retrieve_latents( + components.vae.encode(image[i].unsqueeze(0).unsqueeze(2).to(vae_dtype)), + block_state.generator[i], + ) + for i in range(num_images) + ] + else: + init_latents = [ + retrieve_latents( + components.vae.encode(img.unsqueeze(0).unsqueeze(2).to(vae_dtype)), + block_state.generator, + ) + for img in image + ] + + init_latents = torch.cat(init_latents, dim=0).to(torch.float32) + block_state.image_latents = _normalize_latents( + init_latents, components.vae.latents_mean, components.vae.latents_std + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py b/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py new file mode 100644 index 0000000000000000000000000000000000000000..828c79e1c72df0d0b97763bd7fd99eca1b22a64a --- /dev/null +++ b/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py @@ -0,0 +1,487 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import ( + LTXImage2VideoPrepareLatentsStep, + LTXPrepareLatentsStep, + LTXSetTimestepsStep, + LTXTextInputStep, +) +from .decoders import LTXVaeDecoderStep +from .denoise import LTXDenoiseStep, LTXImage2VideoDenoiseStep +from .encoders import LTXTextEncoderStep, LTXVaeEncoderStep + + +logger = logging.get_logger(__name__) + + +# auto_docstring +class LTXCoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block that takes encoded conditions and runs the denoising process. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) pachifier (`LTXVideoPachifier`) guider + (`ClassifierFreeGuidance`) transformer (`LTXVideoTransformer3DModel`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_attention_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_attention_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`, *optional*, defaults to 161): + TODO: Add description. + frame_rate (`int`, *optional*, defaults to 25): + TODO: Add description. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "ltx" + block_classes = [ + LTXTextInputStep, + LTXSetTimestepsStep, + LTXPrepareLatentsStep, + LTXDenoiseStep, + ] + block_names = ["input", "set_timesteps", "prepare_latents", "denoise"] + + @property + def description(self): + return "Denoise block that takes encoded conditions and runs the denoising process." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# auto_docstring +class LTXImage2VideoCoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block for image-to-video that takes encoded conditions and image latents, and runs the denoising process. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) pachifier (`LTXVideoPachifier`) guider + (`ClassifierFreeGuidance`) transformer (`LTXVideoTransformer3DModel`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_attention_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_attention_mask (`Tensor`, *optional*): + mask for the negative text embeddings. Can be generated from text_encoder step. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`, *optional*, defaults to 161): + TODO: Add description. + frame_rate (`int`, *optional*, defaults to 25): + TODO: Add description. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image_latents (`Tensor`): + TODO: Add description. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "ltx" + block_classes = [ + LTXTextInputStep, + LTXSetTimestepsStep, + LTXPrepareLatentsStep, + LTXImage2VideoPrepareLatentsStep, + LTXImage2VideoDenoiseStep, + ] + block_names = ["input", "set_timesteps", "prepare_latents", "prepare_i2v_latents", "denoise"] + + @property + def description(self): + return "Denoise block for image-to-video that takes encoded conditions and image latents, and runs the denoising process." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# auto_docstring +class LTXBlocks(SequentialPipelineBlocks): + """ + Modular pipeline blocks for LTX Video text-to-video. + + Components: + text_encoder (`T5EncoderModel`) tokenizer (`T5Tokenizer`) guider (`ClassifierFreeGuidance`) scheduler + (`FlowMatchEulerDiscreteScheduler`) pachifier (`LTXVideoPachifier`) transformer + (`LTXVideoTransformer3DModel`) vae (`AutoencoderKLLTXVideo`) video_processor (`VideoProcessor`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 128): + Maximum sequence length for prompt encoding. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`, *optional*, defaults to 161): + TODO: Add description. + frame_rate (`int`, *optional*, defaults to 25): + TODO: Add description. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to np): + Output format: 'pil', 'np', 'pt'. + decode_timestep (`None`, *optional*, defaults to 0.0): + TODO: Add description. + decode_noise_scale (`None`, *optional*): + TODO: Add description. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "ltx" + block_classes = [ + LTXTextEncoderStep, + LTXCoreDenoiseStep, + LTXVaeDecoderStep, + ] + block_names = ["text_encoder", "denoise", "decode"] + + @property + def description(self): + return "Modular pipeline blocks for LTX Video text-to-video." + + @property + def outputs(self): + return [OutputParam.template("videos")] + + +# auto_docstring +class LTXAutoVaeEncoderStep(AutoPipelineBlocks): + """ + VAE encoder step that encodes the image input into its latent representation. + This is an auto pipeline block that works for image-to-video tasks. + - `LTXVaeEncoderStep` is used when `image` is provided. + - If `image` is not provided, step will be skipped. + + Components: + vae (`AutoencoderKLLTXVideo`) video_processor (`VideoProcessor`) + + Inputs: + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + image_latents (`Tensor`): + Encoded image latents from the VAE encoder + """ + + model_name = "ltx" + block_classes = [LTXVaeEncoderStep] + block_names = ["vae_encoder"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "VAE encoder step that encodes the image input into its latent representation.\n" + "This is an auto pipeline block that works for image-to-video tasks.\n" + " - `LTXVaeEncoderStep` is used when `image` is provided.\n" + " - If `image` is not provided, step will be skipped." + ) + + +# auto_docstring +class LTXAutoCoreDenoiseStep(AutoPipelineBlocks): + """ + Auto denoise block that selects the appropriate denoise pipeline based on inputs. + - `LTXImage2VideoCoreDenoiseStep` is used when `image_latents` is provided. + - `LTXCoreDenoiseStep` is used otherwise (text-to-video). + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) pachifier (`LTXVideoPachifier`) guider + (`ClassifierFreeGuidance`) transformer (`LTXVideoTransformer3DModel`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_attention_mask (`Tensor`): + mask for the text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_prompt_attention_mask (`Tensor`): + mask for the negative text embeddings. Can be generated from text_encoder step. + num_inference_steps (`int`): + The number of denoising steps. + timesteps (`Tensor`): + Timesteps for the denoising process. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`, *optional*, defaults to 161): + TODO: Add description. + frame_rate (`int`, *optional*, defaults to 25): + TODO: Add description. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image_latents (`Tensor`, *optional*): + TODO: Add description. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "ltx" + block_classes = [LTXImage2VideoCoreDenoiseStep, LTXCoreDenoiseStep] + block_names = ["image2video", "text2video"] + block_trigger_inputs = ["image_latents", None] + + @property + def description(self): + return ( + "Auto denoise block that selects the appropriate denoise pipeline based on inputs.\n" + " - `LTXImage2VideoCoreDenoiseStep` is used when `image_latents` is provided.\n" + " - `LTXCoreDenoiseStep` is used otherwise (text-to-video)." + ) + + +# auto_docstring +class LTXAutoBlocks(SequentialPipelineBlocks): + """ + Auto blocks for LTX Video that support both text-to-video and image-to-video workflows. + + Supported workflows: + - `text2video`: requires `prompt` + - `image2video`: requires `image`, `prompt` + + Components: + text_encoder (`T5EncoderModel`) tokenizer (`T5Tokenizer`) guider (`ClassifierFreeGuidance`) vae + (`AutoencoderKLLTXVideo`) video_processor (`VideoProcessor`) scheduler (`FlowMatchEulerDiscreteScheduler`) + pachifier (`LTXVideoPachifier`) transformer (`LTXVideoTransformer3DModel`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 128): + Maximum sequence length for prompt encoding. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + num_inference_steps (`int`): + The number of denoising steps. + timesteps (`Tensor`): + Timesteps for the denoising process. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + num_frames (`int`, *optional*, defaults to 161): + TODO: Add description. + frame_rate (`int`, *optional*, defaults to 25): + TODO: Add description. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + image_latents (`Tensor`, *optional*): + TODO: Add description. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to np): + Output format: 'pil', 'np', 'pt'. + decode_timestep (`None`, *optional*, defaults to 0.0): + TODO: Add description. + decode_noise_scale (`None`, *optional*): + TODO: Add description. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "ltx" + block_classes = [ + LTXTextEncoderStep, + LTXAutoVaeEncoderStep, + LTXAutoCoreDenoiseStep, + LTXVaeDecoderStep, + ] + block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] + _workflow_map = { + "text2video": {"prompt": True}, + "image2video": {"image": True, "prompt": True}, + } + + @property + def description(self): + return "Auto blocks for LTX Video that support both text-to-video and image-to-video workflows." + + @property + def outputs(self): + return [OutputParam.template("videos")] + + +# auto_docstring +class LTXImage2VideoBlocks(SequentialPipelineBlocks): + """ + Modular pipeline blocks for LTX Video image-to-video. + + Components: + text_encoder (`T5EncoderModel`) tokenizer (`T5Tokenizer`) guider (`ClassifierFreeGuidance`) vae + (`AutoencoderKLLTXVideo`) video_processor (`VideoProcessor`) scheduler (`FlowMatchEulerDiscreteScheduler`) + pachifier (`LTXVideoPachifier`) transformer (`LTXVideoTransformer3DModel`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 128): + Maximum sequence length for prompt encoding. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + num_frames (`int`, *optional*, defaults to 161): + TODO: Add description. + frame_rate (`int`, *optional*, defaults to 25): + TODO: Add description. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + image_latents (`Tensor`): + TODO: Add description. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to np): + Output format: 'pil', 'np', 'pt'. + decode_timestep (`None`, *optional*, defaults to 0.0): + TODO: Add description. + decode_noise_scale (`None`, *optional*): + TODO: Add description. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "ltx" + block_classes = [ + LTXTextEncoderStep, + LTXAutoVaeEncoderStep, + LTXImage2VideoCoreDenoiseStep, + LTXVaeDecoderStep, + ] + block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] + + @property + def description(self): + return "Modular pipeline blocks for LTX Video image-to-video." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/diffusers/modular_pipelines/ltx/modular_pipeline.py b/diffusers/modular_pipelines/ltx/modular_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..a5771e376cd764368ae2b51a36e736f2cbc4be77 --- /dev/null +++ b/diffusers/modular_pipelines/ltx/modular_pipeline.py @@ -0,0 +1,95 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import LTXVideoLoraLoaderMixin +from ...utils import logging +from ..modular_pipeline import ModularPipeline + + +logger = logging.get_logger(__name__) + + +class LTXVideoPachifier(ConfigMixin): + """ + A class to pack and unpack latents for LTX Video. + """ + + config_name = "config.json" + + @register_to_config + def __init__(self, patch_size: int = 1, patch_size_t: int = 1): + super().__init__() + + def pack_latents(self, latents: torch.Tensor) -> torch.Tensor: + batch_size, _, num_frames, height, width = latents.shape + patch_size = self.config.patch_size + patch_size_t = self.config.patch_size_t + post_patch_num_frames = num_frames // patch_size_t + post_patch_height = height // patch_size + post_patch_width = width // patch_size + latents = latents.reshape( + batch_size, + -1, + post_patch_num_frames, + patch_size_t, + post_patch_height, + patch_size, + post_patch_width, + patch_size, + ) + latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + return latents + + def unpack_latents(self, latents: torch.Tensor, num_frames: int, height: int, width: int) -> torch.Tensor: + batch_size = latents.size(0) + patch_size = self.config.patch_size + patch_size_t = self.config.patch_size_t + latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) + latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) + return latents + + +class LTXModularPipeline( + ModularPipeline, + LTXVideoLoraLoaderMixin, +): + """ + A ModularPipeline for LTX Video. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "LTXAutoBlocks" + + @property + def vae_spatial_compression_ratio(self): + if getattr(self, "vae", None) is not None: + return self.vae.spatial_compression_ratio + return 32 + + @property + def vae_temporal_compression_ratio(self): + if getattr(self, "vae", None) is not None: + return self.vae.temporal_compression_ratio + return 8 + + @property + def requires_unconditional_embeds(self): + if hasattr(self, "guider") and self.guider is not None: + return self.guider._enabled and self.guider.num_conditions > 1 + return False diff --git a/diffusers/modular_pipelines/mellon_node_utils.py b/diffusers/modular_pipelines/mellon_node_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f65459dfc99023c72d250df35f3eee7554b81ecc --- /dev/null +++ b/diffusers/modular_pipelines/mellon_node_utils.py @@ -0,0 +1,1101 @@ +import copy +import json +import logging +import os + +# Simple typed wrapper for parameter overrides +from dataclasses import asdict, dataclass +from typing import Any + +from huggingface_hub import create_repo, hf_hub_download, upload_file +from huggingface_hub.utils import ( + EntryNotFoundError, + HfHubHTTPError, + RepositoryNotFoundError, + RevisionNotFoundError, +) + +from ..utils import HUGGINGFACE_CO_RESOLVE_ENDPOINT +from .modular_pipeline_utils import InputParam, OutputParam + + +logger = logging.getLogger(__name__) + + +def _name_to_label(name: str) -> str: + """Convert snake_case name to Title Case label.""" + return name.replace("_", " ").title() + + +# Template definitions for standard diffuser pipeline parameters +MELLON_PARAM_TEMPLATES = { + # Image I/O + "image": {"label": "Image", "type": "image", "display": "input", "required_block_params": ["image"]}, + "images": {"label": "Images", "type": "image", "display": "output", "required_block_params": ["images"]}, + "control_image": { + "label": "Control Image", + "type": "image", + "display": "input", + "required_block_params": ["control_image"], + }, + # Latents + "latents": {"label": "Latents", "type": "latents", "display": "input", "required_block_params": ["latents"]}, + "image_latents": { + "label": "Image Latents", + "type": "latents", + "display": "input", + "required_block_params": ["image_latents"], + }, + "first_frame_latents": { + "label": "First Frame Latents", + "type": "latents", + "display": "input", + "required_block_params": ["first_frame_latents"], + }, + "latents_preview": {"label": "Latents Preview", "type": "latent", "display": "output"}, + # Image Latents with Strength + "image_latents_with_strength": { + "name": "image_latents", # name is not same as template key + "label": "Image Latents", + "type": "latents", + "display": "input", + "onChange": {"false": ["height", "width"], "true": ["strength"]}, + "required_block_params": ["image_latents", "strength"], + }, + # Embeddings + "embeddings": {"label": "Text Embeddings", "type": "embeddings", "display": "output"}, + "image_embeds": { + "label": "Image Embeddings", + "type": "image_embeds", + "display": "output", + "required_block_params": ["image_embeds"], + }, + # Text inputs + "prompt": { + "label": "Prompt", + "type": "string", + "display": "textarea", + "default": "", + "required_block_params": ["prompt"], + }, + "negative_prompt": { + "label": "Negative Prompt", + "type": "string", + "display": "textarea", + "default": "", + "required_block_params": ["negative_prompt"], + }, + # Numeric params + "guidance_scale": { + "label": "Guidance Scale", + "type": "float", + "display": "slider", + "default": 5.0, + "min": 1.0, + "max": 30.0, + "step": 0.1, + }, + "strength": { + "label": "Strength", + "type": "float", + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "required_block_params": ["strength"], + }, + "height": { + "label": "Height", + "type": "int", + "default": 1024, + "min": 64, + "step": 8, + "required_block_params": ["height"], + }, + "width": { + "label": "Width", + "type": "int", + "default": 1024, + "min": 64, + "step": 8, + "required_block_params": ["width"], + }, + "seed": { + "label": "Seed", + "type": "int", + "default": 0, + "min": 0, + "max": 4294967295, + "display": "random", + "required_block_params": ["generator"], + }, + "num_inference_steps": { + "label": "Steps", + "type": "int", + "default": 25, + "min": 1, + "max": 100, + "display": "slider", + "required_block_params": ["num_inference_steps"], + }, + "num_frames": { + "label": "Frames", + "type": "int", + "default": 81, + "min": 1, + "max": 480, + "display": "slider", + "required_block_params": ["num_frames"], + }, + "layers": { + "label": "Layers", + "type": "int", + "default": 4, + "min": 1, + "max": 10, + "display": "slider", + "required_block_params": ["layers"], + }, + "output_type": { + "label": "Output Type", + "type": "dropdown", + "default": "np", + "options": ["np", "pil", "pt"], + }, + # ControlNet + "controlnet_conditioning_scale": { + "label": "Controlnet Conditioning Scale", + "type": "float", + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "required_block_params": ["controlnet_conditioning_scale"], + }, + "control_guidance_start": { + "label": "Control Guidance Start", + "type": "float", + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "required_block_params": ["control_guidance_start"], + }, + "control_guidance_end": { + "label": "Control Guidance End", + "type": "float", + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "required_block_params": ["control_guidance_end"], + }, + # Video + "videos": {"label": "Videos", "type": "video", "display": "output", "required_block_params": ["videos"]}, + # Models + "vae": {"label": "VAE", "type": "diffusers_auto_model", "display": "input", "required_block_params": ["vae"]}, + "image_encoder": { + "label": "Image Encoder", + "type": "diffusers_auto_model", + "display": "input", + "required_block_params": ["image_encoder"], + }, + "unet": {"label": "Denoise Model", "type": "diffusers_auto_model", "display": "input"}, + "scheduler": {"label": "Scheduler", "type": "diffusers_auto_model", "display": "input"}, + "controlnet": { + "label": "ControlNet Model", + "type": "diffusers_auto_model", + "display": "input", + "required_block_params": ["controlnet"], + }, + "text_encoders": { + "label": "Text Encoders", + "type": "diffusers_auto_models", + "display": "input", + "required_block_params": ["text_encoder"], + }, + # Bundles/Custom + "controlnet_bundle": { + "label": "ControlNet", + "type": "custom_controlnet", + "display": "input", + "required_block_params": "controlnet_image", + }, + "ip_adapter": {"label": "IP Adapter", "type": "custom_ip_adapter", "display": "input"}, + "guider": { + "label": "Guider", + "type": "custom_guider", + "display": "input", + "onChange": {False: ["guidance_scale"], True: []}, + }, + "doc": {"label": "Doc", "type": "string", "display": "output"}, +} + + +class MellonParamMeta(type): + """Metaclass that enables MellonParam.template_name(**overrides) syntax.""" + + def __getattr__(cls, name: str): + if name in MELLON_PARAM_TEMPLATES: + + def factory(default=None, **overrides): + template = MELLON_PARAM_TEMPLATES[name] + # Use template's name if specified, otherwise use the key + params = {"name": template.get("name", name), **template, **overrides} + if default is not None: + params["default"] = default + return cls(**params) + + return factory + + raise AttributeError(f"type object 'MellonParam' has no attribute '{name}'") + + +@dataclass(frozen=True) +class MellonParam(metaclass=MellonParamMeta): + """ + Parameter definition for Mellon nodes. + + Usage: + ```python + # From template (standard diffuser params) + MellonParam.seed() + MellonParam.prompt(default="a cat") + MellonParam.latents(display="output") + + # Generic inputs (for custom blocks) + MellonParam.Input.slider("my_scale", default=1.0, min=0.0, max=2.0) + MellonParam.Input.dropdown("mode", options=["fast", "slow"]) + + # Generic outputs + MellonParam.Output.image("result_images") + + # Fully custom + MellonParam(name="custom", label="Custom", type="float", default=0.5) + ``` + """ + + name: str + label: str + type: str + display: str | None = None + default: Any = None + min: float | None = None + max: float | None = None + step: float | None = None + options: Any = None + value: Any = None + fieldOptions: dict[str, Any] | None = None + onChange: Any = None + onSignal: Any = None + required_block_params: str | list[str] | None = None + + def to_dict(self) -> dict[str, Any]: + """Convert to dict for Mellon schema, excluding None values and internal fields.""" + data = asdict(self) + return {k: v for k, v in data.items() if v is not None and k not in ("name", "required_block_params")} + + # ========================================================================= + # Input: Generic input parameter factories (for custom blocks) + # ========================================================================= + class Input: + """input UI elements for custom blocks.""" + + @classmethod + def image(cls, name: str) -> "MellonParam": + """image input.""" + return MellonParam(name=name, label=_name_to_label(name), type="image", display="input") + + @classmethod + def textbox(cls, name: str, default: str = "") -> "MellonParam": + """text input as textarea.""" + return MellonParam( + name=name, label=_name_to_label(name), type="string", display="textarea", default=default + ) + + @classmethod + def dropdown(cls, name: str, options: list[str] = None, default: str = None) -> "MellonParam": + """dropdown selection.""" + if options and not default: + default = options[0] + if not default: + default = "" + if not options: + options = [default] + return MellonParam(name=name, label=_name_to_label(name), type="string", options=options, value=default) + + @classmethod + def slider( + cls, name: str, default: float = 0, min: float = None, max: float = None, step: float = None + ) -> "MellonParam": + """slider input.""" + is_float = isinstance(default, float) or (step is not None and isinstance(step, float)) + param_type = "float" if is_float else "int" + if min is None: + min = default + if max is None: + max = default + if step is None: + step = 0.01 if is_float else 1 + return MellonParam( + name=name, + label=_name_to_label(name), + type=param_type, + display="slider", + default=default, + min=min, + max=max, + step=step, + ) + + @classmethod + def number( + cls, name: str, default: float = 0, min: float = None, max: float = None, step: float = None + ) -> "MellonParam": + """number input (no slider).""" + is_float = isinstance(default, float) or (step is not None and isinstance(step, float)) + param_type = "float" if is_float else "int" + return MellonParam( + name=name, label=_name_to_label(name), type=param_type, default=default, min=min, max=max, step=step + ) + + @classmethod + def seed(cls, name: str = "seed", default: int = 0) -> "MellonParam": + """seed input with randomize button.""" + return MellonParam( + name=name, + label=_name_to_label(name), + type="int", + display="random", + default=default, + min=0, + max=4294967295, + ) + + @classmethod + def checkbox(cls, name: str, default: bool = False) -> "MellonParam": + """boolean checkbox.""" + return MellonParam(name=name, label=_name_to_label(name), type="boolean", value=default) + + @classmethod + def custom_type(cls, name: str, type: str) -> "MellonParam": + """custom type input for node connections.""" + return MellonParam(name=name, label=_name_to_label(name), type=type, display="input") + + @classmethod + def model(cls, name: str) -> "MellonParam": + """model input for diffusers components.""" + return MellonParam(name=name, label=_name_to_label(name), type="diffusers_auto_model", display="input") + + # ========================================================================= + # Output: Generic output parameter factories (for custom blocks) + # ========================================================================= + class Output: + """output UI elements for custom blocks.""" + + @classmethod + def image(cls, name: str) -> "MellonParam": + """image output.""" + return MellonParam(name=name, label=_name_to_label(name), type="image", display="output") + + @classmethod + def video(cls, name: str) -> "MellonParam": + """video output.""" + return MellonParam(name=name, label=_name_to_label(name), type="video", display="output") + + @classmethod + def text(cls, name: str) -> "MellonParam": + """text output.""" + return MellonParam(name=name, label=_name_to_label(name), type="string", display="output") + + @classmethod + def custom_type(cls, name: str, type: str) -> "MellonParam": + """custom type output for node connections.""" + return MellonParam(name=name, label=_name_to_label(name), type=type, display="output") + + @classmethod + def model(cls, name: str) -> "MellonParam": + """model output for diffusers components.""" + return MellonParam(name=name, label=_name_to_label(name), type="diffusers_auto_model", display="output") + + +def input_param_to_mellon_param(input_param: "InputParam") -> MellonParam: + """ + Convert an InputParam to a MellonParam using metadata. + + Args: + input_param: An InputParam with optional metadata containing either: + - {"mellon": ""} for simple types (image, textbox, slider, etc.) + - {"mellon": MellonParam(...)} for full control over UI configuration + + Returns: + MellonParam instance + """ + name = input_param.name + metadata = input_param.metadata + mellon_value = metadata.get("mellon") if metadata else None + default = input_param.default + + # If it's already a MellonParam, return it directly + if isinstance(mellon_value, MellonParam): + return mellon_value + + mellon_type = mellon_value + + if mellon_type == "image": + return MellonParam.Input.image(name) + elif mellon_type == "textbox": + return MellonParam.Input.textbox(name, default=default or "") + elif mellon_type == "dropdown": + return MellonParam.Input.dropdown(name, default=default or "") + elif mellon_type == "slider": + return MellonParam.Input.slider(name, default=default or 0) + elif mellon_type == "number": + return MellonParam.Input.number(name, default=default or 0) + elif mellon_type == "seed": + return MellonParam.Input.seed(name, default=default or 0) + elif mellon_type == "checkbox": + return MellonParam.Input.checkbox(name, default=default or False) + elif mellon_type == "model": + return MellonParam.Input.model(name) + else: + # None or unknown -> custom + return MellonParam.Input.custom_type(name, type="custom") + + +def output_param_to_mellon_param(output_param: "OutputParam") -> MellonParam: + """ + Convert an OutputParam to a MellonParam using metadata. + + Args: + output_param: An OutputParam with optional metadata={"mellon": ""} where type is one of: + image, video, text, model. If metadata is None or unknown, maps to "custom". + + Returns: + MellonParam instance + """ + name = output_param.name + metadata = output_param.metadata + mellon_type = metadata.get("mellon") if metadata else None + + if mellon_type == "image": + return MellonParam.Output.image(name) + elif mellon_type == "video": + return MellonParam.Output.video(name) + elif mellon_type == "text": + return MellonParam.Output.text(name) + elif mellon_type == "model": + return MellonParam.Output.model(name) + else: + # None or unknown -> custom + return MellonParam.Output.custom_type(name, type="custom") + + +DEFAULT_NODE_SPECS = { + "controlnet": None, + "denoise": { + "inputs": [ + MellonParam.embeddings(display="input"), + MellonParam.width(), + MellonParam.height(), + MellonParam.seed(), + MellonParam.num_inference_steps(), + MellonParam.num_frames(), + MellonParam.guidance_scale(), + MellonParam.strength(), + MellonParam.image_latents_with_strength(), + MellonParam.image_latents(), + MellonParam.first_frame_latents(), + MellonParam.controlnet_bundle(display="input"), + ], + "model_inputs": [ + MellonParam.unet(), + MellonParam.guider(), + MellonParam.scheduler(), + ], + "outputs": [ + MellonParam.latents(display="output"), + MellonParam.latents_preview(), + MellonParam.doc(), + ], + "required_inputs": ["embeddings"], + "required_model_inputs": ["unet", "scheduler"], + "block_name": "denoise", + }, + "vae_encoder": { + "inputs": [ + MellonParam.image(), + ], + "model_inputs": [ + MellonParam.vae(), + ], + "outputs": [ + MellonParam.image_latents(display="output"), + MellonParam.doc(), + ], + "required_inputs": ["image"], + "required_model_inputs": ["vae"], + "block_name": "vae_encoder", + }, + "text_encoder": { + "inputs": [ + MellonParam.prompt(), + MellonParam.negative_prompt(), + ], + "model_inputs": [ + MellonParam.text_encoders(), + ], + "outputs": [ + MellonParam.embeddings(display="output"), + MellonParam.doc(), + ], + "required_inputs": ["prompt"], + "required_model_inputs": ["text_encoders"], + "block_name": "text_encoder", + }, + "decoder": { + "inputs": [ + MellonParam.latents(display="input"), + ], + "model_inputs": [ + MellonParam.vae(), + ], + "outputs": [ + MellonParam.images(), + MellonParam.videos(), + MellonParam.doc(), + ], + "required_inputs": ["latents"], + "required_model_inputs": ["vae"], + "block_name": "decode", + }, +} + + +def mark_required(label: str, marker: str = " *") -> str: + """Add required marker to label if not already present.""" + if label.endswith(marker): + return label + return f"{label}{marker}" + + +def node_spec_to_mellon_dict(node_spec: dict[str, Any], node_type: str) -> dict[str, Any]: + """ + Convert a node spec dict into Mellon format. + + A node spec is how we define a Mellon diffusers node in code. This function converts it into the `params` map + format that Mellon UI expects. + + The `params` map is a dict where keys are parameter names and values are UI configuration: + ```python + {"seed": {"label": "Seed", "type": "int", "default": 0}} + ``` + + For Modular Mellon nodes, we need to distinguish: + - `inputs`: Pipeline inputs (e.g., seed, prompt, image) + - `model_inputs`: Model components (e.g., unet, vae, scheduler) + - `outputs`: Node outputs (e.g., latents, images) + + The node spec also includes: + - `required_inputs` / `required_model_inputs`: Which params are required (marked with *) + - `block_name`: The modular pipeline block this node corresponds to on backend + + We provide factory methods for common parameters (e.g., `MellonParam.seed()`, `MellonParam.unet()`) so you don't + have to manually specify all the UI configuration. + + Args: + node_spec: Dict with `inputs`, `model_inputs`, `outputs` (lists of MellonParam), + plus `required_inputs`, `required_model_inputs`, `block_name`. + node_type: The node type string (e.g., "denoise", "controlnet") + + Returns: + Dict with: + - `params`: Flat dict of all params in Mellon UI format + - `input_names`: List of input parameter names + - `model_input_names`: List of model input parameter names + - `output_names`: List of output parameter names + - `block_name`: The backend block name + - `node_type`: The node type + + Example: + ```python + node_spec = { + "inputs": [MellonParam.seed(), MellonParam.prompt()], + "model_inputs": [MellonParam.unet()], + "outputs": [MellonParam.latents(display="output")], + "required_inputs": ["prompt"], + "required_model_inputs": ["unet"], + "block_name": "denoise", + } + + result = node_spec_to_mellon_dict(node_spec, "denoise") + # Returns: + # { + # "params": { + # "seed": {"label": "Seed", "type": "int", "default": 0}, + # "prompt": {"label": "Prompt *", "type": "string", "default": ""}, # * marks required + # "unet": {"label": "Denoise Model *", "type": "diffusers_auto_model", "display": "input"}, + # "latents": {"label": "Latents", "type": "latents", "display": "output"}, + # }, + # "input_names": ["seed", "prompt"], + # "model_input_names": ["unet"], + # "output_names": ["latents"], + # "block_name": "denoise", + # "node_type": "denoise", + # } + ``` + """ + params = {} + input_names = [] + model_input_names = [] + output_names = [] + + required_inputs = node_spec.get("required_inputs", []) + required_model_inputs = node_spec.get("required_model_inputs", []) + + # Process inputs + for p in node_spec.get("inputs", []): + param_dict = p.to_dict() + if p.name in required_inputs: + param_dict["label"] = mark_required(param_dict["label"]) + params[p.name] = param_dict + input_names.append(p.name) + + # Process model_inputs + for p in node_spec.get("model_inputs", []): + param_dict = p.to_dict() + if p.name in required_model_inputs: + param_dict["label"] = mark_required(param_dict["label"]) + params[p.name] = param_dict + model_input_names.append(p.name) + + # Process outputs: add a prefix to the output name if it already exists as an input + for p in node_spec.get("outputs", []): + if p.name in input_names: + # rename to out_ + output_name = f"out_{p.name}" + else: + output_name = p.name + params[output_name] = p.to_dict() + output_names.append(output_name) + + return { + "params": params, + "input_names": input_names, + "model_input_names": model_input_names, + "output_names": output_names, + "block_name": node_spec.get("block_name"), + "node_type": node_type, + } + + +class MellonPipelineConfig: + """ + Configuration for an entire Mellon pipeline containing multiple nodes. + + Accepts node specs as dicts with inputs/model_inputs/outputs lists of MellonParam, converts them to Mellon-ready + format, and handles save/load to Hub. + + Example: + ```python + config = MellonPipelineConfig( + node_specs={ + "denoise": { + "inputs": [MellonParam.seed(), MellonParam.prompt()], + "model_inputs": [MellonParam.unet()], + "outputs": [MellonParam.latents(display="output")], + "required_inputs": ["prompt"], + "required_model_inputs": ["unet"], + "block_name": "denoise", + }, + "decoder": { + "inputs": [MellonParam.latents(display="input")], + "outputs": [MellonParam.images()], + "block_name": "decoder", + }, + }, + label="My Pipeline", + default_repo="user/my-pipeline", + default_dtype="float16", + ) + + # Access Mellon format dict + denoise = config.node_params["denoise"] + input_names = denoise["input_names"] + params = denoise["params"] + + # Save to Hub + config.save("./my_config", push_to_hub=True, repo_id="user/my-pipeline") + + # Load from Hub + loaded = MellonPipelineConfig.load("user/my-pipeline") + ``` + """ + + config_name = "mellon_pipeline_config.json" + + def __init__( + self, + node_specs: dict[str, dict[str, Any] | None], + label: str = "", + default_repo: str = "", + default_dtype: str = "", + ): + """ + Args: + node_specs: Dict mapping node_type to node spec or None. + Node spec has: inputs, model_inputs, outputs, required_inputs, required_model_inputs, + block_name (all optional) + label: Human-readable label for the pipeline + default_repo: Default HuggingFace repo for this pipeline + default_dtype: Default dtype (e.g., "float16", "bfloat16") + """ + # Convert all node specs to Mellon format immediately + self.node_specs = node_specs + + self.label = label + self.default_repo = default_repo + self.default_dtype = default_dtype + + @property + def node_params(self) -> dict[str, Any]: + """Lazily compute node_params from node_specs.""" + if self.node_specs is None: + return self._node_params + + params = {} + for node_type, spec in self.node_specs.items(): + if spec is None: + params[node_type] = None + else: + params[node_type] = node_spec_to_mellon_dict(spec, node_type) + return params + + def __repr__(self) -> str: + lines = [ + f"MellonPipelineConfig(label={self.label!r}, default_repo={self.default_repo!r}, default_dtype={self.default_dtype!r})" + ] + for node_type, spec in self.node_specs.items(): + if spec is None: + lines.append(f" {node_type}: None") + else: + inputs = [p.name for p in spec.get("inputs", [])] + model_inputs = [p.name for p in spec.get("model_inputs", [])] + outputs = [p.name for p in spec.get("outputs", [])] + lines.append(f" {node_type}:") + lines.append(f" inputs: {inputs}") + lines.append(f" model_inputs: {model_inputs}") + lines.append(f" outputs: {outputs}") + return "\n".join(lines) + + def to_dict(self) -> dict[str, Any]: + """Convert to a JSON-serializable dictionary.""" + return { + "label": self.label, + "default_repo": self.default_repo, + "default_dtype": self.default_dtype, + "node_params": self.node_params, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "MellonPipelineConfig": + """ + Create from a dictionary (loaded from JSON). + + Note: The mellon_params are already in Mellon format when loading from JSON. + """ + instance = cls.__new__(cls) + instance.node_specs = None + instance._node_params = data.get("node_params", {}) + instance.label = data.get("label", "") + instance.default_repo = data.get("default_repo", "") + instance.default_dtype = data.get("default_dtype", "") + return instance + + def to_json_string(self) -> str: + """Serialize to JSON string.""" + return json.dumps(self.to_dict(), indent=2, sort_keys=False) + "\n" + + def to_json_file(self, json_file_path: str | os.PathLike): + """Save to a JSON file.""" + with open(json_file_path, "w", encoding="utf-8") as writer: + writer.write(self.to_json_string()) + + @classmethod + def from_json_file(cls, json_file_path: str | os.PathLike) -> "MellonPipelineConfig": + """Load from a JSON file.""" + with open(json_file_path, "r", encoding="utf-8") as reader: + data = json.load(reader) + return cls.from_dict(data) + + def save(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs): + """Save the mellon pipeline config to a directory.""" + if os.path.isfile(save_directory): + raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file") + + os.makedirs(save_directory, exist_ok=True) + output_path = os.path.join(save_directory, self.config_name) + self.to_json_file(output_path) + logger.info(f"Pipeline config saved to {output_path}") + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + private = kwargs.pop("private", None) + create_pr = kwargs.pop("create_pr", False) + token = kwargs.pop("token", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id + + upload_file( + path_or_fileobj=output_path, + path_in_repo=self.config_name, + repo_id=repo_id, + token=token, + commit_message=commit_message or "Upload MellonPipelineConfig", + create_pr=create_pr, + ) + logger.info(f"Pipeline config pushed to hub: {repo_id}") + + @classmethod + def load( + cls, + pretrained_model_name_or_path: str | os.PathLike, + **kwargs, + ) -> "MellonPipelineConfig": + """Load a pipeline config from a local path or Hugging Face Hub.""" + cache_dir = kwargs.pop("cache_dir", None) + local_dir = kwargs.pop("local_dir", None) + local_dir_use_symlinks = kwargs.pop("local_dir_use_symlinks", "auto") + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + token = kwargs.pop("token", None) + local_files_only = kwargs.pop("local_files_only", False) + revision = kwargs.pop("revision", None) + subfolder = kwargs.pop("subfolder", None) + + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + + if os.path.isfile(pretrained_model_name_or_path): + config_file = pretrained_model_name_or_path + elif os.path.isdir(pretrained_model_name_or_path): + config_file = os.path.join(pretrained_model_name_or_path, cls.config_name) + if not os.path.isfile(config_file): + raise EnvironmentError(f"No file named {cls.config_name} found in {pretrained_model_name_or_path}") + else: + try: + config_file = hf_hub_download( + pretrained_model_name_or_path, + filename=cls.config_name, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder=subfolder, + local_dir=local_dir, + local_dir_use_symlinks=local_dir_use_symlinks, + ) + except RepositoryNotFoundError: + raise EnvironmentError( + f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier" + " listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a" + " token having permission to this repo with `token` or log in with `hf auth login`." + ) + except RevisionNotFoundError: + raise EnvironmentError( + f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for" + " this model name. Check the model page at" + f" 'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions." + ) + except EntryNotFoundError: + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named {cls.config_name}." + ) + except HfHubHTTPError as err: + raise EnvironmentError( + "There was a specific connection error when trying to load" + f" {pretrained_model_name_or_path}:\n{err}" + ) + except ValueError: + raise EnvironmentError( + f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it" + f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a" + f" directory containing a {cls.config_name} file.\nCheckout your internet connection or see how to" + " run the library in offline mode at" + " 'https://huggingface.co/docs/diffusers/installation#offline-mode'." + ) + except EnvironmentError: + raise EnvironmentError( + f"Can't load config for '{pretrained_model_name_or_path}'. If you were trying to load it from " + "'https://huggingface.co/models', make sure you don't have a local directory with the same name. " + f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory " + f"containing a {cls.config_name} file" + ) + + try: + return cls.from_json_file(config_file) + except (json.JSONDecodeError, UnicodeDecodeError): + raise EnvironmentError(f"The config file at '{config_file}' is not a valid JSON file.") + + @classmethod + def from_blocks( + cls, + blocks, + template: dict[str, dict[str, Any]] | None = None, + label: str = "", + default_repo: str = "", + default_dtype: str = "bfloat16", + ) -> "MellonPipelineConfig": + """ + Create MellonPipelineConfig by matching template against actual pipeline blocks. + """ + if template is None: + template = DEFAULT_NODE_SPECS + + sub_block_map = dict(blocks.sub_blocks) + + def filter_spec_for_block(template_spec: dict[str, Any], block) -> dict[str, Any] | None: + """Filter template spec params based on what the block actually supports.""" + block_input_names = set(block.input_names) + block_output_names = set(block.intermediate_output_names) + block_component_names = set(block.component_names) + + filtered_inputs = [ + p + for p in template_spec.get("inputs", []) + if p.required_block_params is None + or all(name in block_input_names for name in p.required_block_params) + ] + filtered_model_inputs = [ + p + for p in template_spec.get("model_inputs", []) + if p.required_block_params is None + or all(name in block_component_names for name in p.required_block_params) + ] + filtered_outputs = [ + p + for p in template_spec.get("outputs", []) + if p.required_block_params is None + or all(name in block_output_names for name in p.required_block_params) + ] + + filtered_input_names = {p.name for p in filtered_inputs} + filtered_model_input_names = {p.name for p in filtered_model_inputs} + + filtered_required_inputs = [ + r for r in template_spec.get("required_inputs", []) if r in filtered_input_names + ] + filtered_required_model_inputs = [ + r for r in template_spec.get("required_model_inputs", []) if r in filtered_model_input_names + ] + + return { + "inputs": filtered_inputs, + "model_inputs": filtered_model_inputs, + "outputs": filtered_outputs, + "required_inputs": filtered_required_inputs, + "required_model_inputs": filtered_required_model_inputs, + "block_name": template_spec.get("block_name"), + } + + # Build node specs + node_specs = {} + for node_type, template_spec in template.items(): + if template_spec is None: + node_specs[node_type] = None + continue + + block_name = template_spec.get("block_name") + if block_name is None or block_name not in sub_block_map: + node_specs[node_type] = None + continue + + node_specs[node_type] = filter_spec_for_block(template_spec, sub_block_map[block_name]) + + return cls( + node_specs=node_specs, + label=label or getattr(blocks, "model_name", ""), + default_repo=default_repo, + default_dtype=default_dtype, + ) + + @classmethod + def from_custom_block( + cls, + block, + node_label: str = None, + input_types: dict[str, Any] | None = None, + output_types: dict[str, Any] | None = None, + ) -> "MellonPipelineConfig": + """ + Create a MellonPipelineConfig from a custom block. + + Args: + block: A block instance with `inputs`, `outputs`, and `expected_components`/`component_names` properties. + Each InputParam/OutputParam should have metadata={"mellon": ""} where type is one of: image, + video, text, checkbox, number, slider, dropdown, model. If metadata is None, maps to "custom". + node_label: The display label for the node. Defaults to block class name with spaces. + input_types: + Optional dict mapping input param names to mellon types. Overrides the block's metadata if provided. + Example: {"prompt": "textbox", "image": "image"} + output_types: + Optional dict mapping output param names to mellon types. Overrides the block's metadata if provided. + Example: {"prompt": "text", "images": "image"} + + Returns: + MellonPipelineConfig instance + """ + if node_label is None: + class_name = block.__class__.__name__ + node_label = "".join([" " + c if c.isupper() else c for c in class_name]).strip() + + if input_types is None: + input_types = {} + if output_types is None: + output_types = {} + + inputs = [] + model_inputs = [] + outputs = [] + + # Process block inputs + for input_param in block.inputs: + if input_param.name is None: + continue + if input_param.name in input_types: + input_param = copy.copy(input_param) + input_param.metadata = {"mellon": input_types[input_param.name]} + print(f" processing input: {input_param.name}, metadata: {input_param.metadata}") + inputs.append(input_param_to_mellon_param(input_param)) + + # Process block outputs + for output_param in block.outputs: + if output_param.name is None: + continue + if output_param.name in output_types: + output_param = copy.copy(output_param) + output_param.metadata = {"mellon": output_types[output_param.name]} + outputs.append(output_param_to_mellon_param(output_param)) + + # Process expected components (all map to model inputs) + component_names = block.component_names + for component_name in component_names: + model_inputs.append(MellonParam.Input.model(component_name)) + + # Always add doc output + outputs.append(MellonParam.doc()) + + node_spec = { + "inputs": inputs, + "model_inputs": model_inputs, + "outputs": outputs, + "required_inputs": [], + "required_model_inputs": [], + "block_name": "custom", + } + + return cls( + node_specs={"custom": node_spec}, + label=node_label, + ) diff --git a/diffusers/modular_pipelines/minimax_h3/__init__.py b/diffusers/modular_pipelines/minimax_h3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6f492f17eff1edd7b6590282802e27b51ee45cf0 --- /dev/null +++ b/diffusers/modular_pipelines/minimax_h3/__init__.py @@ -0,0 +1,49 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_minimax_h3"] = ["MiniMaxH3Blocks", "MiniMaxH3Ref2VABlocks"] + _import_structure["modular_pipeline"] = ["MiniMaxH3ModularPipeline", "MiniMaxH3Ref2VAModularPipeline"] + _import_structure["packing_ref2va"] = ["MiniMaxH3Reference"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_minimax_h3 import MiniMaxH3Blocks, MiniMaxH3Ref2VABlocks + from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline + from .packing_ref2va import MiniMaxH3Reference +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/diffusers/modular_pipelines/minimax_h3/before_denoise.py b/diffusers/modular_pipelines/minimax_h3/before_denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..ef874b1e0100dd770ee42fefcd4e51b5414fbd24 --- /dev/null +++ b/diffusers/modular_pipelines/minimax_h3/before_denoise.py @@ -0,0 +1,425 @@ +# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...schedulers import MiniMaxH3Scheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline +from .packing import ( + MINIMAX_H3_AUDIO_CHANNELS, + MINIMAX_H3_KEYFRAME_NOISE_AUG, + MiniMaxH3PackedSequence, + build_packed_sequence, + build_row_timesteps, + patchify_video_latents, +) +from .packing_ref2va import MiniMaxH3PreparedReference, build_ref2va_packed_sequence + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _layout_inputs() -> list[InputParam]: + r"""What both packed layouts are built from, beyond the conditioning of the task itself.""" + return [ + InputParam( + name="text_token_tags", + type_hint=torch.Tensor, + required=True, + description="The per-row modality tag of every row of `prompt_embeds`.", + ), + InputParam( + name="num_latent_frames", type_hint=int, required=True, description="Number of video latent frames." + ), + InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."), + InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."), + InputParam( + name="num_audio_latents", + type_hint=int, + required=True, + description="Number of audio latents per channel.", + ), + ] + + +def _layout_outputs() -> list[OutputParam]: + r"""The row layout of the packed sequence, shared by the two tasks.""" + return [ + OutputParam( + "layout", + type_hint=MiniMaxH3PackedSequence, + description="The structural description of the packed sequence.", + ), + OutputParam( + "position_ids", + type_hint=torch.Tensor, + description="The `(t, h, w)` rotary coordinate of every row, in float64.", + ), + OutputParam("token_tags", type_hint=torch.Tensor, description="The modality tag of every row."), + OutputParam( + "video_indices", + type_hint=torch.Tensor, + description="Sequence positions of the video rows, conditioning rows first.", + ), + OutputParam( + "audio_indices", + type_hint=torch.Tensor, + description="Sequence positions of the audio rows, reference rows first.", + ), + OutputParam("text_indices", type_hint=torch.Tensor, description="Sequence positions of the text rows."), + OutputParam( + "num_condition_video_rows", + type_hint=int, + description="How many leading video rows are conditioning rows rather than generated rows.", + ), + OutputParam( + "num_condition_audio_rows", + type_hint=int, + description="How many leading audio rows are reference rows rather than generated rows.", + ), + ] + + +def _set_layout_state(block_state, layout: MiniMaxH3PackedSequence, device: torch.device) -> None: + block_state.layout = layout + block_state.position_ids = layout.position_ids.to(device) + block_state.token_tags = layout.token_tags.to(device) + block_state.video_indices = layout.video_indices.to(device) + block_state.audio_indices = layout.audio_indices.to(device) + block_state.text_indices = layout.text_indices.to(device) + block_state.num_condition_video_rows = layout.num_condition_video_rows + block_state.num_condition_audio_rows = layout.num_condition_audio_rows + + +class MiniMaxH3PrepareLayoutStep(ModularPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return ( + "Builds the packed layout of a `t2va` / `fl2va` request — `[text | keyframe conditions | target audio | " + "target video]` — and its fp64 rotary grid. MiniMax-H3 runs full self-attention over this one sequence, " + "so the layout is what every later block addresses rows through." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + *_layout_inputs(), + InputParam( + name="keyframe_anchors", + type_hint=tuple, + default=(), + description="Which end of the video every keyframe is anchored to, in packed order.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return _layout_outputs() + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + layout = build_packed_sequence( + block_state.text_token_tags, + block_state.num_latent_frames, + block_state.latent_height, + block_state.latent_width, + block_state.num_audio_latents, + components.patch_size, + block_state.keyframe_anchors, + ) + _set_layout_state(block_state, layout, components._execution_device) + + self.set_block_state(state, block_state) + return components, state + + +class MiniMaxH3Ref2VAPrepareLayoutStep(ModularPipelineBlocks): + model_name = "minimax-h3-ref2va" + + @property + def description(self) -> str: + return ( + "Builds the packed layout of a `ref2va` request — `[text | reference blocks | target audio | target " + "video]` — and its fp64 rotary grid. The reference order advances the shared audio/video rotary clock, so " + "it is part of the layout rather than a detail of the presentation." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + *_layout_inputs(), + InputParam( + name="prepared_references", + type_hint=list[MiniMaxH3PreparedReference], + required=True, + description="The prepared references, in packed order, with their latent geometry filled in.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return _layout_outputs() + + @torch.no_grad() + def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + layout = build_ref2va_packed_sequence( + block_state.text_token_tags, + block_state.prepared_references, + block_state.num_latent_frames, + block_state.latent_height, + block_state.latent_width, + block_state.num_audio_latents, + components.patch_size, + ) + _set_layout_state(block_state, layout, components._execution_device) + + self.set_block_state(state, block_state) + return components, state + + +class MiniMaxH3PrepareLatentsStep(ModularPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return ( + "Draws the initial noise of the generated rows and prepends the conditioning rows. MiniMax-H3 draws the " + "video noise as a latent tensor and patchifies it afterwards, then the audio noise directly in row " + "layout — both off the request's generator, after the conditioning noise of the encoder step." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="num_latent_frames", type_hint=int, required=True, description="Number of video latent frames." + ), + InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."), + InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."), + InputParam( + name="num_audio_latents", + type_hint=int, + required=True, + description="Number of audio latents per channel.", + ), + InputParam.template( + "generator", + description=( + "The generator of the request. The video noise is drawn from it first, then the audio noise." + ), + ), + InputParam( + name="latents", + type_hint=torch.Tensor, + description=( + "Pre-generated video noise of shape `(1, 24, num_latent_frames, latent_height, latent_width)`, " + "used instead of the draw." + ), + ), + InputParam( + name="audio_latents", + type_hint=torch.Tensor, + description="Pre-generated audio noise of shape `(2, 32, num_audio_latents)`.", + ), + InputParam( + name="condition_latents", + type_hint=torch.Tensor, + description="The video conditioning rows to prepend, or None for a request that has none.", + ), + InputParam( + name="audio_condition_latents", + type_hint=torch.Tensor, + description="The audio conditioning rows to prepend, or None for a request that has none.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="The video rows of the packed sequence, conditioning rows first.", + ), + OutputParam( + "audio_latents", + type_hint=torch.Tensor, + description="The channel-major audio rows of the packed sequence, reference rows first.", + ), + ] + + @staticmethod + def prepare_latents( + components, + num_latent_frames: int, + latent_height: int, + latent_width: int, + num_audio_latents: int, + device: torch.device, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Draw the initial noise of both modalities and pack it into transformer rows. + + A request draws every stream from the one generator it is given, and the order is part of what that generator + reproduces: the conditioning noise of the keyframes or references first (one draw per condition, in + [`~modular_pipelines.minimax_h3.packing.keyframe_condition_noise`]), then the video noise here, as a latent tensor + that is patchified afterwards, then the audio noise, directly in row layout. Passing `latents` or + `audio_latents` skips its draw and shifts the ones after it. + + Args: + num_latent_frames (`int`): Number of video latent frames. + latent_height (`int`): Latent height. + latent_width (`int`): Latent width. + num_audio_latents (`int`): Number of audio latents per channel. + device (`torch.device`): The device the rows are drawn on. + generator (`torch.Generator`, *optional*): The generator of the request. + latents (`torch.Tensor`, *optional*): + Pre-generated video noise of shape `(1, latent_channels, num_latent_frames, latent_height, + latent_width)`, used instead of the draw. + audio_latents (`torch.Tensor`, *optional*): + Pre-generated audio noise of shape `(2, audio_latent_channels, num_audio_latents)`. + + Returns: + `tuple[torch.Tensor, torch.Tensor]`: the video rows and the channel-major audio rows. + """ + if latents is None: + latents = randn_tensor( + (1, components.vae_latent_channels, num_latent_frames, latent_height, latent_width), + generator=generator, + device=device, + dtype=torch.float32, + ) + video_rows = patchify_video_latents(latents.to(torch.float32), components.patch_size) + + if audio_latents is None: + audio_rows = randn_tensor( + (num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS, components.audio_latent_channels), + generator=generator, + device=device, + dtype=torch.float32, + ) + else: + audio_rows = audio_latents.to(torch.float32).permute(0, 2, 1).reshape(-1, components.audio_latent_channels) + return video_rows.to(device), audio_rows.to(device) + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + latents, audio_latents = self.prepare_latents( + components, + block_state.num_latent_frames, + block_state.latent_height, + block_state.latent_width, + block_state.num_audio_latents, + components._execution_device, + block_state.generator, + block_state.latents, + block_state.audio_latents, + ) + if block_state.condition_latents is not None: + latents = torch.cat([block_state.condition_latents, latents]) + if block_state.audio_condition_latents is not None: + audio_latents = torch.cat([block_state.audio_condition_latents, audio_latents]) + block_state.latents, block_state.audio_latents = latents, audio_latents + + self.set_block_state(state, block_state) + return components, state + + +class MiniMaxH3SetTimestepsStep(ModularPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return ( + "Initializes the two schedules — `shift = 12.0` for video, `shift = 3.0` for audio — and stages the " + "row-to-timestep plan of every step. One forward serves every modality and every noise level at once: " + "the generated rows step down their own schedule while the conditioning rows stay pinned at their " + "noise-augmentation level, and that assignment is static per step." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", MiniMaxH3Scheduler), + ComponentSpec("audio_scheduler", MiniMaxH3Scheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps", required=True), + InputParam( + name="layout", + type_hint=MiniMaxH3PackedSequence, + required=True, + description="The structural description of the packed sequence.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("timesteps", type_hint=torch.Tensor, description="Timesteps of the video schedule."), + OutputParam("audio_timesteps", type_hint=torch.Tensor, description="Timesteps of the audio schedule."), + OutputParam( + "row_timestep_plan", + type_hint=list, + description=( + "One `(timestep, timestep_indices)` pair per step: the distinct timesteps of the sequence and the " + "index of every row into them." + ), + ), + ] + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + components.scheduler.set_timesteps(block_state.num_inference_steps, device=device) + components.audio_scheduler.set_timesteps(block_state.num_inference_steps, device=device) + block_state.timesteps = components.scheduler.timesteps + block_state.audio_timesteps = components.audio_scheduler.timesteps + + block_state.row_timestep_plan = [ + tuple( + tensor.to(device) + for tensor in build_row_timesteps( + block_state.layout, + float(timestep), + float(audio_timestep), + max(float(timestep), MINIMAX_H3_KEYFRAME_NOISE_AUG), + 1.0, + ) + ) + for timestep, audio_timestep in zip(block_state.timesteps, block_state.audio_timesteps) + ] + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/minimax_h3/before_encoder.py b/diffusers/modular_pipelines/minimax_h3/before_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..88978bb6ff773faf0daa48faa5883eb5dd30baaf --- /dev/null +++ b/diffusers/modular_pipelines/minimax_h3/before_encoder.py @@ -0,0 +1,408 @@ +# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import PIL +import torch +from PIL import Image, ImageOps + +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import InputParam, OutputParam +from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline +from .packing import ( + MINIMAX_H3_CANVAS_MULTIPLE, + MINIMAX_H3_FPS, + MINIMAX_H3_MAX_DURATION, + MINIMAX_H3_MIN_DURATION, + align_num_frames, + audio_latent_num_frames, + prepare_keyframe_image, + resolve_canvas_size, + video_latent_num_frames, +) +from .packing_ref2va import ( + MINIMAX_H3_MAX_REFERENCE_AUDIOS, + MINIMAX_H3_MAX_REFERENCE_IMAGES, + MINIMAX_H3_MAX_REFERENCE_VIDEOS, + MINIMAX_H3_MAX_REFERENCES, + MiniMaxH3PreparedReference, + MiniMaxH3Reference, + prepare_reference_frames, + prepare_reference_image, + prepare_reference_waveform, + reference_kind, + reference_media_to_uint8, + resample_reference_frames, + resolve_reference_image_size, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _latent_geometry(components, height: int, width: int, num_frames: int) -> tuple[int, int, int, int]: + r"""The latent geometry the packed layout, the noise draws and the decoders all key off.""" + ratio = components.vae_spatial_compression_ratio + return video_latent_num_frames(num_frames), height // ratio, width // ratio, audio_latent_num_frames(num_frames) + + +def _latent_geometry_outputs() -> list[OutputParam]: + r"""The declaration of what [`_latent_geometry`] resolves, shared by the two setup blocks.""" + return [ + OutputParam("num_latent_frames", type_hint=int, description="Number of generated video latent frames."), + OutputParam("latent_height", type_hint=int, description="Height of the generated video latents."), + OutputParam("latent_width", type_hint=int, description="Width of the generated video latents."), + OutputParam("num_audio_latents", type_hint=int, description="Number of generated audio latents per channel."), + ] + + +class MiniMaxH3SetupStep(ModularPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return ( + "Resolves the plan shared by the `t2va` and `fl2va` tasks: the canvas (MiniMax-H3's own 768-short-edge " + "geometry for the aspect ratio of the first keyframe, or 16:9 without keyframes), the `17 * n + 5` frame " + "count the video VAE can decode, the latent geometry every later block keys off, and the keyframes put " + "onto that canvas." + ) + + @staticmethod + def _check_inputs(block_state) -> None: + if (block_state.height is None) != (block_state.width is None): + raise ValueError("`height` and `width` have to be passed together, or neither of them.") + if block_state.height is not None and ( + block_state.height % MINIMAX_H3_CANVAS_MULTIPLE or block_state.width % MINIMAX_H3_CANVAS_MULTIPLE + ): + raise ValueError( + f"`height` and `width` must be multiples of {MINIMAX_H3_CANVAS_MULTIPLE}, got " + f"{block_state.height}x{block_state.width}." + ) + # The duration the request generates is the one of the *aligned* frame count, so that is what the ceiling has + # to hold for: 346 frames would otherwise pass the check and then be rounded up to 362, i.e. 15.083 seconds. + aligned_num_frames = align_num_frames(block_state.num_frames) + duration = aligned_num_frames / MINIMAX_H3_FPS + if not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION: + raise ValueError( + f"MiniMax-H3 generates between {MINIMAX_H3_MIN_DURATION} and {MINIMAX_H3_MAX_DURATION} seconds at " + f"{MINIMAX_H3_FPS} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE can " + f"encode, must be between {int(MINIMAX_H3_MIN_DURATION * MINIMAX_H3_FPS)} and " + f"{int(MINIMAX_H3_MAX_DURATION * MINIMAX_H3_FPS)}, got {block_state.num_frames} (rounded up to " + f"{aligned_num_frames})." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="image", + type_hint=PIL.Image.Image, + description=( + "Keyframe the video starts from. It is *stretched* onto the target canvas, which by default is " + "derived from its own aspect ratio." + ), + ), + InputParam( + name="last_image", + type_hint=PIL.Image.Image, + description=( + "Keyframe the video ends on. Can be passed on its own to generate *up to* a frame. Combined with " + "`image` it is the follower of the two and is cover-cropped onto the canvas." + ), + ), + InputParam.template("height", description="Height of the generated video in pixels, a multiple of 32."), + InputParam.template("width", description="Width of the generated video in pixels, a multiple of 32."), + InputParam( + name="num_frames", + type_hint=int, + default=124, + description=( + "Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video " + "VAE can decode; the resulting duration must stay between 5 and 15 seconds." + ), + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."), + OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."), + OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."), + *_latent_geometry_outputs(), + OutputParam( + "keyframes", + type_hint=list, + description="The keyframes put onto the target canvas, in packed order (empty for `t2va`).", + ), + OutputParam( + "keyframe_anchors", + type_hint=tuple, + description="Which end of the video every keyframe is anchored to, in packed order.", + ), + ] + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self._check_inputs(block_state) + + keyframes = [ + ImageOps.exif_transpose(keyframe).convert("RGB") + for keyframe in (block_state.image, block_state.last_image) + if keyframe is not None + ] + block_state.keyframe_anchors = tuple( + anchor + for anchor, keyframe in (("first", block_state.image), ("last", block_state.last_image)) + if keyframe is not None + ) + if block_state.height is None: + block_state.height, block_state.width = resolve_canvas_size(*(keyframes[0].size if keyframes else (16, 9))) + + aligned_num_frames = align_num_frames(block_state.num_frames) + if aligned_num_frames != block_state.num_frames: + logger.warning( + f"`num_frames` has to be of the form 17 * n + 5 for the video VAE; rounding {block_state.num_frames} " + f"up to {aligned_num_frames}." + ) + block_state.num_frames = aligned_num_frames + + ( + block_state.num_latent_frames, + block_state.latent_height, + block_state.latent_width, + block_state.num_audio_latents, + ) = _latent_geometry(components, block_state.height, block_state.width, block_state.num_frames) + + block_state.keyframes = [ + prepare_keyframe_image(keyframe, block_state.height, block_state.width, stretch=index == 0) + for index, keyframe in enumerate(keyframes) + ] + self.set_block_state(state, block_state) + return components, state + + +class MiniMaxH3Ref2VASetupStep(ModularPipelineBlocks): + model_name = "minimax-h3-ref2va" + + @property + def description(self) -> str: + return ( + "Resolves the `ref2va` plan: the canvas (MiniMax-H3's own 16:9 unless asked otherwise — references never " + "bind the generated geometry), the references prepared at their own resolutions, the frame count they " + "imply when it was left open, and the latent geometry every later block keys off." + ) + + @staticmethod + def _check_inputs(components, block_state) -> None: + if (block_state.height is None) != (block_state.width is None): + raise ValueError("`height` and `width` have to be passed together, or neither of them.") + if block_state.height is not None and ( + block_state.height % MINIMAX_H3_CANVAS_MULTIPLE or block_state.width % MINIMAX_H3_CANVAS_MULTIPLE + ): + raise ValueError( + f"`height` and `width` must be multiples of {MINIMAX_H3_CANVAS_MULTIPLE}, got " + f"{block_state.height}x{block_state.width}." + ) + # The duration the request generates is the one of the *aligned* frame count, so that is what the ceiling has + # to hold for: 346 frames would otherwise pass the check and then be rounded up to 362, i.e. 15.083 seconds. + aligned_num_frames = None if block_state.num_frames is None else align_num_frames(block_state.num_frames) + duration = None if aligned_num_frames is None else aligned_num_frames / MINIMAX_H3_FPS + if duration is not None and not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION: + raise ValueError( + f"MiniMax-H3 generates between {MINIMAX_H3_MIN_DURATION} and {MINIMAX_H3_MAX_DURATION} seconds at " + f"{MINIMAX_H3_FPS} fps, so `num_frames`, rounded up to the next `17 * n + 5` the video VAE can " + f"encode, must be between {int(MINIMAX_H3_MIN_DURATION * MINIMAX_H3_FPS)} and " + f"{int(MINIMAX_H3_MAX_DURATION * MINIMAX_H3_FPS)}, got {block_state.num_frames} (rounded up to " + f"{aligned_num_frames})." + ) + + if not block_state.references: + raise ValueError( + "`ref2va` needs at least one reference; use `MiniMaxH3ModularPipeline` for text-only requests." + ) + kinds = [reference_kind(index, entry) for index, entry in enumerate(block_state.references)] + for kind, limit in ( + ("image", MINIMAX_H3_MAX_REFERENCE_IMAGES), + ("video", MINIMAX_H3_MAX_REFERENCE_VIDEOS), + ("audio", MINIMAX_H3_MAX_REFERENCE_AUDIOS), + ): + if kinds.count(kind) > limit: + raise ValueError(f"MiniMax-H3 accepts at most {limit} {kind} references, got {kinds.count(kind)}.") + if len(kinds) > MINIMAX_H3_MAX_REFERENCES: + raise ValueError( + f"MiniMax-H3 accepts at most {MINIMAX_H3_MAX_REFERENCES} references in total, got {len(kinds)}." + ) + if set(kinds) == {"audio"}: + raise ValueError( + "An audio reference has to be paired with at least one image or video reference and cannot be used " + "on its own." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="references", + type_hint=list[MiniMaxH3Reference], + required=True, + description=( + "The references to condition on, **in the order the model should read them**: the order labels " + "them in the prompt presentation and lays them out on the shared rotary clock, so a different " + "order is a different request. Every [`MiniMaxH3Reference`] carries exactly one medium, a path or " + "in-memory media — `image` (at most 9), `video` at its own `fps` (at most 3, whose `audio` " + "soundtrack is conditioned on as well), or `audio` at its own `sample_rate` (at most 3) — for at " + "most 12 references in total, and audio references cannot be the only ones. A path is decoded " + "when the reference is built, so these blocks only ever see pixels and samples." + ), + ), + InputParam.template("height", description="Height of the generated video in pixels, a multiple of 32."), + InputParam.template("width", description="Width of the generated video in pixels, a multiple of 32."), + InputParam( + name="num_frames", + type_hint=int, + description=( + "Number of frames to generate, at the fixed 24 fps. Snapped up to the next `17 * n + 5` the video " + "VAE can decode. May be left out, but only when exactly one reference carries audio, in which " + "case the duration is that soundtrack's." + ), + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."), + OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."), + OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."), + *_latent_geometry_outputs(), + OutputParam( + "prepared_references", + type_hint=list[MiniMaxH3PreparedReference], + description="The references prepared at their own resolutions, in packed order.", + ), + ] + + @staticmethod + def prepare_references( + components, references: list[MiniMaxH3Reference], num_frames: int | None + ) -> tuple[list[MiniMaxH3PreparedReference], int]: + r""" + Resolve the references and, if it was left open, the duration they imply. + + Every reference is prepared at its own resolution: an image is resized to a 2048 pixel short edge, a video is + resampled onto MiniMax-H3's own 24 fps, rescaled onto the 768 pixel canvas of *its own* aspect ratio and + truncated to the generated frame count, and a soundtrack is put on the audio VAE's sample rate and truncated to + the generated duration. None of this touches the target canvas. + + A reference that left its `fps` or its `sample_rate` out is taken to already be at MiniMax-H3's own rate, and + its frames or its samples then flow through untouched. + + A video reference goes through the two passes the reference implementation's `ffmpeg` decode applied, in the + same order: the constant frame rate resample of `resample_reference_frames` and the LANCZOS rescale of + `prepare_reference_frames`. Frames handed over at 24 fps and already at the canvas their own aspect ratio + resolves to therefore reach the VAE untouched, which is the parity-exact route. + + Args: + references (`list[MiniMaxH3Reference]`): + The `references` input of a [`MiniMaxH3Ref2VABlocks`] request. + num_frames (`int`, *optional*): + The requested frame count, or `None` to derive it from the single audio-bearing reference. + + Returns: + `tuple[list[MiniMaxH3PreparedReference], int]`: the prepared references, in packed order, and the frame + count. + """ + resolved = [ + MiniMaxH3PreparedReference(kind=reference_kind(index, entry), has_audio=entry.has_audio) + for index, entry in enumerate(references) + ] + + # The duration may be left open, but then exactly one reference may carry audio, or the request is ambiguous. + if num_frames is None: + audio_bearing = [index for index, reference in enumerate(resolved) if reference.has_audio] + if len(audio_bearing) != 1: + raise ValueError( + "`num_frames` may only be left to the references when exactly one of them carries audio, got " + f"{len(audio_bearing)}." + ) + index = audio_bearing[0] + sample_rate = references[index].sample_rate or components.audio_sampling_rate + duration = references[index].audio.shape[-1] / sample_rate + if not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION: + raise ValueError( + f"`references[{index}]` is {duration:g} seconds long, outside the " + f"{MINIMAX_H3_MIN_DURATION} to {MINIMAX_H3_MAX_DURATION} seconds MiniMax-H3 generates." + ) + num_frames = align_num_frames(round(duration * MINIMAX_H3_FPS)) + # The duration the request generates is the one of the *aligned* frame count, so that is what the + # ceiling has to hold for: a 14.99 second soundtrack rounds up to 362 frames, i.e. 15.083 seconds. + if num_frames / MINIMAX_H3_FPS > MINIMAX_H3_MAX_DURATION: + raise ValueError( + f"`references[{index}]` is {duration:g} seconds long, which rounds up to {num_frames} frames " + f"(`17 * n + 5`), i.e. {num_frames / MINIMAX_H3_FPS:g} seconds — past the " + f"{MINIMAX_H3_MAX_DURATION} seconds MiniMax-H3 generates. Pass `num_frames` to generate a " + "shorter video from this soundtrack." + ) + num_frames = align_num_frames(num_frames) + + for reference, entry in zip(resolved, references): + if reference.kind == "image": + image = entry.image + if not isinstance(image, Image.Image): + image = Image.fromarray(reference_media_to_uint8(image)) + image = ImageOps.exif_transpose(image).convert("RGB") + height, width = resolve_reference_image_size(*image.size) + reference.image = prepare_reference_image(image, height, width) + elif reference.kind == "video": + frames = resample_reference_frames(reference_media_to_uint8(entry.video), float(entry.fps)) + reference.frames = prepare_reference_frames(frames, num_frames) + if reference.has_audio: + reference.waveform = prepare_reference_waveform( + entry.audio, + entry.sample_rate or components.audio_sampling_rate, + components.audio_sampling_rate, + max_duration=num_frames / MINIMAX_H3_FPS, + ) + return resolved, num_frames + + @torch.no_grad() + def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self._check_inputs(components, block_state) + + if block_state.height is None: + block_state.height, block_state.width = resolve_canvas_size(16, 9) + + requested_num_frames = block_state.num_frames + block_state.prepared_references, block_state.num_frames = self.prepare_references( + components, block_state.references, block_state.num_frames + ) + if requested_num_frames is not None and requested_num_frames != block_state.num_frames: + logger.warning( + f"`num_frames` has to be of the form 17 * n + 5 for the video VAE; rounding {requested_num_frames} up " + f"to {block_state.num_frames}." + ) + + ( + block_state.num_latent_frames, + block_state.latent_height, + block_state.latent_width, + block_state.num_audio_latents, + ) = _latent_geometry(components, block_state.height, block_state.width, block_state.num_frames) + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/minimax_h3/decoders.py b/diffusers/modular_pipelines/minimax_h3/decoders.py new file mode 100644 index 0000000000000000000000000000000000000000..fc4ee359267b790a8a52c9301e055c7d5369af60 --- /dev/null +++ b/diffusers/modular_pipelines/minimax_h3/decoders.py @@ -0,0 +1,198 @@ +# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...configuration_utils import FrozenDict +from ...models import AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio +from ...utils import logging +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import MiniMaxH3ModularPipeline +from .packing import ( + MINIMAX_H3_PIXEL_MEAN, + MINIMAX_H3_PIXEL_STD, + unpack_audio_tokens, + unpatchify_video_tokens, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class MiniMaxH3VideoDecodeStep(ModularPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return ( + "Unpacks the generated video rows back into latents, denormalizes them and decodes them into video. The " + "spatial tiling of the video VAE covers the canvas exactly, so the decoded frames need no crop back, but " + "the decode itself runs under float16 autocast even though the VAE weights are float32, and the VAE " + "produces ImageNet-normalized RGB that is reverted here." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLMiniMaxH3), + ComponentSpec( + "video_processor", + VideoProcessor, + # The video VAE decodes into ImageNet-normalized RGB over a [0, 1] base range, which this block + # reverts itself, so the processor must not denormalize a second time. + config=FrozenDict({"vae_scale_factor": 16, "do_normalize": False}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", + type_hint=torch.Tensor, + required=True, + description="The denoised video rows of the packed sequence, conditioning rows first.", + ), + InputParam( + name="num_condition_video_rows", + type_hint=int, + default=0, + description="How many leading video rows are conditioning rows and are dropped here.", + ), + InputParam( + name="num_latent_frames", type_hint=int, required=True, description="Number of video latent frames." + ), + InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."), + InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."), + InputParam.template( + "output_type", description="Output format: 'pil', 'np', 'pt' or 'latent' for the raw latents." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("videos", description="The generated video.")] + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + latents = unpatchify_video_tokens( + block_state.latents[block_state.num_condition_video_rows :], + block_state.num_latent_frames, + block_state.latent_height, + block_state.latent_width, + components.vae_latent_channels, + components.patch_size, + ) + latents_mean = torch.tensor(components.vae.config.latents_mean, device=device).view(1, -1, 1, 1, 1) + latents_std = torch.tensor(components.vae.config.latents_std, device=device).view(1, -1, 1, 1, 1) + latents = latents * latents_std + latents_mean + + if block_state.output_type == "latent": + block_state.videos = latents + else: + with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"): + video = components.vae.decode(latents, return_dict=False)[0] + pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device=device).view(1, -1, 1, 1, 1) + pixel_std = torch.tensor(MINIMAX_H3_PIXEL_STD, device=device).view(1, -1, 1, 1, 1) + video = (video.float() * pixel_std + pixel_mean).clamp(0, 1) + block_state.videos = components.video_processor.postprocess_video( + video, output_type=block_state.output_type + ) + + self.set_block_state(state, block_state) + return components, state + + +class MiniMaxH3AudioDecodeStep(ModularPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return ( + "Unpacks the generated audio rows back into latents, denormalizes them and decodes them into a stereo " + "waveform. The audio VAE is mono and takes the two stereo channels as two batch items." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("audio_vae", AutoencoderKLMiniMaxH3Audio)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="audio_latents", + type_hint=torch.Tensor, + required=True, + description="The denoised audio rows of the packed sequence, reference rows first.", + ), + InputParam( + name="num_condition_audio_rows", + type_hint=int, + default=0, + description="How many leading audio rows are reference rows and are dropped here.", + ), + InputParam( + name="num_audio_latents", + type_hint=int, + required=True, + description="Number of audio latents per channel.", + ), + InputParam.template( + "output_type", description="Output format: 'pil', 'np', 'pt' or 'latent' for the raw latents." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "audio", + type_hint=torch.Tensor, + description="The generated soundtrack, of shape `(1, 2, num_samples)`.", + ), + OutputParam( + "sampling_rate", + type_hint=int, + description="Sample rate of the generated soundtrack in Hz.", + ), + ] + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + audio_latents = unpack_audio_tokens( + block_state.audio_latents[block_state.num_condition_audio_rows :], block_state.num_audio_latents + ) + audio_latents_mean = torch.tensor(components.audio_vae.config.latents_mean, device=device).view(1, -1, 1) + audio_latents_std = torch.tensor(components.audio_vae.config.latents_std, device=device).view(1, -1, 1) + audio_latents = audio_latents * audio_latents_std + audio_latents_mean + + if block_state.output_type == "latent": + block_state.audio = audio_latents + else: + audio = components.audio_vae.decode(audio_latents, return_dict=False)[0] + block_state.audio = audio.float().permute(1, 0, 2) + block_state.sampling_rate = components.audio_sampling_rate + + self.set_block_state(state, block_state) + return components, state diff --git a/diffusers/modular_pipelines/minimax_h3/denoise.py b/diffusers/modular_pipelines/minimax_h3/denoise.py new file mode 100644 index 0000000000000000000000000000000000000000..d149273371bb33b72f20433f6ff9dc4885686652 --- /dev/null +++ b/diffusers/modular_pipelines/minimax_h3/denoise.py @@ -0,0 +1,325 @@ +# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...models import MiniMaxH3Transformer3DModel +from ...schedulers import MiniMaxH3Scheduler +from ...utils import logging +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + PipelineState, +) +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _denoiser_inputs() -> list[InputParam]: + r"""Everything one MiniMax-H3 forward reads, beyond the transformer itself.""" + return [ + InputParam( + name="latents", + type_hint=torch.Tensor, + required=True, + description="The video rows of the packed sequence, conditioning rows first.", + ), + InputParam( + name="audio_latents", + type_hint=torch.Tensor, + required=True, + description="The channel-major audio rows of the packed sequence, reference rows first.", + ), + InputParam.template("prompt_embeds"), + InputParam( + name="row_timestep_plan", + type_hint=list, + required=True, + description="One `(timestep, timestep_indices)` pair per step.", + ), + InputParam( + name="token_tags", type_hint=torch.Tensor, required=True, description="The modality tag of every row." + ), + InputParam( + name="position_ids", + type_hint=torch.Tensor, + required=True, + description="The `(t, h, w)` rotary coordinate of every row.", + ), + InputParam( + name="video_indices", + type_hint=torch.Tensor, + required=True, + description="Sequence positions of the video rows.", + ), + InputParam( + name="audio_indices", + type_hint=torch.Tensor, + required=True, + description="Sequence positions of the audio rows.", + ), + InputParam( + name="text_indices", + type_hint=torch.Tensor, + required=True, + description="Sequence positions of the text rows.", + ), + InputParam.template("attention_kwargs"), + ] + + +def _denoiser_outputs() -> list[OutputParam]: + return [ + OutputParam( + "noise_pred", type_hint=torch.Tensor, description="Predicted velocity of the video rows of the sequence." + ), + OutputParam( + "audio_noise_pred", + type_hint=torch.Tensor, + description="Predicted velocity of the audio rows of the sequence.", + ), + ] + + +def _predict_velocity(transformer: MiniMaxH3Transformer3DModel, block_state: BlockState, i: int): + r"""One MiniMax-H3 forward pass: every row of the packed sequence, at its own noise level, at once.""" + unique_timesteps, timestep_indices = block_state.row_timestep_plan[i] + return transformer( + hidden_states=block_state.latents[None], + audio_hidden_states=block_state.audio_latents[None], + encoder_hidden_states=block_state.prompt_embeds, + timestep=unique_timesteps, + timestep_indices=timestep_indices, + token_tags=block_state.token_tags, + position_ids=block_state.position_ids, + video_indices=block_state.video_indices, + audio_indices=block_state.audio_indices, + text_indices=block_state.text_indices, + attention_kwargs=block_state.attention_kwargs, + return_dict=False, + ) + + +class MiniMaxH3LoopDenoiser(ModularPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return ( + "Runs the one MiniMax-H3 forward pass of a denoising iteration, which predicts the velocity of every row " + "of the packed sequence at once. The checkpoint is guidance-distilled, so there is no unconditional pass " + "and no guider." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", MiniMaxH3Transformer3DModel)] + + @property + def inputs(self) -> list[InputParam]: + return _denoiser_inputs() + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return _denoiser_outputs() + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + block_state.noise_pred, block_state.audio_noise_pred = _predict_velocity( + components.transformer, block_state, i + ) + return components, block_state + + +class MiniMaxH3Ref2VALoopDenoiser(ModularPipelineBlocks): + model_name = "minimax-h3-ref2va" + + @property + def description(self) -> str: + return ( + "Runs the one MiniMax-H3 forward pass of a `ref2va` denoising iteration, against the `transformer_ref` " + "partition of the checkpoint." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer_ref", MiniMaxH3Transformer3DModel)] + + @property + def inputs(self) -> list[InputParam]: + return _denoiser_inputs() + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return _denoiser_outputs() + + @torch.no_grad() + def __call__(self, components: MiniMaxH3Ref2VAModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + block_state.noise_pred, block_state.audio_noise_pred = _predict_velocity( + components.transformer_ref, block_state, i + ) + return components, block_state + + +class MiniMaxH3LoopSchedulerStep(ModularPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return ( + "Steps the generated video and audio rows down their own schedule. The conditioning rows are re-imposed " + "by construction: only the generated rows are ever written, so the anchors survive the whole loop." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", MiniMaxH3Scheduler), + ComponentSpec("audio_scheduler", MiniMaxH3Scheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", + type_hint=torch.Tensor, + required=True, + description="The video rows of the packed sequence, conditioning rows first.", + ), + InputParam( + name="audio_latents", + type_hint=torch.Tensor, + required=True, + description="The channel-major audio rows of the packed sequence, reference rows first.", + ), + InputParam( + name="noise_pred", + type_hint=torch.Tensor, + required=True, + description="Predicted velocity of the video rows.", + ), + InputParam( + name="audio_noise_pred", + type_hint=torch.Tensor, + required=True, + description="Predicted velocity of the audio rows.", + ), + InputParam( + name="audio_timesteps", + type_hint=torch.Tensor, + required=True, + description="Timesteps of the audio schedule.", + ), + InputParam( + name="num_condition_video_rows", + type_hint=int, + default=0, + description="How many leading video rows are conditioning rows.", + ), + InputParam( + name="num_condition_audio_rows", + type_hint=int, + default=0, + description="How many leading audio rows are reference rows.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="The video rows of the packed sequence after one step.", + ), + OutputParam( + "audio_latents", + type_hint=torch.Tensor, + description="The audio rows of the packed sequence after one step.", + ), + ] + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + num_condition_video_rows = block_state.num_condition_video_rows + num_condition_audio_rows = block_state.num_condition_audio_rows + + block_state.latents[num_condition_video_rows:] = components.scheduler.step( + block_state.noise_pred[0, num_condition_video_rows:].float(), + t, + block_state.latents[num_condition_video_rows:], + return_dict=False, + )[0] + block_state.audio_latents[num_condition_audio_rows:] = components.audio_scheduler.step( + block_state.audio_noise_pred[0, num_condition_audio_rows:].float(), + block_state.audio_timesteps[i], + block_state.audio_latents[num_condition_audio_rows:], + return_dict=False, + )[0] + return components, block_state + + +class MiniMaxH3DenoiseLoopWrapper(LoopSequentialPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return "Iteratively denoises the packed MiniMax-H3 sequence over the two schedules." + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", MiniMaxH3Scheduler), + ComponentSpec("audio_scheduler", MiniMaxH3Scheduler), + ] + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam.template("timesteps", required=True, description="Timesteps of the video schedule."), + ] + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + with self.progress_bar(total=len(block_state.timesteps)) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, block_state = self.loop_step(components, block_state, i=i, t=t) + progress_bar.update() + self.set_block_state(state, block_state) + return components, state + + +class MiniMaxH3DenoiseStep(MiniMaxH3DenoiseLoopWrapper): + block_classes = [MiniMaxH3LoopDenoiser, MiniMaxH3LoopSchedulerStep] + block_names = ["denoiser", "update"] + + @property + def description(self) -> str: + return "Runs the `t2va` / `fl2va` MiniMax-H3 denoising loop, one forward pass per step." + + +class MiniMaxH3Ref2VADenoiseStep(MiniMaxH3DenoiseLoopWrapper): + model_name = "minimax-h3-ref2va" + block_classes = [MiniMaxH3Ref2VALoopDenoiser, MiniMaxH3LoopSchedulerStep] + block_names = ["denoiser", "update"] + + @property + def description(self) -> str: + return "Runs the `ref2va` MiniMax-H3 denoising loop, one forward pass per step." diff --git a/diffusers/modular_pipelines/minimax_h3/encoders.py b/diffusers/modular_pipelines/minimax_h3/encoders.py new file mode 100644 index 0000000000000000000000000000000000000000..da2d611eaba0963cfc70e379735f0ea7824af370 --- /dev/null +++ b/diffusers/modular_pipelines/minimax_h3/encoders.py @@ -0,0 +1,638 @@ +# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import torch +from transformers import Qwen2TokenizerFast, Qwen3VLForConditionalGeneration, Qwen3VLProcessor + +from ...models import AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio +from ...models.autoencoders.vae import DiagonalGaussianDistribution +from ...schedulers import MiniMaxH3Scheduler +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline +from .packing import ( + MINIMAX_H3_KEYFRAME_ENCODE_SEED, + MINIMAX_H3_KEYFRAME_NOISE_AUG, + MINIMAX_H3_PIXEL_MEAN, + MINIMAX_H3_PIXEL_STD, + MINIMAX_H3_TEXT_ENCODER_LAYER, + MINIMAX_H3_TEXT_TAG, + MINIMAX_H3_VIDEO_TAG, + keyframe_condition_noise, + patchify_video_latents, +) +from .packing_ref2va import ( + MiniMaxH3PreparedReference, + build_ref2va_presentation, + sample_reference_video_frames, + trim_reference_num_frames, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _check_prompt(prompt) -> None: + r"""MiniMax-H3 packs one request into one sequence, so a batch of prompts is not a thing.""" + if not isinstance(prompt, str): + raise ValueError( + f"MiniMax-H3 packs one request into one sequence, so `prompt` must be a single string, got {type(prompt)}." + ) + + +def _conditioner_components() -> list[ComponentSpec]: + r"""MiniMax-H3's conditioner: a Qwen3-VL read at its 50th decoder layer, with its language-model head unused.""" + return [ + ComponentSpec("text_encoder", Qwen3VLForConditionalGeneration), + ComponentSpec("tokenizer", Qwen2TokenizerFast), + ComponentSpec("processor", Qwen3VLProcessor), + ] + + +def _conditioner_outputs() -> list[OutputParam]: + return [ + OutputParam.template( + "prompt_embeds", + description=( + "The hidden state MiniMax-H3 conditions on, of shape `(1, num_text_tokens, 5120)`, read after the " + "50th decoder layer of the Qwen3-VL conditioner." + ), + ), + OutputParam( + "text_token_tags", + type_hint=torch.Tensor, + description="The per-row modality tag of every row of `prompt_embeds`; a vision block is tagged as video.", + ), + ] + + +class MiniMaxH3TextEncoderStep(ModularPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return ( + "Encodes MiniMax-H3's presentation of a `t2va` / `fl2va` request: the prompt verbatim, preceded by a " + '`": "` label and a vision block per keyframe, with no chat template and no special tokens. ' + "The checkpoint is guidance-distilled, so there is no negative prompt and no unconditional branch." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return _conditioner_components() + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", description="The prompt to guide generation, a single string."), + InputParam( + name="keyframes", + type_hint=list, + description="The keyframes put onto the target canvas, in packed order (empty or None for `t2va`).", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return _conditioner_outputs() + + @staticmethod + def encode_prompt( + components, + prompt: str, + images: list | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Build MiniMax-H3's presentation of a request and encode it. + + The presentation is the verbatim prompt for `t2va`. Every keyframe prepends a `": "` label and a + vision block (`<|vision_start|>`, one `<|image_pad|>` per vision patch, `<|vision_end|>`) — no chat template + and no special tokens. The rows of a vision block are tagged as *video* rather than text, which is what the + transformer's AdaLN modulation keys off. + + Args: + prompt (`str`): The prompt to encode. + images (`list[PIL.Image.Image]`, *optional*): + The keyframes, already prepared onto the target canvas, in packed order. + device (`torch.device`, *optional*): The device to run the conditioner on. + dtype (`torch.dtype`, *optional*): The dtype of the returned embeddings. + + Returns: + `tuple[torch.Tensor, torch.Tensor]`: the `(1, num_text_tokens, 5120)` hidden states and the + `(num_text_tokens,)` per-row modality tags. + """ + device = device or components._execution_device + dtype = dtype or components.transformer.dtype + + num_layers = components.text_encoder.config.text_config.num_hidden_layers + if num_layers <= MINIMAX_H3_TEXT_ENCODER_LAYER: + raise ValueError( + f"MiniMax-H3 conditions on `hidden_states[{MINIMAX_H3_TEXT_ENCODER_LAYER}]` of its Qwen3-VL " + f"conditioner, which needs more than {MINIMAX_H3_TEXT_ENCODER_LAYER} decoder layers, but " + f"`text_encoder` has {num_layers}. The last hidden state of a stack truncated to exactly " + f"{MINIMAX_H3_TEXT_ENCODER_LAYER} layers is post-norm and is not the conditioning MiniMax-H3 expects." + ) + + pixel_values, image_grid_thw = None, None + token_ids, token_tags = [], [] + if images: + vision = components.processor.image_processor(images=images, return_tensors="pt") + pixel_values, image_grid_thw = vision["pixel_values"], vision["image_grid_thw"] + merge_size = components.processor.image_processor.merge_size**2 + for index in range(len(images)): + num_image_tokens = int(image_grid_thw[index].prod()) // merge_size + label_ids = components.tokenizer(f": ", add_special_tokens=False)["input_ids"] + vision_ids = ( + [components.tokenizer.convert_tokens_to_ids("<|vision_start|>")] + + [components.tokenizer.convert_tokens_to_ids("<|image_pad|>")] * num_image_tokens + + [components.tokenizer.convert_tokens_to_ids("<|vision_end|>")] + ) + token_ids += label_ids + vision_ids + token_tags += [MINIMAX_H3_TEXT_TAG] * len(label_ids) + [MINIMAX_H3_VIDEO_TAG] * len(vision_ids) + prompt_ids = components.tokenizer(prompt, add_special_tokens=False)["input_ids"] + token_ids += prompt_ids + token_tags += [MINIMAX_H3_TEXT_TAG] * len(prompt_ids) + + input_ids = torch.tensor([token_ids], dtype=torch.long, device=device) + # Qwen3-VL lays its 3D rotary positions out per modality run, which it reads off the token type ids the + # processor derives from the vision pad ids (`0` text, `1` image, `2` video). + mm_token_type_ids = torch.tensor( + components.processor.create_mm_token_type_ids([token_ids]), dtype=torch.long, device=device + ) + # `text_encoder.model` is a submodule, and a CPU-offload hook — accelerate's or the one the + # `ComponentsManager` attaches — wraps the *top-level* module's `forward` alone, so calling the submodule + # directly would leave the conditioner on the CPU. Fire the hook by hand instead of routing through + # `text_encoder(...)`: MiniMax-H3 reads `hidden_states[50]` and never uses the language-model head, whose + # vocabulary-wide projection over every token is all the top-level forward would add. + hook = getattr(components.text_encoder, "_hf_hook", None) + if hook is not None and hasattr(hook, "pre_forward"): + hook.pre_forward(components.text_encoder) + outputs = components.text_encoder.model( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + mm_token_type_ids=mm_token_type_ids, + pixel_values=None if pixel_values is None else pixel_values.to(device, components.text_encoder.dtype), + image_grid_thw=None if image_grid_thw is None else image_grid_thw.to(device), + use_cache=False, + output_hidden_states=True, + ) + prompt_embeds = outputs.hidden_states[MINIMAX_H3_TEXT_ENCODER_LAYER].to(device=device, dtype=dtype) + return prompt_embeds, torch.tensor(token_tags, dtype=torch.long) + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + _check_prompt(block_state.prompt) + + # `encode_prompt` defaults the embedding dtype to the denoiser's; a text encoder block has no denoiser of + # its own — it is meant to run on its own — so it emits the conditioner's dtype, as every other model does. + block_state.prompt_embeds, block_state.text_token_tags = self.encode_prompt( + components, + block_state.prompt, + block_state.keyframes, + device=components._execution_device, + dtype=components.text_encoder.dtype, + ) + + self.set_block_state(state, block_state) + return components, state + + +class MiniMaxH3KeyframeVaeEncoderStep(ModularPipelineBlocks): + model_name = "minimax-h3" + + @property + def description(self) -> str: + return ( + "Encodes the `fl2va` keyframes into packed conditioning rows and noises them to MiniMax-H3's " + "conditioning level. The rows are the anchors of the whole denoising loop: the loop only ever writes the " + "generated rows, so they are never updated again." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLMiniMaxH3), + ComponentSpec("scheduler", MiniMaxH3Scheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="keyframes", + type_hint=list, + required=True, + description="The keyframes put onto the target canvas, in packed order.", + ), + InputParam(name="latent_height", type_hint=int, required=True, description="Height of the video latents."), + InputParam(name="latent_width", type_hint=int, required=True, description="Width of the video latents."), + InputParam.template( + "generator", + description=( + "The generator of the request. The conditioning noise is drawn from it before the target noise " + "of the prepare-latents step." + ), + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "condition_latents", + type_hint=torch.Tensor, + description="The noise-augmented video conditioning rows, in packed order.", + ) + ] + + @staticmethod + def encode_keyframes(components, images: list, device: torch.device | None = None) -> torch.Tensor: + r""" + Encode the `fl2va` keyframes into packed conditioning rows. + + The keyframes go through the video VAE's spatial encoder only — they are single frames, so none of its + 17-frame temporal chunking applies — and the posterior is *sampled*, under a generator seeded with 42 + independently of the request seed. The sampled latent is rounded to float16 before being normalized, as in the + reference implementation; both are part of reproducing the released model's conditioning. + + Args: + images (`list[PIL.Image.Image]`): + The keyframes, already prepared onto the target canvas, in packed order. + device (`torch.device`, *optional*): The device to run the VAE on. + + Returns: + `torch.Tensor` of shape `(num_condition_rows, latent_channels * prod(patch_size))`: the float32 + conditioning rows. + """ + device = device or components._execution_device + latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, -1, 1, 1, 1) + latents_std = torch.tensor(components.vae.config.latents_std).view(1, -1, 1, 1, 1) + pixel_mean = torch.tensor(MINIMAX_H3_PIXEL_MEAN, device=device).view(1, -1, 1, 1, 1) + pixel_std = torch.tensor(MINIMAX_H3_PIXEL_STD, device=device).view(1, -1, 1, 1, 1) + + rows = [] + for image in images: + pixels = torch.from_numpy(np.array(image)).to(device).permute(2, 0, 1)[None, :, None] + pixels = (pixels.to(torch.float32).div(255.0) - pixel_mean) / pixel_std + # `vae.encode` chunks along time for videos; a keyframe is one frame and is encoded by the (tiled) + # spatial encoder alone, which is what the released model conditions on. + moments = components.vae._encode_clip(pixels) + posterior = DiagonalGaussianDistribution(moments) + latents = posterior.sample(generator=torch.Generator().manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)) + # The sampled latent is rounded to float16 before it is normalized: ~11 bits of every conditioning + # latent, so the released model's conditioning cannot be reproduced without it. + latents = latents.to(torch.float16).float().cpu() + rows.append(patchify_video_latents((latents - latents_mean) / latents_std, components.patch_size)) + return torch.cat(rows) + + @torch.no_grad() + def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + condition_latents = self.encode_keyframes(components, block_state.keyframes, device=device) + noise = keyframe_condition_noise( + ((1, block_state.latent_height, block_state.latent_width),) * len(block_state.keyframes), + components.patch_size, + components.vae_latent_channels, + generator=block_state.generator, + device=device, + ) + block_state.condition_latents = components.scheduler.scale_noise( + condition_latents.to(device), MINIMAX_H3_KEYFRAME_NOISE_AUG, noise + ) + + self.set_block_state(state, block_state) + return components, state + + +class MiniMaxH3Ref2VATextEncoderStep(ModularPipelineBlocks): + model_name = "minimax-h3-ref2va" + + @property + def description(self) -> str: + return ( + "Encodes MiniMax-H3's presentation of a `ref2va` request: a label per reference, numbered per modality " + '(`": "` plus a vision block, `"