src/backends/zerogpu.py

"""ZeroGPU backend for Qwen Image Edit (real multi-angle LoRA edits).

On Hugging Face **ZeroGPU** Spaces a real GPU is attached only for the duration
of a function decorated with ``@spaces.GPU``. This module therefore:

1. Loads the Qwen-Image-Edit-2509 pipeline + Rapid-AIO transformer + dx8152's
   multiple-angles LoRA **at import time** (guarded by the ``SPACES_ZERO_GPU``
   env var so it only happens on an actual ZeroGPU Space).
2. Exposes a module-level ``@spaces.GPU`` inference function so ZeroGPU can
   schedule it on a GPU.

Mirrors ``linoyts/Qwen-Image-Edit-Angles``. On any non-ZeroGPU machine this
module stays inert and ``ZeroGpuQwenBackend.prepare()`` raises, letting
``select_backend`` fall through to the next option.
"""

from __future__ import annotations

import os
import traceback

from PIL import Image

from ..config import (
    ANGLES_LORA_REPO,
    ANGLES_LORA_WEIGHT,
    QWEN_BASE_MODEL,
    QWEN_RAPID_TRANSFORMER,
)
from ..images import fit_image
from .base import ImageEditBackend


def on_zerogpu() -> bool:
    return str(os.environ.get("SPACES_ZERO_GPU", "")).lower() in ("1", "true", "yes")


# Populated at import time when running on ZeroGPU.
_PIPE = None
_LOAD_ERROR: Exception | None = None
_LOAD_ERROR_TB: str = ""
_LOAD_STEP: str = "not started"
run_zero_edit = None  # module-level @spaces.GPU function (or None off-ZeroGPU)
_LORA_SCALE = 1.25


def _load_pipeline():
    """Load and fuse the Qwen Image Edit pipeline. Runs once at import."""
    global _PIPE, _LOAD_ERROR, _LOAD_ERROR_TB, _LOAD_STEP
    try:
        _LOAD_STEP = "import torch"
        import torch

        # Vendored Qwen classes (mirrors linoyts/Qwen-Image-Edit-Angles). These
        # depend on bleeding-edge diffusers internals, so they are shipped in
        # the repo rather than imported from a released ``diffusers``.
        _LOAD_STEP = "import vendored qwenimage classes"
        from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
        from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel

        dtype = torch.bfloat16
        _LOAD_STEP = f"load transformer {QWEN_RAPID_TRANSFORMER}"
        transformer = QwenImageTransformer2DModel.from_pretrained(
            QWEN_RAPID_TRANSFORMER,
            subfolder="transformer",
            torch_dtype=dtype,
            device_map="cuda",
        )
        _LOAD_STEP = f"load pipeline {QWEN_BASE_MODEL}"
        pipe = QwenImageEditPlusPipeline.from_pretrained(
            QWEN_BASE_MODEL,
            transformer=transformer,
            torch_dtype=dtype,
        ).to("cuda")
        _LOAD_STEP = f"load LoRA {ANGLES_LORA_REPO}/{ANGLES_LORA_WEIGHT}"
        pipe.load_lora_weights(
            ANGLES_LORA_REPO,
            weight_name=ANGLES_LORA_WEIGHT,
            adapter_name="angles",
        )
        pipe.set_adapters(["angles"], adapter_weights=[1.0])
        _LOAD_STEP = "fuse LoRA"
        pipe.fuse_lora(adapter_names=["angles"], lora_scale=_LORA_SCALE)
        pipe.unload_lora_weights()

        # Optional AOT-compiled transformer blocks for faster inference.
        try:
            import spaces

            spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/Qwen-Image", variant="fa3")
        except Exception as exc:  # noqa: BLE001 - optimisation only
            print(f"[zerogpu] AOTI blocks not loaded ({exc}); continuing without.")

        _PIPE = pipe
        _LOAD_STEP = "loaded"
    except Exception as exc:  # noqa: BLE001 - surfaced via prepare()
        _LOAD_ERROR = exc
        _LOAD_ERROR_TB = traceback.format_exc()
        print(f"[zerogpu] Pipeline load failed at [{_LOAD_STEP}]: {exc}")
        print(_LOAD_ERROR_TB)


if on_zerogpu():
    try:
        import spaces  # noqa: WPS433

        _load_pipeline()

        @spaces.GPU(duration=60)
        def run_zero_edit(  # noqa: F811 - intentional module-level assignment
            image: Image.Image,
            prompt: str,
            seed: int,
            num_inference_steps: int,
            true_guidance_scale: float,
            width: int,
            height: int,
        ) -> Image.Image:
            import torch

            if _PIPE is None:
                raise RuntimeError("Qwen pipeline unavailable on ZeroGPU.")
            generator = torch.Generator(device="cuda").manual_seed(int(seed))
            return _PIPE(
                image=[image],
                prompt=prompt,
                width=width,
                height=height,
                num_inference_steps=num_inference_steps,
                generator=generator,
                true_cfg_scale=true_guidance_scale,
                num_images_per_prompt=1,
            ).images[0]

    except Exception as exc:  # noqa: BLE001
        _LOAD_ERROR = exc
        _LOAD_ERROR_TB = traceback.format_exc()
        print(f"[zerogpu] spaces unavailable: {exc}")


def diagnostics() -> str:
    """Human-readable status of the ZeroGPU pipeline (for surfacing in the UI)."""
    lines = [
        f"on_zerogpu: {on_zerogpu()}",
        f"SPACES_ZERO_GPU env: {os.environ.get('SPACES_ZERO_GPU')!r}",
        f"pipeline loaded: {_PIPE is not None}",
        f"run_zero_edit ready: {run_zero_edit is not None}",
        f"last load step: {_LOAD_STEP}",
    ]
    if _LOAD_ERROR is not None:
        lines.append(f"load error: {type(_LOAD_ERROR).__name__}: {_LOAD_ERROR}")
        if _LOAD_ERROR_TB:
            lines.append("traceback:\n" + _LOAD_ERROR_TB.strip())
    return "\n".join(lines)


class ZeroGpuQwenBackend(ImageEditBackend):
    """Runs the real Qwen multi-angle pipeline on a ZeroGPU Space."""

    source = "zerogpu_qwen_image_edit"

    def __init__(self, image_size: int) -> None:
        self.image_size = image_size

    def prepare(self) -> None:
        if not on_zerogpu():
            raise RuntimeError("Not running on a ZeroGPU Space.")
        if _LOAD_ERROR is not None:
            raise RuntimeError(f"ZeroGPU pipeline failed to load: {_LOAD_ERROR}")
        if run_zero_edit is None or _PIPE is None:
            raise RuntimeError("ZeroGPU pipeline not initialised.")

    def edit(
        self,
        image: Image.Image,
        prompt: str,
        seed: int,
        num_inference_steps: int,
        true_guidance_scale: float,
    ) -> Image.Image:
        base = fit_image(image.convert("RGB"), self.image_size)
        if not prompt.strip():
            return base
        result = run_zero_edit(
            base,
            prompt,
            int(seed),
            int(num_inference_steps),
            float(true_guidance_scale),
            base.width,
            base.height,
        )
        return fit_image(result.convert("RGB"), self.image_size)