utils.py

"""Shared spectral expansion and transition scheduling utilities."""
from __future__ import annotations

import csv
import math
import os
import re
from pathlib import Path
from typing import Iterable, List, Sequence, Tuple

import numpy as np
import pywt
import torch
import yaml
from scipy.fft import dctn, idctn


def power_spectrum(omega: float, A: float, beta: float) -> float:
    """Radial power-law spectrum ``P(omega) = A * |omega|**(-beta)``.

    Args:
    - omega: Radial spatial frequency.
    - A: Power-law amplitude (fitted per VAE model in ``configs.yaml``).
    - beta: Power-law decay exponent (fitted per VAE model in ``configs.yaml``).

    Returns:
    - The power-spectrum value ``P(omega)``.
    """
    return A * abs(omega) ** (-beta)


def activation_time(P_omega: float, delta: float) -> float:
    """Return the activation time for one radial frequency ``omega``. This matches Eq. 9 in the paper.

    Args:
    - P_omega: Power-spectrum value ``P(omega)`` at the frequency of interest.
    - delta: Noise-dominated tolerance; smaller ``delta`` delays activation.

    Returns:
    - The activation time ``t_omega`` in ``(0, 1)``.
    """
    if delta >= 1.0:
        raise ValueError(
            f"delta={delta} >= 1, but we assume the error threshold is < 1."
        )
    return 1.0 / (1.0 + math.sqrt(delta / (P_omega * (1.0 + P_omega - delta))))


def delta_optimal_transitions(
    scales: Sequence[float],
    delta: float,
    A: float,
    beta: float,
    H: int,
    W: int,
) -> List[float]:
    """Return transition times for adjacent scales. This matches Eq. 10 from the paper.

    Args:
    - scales: Strictly increasing scale list ending at 1.0. Example: ``[0.25, 0.5, 1.0]``.
    - delta: Noise-dominated tolerance passed to ``activation_time``.
    - A: Power-law amplitude.
    - beta: Power-law decay exponent.
    - H: Full-resolution latent height (sets ``omega_max = min(H, W) / 2``).
    - W: Full-resolution latent width (sets ``omega_max = min(H, W) / 2``).

    Returns:
    - List of transition times ``t*_i`` (length ``len(scales) - 1``).
    """
    validate_scales(scales)
    # Nyquist maximum representable frequency limit
    omega_max = min(H, W) / 2.0
    transitions: List[float] = []
    for i in range(len(scales) - 1):
        omega_i = scales[i] * omega_max
        transitions.append(activation_time(power_spectrum(omega_i, A, beta), delta))
    return transitions


def align_timestep(t: float, r: float) -> float:
    """Return the aligned flow-matching time after spectral noise expansion. This matches Eq. 6 of the paper.

    Args:
    - t: Flow-matching time at the resolution transition.
    - r: Resolution scale ratio ``s_{i + 1} / s_i`` of the transition.

    Returns:
    - The aligned flow-matching time ``t_tilde``.
    """
    return t * kappa(t, r)


def kappa(t: float, r: float) -> float:
    """Return the state-rescaling factor after spectral noise expansion. This matches Eq. 5 of the paper.

    Args:
    - t: Flow-matching time at the resolution transition.
    - r: Scale ratio ``s_{i + 1} / s_i`` of the transition.

    Returns:
    - The state-rescaling factor ``kappa``.
    """
    return r / (1.0 + (r - 1.0) * t)


def _dct_expand_np(
    x_np: np.ndarray, target_hw: Tuple[int, int], t: float, seed: int,
) -> np.ndarray:
    """DCT spectral noise expansion.

    Args:
    - x_np: Source array; trailing two axes are the spatial grid to expand.
    - target_hw: Target ``(height, width)`` of the expanded grid.
    - t: Noise amplitude for the high-frequency coefficients, which is also the current flow matching time. 
    - seed: Seed for the per-call random generator.

    Returns:
    - The expanded array at ``target_hw`` (float32, same leading axes as ``x_np``).
    """
    H_tgt, W_tgt = target_hw
    H_src, W_src = x_np.shape[-2], x_np.shape[-1]
    if H_tgt < H_src or W_tgt < W_src:
        raise ValueError(
            f"DCT expand: cannot expand to target {target_hw} smaller than source ({H_src}, {W_src})."
        )
    rng = np.random.default_rng(seed)
    out = np.empty(x_np.shape[:-2] + (H_tgt, W_tgt), dtype=np.float32)
    for idx in np.ndindex(*x_np.shape[:-2]):
        coeffs_src = dctn(x_np[idx], type=2, norm="ortho")
        big = t * rng.standard_normal((H_tgt, W_tgt)).astype(np.float32)
        big[:H_src, :W_src] = coeffs_src
        out[idx] = idctn(big, type=2, norm="ortho").astype(np.float32)
    return out


def _dwt_expand_np(x_np: np.ndarray, t: float, seed: int) -> np.ndarray:
    """Haar wavelet spectral noise expansion. The target H, W is automatically two times the source H, W.

    Args:
    - x_np: Source array treated as the LL band; trailing two axes are spatial.
    - t: Noise amplitude for the LH/HL/HH detail bands, which is also the current flow matching time.
    - seed: Seed for the per-call random generator.

    Returns:
    - The expanded array at twice the source resolution (float32).
    """
    H_src, W_src = x_np.shape[-2], x_np.shape[-1]
    H_tgt, W_tgt = H_src * 2, W_src * 2
    rng = np.random.default_rng(seed)
    out = np.empty(x_np.shape[:-2] + (H_tgt, W_tgt), dtype=np.float32)
    for idx in np.ndindex(*x_np.shape[:-2]):
        LL = x_np[idx]
        LH = t * rng.standard_normal(LL.shape).astype(np.float32)
        HL = t * rng.standard_normal(LL.shape).astype(np.float32)
        HH = t * rng.standard_normal(LL.shape).astype(np.float32)
        out[idx] = pywt.waverec2(
            [LL, (LH, HL, HH)], "haar", mode="periodization"
        ).astype(np.float32)
    return out


def _fft_expand_np(
    x_np: np.ndarray, target_hw: Tuple[int, int], t: float, seed: int,
) -> np.ndarray:
    """FFT spectral noise expansion.

    Args:
    - x_np: Source array; trailing two axes are the spatial grid to expand.
    - target_hw: Target ``(height, width)`` of the expanded grid.
    - t: Noise amplitude for the outer (high-frequency) spectrum, which is also the current flow matching time.
    - seed: Seed for the per-call random generator.

    Returns:
    - The expanded array at ``target_hw`` (float32, same leading axes as ``x_np``).
    """
    H_tgt, W_tgt = target_hw
    H_src, W_src = x_np.shape[-2], x_np.shape[-1]
    if H_tgt < H_src or W_tgt < W_src:
        raise ValueError(
            f"FFT expand: cannot expand to target {target_hw} smaller than source ({H_src}, {W_src})."
        )
    rng = np.random.default_rng(seed)
    pad_h, pad_w = (H_tgt - H_src) // 2, (W_tgt - W_src) // 2
    out = np.empty(x_np.shape[:-2] + (H_tgt, W_tgt), dtype=np.float32)
    for idx in np.ndindex(*x_np.shape[:-2]):
        X_src = np.fft.fftshift(np.fft.fft2(x_np[idx], norm="ortho"))
        nr = rng.standard_normal((H_tgt, W_tgt)).astype(np.float32)
        ni = rng.standard_normal((H_tgt, W_tgt)).astype(np.float32)
        X_big = np.fft.fftshift(t * (nr + 1j * ni) / np.sqrt(2.0))
        X_big[pad_h:pad_h + H_src, pad_w:pad_w + W_src] = X_src
        out[idx] = np.fft.ifft2(np.fft.ifftshift(X_big), norm="ortho").real.astype(np.float32)
    return out


def spectral_expand_and_align(
    x: torch.Tensor,
    target_hw: Tuple[int, int],
    r: float,
    t: float,
    transform: str,
    seed: int,
) -> Tuple[torch.Tensor, float]:
    """Expand ``x`` to ``target_hw`` and return the rescaled latent and aligned time.

    The caller is the single source of truth for spatial dimensions: it passes
    the already-resolved ``target_hw`` (e.g. snapped to even dims for FLUX) and
    the user-intended scale ratio ``r = s_{i + 1} / s_i`` used by the alignment
    maths (which must not drift with integer rounding of the dims).

    Args:
    - x: Current latent at the source resolution; trailing two axes are spatial.
    - target_hw: Target ``(height, width)`` to expand to (resolved by the caller).
    - r: Scale ratio ``s_{i + 1} / s_i`` of the transition.
    - t: Flow-matching time at the transition.
    - transform: Spectral basis, one of ``'dct'``, ``'dwt'``, ``'fft'``.
    - seed: Seed for the per-call random generator.

    Returns:
    - A ``(x_tilde, t_tilde)`` tuple: the rescaled and expanded latent and the
      aligned flow-matching time.
    """
    if transform not in ("dct", "dwt", "fft"):
        raise ValueError(f"transform must be 'dct'|'dwt'|'fft', got {transform!r}")

    H_tgt, W_tgt = target_hw
    H_src, W_src = x.shape[-2], x.shape[-1]
    x_np = x.detach().cpu().float().numpy()
    if transform == "dwt":
        if abs(r - 2.0) > 1e-6:
            raise ValueError(
                f"DWT requires a 2x scale ratio between consecutive scales; "
                f"got r = s_next/s_i = {r:.4f}. "
                f"Use --transform dct or --transform fft for non-dyadic scales."
            )
        # Haar IDWT is structurally 2x
        if H_tgt != 2 * H_src or W_tgt != 2 * W_src:
            raise ValueError(
                f"DWT requires H_tgt = 2·H_src and W_tgt = 2·W_src; got "
                f"({H_src}, {W_src}) -> ({H_tgt}, {W_tgt})."
            )
        expanded = _dwt_expand_np(x_np, t, seed)
    elif transform == "dct":
        expanded = _dct_expand_np(x_np, (H_tgt, W_tgt), t, seed)
    else:  # fft
        expanded = _fft_expand_np(x_np, (H_tgt, W_tgt), t, seed)

    rescaled = (kappa(t, r) * expanded).astype(np.float32)
    x_tilde = torch.from_numpy(rescaled).to(device=x.device, dtype=x.dtype)
    return x_tilde, align_timestep(t, r)


def find_first_step_below(sigmas: Iterable[float], threshold: float) -> int:
    """Return the first step index whose sigma (step timestamp) is below ``threshold (transition timestep)``.

    Args:
    - sigmas: Scheduler sigma sequence (length ``n_steps + 1``).
    - threshold: Sigma threshold; the first step at or below it is returned.

    Returns:
    - The first step index with ``sigma <= threshold``, or ``n_steps`` if none.
    """
    sigmas = list(sigmas)
    n_steps = len(sigmas) - 1
    for i in range(n_steps):
        s = sigmas[i].item() if hasattr(sigmas[i], "item") else float(sigmas[i])
        if s <= threshold:
            return i
    return n_steps


def reset_scheduler_state(scheduler, step_index: int) -> None:
    """Reset solver buffers after a transition.

    Args:
    - scheduler: Diffusers scheduler whose solver state is reset in place.
    - step_index: Step index to set ``scheduler._step_index`` to.

    Returns:
    - None; the scheduler is mutated in place.
    """
    if hasattr(scheduler, "model_outputs"):
        order = getattr(scheduler.config, "solver_order", 1)
        scheduler.model_outputs = [None] * order
    if hasattr(scheduler, "lower_order_nums"):
        scheduler.lower_order_nums = 0
    if hasattr(scheduler, "last_sample"):
        scheduler.last_sample = None
    scheduler._step_index = step_index


def validate_scales(scales: Sequence[float]) -> None:
    """Validate a strictly increasing resolution scale list ending at 1.0.

    Args:
    - scales: Scale list to validate; each value in ``(0, 1]``, strictly
      increasing, ending at ``1.0``.

    Returns:
    - None; raises ``ValueError`` if the scales are invalid.
    """
    if len(scales) == 0:
        raise ValueError("list of resolution scales is empty; supply at least one value.")
    if any(s <= 0.0 or s > 1.0 for s in scales):
        raise ValueError(f"every scale must be in (0, 1]; got {list(scales)}")
    if abs(scales[-1] - 1.0) > 1e-6:
        raise ValueError(f"last scale must equal 1.0 (full resolution); got {scales[-1]}")
    for a, b in zip(scales[:-1], scales[1:]):
        if not (a < b):
            raise ValueError(f"scales must be strictly increasing; got {list(scales)}")


def parse_scales(tokens, full_height: int) -> List[float]:
    """Parse a scale spec into ``[s_1, ..., s_S = 1.0]``.

    Tokens may be space- or comma-separated and take any of three forms,
    freely mixed:
      - decimal scales -- ``0.5`` ``0.37``
      - fraction scales -- ``1/2`` ``2/3``
      - stage heights in pixels -- ``480 720`` (triggered when any value
        exceeds 1); the last must equal ``full_height`` and every value is
        divided by it to recover the scale.

    Args:
    - tokens: Scale tokens (strings or numbers); each may itself be a
      comma-separated group.
    - full_height: Configured full-resolution height, used to convert the
      pixel-height form to scales and to validate its last value.

    Returns:
    - The validated scale list ``[s_1, ..., s_S = 1.0]``.
    """
    flat: List[str] = []
    for t in tokens:
        flat.extend(str(t).split(","))

    def _one(tok: str) -> float:
        tok = tok.strip()
        if "/" in tok:
            num, den = tok.split("/", 1)
            return float(num) / float(den)
        return float(tok)

    values = [_one(t) for t in flat if t.strip()]
    if not values:
        raise ValueError("--scales is empty; supply at least one value.")
    if max(values) > 1.0 + 1e-6:
        # Resolution mode: tokens are the pixel height of each stage.
        if abs(values[-1] - full_height) > 1e-6:
            raise ValueError(
                f"the last stage resolution {values[-1]:g} must equal the "
                f"configured height {full_height}; otherwise pass scales in (0, 1]."
            )
        values = [v / full_height for v in values]
    validate_scales(values)
    return values


_ENV_PATTERN = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}")


def _expand_env(value):
    """Recursively expand ``${ENV_VAR}`` placeholders in a config value.

    Args:
    - value: A str, dict, or list parsed from YAML; other types pass through.

    Returns:
    - The value with every ``${VAR}`` placeholder expanded (same structure
      as the input).
    """
    if isinstance(value, str):
        def repl(m: re.Match) -> str:
            """Substitute one ``${VAR}`` match with its environment value.

            Args:
            - m: Regex match whose group 1 is the environment variable name.

            Returns:
            - The environment-variable value for the matched name.
            """
            var = m.group(1)
            if var not in os.environ:
                raise KeyError(
                    f"Environment variable {var!r} referenced in configs.yaml "
                    "is not set."
                )
            return os.environ[var]
        return _ENV_PATTERN.sub(repl, value)
    if isinstance(value, dict):
        return {k: _expand_env(v) for k, v in value.items()}
    if isinstance(value, list):
        return [_expand_env(v) for v in value]
    return value


def load_config(yaml_path: str | Path, model_key: str) -> dict:
    """Load a model config and expand ``${ENV_VAR}`` placeholders.

    Args:
    - yaml_path: Path to the ``configs.yaml`` file.
    - model_key: Top-level model key to select (e.g. ``'flux'``, ``'wan21'``).

    Returns:
    - The selected model config dict with ``${ENV_VAR}`` placeholders expanded.
    """
    with open(yaml_path, "r") as f:
        data = yaml.safe_load(f)
    if model_key not in data:
        raise KeyError(f"model {model_key!r} not in {yaml_path}; have {list(data)}")
    return _expand_env(data[model_key])


def load_prompts(args) -> List[str]:
    """Collect prompts from ``--prompts``, ``--prompt_txt``, or ``--prompt_csv``.

    Args:
    - args: Parsed argparse namespace.

    Returns:
    - The list of prompt strings.
    """
    if args.prompts:
        return list(args.prompts)
    if args.prompt_txt:
        return [ln.strip() for ln in open(args.prompt_txt) if ln.strip()]
    return [row["prompt"] for row in csv.DictReader(open(args.prompt_csv))]