pixart.py

"""Strict local PixArt-Sigma generation and ROCm benchmarking."""

from __future__ import annotations

from functools import lru_cache
import json
from pathlib import Path
import subprocess
import sys
import threading
from time import perf_counter
from typing import Any


MODEL_ID = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"
MODEL_CARD_PARAMETER_COUNT = 600_000_000
MODEL_LICENSE = "CreativeML Open RAIL++-M"

class RocmUnavailableError(RuntimeError):
    """Raised instead of silently falling back to CPU or an unsupported GPU runtime."""


def _torch() -> Any:
    try:
        import torch
    except ImportError as exc:
        raise RocmUnavailableError(
            "PyTorch is not installed. Install a ROCm PyTorch wheel on a supported AMD/Linux host first."
        ) from exc
    return torch


def runtime_info(*, require_rocm: bool = False) -> dict[str, Any]:
    try:
        torch = _torch()
    except RocmUnavailableError as exc:
        if require_rocm:
            raise
        return {"ready": False, "reason": str(exc)}

    gpu_available = bool(torch.cuda.is_available())
    hip_version = getattr(torch.version, "hip", None)
    is_rocm = gpu_available and bool(hip_version)
    info: dict[str, Any] = {
        "ready": is_rocm,
        "rocm": is_rocm,
        "torch_version": torch.__version__,
        "rocm_version": hip_version,
        "cuda_version": getattr(torch.version, "cuda", None),
        "gpu_available": gpu_available,
        "device_count": torch.cuda.device_count() if gpu_available else 0,
    }
    if gpu_available:
        properties = torch.cuda.get_device_properties(0)
        info.update(
            {
                "device": torch.cuda.get_device_name(0),
                "total_vram_mib": round(properties.total_memory / (1024**2), 2),
            }
        )
    if not info["ready"]:
        info["reason"] = "PixArt generation requires an AMD ROCm PyTorch runtime with a detected ROCm/HIP device."
    if require_rocm and not info["ready"]:
        raise RocmUnavailableError(info["reason"])
    return info


def _disable_broken_apex_t5_layernorm() -> None:
    """Block Apex's T5 LayerNorm patch before transformers can apply it.

    The Radeon Cloud image preinstalls an Apex build whose FusedRMSNorm assumes
    float32 input. transformers auto-swaps T5's LayerNorm for that fused kernel
    whenever apex is importable, which is what previously crashed PixArt's T5 text
    encoder in float16 with "expected scalar type Float but found Half" (see
    RADEON_CLOUD_TROUBLESHOOTING.md). Poisoning sys.modules makes any `import apex...`
    raise ImportError so transformers falls back to its native, dtype-safe LayerNorm,
    letting the pipeline run in the float16 the PixArt-Sigma model card specifies
    instead of the slower/heavier float32 fallback.
    """
    for name in [mod for mod in sys.modules if mod == "apex" or mod.startswith("apex.")]:
        del sys.modules[name]
    sys.modules["apex"] = None  # type: ignore[assignment]


_GPU_LOCK = threading.Lock()  # Serializes generation; concurrent jobs on one GPU risk OOM/thrash.


@lru_cache(maxsize=2)
def _load_pipeline(model_id: str, compile_transformer: bool) -> Any:
    torch = _torch()
    runtime_info(require_rocm=True)
    _disable_broken_apex_t5_layernorm()
    try:
        from diffusers import PixArtSigmaPipeline
    except ImportError as exc:
        raise RocmUnavailableError(
            "diffusers is not installed. Install requirements after the ROCm PyTorch wheel."
        ) from exc

    # Every request generates at a fixed 1024x1024 shape, so MIOpen can safely
    # autotune conv/attention algorithms once and reuse them for the rest of the process.
    torch.backends.cudnn.benchmark = True

    pipeline = PixArtSigmaPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.float16,
        use_safetensors=True,
    )
    pipeline.to("cuda")  # PyTorch addresses ROCm/HIP GPUs through this same device string.
    pipeline.vae.enable_slicing()  # PixArtSigmaPipeline has no enable_vae_slicing() convenience wrapper.
    pipeline.vae.to(memory_format=torch.channels_last)  # VAE decode is conv-heavy; helps MIOpen throughput.
    if compile_transformer:
        pipeline.transformer = torch.compile(pipeline.transformer, mode="reduce-overhead", fullgraph=True)
    return pipeline


def _parameter_counts(pipeline: Any) -> dict[str, int]:
    counts: dict[str, int] = {}
    for component_name in ("transformer", "text_encoder", "vae"):
        component = getattr(pipeline, component_name, None)
        if component is not None and hasattr(component, "parameters"):
            counts[component_name] = sum(int(parameter.numel()) for parameter in component.parameters())
    counts["loaded_total"] = sum(counts.values())
    return counts


def _percentile(values: list[float], percentile: float) -> float:
    ordered = sorted(values)
    if len(ordered) == 1:
        return ordered[0]
    position = (len(ordered) - 1) * percentile
    lower = int(position)
    upper = min(lower + 1, len(ordered) - 1)
    weight = position - lower
    return ordered[lower] * (1.0 - weight) + ordered[upper] * weight


def _run_pipeline(
    pipeline: Any,
    *,
    prompt: str,
    height: int,
    width: int,
    inference_steps: int,
    guidance_scale: float,
    seed: int,
    device: str = "cuda",
) -> Any:
    torch = _torch()
    generator = torch.Generator(device=device).manual_seed(seed)
    with torch.inference_mode():
        return pipeline(
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=inference_steps,
            guidance_scale=guidance_scale,
            generator=generator,
    ).images[0]


def generate_hero_image(
    *,
    prompt: str,
    output_path: Path,
    seed: int,
    height: int = 1024,
    width: int = 1024,
    inference_steps: int = 20,
    guidance_scale: float = 4.5,
    compile_transformer: bool = False,
) -> dict[str, Any]:
    if height % 8 or width % 8:
        raise ValueError("PixArt image width and height must be multiples of 8.")
    pipeline = _load_pipeline(MODEL_ID, compile_transformer)
    torch = _torch()
    with _GPU_LOCK:
        torch.cuda.synchronize()
        started = perf_counter()
        image = _run_pipeline(
            pipeline,
            prompt=prompt,
            height=height,
            width=width,
            inference_steps=inference_steps,
            guidance_scale=guidance_scale,
            seed=seed,
        )
        torch.cuda.synchronize()
        elapsed_ms = (perf_counter() - started) * 1000.0
    output_path.parent.mkdir(parents=True, exist_ok=True)
    image.save(output_path)
    return {
        "model_id": MODEL_ID,
        "model_card_parameter_count": MODEL_CARD_PARAMETER_COUNT,
        "license": MODEL_LICENSE,
        "runtime": runtime_info(require_rocm=True),
        "parameter_counts": _parameter_counts(pipeline),
        "image_path": str(output_path),
        "width": width,
        "height": height,
        "inference_steps": inference_steps,
        "guidance_scale": guidance_scale,
        "seed": seed,
        "latency_ms": round(elapsed_ms, 2),
        "compiled_transformer": compile_transformer,
        "image_label": "AI-generated representative image; not a factual depiction of the venue or event.",
    }


def _rocm_smi_snapshot() -> dict[str, Any]:
    try:
        completed = subprocess.run(
            ["rocm-smi", "--showproductname", "--showmeminfo", "vram", "--json"],
            check=False,
            capture_output=True,
            text=True,
            timeout=10,
        )
    except (FileNotFoundError, subprocess.TimeoutExpired) as exc:
        return {"available": False, "reason": str(exc)}
    if completed.returncode != 0:
        return {"available": False, "reason": completed.stderr.strip() or "rocm-smi failed"}
    try:
        return {"available": True, "data": json.loads(completed.stdout)}
    except json.JSONDecodeError:
        return {"available": True, "raw": completed.stdout.strip()}