mlx_runtime ¶

def add_memory_limit_args(
    parser: argparse.ArgumentParser,
    *,
    mlx_memory_limit_gib: float | None = None,
    mlx_cache_limit_gib: float | None = None,
    mlx_disable_cache: bool = False,
    mlx_wired_limit_gib: float | None = None,
    torch_mps_high_watermark_ratio: float | None = None,
    torch_mps_low_watermark_ratio: float | None = None,
) -> None:
    """
    Add configurable Apple Silicon memory-limit options to an argument parser.

    Parameters:
        parser (argparse.ArgumentParser): Parser to which the options are added.
        mlx_memory_limit_gib (float | None): Default MLX memory limit in GiB.
        mlx_cache_limit_gib (float | None): Default MLX cache limit in GiB.
        mlx_disable_cache (bool): Whether the cache limit defaults to zero.
        mlx_wired_limit_gib (float | None): Default MLX wired-memory limit in GiB.
        torch_mps_high_watermark_ratio (float | None): Default PyTorch MPS high-watermark ratio.
        torch_mps_low_watermark_ratio (float | None): Default PyTorch MPS low-watermark ratio.
    """
    parser.add_argument("--mlx-memory-limit-gib",
                        type=float,
                        default=mlx_memory_limit_gib,
                        help="Set MLX memory limit in GiB for memory-tier testing (DiT path).")
    parser.add_argument("--mlx-cache-limit-gib",
                        type=float,
                        default=mlx_cache_limit_gib,
                        help="Set MLX cache limit in GiB. Use --mlx-disable-cache to force 0.")
    parser.add_argument("--mlx-disable-cache",
                        action="store_true",
                        default=mlx_disable_cache,
                        help="Set MLX cache limit to 0 for stricter memory-tier tests.")
    parser.add_argument("--mlx-wired-limit-gib",
                        type=float,
                        default=mlx_wired_limit_gib,
                        help="Set MLX wired-memory limit in GiB where supported by macOS/MLX.")
    parser.add_argument("--torch-mps-high-watermark-ratio",
                        type=float,
                        default=torch_mps_high_watermark_ratio,
                        help="Set PYTORCH_MPS_HIGH_WATERMARK_RATIO before importing torch.")
    parser.add_argument("--torch-mps-low-watermark-ratio",
                        type=float,
                        default=torch_mps_low_watermark_ratio,
                        help="Set PYTORCH_MPS_LOW_WATERMARK_RATIO before importing torch.")

def apply_fast_spatial_upsample(
    frames: Iterable[np.ndarray],
    spatial: FastSpatialPlan,
) -> list[np.ndarray]:
    """Resample decoded stage-1 frames up to the target resolution.

    This runs on decoded RGB frames, *not* on latents: see the module
    docstring for why the latent-space version produced a blurred veil.

    Parameters:
        frames (Iterable[np.ndarray]): Decoded HxWx3 uint8 RGB frames, produced
            by decoding at the stage-one resolution.
        spatial (FastSpatialPlan): Plan defining the target size, interpolation
            kernel, and unsharp strength.

    Returns:
        list[np.ndarray]: Frames at the target resolution. When spatial scaling
            is disabled the frames are returned unchanged, as a list.
    """
    if not spatial.enabled:
        return list(frames)
    return upsample_frames(
        frames,
        width=spatial.target_width,
        height=spatial.target_height,
        mode=spatial.upsample_mode,
        sharpen=spatial.sharpen,
    )

def apply_memory_limits(
    *,
    mlx_memory_limit_gib: float | None = None,
    mlx_cache_limit_gib: float | None = None,
    mlx_disable_cache: bool = False,
    mlx_wired_limit_gib: float | None = None,
    torch_mps_high_watermark_ratio: float | None = None,
    torch_mps_low_watermark_ratio: float | None = None,
    mx_module: Any | None = None,
) -> AppliedMemoryLimits:
    """Apply optional MLX allocator limits and PyTorch MPS watermarks.

    PyTorch reads MPS watermark variables when the MPS backend initializes, so
    call this before importing PyTorch. Specifying only a high watermark sets the
    low watermark to ``0.0``. MLX limit-setting failures are recorded in the
    result and do not prevent other limits from being applied.

    Parameters:
        mlx_memory_limit_gib (float | None): Maximum MLX memory in GiB.
        mlx_cache_limit_gib (float | None): Maximum MLX cache size in GiB.
        mlx_disable_cache (bool): Whether to disable the MLX cache.
        mlx_wired_limit_gib (float | None): Maximum MLX wired memory in GiB.
        torch_mps_high_watermark_ratio (float | None): PyTorch MPS high watermark
            ratio.
        torch_mps_low_watermark_ratio (float | None): PyTorch MPS low watermark
            ratio.

    Returns:
        AppliedMemoryLimits: Configured values, applied and previous MLX byte
            limits, MPS watermark values, and per-limit errors.
    """
    if torch_mps_high_watermark_ratio is not None and torch_mps_low_watermark_ratio is None:
        torch_mps_low_watermark_ratio = 0.0

    high = _set_mps_env("PYTORCH_MPS_HIGH_WATERMARK_RATIO", torch_mps_high_watermark_ratio)
    low = _set_mps_env("PYTORCH_MPS_LOW_WATERMARK_RATIO", torch_mps_low_watermark_ratio)

    memory_bytes = gib_to_bytes(mlx_memory_limit_gib)
    cache_bytes = 0 if mlx_disable_cache else gib_to_bytes(mlx_cache_limit_gib)
    wired_bytes = gib_to_bytes(mlx_wired_limit_gib)

    applied: dict[str, int] = {}
    previous: dict[str, int] = {}
    errors: dict[str, str] = {}
    if memory_bytes is not None or cache_bytes is not None or wired_bytes is not None:
        if mx_module is None:
            import mlx.core as mx

            mx_module = mx

        # Apply each limit independently; record failures without stopping.
        limits = [
            ("mlx_memory_limit", memory_bytes, mx_module.set_memory_limit),
            ("mlx_cache_limit", cache_bytes, mx_module.set_cache_limit),
            ("mlx_wired_limit", wired_bytes, mx_module.set_wired_limit),
        ]
        for name, value, setter in limits:
            if value is not None:
                try:
                    previous[name] = int(setter(value))
                    applied[name] = value
                except Exception as exc:  # noqa: BLE001 - macOS/system-limit dependent.
                    errors[name] = f"{type(exc).__name__}: {exc}"

    return AppliedMemoryLimits(
        mlx_memory_limit_gib=mlx_memory_limit_gib,
        mlx_cache_limit_gib=mlx_cache_limit_gib,
        mlx_disable_cache=mlx_disable_cache,
        mlx_wired_limit_gib=mlx_wired_limit_gib,
        torch_mps_high_watermark_ratio=high,
        torch_mps_low_watermark_ratio=low,
        applied_bytes=applied,
        previous_bytes=previous,
        errors=errors,
    )

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 enhance_prompt(
    prompt: str,
    *,
    backend: str = "auto",
    model: str | None = None,
    system_prompt: str = DEFAULT_ENHANCE_SYSTEM_PROMPT,
    max_tokens: int = 128,
) -> EnhanceResult:
    """Enhance a prompt using the selected backend, falling back to a deterministic template when configured for automatic selection.

    Parameters:
        prompt (str): The prompt to enhance.
        backend (str): The enhancement backend: ``"auto"``, ``"mlx-lm"``, or ``"template"``.
        model (str | None): The MLX language model to use.
        system_prompt (str): Instructions provided to the MLX language model.
        max_tokens (int): Maximum number of tokens generated by the MLX language model.

    Returns:
        EnhanceResult: The original and enhanced prompts, selected backend, timing information, and model metadata.

    Raises:
        ValueError: If the prompt is empty or the backend is unsupported.
        Exception: If the explicitly selected ``"mlx-lm"`` backend fails.
    """
    text = _normalize_user_prompt(prompt)
    backend_norm = (backend or "auto").lower()
    if backend_norm not in {"auto", "mlx-lm", "template"}:
        raise ValueError(f"Unknown enhance backend: {backend}")

    start = time.perf_counter()
    used_model: str | None = None

    if backend_norm in {"auto", "mlx-lm"}:
        try:
            used_model = model or DEFAULT_MLX_LM_MODEL
            enhanced = enhance_prompt_mlx_lm(
                text,
                model=used_model,
                system_prompt=system_prompt,
                max_tokens=max_tokens,
            )
            return EnhanceResult(
                original=text,
                enhanced=enhanced,
                backend="mlx-lm",
                elapsed_s=time.perf_counter() - start,
                model=used_model,
            )
        except Exception as exc:
            if backend_norm == "mlx-lm":
                raise
            logger.info(
                "[MLX enhance] mlx-lm unavailable (%s); using template backend",
                exc,
            )

    enhanced = enhance_prompt_template(text)
    return EnhanceResult(
        original=text,
        enhanced=enhanced,
        backend="template",
        elapsed_s=time.perf_counter() - start,
        model=None,
    )

def enhance_prompt_template(prompt: str) -> str:
    """
    Expand a prompt with cinematic camera, lighting, motion, and visual-quality details.

    Rich prompts are preserved, while thinner prompts receive deterministic enhancements
    without changing their subject.

    Returns:
        str: The original or expanded prompt with normalized whitespace and punctuation.
    """
    text = _normalize_user_prompt(prompt)
    if _already_rich(text):
        return text

    lower = text.lower()
    parts = [text.rstrip(".")]

    if not any(c in lower for c in _CAMERA_CUES):
        parts.append("shot on a 35mm anamorphic lens, gentle handheld micro-movement, "
                     "shallow depth of field")
    if not any(c in lower for c in _LIGHT_CUES):
        parts.append("natural cinematic lighting with soft volumetric haze and subtle "
                     "rim light separating subject from background")
    if not any(c in lower for c in _MOTION_CUES):
        parts.append("smooth continuous motion with grounded physics")

    parts.append("highly detailed, coherent temporal continuity, film grain, "
                 "color graded like a contemporary drama")
    enhanced = ", ".join(parts)
    # Single trailing period; collapse duplicate whitespace.
    enhanced = re.sub(r"\s+", " ", enhanced).strip()
    if not enhanced.endswith("."):
        enhanced += "."
    return enhanced

def enhance_result_as_metrics(result: EnhanceResult | None) -> dict[str, Any]:
    """
    Convert prompt enhancement results into metrics fields.

    Parameters:
        result (EnhanceResult | None): The enhancement result, or `None` when no enhancement was performed.

    Returns:
        dict[str, Any]: A metrics mapping containing enhancement status, backend metadata, timing, and original and enhanced prompts.
    """
    if result is None:
        return {
            "enhance_prompt": False,
            "enhance_backend": None,
            "enhance_model": None,
            "enhance_elapsed_s": None,
            "prompt_original": None,
            "prompt_enhanced": None,
        }
    return {
        "enhance_prompt": True,
        "enhance_backend": result.backend,
        "enhance_model": result.model,
        "enhance_elapsed_s": result.elapsed_s,
        "prompt_original": result.original,
        "prompt_enhanced": result.enhanced,
    }

def ensure_quantization_supported(spec: MLXQuantizationSpec | None) -> None:
    """Raise :class:`UnsupportedMLXQuantizationError` if ``spec`` cannot run here."""
    if spec is None:
        return
    error = quantization_support_error(spec)
    if error is None:
        return
    import mlx.core as mx

    mlx_version = getattr(mx, "__version__", "unknown")
    raise UnsupportedMLXQuantizationError(f"MLX quantization mode '{spec.label}' is not supported by the installed mlx "
                                          f"({mlx_version}): {error}. Upgrade mlx or pick a supported mode "
                                          f"(int8 is currently the most reliable quality/memory target).")

def fastwan_shape(
        *,
        height: int,
        width: int,
        num_frames: int,
        vae_temporal_compression: int = 4,
        vae_spatial_compression: int = 8,
        patch_size: tuple[int, int, int] = (1, 2, 2),
        num_heads: int = 12,
        head_dim: int = 128,
) -> FastWanShape:
    """Return the approximate DiT token shape for Wan/FastWan T2V inference."""
    latent_frames = (num_frames - 1) // vae_temporal_compression + 1
    latent_height = height // vae_spatial_compression
    latent_width = width // vae_spatial_compression
    patch_frames = latent_frames // patch_size[0]
    patch_height = latent_height // patch_size[1]
    patch_width = latent_width // patch_size[2]
    tokens = patch_frames * patch_height * patch_width
    return FastWanShape(
        height=height,
        width=width,
        num_frames=num_frames,
        latent_frames=latent_frames,
        latent_height=latent_height,
        latent_width=latent_width,
        patch_frames=patch_frames,
        patch_height=patch_height,
        patch_width=patch_width,
        tokens=tokens,
        hidden_size=num_heads * head_dim,
        num_heads=num_heads,
        head_dim=head_dim,
    )

def gib_to_bytes(value: float | None) -> int | None:
    """
    Convert a positive memory limit from GiB to bytes.

    Parameters:
        value (float | None): Memory limit in GiB, or `None` when unset.

    Returns:
        int | None: The memory limit in bytes, or `None` when no limit is provided.

    Raises:
        ValueError: If `value` is zero or negative.
    """
    if value is None:
        return None
    if value <= 0:
        raise ValueError(f"Memory limit must be positive GiB, got {value}")
    return int(value * GIB)

def load_mlx_dit_checkpoint(checkpoint_dir: str | Path, *, compile: bool = False) -> MLXWanDiT:
    """
    Reconstruct an MLXWanDiT model from a versioned checkpoint.

    Parameters:
        checkpoint_dir (str | Path): Directory containing the checkpoint manifest and weights.
        compile (bool): Whether to configure the reconstructed model for compilation.

    Returns:
        MLXWanDiT: The reconstructed model.

    Raises:
        FileNotFoundError: If the checkpoint manifest or weights file is missing.
        ValueError: If the checkpoint format is unsupported or block weights are incomplete.
    """
    import mlx.core as mx

    checkpoint_dir = Path(checkpoint_dir)
    manifest_path = checkpoint_dir / MANIFEST_FILENAME
    weights_path = checkpoint_dir / WEIGHTS_FILENAME
    if not manifest_path.exists() or not weights_path.exists():
        raise FileNotFoundError(f"Not an MLX DiT checkpoint directory: {checkpoint_dir} "
                                f"(expected {MANIFEST_FILENAME} and {WEIGHTS_FILENAME}).")

    manifest = json.loads(manifest_path.read_text())
    version = manifest.get("format_version")
    if version != FORMAT_VERSION:
        raise ValueError(f"MLX DiT checkpoint {checkpoint_dir} has format_version={version}; "
                         f"this FastVideo build reads version {FORMAT_VERSION}. Re-export the checkpoint.")

    spec = None
    if manifest["quantization"] is not None:
        spec = MLXQuantizationSpec(**manifest["quantization"])
        # The packed layout of mx.quantize output is mode-specific, so a build
        # that cannot run the mode cannot use these arrays at all.
        ensure_quantization_supported(spec)

    arrays = mx.load(str(weights_path))
    quantized_keys: dict[str, dict[str, Any]] = manifest["quantized_keys"]

    def rebuild(key: str):
        """
        Reconstructs a weight array or quantized matrix from checkpoint data.

        Parameters:
            key (str): The weight key to rebuild.

        Returns:
            The stored array for an unquantized weight or a reconstructed quantized matrix.
        """
        if key not in quantized_keys:
            return arrays[key]
        info = quantized_keys[key]
        assert spec is not None, f"Quantized key '{key}' in a checkpoint without a quantization spec"
        return QuantizedMatrix(
            weight=arrays[key],
            scales=arrays[f"{key}.scales"],
            biases=arrays[f"{key}.biases"] if info["has_biases"] else None,
            spec=spec,
            dequantized_dtype=_name_to_dtype(info["dequantized_dtype"]),
        )

    config = manifest["config"]
    block_keys: dict[int, list[str]] = {}
    top_level_keys: list[str] = []
    for key in arrays:
        if key.endswith(".scales") or key.endswith(".biases"):
            continue
        if key.startswith(f"{_BLOCK_PREFIX}."):
            index_str, _, _ = key[len(_BLOCK_PREFIX) + 1:].partition(".")
            block_keys.setdefault(int(index_str), []).append(key)
        else:
            top_level_keys.append(key)

    weights = {key: rebuild(key) for key in top_level_keys}

    num_blocks = int(manifest["num_blocks"])
    if sorted(block_keys) != list(range(num_blocks)):
        raise ValueError(f"MLX DiT checkpoint {checkpoint_dir} is missing block weights: "
                         f"manifest says {num_blocks} blocks, found indices {sorted(block_keys)}.")

    inner_dim = int(config["num_attention_heads"]) * int(config["attention_head_dim"])
    blocks = []
    for index in range(num_blocks):
        prefix = f"{_BLOCK_PREFIX}.{index}."
        block_weights = {key[len(prefix):]: rebuild(key) for key in block_keys[index]}
        blocks.append(
            MLXWanTransformerBlock(
                block_weights,
                dim=inner_dim,
                ffn_dim=int(config["ffn_dim"]),
                num_heads=int(config["num_attention_heads"]),
                eps=float(config["eps"]),
            ))
    return MLXWanDiT(weights, blocks, config, compile=compile)

def load_or_enhance_prompt(
    prompt: str,
    *,
    backend: str = "auto",
    model: str | None = None,
    system_prompt: str = DEFAULT_ENHANCE_SYSTEM_PROMPT,
    max_tokens: int = 128,
    cache: bool = True,
    cache_dir: Path | None = None,
) -> EnhanceResult:
    """
    Enhance a prompt, reusing a cached result when available.

    Parameters:
        prompt (str): The prompt to enhance.
        backend (str): Enhancement backend to use.
        model (str | None): Optional model identifier.
        system_prompt (str): System prompt for model-based enhancement.
        max_tokens (int): Maximum number of tokens generated by the model.
        cache (bool): Whether to read and write the on-disk cache.
        cache_dir (Path | None): Optional directory for cached results.

    Returns:
        EnhanceResult: The enhanced prompt and backend metadata. Cached results are marked with the ``"cache"`` backend.
    """
    text = _normalize_user_prompt(prompt)
    path = enhance_cache_path(text, backend=backend, model=model, cache_dir=cache_dir)
    if cache and path.is_file():
        try:
            payload = json.loads(path.read_text())
            return EnhanceResult(
                original=str(payload.get("original", text)),
                enhanced=str(payload["enhanced"]),
                # Mark cache hits explicitly so metrics/logs can distinguish
                # a free replay from a fresh template/mlx-lm call.
                backend="cache",
                elapsed_s=0.0,
                model=payload.get("model"),
            )
        except (OSError, KeyError, json.JSONDecodeError):
            pass

    result = enhance_prompt(
        text,
        backend=backend,
        model=model,
        system_prompt=system_prompt,
        max_tokens=max_tokens,
    )
    if cache:
        try:
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(
                json.dumps(
                    {
                        "original": result.original,
                        "enhanced": result.enhanced,
                        "backend": result.backend,
                        "model": result.model,
                    },
                    indent=2,
                ))
        except OSError as exc:  # pragma: no cover - cache is best-effort
            logger.info("[MLX enhance] cache write skipped: %s", exc)
    return result

def mlx_block_weights_from_diffusers_safetensors(
    checkpoint_path: str | Path,
    *,
    block_index: int = 0,
    quantization: str | MLXQuantizationSpec | None = None,
    dtype=None,
) -> dict[str, mx.array]:
    """Load one Diffusers-format Wan block into the MLX dense-block key layout."""
    from safetensors import safe_open

    prefix = f"blocks.{block_index}."
    key_map = _WAN_BLOCK_KEY_MAP

    spec = MLXQuantizationSpec.from_name(quantization) if (quantization is None
                                                           or isinstance(quantization, str)) else quantization
    ensure_quantization_supported(spec)
    matrix_targets = {target for target in key_map.values() if target.endswith(".weight") and "norm" not in target}
    weights = {}
    with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle:
        available = set(handle.keys())
        for source_name, target_name in key_map.items():
            full = prefix + source_name
            if full not in available:
                # Biases are optional: e.g. Wan2.1-14B has bias-free attention/FFN.
                # The block forward already fetches biases via ``.get(...)``.
                if source_name.endswith(".bias"):
                    continue
                raise KeyError(f"missing required block weight: {full}")
            array = _load_mx_array_from_safetensor(handle, full, dtype)
            loaded = quantize_matrix(array, spec) if target_name in matrix_targets else array
            _eval_loaded_weight(loaded)
            weights[target_name] = loaded
            del array
    return weights

def plan_fast_spatial(
    *,
    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),
    upsample_mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE,
    sharpen: float = DEFAULT_FAST_SPATIAL_SHARPEN,
    enabled: bool = True,
) -> FastSpatialPlan:
    """
    Build a plan for reduced-resolution denoising followed by pixel-space upsampling.

    Parameters:
        upsample_mode (str): Pixel interpolation kernel, one of
            :data:`~fastvideo.mlx_runtime.frame_upsample.PIXEL_UPSAMPLE_MODES`.
        sharpen (float): Unsharp strength applied after the resize.

    Returns:
        FastSpatialPlan: The validated spatial-fast processing plan.

    Raises:
        ValueError: If the upsample mode is unsupported or ``sharpen`` is negative.
    """
    if upsample_mode not in PIXEL_UPSAMPLE_MODES:
        raise ValueError(f"Unsupported upsample mode: {upsample_mode!r} "
                         f"(expected one of {', '.join(PIXEL_UPSAMPLE_MODES)})")
    if sharpen < 0.0:
        raise ValueError(f"sharpen must be >= 0, got {sharpen}")
    plan = plan_refine_resolutions(
        height=height,
        width=width,
        num_frames=num_frames,
        spatial_scale=spatial_scale,
        vae_spatial_compression=vae_spatial_compression,
        vae_temporal_compression=vae_temporal_compression,
        patch_size=patch_size,
        enabled=enabled,
        mode_label="fast-spatial",
    )
    if plan.spatial_scale > 1:
        logger.info(
            "[MLX fast-spatial] denoise+decode %dx%d → upsample %dx to %dx%d (%s, sharpen=%.2f)",
            plan.stage1_width,
            plan.stage1_height,
            plan.spatial_scale,
            plan.target_width,
            plan.target_height,
            upsample_mode,
            sharpen,
        )
    return FastSpatialPlan(plan=plan, upsample_mode=upsample_mode, sharpen=sharpen)