# Copyright 2026 OpenMOSS and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0. See the LICENSE file in this repository. import numpy as np 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 latent patch size (2) SPATIAL_COMPRESSION = 16 # 32 VAE channels packed in 2x2 patches LATENT_CHANNELS = 128 class MOSSImageTextInputStep(ModularPipelineBlocks): model_name = "moss-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="Text-encoder states, one tensor per prompt. Can be generated in the text_encoder step.", ), InputParam( "negative_prompt_embeds", type_hint=list[torch.Tensor], description="Negative text-encoder states, one tensor per prompt. 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="Text-encoder states repeated per image.", ), OutputParam( "negative_prompt_embeds", type_hint=list[torch.Tensor], kwargs_type="denoiser_input_fields", description="Negative text-encoder states repeated per image.", ), ] @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) negative_prompt_embeds = block_state.negative_prompt_embeds if negative_prompt_embeds is not None and len(negative_prompt_embeds) != len(block_state.prompt_embeds): raise ValueError( f"`negative_prompt_embeds` ({len(negative_prompt_embeds)}) and `prompt_embeds` " f"({len(block_state.prompt_embeds)}) 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.negative_prompt_embeds = repeat(negative_prompt_embeds) self.set_block_state(state, block_state) return components, state class MOSSImagePrepareLatentsStep(ModularPipelineBlocks): model_name = "moss-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=1024), InputParam.template("width", default=1024), InputParam( "latents", type_hint=torch.Tensor, description="Initial noise of shape `(batch_size * num_images_per_prompt, 128, height // 16, " "width // 16)`.", ), 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, 128, height // 16, width // 16)`.", ) ] @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 // SPATIAL_COMPRESSION, block_state.width // SPATIAL_COMPRESSION, ) # Sample in float32: noise cast from bfloat16 would have lost its precision 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 MOSSImageSetTimestepsStep(ModularPipelineBlocks): model_name = "moss-image" @property def description(self) -> str: return ( "Step that sets the scheduler timesteps. By default the sigmas run linearly from 1 down to " "`shift / (num_train_timesteps + shift - 1)` and the scheduler applies its static `shift`, so the last " "step lands at `shift**2 / (num_train_timesteps - 1 + shift**2)` before the final step 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=28), InputParam.template("sigmas"), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps of the denoising loop."), OutputParam("num_inference_steps", type_hint=int, description="The number of denoising steps."), ] @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) scheduler = components.scheduler sigmas = block_state.sigmas if sigmas is None: shift = scheduler.config.shift lowest = shift / (scheduler.config.num_train_timesteps + shift - 1) sigmas = np.linspace(1.0, lowest, block_state.num_inference_steps) scheduler.set_timesteps(sigmas=sigmas, device=components._execution_device) scheduler.set_begin_index(0) block_state.timesteps = scheduler.timesteps block_state.num_inference_steps = len(block_state.timesteps) self.set_block_state(state, block_state) return components, state