runtime/portable_turbo_runtime.py

"""Portable three-file Mage-Flow XPO3 runtime for standalone editing."""

from __future__ import annotations

from contextlib import contextmanager
import json
from pathlib import Path
from typing import Any, Iterator

from safetensors import safe_open


DEFAULT_BRIDGE_BLOCK_SCALES = {
    0: 2.575,
    1: 4.15,
    2: 5.025,
    3: 8.1,
    4: 8.0,
    5: 6.8,
    6: 6.275,
    7: 5.6,
    8: 6.65,
}


def parse_index_spec(
    value: str | list[int] | tuple[int, ...] | set[int],
    *,
    label: str,
    minimum: int,
    maximum: int,
) -> set[int]:
    """Parse comma-separated indices and inclusive ranges such as ``0-3,7``."""

    if isinstance(value, str):
        text = value.strip()
        if not text:
            raise ValueError(f"{label} must not be empty")
        parsed: set[int] = set()
        for raw_part in text.split(","):
            part = raw_part.strip()
            if not part:
                raise ValueError(f"{label} contains an empty entry")
            if "-" in part:
                pieces = part.split("-")
                if len(pieces) != 2:
                    raise ValueError(f"invalid {label} range: {part!r}")
                try:
                    start, end = (int(piece.strip()) for piece in pieces)
                except ValueError as exc:
                    raise ValueError(
                        f"invalid {label} range: {part!r}"
                    ) from exc
                if end < start:
                    raise ValueError(
                        f"{label} range runs backwards: {part!r}"
                    )
                parsed.update(range(start, end + 1))
            else:
                try:
                    parsed.add(int(part))
                except ValueError as exc:
                    raise ValueError(
                        f"invalid {label} index: {part!r}"
                    ) from exc
    else:
        parsed = {int(item) for item in value}
    if not parsed:
        raise ValueError(f"{label} must select at least one index")
    invalid = sorted(
        item for item in parsed if item < minimum or item > maximum
    )
    if invalid:
        raise ValueError(
            f"{label} indices must be in [{minimum}, {maximum}]; "
            f"got {invalid}"
        )
    return parsed


def transformer_config(checkpoint_path: str | Path) -> dict[str, Any]:
    with safe_open(
        Path(checkpoint_path).resolve(),
        framework="pt",
        device="cpu",
    ) as handle:
        metadata = handle.metadata() or {}
    try:
        config = json.loads(metadata["mage_flow.transformer_config"])
    except (KeyError, json.JSONDecodeError) as exc:
        raise RuntimeError(
            "diffusion model has no valid Mage-Flow transformer config"
        ) from exc
    if not isinstance(config, dict):
        raise RuntimeError("embedded Mage-Flow transformer config is invalid")
    return config


def structure_from_config(config: dict[str, Any]) -> dict[str, Any]:
    metadata_keys = {
        "_class_name",
        "txt_max_length",
        "max_sequence_length",
        "param_dtype",
        "packing",
        "schedule_mode",
        "static_shift",
        "use_time_shift",
        "rope_type",
        "apply_text_rotary_emb",
        "mlp_ratio",
        "depth_single_blocks",
        "theta",
        "qkv_bias",
        "guidance_embed",
        "vec_in_dim",
        "vec_type",
        "time_type",
        "double_block_type",
        "quantization_config",
    }
    return {
        key: value
        for key, value in config.items()
        if key not in metadata_keys
    }


def load_pipeline_from_files(
    *,
    diffusion_model: str | Path,
    text_encoder: str | Path,
    vae: str | Path,
    support_root: str | Path,
    fused_gelu_library: str | Path,
    bridge_up_library: str | Path,
    bridge_down_library: str | Path,
    torch: Any,
) -> tuple[Any, dict[str, Any]]:
    import torch.nn as nn
    from diffusers import FlowMatchEulerDiscreteScheduler
    from fp4_bridge_runtime import install_selected_img_mlp_bridges
    from fused_gelu_up_runtime import install_fused_gelu_up
    from mage_flow.models.mage_flow import MageFlowModel, ModelConfig
    from mage_flow.models.modules._attn_backend import set_attn_backend
    from mage_flow.pipeline import MageFlowPipeline
    from single_file_transformer import (
        load_single_file_native_transformer,
    )
    from text_encoder_variants import load_scaled_fp8_text_encoder

    diffusion_model = Path(diffusion_model).resolve()
    text_encoder = Path(text_encoder).resolve()
    vae = Path(vae).resolve()
    support_root = Path(support_root).resolve()
    fused_gelu_library = Path(fused_gelu_library).resolve()
    bridge_up_library = Path(bridge_up_library).resolve()
    bridge_down_library = Path(bridge_down_library).resolve()
    for label, path in (
        ("diffusion model", diffusion_model),
        ("text encoder", text_encoder),
        ("VAE", vae),
        ("fused GELU library", fused_gelu_library),
        ("FP4 bridge-up library", bridge_up_library),
        ("FP4 bridge-down library", bridge_down_library),
    ):
        if not path.is_file():
            raise RuntimeError(f"{label} is missing: {path}")

    config_data = transformer_config(diffusion_model)
    quantization_config = config_data.get("quantization_config", {})
    if not isinstance(quantization_config, dict):
        raise RuntimeError("diffusion model quantization config is invalid")
    runtime_profile = quantization_config.get("xpo3_runtime_profile", {})
    if not isinstance(runtime_profile, dict):
        raise RuntimeError("diffusion model XPO3 runtime profile is invalid")
    raw_bridge_scales = runtime_profile.get(
        "fp4_bridge_scales",
        DEFAULT_BRIDGE_BLOCK_SCALES,
    )
    if not isinstance(raw_bridge_scales, dict):
        raise RuntimeError("diffusion model bridge scale profile is invalid")
    bridge_block_scales = {
        int(block): float(scale)
        for block, scale in raw_bridge_scales.items()
    }
    fused_streams = tuple(
        str(stream)
        for stream in runtime_profile.get(
            "fused_gelu_streams",
            ("img_mlp", "txt_mlp"),
        )
    )
    transformer_fallback_attention_backend = str(
        runtime_profile.get(
            "transformer_fallback_attention_backend",
            config_data.get("attn_type", "flash2"),
        )
    )
    text_encoder_attention_backend = str(
        runtime_profile.get(
            "text_encoder_attention_backend",
            config_data.get("attn_type", "flash2"),
        )
    )
    structure = structure_from_config(config_data)
    text_support = support_root / "text_encoder"
    scheduler_support = support_root / "scheduler"
    config = ModelConfig(
        vae_path=str(vae),
        txt_enc_path=str(text_support),
        model_structure=structure,
        txt_max_length=int(config_data.get("txt_max_length", 2048)),
        packing=bool(config_data.get("packing", True)),
        static_shift=float(config_data.get("static_shift", 6.0)),
    )

    transformer, transformer_report = (
        load_single_file_native_transformer(
            diffusion_model,
            support_root=support_root,
            device=torch.device("cuda:0"),
        )
    )
    fused_runtime = None
    bridge_runtime = None
    try:
        # Install toggleable fused wrappers first. The bridge wrappers then
        # retain those modules as their exact off-path fallback, allowing
        # fused GELU and bridge routing to be controlled independently.
        fused_runtime, fused_report = install_fused_gelu_up(
            transformer,
            library_path=fused_gelu_library,
            torch=torch,
            stream_names=fused_streams,
        )
        bridge_runtime, bridge_report = install_selected_img_mlp_bridges(
            transformer,
            bridge_up_library_path=bridge_up_library,
            bridge_down_library_path=bridge_down_library,
            block_tensor_scales=bridge_block_scales,
            torch=torch,
            enabled=True,
        )

        model = MageFlowModel.__new__(MageFlowModel)
        nn.Module.__init__(model)
        model.config = config
        set_attn_backend(transformer_fallback_attention_backend)
        model.patch_text_encoder_forward()
        model.vae = model.load_vae()
        model.transformer = transformer
        model.txt_enc, text_report = load_scaled_fp8_text_encoder(
            text_encoder_dir=text_support,
            artifact_path=text_encoder,
            tokenizer_max_length=config.txt_max_length,
            dit_structure=structure,
            use_packed_text_infer=config.packing,
            attn_type=text_encoder_attention_backend,
        )
        model.vae.requires_grad_(False).to(torch.bfloat16)
        model.txt_enc.requires_grad_(False)
        model.eval()
        model.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
            scheduler_support
        )
        # Keep native contexts alive for as long as the pipeline is alive.
        model._mage_fused_gelu_runtime = fused_runtime
        model._xpo3_fused_gelu_runtime = fused_runtime
        model._xpo3_fp4_bridge_runtime = bridge_runtime
    except BaseException as primary_error:
        cleanup_errors: list[str] = []
        for label, runtime in (
            ("fp4 image-MLP bridge", bridge_runtime),
            ("fused GELU-up", fused_runtime),
        ):
            if runtime is None:
                continue
            try:
                runtime.close()
            except BaseException as cleanup_error:
                cleanup_errors.append(
                    f"{label}: {type(cleanup_error).__name__}: "
                    f"{cleanup_error}"
                )
        if cleanup_errors:
            primary_error.add_note(
                "XPO3 loader cleanup also failed: "
                + "; ".join(cleanup_errors)
            )
        raise
    return (
        MageFlowPipeline(model, device="cuda:0"),
        {
            "diffusion_model": str(diffusion_model),
            "text_encoder": str(text_encoder),
            "vae": str(vae),
            "transformer": transformer_report,
            "fused_gelu_up": fused_report,
            "fp4_image_mlp_bridge": bridge_report,
            "text_encoder_load": text_report,
            "runtime_toggle_contract": {
                "fused_gelu_up": True,
                "fp4_image_mlp_bridge": True,
                "accelerated_attention": True,
                "direct_hnd": True,
                "bridge_blocks": sorted(bridge_block_scales),
                "attention_steps": list(
                    runtime_profile.get("attention_steps", [1, 2])
                ),
                "attention_blocks": list(
                    runtime_profile.get(
                        "attention_blocks",
                        range(12),
                    )
                ),
                "profile": runtime_profile,
                "transformer_fallback_attention_backend": (
                    transformer_fallback_attention_backend
                ),
                "text_encoder_attention_backend": (
                    text_encoder_attention_backend
                ),
            },
        },
    )


@contextmanager
def generation_optimization_context(
    *,
    pipe: Any,
    torch: Any,
    enable_fused_gelu_up: bool,
    enable_fp4_bridge: bool,
    bridge_blocks: str | list[int] | tuple[int, ...] | set[int],
    enable_attention_accel: bool,
    enable_direct_hnd: bool,
    attention_steps: str | list[int] | tuple[int, ...] | set[int],
    attention_blocks: str | list[int] | tuple[int, ...] | set[int],
    steps: int,
    static_shift: float,
    cfg: float,
    required_cfg: float,
    expected_steps: int,
) -> Iterator[dict[str, Any]]:
    """Apply one generation's independently configurable optimization policy."""

    from xpo3_attention_runtime import xpo3_attention_runtime

    model = pipe.model
    fused_runtime = getattr(model, "_xpo3_fused_gelu_runtime", None)
    bridge_runtime = getattr(model, "_xpo3_fp4_bridge_runtime", None)
    if fused_runtime is None or bridge_runtime is None:
        raise RuntimeError("XPO3 optimization runtimes were not installed")

    installed_bridge_blocks = set(
        int(value) for value in bridge_runtime.enabled_block_indices
    )
    requested_bridge_blocks = parse_index_spec(
        bridge_blocks,
        label="bridge blocks",
        minimum=0,
        maximum=max(installed_bridge_blocks),
    )
    unknown_bridge_blocks = (
        requested_bridge_blocks - installed_bridge_blocks
    )
    if unknown_bridge_blocks:
        raise ValueError(
            "bridge blocks were not installed: "
            f"{sorted(unknown_bridge_blocks)}"
        )
    requested_attention_steps = parse_index_spec(
        attention_steps,
        label="attention steps",
        minimum=0,
        maximum=int(steps) - 1,
    )
    requested_attention_blocks = parse_index_spec(
        attention_blocks,
        label="attention blocks",
        minimum=0,
        maximum=11,
    )
    previous = {
        "fused_enabled": bool(fused_runtime.enabled),
        "bridge_enabled": bool(bridge_runtime.enabled),
        "bridge_blocks": list(bridge_runtime.enabled_block_indices),
    }
    manifest: dict[str, Any] = {
        "schema_version": "xpo3-runtime-feature-manifest-v1",
        "requested": {
            "fused_gelu_up": bool(enable_fused_gelu_up),
            "fp4_image_mlp_bridge": bool(enable_fp4_bridge),
            "bridge_blocks": sorted(requested_bridge_blocks),
            "accelerated_attention": bool(enable_attention_accel),
            "direct_hnd": bool(enable_direct_hnd),
            "attention_steps": sorted(requested_attention_steps),
            "attention_blocks": sorted(requested_attention_blocks),
        },
        "accelerated_attention": None,
        "fused_gelu_up": None,
        "fp4_image_mlp_bridge": None,
        "restoration": {
            "fused_state_restored": None,
            "bridge_global_state_restored": None,
            "bridge_block_state_restored": None,
            "attention_patches_restored": None,
            "all_restored": None,
        },
    }
    attention_report = None
    primary_error: BaseException | None = None
    try:
        fused_runtime.set_enabled(bool(enable_fused_gelu_up))
        bridge_runtime.set_active_blocks(requested_bridge_blocks)
        bridge_runtime.set_enabled(bool(enable_fp4_bridge))
        bridge_runtime.reset_telemetry()
        with xpo3_attention_runtime(
            pipe=pipe,
            torch=torch,
            enabled=bool(enable_attention_accel),
            direct_hnd=bool(enable_direct_hnd),
            steps=int(steps),
            static_shift=float(static_shift),
            cfg=float(cfg),
            selected_steps=requested_attention_steps,
            selected_blocks=requested_attention_blocks,
            required_cfg=float(required_cfg),
            expected_steps=int(expected_steps),
        ) as attention_report:
            manifest["accelerated_attention"] = attention_report
            manifest["fused_gelu_up"] = {
                "enabled": bool(fused_runtime.enabled),
                "installed_modules": list(
                    fused_runtime.installed_modules
                ),
            }
            manifest["fp4_image_mlp_bridge"] = bridge_runtime.report()
            yield manifest
    except BaseException as error:
        primary_error = error
        raise
    finally:
        restore_errors: list[dict[str, str]] = []

        def attempt_restore(label: str, operation: Any) -> None:
            try:
                operation()
            except BaseException as error:
                restore_errors.append(
                    {
                        "operation": label,
                        "type": type(error).__name__,
                        "message": str(error),
                    }
                )

        attempt_restore(
            "restore_fused_enabled",
            lambda: fused_runtime.set_enabled(previous["fused_enabled"]),
        )
        attempt_restore(
            "restore_bridge_blocks",
            lambda: bridge_runtime.set_active_blocks(
                previous["bridge_blocks"]
            ),
        )
        attempt_restore(
            "restore_bridge_enabled",
            lambda: bridge_runtime.set_enabled(previous["bridge_enabled"]),
        )
        manifest["fused_gelu_up"] = {
            "enabled_during_generation": bool(enable_fused_gelu_up),
            "installed_modules": list(fused_runtime.installed_modules),
        }
        manifest["fp4_image_mlp_bridge"] = {
            **bridge_runtime.report(),
            "enabled_during_generation": bool(enable_fp4_bridge),
            "active_blocks_during_generation": sorted(
                requested_bridge_blocks
            ),
        }
        manifest["accelerated_attention"] = attention_report
        restoration = manifest["restoration"]
        restoration["errors"] = restore_errors
        restoration["fused_state_restored"] = (
            bool(fused_runtime.enabled) == previous["fused_enabled"]
        )
        restoration["bridge_global_state_restored"] = (
            bool(bridge_runtime.enabled) == previous["bridge_enabled"]
        )
        restoration["bridge_block_state_restored"] = (
            list(bridge_runtime.enabled_block_indices)
            == previous["bridge_blocks"]
        )
        restoration["attention_patches_restored"] = (
            True
            if attention_report is None
            else attention_report.get("restoration", {}).get("all_restored")
            in (True, "not_applicable")
        )
        restoration["all_restored"] = all(
            bool(value)
            for key, value in restoration.items()
            if key not in {"all_restored", "errors"}
        ) and not restore_errors
        if restore_errors:
            detail = "; ".join(
                f"{row['operation']}: {row['type']}: {row['message']}"
                for row in restore_errors
            )
            if primary_error is not None:
                primary_error.add_note(
                    "XPO3 feature restoration also failed: " + detail
                )
            else:
                raise RuntimeError(
                    "XPO3 feature restoration failed: " + detail
                )


def close_pipeline_optimization_runtimes(pipe: Any) -> dict[str, Any]:
    """Close the bridge and fused native contexts in dependency-safe order."""

    report: dict[str, Any] = {"attempts": {}, "errors": []}
    for label, attribute in (
        ("fp4_bridge", "_xpo3_fp4_bridge_runtime"),
        ("fused_gelu_up", "_xpo3_fused_gelu_runtime"),
    ):
        runtime = getattr(pipe.model, attribute, None)
        if runtime is None:
            report["attempts"][label] = "not_applicable"
            continue
        try:
            runtime.close()
            report["attempts"][label] = "closed"
        except Exception as exc:  # noqa: BLE001
            report["attempts"][label] = "error"
            report["errors"].append(
                {
                    "runtime": label,
                    "type": type(exc).__name__,
                    "message": str(exc),
                }
            )
    report["all_closed_without_error"] = not report["errors"]
    return report


__all__ = [
    "DEFAULT_BRIDGE_BLOCK_SCALES",
    "close_pipeline_optimization_runtimes",
    "generation_optimization_context",
    "load_pipeline_from_files",
    "parse_index_spec",
    "structure_from_config",
    "transformer_config",
]