| 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" |
| |
| _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 |
|
|
