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()