app.py

"""Gradio/ZeroGPU entry point for the LayerD image decomposition pipeline."""

from __future__ import annotations

import logging
import math
import os
import shutil
import tempfile
import time
import uuid
from importlib.metadata import PackageNotFoundError, version
from pathlib import Path
from typing import NoReturn

import gradio as gr
import spaces
import torch
from PIL import Image, ImageOps, UnidentifiedImageError


logging.basicConfig(
    level=os.getenv("LAYERD_LOG_LEVEL", "INFO").upper(),
    format="%(asctime)s %(levelname)s %(name)s - %(message)s",
)
LOGGER = logging.getLogger("layerd_api")


def _package_version(package_name: str) -> str:
    """Return a package version for diagnostics without masking the real error."""

    try:
        return version(package_name)
    except PackageNotFoundError:
        return "unknown"


def _env_int(name: str, default: int, *, minimum: int, maximum: int) -> int:
    """Read a bounded integer setting without making a bad env var fatal."""

    raw_value = os.getenv(name)
    if raw_value is None:
        return default

    try:
        value = int(raw_value)
    except ValueError:
        LOGGER.warning("Ignoring invalid %s=%r; using %d", name, raw_value, default)
        return default

    bounded = max(minimum, min(value, maximum))
    if bounded != value:
        LOGGER.warning("Clamped %s=%d to the supported range [%d, %d]", name, value, minimum, maximum)
    return bounded


def _env_float(name: str, default: float, *, minimum: float, maximum: float) -> float:
    """Read a bounded floating-point setting without making startup fragile."""

    raw_value = os.getenv(name)
    if raw_value is None:
        return default

    try:
        value = float(raw_value)
    except ValueError:
        LOGGER.warning("Ignoring invalid %s=%r; using %s", name, raw_value, default)
        return default

    bounded = max(minimum, min(value, maximum))
    if bounded != value:
        LOGGER.warning("Clamped %s=%s to the supported range [%s, %s]", name, value, minimum, maximum)
    return bounded


def _resolve_device() -> str:
    """Select CUDA for Spaces and retain a CPU escape hatch for local runs.

    ZeroGPU exposes a CUDA emulation layer while the module is importing, so
    loading the model on ``cuda`` here is intentional and required for efficient
    placement when a request receives a real GPU.
    """

    requested = os.getenv("LAYERD_DEVICE", "").strip().lower()
    if requested:
        if requested not in {"cpu", "cuda"}:
            raise ValueError("LAYERD_DEVICE must be either 'cpu' or 'cuda'")
        return requested

    # SPACE_ID is set by Hugging Face Spaces.  The configured Space uses
    # ZeroGPU, where torch.cuda.is_available() is not a reliable startup test.
    if os.getenv("SPACE_ID"):
        return "cuda"
    return "cuda" if torch.cuda.is_available() else "cpu"


def _assert_transformers_compatibility() -> None:
    """Fail early with an actionable message if dependency resolution drifts."""

    installed = _package_version("transformers")
    if installed == "unknown":
        return

    try:
        major_version = int(installed.split(".", 1)[0])
    except (TypeError, ValueError):
        LOGGER.warning("Could not parse Transformers version %r", installed)
        return

    if major_version >= 5:
        raise RuntimeError(
            "LayerD's BiRefNet remote model requires transformers<5.0.0; "
            f"found transformers=={installed}. Pin transformers==4.57.5 and rebuild the Space."
        )


_assert_transformers_compatibility()

from layerd import LayerDPipeline  # noqa: E402  (check dependency before importing LayerD)


MODEL_CARD = os.getenv("LAYERD_MATTING_MODEL", "cyberagent/layerd-birefnet")
DEVICE = _resolve_device()
MATTING_SIZE = _env_int("LAYERD_MATTING_SIZE", 1024, minimum=256, maximum=2048)
MAX_ITERATIONS = _env_int("LAYERD_MAX_ITERATIONS", 3, minimum=1, maximum=5)
MAX_SIDE = _env_int("LAYERD_MAX_SIDE", 2048, minimum=256, maximum=4096)
MAX_PIXELS = _env_int("LAYERD_MAX_PIXELS", 4_194_304, minimum=65_536, maximum=16_777_216)
OUTPUT_TTL_SECONDS = _env_float("LAYERD_OUTPUT_TTL_SECONDS", 7_200.0, minimum=300.0, maximum=86_400.0)
OUTPUT_ROOT = Path(os.getenv("LAYERD_OUTPUT_DIR", tempfile.gettempdir())) / "layerd-api"


def _build_pipeline() -> LayerDPipeline:
    LOGGER.info(
        "Loading LayerD model card=%s device=%s matting_size=%sx%s max_iterations=%s",
        MODEL_CARD,
        DEVICE,
        MATTING_SIZE,
        MATTING_SIZE,
        MAX_ITERATIONS,
    )
    pipeline = LayerDPipeline(
        matting_hf_card=MODEL_CARD,
        matting_process_size=(MATTING_SIZE, MATTING_SIZE),
        device=DEVICE,
    )
    LOGGER.info("LayerD model is ready (transformers=%s)", _package_version("transformers"))
    return pipeline


# Load at module level so ZeroGPU can place weights efficiently during startup.
pipeline = _build_pipeline()


def _cleanup_expired_outputs(now: float | None = None) -> None:
    """Remove old request artifacts while preserving files still offered to users."""

    if not OUTPUT_ROOT.exists():
        return

    current_time = time.time() if now is None else now
    cutoff = current_time - OUTPUT_TTL_SECONDS
    for child in OUTPUT_ROOT.iterdir():
        try:
            if child.is_dir() and child.stat().st_mtime < cutoff:
                shutil.rmtree(child)
        except FileNotFoundError:
            # Another queued request may have cleaned the same directory.
            continue
        except OSError:
            LOGGER.warning("Could not clean old output directory %s", child, exc_info=True)


def _prepare_image(input_image: Image.Image) -> Image.Image:
    """Normalize orientation/color and bound work for the shared GPU."""

    if not isinstance(input_image, Image.Image):
        raise ValueError("请上传一张有效的 PNG、JPG 或 WEBP 图片。")

    try:
        image = ImageOps.exif_transpose(input_image)
    except (AttributeError, ValueError, UnidentifiedImageError) as exc:
        raise ValueError("无法读取图片元数据，请重新上传图片。") from exc

    width, height = image.size
    if width < 16 or height < 16:
        raise ValueError("图片尺寸至少需要 16 x 16 像素。")

    # LayerD ultimately works on RGB.  Converting once here also avoids carrying
    # an unexpected palette/alpha mode through the matting and inpainting code.
    image = image.convert("RGB")
    scale = min(
        1.0,
        MAX_SIDE / max(width, height),
        math.sqrt(MAX_PIXELS / (width * height)),
    )
    if scale < 1.0:
        # Floor rather than round so the configured pixel/side limits remain
        # hard upper bounds even for dimensions that land on a fractional pixel.
        resized = (max(16, math.floor(width * scale)), max(16, math.floor(height * scale)))
        LOGGER.info("Resizing input from %sx%s to %sx%s", width, height, resized[0], resized[1])
        image = image.resize(resized, Image.Resampling.LANCZOS)

    return image


def _new_output_dir(request_id: str) -> Path:
    """Create an isolated output directory for one request."""

    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
    output_dir = OUTPUT_ROOT / request_id
    output_dir.mkdir()
    return output_dir


def _raise_gradio_error(message: str) -> NoReturn:
    """Raise a user-visible Gradio error across supported Gradio versions."""

    error_type = getattr(gr, "Error", RuntimeError)
    raise error_type(message)


@spaces.GPU(duration=120)
def process_image(input_image: Image.Image | None) -> tuple[str | None, str | None]:
    """Decompose one image and return downloadable PSD and self-contained SVG files."""

    request_id = uuid.uuid4().hex[:12]
    if input_image is None:
        _raise_gradio_error("请先上传一张图片。")

    output_dir: Path | None = None
    started_at = time.perf_counter()
    try:
        _cleanup_expired_outputs()
        image = _prepare_image(input_image)
        output_dir = _new_output_dir(request_id)

        # LayerD and LaMa already use inference mode internally; the outer guard
        # also covers classification and any future model components.
        with torch.inference_mode():
            result = pipeline(image, max_iterations=MAX_ITERATIONS)

        psd_path = output_dir / "layerd-output.psd"
        svg_path = output_dir / "layerd-output.svg"
        result.save(str(psd_path), format="psd")
        result.save(str(svg_path), format="svg")

        elapsed = time.perf_counter() - started_at
        LOGGER.info(
            "request=%s complete input=%sx%s elements=%s psd=%s svg=%s elapsed=%.2fs",
            request_id,
            image.width,
            image.height,
            len(result.elements),
            psd_path.stat().st_size,
            svg_path.stat().st_size,
            elapsed,
        )
        return str(psd_path), str(svg_path)
    except (ValueError, UnidentifiedImageError) as exc:
        LOGGER.warning("request=%s rejected: %s", request_id, exc)
        if output_dir is not None:
            shutil.rmtree(output_dir, ignore_errors=True)
        _raise_gradio_error(str(exc))
    except Exception:
        LOGGER.exception("request=%s failed", request_id)
        if output_dir is not None:
            shutil.rmtree(output_dir, ignore_errors=True)
        _raise_gradio_error(f"处理失败（请求 ID: {request_id}），请稍后重试。")


demo = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil", label="上传海报/设计图"),
    outputs=[
        gr.File(label="下载 PSD 分层文件"),
        gr.File(label="下载 SVG 矢量文件"),
    ],
    title="LayerD GPU 分层服务",
    description="上传一张海报或设计图，生成可编辑的 PSD 分层文件和自包含 SVG 文件。",
    api_name="process",
)

# Serialize GPU jobs to avoid concurrent LaMa allocations on the shared device.
demo.queue(max_size=8, default_concurrency_limit=1)


if __name__ == "__main__":
    demo.launch()