691 lines
26 KiB
Python
691 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
"""Two-pass spatial refine for the MLX Wan runtime (H3 / LTX-2 pattern).
|
||
|
||
Biggest quality lever on Apple Silicon without a new model or training:
|
||
generate at base resolution, then run a second denoising pass with the
|
||
*same* DiT at a higher resolution.
|
||
|
||
This is the MLX-side port of the CUDA refine template in
|
||
``fastvideo/pipelines/basic/ltx2/stages/ltx2_refine.py`` and the H3
|
||
"base + regenerate" pattern documented in
|
||
``docs/design/mac_qad_two_product_strategy.md``:
|
||
|
||
1. :func:`plan_refine_resolutions` — split the request into stage-1
|
||
(base) and stage-2 (target) pixel sizes, validating VAE / patch
|
||
alignment the way :class:`LTX2RefineInitStage` does.
|
||
2. :func:`upsample_latents_spatial` — 2× (or N×) spatial upsample of
|
||
clean latents. Wan has no learned latent upsampler on Mac, so this
|
||
is bilinear over the H×W plane (temporal axis untouched) — same
|
||
role as LTX-2's ``upsample_video`` hand-off, without the learned
|
||
residual.
|
||
3. :func:`prepare_refine_latents` — upsample + re-noise the clean
|
||
stage-1 latents to the stage-2 sigma so the second denoise has
|
||
something to refine (mirrors :class:`LTX2UpsampleStage` +
|
||
``apply_ltx2_gaussian_noiser``).
|
||
4. :func:`run_two_pass_dmd` — orchestrate stage-1 denoise → refine
|
||
hand-off → stage-2 denoise with the same model / prompt embeds.
|
||
|
||
No LoRA swap, no dedicated SR weights, no new training — pure pipeline
|
||
work reusable by Wan2.1-14B and Wan2.2-5B on Apple Silicon.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from typing import TYPE_CHECKING, Any
|
||
from collections.abc import Callable, Sequence
|
||
|
||
import numpy as np
|
||
|
||
from fastvideo.logger import init_logger
|
||
from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, add_noise, dmd_step
|
||
|
||
if TYPE_CHECKING: # pragma: no cover - typing only
|
||
import mlx.core as mx
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
# Default stage-2 noise level when the caller does not supply a schedule.
|
||
# Matches the first entry of LTX-2's STAGE_2_DISTILLED_SIGMA_VALUES in spirit
|
||
# (start the refine denoise from a high-noise level) without hard-wiring the
|
||
# LTX-2 distilled grid onto Wan's flow-match schedule.
|
||
DEFAULT_REFINE_SIGMA = 0.909375
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class RefinePlan:
|
||
"""Resolved stage-1 / stage-2 geometry for a two-pass refine run."""
|
||
|
||
target_height: int
|
||
target_width: int
|
||
stage1_height: int
|
||
stage1_width: int
|
||
spatial_scale: int
|
||
vae_spatial_compression: int
|
||
vae_temporal_compression: int
|
||
num_frames: int
|
||
|
||
@property
|
||
def stage1_latent_height(self) -> int:
|
||
"""Return the stage-1 latent height after VAE spatial compression."""
|
||
return self.stage1_height // self.vae_spatial_compression
|
||
|
||
@property
|
||
def stage1_latent_width(self) -> int:
|
||
"""Return the stage-one latent width after VAE spatial compression."""
|
||
return self.stage1_width // self.vae_spatial_compression
|
||
|
||
@property
|
||
def stage2_latent_height(self) -> int:
|
||
"""Calculate the target-resolution latent height.
|
||
|
||
Returns:
|
||
int: The target height divided by the VAE spatial compression factor.
|
||
"""
|
||
return self.target_height // self.vae_spatial_compression
|
||
|
||
@property
|
||
def stage2_latent_width(self) -> int:
|
||
"""Return the target image width in latent-space units."""
|
||
return self.target_width // self.vae_spatial_compression
|
||
|
||
@property
|
||
def latent_frames(self) -> int:
|
||
"""Calculate the number of latent frames after VAE temporal compression.
|
||
|
||
Returns:
|
||
int: The compressed latent frame count.
|
||
"""
|
||
return (self.num_frames - 1) // self.vae_temporal_compression + 1
|
||
|
||
|
||
def plan_refine_resolutions(
|
||
*,
|
||
height: int,
|
||
width: int,
|
||
num_frames: int,
|
||
spatial_scale: int = 2,
|
||
vae_spatial_compression: int = 8,
|
||
vae_temporal_compression: int = 4,
|
||
patch_size: tuple[int, int, int] = (1, 2, 2),
|
||
enabled: bool = True,
|
||
mode_label: str = "Refine",
|
||
) -> RefinePlan:
|
||
"""
|
||
Validate the requested dimensions and create the stage-1 and target-resolution refinement plan.
|
||
|
||
Parameters:
|
||
height (int): Target image height in pixels.
|
||
width (int): Target image width in pixels.
|
||
num_frames (int): Number of frames in the input sequence.
|
||
spatial_scale (int): Factor used to reduce spatial dimensions for stage 1.
|
||
vae_spatial_compression (int): Spatial compression factor of the VAE.
|
||
vae_temporal_compression (int): Temporal compression factor of the VAE.
|
||
patch_size (tuple[int, int, int]): Temporal and spatial patch dimensions used to validate latent-grid alignment.
|
||
enabled (bool): Whether to use two-pass refinement.
|
||
mode_label (str): Name of the calling mode, used to prefix validation
|
||
errors so ``--fast-spatial`` failures do not read as refine failures.
|
||
|
||
Returns:
|
||
RefinePlan: The validated stage-1 and target-resolution plan.
|
||
"""
|
||
if height <= 0 or width <= 0:
|
||
raise ValueError(f"height/width must be positive, got {height}x{width}")
|
||
if spatial_scale < 1:
|
||
raise ValueError(f"spatial_scale must be >= 1, got {spatial_scale}")
|
||
if num_frames <= 0:
|
||
raise ValueError(f"num_frames must be positive, got {num_frames}")
|
||
if vae_spatial_compression < 1 or vae_temporal_compression < 1:
|
||
raise ValueError("VAE compression factors must be positive")
|
||
if height % vae_spatial_compression != 0 or width % vae_spatial_compression != 0:
|
||
raise ValueError(f"height/width must be divisible by vae_spatial_compression={vae_spatial_compression} "
|
||
f"(got {height}x{width}).")
|
||
if (num_frames - 1) % vae_temporal_compression != 0:
|
||
raise ValueError(f"num_frames must be 1 modulo vae_temporal_compression={vae_temporal_compression} "
|
||
f"(got {num_frames}).")
|
||
|
||
if not enabled or spatial_scale == 1:
|
||
plan = RefinePlan(
|
||
target_height=height,
|
||
target_width=width,
|
||
stage1_height=height,
|
||
stage1_width=width,
|
||
spatial_scale=1,
|
||
vae_spatial_compression=vae_spatial_compression,
|
||
vae_temporal_compression=vae_temporal_compression,
|
||
num_frames=num_frames,
|
||
)
|
||
_validate_plan(plan, patch_size=patch_size, mode_label=mode_label)
|
||
return plan
|
||
|
||
if height % spatial_scale != 0 or width % spatial_scale != 0:
|
||
raise ValueError(f"{mode_label} requires height/width divisible by spatial_scale={spatial_scale} "
|
||
f"(got {height}x{width}).")
|
||
|
||
stage1_height = height // spatial_scale
|
||
stage1_width = width // spatial_scale
|
||
# Stage-1 must land on a VAE-aligned grid so the first denoise produces
|
||
# valid latents; the LTX-2 init stage enforces the same constraint.
|
||
if (stage1_height % vae_spatial_compression != 0 or stage1_width % vae_spatial_compression != 0):
|
||
raise ValueError(f"{mode_label} requires height/width divisible by "
|
||
f"{spatial_scale * vae_spatial_compression} "
|
||
f"(got {height}x{width}, vae_spatial={vae_spatial_compression}).")
|
||
|
||
plan = RefinePlan(
|
||
target_height=height,
|
||
target_width=width,
|
||
stage1_height=stage1_height,
|
||
stage1_width=stage1_width,
|
||
spatial_scale=spatial_scale,
|
||
vae_spatial_compression=vae_spatial_compression,
|
||
vae_temporal_compression=vae_temporal_compression,
|
||
num_frames=num_frames,
|
||
)
|
||
_validate_plan(plan, patch_size=patch_size, mode_label=mode_label)
|
||
logger.info(
|
||
"[MLX refine] enabled: stage1=%dx%d stage2=%dx%d scale=%dx",
|
||
stage1_width,
|
||
stage1_height,
|
||
width,
|
||
height,
|
||
spatial_scale,
|
||
)
|
||
return plan
|
||
|
||
|
||
def _validate_plan(plan: RefinePlan, *, patch_size: tuple[int, int, int], mode_label: str = "Refine") -> None:
|
||
"""
|
||
Validate that both refinement stages have latent dimensions aligned to the patch grid.
|
||
|
||
Parameters:
|
||
patch_size (tuple[int, int, int]): Temporal, height, and width patch dimensions.
|
||
|
||
Raises:
|
||
ValueError: If a stage's spatial latent dimensions or the temporal latent
|
||
dimension is not divisible by the corresponding patch dimension.
|
||
"""
|
||
pt, ph, pw = patch_size
|
||
for label, lh, lw in (
|
||
("stage1", plan.stage1_latent_height, plan.stage1_latent_width),
|
||
("stage2", plan.stage2_latent_height, plan.stage2_latent_width),
|
||
):
|
||
if lh % ph != 0 or lw % pw != 0:
|
||
raise ValueError(f"{mode_label} {label} latent grid {lh}x{lw} is not divisible by "
|
||
f"patch spatial size {ph}x{pw}.")
|
||
if plan.latent_frames % pt != 0:
|
||
raise ValueError(f"{mode_label} latent_frames={plan.latent_frames} is not divisible by "
|
||
f"patch temporal size {pt}.")
|
||
|
||
|
||
def upsample_latents_spatial(
|
||
latents: Any,
|
||
*,
|
||
scale: int = 2,
|
||
mode: str = "bilinear",
|
||
) -> Any:
|
||
"""
|
||
Upsample the spatial dimensions of 5-D latent arrays while preserving the batch, channel, and temporal dimensions.
|
||
|
||
Parameters:
|
||
latents (Any): Latents with shape ``(B, C, T, H, W)``.
|
||
scale (int): Integer factor for enlarging the spatial dimensions.
|
||
mode (str): Interpolation mode, either ``"nearest"`` or ``"bilinear"``.
|
||
|
||
Returns:
|
||
Any: Latents with shape ``(B, C, T, H * scale, W * scale)``.
|
||
"""
|
||
if scale < 1:
|
||
raise ValueError(f"scale must be >= 1, got {scale}")
|
||
if scale == 1:
|
||
return latents
|
||
|
||
# Accept both mx.array and np.ndarray so unit tests can run without MLX.
|
||
is_mlx = hasattr(latents, "dtype") and type(latents).__module__.startswith("mlx")
|
||
if is_mlx:
|
||
return _upsample_latents_mlx(latents, scale=scale, mode=mode)
|
||
return _upsample_latents_numpy(np.asarray(latents), scale=scale, mode=mode)
|
||
|
||
|
||
def _upsample_latents_numpy(
|
||
latents: np.ndarray,
|
||
*,
|
||
scale: int,
|
||
mode: str,
|
||
) -> np.ndarray:
|
||
"""Upsample 5-D latent arrays spatially using nearest-neighbor or bilinear interpolation.
|
||
|
||
Parameters:
|
||
latents (np.ndarray): Latents with shape ``(B, C, T, H, W)``.
|
||
scale (int): Spatial upsampling factor.
|
||
mode (str): Interpolation mode, either ``"nearest"`` or ``"bilinear"``.
|
||
|
||
Returns:
|
||
np.ndarray: Spatially upsampled latents with preserved batch, channel, and temporal dimensions.
|
||
|
||
Raises:
|
||
ValueError: If the latents are not five-dimensional or the interpolation mode is unsupported.
|
||
"""
|
||
if latents.ndim != 5:
|
||
raise ValueError(f"Expected 5-D latents (B,C,T,H,W), got shape {latents.shape}")
|
||
b, c, t, h, w = latents.shape
|
||
if mode == "nearest":
|
||
# (B,C,T,H,1,W,1) -> broadcast to (B,C,T,H,scale,W,scale) -> merge.
|
||
out = np.repeat(np.repeat(latents, scale, axis=3), scale, axis=4)
|
||
return out
|
||
|
||
if mode != "bilinear":
|
||
raise ValueError(f"Unsupported upsample mode: {mode}")
|
||
|
||
# Bilinear over the spatial plane. Flatten (B,C,T) into a batch of 2-D
|
||
# maps so a single vectorized gather covers every frame/channel.
|
||
src = latents.reshape(b * c * t, h, w).astype(np.float32, copy=False)
|
||
out_h, out_w = h * scale, w * scale
|
||
# Map output pixel centers onto the input grid (align_corners=False).
|
||
ys = (np.arange(out_h, dtype=np.float32) + 0.5) * (h / out_h) - 0.5
|
||
xs = (np.arange(out_w, dtype=np.float32) + 0.5) * (w / out_w) - 0.5
|
||
ys = np.clip(ys, 0.0, h - 1.0)
|
||
xs = np.clip(xs, 0.0, w - 1.0)
|
||
y0 = np.floor(ys).astype(np.int64)
|
||
x0 = np.floor(xs).astype(np.int64)
|
||
y1 = np.minimum(y0 + 1, h - 1)
|
||
x1 = np.minimum(x0 + 1, w - 1)
|
||
wy = (ys - y0.astype(np.float32))[:, None]
|
||
wx = (xs - x0.astype(np.float32))[None, :]
|
||
# Gather the four corners: shape (N, out_h, out_w).
|
||
Ia = src[:, y0[:, None], x0[None, :]]
|
||
Ib = src[:, y0[:, None], x1[None, :]]
|
||
Ic = src[:, y1[:, None], x0[None, :]]
|
||
Id = src[:, y1[:, None], x1[None, :]]
|
||
wa = (1.0 - wy) * (1.0 - wx)
|
||
wb = (1.0 - wy) * wx
|
||
wc = wy * (1.0 - wx)
|
||
wd = wy * wx
|
||
out = wa * Ia + wb * Ib + wc * Ic + wd * Id
|
||
return out.reshape(b, c, t, out_h, out_w).astype(latents.dtype, copy=False)
|
||
|
||
|
||
def _upsample_latents_mlx(latents: mx.array, *, scale: int, mode: str) -> mx.array:
|
||
"""
|
||
Upsample MLX latent tensors along their spatial dimensions.
|
||
|
||
Parameters:
|
||
latents (mx.array): A latent tensor with shape `(B, C, T, H, W)`.
|
||
scale (int): The integer spatial upsampling factor.
|
||
mode (str): The interpolation mode, such as `"nearest"` or `"bilinear"`.
|
||
|
||
Returns:
|
||
mx.array: The spatially upsampled latent tensor with its original data type.
|
||
"""
|
||
import mlx.core as mx
|
||
|
||
# Route through NumPy for the interpolation math. Latent tensors at Mac
|
||
# resolutions are small (e.g. 1×16×21×30×52 ≈ 1 MB) so the host hop is
|
||
# cheaper than carrying a bespoke Metal bilinear kernel, and it keeps
|
||
# the CPU-only unit tests and the MLX path on one implementation.
|
||
np_latents = np.array(latents.astype(mx.float32))
|
||
up = _upsample_latents_numpy(np_latents, scale=scale, mode=mode)
|
||
return mx.array(up).astype(latents.dtype)
|
||
|
||
|
||
def prepare_refine_latents(
|
||
clean_latents: Any,
|
||
*,
|
||
scale: int = 2,
|
||
sigma: float = DEFAULT_REFINE_SIGMA,
|
||
noise: Any | None = None,
|
||
add_noise_flag: bool = True,
|
||
upsample_mode: str = "bilinear",
|
||
seed: int | None = None,
|
||
) -> Any:
|
||
"""
|
||
Upsample clean latents spatially and optionally mix them with Gaussian noise.
|
||
|
||
Parameters:
|
||
clean_latents: The stage-1 latent tensor.
|
||
sigma: Noise mixing factor between 0 and 1.
|
||
noise: Optional noise tensor to mix with the upsampled latents.
|
||
add_noise_flag: Whether to apply noise mixing.
|
||
upsample_mode: Spatial interpolation mode.
|
||
seed: Optional seed for generated noise.
|
||
|
||
Returns:
|
||
The upsampled latents, optionally mixed with noise.
|
||
|
||
Raises:
|
||
ValueError: If sigma is outside the range from 0 to 1.
|
||
"""
|
||
if sigma < 0.0 or sigma > 1.0:
|
||
raise ValueError(f"sigma must be in [0, 1], got {sigma}")
|
||
|
||
upsampled = upsample_latents_spatial(clean_latents, scale=scale, mode=upsample_mode)
|
||
if not add_noise_flag or sigma == 0.0:
|
||
return upsampled
|
||
|
||
is_mlx = hasattr(upsampled, "dtype") and type(upsampled).__module__.startswith("mlx")
|
||
if noise is None:
|
||
noise = _draw_noise_like(upsampled, seed=seed, is_mlx=is_mlx)
|
||
return add_noise(upsampled, noise, float(sigma))
|
||
|
||
|
||
def refine_sigma_from_schedule(
|
||
schedule: MLXDMDSchedule,
|
||
timesteps: Sequence[float | int],
|
||
) -> float:
|
||
"""Derive the refinement noise level from the first refinement timestep.
|
||
|
||
Parameters:
|
||
schedule (MLXDMDSchedule): Schedule used to map timesteps to noise levels.
|
||
timesteps (Sequence[float | int]): Refinement timesteps, whose first value determines the sigma.
|
||
|
||
Returns:
|
||
float: Sigma corresponding to the first refinement timestep.
|
||
|
||
Raises:
|
||
ValueError: If `timesteps` is empty.
|
||
"""
|
||
if not timesteps:
|
||
raise ValueError("timesteps must be non-empty to derive a refine sigma")
|
||
return float(schedule.sigma_for(float(timesteps[0])))
|
||
|
||
|
||
def default_refine_timesteps(
|
||
schedule: MLXDMDSchedule,
|
||
timesteps: Sequence[float | int],
|
||
) -> list[float]:
|
||
"""Derive stage-2 timesteps from the stage-1 DMD grid.
|
||
|
||
The stage-2 pass must start *below* full noise, otherwise the hand-off
|
||
``(1 - sigma) * upsampled + sigma * noise`` weights stage 1 at zero and
|
||
the refine pass silently becomes a plain full-resolution generation at
|
||
twice the cost. FastWan's stage-1 grid opens at ``t=1000`` (``sigma``
|
||
exactly 1.0), so reusing it verbatim — which is what happens when
|
||
``--refine-dmd-denoising-steps`` is left unset — discards stage 1.
|
||
|
||
Dropping the leading full-noise entries keeps the pass on timesteps the
|
||
distilled student was actually trained on (no off-grid ``t`` the DiT has
|
||
never seen) while letting the stage-1 structure through.
|
||
|
||
Parameters:
|
||
schedule (MLXDMDSchedule): Schedule used to map timesteps to noise levels.
|
||
timesteps (Sequence[float | int]): The stage-1 DMD timestep grid.
|
||
|
||
Returns:
|
||
list[float]: The stage-1 grid with leading full-noise timesteps removed.
|
||
|
||
Raises:
|
||
ValueError: If every timestep in the grid is at full noise, leaving no
|
||
usable refine step.
|
||
"""
|
||
steps = [float(step) for step in timesteps]
|
||
first = 0
|
||
while first < len(steps) and schedule.sigma_for(steps[first]) >= 1.0:
|
||
first += 1
|
||
if first == len(steps):
|
||
raise ValueError(f"No usable refine timesteps in {steps}: every entry is at sigma >= 1 "
|
||
"(full noise), which would discard the stage-1 result. Pass "
|
||
"explicit stage-2 timesteps below the full-noise step.")
|
||
return steps[first:]
|
||
|
||
|
||
def run_dmd_loop(
|
||
*,
|
||
dit: Any,
|
||
latents: Any,
|
||
encoder_hidden_states: Any,
|
||
freqs_cis: tuple[Any, Any],
|
||
timesteps: Sequence[float | int],
|
||
schedule: MLXDMDSchedule,
|
||
mx_dtype: Any,
|
||
seed: int | None = None,
|
||
step_callback: Callable[[int, int], None] | None = None,
|
||
label: str = "denoise",
|
||
) -> Any:
|
||
"""
|
||
Denoise latents over the supplied timesteps using the DMD schedule.
|
||
|
||
Parameters:
|
||
timesteps (Sequence[float | int]): Denoising timesteps in execution order.
|
||
seed (int | None): Seed for reproducible intermediate noise generation.
|
||
step_callback (Callable[[int, int], None] | None): Callback receiving the
|
||
completed step number and total step count.
|
||
label (str): Label used for progress output when no callback is provided.
|
||
|
||
Returns:
|
||
Any: The denoised latents.
|
||
"""
|
||
import mlx.core as mx
|
||
|
||
renoise_rng = np.random.default_rng(seed) if seed is not None else None
|
||
latents_out = latents
|
||
n_steps = len(timesteps)
|
||
for step_index, timestep in enumerate(timesteps):
|
||
noise_input = latents_out
|
||
ts_val = float(timestep)
|
||
timestep_mx = mx.array([ts_val]).astype(mx.float32)
|
||
noise_pred = dit(
|
||
latents_out.astype(mx_dtype),
|
||
encoder_hidden_states,
|
||
timestep_mx,
|
||
freqs_cis,
|
||
)
|
||
noise_input_f32 = noise_input.astype(mx.float32)
|
||
pred_noise_f32 = noise_pred.astype(mx.float32)
|
||
if step_index < n_steps - 1:
|
||
next_ts: float | None = float(timesteps[step_index + 1])
|
||
if renoise_rng is not None:
|
||
renoise = mx.array(renoise_rng.standard_normal(tuple(noise_input_f32.shape)).astype(np.float32))
|
||
else:
|
||
renoise = mx.random.normal(noise_input_f32.shape).astype(mx.float32)
|
||
else:
|
||
next_ts, renoise = None, None
|
||
latents_out = dmd_step(
|
||
latents=noise_input_f32,
|
||
noise_input_latent=noise_input_f32,
|
||
pred_noise=pred_noise_f32,
|
||
schedule=schedule,
|
||
timestep=ts_val,
|
||
next_timestep=next_ts,
|
||
noise=renoise,
|
||
).astype(mx_dtype)
|
||
mx.eval(latents_out)
|
||
if step_callback is not None:
|
||
step_callback(step_index + 1, n_steps)
|
||
else:
|
||
print(f"{label} step {step_index + 1}/{n_steps} complete")
|
||
return latents_out
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class TwoPassResult:
|
||
"""Outputs of :func:`run_two_pass_dmd`."""
|
||
|
||
latents: Any
|
||
stage1_latents: Any
|
||
plan: RefinePlan
|
||
refine_sigma: float
|
||
|
||
|
||
def run_two_pass_dmd(
|
||
*,
|
||
dit: Any,
|
||
encoder_hidden_states: Any,
|
||
noise_latents_stage1: Any,
|
||
freqs_cis_stage1: tuple[Any, Any],
|
||
freqs_cis_stage2: tuple[Any, Any] | None,
|
||
plan: RefinePlan,
|
||
schedule: MLXDMDSchedule,
|
||
timesteps: Sequence[float | int],
|
||
refine_timesteps: Sequence[float | int] | None = None,
|
||
mx_dtype: Any,
|
||
seed: int = 0,
|
||
add_noise_flag: bool = True,
|
||
upsample_mode: str = "bilinear",
|
||
refine_sigma: float | None = None,
|
||
step_callback: Callable[[str, int, int], None] | None = None,
|
||
) -> TwoPassResult:
|
||
"""
|
||
Run base denoising and, when enabled, spatial refinement denoising.
|
||
|
||
Parameters:
|
||
dit: DiT callable used for both denoising passes.
|
||
encoder_hidden_states: Prompt embeddings shared across both passes.
|
||
noise_latents_stage1: Initial stage-1 noise latents.
|
||
freqs_cis_stage1: RoPE tables for the stage-1 resolution.
|
||
freqs_cis_stage2: RoPE tables for the stage-2 resolution, required when refinement is enabled.
|
||
plan: Refinement geometry and configuration.
|
||
schedule: Flow-matching schedule used by both passes.
|
||
timesteps: Stage-1 denoising timesteps.
|
||
refine_timesteps: Stage-2 denoising timesteps. Uses `timesteps` when omitted.
|
||
mx_dtype: MLX dtype used for DiT inputs and outputs.
|
||
seed: Base seed for reproducible noise generation.
|
||
add_noise_flag: Whether to add noise to the upsampled stage-1 latents.
|
||
upsample_mode: Spatial upsampling mode, either `"bilinear"` or `"nearest"`.
|
||
refine_sigma: Stage-2 starting noise level. Derived from the first refinement timestep when omitted.
|
||
step_callback: Optional callback receiving the phase name, step index, and total step count.
|
||
|
||
Returns:
|
||
TwoPassResult containing the final latents, stage-1 latents, refinement plan, and applied refinement sigma.
|
||
|
||
Raises:
|
||
ValueError: If refinement is enabled without stage-2 RoPE tables, without refinement timesteps, or if upsampled latents do not match the planned stage-2 dimensions.
|
||
"""
|
||
stage1_cb = None
|
||
stage2_cb = None
|
||
if step_callback is not None:
|
||
stage1_cb = lambda i, n: step_callback("stage1", i, n) # noqa: E731
|
||
stage2_cb = lambda i, n: step_callback("stage2", i, n) # noqa: E731
|
||
|
||
stage1_latents = run_dmd_loop(
|
||
dit=dit,
|
||
latents=noise_latents_stage1,
|
||
encoder_hidden_states=encoder_hidden_states,
|
||
freqs_cis=freqs_cis_stage1,
|
||
timesteps=timesteps,
|
||
schedule=schedule,
|
||
mx_dtype=mx_dtype,
|
||
seed=seed,
|
||
step_callback=stage1_cb,
|
||
label="stage1 denoise",
|
||
)
|
||
|
||
if plan.spatial_scale == 1:
|
||
return TwoPassResult(
|
||
latents=stage1_latents,
|
||
stage1_latents=stage1_latents,
|
||
plan=plan,
|
||
refine_sigma=0.0,
|
||
)
|
||
|
||
if freqs_cis_stage2 is None:
|
||
raise ValueError("freqs_cis_stage2 is required when refine spatial_scale > 1")
|
||
|
||
if refine_timesteps is not None:
|
||
stage2_timesteps = [float(step) for step in refine_timesteps]
|
||
if not stage2_timesteps:
|
||
raise ValueError("refine_timesteps must be non-empty when refine is enabled")
|
||
else:
|
||
# Not `list(timesteps)`: the stage-1 grid opens at full noise, which
|
||
# would weight the stage-1 result at zero. See default_refine_timesteps.
|
||
stage2_timesteps = default_refine_timesteps(schedule, timesteps)
|
||
grid_sigma = refine_sigma_from_schedule(schedule, stage2_timesteps)
|
||
sigma = float(refine_sigma) if refine_sigma is not None else grid_sigma
|
||
if refine_sigma is not None and abs(sigma - grid_sigma) > 1e-6:
|
||
# The loop tells the DiT `stage2_timesteps[0]`, which implies grid_sigma.
|
||
# Overriding the hand-off noise level breaks that correspondence, so the
|
||
# model is denoising from a level it was not told about. Useful for
|
||
# exploring schedules that bottom out too high, but say so out loud.
|
||
logger.warning(
|
||
"[MLX refine] refine_sigma=%.4f overrides the schedule's %.4f for timestep %g; "
|
||
"the DiT is told a timestep that no longer matches the noise it receives.",
|
||
sigma,
|
||
grid_sigma,
|
||
stage2_timesteps[0],
|
||
)
|
||
|
||
# A hand-off at sigma >= 1 is `0 * upsampled + 1 * noise`: stage 1 is
|
||
# thrown away and refine degrades to a plain full-res run at 2x the cost.
|
||
# Fail loudly rather than silently burning the first pass.
|
||
if add_noise_flag and sigma >= 1.0:
|
||
raise ValueError(f"Refine hand-off sigma={sigma:.4f} (from stage-2 timestep "
|
||
f"{stage2_timesteps[0]:g}) discards the stage-1 result entirely: "
|
||
"the upsampled latents are weighted (1 - sigma) = 0. Start the "
|
||
"stage-2 grid below the full-noise timestep, or pass "
|
||
"add_noise_flag=False to hand off the clean upsample.")
|
||
|
||
stage2_input = prepare_refine_latents(
|
||
stage1_latents,
|
||
scale=plan.spatial_scale,
|
||
sigma=sigma,
|
||
add_noise_flag=add_noise_flag,
|
||
upsample_mode=upsample_mode,
|
||
seed=seed + 1,
|
||
)
|
||
|
||
# Shape guard: upsampled latents must match the stage-2 RoPE grid.
|
||
expected_h = plan.stage2_latent_height
|
||
expected_w = plan.stage2_latent_width
|
||
got_h, got_w = int(stage2_input.shape[-2]), int(stage2_input.shape[-1])
|
||
if got_h != expected_h or got_w != expected_w:
|
||
raise ValueError(f"Refine upsample produced {got_h}x{got_w} latents, expected "
|
||
f"{expected_h}x{expected_w} for target "
|
||
f"{plan.target_height}x{plan.target_width}.")
|
||
|
||
logger.info(
|
||
"[MLX refine] stage2 start: latent=%dx%d sigma=%.4f steps=%d",
|
||
expected_w,
|
||
expected_h,
|
||
sigma,
|
||
len(stage2_timesteps),
|
||
)
|
||
|
||
stage2_latents = run_dmd_loop(
|
||
dit=dit,
|
||
latents=stage2_input,
|
||
encoder_hidden_states=encoder_hidden_states,
|
||
freqs_cis=freqs_cis_stage2,
|
||
timesteps=stage2_timesteps,
|
||
schedule=schedule,
|
||
mx_dtype=mx_dtype,
|
||
seed=seed + 2,
|
||
step_callback=stage2_cb,
|
||
label="stage2 refine",
|
||
)
|
||
return TwoPassResult(
|
||
latents=stage2_latents,
|
||
stage1_latents=stage1_latents,
|
||
plan=plan,
|
||
refine_sigma=sigma,
|
||
)
|
||
|
||
|
||
__all__ = [
|
||
"DEFAULT_REFINE_SIGMA",
|
||
"RefinePlan",
|
||
"TwoPassResult",
|
||
"default_refine_timesteps",
|
||
"plan_refine_resolutions",
|
||
"prepare_refine_latents",
|
||
"refine_sigma_from_schedule",
|
||
"run_dmd_loop",
|
||
"run_two_pass_dmd",
|
||
"upsample_latents_spatial",
|
||
]
|
||
|
||
|
||
def _draw_noise_like(like: Any, *, seed: int | None, is_mlx: bool) -> Any:
|
||
"""Generate Gaussian noise with the shape and array type of the input."""
|
||
shape = tuple(int(s) for s in like.shape)
|
||
if seed is not None:
|
||
rng = np.random.default_rng(seed)
|
||
noise_np = rng.standard_normal(shape).astype(np.float32)
|
||
if is_mlx:
|
||
import mlx.core as mx
|
||
|
||
return mx.array(noise_np).astype(mx.float32)
|
||
return noise_np.astype(np.asarray(like).dtype, copy=False)
|
||
if is_mlx:
|
||
import mlx.core as mx
|
||
|
||
return mx.random.normal(shape).astype(mx.float32)
|
||
return np.random.standard_normal(shape).astype(np.asarray(like).dtype, copy=False)
|