moss_image_custom_blocks / before_denoise.py

OzzyGT HF Staff

added moss image custom blocks

2107900 3 days ago

8.5 kB

# 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

	# 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 `shift2 / (num_train_timesteps - 1 + shift2)` 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