refine ¶

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

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,
    )

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)