File size: 7,576 Bytes
2107900 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | # 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 torch
from diffusers import AutoModel
from diffusers.configuration_utils import FrozenDict
from diffusers.guiders import ClassifierFreeGuidance
from diffusers.modular_pipelines import (
BlockState,
LoopSequentialPipelineBlocks,
ModularPipelineBlocks,
PipelineState,
)
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
class MOSSImageLoopBeforeDenoiser(ModularPipelineBlocks):
model_name = "moss-image"
@property
def description(self) -> str:
return "Step within the denoising loop that casts the float32 latents to the transformer dtype."
@property
def inputs(self) -> list[InputParam]:
return [
InputParam(
"latents",
required=True,
type_hint=torch.Tensor,
description="The latents being denoised. Can be generated in the prepare latents step.",
),
InputParam(
"dtype",
required=True,
type_hint=torch.dtype,
description="Dtype of the transformer inputs. Can be generated in the input step.",
),
]
@torch.no_grad()
def __call__(self, components, 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 MOSSImageLoopDenoiser(ModularPipelineBlocks):
model_name = "moss-image"
# transformer argument -> (conditional, unconditional) block_state fields
guider_input_fields = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")}
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"guider",
ClassifierFreeGuidance,
config=FrozenDict({"guidance_scale": 3.0}),
default_creation_method="from_config",
),
ComponentSpec("transformer", AutoModel),
]
@property
def description(self) -> str:
return (
"Step within the denoising loop that runs the transformer for each guidance batch and combines the "
"predictions with the guider. The unconditional prediction skips the transformer blocks in "
"`uncond_skip_layers`."
)
@property
def inputs(self) -> list[InputParam]:
return [
InputParam(
"num_inference_steps",
required=True,
type_hint=int,
description="The number of denoising steps. Can be generated in the set timesteps step.",
),
InputParam(
"uncond_skip_layers",
type_hint=list[int],
default=[22],
description="Zero-based transformer blocks skipped in the unconditional prediction. Pass `[]` to "
"run the full model for it.",
),
InputParam.template("attention_kwargs"),
InputParam("prompt_embeds", required=True, type_hint=list[torch.Tensor]),
InputParam("negative_prompt_embeds", type_hint=list[torch.Tensor]),
]
@torch.no_grad()
def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor):
guider = components.guider
guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
guider_state = guider.prepare_inputs_from_block_state(block_state, self.guider_input_fields)
for guider_state_batch in guider_state:
guider.prepare_models(components.transformer)
is_uncond = getattr(guider_state_batch, guider._identifier_key) == "pred_uncond"
skip_layers = block_state.uncond_skip_layers if is_uncond and block_state.uncond_skip_layers else None
model_out = components.transformer(
hidden_states=block_state.latent_model_input,
timestep=t,
encoder_hidden_states=[e.to(block_state.dtype) for e in guider_state_batch.encoder_hidden_states],
skip_layers=skip_layers,
attention_kwargs=block_state.attention_kwargs,
return_dict=False,
)[0]
# guidance and the Euler update run in float32
guider_state_batch.noise_pred = model_out.float()
guider.cleanup_models(components.transformer)
block_state.noise_pred = guider(guider_state)[0]
return components, block_state
class MOSSImageLoopAfterDenoiser(ModularPipelineBlocks):
model_name = "moss-image"
@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 float32 latents with the scheduler."
@torch.no_grad()
def __call__(self, components, 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
class MOSSImageDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
model_name = "moss-image"
@property
def description(self) -> str:
return (
"Pipeline block that iteratively denoises the latents over `timesteps`. "
"The steps of each iteration are defined by `sub_blocks`."
)
@property
def loop_expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
@property
def loop_inputs(self) -> list[InputParam]:
return [
InputParam(
"timesteps",
required=True,
type_hint=torch.Tensor,
description="The timesteps of the denoising loop. Can be generated in the set timesteps step.",
),
InputParam(
"num_inference_steps",
required=True,
type_hint=int,
description="The number of denoising steps. Can be generated in the set timesteps step.",
),
]
@torch.no_grad()
def __call__(self, components, 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
class MOSSImageDenoiseStep(MOSSImageDenoiseLoopWrapper):
block_classes = [MOSSImageLoopBeforeDenoiser, MOSSImageLoopDenoiser, MOSSImageLoopAfterDenoiser]
block_names = ["before_denoiser", "denoiser", "after_denoiser"]
@property
def description(self) -> str:
return (
"Denoise step that iteratively denoises the latents. At each iteration it runs, in order: "
"`MOSSImageLoopBeforeDenoiser`, `MOSSImageLoopDenoiser`, `MOSSImageLoopAfterDenoiser`."
)
|