# Copyright 2026 inclusionAI and The HuggingFace Team. All rights reserved. # # Licensed under the MIT License. See the LICENSE file in this repository. import torch from diffusers import AutoModel from diffusers.modular_pipelines import ModularPipelineBlocks, PipelineState from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from diffusers.utils.torch_utils import randn_tensor # VAE spatial downsampling (8) times the transformer patch size (2) SPATIAL_COMPRESSION = 16 LATENT_CHANNELS = 16 class MingImageTextInputStep(ModularPipelineBlocks): model_name = "ming-image" @property def description(self) -> str: return ( "Input step that determines `batch_size` and `dtype`, and repeats each prompt's conditioning " "`num_images_per_prompt` times." ) @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", AutoModel)] @property def inputs(self) -> list[InputParam]: return [ InputParam.template("num_images_per_prompt"), InputParam( "prompt_embeds", required=True, type_hint=list[torch.Tensor], description="Query-token conditioning, one tensor per prompt. Can be generated in the text_encoder step.", ), InputParam( "prompt_embeds_2", required=True, type_hint=list[torch.Tensor], description="Text-token conditioning, one tensor per prompt. Can be generated in the text_encoder step.", ), InputParam( "negative_prompt_embeds", type_hint=list[torch.Tensor], description="Unconditional query-token conditioning. Can be generated in the text_encoder step.", ), InputParam( "negative_prompt_embeds_2", type_hint=list[torch.Tensor], description="Unconditional text-token conditioning. Can be generated in the text_encoder step.", ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "batch_size", type_hint=int, description="Number of prompts; the model batch is `batch_size * num_images_per_prompt`.", ), OutputParam("dtype", type_hint=torch.dtype, description="Dtype of the transformer inputs."), OutputParam( "prompt_embeds", type_hint=list[torch.Tensor], kwargs_type="denoiser_input_fields", description="Query-token conditioning repeated per image.", ), OutputParam( "prompt_embeds_2", type_hint=list[torch.Tensor], kwargs_type="denoiser_input_fields", description="Text-token conditioning repeated per image.", ), OutputParam( "negative_prompt_embeds", type_hint=list[torch.Tensor], kwargs_type="denoiser_input_fields", description="Unconditional query-token conditioning repeated per image.", ), OutputParam( "negative_prompt_embeds_2", type_hint=list[torch.Tensor], kwargs_type="denoiser_input_fields", description="Unconditional text-token conditioning repeated per image.", ), ] @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) if len(block_state.prompt_embeds) != len(block_state.prompt_embeds_2): raise ValueError( f"`prompt_embeds` ({len(block_state.prompt_embeds)}) and `prompt_embeds_2` " f"({len(block_state.prompt_embeds_2)}) must hold one tensor per prompt" ) block_state.batch_size = len(block_state.prompt_embeds) block_state.dtype = components.transformer.dtype def repeat(embeds): if embeds is None: return None return [e for e in embeds for _ in range(block_state.num_images_per_prompt)] block_state.prompt_embeds = repeat(block_state.prompt_embeds) block_state.prompt_embeds_2 = repeat(block_state.prompt_embeds_2) block_state.negative_prompt_embeds = repeat(block_state.negative_prompt_embeds) block_state.negative_prompt_embeds_2 = repeat(block_state.negative_prompt_embeds_2) self.set_block_state(state, block_state) return components, state class MingImagePrepareLatentsStep(ModularPipelineBlocks): model_name = "ming-image" @property def description(self) -> str: return "Prepare latents step that creates the initial float32 noise for text-to-image generation." @property def inputs(self) -> list[InputParam]: return [ InputParam.template("height", default=2048), InputParam.template("width", default=2048), InputParam( "latents", type_hint=torch.Tensor, description="Initial noise of shape `(batch_size * num_images_per_prompt, 16, height // 8, width // 8)`.", ), InputParam.template("num_images_per_prompt"), InputParam.template("generator"), InputParam( "batch_size", required=True, type_hint=int, description="Number of prompts. Can be generated in the input step.", ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "latents", type_hint=torch.Tensor, description="The initial noisy latents, float32, of shape `(batch, 16, height // 8, width // 8)`.", ) ] @staticmethod def check_inputs(height, width): if height % SPATIAL_COMPRESSION != 0 or width % SPATIAL_COMPRESSION != 0: raise ValueError( f"`height` and `width` have to be divisible by {SPATIAL_COMPRESSION} but are {height} and {width}." ) @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) self.check_inputs(block_state.height, block_state.width) shape = ( block_state.batch_size * block_state.num_images_per_prompt, LATENT_CHANNELS, block_state.height // 8, block_state.width // 8, ) if block_state.latents is None: block_state.latents = randn_tensor( shape, generator=block_state.generator, device=components._execution_device, dtype=torch.float32 ) elif block_state.latents.shape != shape: raise ValueError(f"Unexpected latents shape, got {tuple(block_state.latents.shape)}, expected {shape}") else: block_state.latents = block_state.latents.to(components._execution_device, torch.float32) self.set_block_state(state, block_state) return components, state class MingImageSetTimestepsStep(ModularPipelineBlocks): model_name = "ming-image" @property def description(self) -> str: return ( "Step that sets the scheduler timesteps: `num_inference_steps` sigmas from 1 down to 0, shifted by the " "scheduler's static `shift`. The schedule does not depend on the resolution. Sets the scheduler's " "`sigma_min` to 0." ) @property def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] @property def inputs(self) -> list[InputParam]: return [ InputParam.template("num_inference_steps", default=12), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps of the denoising loop."), ] @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) # The original pipeline runs the schedule down to sigma 0 by setting `sigma_min` on the scheduler and letting # it build the sigmas; doing the same reproduces its float64 schedule exactly (custom `sigmas` would be cast to # float32 before shifting). components.scheduler.sigma_min = 0.0 components.scheduler.set_timesteps(block_state.num_inference_steps, device=components._execution_device) block_state.timesteps = components.scheduler.timesteps self.set_block_state(state, block_state) return components, state