app.py

try:
    import spaces  # type: ignore
except ImportError:  # Local runs do not require the HF Spaces helper package.
    class _SpacesFallback:
        @staticmethod
        def GPU(func=None, **_kwargs):
            if func is None:
                return lambda wrapped: wrapped
            return func

    spaces = _SpacesFallback()

import argparse
import os
import time
from typing import Optional, Tuple

import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageOps

from transnormal2 import TransNormal2Pipeline


DEFAULT_BASE_MODEL = "black-forest-labs/FLUX.2-klein-base-9B"
DEFAULT_WEIGHTS = "Longxiang-ai/TransNormal-2"
IMAGE_EXAMPLES = [
    ["examples/input/houses_unsplash.jpg", "Opaque", 768, False],
    ["examples/input/glass_vase.jpg", "Transparent", 768, False],
]

PIPELINE = None
APP_ARGS = None
LOAD_ERROR: Optional[str] = None


def parse_args():
    parser = argparse.ArgumentParser(description="TransNormal-2 Gradio demo")
    parser.add_argument("--weights", type=str, default=os.getenv("TRANSNORMAL2_WEIGHTS", DEFAULT_WEIGHTS))
    parser.add_argument("--base_model", type=str, default=os.getenv("TRANSNORMAL2_BASE_MODEL", DEFAULT_BASE_MODEL))
    parser.add_argument("--server_name", type=str, default=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"))
    parser.add_argument("--server_port", type=int, default=int(os.getenv("GRADIO_SERVER_PORT", "7860")))
    parser.add_argument("--share", action="store_true", default=os.getenv("GRADIO_SHARE", "0") == "1")
    parser.add_argument("--cpu_offload", action="store_true", default=os.getenv("TRANSNORMAL2_CPU_OFFLOAD", "0") == "1")
    parser.add_argument("--preload", action="store_true", default=os.getenv("TRANSNORMAL2_PRELOAD", "1") == "1")
    parser.add_argument("--device", type=str, default=os.getenv("TRANSNORMAL2_DEVICE", "cuda"))
    parser.add_argument("--dtype", choices=["bf16", "fp32"], default=os.getenv("TRANSNORMAL2_DTYPE", "bf16"))
    return parser.parse_args()


def resolve_weights(weights: str) -> str:
    if os.path.isdir(weights):
        return weights
    from huggingface_hub import snapshot_download

    return snapshot_download(repo_id=weights, token=os.getenv("HF_TOKEN"))


def get_pipeline() -> TransNormal2Pipeline:
    global PIPELINE
    if PIPELINE is not None:
        return PIPELINE
    if APP_ARGS is None:
        raise RuntimeError("Application arguments are not initialized.")

    dtype = torch.bfloat16 if APP_ARGS.dtype == "bf16" else torch.float32
    weights_dir = resolve_weights(APP_ARGS.weights)
    device = None if APP_ARGS.cpu_offload else APP_ARGS.device

    pipe = TransNormal2Pipeline.from_pretrained_transnormal2(
        base_model=APP_ARGS.base_model,
        weights_dir=weights_dir,
        torch_dtype=dtype,
        device=device,
    )
    pipe.set_progress_bar_config(disable=True)
    if APP_ARGS.cpu_offload:
        pipe.enable_model_cpu_offload(device=APP_ARGS.device)
        pipe.local_continuity_module.to(APP_ARGS.device)
        pipe.gsm.to(APP_ARGS.device)

    PIPELINE = pipe
    return PIPELINE


def pil_to_tensor(image: Image.Image) -> torch.Tensor:
    image = ImageOps.exif_transpose(image).convert("RGB")
    arr = np.asarray(image).astype(np.float32)
    tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
    return tensor / 127.5 - 1.0


def normal_to_pil(normal: torch.Tensor) -> Image.Image:
    arr = normal.detach().float().clamp(0, 1)[0].permute(1, 2, 0).cpu().numpy()
    return Image.fromarray((arr * 255).round().astype(np.uint8))


def process_res_from_choice(choice) -> Optional[int]:
    if isinstance(choice, str) and choice.lower() == "native":
        return None
    return int(choice)


def is_transparent_mode(mode: str) -> bool:
    return mode == "Transparent"


def _gpu_kwargs():
    kwargs = {
        "duration": int(os.getenv("TRANSNORMAL2_GPU_DURATION", "240")),
    }
    size = os.getenv("TRANSNORMAL2_ZERO_GPU_SIZE", "xlarge")
    if size:
        kwargs["size"] = size
    return kwargs


def _zero_gpu_decorator(func):
    kwargs = _gpu_kwargs()
    try:
        return spaces.GPU(**kwargs)(func)
    except TypeError:
        kwargs.pop("size", None)
        return spaces.GPU(**kwargs)(func)


@_zero_gpu_decorator
def predict(
    image: Image.Image,
    scene_mode: str,
    max_edge,
    show_raw_prediction: bool,
) -> Tuple[Optional[Image.Image], Optional[Image.Image], str]:
    if image is None:
        return None, None, "Upload an RGB image first."

    start = time.time()
    pipe = get_pipeline()
    image_tensor = pil_to_tensor(image)
    process_res = process_res_from_choice(max_edge)
    domain_is_transparent = is_transparent_mode(scene_mode)

    normal = pipe(
        image_tensor,
        domain_is_transparent=domain_is_transparent,
        process_res=process_res,
        apply_gsm=True,
        output_type="pt",
    )
    final_image = normal_to_pil(normal)

    raw_image = None
    if show_raw_prediction:
        raw_normal = pipe(
            image_tensor,
            domain_is_transparent=domain_is_transparent,
            process_res=process_res,
            apply_gsm=False,
            output_type="pt",
        )
        raw_image = normal_to_pil(raw_normal)

    if APP_ARGS and APP_ARGS.device.startswith("cuda") and torch.cuda.is_available():
        torch.cuda.synchronize()

    mode_note = scene_mode
    if scene_mode == "Auto":
        mode_note = "Auto (uses opaque anchoring by default; choose Transparent for glass/liquids)"
    elapsed = time.time() - start
    status = f"Done in {elapsed:.2f}s. Mode: {mode_note}. Output is camera-space normal encoded as RGB."
    return final_image, raw_image, status


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="TransNormal-2 Demo") as demo:
        gr.Markdown(
            """
            # TransNormal-2
            Upload an RGB image to estimate a camera-space surface normal map.
            Use **Transparent** for glass, liquids, or transparent-object scenes.
            """
        )
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(label="Input RGB", type="pil")
                scene_mode = gr.Radio(
                    choices=["Opaque", "Transparent", "Auto"],
                    value="Opaque",
                    label="Scene mode",
                )
                max_edge = gr.Dropdown(
                    choices=[576, 768, 1024, "Native"],
                    value=768,
                    label="Max processing edge",
                )
                show_raw = gr.Checkbox(
                    value=False,
                    label="Also show raw prediction before GSM",
                )
                run_btn = gr.Button("Run TransNormal-2", variant="primary")
            with gr.Column(scale=1):
                final_output = gr.Image(label="Predicted normal (GSM refined)", type="pil")
                raw_output = gr.Image(label="Raw prediction before GSM", type="pil")
                status = gr.Textbox(label="Status", interactive=False)

        gr.Examples(
            examples=[e for e in IMAGE_EXAMPLES if os.path.exists(e[0])],
            inputs=[input_image, scene_mode, max_edge, show_raw],
        )
        gr.Markdown(
            """
            The normal map is encoded as `(n + 1) / 2`: X in red, Y in green, Z in blue.
            First launch on Hugging Face Spaces may take several minutes because FLUX.2 [klein] is large.
            """
        )
        run_btn.click(
            fn=predict,
            inputs=[input_image, scene_mode, max_edge, show_raw],
            outputs=[final_output, raw_output, status],
        )
    return demo


def main():
    global APP_ARGS, LOAD_ERROR
    APP_ARGS = parse_args()
    if APP_ARGS.preload:
        try:
            get_pipeline()
        except Exception as exc:  # Keep the UI visible and retry on first request.
            LOAD_ERROR = str(exc)
            print(f"[TransNormal-2 demo] Preload failed, will retry on demand: {LOAD_ERROR}")

    demo = build_demo()
    demo.queue(max_size=8).launch(
        server_name=APP_ARGS.server_name,
        server_port=APP_ARGS.server_port,
        share=APP_ARGS.share,
    )


if __name__ == "__main__":
    main()