File size: 9,876 Bytes
92011e4 | 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 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 | from dataclasses import dataclass, field
import torch
from diffusers import ModelMixin, SchedulerMixin
from diffusers.modular_pipelines import LoopSequentialPipelineBlocks
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
from diffusers.utils.torch_utils import randn_tensor
from .guidance import Pixal3DGuidance
from .resolvable import ResolvableBlock
@dataclass(frozen=True)
class Pixal3DStage:
"""Components and guidance settings of one flow stage."""
transformer: str
scheduler: str
guider: str
guider_config: dict = field(default_factory=dict)
sparse: bool = True
concat_cond: bool = False
STRUCTURE_STAGE = Pixal3DStage(
"sparse_structure_flow_model",
"structure_scheduler",
"structure_guider",
{"guidance_scale": 7.5, "guidance_rescale": 0.7, "guidance_interval": (0.6, 1.0)},
sparse=False,
)
SHAPE_GUIDANCE = {"guidance_scale": 7.5, "guidance_rescale": 0.5, "guidance_interval": (0.6, 1.0)}
SHAPE_LOW_STAGE = Pixal3DStage("shape_slat_flow_model_512", "shape_scheduler", "shape_guider", SHAPE_GUIDANCE)
SHAPE_HIGH_STAGE = Pixal3DStage("shape_slat_flow_model_1024", "shape_scheduler", "shape_guider", SHAPE_GUIDANCE)
TEXTURE_STAGE = Pixal3DStage(
"tex_slat_flow_model_1024",
"texture_scheduler",
"texture_guider",
{"guidance_scale": 1.0, "guidance_rescale": 0.0, "guidance_interval": (0.6, 0.9)},
concat_cond=True,
)
class Pixal3DPrepareLatentsStep(ResolvableBlock):
model_name = "pixal3d"
def __init__(self, stage: Pixal3DStage):
self.stage = stage
super().__init__()
@property
def description(self):
return "Draw the initial noise for a dense structure volume or for the sparse shape/texture tokens."
@property
def expected_components(self):
return self.resolve([ComponentSpec(self.stage.transformer, ModelMixin)])
@property
def inputs(self):
inputs = [InputParam.template("generator")]
if self.stage.sparse:
inputs.append(
InputParam("coords", required=True, type_hint=torch.Tensor, description="Sparse (N, 4) coordinates.")
)
if self.stage.concat_cond:
inputs.append(
InputParam("concat_cond", required=True, type_hint=torch.Tensor, description="Shape latents.")
)
return inputs
@property
def intermediate_outputs(self):
return [OutputParam("latents", type_hint=torch.Tensor, description="Initial noise for this stage.")]
@torch.no_grad()
def __call__(self, components, state):
b = self.get_block_state(state)
config = getattr(components, self.stage.transformer).config
if not self.stage.sparse:
shape = (1, config.in_channels, *([config.resolution] * 3))
else:
if len(b.coords) == 0:
raise ValueError("Nonempty sparse coordinates are required")
channels = config.in_channels - (b.concat_cond.shape[-1] if self.stage.concat_cond else 0)
shape = (1, len(b.coords), channels)
b.latents = randn_tensor(
shape, generator=b.generator, device=components._execution_device, dtype=torch.float32
)
self.set_block_state(state, b)
return components, state
class Pixal3DSetTimestepsStep(ResolvableBlock):
model_name = "pixal3d"
def __init__(self, stage: Pixal3DStage):
self.stage = stage
super().__init__()
@property
def description(self):
return "Set the shifted flow-matching sigmas for this stage."
@property
def expected_components(self):
return self.resolve([ComponentSpec(self.stage.scheduler, SchedulerMixin)])
@property
def inputs(self):
return [InputParam.template("num_inference_steps", default=12)]
@property
def intermediate_outputs(self):
return [OutputParam("timesteps", type_hint=list, description="Sigmas in [0, 1], one per step.")]
def __call__(self, components, state):
b = self.get_block_state(state)
scheduler = getattr(components, self.stage.scheduler)
scheduler.set_timesteps(b.num_inference_steps)
b.timesteps = scheduler.timesteps
self.set_block_state(state, b)
return components, state
class Pixal3DLoopDenoiser(ResolvableBlock):
model_name = "pixal3d"
# Conditional batches use the image features; the unconditional batch uses zeros of the same shape.
_guider_input_fields = {
"global_features": ("global_features", "negative_global_features"),
"projected_features": ("projected_features", "negative_projected_features"),
"latents": "latents",
}
def __init__(self, stage: Pixal3DStage):
self.stage = stage
super().__init__()
@property
def description(self):
return "Step within the denoising loop: run the flow model for each guider batch and combine them."
@property
def expected_components(self):
return self.resolve(
[
ComponentSpec(self.stage.transformer, ModelMixin),
ComponentSpec(
self.stage.guider,
Pixal3DGuidance,
default_creation_method="from_config",
config=self.stage.guider_config,
),
]
)
@property
def inputs(self):
return [
InputParam("global_features", required=True, type_hint=torch.Tensor, description="DINOv3 global tokens."),
InputParam(
"projected_features", required=True, type_hint=torch.Tensor, description="Camera-projected features."
),
]
@torch.no_grad()
def __call__(self, components, block_state, i, t):
model = getattr(components, self.stage.transformer)
guider = getattr(components, self.stage.guider)
guider.set_state(step=i, num_inference_steps=len(block_state.timesteps), timestep=t)
batches = guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)
latents = block_state.latents
timestep = torch.full((latents.shape[0],), 1000 * t, device=latents.device, dtype=torch.float32)
for batch in batches:
guider.prepare_models(model)
batch.noise_pred = model(
latents,
timestep,
batch.global_features,
batch.projected_features,
coords=block_state.model_coords,
concat_cond=block_state.model_concat_cond,
)
guider.cleanup_models(model)
block_state.noise_pred = guider(batches).pred
return components, block_state
class Pixal3DLoopAfterDenoiser(ResolvableBlock):
model_name = "pixal3d"
def __init__(self, stage: Pixal3DStage):
self.stage = stage
super().__init__()
@property
def description(self):
return "Step within the denoising loop: take one Euler step with the guided velocity."
@property
def expected_components(self):
return self.resolve([ComponentSpec(self.stage.scheduler, SchedulerMixin)])
@torch.no_grad()
def __call__(self, components, block_state, i, t):
scheduler = getattr(components, self.stage.scheduler)
prediction = block_state.noise_pred
block_state.latents = scheduler.step(prediction, t, block_state.latents.to(prediction.device)).prev_sample
return components, block_state
class Pixal3DDenoiseLoop(LoopSequentialPipelineBlocks):
model_name = "pixal3d"
def __init__(self, stage: Pixal3DStage):
self.stage = stage
self.block_classes = [Pixal3DLoopDenoiser(stage), Pixal3DLoopAfterDenoiser(stage)]
self.block_names = ["denoiser", "after_denoiser"]
super().__init__()
@property
def description(self):
return (
"Iteratively denoise one stage's latents over `timesteps`. Each iteration runs:\n"
" - `Pixal3DLoopDenoiser`: flow model per guider batch, then guidance\n"
" - `Pixal3DLoopAfterDenoiser`: Euler step"
)
@property
def loop_inputs(self):
inputs = [
InputParam("timesteps", required=True, type_hint=list, description="From the set-timesteps step."),
InputParam("latents", required=True, type_hint=torch.Tensor, description="From the prepare-latents step."),
]
if self.stage.sparse:
inputs.append(InputParam("coords", required=True, type_hint=torch.Tensor))
if self.stage.concat_cond:
inputs.append(InputParam("concat_cond", required=True, type_hint=torch.Tensor))
return inputs
@property
def loop_intermediate_outputs(self):
return [OutputParam("latents", type_hint=torch.Tensor, description="Denoised latents of this stage.")]
@torch.no_grad()
def __call__(self, components, state):
b = self.get_block_state(state)
device = components._execution_device
b.latents = b.latents.to(device)
b.global_features = b.global_features.to(device)
b.projected_features = b.projected_features.to(device)
b.negative_global_features = torch.zeros_like(b.global_features)
b.negative_projected_features = torch.zeros_like(b.projected_features)
b.model_coords = b.coords[None, :, 1:].to(device) if self.stage.sparse else None
b.model_concat_cond = b.concat_cond.to(device) if self.stage.concat_cond else None
with self.progress_bar(total=len(b.timesteps)) as progress_bar:
for i, t in enumerate(b.timesteps):
components, b = self.loop_step(components, b, i=i, t=t)
progress_bar.update()
self.set_block_state(state, b)
return components, state
|