app.py

"""Public, revision-pinned general + LoRA Image-to-Image backend for hy3d.dev."""

from __future__ import annotations

import base64
import gc
import logging
import math
import os
import random
import threading
from io import BytesIO

import gradio as gr
import spaces
import torch
from fastapi import HTTPException
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
from gradio import Server
from PIL import Image, ImageOps

from i2i_contract import (
    ALLOWED_MIME_FORMATS,
    API_VERSION,
    BASE_EDIT_MODE_ID,
    CONTENT_MODERATION_ENABLED,
    ContractError,
    MAX_GUIDANCE,
    MAX_IMAGE_BYTES,
    MAX_IMAGE_PIXELS,
    MAX_INPUT_IMAGES,
    MAX_OUTPUT_SIDE,
    MAX_PROMPT_CHARS,
    MAX_REQUEST_JSON_CHARS,
    MAX_SEED,
    MAX_STEPS,
    MAX_TOTAL_IMAGE_BYTES,
    MAX_TOTAL_IMAGE_PIXELS,
    MIN_GUIDANCE,
    MIN_INPUT_IMAGES,
    MIN_IMAGE_SIDE,
    MIN_STEPS,
    SingleResidentAdapterManager,
    ValidatedEditRequest,
    decode_validated_images,
    execute_validated_edit,
    output_dimensions,
    parse_example_index,
    public_error_text,
    serialize_png_candidate,
)
from provenance import (
    BASE_MODEL_ID,
    BASE_MODEL_REVISION,
    ENABLED_ADAPTER_SPECS,
    TRANSFORMER_MODEL_ID,
    TRANSFORMER_MODEL_REVISION,
    UPSTREAM_SPACE_ID,
    UPSTREAM_SPACE_REVISION,
)


LOGGER = logging.getLogger("hy3dlab.image_to_image")
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s %(message)s")

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.bfloat16
PIPELINE_LOCK = threading.Lock()
LANCZOS = getattr(Image, "Resampling", Image).LANCZOS

NEGATIVE_PROMPT = (
    "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, "
    "extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"
)

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
EXAMPLES_DIR = os.path.realpath(os.path.join(BASE_DIR, "examples"))

# Fixed, index-addressed manifest copied from the pinned upstream Space. The
# Upscaler text is deliberately accurate for this API's 1024px output limit.
EXAMPLES_CONFIG = (
    {"images": ("examples/B.jpg",), "prompt": "Transform into anime.", "lora": "Photo-to-Anime"},
    {"images": ("examples/HRP.jpg",), "prompt": "Transform into a hyper-realistic face portrait.", "lora": "Hyper-Realistic-Portrait"},
    {"images": ("examples/A.jpeg",), "prompt": "Rotate the camera 45 degrees to the right.", "lora": "Multiple-Angles"},
    {"images": ("examples/U.jpg",), "prompt": "Upscale and enhance image detail.", "lora": "Upscaler"},
    {"images": ("examples/L1.jpg", "examples/L2.jpg"), "prompt": "Apply the lighting from image 2 to image 1.", "lora": "Any-light"},
    {"images": ("examples/PP1.jpg",), "prompt": "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed preserving realistic texture and details.", "lora": "Polaroid-Photo"},
    {"images": ("examples/Z1.jpg",), "prompt": "Front-right quarter view.", "lora": "Fal-Multiple-Angles"},
    {"images": ("examples/URP.jpg",), "prompt": "Transform into a cinematic flat log.", "lora": "Cinematic-FlatLog"},
    {"images": ("examples/SL.jpg",), "prompt": "Neutral uniform lighting. Preserve identity and composition.", "lora": "Studio-DeLight"},
    {"images": ("examples/PI.jpg",), "prompt": "Transform it into Pixar-inspired 3D.", "lora": "Pixar-Inspired-3D"},
    {"images": ("examples/MT.jpg",), "prompt": "Paint with manga tone.", "lora": "Manga-Tone"},
    {"images": ("examples/NCB.jpg",), "prompt": "Transform into a noir comic book style.", "lora": "Noir-Comic-Book"},
    {"images": ("examples/URP.jpg",), "prompt": "Ultra-realistic portrait.", "lora": "Ultra-Realistic-Portrait"},
    {"images": ("examples/MN.jpg",), "prompt": "Transform into Midnight Noir Eyes Spotlight.", "lora": "Midnight-Noir-Eyes-Spotlight"},
    {"images": ("examples/ST1.jpg", "examples/ST2.jpg"), "prompt": "Convert Image 1 to the style of Image 2.", "lora": "Style-Transfer"},
    {"images": ("examples/R1.jpg",), "prompt": "Change the picture to realistic photograph.", "lora": "Anything2Real"},
    {"images": ("examples/UA.jpeg",), "prompt": "Unblur and upscale.", "lora": "Unblur-Anything"},
    {"images": ("examples/L1.jpg", "examples/L2.jpg"), "prompt": "Refer to the color tone, remove the original lighting from Image 1, and relight Image 1 based on the lighting and color tone of Image 2.", "lora": "Light-Migration"},
    {"images": ("examples/P1.jpg",), "prompt": "Transform into anime (while preserving the background and remaining elements maintaining realism and original details.)", "lora": "Anime-V2"},
)

if len(EXAMPLES_CONFIG) != 19 or any(example["lora"] not in ENABLED_ADAPTER_SPECS for example in EXAMPLES_CONFIG):
    raise RuntimeError("The fixed example manifest must reference all approved runtime assets")

EXAMPLE_ASSET_PATHS = frozenset(
    relative_path
    for example in EXAMPLES_CONFIG
    for relative_path in example["images"]
)


def _fixed_example_path(relative_path: str) -> str:
    """Resolve only paths present in the static example manifest."""

    if relative_path not in EXAMPLE_ASSET_PATHS:
        raise ValueError("Unknown example asset")
    resolved = os.path.realpath(os.path.join(BASE_DIR, relative_path))
    if os.path.commonpath((EXAMPLES_DIR, resolved)) != EXAMPLES_DIR:
        raise ValueError("Invalid example asset path")
    return resolved


def _make_example_thumbnail(relative_path: str, max_side: int = 220) -> str:
    try:
        with Image.open(_fixed_example_path(relative_path)) as source:
            source.load()
            image = ImageOps.exif_transpose(source).convert("RGB")
            image.thumbnail((max_side, max_side), LANCZOS)
            buffer = BytesIO()
            image.save(buffer, format="JPEG", quality=65, optimize=True)
        return f"data:image/jpeg;base64,{base64.b64encode(buffer.getvalue()).decode('ascii')}"
    except (OSError, SyntaxError, ValueError) as exc:
        LOGGER.warning("event=example_thumbnail_failed error_class=%s", type(exc).__name__)
        return ""


def _encode_example_image(relative_path: str) -> str:
    """Encode one fixed example, normalizing it to the public input limits."""

    with Image.open(_fixed_example_path(relative_path)) as source:
        source.load()
        image = ImageOps.exif_transpose(source).convert("RGB")
        if image.width * image.height > MAX_IMAGE_PIXELS:
            scale = math.sqrt(MAX_IMAGE_PIXELS / (image.width * image.height))
            target = (max(MIN_IMAGE_SIDE, int(image.width * scale)), max(MIN_IMAGE_SIDE, int(image.height * scale)))
            image = image.resize(target, LANCZOS)
        buffer = BytesIO()
        image.save(buffer, format="JPEG", quality=90, optimize=True)
    encoded = buffer.getvalue()
    if len(encoded) > MAX_IMAGE_BYTES:
        raise ValueError("Normalized example exceeds the public image limit")
    return f"data:image/jpeg;base64,{base64.b64encode(encoded).decode('ascii')}"


def _build_example_cards() -> list[dict]:
    cards = []
    for index, example in enumerate(EXAMPLES_CONFIG):
        cards.append(
            {
                "idx": index,
                "thumbs": [_make_example_thumbnail(path) for path in example["images"]],
                "n_images": len(example["images"]),
                "lora": example["lora"],
                "prompt": example["prompt"],
            }
        )
    return cards


EXAMPLE_CARDS = _build_example_cards()
EXAMPLE_ASSETS_READY = all(card["thumbs"] and all(card["thumbs"]) for card in EXAMPLE_CARDS)


# Importing spaces before torch and moving the model at module scope follows the
# Hugging Face ZeroGPU model-placement contract. All Hub assets are pinned to
# immutable 40-character revisions in provenance.py.
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel


pipe = QwenImageEditPlusPipeline.from_pretrained(
    BASE_MODEL_ID,
    revision=BASE_MODEL_REVISION,
    transformer=QwenImageTransformer2DModel.from_pretrained(
        TRANSFORMER_MODEL_ID,
        revision=TRANSFORMER_MODEL_REVISION,
        torch_dtype=DTYPE,
        device_map="cuda",
    ),
    torch_dtype=DTYPE,
).to(DEVICE)

try:
    pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
    LOGGER.info("event=attention_processor_ready processor=fa3_or_sdpa")
except Exception as exc:  # The pipeline can run with its default processor.
    LOGGER.warning("event=attention_processor_fallback error_class=%s", type(exc).__name__)


def _load_adapter(edit_mode: str) -> None:
    spec = ENABLED_ADAPTER_SPECS[edit_mode]
    pipe.load_lora_weights(
        spec.repo,
        revision=spec.revision,
        weight_name=spec.weights,
        adapter_name=spec.adapter_name,
    )


def _activate_adapter(edit_mode: str) -> None:
    pipe.set_adapters(
        [ENABLED_ADAPTER_SPECS[edit_mode].adapter_name],
        adapter_weights=[1.0],
    )


def _unload_adapters() -> None:
    pipe.unload_lora_weights()


ADAPTER_MANAGER = SingleResidentAdapterManager(
    adapter_modes=ENABLED_ADAPTER_SPECS,
    load_adapter=_load_adapter,
    activate_adapter=_activate_adapter,
    unload_adapters=_unload_adapters,
)


def _public_gradio_error(code: str) -> gr.Error:
    return gr.Error(public_error_text(code), print_exception=False)


def _is_public_error(error: gr.Error) -> bool:
    message = str(getattr(error, "message", error))
    return message in {
        public_error_text(code)
        for code in (
            "I2I_BAD_REQUEST",
            "I2I_UNSUPPORTED_TYPE",
            "I2I_FILE_TOO_LARGE",
            "I2I_IMAGE_DECODE_FAILED",
            "I2I_IMAGE_DIMENSIONS_INVALID",
            "I2I_PROMPT_REQUIRED",
            "I2I_PROMPT_TOO_LONG",
            "I2I_LORA_NOT_ALLOWED",
            "I2I_PROVIDER_BUSY",
            "I2I_INFERENCE_FAILED",
            "I2I_INVALID_OUTPUT",
        )
    }


def _provider_error_code(error: gr.Error) -> str:
    """Map scheduler/capacity failures without returning their raw text."""

    message = str(getattr(error, "message", error)).lower()
    busy_markers = ("quota", "queue", "busy", "capacity", "no gpu", "rate limit")
    return "I2I_PROVIDER_BUSY" if any(marker in message for marker in busy_markers) else "I2I_INFERENCE_FAILED"


def _clear_cuda_cache() -> None:
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


@spaces.GPU(size="xlarge")
def _run_gpu_edit(request: ValidatedEditRequest) -> tuple[bytes, int]:
    """Activate one allowed mode and run inference under one pipeline lock."""

    _clear_cuda_cache()
    try:
        with PIPELINE_LOCK:
            try:
                ADAPTER_MANAGER.activate_mode(request.edit_mode)
            except ContractError:
                raise
            except Exception as exc:
                LOGGER.warning(
                    "event=adapter_transition_failed edit_mode=%s error_class=%s",
                    request.edit_mode,
                    type(exc).__name__,
                )
                raise _public_gradio_error("I2I_INFERENCE_FAILED") from None

            used_seed = random.randint(0, MAX_SEED) if request.randomize_seed else request.seed
            generator = torch.Generator(device=DEVICE).manual_seed(used_seed)
            width, height = output_dimensions(request.images[0])

            try:
                decoded_images = decode_validated_images(request.images)
                result = pipe(
                    image=decoded_images,
                    prompt=request.prompt,
                    negative_prompt=NEGATIVE_PROMPT,
                    height=height,
                    width=width,
                    num_inference_steps=request.steps,
                    generator=generator,
                    true_cfg_scale=request.guidance_scale,
                )
                if len(result.images) != 1:
                    raise ContractError("I2I_INVALID_OUTPUT")
                return serialize_png_candidate(result.images[0]), used_seed
            except ContractError as exc:
                LOGGER.warning("event=invalid_model_output code=%s", exc.code)
                raise _public_gradio_error("I2I_INVALID_OUTPUT") from None
            except gr.Error:
                raise
            except Exception as exc:
                LOGGER.warning("event=inference_failed error_class=%s", type(exc).__name__)
                raise _public_gradio_error("I2I_INFERENCE_FAILED") from None
    finally:
        _clear_cuda_cache()


app = Server(title="hy3dlab Image-to-Image API", debug=False)


@app.api(name="edit_image", concurrency_limit=1, concurrency_id="i2i-pipeline")
def edit_image(
    images_b64_json: str,
    prompt: str,
    lora_adapter: str,
    seed: int,
    randomize_seed: bool,
    guidance_scale: float,
    steps: int,
) -> dict:
    """Validate on CPU, then call the internal ZeroGPU edit function.

    The public seven-field contract is retained for compatibility. General mode
    uses the reserved ``__base__`` value; 19 provenance-pinned LoRA IDs are also
    accepted. Validation is technical only, with no semantic moderation.
    """

    try:
        return execute_validated_edit(
            images_b64_json=images_b64_json,
            prompt=prompt,
            lora_adapter=lora_adapter,
            seed=seed,
            randomize_seed=randomize_seed,
            guidance_scale=guidance_scale,
            steps=steps,
            allowed_edit_modes={BASE_EDIT_MODE_ID, *ENABLED_ADAPTER_SPECS},
            runner=_run_gpu_edit,
        )
    except ContractError as exc:
        raise _public_gradio_error(exc.code) from None

    except gr.Error as exc:
        if _is_public_error(exc):
            raise
        code = _provider_error_code(exc)
        LOGGER.warning("event=provider_error code=%s error_class=%s", code, type(exc).__name__)
        raise _public_gradio_error(code) from None
    except Exception as exc:
        LOGGER.warning("event=provider_error error_class=%s", type(exc).__name__)
        raise _public_gradio_error("I2I_INFERENCE_FAILED") from None


def _service_metadata() -> dict:
    deploy_sha = os.environ.get("APP_DEPLOY_SHA", "unknown")
    lora_ids = list(ENABLED_ADAPTER_SPECS)
    return {
        "service": "hy3dlab-image-to-image",
        "api_version": API_VERSION,
        "target_commit": deploy_sha,
        "upstream_space": UPSTREAM_SPACE_ID,
        "upstream_revision": UPSTREAM_SPACE_REVISION,
        "base_model_revision": BASE_MODEL_REVISION,
        "transformer_revision": TRANSFORMER_MODEL_REVISION,
        "content_moderation_enabled": CONTENT_MODERATION_ENABLED,
        "edit_mode": "general-and-lora",
        "uses_lora": True,
        "supports_general_mode": True,
        "general_mode_value": BASE_EDIT_MODE_ID,
        "required_lora_adapter_value": None,
        "default_edit_mode": BASE_EDIT_MODE_ID,
        "enabled_loras": lora_ids,
        "enabled_edit_modes": [BASE_EDIT_MODE_ID, *lora_ids],
        "adapter_residency": "single",
        "examples_count": len(EXAMPLE_CARDS),
        "example_assets_ready": EXAMPLE_ASSETS_READY,
        "ready": deploy_sha != "unknown" and EXAMPLE_ASSETS_READY,
        "limits": {
            "input_images": MAX_INPUT_IMAGES,
            "min_input_images": MIN_INPUT_IMAGES,
            "max_input_images": MAX_INPUT_IMAGES,
            "mime_types": sorted(ALLOWED_MIME_FORMATS),
            "max_image_bytes": MAX_IMAGE_BYTES,
            "max_total_image_bytes": MAX_TOTAL_IMAGE_BYTES,
            "max_request_json_chars": MAX_REQUEST_JSON_CHARS,
            "max_image_pixels": MAX_IMAGE_PIXELS,
            "max_total_image_pixels": MAX_TOTAL_IMAGE_PIXELS,
            "min_image_side": MIN_IMAGE_SIDE,
            "max_prompt_chars": MAX_PROMPT_CHARS,
            "seed": [0, MAX_SEED],
            "guidance": [MIN_GUIDANCE, MAX_GUIDANCE],
            "steps": [MIN_STEPS, MAX_STEPS],
            "max_output_side": MAX_OUTPUT_SIDE,
        },
    }


@app.get("/healthz")
def healthz() -> JSONResponse:
    """Non-inference configuration and revision check."""

    return JSONResponse(content=_service_metadata(), headers=NO_STORE_HEADERS)


def _example_error_response() -> dict:
    return {"images": [], "prompt": "", "lora": "", "names": [], "status": "error"}


@app.api(name="load_example", queue=False)
def load_example(idx: float) -> dict:
    """Return one fixed example payload without allocating a GPU."""

    index = parse_example_index(idx, len(EXAMPLES_CONFIG))
    if index is None:
        return _example_error_response()

    example = EXAMPLES_CONFIG[index]
    try:
        images = [_encode_example_image(path) for path in example["images"]]
    except (OSError, SyntaxError, ValueError) as exc:
        LOGGER.warning("event=example_load_failed error_class=%s", type(exc).__name__)
        return _example_error_response()

    if len(images) != len(example["images"]) or not all(images):
        return _example_error_response()
    return {
        "images": images,
        "prompt": example["prompt"],
        "lora": example["lora"],
        "names": [os.path.basename(path) for path in example["images"]],
        "status": "ok",
    }


@app.get("/api/config")
def client_config() -> JSONResponse:
    """Public, non-secret configuration for the first-party frontend."""

    metadata = _service_metadata()
    config = {
        "api_version": metadata["api_version"],
        "content_moderation_enabled": metadata["content_moderation_enabled"],
        "edit_mode": metadata["edit_mode"],
        "uses_lora": metadata["uses_lora"],
        "supports_general_mode": metadata["supports_general_mode"],
        "general_mode_value": metadata["general_mode_value"],
        "required_lora_adapter_value": metadata["required_lora_adapter_value"],
        "default_edit_mode": metadata["default_edit_mode"],
        "enabled_edit_modes": metadata["enabled_edit_modes"],
        "loras": metadata["enabled_loras"],
        "default_lora": None,
        "examples": EXAMPLE_CARDS,
        "limits": metadata["limits"],
    }
    return JSONResponse(content=config, headers=NO_STORE_HEADERS)


NO_STORE_HEADERS = {
    "Cache-Control": "no-store, max-age=0, must-revalidate",
    "Pragma": "no-cache",
    "Expires": "0",
}


@app.get("/", response_class=HTMLResponse)
async def homepage() -> HTMLResponse:
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as file:
        return HTMLResponse(content=file.read(), headers=NO_STORE_HEADERS)


@app.get("/hy3d-assets/{asset_name}", include_in_schema=False)
async def ui_asset(asset_name: str) -> FileResponse:
    """Serve the two allowlisted, revision-controlled browser assets."""

    allowed_assets = {"app-ui.js", "gradio-client-2.3.1.js"}
    if asset_name not in allowed_assets:
        raise HTTPException(status_code=404, detail="Not found")
    asset_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", asset_name)
    return FileResponse(asset_path, media_type="application/javascript", headers=NO_STORE_HEADERS)


if __name__ == "__main__":
    examples_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
    app.launch(
        show_error=False,
        mcp_server=False,
        blocked_paths=[examples_path],
        footer_links=["api"],
    )