diff --git a/docs/source/en/api/pipelines/chroma.md b/docs/source/en/api/pipelines/chroma.md index 9ca1ff4dee68..7ec1e10849ba 100644 --- a/docs/source/en/api/pipelines/chroma.md +++ b/docs/source/en/api/pipelines/chroma.md @@ -105,3 +105,11 @@ image.save("chroma-single-file.png") [[autodoc]] ChromaInpaintPipeline - all - __call__ + +## ChromaModularPipeline + +[[autodoc]] ChromaModularPipeline + +## ChromaAutoBlocks + +[[autodoc]] ChromaAutoBlocks diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 2deff4b71cb2..7dd14152d6b6 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -503,6 +503,8 @@ [ "AnimaAutoBlocks", "AnimaModularPipeline", + "ChromaAutoBlocks", + "ChromaModularPipeline", "Cosmos3DistilledBlocks", "Cosmos3DistilledModularPipeline", "Cosmos3OmniBlocks", @@ -1329,6 +1331,8 @@ from .modular_pipelines import ( AnimaAutoBlocks, AnimaModularPipeline, + ChromaAutoBlocks, + ChromaModularPipeline, Cosmos3DistilledBlocks, Cosmos3DistilledModularPipeline, Cosmos3OmniBlocks, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 008a654c3fa3..7b7adbfcb57d 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -103,6 +103,10 @@ "AnimaAutoBlocks", "AnimaModularPipeline", ] + _import_structure["chroma"] = [ + "ChromaAutoBlocks", + "ChromaModularPipeline", + ] _import_structure["cosmos"] = [ "Cosmos3DistilledBlocks", "Cosmos3DistilledModularPipeline", @@ -139,6 +143,7 @@ from ..utils.dummy_pt_objects import * # noqa F403 else: from .anima import AnimaAutoBlocks, AnimaModularPipeline + from .chroma import ChromaAutoBlocks, ChromaModularPipeline from .components_manager import ComponentsManager from .cosmos import ( Cosmos3DistilledBlocks, diff --git a/src/diffusers/modular_pipelines/chroma/__init__.py b/src/diffusers/modular_pipelines/chroma/__init__.py new file mode 100644 index 000000000000..6ebeae8359fd --- /dev/null +++ b/src/diffusers/modular_pipelines/chroma/__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_chroma"] = ["ChromaAutoBlocks"] + _import_structure["modular_pipeline"] = ["ChromaModularPipeline"] + +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_chroma import ChromaAutoBlocks + from .modular_pipeline import ChromaModularPipeline +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/src/diffusers/modular_pipelines/chroma/before_denoise.py b/src/diffusers/modular_pipelines/chroma/before_denoise.py new file mode 100644 index 000000000000..9e079bd6c591 --- /dev/null +++ b/src/diffusers/modular_pipelines/chroma/before_denoise.py @@ -0,0 +1,411 @@ +# 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 ...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 ChromaModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-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 + + +# 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 + + +def _prepare_latent_image_ids(height, width, device, dtype): + latent_image_ids = torch.zeros(height, width, 3) + latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None] + latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :] + + latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape + + latent_image_ids = latent_image_ids.reshape( + latent_image_id_height * latent_image_id_width, latent_image_id_channels + ) + + return latent_image_ids.to(device=device, dtype=dtype) + + +def _pack_latents(latents, batch_size, num_channels_latents, height, width): + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 2, 4, 1, 3, 5) + latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4) + + return latents + + +class ChromaPrepareLatentsStep(ModularPipelineBlocks): + model_name = "chroma" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [] + + @property + def description(self) -> str: + return "Prepare latents step that prepares the latents for the text-to-image generation process" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("height"), + InputParam.template("width"), + InputParam.template("latents"), + InputParam.template("num_images_per_prompt"), + InputParam.template("generator"), + InputParam( + "batch_size", + required=True, + type_hint=int, + description="Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. Can be generated in input step.", + ), + InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="The initial noisy latents (patchified) to use for the denoising process", + ), + ] + + @staticmethod + def check_inputs(components, block_state): + if (block_state.height is not None and block_state.height % (components.vae_scale_factor * 2) != 0) or ( + block_state.width is not None and block_state.width % (components.vae_scale_factor * 2) != 0 + ): + logger.warning( + f"`height` and `width` have to be divisible by {components.vae_scale_factor * 2} but are {block_state.height} and {block_state.width}." + ) + + @staticmethod + def prepare_latents( + comp, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (comp.vae_scale_factor * 2)) + width = 2 * (int(width) // (comp.vae_scale_factor * 2)) + + shape = (batch_size, num_channels_latents, height, width) + + if latents is not None: + return latents.to(device=device, dtype=dtype) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(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=generator, device=device, dtype=dtype) + latents = _pack_latents(latents, batch_size, num_channels_latents, height, width) + + return latents + + @torch.no_grad() + def __call__(self, components: ChromaModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.height = block_state.height or components.default_height + block_state.width = block_state.width or components.default_width + device = components._execution_device + + self.check_inputs(components, block_state) + batch_size = block_state.batch_size * block_state.num_images_per_prompt + block_state.latents = self.prepare_latents( + components, + batch_size, + components.num_channels_latents, + block_state.height, + block_state.width, + block_state.dtype, + device, + block_state.generator, + block_state.latents, + ) + + self.set_block_state(state, block_state) + + return components, state + + +class ChromaSetTimestepsStep(ModularPipelineBlocks): + model_name = "chroma" + + @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. Should be run after the prepare latents step." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True, note="patchified, can be generated in prepare_latents step"), + InputParam.template("num_inference_steps", default=35), + InputParam.template("sigmas"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"), + OutputParam( + "num_inference_steps", + type_hint=int, + description="The number of denoising steps to perform at inference time", + ), + ] + + @torch.no_grad() + def __call__(self, components: ChromaModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + num_inference_steps = block_state.num_inference_steps + sigmas = ( + np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + if block_state.sigmas is None + else block_state.sigmas + ) + image_seq_len = block_state.latents.shape[1] + 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), + ) + block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps( + components.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + mu=mu, + ) + + self.set_block_state(state, block_state) + return components, state + + +class ChromaPrepareAttentionMaskStep(ModularPipelineBlocks): + model_name = "chroma" + + @property + def description(self) -> str: + return ( + "Step that extends the prompt attention masks to cover the image tokens in the final sequence. " + "Should be run after the text input and prepare latents steps." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True, note="patchified, can be generated in prepare_latents step"), + InputParam( + "prompt_attention_mask", + required=True, + type_hint=torch.Tensor, + description="Attention mask for the prompt embeddings. Can be generated from text_encoder step.", + ), + InputParam( + "negative_prompt_attention_mask", + type_hint=torch.Tensor, + description="Attention mask for the negative prompt embeddings. Can be generated from text_encoder step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "attention_mask", + type_hint=torch.Tensor, + description="Attention mask covering the text and image tokens of the final sequence.", + ), + OutputParam( + "negative_attention_mask", + type_hint=torch.Tensor, + description="Negative attention mask covering the text and image tokens of the final sequence.", + ), + ] + + @staticmethod + def prepare_attention_mask(batch_size, sequence_length, attention_mask): + if attention_mask is None: + return attention_mask + + # Extend the prompt attention mask to account for image tokens in the final sequence + attention_mask = torch.cat( + [attention_mask, torch.ones(batch_size, sequence_length, device=attention_mask.device, dtype=torch.bool)], + dim=1, + ) + + return attention_mask + + @torch.no_grad() + def __call__(self, components: ChromaModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + batch_size, image_seq_len = block_state.latents.shape[0], block_state.latents.shape[1] + + block_state.attention_mask = self.prepare_attention_mask( + batch_size, image_seq_len, block_state.prompt_attention_mask + ) + block_state.negative_attention_mask = self.prepare_attention_mask( + batch_size, image_seq_len, block_state.negative_prompt_attention_mask + ) + + self.set_block_state(state, block_state) + return components, state + + +class ChromaRoPEInputsStep(ModularPipelineBlocks): + model_name = "chroma" + + @property + def description(self) -> str: + return "Step that prepares the RoPE inputs for the denoising process. Should be placed after text encoder and latent preparation steps." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam.template("prompt_embeds"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="txt_ids", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="IDs computed from the text sequence, used for RoPE calculation.", + ), + OutputParam( + name="img_ids", + kwargs_type="denoiser_input_fields", + type_hint=torch.Tensor, + description="IDs computed from the latent sequence, used for RoPE calculation.", + ), + ] + + def __call__(self, components: ChromaModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + prompt_embeds = block_state.prompt_embeds + device, dtype = prompt_embeds.device, prompt_embeds.dtype + block_state.txt_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype) + + height = 2 * (int(block_state.height) // (components.vae_scale_factor * 2)) + width = 2 * (int(block_state.width) // (components.vae_scale_factor * 2)) + block_state.img_ids = _prepare_latent_image_ids(height // 2, width // 2, device, dtype) + + self.set_block_state(state, block_state) + + return components, state diff --git a/src/diffusers/modular_pipelines/chroma/decoders.py b/src/diffusers/modular_pipelines/chroma/decoders.py new file mode 100644 index 000000000000..a524f2b22234 --- /dev/null +++ b/src/diffusers/modular_pipelines/chroma/decoders.py @@ -0,0 +1,102 @@ +# 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 ...image_processor import VaeImageProcessor +from ...models import AutoencoderKL +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 + + +def _unpack_latents(latents, height, width, vae_scale_factor): + batch_size, num_patches, channels = latents.shape + + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (vae_scale_factor * 2)) + width = 2 * (int(width) // (vae_scale_factor * 2)) + + latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5) + + latents = latents.reshape(batch_size, channels // (2 * 2), height, width) + + return latents + + +class ChromaDecodeStep(ModularPipelineBlocks): + model_name = "chroma" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKL), + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def description(self) -> str: + return "Step that decodes the denoised latents into images" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("output_type"), + InputParam.template("height", default=1024), + InputParam.template("width", default=1024), + InputParam.template("latents", required=True, note="denoised latents from the denoise step"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "images", + type_hint=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", + ) + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + vae = components.vae + + if not block_state.output_type == "latent": + latents = block_state.latents + latents = _unpack_latents(latents, block_state.height, block_state.width, components.vae_scale_factor) + latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor + 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 + ) + else: + block_state.images = block_state.latents + + self.set_block_state(state, block_state) + + return components, state diff --git a/src/diffusers/modular_pipelines/chroma/denoise.py b/src/diffusers/modular_pipelines/chroma/denoise.py new file mode 100644 index 000000000000..a888d32d707b --- /dev/null +++ b/src/diffusers/modular_pipelines/chroma/denoise.py @@ -0,0 +1,279 @@ +# 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 ChromaTransformer2DModel +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 ChromaModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class ChromaLoopDenoiser(ModularPipelineBlocks): + model_name = "chroma" + + def __init__( + self, + guider_input_fields: dict[str, Any] | None = None, + ): + """Initialize a denoiser block that calls the denoiser model with guidance. This block is used in Chroma. + + Args: + guider_input_fields: A dictionary that maps each argument expected by the denoiser model + (for example, "encoder_hidden_states") to data stored on 'block_state'. The value can be either: + + - A tuple of strings. For instance, {"encoder_hidden_states": ("prompt_embeds", + "negative_prompt_embeds")} tells the guider to read `block_state.prompt_embeds` and + `block_state.negative_prompt_embeds` and pass them as the conditional and unconditional batches of + 'encoder_hidden_states'. + - A string. For example, {"encoder_hidden_image": "image_embeds"} makes the guider forward + `block_state.image_embeds` for both conditional and unconditional batches. + """ + if guider_input_fields is None: + guider_input_fields = { + "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"), + "attention_mask": ("attention_mask", "negative_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": 5.0}), + default_creation_method="from_config", + ), + ComponentSpec("transformer", ChromaTransformer2DModel), + ] + + @property + def description(self) -> str: + return ( + "Step within the denoising loop that denoise the latents with guidance. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `ChromaDenoiseLoopWrapper`)" + ) + + @property + def inputs(self) -> list[InputParam]: + inputs = [ + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The initial latents to use for the denoising process. Can be generated in prepare_latents step.", + ), + InputParam( + "num_inference_steps", + required=True, + type_hint=int, + description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.", + ), + InputParam( + "txt_ids", + required=True, + type_hint=torch.Tensor, + description="IDs computed from text sequence needed for RoPE", + ), + InputParam( + "img_ids", + required=True, + type_hint=torch.Tensor, + description="IDs computed from latent sequence needed for RoPE", + ), + InputParam( + "joint_attention_kwargs", + type_hint=dict, + description="Additional kwargs passed along to the attention processors.", + ), + InputParam.template("denoiser_input_fields"), + ] + + guider_input_names = [] + uncond_guider_input_names = [] + for value in self._guider_input_fields.values(): + if isinstance(value, tuple): + guider_input_names.append(value[0]) + uncond_guider_input_names.append(value[1]) + else: + guider_input_names.append(value) + + for name in guider_input_names: + inputs.append(InputParam(name=name, required=True)) + for name in uncond_guider_input_names: + inputs.append(InputParam(name=name)) + return inputs + + @torch.no_grad() + def __call__( + self, components: ChromaModularPipeline, 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) + guider_state = components.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields) + + latents = block_state.latents + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = {key: getattr(guider_state_batch, key) for key in self._guider_input_fields.keys()} + + guider_state_batch.noise_pred = components.transformer( + hidden_states=latents, + timestep=timestep / 1000, + txt_ids=block_state.txt_ids, + img_ids=block_state.img_ids, + joint_attention_kwargs=block_state.joint_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 ChromaLoopAfterDenoiser(ModularPipelineBlocks): + model_name = "chroma" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def description(self) -> str: + return ( + "step within the denoising loop that update the latents. " + "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` " + "object (e.g. `ChromaDenoiseLoopWrapper`)" + ) + + @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: ChromaModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + # Perform scheduler step using the predicted output + 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 ChromaDenoiseLoopWrapper(LoopSequentialPipelineBlocks): + model_name = "chroma" + + @property + def description(self) -> str: + return ( + "Pipeline block that iteratively denoise the latents over `timesteps`. " + "The specific steps with each iteration can be customized with `sub_blocks` attributes" + ) + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ComponentSpec("transformer", ChromaTransformer2DModel), + ] + + @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. Can be generated in set_timesteps step.", + ), + InputParam( + "num_inference_steps", + required=True, + type_hint=int, + description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.", + ), + ] + + @torch.no_grad() + def __call__(self, components: ChromaModularPipeline, 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 ChromaDenoiseStep(ChromaDenoiseLoopWrapper): + block_classes = [ + ChromaLoopDenoiser( + guider_input_fields={ + "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"), + "attention_mask": ("attention_mask", "negative_attention_mask"), + } + ), + ChromaLoopAfterDenoiser, + ] + block_names = ["denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step that iteratively denoise the latents. \n" + "Its loop logic is defined in `ChromaDenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `ChromaLoopDenoiser`\n" + " - `ChromaLoopAfterDenoiser`\n" + "This block supports the text2image task." + ) diff --git a/src/diffusers/modular_pipelines/chroma/encoders.py b/src/diffusers/modular_pipelines/chroma/encoders.py new file mode 100644 index 000000000000..5f9d76b8100b --- /dev/null +++ b/src/diffusers/modular_pipelines/chroma/encoders.py @@ -0,0 +1,247 @@ +# 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 ...loaders import FluxLoraLoaderMixin, TextualInversionLoaderMixin +from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import ChromaModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def get_t5_prompt_embeds( + components, + prompt: str | list[str], + max_sequence_length: int, + device: torch.device, +): + dtype = components.text_encoder.dtype + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + if isinstance(components, TextualInversionLoaderMixin): + prompt = components.maybe_convert_prompt(prompt, components.tokenizer) + + text_inputs = components.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_length=False, + return_overflowing_tokens=False, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + tokenizer_mask = text_inputs.attention_mask + + tokenizer_mask_device = tokenizer_mask.to(device) + + # unlike FLUX, Chroma uses the attention mask when generating the T5 embedding + prompt_embeds = components.text_encoder( + text_input_ids.to(device), + output_hidden_states=False, + attention_mask=tokenizer_mask_device, + )[0] + + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + # for the text tokens, chroma requires that all except the first padding token are masked out during the forward + # pass through the transformer + seq_lengths = tokenizer_mask_device.sum(dim=1) + mask_indices = torch.arange(tokenizer_mask_device.size(1), device=device).unsqueeze(0).expand(batch_size, -1) + attention_mask = (mask_indices <= seq_lengths.unsqueeze(1)).to(dtype=dtype, device=device) + + return prompt_embeds, attention_mask + + +class ChromaTextEncoderStep(ModularPipelineBlocks): + model_name = "chroma" + + @property + def description(self) -> str: + return ( + "Text Encoder step that generates T5 text embeddings and the attention masks Chroma uses to mask out " + "padding tokens (keeping one padding token unmasked, as required by Chroma)" + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", T5EncoderModel), + ComponentSpec("tokenizer", T5TokenizerFast), + 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"), + InputParam( + "joint_attention_kwargs", + type_hint=dict, + description="Additional kwargs for attention processors; `scale` is used as the text encoder LoRA scale.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("negative_prompt_embeds"), + OutputParam( + "prompt_attention_mask", + type_hint=torch.Tensor, + description="Attention mask for the prompt embeddings, with all padding tokens except the first one masked out.", + ), + OutputParam( + "negative_prompt_attention_mask", + type_hint=torch.Tensor, + description="Attention mask for the negative prompt embeddings, with all padding tokens except the first one masked out.", + ), + ] + + @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)}") + if block_state.max_sequence_length is not None and block_state.max_sequence_length > 512: + raise ValueError( + f"`max_sequence_length` cannot be greater than 512 but is {block_state.max_sequence_length}" + ) + + @staticmethod + def encode_prompt( + components, + prompt: str | list[str], + device: torch.device | None = None, + prepare_unconditional_embeds: bool = True, + negative_prompt: str | list[str] | None = None, + max_sequence_length: int = 512, + lora_scale: float | None = None, + ): + r""" + Encodes the prompt into T5 hidden states and builds the Chroma attention masks. + + Args: + prompt (`str` or `list[str]`): + prompt to be encoded + device: (`torch.device`): + torch device + prepare_unconditional_embeds (`bool`): + whether to prepare unconditional embeddings or not + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, defaults to `512`): + The maximum number of text tokens to be used for the generation process. + lora_scale (`float`, *optional*): + A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. + """ + device = device or components._execution_device + + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(components, FluxLoraLoaderMixin): + components._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + if components.text_encoder is not None and USE_PEFT_BACKEND: + scale_lora_layers(components.text_encoder, lora_scale) + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + prompt_embeds, prompt_attention_mask = get_t5_prompt_embeds( + components, + prompt=prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + + 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 prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif 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, + prompt=negative_prompt, + max_sequence_length=max_sequence_length, + device=device, + ) + + if components.text_encoder is not None: + if isinstance(components, FluxLoraLoaderMixin) and USE_PEFT_BACKEND: + # Retrieve the original scale by scaling back the LoRA layers + unscale_lora_layers(components.text_encoder, lora_scale) + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + @torch.no_grad() + def __call__(self, components: ChromaModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + + device = components._execution_device + + lora_scale = ( + block_state.joint_attention_kwargs.get("scale", None) + if block_state.joint_attention_kwargs is not None + else None + ) + ( + block_state.prompt_embeds, + block_state.prompt_attention_mask, + block_state.negative_prompt_embeds, + block_state.negative_prompt_attention_mask, + ) = self.encode_prompt( + components, + prompt=block_state.prompt, + device=device, + prepare_unconditional_embeds=components.requires_unconditional_embeds, + negative_prompt=block_state.negative_prompt, + max_sequence_length=block_state.max_sequence_length, + lora_scale=lora_scale, + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/chroma/inputs.py b/src/diffusers/modular_pipelines/chroma/inputs.py new file mode 100644 index 000000000000..6c9cdeb3ca5b --- /dev/null +++ b/src/diffusers/modular_pipelines/chroma/inputs.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. + + +import torch + +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import InputParam, OutputParam +from .modular_pipeline import ChromaModularPipeline + + +logger = logging.get_logger(__name__) + + +class ChromaTextInputStep(ModularPipelineBlocks): + model_name = "chroma" + + @property + def description(self) -> str: + return ( + "Text input processing step that standardizes text embeddings and attention masks for the pipeline.\n" + "This step:\n" + " 1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n" + " 2. Ensures all text embeddings and attention masks have consistent batch sizes (batch_size * num_images_per_prompt)" + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_images_per_prompt"), + InputParam.template("prompt_embeds"), + InputParam.template("negative_prompt_embeds"), + InputParam( + "prompt_attention_mask", + type_hint=torch.Tensor, + description="Attention mask for the prompt embeddings. Can be generated from text_encoder step.", + ), + InputParam( + "negative_prompt_attention_mask", + type_hint=torch.Tensor, + description="Attention mask for the negative prompt embeddings. Can be generated from text_encoder step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + 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.template("prompt_embeds"), + OutputParam.template("negative_prompt_embeds"), + OutputParam( + "prompt_attention_mask", + type_hint=torch.Tensor, + description="Attention mask for the prompt embeddings, expanded to the final batch size.", + ), + OutputParam( + "negative_prompt_attention_mask", + type_hint=torch.Tensor, + description="Attention mask for the negative prompt embeddings, expanded to the final batch size.", + ), + ] + + def check_inputs(self, components, block_state): + if block_state.prompt_embeds is not None and block_state.prompt_attention_mask is None: + raise ValueError("Cannot provide `prompt_embeds` without also providing `prompt_attention_mask`") + + if block_state.negative_prompt_embeds is not None and block_state.negative_prompt_attention_mask is None: + raise ValueError( + "Cannot provide `negative_prompt_embeds` without also providing `negative_prompt_attention_mask`" + ) + + if block_state.negative_prompt_embeds is not None: + if block_state.prompt_embeds.shape[0] != block_state.negative_prompt_embeds.shape[0]: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same batch size, but got:" + f" `prompt_embeds` {block_state.prompt_embeds.shape} != `negative_prompt_embeds`" + f" {block_state.negative_prompt_embeds.shape}." + ) + + @staticmethod + def expand_text_tensor(text_tensor: torch.Tensor, num_images_per_prompt: int) -> torch.Tensor: + # duplicate text embeddings/attention masks for each generation per prompt, using mps friendly method + batch_size, seq_len = text_tensor.shape[:2] + if text_tensor.ndim == 2: + text_tensor = text_tensor.repeat(1, num_images_per_prompt) + return text_tensor.view(batch_size * num_images_per_prompt, seq_len) + text_tensor = text_tensor.repeat(1, num_images_per_prompt, 1) + return text_tensor.view(batch_size * num_images_per_prompt, seq_len, -1) + + @torch.no_grad() + def __call__(self, components: ChromaModularPipeline, 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 + + for field_name in [ + "prompt_embeds", + "negative_prompt_embeds", + "prompt_attention_mask", + "negative_prompt_attention_mask", + ]: + text_tensor = getattr(block_state, field_name) + if text_tensor is None: + continue + setattr(block_state, field_name, self.expand_text_tensor(text_tensor, block_state.num_images_per_prompt)) + + self.set_block_state(state, block_state) + + return components, state diff --git a/src/diffusers/modular_pipelines/chroma/modular_blocks_chroma.py b/src/diffusers/modular_pipelines/chroma/modular_blocks_chroma.py new file mode 100644 index 000000000000..87e395c71d31 --- /dev/null +++ b/src/diffusers/modular_pipelines/chroma/modular_blocks_chroma.py @@ -0,0 +1,180 @@ +# 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 SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import ( + ChromaPrepareAttentionMaskStep, + ChromaPrepareLatentsStep, + ChromaRoPEInputsStep, + ChromaSetTimestepsStep, +) +from .decoders import ChromaDecodeStep +from .denoise import ChromaDenoiseStep +from .encoders import ChromaTextEncoderStep +from .inputs import ChromaTextInputStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# auto_docstring +class ChromaCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core step that performs the denoising process for Chroma. + This step takes the encoded conditions (prompt embeddings and attention masks) and runs the text-to-image + denoising process. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer + (`ChromaTransformer2DModel`) + + Inputs: + num_images_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. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + prompt_attention_mask (`Tensor`, *optional*): + Attention mask for the prompt embeddings. Can be generated from text_encoder step. + negative_prompt_attention_mask (`Tensor`, *optional*): + Attention mask for the negative prompt embeddings. Can be generated from text_encoder step. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 35): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + joint_attention_kwargs (`dict`, *optional*): + Additional kwargs passed along to the attention processors. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "chroma" + block_classes = [ + ChromaTextInputStep, + ChromaPrepareLatentsStep, + ChromaSetTimestepsStep, + ChromaPrepareAttentionMaskStep, + ChromaRoPEInputsStep, + ChromaDenoiseStep, + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "prepare_attention_mask", + "prepare_rope_inputs", + "denoise", + ] + + @property + def description(self): + return ( + "Core step that performs the denoising process for Chroma.\n" + + "This step takes the encoded conditions (prompt embeddings and attention masks) and runs the " + + "text-to-image denoising process." + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + ] + + +TEXT2IMAGE_BLOCKS = InsertableDict( + [ + ("text_encoder", ChromaTextEncoderStep()), + ("denoise", ChromaCoreDenoiseStep()), + ("decode", ChromaDecodeStep()), + ] +) + + +# auto_docstring +class ChromaAutoBlocks(SequentialPipelineBlocks): + """ + Auto Modular pipeline for text-to-image using Chroma. + + Supported workflows: + - `text2image`: requires `prompt` + + Components: + text_encoder (`T5EncoderModel`) tokenizer (`T5Tokenizer`) guider (`ClassifierFreeGuidance`) scheduler + (`FlowMatchEulerDiscreteScheduler`) transformer (`ChromaTransformer2DModel`) vae (`AutoencoderKL`) + image_processor (`VaeImageProcessor`) + + 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. + joint_attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors; `scale` is used as the text encoder LoRA scale. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 35): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "chroma" + + block_classes = TEXT2IMAGE_BLOCKS.values() + block_names = TEXT2IMAGE_BLOCKS.keys() + + _workflow_map = { + "text2image": {"prompt": True}, + } + + @property + def description(self): + return "Auto Modular pipeline for text-to-image using Chroma." + + @property + def outputs(self): + return [OutputParam.template("images")] diff --git a/src/diffusers/modular_pipelines/chroma/modular_pipeline.py b/src/diffusers/modular_pipelines/chroma/modular_pipeline.py new file mode 100644 index 000000000000..f0c2079668c4 --- /dev/null +++ b/src/diffusers/modular_pipelines/chroma/modular_pipeline.py @@ -0,0 +1,66 @@ +# 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 FluxLoraLoaderMixin, TextualInversionLoaderMixin +from ...utils import logging +from ..modular_pipeline import ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class ChromaModularPipeline(ModularPipeline, FluxLoraLoaderMixin, TextualInversionLoaderMixin): + """ + A ModularPipeline for Chroma. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "ChromaAutoBlocks" + + @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 = 16 + if getattr(self, "transformer", None) is not None: + num_channels_latents = self.transformer.config.in_channels // 4 + return num_channels_latents + + @property + def requires_unconditional_embeds(self): + requires_unconditional_embeds = True + + 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/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 9bf1ddca3b98..378a5fcf5031 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -129,6 +129,7 @@ def _helios_pyramid_map_fn(config_dict=None): MODULAR_PIPELINE_MAPPING = OrderedDict( [ + ("chroma", _create_default_map_fn("ChromaModularPipeline")), ("stable-diffusion-xl", _create_default_map_fn("StableDiffusionXLModularPipeline")), ("stable-diffusion-3", _create_default_map_fn("StableDiffusion3ModularPipeline")), ("wan", _wan_map_fn), diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 8058937aa2a8..611f0c36848a 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -32,6 +32,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class ChromaAutoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class ChromaModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class Cosmos3DistilledBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/modular_pipelines/chroma/__init__.py b/tests/modular_pipelines/chroma/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/modular_pipelines/chroma/test_modular_pipeline_chroma.py b/tests/modular_pipelines/chroma/test_modular_pipeline_chroma.py new file mode 100644 index 000000000000..fc55ef1d2f34 --- /dev/null +++ b/tests/modular_pipelines/chroma/test_modular_pipeline_chroma.py @@ -0,0 +1,59 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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 diffusers.modular_pipelines import ChromaAutoBlocks, ChromaModularPipeline + +from ..test_modular_pipelines_common import ModularPipelineTesterMixin + + +CHROMA_WORKFLOWS = { + "text2image": [ + ("text_encoder", "ChromaTextEncoderStep"), + ("denoise.input", "ChromaTextInputStep"), + ("denoise.prepare_latents", "ChromaPrepareLatentsStep"), + ("denoise.set_timesteps", "ChromaSetTimestepsStep"), + ("denoise.prepare_attention_mask", "ChromaPrepareAttentionMaskStep"), + ("denoise.prepare_rope_inputs", "ChromaRoPEInputsStep"), + ("denoise.denoise", "ChromaDenoiseStep"), + ("decode", "ChromaDecodeStep"), + ], +} + + +class TestChromaModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = ChromaModularPipeline + pipeline_blocks_class = ChromaAutoBlocks + pretrained_model_name_or_path = "charchits7/tiny-chroma-modular-pipe" + + params = frozenset(["prompt", "negative_prompt", "height", "width"]) + batch_params = frozenset(["prompt", "negative_prompt"]) + expected_workflow_blocks = CHROMA_WORKFLOWS + + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 2, + "height": 32, + "width": 32, + "max_sequence_length": 16, + "output_type": "pt", + } + return inputs + + def test_inference_batch_single_identical(self): + super().test_inference_batch_single_identical(expected_max_diff=5e-3)