Add TransNormal-2 Gradio demo

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/input/houses_unsplash.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,28 @@
+Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0)
+
+Copyright (c) 2026 TransNormal-2 Authors
+
+This work (the TransNormal-2 code and released model weights) is licensed
+under the Creative Commons Attribution-NonCommercial 4.0 International
+License.
+
+You are free to:
+  - Share: copy and redistribute the material in any medium or format
+  - Adapt: remix, transform, and build upon the material
+
+Under the following terms:
+  - Attribution: You must give appropriate credit, provide a link to the
+    license, and indicate if changes were made.
+  - NonCommercial: You may not use the material for commercial purposes.
+
+No additional restrictions: You may not apply legal terms or technological
+measures that legally restrict others from doing anything the license
+permits.
+
+Full license text: https://creativecommons.org/licenses/by-nc/4.0/legalcode
+
+NOTE: The released LoRA / LCM / GSM weights are derivatives of FLUX.2 [klein]
+by Black Forest Labs and are additionally subject to the FLUX [dev]
+Non-Commercial License v2.1. See the NOTICE file. The FLUX.2-klein-base-9B
+base model is NOT distributed by this repository; download it from Black
+Forest Labs under their license terms.

NOTICE

@@ -0,0 +1,22 @@
+NOTICE
+
+The TransNormal-2 released weights (core-predictor LoRA, Local Continuity
+Module, and Geometric Sharpening Module) are "Derivatives" of:
+
+    FLUX.2 [klein] (black-forest-labs/FLUX.2-klein-base-9B)
+    Copyright (c) Black Forest Labs Inc.
+
+and are distributed in accordance with the FLUX [dev] Non-Commercial License
+v2.1 (https://bfl.ai/legal/non-commercial-license):
+
+  1. The weights in this distribution are a MODIFIED derivative work of the
+     FLUX.2 [klein] base model (LoRA adaptation plus auxiliary heads trained
+     for surface normal estimation). They are not the original model released
+     by Black Forest Labs.
+  2. Use is limited to NON-COMMERCIAL purposes, as defined by the FLUX [dev]
+     Non-Commercial License.
+  3. The FLUX.2 [klein] base model itself is not redistributed here. Obtain
+     it directly from Black Forest Labs / Hugging Face under their license.
+
+This repository's own code and weights are otherwise licensed under
+CC BY-NC 4.0 (see LICENSE).

README.md

@@ -1,13 +1,40 @@
 ---
-title: TransNormal 2 Demo
-emoji: 🦀
-colorFrom: blue
-colorTo: green
+title: TransNormal-2 Demo
 sdk: gradio
-sdk_version: 6.18.0
-python_version: '3.13'
+sdk_version: 5.49.1
 app_file: app.py
-pinned: false
+python_version: 3.10.13
+startup_duration_timeout: 1h
+models:
+  - Longxiang-ai/TransNormal-2
+  - black-forest-labs/FLUX.2-klein-base-9B
+tags:
+  - computer-vision
+  - monocular-normal-estimation
+  - surface-normal-estimation
+  - gradio
+  - zerogpu
+license: cc-by-nc-4.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# TransNormal-2 Demo
+
+Interactive Gradio demo for TransNormal-2 surface normal estimation.
+
+Upload an RGB image, choose **Opaque** or **Transparent** scene mode, and run
+the one-step FLUX.2 [klein] predictor followed by the Geometric Sharpening
+Module (GSM).
+
+## Runtime Notes
+
+- This Space is intended to run on Hugging Face ZeroGPU first.
+- Use ZeroGPU `xlarge` if available. The FLUX.2 [klein] 9B base model is large.
+- Set `HF_TOKEN` in Space secrets after accepting the FLUX.2 [klein] license.
+- First launch can take several minutes because the base model and TransNormal-2
+  weights need to be downloaded and cached.
+
+## Links
+
+- Project page: https://longxiang-ai.github.io/TransNormal-2
+- Code and README: https://github.com/longxiang-ai/TransNormal-2
+- Model weights: https://huggingface.co/Longxiang-ai/TransNormal-2

app.py

@@ -0,0 +1,244 @@
+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()

requirements.txt

@@ -0,0 +1,11 @@
+torch>=2.5
+diffusers>=0.37.1
+transformers>=4.55
+peft>=0.17
+safetensors>=0.4
+accelerate>=1.0
+huggingface_hub>=0.34
+numpy
+Pillow
+gradio>=5.0
+spaces

transnormal2/__init__.py

@@ -0,0 +1,16 @@
+"""TransNormal-2: one-step rectified-flow surface normal estimation."""
+
+from .gsm import GSM
+from .lcm import LocalContinuityModule
+from .pipeline import TransNormal2Pipeline
+from .utils import load_image, save_normal_map
+
+__all__ = [
+    "GSM",
+    "LocalContinuityModule",
+    "TransNormal2Pipeline",
+    "load_image",
+    "save_normal_map",
+]
+
+__version__ = "1.0.0"

transnormal2/gsm.py

@@ -0,0 +1,279 @@
+"""Geometric Sharpening Module (GSM).
+
+A lightweight (~0.4M param) edge-aware refinement head applied to the
+VAE-decoded normal map. The VAE encoder/decoder bottleneck (8x downsample)
+inherently blurs sharp geometric edges; the GSM compensates by predicting a
+small residual on top of an explicit anchor:
+
+- opaque scenes use a full-RGB guided-filter anchor (He et al. 2010);
+- transparent scenes use the identity anchor (no RGB-guided smoothing, since
+  RGB edges behind glass do not correspond to geometry edges).
+
+The residual head is zero-initialized at training time, so the module starts
+exactly at the anchor. At inference the residual is scaled by
+``residual_scale`` and clipped to ``residual_clip_deg`` degrees before the
+final L2 normalization.
+
+This file keeps only the inference path of the released model (the
+``classical_anchor`` architecture); training losses and unused architecture
+variants are intentionally not part of the release. State-dict keys are fully
+compatible with the research checkpoints.
+"""
+
+from typing import List, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ResBlock(nn.Module):
+    """Residual conv block with optional dilation."""
+
+    def __init__(self, channels: int, dilation: int = 1):
+        super().__init__()
+        self.conv1 = nn.Conv2d(channels, channels, 3, padding=dilation, dilation=dilation)
+        self.conv2 = nn.Conv2d(channels, channels, 3, padding=dilation, dilation=dilation)
+        self.norm1 = nn.GroupNorm(8, channels)
+        self.norm2 = nn.GroupNorm(8, channels)
+        self.act = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        out = self.act(self.norm1(self.conv1(x)))
+        out = self.norm2(self.conv2(out))
+        return self.act(out + residual)
+
+
+def _build_edge_encoder(edge_channels: int) -> nn.Sequential:
+    return nn.Sequential(
+        nn.Conv2d(3, edge_channels, 3, padding=1),
+        nn.GELU(),
+        nn.Conv2d(edge_channels, edge_channels, 3, padding=1),
+        nn.GELU(),
+    )
+
+
+def _resolve_dilations(num_blocks: int, dilations: Optional[List[int]]) -> List[int]:
+    if dilations is None or len(dilations) == 0:
+        return [1] * num_blocks
+    if len(dilations) != num_blocks:
+        raise ValueError(
+            f"dilations length ({len(dilations)}) must equal num_blocks ({num_blocks})"
+        )
+    return [int(d) for d in dilations]
+
+
+def _box_mean(x: torch.Tensor, radius: int) -> torch.Tensor:
+    kernel = 2 * int(radius) + 1
+    return F.avg_pool2d(
+        x.float(), kernel_size=kernel, stride=1, padding=int(radius),
+        count_include_pad=False,
+    )
+
+
+def guided_filter_torch(
+    coarse_normal: torch.Tensor,
+    rgb: torch.Tensor,
+    radius: int = 4,
+    eps: float = 0.01,
+) -> torch.Tensor:
+    """Full-RGB guided filter anchor for normal maps.
+
+    Follows the multi-channel guided-filter solve used by OpenCV's ximgproc
+    guidedFilter: a local linear model maps RGB guide values to each normal
+    channel, then the filtered normal is L2-normalized.
+    """
+    if coarse_normal.shape[1] != 3 or rgb.shape[1] != 3:
+        raise ValueError("guided_filter_torch expects 3-channel normal and RGB")
+
+    out_dtype = coarse_normal.dtype
+    I = rgb.float()
+    p = coarse_normal.float()
+    B, p_ch, H, W = p.shape
+    i_ch = I.shape[1]
+
+    mean_I = _box_mean(I, radius)
+    mean_p = _box_mean(p, radius)
+
+    cov_Ip = []
+    for pc in range(p_ch):
+        row = []
+        for ic in range(i_ch):
+            mean_Ip = _box_mean(I[:, ic:ic + 1] * p[:, pc:pc + 1], radius)
+            row.append(mean_Ip - mean_I[:, ic:ic + 1] * mean_p[:, pc:pc + 1])
+        cov_Ip.append(torch.cat(row, dim=1))
+    cov_Ip = torch.stack(cov_Ip, dim=1)  # (B, 3 normal, 3 rgb, H, W)
+
+    var_rows = []
+    for i in range(i_ch):
+        cols = []
+        for j in range(i_ch):
+            mean_II = _box_mean(I[:, i:i + 1] * I[:, j:j + 1], radius)
+            cols.append(mean_II - mean_I[:, i:i + 1] * mean_I[:, j:j + 1])
+        var_rows.append(torch.cat(cols, dim=1))
+    var_I = torch.stack(var_rows, dim=1)  # (B, 3 row, 3 col, H, W)
+
+    sigma = var_I.permute(0, 3, 4, 1, 2).reshape(-1, i_ch, i_ch)
+    eye = torch.eye(i_ch, device=sigma.device, dtype=sigma.dtype).unsqueeze(0)
+    sigma = sigma + float(eps) * eye
+    rhs = cov_Ip.permute(0, 3, 4, 2, 1).reshape(-1, i_ch, p_ch)
+
+    coeff = torch.linalg.solve(sigma, rhs)
+    a = coeff.reshape(B, H, W, i_ch, p_ch).permute(0, 4, 3, 1, 2)
+    b = mean_p - (a * mean_I.unsqueeze(1)).sum(dim=2)
+
+    mean_a = _box_mean(a.reshape(B, p_ch * i_ch, H, W), radius)
+    mean_a = mean_a.reshape(B, p_ch, i_ch, H, W)
+    mean_b = _box_mean(b, radius)
+    q = (mean_a * I.unsqueeze(1)).sum(dim=2) + mean_b
+    q = F.normalize(q, p=2, dim=1)
+    return q.to(dtype=out_dtype)
+
+
+def _domain_mask(
+    domain_is_transparent: Optional[Union[bool, torch.Tensor]],
+    coarse_normal: torch.Tensor,
+) -> torch.Tensor:
+    B, _, H, W = coarse_normal.shape
+    if domain_is_transparent is None:
+        value = torch.zeros(B, 1, 1, 1, device=coarse_normal.device, dtype=coarse_normal.dtype)
+    elif isinstance(domain_is_transparent, bool):
+        value = torch.full(
+            (B, 1, 1, 1), float(domain_is_transparent),
+            device=coarse_normal.device, dtype=coarse_normal.dtype,
+        )
+    else:
+        value = domain_is_transparent.to(device=coarse_normal.device, dtype=coarse_normal.dtype)
+        if value.ndim == 0:
+            value = value.view(1, 1, 1, 1).expand(B, 1, 1, 1)
+        elif value.ndim in (1, 2):
+            value = value.view(B, 1, 1, 1)
+        elif value.ndim == 3:
+            value = value.unsqueeze(1)
+        elif value.ndim == 4:
+            value = value[:, :1]
+        else:
+            raise ValueError(f"Unsupported domain flag shape: {tuple(value.shape)}")
+    return value.expand(B, 1, H, W)
+
+
+class GSM(nn.Module):
+    """Geometric Sharpening Module (release inference path).
+
+    Interface: ``refined = gsm(coarse_normal, rgb, domain_is_transparent)``
+    where ``coarse_normal`` is the VAE-decoded normal map in [-1, 1] and
+    ``rgb`` is the input image in [-1, 1], both (B, 3, H, W). The output is a
+    unit-normalized normal map in [-1, 1].
+
+    Args mirror the released checkpoint configuration (see weights
+    ``config.json``); state-dict keys match the research checkpoints.
+    """
+
+    def __init__(
+        self,
+        hidden_channels: int = 64,
+        edge_channels: int = 32,
+        num_blocks: int = 4,
+        use_rgb: bool = True,
+        dilations: Optional[List[int]] = None,
+        anchor_radius: int = 4,
+        anchor_eps: float = 0.01,
+        residual_scale: float = 1.0,
+        residual_clip_deg: float = 0.0,
+        initial_alpha: float = 0.5,
+    ):
+        super().__init__()
+        self.use_rgb = bool(use_rgb)
+        self.dilations = _resolve_dilations(num_blocks, dilations)
+        self.anchor_radius = int(anchor_radius)
+        self.anchor_eps = float(anchor_eps)
+        self.residual_scale = float(residual_scale)
+        self.residual_clip_deg = float(residual_clip_deg)
+
+        if self.use_rgb:
+            self.edge_encoder = _build_edge_encoder(edge_channels)
+            main_in_ch = 10  # coarse normal + anchor normal + RGB + domain flag
+            combined_ch = hidden_channels + edge_channels
+        else:
+            self.edge_encoder = None
+            main_in_ch = 7  # coarse normal + anchor normal + domain flag
+            combined_ch = hidden_channels
+
+        self.main_encoder = nn.Sequential(
+            nn.Conv2d(main_in_ch, hidden_channels, 3, padding=1),
+            nn.GELU(),
+        )
+
+        self.blocks = nn.ModuleList()
+        self.projections = nn.ModuleList()
+        for d in self.dilations:
+            self.projections.append(nn.Conv2d(combined_ch, hidden_channels, 1))
+            self.blocks.append(ResBlock(hidden_channels, dilation=d))
+
+        self.feature_proj = nn.Sequential(
+            nn.Conv2d(combined_ch, hidden_channels, 3, padding=1),
+            nn.GELU(),
+        )
+        self.delta_head = nn.Conv2d(hidden_channels, 3, 3, padding=1)
+        self.alpha_head = nn.Conv2d(hidden_channels, 1, 1)
+
+        # Zero-init residual head so an untrained module returns the anchor.
+        nn.init.zeros_(self.delta_head.weight)
+        nn.init.zeros_(self.delta_head.bias)
+        init_alpha = min(max(float(initial_alpha), 1e-4), 1.0 - 1e-4)
+        nn.init.zeros_(self.alpha_head.weight)
+        nn.init.constant_(self.alpha_head.bias, torch.logit(torch.tensor(init_alpha)).item())
+
+    @torch.no_grad()
+    def forward(
+        self,
+        coarse_normal: torch.Tensor,
+        rgb: Optional[torch.Tensor] = None,
+        domain_is_transparent: Optional[Union[bool, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        if self.use_rgb and rgb is None:
+            raise ValueError("GSM requires the RGB input when use_rgb=True.")
+
+        transparent_anchor = F.normalize(coarse_normal.float(), p=2, dim=1).to(coarse_normal.dtype)
+        domain = _domain_mask(domain_is_transparent, coarse_normal)
+        if self.use_rgb:
+            opaque_anchor = guided_filter_torch(
+                coarse_normal, rgb, radius=self.anchor_radius, eps=self.anchor_eps
+            )
+            anchor = torch.where(domain > 0.5, transparent_anchor, opaque_anchor)
+        else:
+            anchor = transparent_anchor
+
+        edge_feat = self.edge_encoder(rgb) if self.use_rgb else None
+        if self.use_rgb:
+            x_in = torch.cat([coarse_normal, anchor, rgb, domain], dim=1)
+        else:
+            x_in = torch.cat([coarse_normal, anchor, domain], dim=1)
+        x = self.main_encoder(x_in)
+
+        for proj, block in zip(self.projections, self.blocks):
+            x = proj(torch.cat([x, edge_feat], dim=1) if edge_feat is not None else x)
+            x = block(x)
+
+        feat = self.feature_proj(torch.cat([x, edge_feat], dim=1) if edge_feat is not None else x)
+        delta = self.delta_head(feat)
+        alpha = torch.sigmoid(self.alpha_head(feat))
+        residual = self.residual_scale * alpha * delta
+
+        if self.residual_clip_deg > 0:
+            # Clip the tangent-plane magnitude of the residual to a max angle.
+            max_norm = torch.tan(torch.deg2rad(torch.tensor(
+                self.residual_clip_deg, device=residual.device, dtype=torch.float32
+            )))
+            residual_norm = torch.linalg.vector_norm(
+                residual.float(), ord=2, dim=1, keepdim=True
+            ).clamp_min(1e-6)
+            clip = (max_norm / residual_norm).clamp(max=1.0).to(residual.dtype)
+            residual = residual * clip
+
+        refined = F.normalize(anchor + residual, p=2, dim=1)
+        return refined
+
+    def count_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters())

transnormal2/lcm.py

@@ -0,0 +1,32 @@
+"""Local Continuity Module (LCM).
+
+A small residual convolutional head applied to the predicted normal latent
+before VAE decoding. It enforces local smoothness in latent space and is
+trained jointly with the core predictor LoRA.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class LocalContinuityModule(nn.Module):
+    """Residual 2-layer conv head operating on raw VAE latents.
+
+    Args:
+        num_channels: Latent channel count (``transformer.in_channels // 4``,
+            i.e. 32 for FLUX.2 [klein]).
+    """
+
+    def __init__(self, num_channels: int):
+        super().__init__()
+        self.lcm = nn.Sequential(
+            nn.Conv2d(num_channels, num_channels * 2, kernel_size=3, padding=1),
+            nn.GELU(),
+            nn.Conv2d(num_channels * 2, num_channels, kernel_size=3, padding=1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        lcm_dtype = next(self.lcm.parameters()).dtype
+        if x.dtype != lcm_dtype:
+            x = x.to(dtype=lcm_dtype)
+        return x + self.lcm(x)

transnormal2/pipeline.py

@@ -0,0 +1,288 @@
+"""TransNormal-2 inference pipeline.
+
+Single-step rectified-flow normal estimation on the FLUX.2 [klein] backbone:
+
+    RGB image -> VAE encode -> core-predictor LoRA (1 transformer step)
+              -> LCM (latent smoothing) -> VAE decode
+              -> GSM (edge-aware refinement) -> unit normal map
+
+The released model predicts the normal latent in a single deterministic
+transformer evaluation (no diffusion sampling loop, no random noise), so the
+output is reproducible bit-for-bit for a fixed input and dtype.
+"""
+
+import json
+import os
+from typing import List, Optional, Union
+
+import torch
+import torch.nn.functional as F
+
+from diffusers import Flux2KleinPipeline
+from diffusers.models import Flux2Transformer2DModel
+
+from .gsm import GSM
+from .lcm import LocalContinuityModule
+from .utils import resize_image_first, tensor_to_output
+
+# LoRA target modules of the FLUX.2 [klein] transformer (must match training).
+FLUX2_LORA_TARGET_MODULES = "|".join([
+    # Double-stream: image attention
+    r".*\.attn\.to_k$",
+    r".*\.attn\.to_q$",
+    r".*\.attn\.to_v$",
+    r".*\.attn\.to_out\.0$",
+    # Double-stream: text attention
+    r".*\.attn\.add_k_proj$",
+    r".*\.attn\.add_q_proj$",
+    r".*\.attn\.add_v_proj$",
+    r".*\.attn\.to_add_out$",
+    # Double-stream: image FF (SwiGLU)
+    r".*\.ff\.linear_in$",
+    r".*\.ff\.linear_out$",
+    # Double-stream: text FF (SwiGLU)
+    r".*\.ff_context\.linear_in$",
+    r".*\.ff_context\.linear_out$",
+    # Single-stream: fused QKV+FF input projection
+    r".*\.attn\.to_qkv_mlp_proj$",
+    # Single-stream: output projection
+    r"single_transformer_blocks\.\d+\.attn\.to_out$",
+    # Context embedder
+    r"context_embedder$",
+])
+
+_WEIGHT_FILES = {
+    "lora": ("lora_core_predictor.safetensors",),
+    "lcm": ("lcm_normal.safetensors", "lcm_normal.pt"),
+    "gsm": ("gsm.safetensors", "gsm.pt"),
+}
+
+
+def _load_state_dict(path: str) -> dict:
+    if path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+        return load_file(path)
+    return torch.load(path, map_location="cpu", weights_only=True)
+
+
+def _find_weight(weights_dir: str, kind: str) -> str:
+    for name in _WEIGHT_FILES[kind]:
+        p = os.path.join(weights_dir, name)
+        if os.path.exists(p):
+            return p
+    raise FileNotFoundError(
+        f"Could not find {kind} weights in {weights_dir} "
+        f"(looked for {', '.join(_WEIGHT_FILES[kind])})"
+    )
+
+
+class TransNormal2Pipeline(Flux2KleinPipeline):
+    """FLUX.2 [klein] pipeline specialized for one-step normal estimation."""
+
+    # Set by from_pretrained_transnormal2().
+    local_continuity_module: Optional[LocalContinuityModule] = None
+    gsm: Optional[GSM] = None
+
+    # ── latent helpers (FLUX.2 BatchNorm-normalized patchified latents) ──
+
+    def _encode_vae_to_raw(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """Encode pixel values to raw VAE latents (before patchify/BN)."""
+        return self.vae.encode(pixel_values.to(dtype=self.vae.dtype)).latent_dist.mode()
+
+    def _raw_to_transformer_input(self, raw_latents: torch.Tensor) -> torch.Tensor:
+        """Raw VAE latents -> patchified + BN-normalized transformer input."""
+        patchified = self._patchify_latents(raw_latents)
+        bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(patchified.device, patchified.dtype)
+        bn_std = torch.sqrt(
+            self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps
+        ).to(patchified.device, patchified.dtype)
+        return (patchified - bn_mean) / bn_std
+
+    def _transformer_output_to_raw(self, latents_norm: torch.Tensor) -> torch.Tensor:
+        """Transformer output (patchified + BN-normalized) -> raw VAE latents."""
+        bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents_norm.device, latents_norm.dtype)
+        bn_std = torch.sqrt(
+            self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps
+        ).to(latents_norm.device, latents_norm.dtype)
+        denorm = latents_norm * bn_std + bn_mean
+        return self._unpatchify_latents(denorm)
+
+    def _decode_raw_latents(self, raw_latents: torch.Tensor) -> torch.Tensor:
+        """Decode raw VAE latents to pixel space ([-1, 1])."""
+        return self.vae.decode(raw_latents.to(dtype=self.vae.dtype), return_dict=False)[0]
+
+    # ── loading ──
+
+    @classmethod
+    def from_pretrained_transnormal2(
+        cls,
+        base_model: str = "black-forest-labs/FLUX.2-klein-base-9B",
+        weights_dir: str = "weights/transnormal2",
+        torch_dtype: torch.dtype = torch.bfloat16,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> "TransNormal2Pipeline":
+        """One-call loader: base model + core-predictor LoRA + LCM + GSM.
+
+        Args:
+            base_model: HF id or local path of ``FLUX.2-klein-base-9B``.
+            weights_dir: Local directory (or HF snapshot) containing
+                ``lora_core_predictor.safetensors``, ``lcm_normal.*``,
+                ``gsm.*`` and ``config.json``.
+            torch_dtype: bf16 strongly recommended (fp16 can produce NaNs).
+            device: Optional device to move the pipeline to (e.g. "cuda").
+        """
+        from peft import LoraConfig, set_peft_model_state_dict
+
+        with open(os.path.join(weights_dir, "config.json")) as f:
+            config = json.load(f)
+
+        transformer = Flux2Transformer2DModel.from_pretrained(base_model, subfolder="transformer")
+        transformer.requires_grad_(False)
+        if device is not None:
+            transformer.to(device=device, dtype=torch_dtype)
+        else:
+            transformer.to(dtype=torch_dtype)
+
+        lora_cfg = config["lora"]
+        adapter_name = lora_cfg.get("adapter_name", "core_predictor")
+        transformer.add_adapter(
+            LoraConfig(
+                r=int(lora_cfg["rank"]),
+                lora_alpha=float(lora_cfg["alpha"]),
+                init_lora_weights="gaussian",
+                target_modules=FLUX2_LORA_TARGET_MODULES,
+            ),
+            adapter_name=adapter_name,
+        )
+        lora_state = _load_state_dict(_find_weight(weights_dir, "lora"))
+        set_peft_model_state_dict(transformer, lora_state, adapter_name=adapter_name)
+        transformer.set_adapter(adapter_name)
+
+        pipe = cls.from_pretrained(
+            base_model,
+            transformer=transformer,
+            torch_dtype=torch_dtype,
+        )
+
+        lcm = LocalContinuityModule(transformer.config.in_channels // 4)
+        lcm.load_state_dict(_load_state_dict(_find_weight(weights_dir, "lcm")))
+        lcm.requires_grad_(False)
+
+        gsm = GSM(**config["gsm"])
+        gsm.load_state_dict(_load_state_dict(_find_weight(weights_dir, "gsm")), strict=True)
+        gsm.requires_grad_(False)
+        gsm.eval()
+
+        pipe.local_continuity_module = lcm
+        pipe.gsm = gsm
+        if device is not None:
+            pipe.to(device)
+            lcm.to(device=device, dtype=torch_dtype)
+            gsm.to(device=device, dtype=torch_dtype)
+        else:
+            lcm.to(dtype=torch_dtype)
+            gsm.to(dtype=torch_dtype)
+        return pipe
+
+    def to(self, *args, **kwargs):
+        result = super().to(*args, **kwargs)
+        # Keep the auxiliary heads on the same device/dtype as the pipeline.
+        for module in (self.local_continuity_module, self.gsm):
+            if module is not None:
+                module.to(*args, **kwargs)
+        return result
+
+    # ── inference ──
+
+    @torch.no_grad()
+    def __call__(  # type: ignore[override]
+        self,
+        image: torch.Tensor,
+        domain_is_transparent: Union[bool, torch.Tensor] = False,
+        process_res: Optional[int] = None,
+        output_type: str = "pt",
+        apply_gsm: bool = True,
+    ) -> Union[torch.Tensor, "np.ndarray", List]:
+        """Predict a surface normal map from an RGB image.
+
+        Args:
+            image: RGB tensor (B, 3, H, W) normalized to [-1, 1].
+            domain_is_transparent: True for transparent-dominant scenes
+                (glass/liquids): the GSM then anchors on the raw prediction
+                instead of the RGB guided filter. Default False (opaque).
+            process_res: Optional max processing edge; None keeps the input
+                resolution (snapped to a multiple of 16 internally).
+            output_type: "pt" (B, 3, H, W) in [0, 1], "np" (B, H, W, 3) in
+                [0, 1], or "pil".
+            apply_gsm: Disable to inspect the raw (pre-GSM) prediction.
+
+        Returns:
+            Normal map(s) encoded as ``(n + 1) / 2`` in [0, 1], resized back
+            to the input resolution.
+        """
+        if image.dim() != 4 or image.shape[1] != 3:
+            raise ValueError(f"Expected image of shape (B, 3, H, W), got {tuple(image.shape)}")
+
+        input_size = image.shape[2:]
+        rgb_in = resize_image_first(image, process_res)
+
+        device = self._execution_device
+        rgb_in = rgb_in.to(device=device, dtype=self.dtype)
+        batch_size = rgb_in.shape[0]
+
+        # 1. Empty-prompt conditioning (the released model is prompt-free).
+        prompt_embeds, text_ids = self.encode_prompt(prompt="", device=device)
+        # The Qwen3 text encoder may produce NaNs for the empty prompt in bf16.
+        if torch.isnan(prompt_embeds).any():
+            prompt_embeds = torch.nan_to_num(prompt_embeds, nan=0.0)
+        if prompt_embeds.shape[0] != batch_size:
+            prompt_embeds = prompt_embeds.expand(batch_size, -1, -1)
+            text_ids = text_ids.expand(batch_size, -1, -1)
+
+        # 2. Encode RGB to transformer-ready latents.
+        raw_rgb_latents = self._encode_vae_to_raw(rgb_in)
+        rgb_latents_norm = self._raw_to_transformer_input(raw_rgb_latents)
+        latent_ids = self._prepare_latent_ids(rgb_latents_norm).to(device)
+        packed_rgb_latents = self._pack_latents(rgb_latents_norm)
+
+        # 3. Klein base has guidance_embeds=False; keep the guard for safety.
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], 1.0, device=device, dtype=torch.float32)
+            guidance = guidance.expand(batch_size)
+        else:
+            guidance = None
+
+        # 4. Single-step core predictor (timestep fixed at 1/1000).
+        timestep = torch.tensor(1, device=device, dtype=self.dtype).expand(batch_size)
+        latents_out = self.transformer(
+            hidden_states=packed_rgb_latents,
+            timestep=timestep / 1000,
+            guidance=guidance,
+            encoder_hidden_states=prompt_embeds,
+            txt_ids=text_ids,
+            img_ids=latent_ids,
+            joint_attention_kwargs={},
+            return_dict=False,
+        )[0]
+        latents_out = latents_out[:, :packed_rgb_latents.size(1)]
+
+        # 5. Unpack -> denormalize -> LCM -> decode.
+        latents_patched = self._unpack_latents_with_ids(latents_out, latent_ids)
+        latents_raw = self._transformer_output_to_raw(latents_patched)
+        latents_raw = self.local_continuity_module(latents_raw)
+        latents_raw = latents_raw.to(dtype=self.dtype)
+        normal = self._decode_raw_latents(latents_raw)
+
+        # 6. GSM edge-aware refinement in pixel space.
+        if apply_gsm and self.gsm is not None:
+            rgb_resized = F.interpolate(
+                rgb_in, size=normal.shape[2:], mode="bilinear", align_corners=False
+            )
+            normal = self.gsm(normal, rgb_resized, domain_is_transparent=domain_is_transparent)
+
+        # 7. [-1, 1] -> [0, 1], resize back to the input resolution.
+        normal = self.image_processor.postprocess(normal, output_type="pt")
+        normal = F.interpolate(normal, size=input_size, mode="bilinear", align_corners=False)
+
+        self.maybe_free_model_hooks()
+        return tensor_to_output(normal, output_type)

transnormal2/utils.py

@@ -0,0 +1,89 @@
+"""Image IO and resize helpers for TransNormal-2 inference."""
+
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+
+def load_image(path: str) -> torch.Tensor:
+    """Load an image file as a (1, 3, H, W) float tensor in [-1, 1]."""
+    img = Image.open(path).convert("RGB")
+    arr = np.asarray(img).astype(np.float32)
+    ts = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
+    return ts / 127.5 - 1.0
+
+
+def resize_to_multiple_of_16(image_tensor: torch.Tensor) -> torch.Tensor:
+    """Rescale (B, C, H, W) so both sides are multiples of 16 (aspect kept)."""
+    h, w = image_tensor.shape[2], image_tensor.shape[3]
+    min_side = min(h, w)
+    scale = (min_side // 16) * 16 / min_side
+
+    new_h = (int(h * scale) // 16) * 16
+    new_w = (int(w * scale) // 16) * 16
+
+    if (new_h, new_w) == (h, w):
+        return image_tensor
+    return F.interpolate(
+        image_tensor, size=(new_h, new_w), mode="bilinear", align_corners=False
+    )
+
+
+def resize_image_first(image_tensor: torch.Tensor, process_res: Optional[int] = None) -> torch.Tensor:
+    """Optionally cap the max edge at ``process_res``, then snap to /16."""
+    if process_res:
+        max_edge = max(image_tensor.shape[2], image_tensor.shape[3])
+        if max_edge > process_res:
+            scale = process_res / max_edge
+            new_height = int(image_tensor.shape[2] * scale)
+            new_width = int(image_tensor.shape[3] * scale)
+            image_tensor = F.interpolate(
+                image_tensor, size=(new_height, new_width), mode="bilinear", align_corners=False
+            )
+    return resize_to_multiple_of_16(image_tensor)
+
+
+def tensor_to_output(
+    normal_01: torch.Tensor, output_type: str = "pt"
+) -> Union[torch.Tensor, np.ndarray, List[Image.Image]]:
+    """Convert a (B, 3, H, W) [0, 1] tensor to the requested output format."""
+    if output_type == "pt":
+        return normal_01
+    arr = normal_01.float().clamp(0, 1).permute(0, 2, 3, 1).cpu().numpy()
+    if output_type == "np":
+        return arr
+    if output_type == "pil":
+        return [Image.fromarray((a * 255).round().astype(np.uint8)) for a in arr]
+    raise ValueError(f"Unsupported output_type: {output_type} (use 'pt', 'np' or 'pil')")
+
+
+def save_normal_map(
+    normal: Union[torch.Tensor, np.ndarray],
+    path: str,
+    save_npy: Optional[str] = None,
+) -> None:
+    """Save a normal map prediction as a PNG (and optionally raw .npy).
+
+    Accepts (3, H, W) / (1, 3, H, W) tensors or (H, W, 3) arrays in [0, 1]
+    (the ``(n + 1) / 2`` encoding: camera-space X right, Y up, Z toward
+    camera mapped to RGB).
+    """
+    if isinstance(normal, torch.Tensor):
+        t = normal.detach().float().cpu()
+        if t.dim() == 4:
+            t = t[0]
+        if t.dim() == 3 and t.shape[0] == 3:
+            t = t.permute(1, 2, 0)
+        arr = t.numpy()
+    else:
+        arr = np.asarray(normal, dtype=np.float32)
+        if arr.ndim == 4:
+            arr = arr[0]
+
+    arr = np.clip(arr, 0.0, 1.0)
+    Image.fromarray((arr * 255).round().astype(np.uint8)).save(path)
+    if save_npy:
+        np.save(save_npy, arr)