File size: 7,774 Bytes
dc9d63c | 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 195 196 | # 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.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 MingImageLoopBeforeDenoiser(ModularPipelineBlocks):
model_name = "ming-image"
@property
def description(self) -> str:
return (
"Step within the denoising loop that prepares the transformer inputs: one `(channels, 1, height, width)` "
"latent per image and the normalized time `(1000 - t) / 1000`."
)
@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):
latents = block_state.latents.to(block_state.dtype).unsqueeze(2)
block_state.latent_model_input = list(latents.unbind(dim=0))
block_state.timestep = (1000 - t.expand(latents.shape[0])) / 1000
return components, block_state
class MingImageLoopDenoiser(ModularPipelineBlocks):
model_name = "ming-image"
# transformer argument -> (conditional, unconditional) block_state fields
guider_input_fields = {
"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
"encoder_hidden_states_2": ("prompt_embeds_2", "negative_prompt_embeds_2"),
}
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec(
"guider",
ClassifierFreeGuidance,
config=FrozenDict({"guidance_scale": 1.0, "use_original_formulation": True, "enabled": False}),
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 original pipeline's `cfg` matches `ClassifierFreeGuidance` with "
"`use_original_formulation=True` and `guidance_scale=cfg`."
)
@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.template("denoiser_input_fields"),
InputParam("prompt_embeds", required=True, type_hint=list[torch.Tensor]),
InputParam("prompt_embeds_2", required=True, type_hint=list[torch.Tensor]),
InputParam("negative_prompt_embeds", type_hint=list[torch.Tensor]),
InputParam("negative_prompt_embeds_2", type_hint=list[torch.Tensor]),
]
@torch.no_grad()
def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor):
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)
for guider_state_batch in guider_state:
components.guider.prepare_models(components.transformer)
cond_kwargs = {
name: [e.to(block_state.dtype) for e in getattr(guider_state_batch, name)]
for name in self.guider_input_fields
}
model_out = components.transformer(
hidden_states=block_state.latent_model_input,
timestep=block_state.timestep,
return_dict=False,
**cond_kwargs,
)[0]
# the transformer predicts the negated flow-matching velocity
guider_state_batch.noise_pred = -torch.stack([o.float() for o in model_out], dim=0).squeeze(2)
components.guider.cleanup_models(components.transformer)
block_state.noise_pred = components.guider(guider_state)[0]
return components, block_state
class MingImageLoopAfterDenoiser(ModularPipelineBlocks):
model_name = "ming-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.float(), t, block_state.latents, return_dict=False
)[0]
return components, block_state
class MingImageDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
model_name = "ming-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 MingImageDenoiseStep(MingImageDenoiseLoopWrapper):
block_classes = [MingImageLoopBeforeDenoiser, MingImageLoopDenoiser, MingImageLoopAfterDenoiser]
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: "
"`MingImageLoopBeforeDenoiser`, `MingImageLoopDenoiser`, `MingImageLoopAfterDenoiser`."
)
|