Initial commit

Copied chord-demo and modified for Minecraft PBR material standards.

.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
+*.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,28 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+.eggs/
+dist/
+build/
+
+# Environment
+.env
+.venv/
+venv/
+
+# IDE
+.vscode/
+.idea/
+
+# Claude Code
+.claude/
+
+# Project specific
+output/
+*.safetensors
+test_*.png
+
+# OS
+.DS_Store
+Thumbs.db

LICENSE

@@ -0,0 +1,99 @@
+Ubisoft Machine Learning License (Research-Only - Copyleft)
+
+This license governs the use, reproduction, and distribution of the Licensed
+Materials, including AI Models and associated source code for the sole purpose
+of scientific research. By accessing, downloading or using the Licensed
+Materials, you hereby accept to be bound by this [Ubisoft Machine Learning
+License (Research-Only - Copyleft)] agreement (hereinafter the “License”).
+
+1. Licensed Materials
+
+- AI Models
+- Source Code
+
+2. Definitions
+
+“Licensed Materials”: Refers to the AI Models and/or Source Code licensed under
+this agreement.
+"Source Code" means the preferred form of the work for making modifications to
+it corresponding to text written using human-readable programming language.
+"Object Code" means any non-source form of a work.
+“AI Model” means any machine learning based assembly or assemblies (including
+checkpoints), consisting of learnt weights, parameters (including optimizer
+states), corresponding to the model architecture as embodied in the Source Code.
+“Output” means the results of operating an AI Model as embodied in
+informational content resulting therefrom.
+“Derivative”: Any work derived from or based upon the Licensed Materials,
+including modifications.
+“Permitted Purpose”: Use for academic or research purposes only. Commercial
+use is strictly prohibited.
+“Distribution”: Any sharing of the Licensed Materials or Derivatives with third
+parties, including hosting as a service.
+“Licensor”: The rights holder or authorized entity granting this License.
+“You”: The individual or entity receiving and exercising rights under this
+License.
+
+3. Grant of Rights
+
+Subject to compliance with the terms of this License, You are granted a
+worldwide, royalty-free, non-exclusive License to use, study, reproduce,
+modify, and distribute the Licensed Materials and Derivatives solely for the
+Permitted Purpose. As between You and Licensor, Licensor claims no rights in
+the Outputs You generate using the AI Models used in accordance with the
+Permitted Purpose.
+
+4. Distribution of Licensed Materials and Derivatives
+
+Any Distribution of the Derivatives of the Licensed Materials, or the Licensed
+Materials shall be licensed under the same exact terms as this License.
+Redistribution shall include this License and retain all notices of author
+attribution and all modifications shall be clearly marked.
+
+5. Use Restrictions
+
+You shall not use the Licensed Materials or its Derivatives for:
+- any other purposes than the Permitted Purpose, including for commercial
+  purposes such as using the Licensed Materials in any activity intended for
+  commercial advantage or monetary compensation directly or indirectly;
+- weaponry, warfare, military applications, surveillance, or any activity that
+  may cause harm or violate human rights;
+- engaging or enabling fully automated decision-making that may adversely
+  impacts a natural person's legal rights;
+- providing medical advice or making clinical decisions;
+- generating content that promotes or incites hatred, violence, discrimination,
+  or harm based on race, ethnicity, religion, gender, sexual orientation, or
+  any other protected characteristic;
+- generating content that includes depictions of sexual abuse, sexual
+  violence, explicit pornography, or any form of non-consensual acts and/or
+  generating content that includes depictions of child nudity, child
+  pornography, or any form of child exploitation;
+
+6. Disclaimer of Warranty
+
+THE LICENSED MATERIALS IS PROVIDED "AS IS" AND “AS AVAILABLE” WITHOUT
+WARRANTIES OF ANY KIND WHETHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION
+THE IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE,
+NON-INFRINGEMENT, CORRECTNESS, ACCURACY, OR RELIABILITY. THE LICENSOR DISCLAIMS
+ALL LIABILITY FOR DAMAGES RESULTING FROM THE USE OR INABILITY TO USE THE
+LICENSED MATERIALS. THE USE OF THE LICENSED MATERIALS AND ANY OUTPUTS YOU MAY
+GENERATE SHALL BE AT YOUR OWN RISK.
+
+7. Termination
+
+This License terminates automatically if You violate any of its terms. Upon
+termination, You shall cease all use and distribution of the Licensed
+Materials and its Derivatives.
+
+8. Governing Law
+
+The validity of this Agreement and any of its terms and provisions, as well as
+the rights and duties of the parties hereunder, shall be governed, interpreted
+and enforced in accordance with the laws of France.
+
+9. Miscellaneous
+
+If any provision of this License is held to be invalid, illegal or
+unenforceable, the remaining provisions shall be unaffected thereby and remain
+valid as if such provision had not been set forth herein.
+
+Copyright (C) 2025 UBISOFT ENTERTAINMENT. All Rights Reserved.

README.md

@@ -1,14 +1,44 @@
 ---
-title: Chord Minecraft Pbr
-emoji: 🏆
-colorFrom: blue
-colorTo: purple
+title: Chord Demo
+emoji: 🐠
+colorFrom: yellow
+colorTo: yellow
 sdk: gradio
-sdk_version: 6.3.0
+sdk_version: 5.49.1
 app_file: app.py
 pinned: false
-license: apache-2.0
-short_description: Chord for Minecraft shader PBR textures (Java and Bedrock)
+python_version: 3.12
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+<div align="center">
+<h1>Chord: Chain of Rendering Decomposition for PBR Material Estimation from Generated Texture Images</h1>
+
+<a href="https://arxiv.org/abs/2509.09952"><img src="https://img.shields.io/badge/arXiv-2509.09952-B31B1B?logo=arxiv&logoColor=white&style=flat-square" alt="arXiv"></a>
+<a href="https://ubisoft-laforge.github.io/world/chord/"><img src="https://img.shields.io/badge/Project-Page-brightgreen?logo=ubisoft&logoColor=white&style=flat-square" alt="Project Page"></a>
+
+[Zhi Ying](https://orcid.org/0009-0008-8390-3366)\*, [Boxiang Rong](https://ribosome-rbx.github.io/)\*, [Jingyu Wang](https://ccetaw.github.io/), [Maoyuan Xu](https://ultraman-blazar.github.io/)
+</div>
+
+Official implementation of the paper "**Chord: Chain of Rendering Decomposition for PBR Material Estimation from Generated Texture Images**".
+
+## License
+
+This project is released under the **Ubisoft Machine Learning License (Research-Only - Copyleft)**. See the full terms in the [LICENSE](LICENSE) file.
+
+## Citation
+
+If you find our work useful, please consider citing:
+
+```
+@misc{ying2025chord,
+    title={Chord: Chain of Rendering Decomposition for PBR Material Estimation from Generated Texture Images}, 
+    author={Zhi Ying and Boxiang Rong and Jingyu Wang and Maoyuan Xu},
+    year={2025},
+    eprint={2509.09952},
+    archivePrefix={arXiv},
+    primaryClass={cs.GR},
+    url={https://arxiv.org/abs/2509.09952}, 
+}
+```
+
+© [2025] Ubisoft Entertainment. All Rights Reserved.

app.py

@@ -0,0 +1,315 @@
+import gradio as gr
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from PIL import Image
+import torch
+import copy
+from omegaconf import OmegaConf
+from torchvision.transforms import v2
+from torchvision.transforms.functional import to_pil_image
+from huggingface_hub import hf_hub_download, login
+import spaces
+
+from chord import ChordModel
+from chord.module import make
+from chord.util import get_positions, rgb_to_srgb
+from chord.io import load_torch_file
+from chord.minecraft_pbr import convert_to_labpbr, convert_to_bedrock, LABPBR_METAL_CHOICES
+
+
+def _load_examples(directory: str) -> list:
+    """Load example images from a directory, returning empty list if not found."""
+    if os.path.isdir(directory):
+        return [[f"{directory}/{f}"] for f in sorted(os.listdir(directory))]
+    return []
+
+
+EXAMPLES_USECASE_1 = _load_examples("examples/generated")
+EXAMPLES_USECASE_2 = _load_examples("examples/in_the_wild")
+EXAMPLES_USECASE_3 = _load_examples("examples/specular")
+
+MODEL_OBJ = None
+
+# Use local model if available, otherwise download from HuggingFace
+LOCAL_MODEL_PATH = "chord_v1.safetensors"
+if os.path.exists(LOCAL_MODEL_PATH):
+    MODEL_CKPT_PATH = LOCAL_MODEL_PATH
+    print(f"Using local model: {MODEL_CKPT_PATH}")
+else:
+    hf_token = os.environ.get("HF_TOKEN")
+    if not hf_token:
+        raise EnvironmentError(
+            "HF_TOKEN environment variable is required to download the model. "
+            "Set it in a .env file or export it in your shell."
+        )
+    login(token=hf_token)
+    MODEL_CKPT_PATH = hf_hub_download(repo_id="Ubisoft/ubisoft-laforge-chord", filename="chord_v1.safetensors")
+    print(f"Downloaded model to: {MODEL_CKPT_PATH}")
+
+
+def load_model(ckpt_path):
+    print("Loading model from:", ckpt_path)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    config = OmegaConf.load("config/chord.yaml")
+    model = ChordModel(config)
+    state_dict = load_torch_file(ckpt_path)
+    model.load_state_dict(state_dict)
+    model.eval()
+    model.to(device)
+    return model
+
+def run_model(model, img: Image.Image):
+    device = next(model.parameters()).device
+    to_tensor = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])
+    image = to_tensor(img).to(device)
+    x = v2.Resize(size=(1024, 1024), antialias=True)(image).unsqueeze(0)
+    with torch.no_grad(), torch.autocast(device_type=device.type):
+        output = model(x)
+    return output
+
+def relit(model, maps):
+    maps['metallic'] = maps.get('metalness', torch.zeros_like(maps['basecolor']))
+    device = next(model.parameters()).device
+    h, w = maps["basecolor"].shape[-2:]
+    light = make("point-light", {"position": [0, 0, 10]}).to(device)
+    pos = get_positions(h, w, 10).to(device)
+    camera = torch.tensor([0, 0, 10.0]).to(device)
+    for key in maps:
+        if maps[key].dim() == 3:
+            maps[key] = maps[key].unsqueeze(0)
+        maps[key] = maps[key].permute(0,2,3,1)  # BxCxHxW -> BxHxWxC
+    rgb = model.model.compute_render(maps, camera, pos, light).squeeze(0).permute(0,3,1,2)  # GxBxHxWxC -> BxCxHxW
+    return torch.clamp(rgb_to_srgb(rgb), 0, 1)
+
+@spaces.GPU
+def inference(
+    img,
+    output_format,
+    seamless,
+    ao_strength,
+    ao_blur,
+    compute_porosity,
+    normalize_porosity,
+    compute_sss,
+    sss_curvature_weight,
+    sss_ao_weight,
+    sss_blur,
+    compute_emission,
+    emission_threshold,
+    emission_knee,
+    emission_bloom,
+    hardcoded_metal,
+):
+    """
+    Run Chord model and output shader-compatible textures.
+
+    Args:
+        output_format: "labpbr" for Java Edition shaders, "bedrock" for Bedrock RTX
+
+    Returns:
+        albedo: RGB albedo/basecolor texture
+        packed: Specular (_s) for LabPBR or MER/MERS for Bedrock
+        normal: Normal texture (_n for LabPBR, _normal for Bedrock)
+        render: Relit preview image
+    """
+    global MODEL_OBJ
+
+    if MODEL_OBJ is None or getattr(MODEL_OBJ, "_ckpt", None) != MODEL_CKPT_PATH:
+        MODEL_OBJ = load_model(MODEL_CKPT_PATH)
+        MODEL_OBJ._ckpt = MODEL_CKPT_PATH  # store path inside object
+
+    if img is None:
+        return None, None, None, None
+
+    ori_h, ori_w = img.size[1], img.size[0]
+    out = run_model(MODEL_OBJ, img)
+    maps = copy.deepcopy(out)
+    rendered = relit(MODEL_OBJ, maps)
+    resize_back = v2.Resize(size=(ori_h, ori_w), antialias=True)
+
+    # Resize all maps to original resolution
+    basecolor = resize_back(out["basecolor"])
+    normal = resize_back(out["normal"])
+    roughness = resize_back(out["roughness"].unsqueeze(0) if out["roughness"].dim() == 2 else out["roughness"])
+    metalness = resize_back(out["metalness"].unsqueeze(0) if out["metalness"].dim() == 2 else out["metalness"])
+
+    if output_format == "bedrock":
+        # Convert to Bedrock RTX format (MER/MERS)
+        result = convert_to_bedrock(
+            basecolor=basecolor,
+            normal=normal,
+            roughness=roughness,
+            metalness=metalness,
+            compute_sss=compute_sss,
+            sss_curvature_weight=sss_curvature_weight,
+            sss_ao_weight=sss_ao_weight,
+            sss_blur=int(sss_blur),
+            compute_emission=compute_emission,
+            emission_threshold=emission_threshold,
+            emission_knee=emission_knee,
+            emission_bloom=int(emission_bloom),
+        )
+        return (
+            result['albedo'],
+            result['mer'],
+            result['normal'],
+            to_pil_image(resize_back(rendered).squeeze(0)),
+        )
+    else:
+        # Convert to LabPBR 1.3 format (default)
+        result = convert_to_labpbr(
+            basecolor=basecolor,
+            normal=normal,
+            roughness=roughness,
+            metalness=metalness,
+            derive_ao_height=True,
+            seamless=seamless,
+            ao_strength=ao_strength,
+            ao_blur=int(ao_blur),
+            compute_porosity=compute_porosity,
+            normalize_porosity=normalize_porosity,
+            compute_sss=compute_sss,
+            sss_curvature_weight=sss_curvature_weight,
+            sss_ao_weight=sss_ao_weight,
+            sss_blur=int(sss_blur),
+            compute_emission=compute_emission,
+            emission_threshold=emission_threshold,
+            emission_knee=emission_knee,
+            emission_bloom=int(emission_bloom),
+            hardcoded_metal=hardcoded_metal,
+        )
+        return (
+            result['albedo'],
+            result['specular'],
+            result['normal'],
+            to_pil_image(resize_back(rendered).squeeze(0)),
+        )
+
+with gr.Blocks(title="Chord - PBR Material Estimation") as demo:
+
+    gr.Markdown("# **Chord: PBR Material Estimation → Minecraft Shader Formats**")
+    gr.Markdown("""
+Upload an image to estimate PBR materials and export for Minecraft shaders.
+
+**Supported Formats:**
+- **LabPBR 1.3** (Java Edition): Specular (_s) + Normal (_n) with AO/height
+- **Bedrock RTX**: MER/MERS (_mer/_mers) + Normal (_normal)
+""")
+
+    with gr.Row():
+        with gr.Column():
+            input_img = gr.Image(type="pil", label="Input Image", height=512)
+
+            with gr.Accordion("Output Settings", open=True):
+                gr.Markdown("#### Output Format")
+                output_format = gr.Radio(
+                    choices=[("LabPBR 1.3 (Java Edition)", "labpbr"), ("Bedrock RTX", "bedrock")],
+                    value="labpbr",
+                    label="Format",
+                    info="Choose shader format for your Minecraft edition"
+                )
+
+                gr.Markdown("#### General (LabPBR only)")
+                seamless = gr.Checkbox(label="Seamless/Tileable", value=False,
+                    info="Enable for textures that should tile seamlessly")
+
+                gr.Markdown("#### AO & Height Derivation (LabPBR only)")
+                ao_strength = gr.Slider(minimum=0.5, maximum=5.0, value=2.0, step=0.1,
+                    label="AO Strength", info="AO contrast multiplier")
+                ao_blur = gr.Slider(minimum=0, maximum=15, value=5, step=1,
+                    label="AO Blur", info="Gaussian blur radius for AO smoothing")
+
+                gr.Markdown("#### Porosity (LabPBR Blue 0-64)")
+                compute_porosity = gr.Checkbox(label="Compute Porosity", value=False,
+                    info="Calculate porosity from AO, smoothness, and F0 (LabPBR only)")
+                normalize_porosity = gr.Checkbox(label="Normalize Porosity", value=True,
+                    info="Normalize to full range before LabPBR scaling")
+
+                gr.Markdown("#### Subsurface Scattering")
+                compute_sss = gr.Checkbox(label="Compute SSS", value=False,
+                    info="Calculate SSS thickness from normal curvature (LabPBR: blue 65-255, Bedrock: MERS alpha)")
+                sss_curvature_weight = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.05,
+                    label="Curvature Weight", info="Weight for curvature contribution")
+                sss_ao_weight = gr.Slider(minimum=0.0, maximum=1.0, value=0.3, step=0.05,
+                    label="AO Weight", info="Weight for inverted AO contribution")
+                sss_blur = gr.Slider(minimum=0, maximum=10, value=2, step=1,
+                    label="SSS Blur", info="Gaussian blur for soft SSS look")
+
+                gr.Markdown("#### Emission")
+                compute_emission = gr.Checkbox(label="Compute Emission", value=False,
+                    info="Extract emission from bright areas of basecolor")
+                emission_threshold = gr.Slider(minimum=0.5, maximum=1.0, value=0.85, step=0.01,
+                    label="Emission Threshold", info="Luminance threshold for detection")
+                emission_knee = gr.Slider(minimum=0.0, maximum=0.3, value=0.1, step=0.01,
+                    label="Emission Knee", info="Soft knee width (0 = hard threshold)")
+                emission_bloom = gr.Slider(minimum=0, maximum=21, value=0, step=1,
+                    label="Emission Bloom", info="Gaussian blur radius for bloom effect (0 = disabled)")
+
+                gr.Markdown("#### Hardcoded Metal (LabPBR only)")
+                hardcoded_metal = gr.Dropdown(
+                    choices=LABPBR_METAL_CHOICES,
+                    value="none",
+                    label="Metal Type",
+                    info="Use predefined metal F0 values (230-237) for metallic areas"
+                )
+
+            gr.Markdown("### Example Inputs — Generated Textures")
+            gr.Examples(
+                examples=EXAMPLES_USECASE_1,
+                inputs=[input_img],
+                label="Examples (Generated Textures)"
+            )
+
+            gr.Markdown("### Example Inputs — In The Wild Photographs")
+            gr.Examples(
+                examples=EXAMPLES_USECASE_2,
+                inputs=[input_img],
+                label="Examples (In The Wild Photographs)"
+            )
+
+            gr.Markdown("### Example Inputs — Specular Textures")
+            gr.Examples(
+                examples=EXAMPLES_USECASE_3,
+                inputs=[input_img],
+                label="Examples (Specular Textures)"
+            )
+
+            run_button = gr.Button("Run Estimation", variant="primary")
+
+        with gr.Column():
+            gr.Markdown("### Output Textures")
+            albedo_out = gr.Image(label="Albedo (basecolor)", height=340, format="png")
+            packed_out = gr.Image(label="Specular/MER — LabPBR: _s (smoothness,F0,porosity,emission) | Bedrock: _mer (M,E,R,[S])", height=340, format="png")
+            normal_out = gr.Image(label="Normal — LabPBR: _n (XY,AO,height) | Bedrock: _normal (XYZ DirectX)", height=340, format="png")
+
+            gr.Markdown("### Preview")
+            render_out = gr.Image(label="Relit Preview (Point Light)", height=340, format="png")
+
+    run_button.click(
+        inference,
+        inputs=[
+            input_img,
+            output_format,
+            seamless,
+            ao_strength,
+            ao_blur,
+            compute_porosity,
+            normalize_porosity,
+            compute_sss,
+            sss_curvature_weight,
+            sss_ao_weight,
+            sss_blur,
+            compute_emission,
+            emission_threshold,
+            emission_knee,
+            emission_bloom,
+            hardcoded_metal,
+        ],
+        outputs=[albedo_out, packed_out, normal_out, render_out]
+    )
+
+if __name__ == "__main__":
+    demo.launch()

chord/__init__.py

@@ -0,0 +1,14 @@
+import torch
+import torch.nn as nn
+from chord.module import make
+from chord.module.chord import post_decoder
+
+class ChordModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.model = make(config.model.name, config.model)
+
+    def forward(self, x: torch.Tensor):
+        x = {"render": x}
+        pred = self.model(x)
+        return post_decoder(pred)

chord/io.py

@@ -0,0 +1,100 @@
+import torch
+import imageio.v3 as imageio
+import numpy as np
+import warnings
+import os
+import safetensors
+
+import torchvision.transforms.functional as F
+
+def read_image(filename: str, out: torch.Tensor=None) -> torch.Tensor:
+    '''
+    Read a local image file into a float tensor (pixel values are normalized to [0, 1], CxHxW)
+
+    Args:
+        filename: Image file path.
+        out: Fill in this tensor rather than return a new tensor if provided.
+
+    Returns:
+        Loaded image tensor.
+    '''
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore") # ignore PIL's user warning that reads fp16 img as fp32
+        img: np.ndarray = imageio.imread(filename)
+
+    # Convert the image array to float tensor according to its data type
+    res = None
+    if img.dtype == np.uint8:
+        img = img.astype(np.float32) / 255.0
+    elif img.dtype == np.uint16 or img.dtype == np.int32:
+        img = img.astype(np.float32) / 65535.0
+    else:
+        raise ValueError(f'Unrecognized image pixel value type: {img.dtype}')
+    if img.ndim == 2:
+        res = torch.from_numpy(img).unsqueeze(0)  # 1xHxW for grayscale images
+    elif img.ndim == 3:
+        res = torch.from_numpy(img).movedim(2, 0)[:3] # HxWxC to CxHxW
+    else:
+        raise ValueError(f'Unrecognized image dimension: {img.shape}')
+    
+    if out is None:
+        return res
+    out.copy_(res)
+
+def create_img(img: torch.Tensor):
+    '''
+    Convert tensor to PIL image
+
+    Args:
+        path: Image tensor CxHxW. Squeeze if BxCxHxW and B==1
+
+    Returns:
+        PIL image
+    '''
+    if img.dim() == 4:
+        assert img.shape[0] == 1
+        img = img.squeeze(0)
+
+    if img.shape[0] == 4:
+        out_img = F.to_pil_image(img, mode="CMYK")
+        out_img = out_img.convert('RGB')
+    elif img.shape[0] == 3:
+        out_img = F.to_pil_image(img, mode="RGB")
+    elif img.shape[0] == 1:
+        out_img = F.to_pil_image(img, mode="L")
+    else:
+        raise ValueError("Unsupported image dimension.")
+    return out_img
+    
+def save_maps(path: str, maps: dict):
+    '''
+    Save SVBRDF maps to a given path.
+
+    Args:
+        path: Output path.   
+        maps: Named maps of tensor images. 
+    '''
+    if not os.path.exists(path):
+        os.makedirs(path)
+    for name, image in maps.items():
+        out_img = create_img(image)
+        out_img.save(os.path.join(path, name+".png"))
+
+def load_torch_file(ckpt, device=None):
+    if device is None:
+        device = torch.device("cpu")
+    if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"):
+        with safetensors.safe_open(ckpt, framework="pt", device=device.type) as f:
+            state_dict = {}
+            for k in f.keys():
+                tensor = f.get_tensor(k)
+                state_dict[k] = tensor
+    else:
+        torch_args = {}
+        ckpt = torch.load(ckpt, map_location=device, weights_only=True, **torch_args)
+
+        if "state_dict" in ckpt:
+            state_dict = ckpt["state_dict"]
+        else:
+            state_dict = ckpt
+    return state_dict

chord/minecraft_pbr.py

@@ -0,0 +1,981 @@
+"""
+PBR Material Format Conversion
+
+Converts Chord model outputs to shader-compatible texture formats.
+
+Supported formats:
+1. LabPBR 1.3 (Minecraft Java Edition shaders)
+   Reference: https://shaderlabs.org/wiki/LabPBR_Material_Standard
+   - Specular (_s): R=smoothness, G=F0/metal, B=porosity/SSS, A=emission
+   - Normal (_n): R=normal.x, G=normal.y (DirectX), B=AO, A=height
+
+2. Bedrock RTX (Minecraft Bedrock Edition shaders)
+   - Normal (_normal): DirectX normal map (RGB)
+   - MER (_mer): R=metalness, G=emissive, B=roughness
+   - MERS (_mers): R=metalness, G=emissive, B=roughness, A=SSS
+     Note: Metalness and SSS are mutually exclusive per-pixel (metalness takes priority)
+"""
+
+import torch
+import torch.nn.functional as Fn
+from PIL import Image
+import numpy as np
+
+from .normal_utils import derive_ao_and_height
+
+
+# LabPBR Hardcoded Metal Values (Green channel of specular texture)
+# Reference: https://shaderlabs.org/wiki/LabPBR_Material_Standard
+LABPBR_METALS = {
+    "none": None,           # Use default behavior (dielectric/custom metal based on metalness)
+    "custom": 255,          # Custom metal - shader uses albedo as F0
+    "iron": 230,
+    "gold": 231,
+    "aluminum": 232,
+    "chrome": 233,
+    "copper": 234,
+    "lead": 235,
+    "platinum": 236,
+    "silver": 237,
+}
+
+# Display names for UI
+LABPBR_METAL_CHOICES = [
+    ("None (Auto)", "none"),
+    ("Custom Metal (255)", "custom"),
+    ("Iron (230)", "iron"),
+    ("Gold (231)", "gold"),
+    ("Aluminum (232)", "aluminum"),
+    ("Chrome (233)", "chrome"),
+    ("Copper (234)", "copper"),
+    ("Lead (235)", "lead"),
+    ("Platinum (236)", "platinum"),
+    ("Silver (237)", "silver"),
+]
+
+
+def roughness_to_smoothness(roughness: torch.Tensor) -> torch.Tensor:
+    """
+    Convert roughness to LabPBR perceptual smoothness.
+
+    LabPBR formula: roughness = pow(1.0 - perceptualSmoothness, 2.0)
+    Inverse: perceptualSmoothness = 1.0 - sqrt(roughness)
+
+    Args:
+        roughness: Roughness tensor in range [0, 1]
+
+    Returns:
+        Perceptual smoothness in range [0, 1]
+    """
+    return 1.0 - torch.sqrt(torch.clamp(roughness, 0.0, 1.0))
+
+
+def metalness_to_f0(
+    metalness: torch.Tensor,
+    threshold: float = 0.5,
+    hardcoded_metal: str = "none",
+) -> torch.Tensor:
+    """
+    Convert metalness to LabPBR F0/metal channel.
+
+    LabPBR encoding (stored LINEAR, no gamma):
+    - 0-229: Dielectric F0 (we use ~10 for typical dielectric F0 of 0.04)
+    - 230-254: Predefined metals (iron=230, gold=231, aluminum=232, etc.)
+    - 255: Custom metal (shader uses albedo as F0)
+
+    F0 values are stored linearly: F0 = green_channel / 255.0
+    No sRGB or gamma correction is applied.
+
+    Args:
+        metalness: Metalness tensor in range [0, 1]
+        threshold: Threshold above which material is considered metal
+        hardcoded_metal: Name of predefined metal type ("none", "custom", "iron", "gold", etc.)
+            When not "none", metallic areas use this specific metal value instead of 255.
+
+    Returns:
+        F0 channel values in range [0, 1] (stored LINEAR, scaled to 0-255 on save)
+    """
+    # F0 values stored linearly (no gamma correction)
+    # Dielectric F0 ≈ 0.04, which maps to ~10/255 ≈ 0.039 (linear)
+    dielectric_f0 = 10.0 / 255.0
+
+    # Determine metal F0 value (linear)
+    metal_value = LABPBR_METALS.get(hardcoded_metal, None)
+    if metal_value is None:
+        # Default: custom metal (255)
+        metal_f0 = 1.0
+    else:
+        metal_f0 = metal_value / 255.0
+
+    # Blend based on metalness (hard threshold for cleaner results)
+    is_metal = (metalness > threshold).float()
+    return torch.lerp(
+        torch.full_like(metalness, dielectric_f0),
+        torch.full_like(metalness, metal_f0),
+        is_metal
+    )
+
+
+def convert_normal_to_directx(
+    normal: torch.Tensor,
+    flip_y: bool = True,
+    swap_xy: bool = False,
+) -> torch.Tensor:
+    """
+    Convert normal map to DirectX convention (Y-down) if needed.
+
+    LabPBR uses DirectX-style normals where:
+    - R (X): Right is positive
+    - G (Y): Down is positive (Y-)
+
+    Chord outputs normals in range [0, 1] where 0.5 is neutral.
+    Note: Chord's internal rendering swaps X/Y channels (see chord.py:208),
+    suggesting the model may output [Y, X, Z]. Use swap_xy=True if needed.
+
+    Args:
+        normal: Normal map tensor, shape (B, 3, H, W) or (3, H, W), range [0, 1]
+        flip_y: Whether to flip Y for DirectX convention (OpenGL Y-up to DirectX Y-down)
+        swap_xy: Whether to swap X and Y channels (if model outputs [Y, X, Z])
+
+    Returns:
+        Normal map with DirectX convention, range [0, 1]
+    """
+    normal_out = normal.clone()
+
+    if normal.dim() == 4:
+        # Swap X and Y if needed (model outputs [Y, X, Z])
+        if swap_xy:
+            normal_out[:, 0, :, :] = normal[:, 1, :, :]
+            normal_out[:, 1, :, :] = normal[:, 0, :, :]
+        # Flip Y channel: new_y = 1.0 - old_y (OpenGL to DirectX)
+        if flip_y:
+            normal_out[:, 1, :, :] = 1.0 - normal_out[:, 1, :, :]
+    else:
+        if swap_xy:
+            normal_out[0, :, :] = normal[1, :, :]
+            normal_out[1, :, :] = normal[0, :, :]
+        if flip_y:
+            normal_out[1, :, :] = 1.0 - normal_out[1, :, :]
+
+    return normal_out
+
+
+def calculate_porosity(
+    ao: torch.Tensor,
+    smoothness: torch.Tensor,
+    f0: torch.Tensor,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """
+    Calculate porosity from AO, smoothness, and F0 maps.
+
+    Porosity is high where: AO is low AND smoothness is low AND F0 is low.
+    This identifies deep, rough, non-reflective areas typical of porous materials.
+
+    Args:
+        ao: Ambient occlusion map, shape (1, H, W) or (B, 1, H, W), range [0, 1]
+            where 0 = full occlusion, 1 = no occlusion
+        smoothness: Smoothness map, shape (1, H, W) or (B, 1, H, W), range [0, 1]
+        f0: F0/reflectance map, shape (1, H, W) or (B, 1, H, W), range [0, 1]
+        normalize: Whether to normalize output to full 0-1 range (default True)
+
+    Returns:
+        Porosity map, shape matching input, range [0, 1]
+    """
+    # Invert AO: deep/occluded areas become high
+    ao_factor = 1.0 - ao
+
+    # Invert smoothness to get roughness factor
+    roughness_factor = 1.0 - smoothness
+
+    # Invert F0: non-reflective areas become high
+    reflection_factor = 1.0 - f0
+
+    # Combine via multiplication (intersection of all three traits)
+    porosity = ao_factor * roughness_factor * reflection_factor
+
+    # Normalize to span full 0-1 range for better contrast
+    if normalize:
+        p_min = porosity.min()
+        p_max = porosity.max()
+        if p_max - p_min > 0:
+            porosity = (porosity - p_min) / (p_max - p_min)
+
+    # Scale to LabPBR porosity range (0-64 out of 255)
+    # LabPBR: 0-64 = porosity, 65-255 = SSS
+    porosity = porosity * (64.0 / 255.0)
+
+    return porosity
+
+
+def calculate_sss_thickness(
+    normal: torch.Tensor,
+    ao: torch.Tensor = None,
+    curvature_weight: float = 0.7,
+    ao_weight: float = 0.3,
+    blur_radius: int = 2,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """
+    Calculate SSS thickness map from normal map curvature and AO.
+
+    High curvature areas (edges, tips) are considered "thin" and allow more
+    light scattering. Inverted AO adds the "gummy glow" in crevices.
+
+    Args:
+        normal: Normal map, shape (3, H, W) or (B, 3, H, W), range [0, 1]
+        ao: Optional AO map, shape (1, H, W) or (B, 1, H, W), range [0, 1]
+            If None, only curvature is used
+        curvature_weight: Weight for curvature contribution (default 0.7)
+        ao_weight: Weight for inverted AO contribution (default 0.3)
+        blur_radius: Gaussian blur radius for soft SSS look (default 2)
+        normalize: Whether to normalize to full 0-1 range (default True)
+
+    Returns:
+        SSS thickness map, shape (1, H, W) or (B, 1, H, W), range [0, 1]
+        where 1.0 = maximum SSS (thin/translucent), 0.0 = no SSS
+    """
+    has_batch = normal.dim() == 4
+
+    if not has_batch:
+        normal = normal.unsqueeze(0)
+        if ao is not None:
+            ao = ao.unsqueeze(0)
+
+    b, _, h, w = normal.shape
+    device = normal.device
+
+    # Unpack normals from [0, 1] to [-1, 1]
+    nx = normal[:, 0:1, :, :] * 2.0 - 1.0
+    ny = normal[:, 1:2, :, :] * 2.0 - 1.0
+    nz = normal[:, 2:3, :, :] * 2.0 - 1.0
+
+    # Calculate curvature via gradient magnitude of normal components
+    # Using Sobel-like gradients
+    sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]],
+                           dtype=normal.dtype, device=device).view(1, 1, 3, 3) / 8.0
+    sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]],
+                           dtype=normal.dtype, device=device).view(1, 1, 3, 3) / 8.0
+
+    # Pad for convolution
+    pad = Fn.pad
+
+    # Gradients of each normal component
+    nx_padded = pad(nx, (1, 1, 1, 1), mode='reflect')
+    ny_padded = pad(ny, (1, 1, 1, 1), mode='reflect')
+    nz_padded = pad(nz, (1, 1, 1, 1), mode='reflect')
+
+    dx_nx = Fn.conv2d(nx_padded, sobel_x)
+    dy_nx = Fn.conv2d(nx_padded, sobel_y)
+    dx_ny = Fn.conv2d(ny_padded, sobel_x)
+    dy_ny = Fn.conv2d(ny_padded, sobel_y)
+    dx_nz = Fn.conv2d(nz_padded, sobel_x)
+    dy_nz = Fn.conv2d(nz_padded, sobel_y)
+
+    # Curvature magnitude
+    curvature = torch.sqrt(
+        dx_nx**2 + dy_nx**2 + dx_ny**2 + dy_ny**2 + dx_nz**2 + dy_nz**2 + 1e-6
+    )
+
+    # Normalize curvature to 0-1
+    if normalize:
+        c_min = curvature.min()
+        c_max = curvature.max()
+        if c_max - c_min > 0:
+            curvature = (curvature - c_min) / (c_max - c_min)
+
+    # Combine curvature with inverted AO
+    if ao is not None and ao_weight > 0:
+        inverted_ao = 1.0 - ao
+        sss = curvature * curvature_weight + inverted_ao * ao_weight
+    else:
+        sss = curvature
+
+    # Apply Gaussian blur for soft SSS look
+    if blur_radius > 0:
+        kernel_size = blur_radius * 2 + 1
+        sigma = blur_radius / 3.0
+
+        # Create 1D Gaussian kernel
+        x = torch.arange(kernel_size, device=device, dtype=normal.dtype)
+        x = x - kernel_size // 2
+        gaussian_1d = torch.exp(-x * x / (2 * sigma * sigma))
+        gaussian_1d = gaussian_1d / gaussian_1d.sum()
+
+        # Separable blur
+        padding = kernel_size // 2
+        sss_padded = pad(sss, (padding, padding, padding, padding), mode='reflect')
+
+        kernel_h = gaussian_1d.view(1, 1, 1, -1)
+        sss_h = Fn.conv2d(sss_padded, kernel_h, padding=0)
+
+        kernel_v = gaussian_1d.view(1, 1, -1, 1)
+        sss = Fn.conv2d(sss_h, kernel_v, padding=0)
+
+    # Final normalization and clamp
+    if normalize:
+        s_min = sss.min()
+        s_max = sss.max()
+        if s_max - s_min > 0:
+            sss = (sss - s_min) / (s_max - s_min)
+
+    sss = torch.clamp(sss, 0.0, 1.0)
+
+    if not has_batch:
+        sss = sss.squeeze(0)
+
+    return sss
+
+
+def sss_to_labpbr(sss: torch.Tensor) -> torch.Tensor:
+    """
+    Convert SSS intensity to LabPBR blue channel encoding.
+
+    LabPBR encoding for blue channel:
+    - 0-64: Porosity
+    - 65-255: SSS (linear)
+
+    Args:
+        sss: SSS intensity, range [0, 1] where 1 = max SSS
+
+    Returns:
+        LabPBR encoded SSS for blue channel, range [65/255, 1.0]
+    """
+    # Map [0, 1] to [65, 255] in normalized form
+    # 65/255 ≈ 0.255, 255/255 = 1.0
+    return 65.0 / 255.0 + sss * (190.0 / 255.0)
+
+
+def extract_emissive(
+    basecolor: torch.Tensor,
+    threshold: float = 0.85,
+    knee: float = 0.1,
+    bloom_radius: int = 0,
+) -> torch.Tensor:
+    """
+    Extract emissive regions from basecolor using luminance thresholding with soft knee.
+
+    Uses a soft knee to create smooth transitions at the threshold boundary,
+    avoiding hard cutoffs that look unnatural.
+
+    Args:
+        basecolor: RGB basecolor map, shape (3, H, W) or (B, 3, H, W), range [0, 1]
+        threshold: Luminance threshold for emission detection (0-1, default 0.85)
+        knee: Soft knee width for smooth transition (0-1, default 0.1)
+            - 0 = hard threshold
+            - Higher values = softer/wider transition
+        bloom_radius: Optional Gaussian blur radius for bloom effect (0 = disabled)
+
+    Returns:
+        Emission intensity map, shape (1, H, W) or (B, 1, H, W), range [0, 1]
+        where 1.0 = maximum emission, 0.0 = no emission
+    """
+    has_batch = basecolor.dim() == 4
+
+    if not has_batch:
+        basecolor = basecolor.unsqueeze(0)
+
+    # Calculate luminance (Rec. 709 coefficients)
+    luminance = (
+        0.2126 * basecolor[:, 0:1, :, :]
+        + 0.7152 * basecolor[:, 1:2, :, :]
+        + 0.0722 * basecolor[:, 2:3, :, :]
+    )
+
+    # Soft knee thresholding
+    # Creates smooth transition: 0 below (threshold - knee), 1 above (threshold + knee)
+    knee_low = threshold - knee
+    knee_high = threshold + knee
+
+    if knee > 0:
+        # Smooth hermite interpolation in the knee region
+        t = torch.clamp((luminance - knee_low) / (knee_high - knee_low + 1e-6), 0.0, 1.0)
+        # Smoothstep: 3t^2 - 2t^3
+        emission = t * t * (3.0 - 2.0 * t)
+    else:
+        # Hard threshold
+        emission = (luminance > threshold).float()
+
+    # Scale by how much luminance exceeds threshold (for intensity variation)
+    excess = torch.clamp((luminance - knee_low) / (1.0 - knee_low + 1e-6), 0.0, 1.0)
+    emission = emission * excess
+
+    # Optional bloom effect using Gaussian blur
+    if bloom_radius > 0:
+        # Ensure odd kernel size
+        kernel_size = bloom_radius * 2 + 1
+        sigma = bloom_radius / 3.0
+
+        # Create 1D Gaussian kernel
+        x = torch.arange(kernel_size, device=basecolor.device, dtype=basecolor.dtype)
+        x = x - kernel_size // 2
+        gaussian_1d = torch.exp(-x * x / (2 * sigma * sigma))
+        gaussian_1d = gaussian_1d / gaussian_1d.sum()
+
+        # Apply separable blur (horizontal then vertical)
+        padding = kernel_size // 2
+        emission_padded = Fn.pad(emission, (padding, padding, padding, padding), mode='reflect')
+
+        # Horizontal pass
+        kernel_h = gaussian_1d.view(1, 1, 1, -1)
+        emission_h = Fn.conv2d(emission_padded, kernel_h, padding=0)
+
+        # Vertical pass
+        kernel_v = gaussian_1d.view(1, 1, -1, 1)
+        emission_blurred = Fn.conv2d(emission_h, kernel_v, padding=0)
+
+        # Blend original with bloom (additive-ish)
+        emission = torch.maximum(emission, emission_blurred)
+
+    # Clamp final result
+    emission = torch.clamp(emission, 0.0, 1.0)
+
+    if not has_batch:
+        emission = emission.squeeze(0)
+
+    return emission
+
+
+def emission_to_labpbr(emission: torch.Tensor) -> torch.Tensor:
+    """
+    Convert emission intensity to LabPBR alpha channel encoding.
+
+    LabPBR emission encoding (LINEAR):
+    - 0 = no emission
+    - 1-254 = emissive levels (1 = min, 254 = max)
+    - 255 = no emission (same as 0, or no alpha channel)
+
+    Args:
+        emission: Emission intensity, range [0, 1] where 1 = max emission
+
+    Returns:
+        LabPBR encoded emission for alpha channel, range [0, 254/255]
+        (will become 0-254 when saved as uint8)
+    """
+    # Map [0, 1] to [0, 254]: 0 = no emission, 254 = max emission
+    # Avoid 255 which also means no emission in LabPBR
+    return torch.clamp(emission, 0.0, 1.0) * (254.0 / 255.0)
+
+
+def create_specular_texture(
+    roughness: torch.Tensor,
+    metalness: torch.Tensor,
+    porosity: torch.Tensor = None,
+    sss: torch.Tensor = None,
+    emission: torch.Tensor = None,
+    hardcoded_metal: str = "none",
+    sss_threshold: float = 0.01,
+) -> torch.Tensor:
+    """
+    Create LabPBR specular texture (_s).
+
+    All channels are stored LINEAR (no gamma correction).
+
+    Channel layout:
+    - R: Perceptual smoothness (1 - sqrt(roughness))
+    - G: F0 / metal ID (LINEAR: 0-229 dielectric, 230-237 metals, 255 custom)
+    - B: Porosity (0-64) or SSS (65-255) - per-pixel, SSS takes priority
+    - A: Emission (0 or 255 = no emission, 1-254 = emissive levels)
+
+    Args:
+        roughness: Roughness map, shape (B, 1, H, W) or (1, H, W)
+        metalness: Metalness map, shape (B, 1, H, W) or (1, H, W)
+        porosity: Optional porosity map, pre-scaled to 0-64/255 range (default: zeros)
+        sss: Optional SSS thickness map, range [0, 1] - will be encoded to 65-255
+            On a per-pixel basis, SSS takes precedence over porosity where SSS > threshold
+        emission: Optional emission intensity map, range [0, 1] where 1 = max emission
+            If provided, output is RGBA; if None, output is RGB
+        hardcoded_metal: Predefined metal type for metallic areas ("none", "iron", "gold", etc.)
+        sss_threshold: SSS intensity threshold for per-pixel priority (default 0.01)
+            Pixels with SSS > threshold use SSS encoding, others use porosity
+
+    Returns:
+        RGB or RGBA specular texture, shape (B, 3/4, H, W) or (3/4, H, W)
+    """
+    has_batch = roughness.dim() == 4
+
+    if not has_batch:
+        roughness = roughness.unsqueeze(0)
+        metalness = metalness.unsqueeze(0)
+        if porosity is not None:
+            porosity = porosity.unsqueeze(0)
+        if sss is not None:
+            sss = sss.unsqueeze(0)
+        if emission is not None:
+            emission = emission.unsqueeze(0)
+
+    b, _, h, w = roughness.shape
+    device = roughness.device
+
+    # R: Smoothness
+    smoothness = roughness_to_smoothness(roughness)
+
+    # G: F0/Metal (stored LINEAR, uses hardcoded metal value for metallic areas if specified)
+    f0 = metalness_to_f0(metalness, hardcoded_metal=hardcoded_metal)
+
+    # B: Porosity (0-64) and/or SSS (65-255) - per-pixel blending
+    # SSS takes priority where it exceeds threshold
+    if sss is not None and porosity is not None:
+        # Both provided: blend per-pixel, SSS takes priority where > threshold
+        sss_encoded = sss_to_labpbr(sss)
+        sss_mask = (sss > sss_threshold).float()
+        blue_channel = torch.where(sss_mask > 0.5, sss_encoded, porosity)
+    elif sss is not None:
+        # Only SSS provided
+        blue_channel = sss_to_labpbr(sss)
+    elif porosity is not None:
+        # Only porosity provided (already scaled to 0-64/255 range)
+        blue_channel = porosity
+    else:
+        # Default: no porosity or SSS
+        blue_channel = torch.zeros(b, 1, h, w, device=device)
+
+    channels = [smoothness, f0, blue_channel]
+
+    # A: Emission (optional)
+    if emission is not None:
+        emission_encoded = emission_to_labpbr(emission)
+        channels.append(emission_encoded)
+
+    specular = torch.cat(channels, dim=1)
+
+    if not has_batch:
+        specular = specular.squeeze(0)
+
+    return specular
+
+
+def create_normal_texture(
+    normal: torch.Tensor,
+    ao: torch.Tensor = None,
+    height: torch.Tensor = None,
+    flip_y: bool = True,
+    swap_xy: bool = False,
+) -> torch.Tensor:
+    """
+    Create LabPBR normal texture (_n).
+
+    Channel layout:
+    - R: Normal X
+    - G: Normal Y (DirectX convention, Y-down)
+    - B: Ambient Occlusion (0=full AO, 255=none)
+    - A: Height/displacement
+
+    Args:
+        normal: Normal map, shape (B, 3, H, W) or (3, H, W), range [0, 1]
+        ao: Optional AO map (default: 1.0 = no occlusion)
+        height: Optional height map (default: 0.5 = neutral)
+        flip_y: Whether to flip Y for DirectX convention
+        swap_xy: Whether to swap X and Y channels (if model outputs [Y, X, Z])
+
+    Returns:
+        RGBA normal texture, shape (B, 4, H, W) or (4, H, W)
+    """
+    has_batch = normal.dim() == 4
+
+    if not has_batch:
+        normal = normal.unsqueeze(0)
+        if ao is not None:
+            ao = ao.unsqueeze(0)
+        if height is not None:
+            height = height.unsqueeze(0)
+
+    b, _, h, w = normal.shape
+    device = normal.device
+
+    # Convert to DirectX convention
+    normal = convert_normal_to_directx(normal, flip_y=flip_y, swap_xy=swap_xy)
+
+    # Extract X and Y channels
+    normal_x = normal[:, 0:1, :, :]
+    normal_y = normal[:, 1:2, :, :]
+
+    # B: Ambient Occlusion
+    # LabPBR: 0 = full occlusion, 255 = no occlusion
+    # When not provided, caller should derive from normal map via derive_ao_and_height()
+    if ao is None:
+        ao = torch.ones(b, 1, h, w, device=device)  # Default: no occlusion
+
+    # A: Height map
+    # When not provided, caller should derive from normal map via derive_ao_and_height()
+    if height is None:
+        height = torch.full((b, 1, h, w), 0.5, device=device)  # Neutral height
+
+    normal_tex = torch.cat([normal_x, normal_y, ao, height], dim=1)
+
+    if not has_batch:
+        normal_tex = normal_tex.squeeze(0)
+
+    return normal_tex
+
+
+def tensor_to_pil(tensor: torch.Tensor, mode: str = "RGBA") -> Image.Image:
+    """
+    Convert a tensor to PIL Image using LINEAR mapping (no gamma correction).
+
+    All LabPBR data channels (F0, smoothness, porosity, SSS, emission, AO, height)
+    are stored linearly. This function performs a direct value * 255 conversion.
+
+    Args:
+        tensor: Image tensor, shape (C, H, W), range [0, 1]
+        mode: PIL image mode ("RGBA", "RGB", etc.)
+
+    Returns:
+        PIL Image with linear channel values
+    """
+    if tensor.dim() == 4:
+        tensor = tensor.squeeze(0)
+
+    # Linear conversion to uint8 (no gamma correction)
+    tensor = torch.clamp(tensor, 0.0, 1.0)
+    array = (tensor.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
+
+    return Image.fromarray(array, mode=mode)
+
+
+def convert_to_labpbr(
+    basecolor: torch.Tensor,
+    normal: torch.Tensor,
+    roughness: torch.Tensor,
+    metalness: torch.Tensor,
+    ao: torch.Tensor = None,
+    height: torch.Tensor = None,
+    porosity: torch.Tensor = None,
+    sss: torch.Tensor = None,
+    emission: torch.Tensor = None,
+    flip_normal_y: bool = True,
+    swap_normal_xy: bool = False,
+    derive_ao_height: bool = True,
+    compute_porosity: bool = False,
+    normalize_porosity: bool = True,
+    compute_sss: bool = False,
+    sss_curvature_weight: float = 0.7,
+    sss_ao_weight: float = 0.3,
+    sss_blur: int = 2,
+    compute_emission: bool = False,
+    emission_threshold: float = 0.85,
+    emission_knee: float = 0.1,
+    emission_bloom: int = 0,
+    hardcoded_metal: str = "none",
+    seamless: bool = False,
+    ao_strength: float = 2.0,
+    ao_blur: int = 5,
+) -> dict:
+    """
+    Convert Chord outputs to LabPBR 1.3 format.
+
+    Args:
+        basecolor: Albedo/basecolor map (B, 3, H, W) or (3, H, W)
+        normal: Normal map (B, 3, H, W) or (3, H, W)
+        roughness: Roughness map (B, 1, H, W) or (1, H, W)
+        metalness: Metalness map (B, 1, H, W) or (1, H, W)
+        ao: Optional AO map (if None and derive_ao_height=True, derived from normal)
+        height: Optional height map (if None and derive_ao_height=True, derived from normal)
+        porosity: Optional porosity map (if None and compute_porosity=True, calculated from AO/smoothness/F0)
+        sss: Optional SSS thickness map (if None and compute_sss=True, calculated from normal curvature)
+            Note: Porosity and SSS can coexist on same texture. Per-pixel, SSS takes priority where > threshold.
+        emission: Optional emission map (if None and compute_emission=True, extracted from basecolor)
+        flip_normal_y: Flip Y channel for DirectX convention (default True)
+        swap_normal_xy: Swap X/Y channels if model outputs [Y,X,Z] (default False)
+        derive_ao_height: If True, derive AO and height from normal map when not provided
+        compute_porosity: If True, calculate porosity from AO, smoothness, and F0 when not provided
+        normalize_porosity: If True, normalize porosity to full 0-1 range before LabPBR scaling
+        compute_sss: If True, calculate SSS thickness from normal curvature and AO when not provided
+        sss_curvature_weight: Weight for curvature in SSS calculation (default 0.7)
+        sss_ao_weight: Weight for inverted AO in SSS calculation (default 0.3)
+        sss_blur: Gaussian blur radius for soft SSS look (default 2)
+        compute_emission: If True, extract emission from basecolor luminance when not provided
+        emission_threshold: Luminance threshold for emission detection (0-1, default 0.85)
+        emission_knee: Soft knee width for smooth emission transition (0-1, default 0.1)
+        emission_bloom: Gaussian blur radius for emission bloom effect (0 = disabled)
+        hardcoded_metal: Predefined metal type for specular G channel ("none", "iron", "gold", etc.)
+            Uses metalness map as mask - metallic areas get this metal ID value.
+        seamless: Whether the texture should tile seamlessly (for height derivation)
+        ao_strength: AO intensity multiplier (higher = more contrast)
+        ao_blur: Gaussian blur radius for AO smoothing
+
+    Returns:
+        Dictionary with PIL Images:
+        - 'albedo': RGB albedo texture
+        - 'specular': RGB or RGBA specular texture (_s) - RGBA if emission enabled
+        - 'normal': RGBA normal texture (_n)
+    """
+    # Ensure tensors have consistent dimensions
+    has_batch = basecolor.dim() == 4
+
+    if has_batch:
+        basecolor = basecolor.squeeze(0)
+        normal = normal.squeeze(0)
+        roughness = roughness.squeeze(0)
+        metalness = metalness.squeeze(0)
+        if ao is not None:
+            ao = ao.squeeze(0)
+        if height is not None:
+            height = height.squeeze(0)
+
+    # Ensure roughness/metalness are (1, H, W)
+    if roughness.dim() == 2:
+        roughness = roughness.unsqueeze(0)
+    if metalness.dim() == 2:
+        metalness = metalness.unsqueeze(0)
+
+    # Derive AO and height from normal map if not provided
+    if derive_ao_height and (ao is None or height is None):
+        derived_ao, derived_height = derive_ao_and_height(
+            normal,
+            seamless=seamless,
+            ao_strength=ao_strength,
+            ao_blur=ao_blur,
+        )
+        if ao is None:
+            ao = derived_ao
+        if height is None:
+            height = derived_height
+
+    # Calculate porosity from AO, smoothness, and F0 if requested
+    # Note: Both porosity and SSS can coexist - SSS takes precedence per-pixel
+    if compute_porosity and porosity is None:
+        smoothness = roughness_to_smoothness(roughness)
+        f0 = metalness_to_f0(metalness)
+        # Use derived AO if available, otherwise use neutral (1.0 = no occlusion)
+        ao_for_porosity = ao if ao is not None else torch.ones_like(roughness)
+        porosity = calculate_porosity(
+            ao_for_porosity,
+            smoothness,
+            f0,
+            normalize=normalize_porosity,
+        )
+
+    # Calculate SSS thickness from normal curvature and AO if requested
+    if compute_sss and sss is None:
+        ao_for_sss = ao if ao is not None else None
+        sss = calculate_sss_thickness(
+            normal,
+            ao=ao_for_sss,
+            curvature_weight=sss_curvature_weight,
+            ao_weight=sss_ao_weight,
+            blur_radius=sss_blur,
+        )
+
+    # Extract emission from basecolor luminance if requested
+    if compute_emission and emission is None:
+        emission = extract_emissive(
+            basecolor,
+            threshold=emission_threshold,
+            knee=emission_knee,
+            bloom_radius=emission_bloom,
+        )
+
+    # Create LabPBR textures
+    specular_tex = create_specular_texture(
+        roughness, metalness, porosity, sss, emission, hardcoded_metal=hardcoded_metal
+    )
+    normal_tex = create_normal_texture(
+        normal, ao, height, flip_y=flip_normal_y, swap_xy=swap_normal_xy
+    )
+
+    # Determine specular output mode based on emission presence
+    specular_mode = "RGBA" if emission is not None else "RGB"
+
+    return {
+        'albedo': tensor_to_pil(basecolor, mode="RGB"),
+        'specular': tensor_to_pil(specular_tex, mode=specular_mode),
+        'normal': tensor_to_pil(normal_tex, mode="RGBA"),
+    }
+
+
+# =============================================================================
+# Bedrock RTX Format Conversion
+# =============================================================================
+
+def create_bedrock_normal_texture(
+    normal: torch.Tensor,
+    flip_y: bool = True,
+    swap_xy: bool = False,
+) -> torch.Tensor:
+    """
+    Create Bedrock RTX normal texture (_normal).
+
+    Bedrock uses standard DirectX normal maps (RGB only, no AO/height packing).
+
+    Args:
+        normal: Normal map, shape (B, 3, H, W) or (3, H, W), range [0, 1]
+        flip_y: Whether to flip Y for DirectX convention
+        swap_xy: Whether to swap X and Y channels
+
+    Returns:
+        RGB normal texture, shape (B, 3, H, W) or (3, H, W)
+    """
+    return convert_normal_to_directx(normal, flip_y=flip_y, swap_xy=swap_xy)
+
+
+def create_mer_texture(
+    metalness: torch.Tensor,
+    roughness: torch.Tensor,
+    emission: torch.Tensor = None,
+    sss: torch.Tensor = None,
+    metal_threshold: float = 0.5,
+) -> torch.Tensor:
+    """
+    Create Bedrock RTX MER/MERS texture.
+
+    All channels are stored LINEAR (no gamma correction).
+
+    Channel layout:
+    - R: Metalness (LINEAR, directly from Chord)
+    - G: Emissive level (LINEAR, 0 = no emission, 255 = max emission)
+    - B: Roughness (LINEAR, directly from Chord)
+    - A (optional): SSS thickness (only in MERS format)
+
+    Note: Metalness and SSS are mutually exclusive per-pixel.
+    Where metalness > threshold, SSS is set to 0.
+
+    Args:
+        metalness: Metalness map, shape (B, 1, H, W) or (1, H, W), range [0, 1]
+        roughness: Roughness map, shape (B, 1, H, W) or (1, H, W), range [0, 1]
+        emission: Optional emission intensity, range [0, 1] (0 = none, 1 = max)
+        sss: Optional SSS thickness, range [0, 1] (creates MERS output)
+        metal_threshold: Threshold for metalness priority over SSS (default 0.5)
+
+    Returns:
+        RGB (MER) or RGBA (MERS) texture, shape (B, 3/4, H, W) or (3/4, H, W)
+    """
+    has_batch = roughness.dim() == 4
+
+    if not has_batch:
+        roughness = roughness.unsqueeze(0)
+        metalness = metalness.unsqueeze(0)
+        if emission is not None:
+            emission = emission.unsqueeze(0)
+        if sss is not None:
+            sss = sss.unsqueeze(0)
+
+    b, _, h, w = roughness.shape
+    device = roughness.device
+
+    # R: Metalness (stored LINEAR)
+    metal_channel = metalness
+
+    # G: Emissive level (stored LINEAR, 0 = no emission, 255 = max)
+    # Note: Bedrock uses full 0-255 range (vs LabPBR where 0 and 255 both mean no emission)
+    if emission is not None:
+        emissive_channel = emission
+    else:
+        emissive_channel = torch.zeros(b, 1, h, w, device=device)
+
+    # B: Roughness (stored LINEAR)
+    rough_channel = roughness
+
+    channels = [metal_channel, emissive_channel, rough_channel]
+
+    # A: SSS (optional, creates MERS format)
+    # Metalness and SSS are mutually exclusive - metalness takes priority
+    if sss is not None:
+        # Zero out SSS where metalness exceeds threshold
+        is_metal = (metalness > metal_threshold).float()
+        sss_masked = sss * (1.0 - is_metal)
+        channels.append(sss_masked)
+
+    mer_tex = torch.cat(channels, dim=1)
+
+    if not has_batch:
+        mer_tex = mer_tex.squeeze(0)
+
+    return mer_tex
+
+
+def convert_to_bedrock(
+    basecolor: torch.Tensor,
+    normal: torch.Tensor,
+    roughness: torch.Tensor,
+    metalness: torch.Tensor,
+    emission: torch.Tensor = None,
+    sss: torch.Tensor = None,
+    flip_normal_y: bool = True,
+    swap_normal_xy: bool = False,
+    compute_sss: bool = False,
+    sss_curvature_weight: float = 0.7,
+    sss_ao_weight: float = 0.3,
+    sss_blur: int = 2,
+    compute_emission: bool = False,
+    emission_threshold: float = 0.85,
+    emission_knee: float = 0.1,
+    emission_bloom: int = 0,
+) -> dict:
+    """
+    Convert Chord outputs to Bedrock RTX format.
+
+    Args:
+        basecolor: Albedo/basecolor map (B, 3, H, W) or (3, H, W)
+        normal: Normal map (B, 3, H, W) or (3, H, W)
+        roughness: Roughness map (B, 1, H, W) or (1, H, W)
+        metalness: Metalness map (B, 1, H, W) or (1, H, W)
+        emission: Optional emission map (if None and compute_emission=True, extracted from basecolor)
+        sss: Optional SSS thickness map (if None and compute_sss=True, calculated from normal curvature)
+        flip_normal_y: Flip Y channel for DirectX convention (default True)
+        swap_normal_xy: Swap X/Y channels if model outputs [Y,X,Z] (default False)
+        compute_sss: If True, calculate SSS thickness from normal curvature (creates MERS output)
+        sss_curvature_weight: Weight for curvature in SSS calculation (default 0.7)
+        sss_ao_weight: Weight for inverted AO in SSS calculation (default 0.3)
+        sss_blur: Gaussian blur radius for soft SSS look (default 2)
+        compute_emission: If True, extract emission from basecolor luminance
+        emission_threshold: Luminance threshold for emission detection (0-1, default 0.85)
+        emission_knee: Soft knee width for smooth emission transition (0-1, default 0.1)
+        emission_bloom: Gaussian blur radius for emission bloom effect (0 = disabled)
+
+    Returns:
+        Dictionary with PIL Images:
+        - 'albedo': RGB albedo texture (basecolor)
+        - 'normal': RGB DirectX normal texture (_normal)
+        - 'mer': RGB MER texture or RGBA MERS texture (_mer / _mers)
+    """
+    # Ensure tensors have consistent dimensions
+    has_batch = basecolor.dim() == 4
+
+    if has_batch:
+        basecolor = basecolor.squeeze(0)
+        normal = normal.squeeze(0)
+        roughness = roughness.squeeze(0)
+        metalness = metalness.squeeze(0)
+
+    # Ensure roughness/metalness are (1, H, W)
+    if roughness.dim() == 2:
+        roughness = roughness.unsqueeze(0)
+    if metalness.dim() == 2:
+        metalness = metalness.unsqueeze(0)
+
+    # Calculate SSS thickness from normal curvature if requested
+    if compute_sss and sss is None:
+        sss = calculate_sss_thickness(
+            normal,
+            ao=None,  # Bedrock doesn't derive AO
+            curvature_weight=sss_curvature_weight,
+            ao_weight=sss_ao_weight,
+            blur_radius=sss_blur,
+        )
+
+    # Extract emission from basecolor luminance if requested
+    if compute_emission and emission is None:
+        emission = extract_emissive(
+            basecolor,
+            threshold=emission_threshold,
+            knee=emission_knee,
+            bloom_radius=emission_bloom,
+        )
+
+    # Create Bedrock textures
+    normal_tex = create_bedrock_normal_texture(
+        normal, flip_y=flip_normal_y, swap_xy=swap_normal_xy
+    )
+    mer_tex = create_mer_texture(
+        metalness, roughness, emission, sss
+    )
+
+    # Determine MER output mode based on SSS presence
+    mer_mode = "RGBA" if sss is not None else "RGB"
+
+    return {
+        'albedo': tensor_to_pil(basecolor, mode="RGB"),
+        'normal': tensor_to_pil(normal_tex, mode="RGB"),
+        'mer': tensor_to_pil(mer_tex, mode=mer_mode),
+    }

chord/module/__init__.py

@@ -0,0 +1,19 @@
+modules = {}
+
+def register(name):
+    def decorator(cls):
+        modules[name] = cls
+        return cls
+    return decorator
+
+
+def make(name, config):
+    model = modules[name](config)
+    return model
+
+
+from . import (
+    light,
+    stable_diffusion,
+    chord,
+)

chord/module/base.py

@@ -0,0 +1,13 @@
+import torch
+import torch.nn as nn
+
+class Base(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.setup()
+    
+    def setup(self):
+        raise NotImplementedError
+    

chord/module/chord.py

@@ -0,0 +1,281 @@
+import copy
+import torch
+from torch import nn
+import torch.nn.functional as Fn
+from torchvision.transforms import v2
+
+from . import register, make
+from .base import Base
+
+from chord.util import fresnelSchlick, GeometrySchlickGGX, DistributionGGX
+from chord.util import srgb_to_rgb, tone_gamma, get_positions, safe_01_div
+
+class dummy_module(nn.Module):
+    def forward(self, x): return x
+
+def post_decoder(out_dict):
+    out = {}
+    for key in out_dict.keys():
+        if key.startswith("approx"): continue
+        elif key == "normal": 
+            out[key] = Fn.normalize(2. * out_dict[key] - 1., dim=1) / 2. + 0.5
+        elif key == "rou_met":
+            out['roughness'], out['metalness'] = out_dict['rou_met'][:,0], out_dict['rou_met'][:,1]
+        else: out[key] = out_dict[key]
+    return out
+
+def process_irradiance(radiance, kernel_size=25, res=64):
+    """
+    Process the irradiance using PyTorch, equivalent to the original OpenCV-based function.
+    
+    Args:
+        radiance (torch.Tensor): Input radiance tensor (H, W).
+        kernel_size (int): Size of the kernel for the median blur.
+        res (int): Target resolution for resizing the image.
+    
+    Returns:
+        torch.Tensor: Processed radiance tensor (res, res).
+    """
+    # Ensure the input radiance is a 4D tensor (B, 1, H, W)
+    assert radiance.shape[1] == 1 and radiance.dim() == 4, f"Invalid radiance shape, got {radiance.shape}"
+    # resize to low resolution
+    resizer = v2.Resize(size=res, antialias=True)
+    radiance = resizer(radiance)
+
+    # Define a 11x11 averaging kernel
+    kernel = torch.ones((1, 1, 11, 11), dtype=torch.float32).to(radiance) / 121.0
+    # Apply convolution (averaging filter)
+    radiance = Fn.pad(radiance, (5,)*4, mode="reflect")  # Pad for edge handling
+    radiance = Fn.conv2d(radiance, kernel, padding=0)  # 'padding=2' to maintain input dimensions
+
+    # Clamp values and scale to [0, 255] for median filtering
+    radiance = torch.clamp(radiance * 255, 0, 255)  # Remove batch/channel dims
+
+    # Apply median filtering
+    paded_radiance = Fn.pad(radiance, (kernel_size // 2,) * 4, mode="reflect")  # Pad for edge handling
+    unfolded = Fn.unfold(paded_radiance, kernel_size)  # Extract patches
+    radiance = torch.median(unfolded, dim=1).values.view(radiance.shape)  # Median of patches
+
+    # Normalize to [0, 1]
+    rad_min, rad_max = radiance.amin([2,3], keepdim=True), radiance.amax([2,3], keepdim=True)
+    radiance = (radiance - rad_min) / (rad_max - rad_min)
+    return radiance
+
+def opt_light_dir(_radiance, _num_samples=6):
+    '''
+        _radiance: (bs, 1, h, w)
+    '''
+    assert _radiance.shape[1] == 1 and _radiance.dim()==4
+    bs, _, h, w = _radiance.shape
+
+    def evenly_sample(_num_samples, min=0, max=2*torch.pi):
+        # returns torch.tensor([1, _num_samples])
+        return torch.tensor(range(_num_samples+1)) * (max - min) / _num_samples + min
+
+    def compute_radiance_diff(angles):
+        num = angles.shape[-1]
+        dirs = torch.cat([torch.cos(angles), torch.sin(angles)]).T
+        pos_dir = grid_pos.repeat(num, 1, 1, 1)
+        pos_mask = torch.einsum("abcd,ad->abc", pos_dir, dirs) > 0
+        neg_mask = torch.einsum("abcd,ad->abc", pos_dir, dirs) < 0
+        samples_radiance = _radiance.repeat(1,num,1,1)
+        radiance_diff = (samples_radiance*pos_mask[None] - samples_radiance*neg_mask[None]).sum([2,3])
+        return radiance_diff
+    
+    angle_min, angle_max = 0, 2*torch.pi
+    grid_pos = Fn.normalize(get_positions(h,w,10)[...,:2], dim=-1, eps=1e-6).to(_radiance)
+    while(((angle_max - angle_min) > (torch.pi/90))):
+        angles = evenly_sample(_num_samples, angle_min, angle_max)[None].to(_radiance)
+        diffs = compute_radiance_diff(angles).mean(0)
+        angle_min = angles[:,diffs.argmax()].item() - (angle_max - angle_min)/_num_samples
+        angle_max = angles[:,diffs.argmax()].item() + (angle_max - angle_min)/_num_samples
+
+    light_angle = angles[:, diffs.argmax()]
+    return torch.tensor([torch.cos(light_angle), torch.sin(light_angle)]).to(_radiance)
+
+
+def find_light_dir(raw_irradiance, light):
+    raw_irradiance = v2.functional.rgb_to_grayscale(raw_irradiance)
+    irradiance = process_irradiance(raw_irradiance)
+    dir = opt_light_dir(irradiance)
+    dir = torch.cat([dir, torch.tensor([0.5**0.5]).to(dir)])
+    _light = copy.deepcopy(light)
+    _light.direction = dir
+    return _light
+
+@register("chord")
+class Chord(Base):
+    def setup(self):
+        # Define forward chain
+        self.chain_type = self.config.get("chain_type", "chord")
+        self.chain = self.config.get("chain_library", {})[self.chain_type]
+        self.prompts = self.config.get("rgbx_prompts", {})
+        self.roughness_step = self.config.get("roughness_step", 10)
+        self.metallic_step = self.config.get("metallic_step", 0.2)
+
+        self.sd = make(self.config.stable_diffusion.name, self.config.stable_diffusion)
+        self.dtype = self.sd.dtype
+        self.device = self.sd.device
+
+        # LEGO-conditioning
+        self.sd.unet.ConvIns = nn.ModuleDict()
+        self.sd.unet.ConvOuts = nn.ModuleDict()
+        self.sd.unet.FirstDownBlocks = nn.ModuleDict()
+        self.sd.unet.LastUpBlocks = nn.ModuleDict()
+        for key in list(set("_".join(self.chain.values()).split("_"))) + ["noise"]:
+            if "0" in key or "1" in key: continue
+            self.sd.unet.ConvIns[key] = nn.Conv2d(4, 320, 3, 1 , 1, device=self.device, dtype=self.dtype)
+            self.sd.unet.ConvIns[key].load_state_dict(self.sd.unet.conv_in.state_dict())
+        for kout in list(set(self.chain.keys())):
+            self.sd.unet.ConvOuts[kout] = nn.Conv2d(320, 4, 3, 1 , 1, device=self.device, dtype=self.dtype)
+            self.sd.unet.ConvOuts[kout].load_state_dict(self.sd.unet.conv_out.state_dict())
+            self.sd.unet.LastUpBlocks[kout] = copy.deepcopy(self.sd.unet.up_blocks[-1]).to(self.device)
+            self.sd.unet.FirstDownBlocks[kout] = copy.deepcopy(self.sd.unet.down_blocks[0]).to(self.device)
+        self.sd.unet.ConvIns.train()
+        self.sd.unet.ConvOuts.train()
+        self.sd.unet.FirstDownBlocks.train()
+        self.sd.unet.LastUpBlocks.train()
+        self.sd.unet.conv_in = dummy_module()
+        self.sd.unet.conv_out = dummy_module()
+
+        # Load Lights
+        if self.config.get("prior_light", None) is None:
+            self.prior_light = make("point-light", {"position": [0, 0, 10]})
+        else:
+            self.prior_light = make(self.config.prior_light.name, self.config.prior_light)
+
+        # Init Embeddings
+        self.text_emb = {}
+    # Eq.3
+    def compute_approxIrr(self, render, basecolor):
+        approxIrr = safe_01_div.apply(srgb_to_rgb(render), srgb_to_rgb(basecolor))
+        return tone_gamma(approxIrr)
+    # Eq.6
+    @torch.no_grad()
+    def compute_approxRouMet(self, render, maps, seperate=False, light=None):
+        render = srgb_to_rgb(render)
+        bs, _, h, w = render.shape
+        light = find_light_dir(maps['approxIrr'], self.prior_light) if light is None else light
+        # light.direction = estimate_light_dir(render, maps)
+        pos = get_positions(h, w, 10).to(self.device)
+        cameras = torch.tensor([0, 0, 10.0]).to(self.device)
+
+        # sample grid
+        r_samples = torch.arange(25, 225+self.roughness_step, self.roughness_step) / 255
+        m_samples = torch.arange(0., 1.+self.metallic_step, self.metallic_step)
+
+        grid_maps = {} # change map size into: gs, bs, h, w, c
+        grid_maps['basecolor'] = maps['basecolor'][None].permute(0,1,3,4,2)
+        grid_maps['normal'] = maps['normal'][None].permute(0,1,3,4,2)
+        r_values = r_samples[:,None].repeat(1,len(m_samples)).reshape(-1,1,1,1,1).to(maps['basecolor'])
+        m_values = m_samples[None].repeat(len(r_samples),1).reshape(-1,1,1,1,1).to(maps['basecolor'])
+        # split into chunks to avoid OOM
+        chunk_size = 25
+        rgb_list, r_list, m_list = [], [], []
+        for _r, _m in zip(torch.split(r_values, chunk_size), torch.split(m_values, chunk_size)):
+            grid_maps['roughness'], grid_maps['metallic'] = _r, _m
+            _rgb = self.compute_render(grid_maps, cameras, pos, light)
+            loss = (render[None].permute(0,1,3,4,2) - _rgb).abs().sum(-1,keepdim=True)
+            min_idx = loss.argmin(dim=0,keepdim=True)
+            r_list.append(torch.gather(grid_maps['roughness'].flatten(), 0, min_idx.flatten()).reshape(min_idx.shape))
+            m_list.append(torch.gather(grid_maps['metallic'].flatten(), 0, min_idx.flatten()).reshape(min_idx.shape))
+            rgb_list.append(torch.gather(_rgb, 0, min_idx.repeat(1,1,1,1,3)))
+        rgb = torch.cat(rgb_list).permute(0,1,4,2,3)
+        roughness = torch.cat(r_list).permute(0,1,4,2,3)
+        metallic = torch.cat(m_list).permute(0,1,4,2,3)
+        loss = (render[None] - rgb).abs().sum(2,keepdim=True)
+        roughness = torch.gather(roughness, 0, loss.argmin(dim=0,keepdim=True))[0]
+        metallic = torch.gather(metallic, 0, loss.argmin(dim=0,keepdim=True))[0]
+        torch.cuda.empty_cache()
+        if seperate:
+            return roughness, metallic
+        else:
+            out = torch.cat([roughness, metallic, torch.zeros_like(roughness)], dim=1)
+            return out
+
+
+    @torch.no_grad()
+    def compute_render(self, maps, camera_position, pos, light):
+        '''
+            maps: gs, bs, h, w, c (gs: the number of grids)
+        '''
+        def cos(x, y): 
+            return torch.clamp((x*y).sum(-1, keepdim=True), min=0, max=1)
+        
+        # pre-process
+        albedo = srgb_to_rgb(maps['basecolor'])
+        normal = maps['normal'].clone()
+        normal[..., :2] = normal[..., [1,0]]
+        N = Fn.normalize((normal - 0.5) * 2.0, dim=-1, eps=1e-6)
+        roughness = maps['roughness']
+        metallic = maps['metallic']
+        V = Fn.normalize(camera_position - pos, dim=-1, eps=1e-6).repeat(1,1,1,1,1).to(self.device)
+        irradiance, L = light(pos)
+        irradiance, L = irradiance.repeat(1,1,1,1,1).to(self.device), L.repeat(1,1,1,1,1).to(self.device)
+        # rendering
+        H = Fn.normalize(L+V, dim=-1, eps=1e-6)
+        f0 = torch.ones_like(albedo).to(self.device) * 0.04
+        F0 = torch.lerp(f0, albedo, metallic)
+        F = fresnelSchlick(cos(H,V), F0)   
+        ks = F
+
+        diffuse = (1-ks) * albedo / torch.pi
+        diffuse *= 1-metallic
+
+        NDF = DistributionGGX(cos(N,H), roughness) 
+        G = GeometrySchlickGGX(cos(N,L), roughness) * GeometrySchlickGGX(cos(N,V), roughness)
+
+        numerator = NDF * G * F
+        denominator = 4.0 * cos(N,V) * cos(N,L) + 1e-3
+        specular = numerator / denominator
+        ambient = 0.3 * albedo
+
+        rgb = (diffuse + specular) * irradiance * cos(N,L) + ambient
+
+        return rgb
+    
+    def forward(self, maps:dict):
+        # prepare
+        bs = maps['render'].shape[0]
+        self.sd.scheduler.set_timesteps(1)
+        t = self.sd.scheduler.timesteps[0]
+        # chain processing
+        pred, pred_latent, arxiv_latent = {}, {}, {}
+        for kout, info in self.chain.items():
+            info = info.split("_")
+            keys, ids = info[:-1], info[-1]
+            # Swap active LEGO blocks
+            self.sd.unet.down_blocks[0] = self.sd.unet.FirstDownBlocks[kout]
+            self.sd.unet.up_blocks[-1] = self.sd.unet.LastUpBlocks[kout]
+            # Eq.2, summing input latents
+            in_latent = 0
+            for k, i in zip(keys, ids):
+                if i=="0":
+                    if not k in arxiv_latent.keys(): arxiv_latent[k] = self.sd.encode_imgs_deterministic(maps[k])
+                    zx = arxiv_latent[k]
+                else:
+                    zx = pred_latent[k]
+                in_latent += self.sd.unet.ConvIns[k](zx)
+            in_latent = in_latent / len(keys)
+            # single-step denoising
+            embs = self.produce_embeddings(kout, bs)
+            out_latent = self.sd.unet(in_latent, t, **embs)[0]
+            out_latent = self.sd.unet.ConvOuts[kout](out_latent)            
+            pred_latent[kout] = self.sd.scheduler.step(out_latent, t, torch.zeros_like(zx)).pred_original_sample
+            pred[kout] = self.sd.decode_latents(pred_latent[kout]).float()
+            # compute intermediate representations
+            if self.chain_type in ["chord"] and kout == "basecolor":
+                pred['approxIrr'] = self.compute_approxIrr(maps['render'], pred['basecolor'])
+                pred_latent['approxIrr'] = self.sd.encode_imgs_deterministic(pred['approxIrr'])
+            if self.chain_type in ["chord"] and kout == "normal":
+                pred['approxRM'] = self.compute_approxRouMet(maps['render'], pred, seperate=False)
+                pred_latent['approxRM'] = self.sd.encode_imgs_deterministic(pred['approxRM'])
+
+        return pred     
+    
+    @torch.no_grad()
+    def produce_embeddings(self, key, batch_size):
+        if key not in self.text_emb.keys():
+            self.text_emb[key] = self.sd.encode_text(self.prompts[key], "max_length")
+        prompt_emb = self.text_emb[key].expand(batch_size, -1, -1)
+        return { "encoder_hidden_states": prompt_emb }

chord/module/light.py

@@ -0,0 +1,96 @@
+import torch 
+from typing import Optional
+import torch.nn.functional as Fn
+import math
+import copy
+
+from . import register
+from .base import Base
+
+class BaseLight(Base):
+    """
+    Base class for light models. 
+    """
+    
+    def setup(self):
+        pass
+
+    def forward(self, x: Optional[torch.Tensor] = None):
+        """
+        Get the light intensity.
+
+        Args:
+        x: positions of shape (..., 3). 
+
+        Returns:
+        color: radiance intensity of shape (..., 3)
+        d: directions of shape (..., 3).
+        """
+        raise NotImplementedError
+    
+
+@register("point-light")
+class PointLight(BaseLight):
+    """Point light definitions
+    """
+    def setup(self):
+        """Initialize point light.
+
+        Args:
+            position (float, float, float): World coordinate of the light.
+            color (float, float, float): Light color in (R, G, B).
+            power (float): Light power, it will be directly multiplied to each color channel.
+        """
+        position = self.config.get("position", [0., 0., 10.])
+        color = self.config.get("color", [23.47, 21.31, 20.79])
+        power = self.config.get("power", 10.) 
+
+        self.register_buffer("position", torch.tensor(position))
+        self.register_buffer("color", torch.tensor(color) * power)
+
+    def forward(self, x: Optional[torch.Tensor] = None):
+        """Compute light radiance and direction.
+
+        Args:
+            x : World coordinate of the interacting surface. [B, H, W, 3]
+        Returns:
+            color: radiance intensity of shape [B, H, W, 3]
+            d: directions of shape [B, H, W, 3], V = (light_pos - world_pos)
+        """
+        distance    = torch.norm(self.position - x, dim=-1, keepdim=True)
+        attenuation = 1.0 / (distance ** 2)
+        radiance    = self.color * attenuation
+        direction = Fn.normalize(self.position - x, dim=-1)
+        return radiance, direction
+    
+@register("distant-light")
+class DistantLight(BaseLight):
+    """Distant light definitions
+    """
+    def setup(self):
+        """Initialize distant light.
+
+        Args:
+            direction (float, float, float):The direction of light vector.
+            color (float, float, float): Light color in (R, G, B).
+            power (float): Light power, it will be directly multiplied to each color channel.
+        """
+        direction = self.config.get("direction", [0., 0., 1.])
+        color = self.config.get("color", [23.47, 21.31, 20.79])
+        power = self.config.get("power", 0.1)
+
+        self.register_buffer("color", torch.tensor(color) * power)
+        self.register_buffer("direction", Fn.normalize(torch.tensor(direction), dim=0))
+
+    def forward(self, x: Optional[torch.Tensor] = None):
+        """Compute light radiance and direction.
+
+        Args:
+            x : World coordinate of the interacting surface. [B, H, W, 3]
+        Returns:
+            color: radiance intensity of shape [B, H, W, 3]
+            d: directions of shape [B, H, W, 3]
+        """
+        radiance = self.color.repeat(*x.shape[:-1], 1)
+        direction = self.direction.repeat(*x.shape[:-1], 1)
+        return radiance, direction

chord/module/stable_diffusion.py

@@ -0,0 +1,105 @@
+import torch
+from torchvision.transforms import v2
+
+from diffusers import UNet2DConditionModel, AutoencoderKL, DDIMScheduler
+from transformers import CLIPTextModel, CLIPTextConfig, CLIPTokenizer
+
+from . import register
+from .base import Base
+
+
+def apply_padding(model, mode):
+    for layer in [layer for _, layer in model.named_modules() if isinstance(layer, torch.nn.Conv2d)]:
+        if mode == 'circular':
+            layer.padding_mode = 'circular'
+        else:
+            layer.padding_mode = 'zeros'
+    return model
+
+def freeze(model):
+    model = model.eval()
+    for param in model.parameters():
+        param.requires_grad = False
+    return model
+
+@register("stable_diffusion")
+class StableDiffusion(Base):
+    def setup(self):
+        hf_key = self.config.get("hf_key", None)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        fp16 = self.config.get("fp16", True)
+        self.dtype = torch.bfloat16 if fp16 else torch.float32
+        vae_padding = self.config.get("vae_padding", "zeros")
+
+        self.sd_version = self.config.get("version", 2.1)
+        local_files_only = False
+        if hf_key is not None:
+            print(f"[INFO] using hugging face custom model key: {hf_key}")
+            model_key = hf_key
+            local_files_only = True
+        elif str(self.sd_version) == "2.1":
+            # model_key = "stabilityai/stable-diffusion-2-1"
+            # StabilityAI deleted the original 2.1 model from HF, use a community version
+            model_key = "RedbeardNZ/stable-diffusion-2-1-base"
+        else:
+            raise ValueError(
+                f"Stable-diffusion version {self.sd_version} not supported."
+            )
+
+        # Load components separately to avoid download unnecessary weights
+        # 1. UNet (diffusion backbone)
+        unet_config = UNet2DConditionModel.load_config(model_key, subfolder="unet")
+        self.unet = UNet2DConditionModel.from_config(unet_config, local_files_only=local_files_only)
+        self.unet.to(self.device, dtype=self.dtype).eval()
+        # 2. VAE (image autoencoder)
+        vae_config = AutoencoderKL.load_config(model_key, subfolder="vae")
+        self.vae = AutoencoderKL.from_config(vae_config, local_files_only=local_files_only)
+        self.vae.to(self.device, dtype=self.dtype).eval()
+        self.vae = apply_padding(freeze(self.vae), vae_padding)
+        # 3. Text encoder (CLIP)
+        text_encoder_config = CLIPTextConfig.from_pretrained(model_key, subfolder="text_encoder", local_files_only=local_files_only)
+        self.text_encoder = CLIPTextModel(text_encoder_config)
+        self.text_encoder.to(self.device, dtype=self.dtype).eval()
+        # 4. Tokenizer (CLIP tokenizer, this one has vocab so from_pretrained is needed)
+        self.tokenizer = CLIPTokenizer.from_pretrained(model_key, subfolder="tokenizer", local_files_only=local_files_only)
+        # 5. Scheduler
+        scheduler_config = DDIMScheduler.load_config(model_key, subfolder="scheduler")
+        scheduler_config["prediction_type"] = "v_prediction"
+        scheduler_config["timestep_spacing"] = "trailing"
+        scheduler_config["rescale_betas_zero_snr"] = True
+        self.scheduler = DDIMScheduler.from_config(scheduler_config)
+
+    def encode_text(self, prompt, padding_mode="do_not_pad"):
+        # prompt: [str]
+        inputs = self.tokenizer(
+            prompt,
+            padding=padding_mode,
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt",
+        )
+        embeddings = self.text_encoder(inputs.input_ids.to(self.device))[0]
+        return embeddings    
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+        imgs = self.vae.decode(latents).sample
+        imgs = (imgs / 2 + 0.5).clamp(0, 1)
+        return imgs
+
+    def encode_imgs(self, imgs):
+        if imgs.shape[1] == 1: # for grayscale maps
+            imgs = v2.functional.grayscale_to_rgb(imgs)
+        imgs = 2 * imgs - 1
+        posterior = self.vae.encode(imgs).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor
+        return latents
+    
+    def encode_imgs_deterministic(self, imgs):
+        if imgs.shape[1] == 1: # for grayscale maps
+            imgs = v2.functional.grayscale_to_rgb(imgs)
+        imgs = 2 * imgs - 1
+        h = self.vae.encoder(imgs)
+        moments = self.vae.quant_conv(h)
+        mean, logvar = torch.chunk(moments, 2, dim=1)
+        latents = mean * self.vae.config.scaling_factor
+        return latents

chord/normal_utils.py

@@ -0,0 +1,289 @@
+"""
+Normal Map Utilities
+
+Derives height and ambient occlusion from normal maps.
+
+Height: Frankot-Chellappa algorithm (FFT-based gradient integration)
+AO: Divergence of normal vectors (concave areas have higher occlusion)
+"""
+
+import torch
+import torch.nn.functional as F
+
+
+def normal_to_gradient(normal: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Extract surface gradients from a normal map.
+
+    Normal maps store (nx, ny, nz) where the surface gradient is:
+    - dz/dx = -nx/nz
+    - dz/dy = -ny/nz
+
+    For normalized normals in [0,1] range (0.5 = neutral):
+    - nx = (R - 0.5) * 2
+    - ny = (G - 0.5) * 2
+
+    Args:
+        normal: Normal map tensor (B, 3, H, W) or (3, H, W), range [0, 1]
+
+    Returns:
+        grad_x, grad_y: Surface gradients
+    """
+    if normal.dim() == 3:
+        normal = normal.unsqueeze(0)
+
+    # Convert from [0,1] to [-1,1]
+    nx = (normal[:, 0:1, :, :] - 0.5) * 2.0
+    ny = (normal[:, 1:2, :, :] - 0.5) * 2.0
+
+    return nx, ny
+
+
+def frankot_chellappa(grad_x: torch.Tensor, grad_y: torch.Tensor) -> torch.Tensor:
+    """
+    Reconstruct height map from gradients using Frankot-Chellappa algorithm.
+
+    This uses FFT-based integration to find a height field Z such that:
+    dZ/dx ≈ grad_x and dZ/dy ≈ grad_y
+
+    Args:
+        grad_x: X gradient (B, 1, H, W)
+        grad_y: Y gradient (B, 1, H, W)
+
+    Returns:
+        Height map (B, 1, H, W), normalized to [0, 1]
+    """
+    b, _, h, w = grad_x.shape
+    device = grad_x.device
+    dtype = grad_x.dtype
+
+    # Create frequency grids
+    rows_scale = (torch.arange(h, device=device, dtype=dtype) - (h // 2 + 1)) / (h - h % 2)
+    cols_scale = (torch.arange(w, device=device, dtype=dtype) - (w // 2 + 1)) / (w - w % 2)
+
+    v_grid, u_grid = torch.meshgrid(rows_scale, cols_scale, indexing='ij')
+    u_grid = torch.fft.ifftshift(u_grid)
+    v_grid = torch.fft.ifftshift(v_grid)
+
+    # FFT of gradients
+    grad_x_F = torch.fft.fft2(grad_x.squeeze(1))
+    grad_y_F = torch.fft.fft2(grad_y.squeeze(1))
+
+    # Integrate in frequency domain
+    # Z_F = (-j*u*Gx - j*v*Gy) / (u^2 + v^2)
+    numerator = (-1j * u_grid * grad_x_F) + (-1j * v_grid * grad_y_F)
+    denominator = (u_grid ** 2) + (v_grid ** 2) + 1e-16
+    Z_F = numerator / denominator
+    Z_F[:, 0, 0] = 0.0  # Set DC component to zero
+
+    # Inverse FFT to get height
+    Z = torch.real(torch.fft.ifft2(Z_F)).unsqueeze(1)
+
+    # Normalize to [0, 1]
+    Z_min = Z.amin(dim=(2, 3), keepdim=True)
+    Z_max = Z.amax(dim=(2, 3), keepdim=True)
+    Z = (Z - Z_min) / (Z_max - Z_min + 1e-8)
+
+    return Z
+
+
+def make_seamless_gradients(grad_x: torch.Tensor, grad_y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Create seamless gradient field by mirroring.
+
+    Concatenates 4 flipped copies to create a seamless tileable field,
+    doubling the dimensions.
+
+    Args:
+        grad_x, grad_y: Gradients (B, 1, H, W)
+
+    Returns:
+        Seamless gradients (B, 1, 2H, 2W)
+    """
+    # Flip operations
+    grad_x_hflip = torch.flip(grad_x, dims=[3])  # Horizontal flip
+    grad_x_vflip = torch.flip(grad_x, dims=[2])  # Vertical flip
+    grad_x_hvflip = torch.flip(grad_x, dims=[2, 3])  # Both
+
+    grad_y_hflip = torch.flip(grad_y, dims=[3])
+    grad_y_vflip = torch.flip(grad_y, dims=[2])
+    grad_y_hvflip = torch.flip(grad_y, dims=[2, 3])
+
+    # Create 2x2 grid with appropriate sign flips for seamless tiling
+    grad_x_top = torch.cat([grad_x, -grad_x_hflip], dim=3)
+    grad_x_bottom = torch.cat([grad_x_vflip, -grad_x_hvflip], dim=3)
+    new_grad_x = torch.cat([grad_x_top, grad_x_bottom], dim=2)
+
+    grad_y_top = torch.cat([grad_y, grad_y_hflip], dim=3)
+    grad_y_bottom = torch.cat([-grad_y_vflip, -grad_y_hvflip], dim=3)
+    new_grad_y = torch.cat([grad_y_top, grad_y_bottom], dim=2)
+
+    return new_grad_x, new_grad_y
+
+
+def normal_to_height(
+    normal: torch.Tensor,
+    seamless: bool = False,
+) -> torch.Tensor:
+    """
+    Convert normal map to height map using Frankot-Chellappa algorithm.
+
+    Args:
+        normal: Normal map (B, 3, H, W) or (3, H, W), range [0, 1]
+        seamless: If True, input is assumed tileable and output will tile seamlessly
+
+    Returns:
+        Height map (B, 1, H, W) or (1, H, W), range [0, 1]
+    """
+    squeeze = normal.dim() == 3
+    if squeeze:
+        normal = normal.unsqueeze(0)
+
+    b, _, h, w = normal.shape
+
+    # Extract gradients from normal map
+    grad_x, grad_y = normal_to_gradient(normal)
+
+    # For non-seamless, use mirror trick to avoid edge artifacts
+    if not seamless:
+        grad_x, grad_y = make_seamless_gradients(grad_x, grad_y)
+
+    # Reconstruct height via FFT integration
+    height = frankot_chellappa(-grad_x, grad_y)
+
+    # Crop back to original size if we expanded
+    if not seamless:
+        height = height[:, :, :h, :w]
+
+    if squeeze:
+        height = height.squeeze(0)
+
+    return height
+
+
+def compute_divergence(normal: torch.Tensor) -> torch.Tensor:
+    """
+    Compute divergence of 2D normal field.
+
+    Divergence = dnx/dx + dny/dy
+
+    High positive divergence = concave (crevices) → more AO
+    Negative divergence = convex (bumps) → less AO
+
+    Args:
+        normal: Normal map (B, 3, H, W) or (3, H, W), range [0, 1]
+
+    Returns:
+        Divergence field (B, 1, H, W) or (1, H, W)
+    """
+    squeeze = normal.dim() == 3
+    if squeeze:
+        normal = normal.unsqueeze(0)
+
+    # Convert to [-1, 1] range
+    nx = (normal[:, 0:1, :, :] - 0.5) * 2.0
+    ny = (normal[:, 1:2, :, :] - 0.5) * 2.0
+
+    # Sobel-like kernels for derivatives
+    # dnx/dx
+    kernel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]],
+                           dtype=normal.dtype, device=normal.device).view(1, 1, 3, 3) / 8.0
+    # dny/dy
+    kernel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]],
+                           dtype=normal.dtype, device=normal.device).view(1, 1, 3, 3) / 8.0
+
+    # Compute derivatives with replicate padding
+    dnx_dx = F.conv2d(nx, kernel_x, padding=1)
+    dny_dy = F.conv2d(ny, kernel_y, padding=1)
+
+    divergence = dnx_dx + dny_dy
+
+    if squeeze:
+        divergence = divergence.squeeze(0)
+
+    return divergence
+
+
+def normal_to_ao(
+    normal: torch.Tensor,
+    strength: float = 1.0,
+    blur_radius: int = 3,
+) -> torch.Tensor:
+    """
+    Approximate ambient occlusion from normal map using divergence.
+
+    Concave areas (crevices) have positive divergence and should be darker.
+    Convex areas (bumps) have negative divergence and should be lighter.
+
+    Args:
+        normal: Normal map (B, 3, H, W) or (3, H, W), range [0, 1]
+        strength: AO intensity multiplier
+        blur_radius: Gaussian blur radius for smoothing
+
+    Returns:
+        AO map (B, 1, H, W) or (1, H, W), range [0, 1]
+        1.0 = no occlusion, 0.0 = full occlusion
+    """
+    squeeze = normal.dim() == 3
+    if squeeze:
+        normal = normal.unsqueeze(0)
+
+    # Compute divergence
+    div = compute_divergence(normal)
+
+    # Apply gaussian blur to smooth the AO
+    if blur_radius > 0:
+        kernel_size = blur_radius * 2 + 1
+        sigma = blur_radius / 2.0
+
+        # Create 1D Gaussian kernel
+        x = torch.arange(kernel_size, dtype=normal.dtype, device=normal.device) - blur_radius
+        gauss_1d = torch.exp(-x ** 2 / (2 * sigma ** 2))
+        gauss_1d = gauss_1d / gauss_1d.sum()
+
+        # Separable 2D convolution
+        gauss_h = gauss_1d.view(1, 1, 1, -1)
+        gauss_v = gauss_1d.view(1, 1, -1, 1)
+
+        div = F.pad(div, (blur_radius, blur_radius, blur_radius, blur_radius), mode='replicate')
+        div = F.conv2d(div, gauss_h)
+        div = F.conv2d(div, gauss_v)
+
+    # Convert divergence to AO
+    # Positive divergence (concave) → darker → lower AO
+    # We want: high div → low AO, low div → high AO
+    # AO = 1 - clamp(div * strength, 0, 1)
+    ao = 1.0 - torch.clamp(div * strength, 0.0, 1.0)
+
+    # Ensure output is in [0, 1]
+    ao = torch.clamp(ao, 0.0, 1.0)
+
+    if squeeze:
+        ao = ao.squeeze(0)
+
+    return ao
+
+
+def derive_ao_and_height(
+    normal: torch.Tensor,
+    seamless: bool = False,
+    ao_strength: float = 2.0,
+    ao_blur: int = 5,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Derive both AO and height from a normal map.
+
+    Args:
+        normal: Normal map (B, 3, H, W) or (3, H, W), range [0, 1]
+        seamless: Whether the texture should tile seamlessly
+        ao_strength: AO intensity multiplier
+        ao_blur: Gaussian blur radius for AO smoothing
+
+    Returns:
+        ao: Ambient occlusion (1.0 = no occlusion)
+        height: Height map (0.0 = low, 1.0 = high)
+    """
+    ao = normal_to_ao(normal, strength=ao_strength, blur_radius=ao_blur)
+    height = normal_to_height(normal, seamless=seamless)
+
+    return ao, height

chord/util.py

@@ -0,0 +1,67 @@
+import torch
+
+def vector_dot(A: torch.Tensor, B: torch.Tensor, min=0.0) -> torch.Tensor:
+    return torch.clamp((A * B).sum(1, keepdim=True), min=min, max=1.0)
+
+def srgb_to_rgb(f: torch.Tensor) -> torch.Tensor:
+    return torch.where(f <= 0.04045, f / 12.92, torch.pow((torch.clamp(f, 0.04045) + 0.055) / 1.055, 2.4)).to(f.dtype)
+    
+def rgb_to_srgb(f: torch.Tensor) -> torch.Tensor:
+    return torch.where(f <= 0.0031308, f * 12.92, torch.pow(torch.clamp(f, 0.0031308), 1.0/2.4)*1.055 - 0.055).to(f.dtype)
+
+def tone_gamma(x: torch.Tensor) -> torch.Tensor:
+    x = 1 - torch.exp(-x)
+    return torch.pow(x, 1.0/2.2)
+
+# safe division for value range 0-1
+class safe_01_div(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, a, b):
+        ctx.save_for_backward(a, b)
+        return torch.div(a, torch.clamp(b, min=1e-4, max=1.0))
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        a, b = ctx.saved_tensors
+        grad_input = grad_output.clone()
+
+        return torch.div(1, torch.clamp(b, min=1e-4, max=1.0)) * grad_input, -1 * torch.div(a, torch.clamp(b, min=1e-2, max=1.0)**2) * grad_input
+
+
+def get_positions(h, w, real_size, use_pixel_centers=True) -> torch.Tensor:
+    pixel_center = 0.5 if use_pixel_centers else 0
+    i, j = torch.meshgrid(
+        torch.arange(h) + pixel_center,
+        torch.arange(w) + pixel_center,
+        indexing='ij'
+    )
+    if not isinstance(real_size, list):
+        real_size = [real_size] * 2
+    pos = torch.stack([(i / h - 0.5) * real_size[0], (j / w - 0.5) * real_size[1], torch.zeros_like(i)], dim=-1)
+    return pos
+    
+# N, H: （Bx3xHxW), roughness: (Bx1xHxW)
+# The "D", facet distribution function in Cook-Torrence model
+def DistributionGGX(cosNH, roughness):
+    a = roughness * roughness
+    a2 = a * a
+    cosNH2 = cosNH * cosNH
+    num = a2
+    denom = cosNH2 * (a2 - 1.0) + 1.0
+    denom = torch.pi * denom * denom
+    return num / denom
+
+# NdotV, roughness: (Bx1xHxW)
+def GeometrySchlickGGX(NdotV: torch.Tensor, roughness: torch.Tensor) -> torch.Tensor:
+    r = (roughness + 1.0)
+    k = (r*r) / 8.0
+
+    num   = NdotV
+    denom = NdotV * (1.0 - k) + k
+	
+    return num / denom
+
+# cosTheta, F0 (Bx1xHxW)
+# The "F"
+def fresnelSchlick(cosTheta: torch.Tensor, F0: torch.Tensor) -> torch.Tensor:
+    return F0 + (1.0 - F0) * torch.pow(1.0 - cosTheta, 5.0)

config/chord.yaml

@@ -0,0 +1,29 @@
+model:
+  name: chord
+  roughness_step: 5.
+  metallic_step: 1.
+  # format: "OutputMapName": ConvInInput1_ConvInInput2_{0/1}
+  # 0/1 stands for using gt/pred image;
+  chain_type: chord
+  chain_library:
+    chord:
+      basecolor: render_0
+      normal: render_approxIrr_01
+      rou_met: render_approxRM_01
+  rgbx_prompts:
+    basecolor: Basecolor
+    normal: Normal
+    roughness: Roughness
+    metallic: Metallic
+    irradiance: Irradiance
+    rou_met: Roughness and Metallic
+  prior_light:
+    name: distant-light
+    direction: [-1.0, -1.0, 1.0]   # Top-left corner towards bottom right
+    color: [23.47, 21.31, 20.79]
+    power: 0.1                 
+  stable_diffusion: 
+    name: stable_diffusion
+    fp16: true
+    vae_padding: circular
+    version: 2.1

requirements.txt

@@ -0,0 +1,13 @@
+--extra-index-url https://download.pytorch.org/whl/cu128
+torch
+torchvision
+huggingface_hub[hf_xet]
+diffusers==0.35.2
+transformers==4.57.1
+tokenizers==0.22.1
+safetensors
+omegaconf
+imageio
+gradio
+spaces
+python-dotenv

test_minecraft_pbr.py

@@ -0,0 +1,57 @@
+"""Quick test of Minecraft PBR conversion without running full model."""
+import torch
+from chord.minecraft_pbr import convert_to_labpbr, convert_to_bedrock
+from chord.normal_utils import derive_ao_and_height
+
+# Create dummy tensors (simulating model output)
+h, w = 256, 256
+basecolor = torch.rand(3, h, w)
+normal = torch.rand(3, h, w) * 0.5 + 0.25  # Center around 0.5
+roughness = torch.rand(1, h, w)
+metalness = torch.rand(1, h, w)
+
+print("Testing AO and height derivation from normal map...")
+ao, height = derive_ao_and_height(normal)
+print(f"  AO shape: {ao.shape}, range: [{ao.min():.3f}, {ao.max():.3f}]")
+print(f"  Height shape: {height.shape}, range: [{height.min():.3f}, {height.max():.3f}]")
+
+print("\nTesting Minecraft PBR (LabPBR format) conversion...")
+result = convert_to_labpbr(
+    basecolor=basecolor,
+    normal=normal,
+    roughness=roughness,
+    metalness=metalness,
+    derive_ao_height=True,
+)
+
+print(f"  Albedo: {result['albedo'].size} {result['albedo'].mode}")
+print(f"  Specular: {result['specular'].size} {result['specular'].mode}")
+print(f"  Normal: {result['normal'].size} {result['normal'].mode}")
+
+# Save LabPBR test outputs
+result['albedo'].save('test_albedo.png')
+result['specular'].save('test_specular_s.png')
+result['normal'].save('test_normal_n.png')
+print("Saved: test_albedo.png, test_specular_s.png, test_normal_n.png")
+
+print("\nTesting Minecraft PBR (Bedrock RTX format) conversion...")
+bedrock_result = convert_to_bedrock(
+    basecolor=basecolor,
+    normal=normal,
+    roughness=roughness,
+    metalness=metalness,
+    compute_sss=True,
+    compute_emission=True,
+)
+
+print(f"  Albedo: {bedrock_result['albedo'].size} {bedrock_result['albedo'].mode}")
+print(f"  MER: {bedrock_result['mer'].size} {bedrock_result['mer'].mode}")
+print(f"  Normal: {bedrock_result['normal'].size} {bedrock_result['normal'].mode}")
+
+# Save Bedrock test outputs
+bedrock_result['albedo'].save('test_bedrock_albedo.png')
+bedrock_result['mer'].save('test_bedrock_mer.png')
+bedrock_result['normal'].save('test_bedrock_normal.png')
+print("Saved: test_bedrock_albedo.png, test_bedrock_mer.png, test_bedrock_normal.png")
+
+print("\nAll tests passed!")