diff --git a/.gitmodules b/.gitmodules
new file mode 100644
index 0000000..cfefbd0
--- /dev/null
+++ b/.gitmodules
@@ -0,0 +1,3 @@
+[submodule "chord"]
+ path = chord
+ url = https://github.com/ubisoft/ubisoft-laforge-chord.git
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..cb6c031
--- /dev/null
+++ b/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.
diff --git a/README.md b/README.md
index 443fa08..7376f7d 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,47 @@
# ComfyUI-Chord
-ComfyUI Custom Node for Chord
+
+
+
+
+
+ComfyUI Custom Node for paper: **Chord: Chain of Rendering Decomposition for PBR Material Estimation from Generated Texture Images**
+
+## Installation
+
+1. Download and install ComfyUI. If you are new to ComfyUI, refer to the original [repository](https://github.com/comfyanonymous/ComfyUI) to get started. **Make sure you have the latest version.**
+
+2. Download the pretrained model from [Hugging Face](Ubisoft/ubisoft-laforge-chord) and place the **chord_v1.ckpt** file into the **ComfyUI/models/checkpoints** folder.
+
+3. Install the custom nodes from the [ComfyUI-Manager](https://github.com/Comfy-Org/ComfyUI-Manager) or by cloning this repository recursively in the custom nodes folder of ComfyUI:
+
+ ```shell
+ cd ./ComfyUI/custom_nodes
+ git clone --recursive https://github.com/ubisoft/ComfyUI-Chord.git
+ ```
+## Example Workflow
+
+
+
+You can load this workflow using the JSON file:
+`examples/chord_example.json`
+
+## 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.
\ No newline at end of file
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..1a05868
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,11 @@
+import os,sys
+sys.path.append(os.path.join(os.path.dirname(__file__), "chord"))
+
+from .nodes import ChordLoadModel, ChordMaterialEstimation, ChordNormalToHeight
+
+NODE_CLASS_MAPPINGS = {
+ "ChordLoadModel": ChordLoadModel,
+ "ChordMaterialEstimation": ChordMaterialEstimation,
+ "ChordNormalToHeight": ChordNormalToHeight,
+}
+__all__ = ["NODE_CLASS_MAPPINGS"]
\ No newline at end of file
diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000..c15f9c4
Binary files /dev/null and b/__pycache__/__init__.cpython-312.pyc differ
diff --git a/__pycache__/node_chord.cpython-312.pyc b/__pycache__/node_chord.cpython-312.pyc
new file mode 100644
index 0000000..50a18b6
Binary files /dev/null and b/__pycache__/node_chord.cpython-312.pyc differ
diff --git a/__pycache__/nodes.cpython-312.pyc b/__pycache__/nodes.cpython-312.pyc
new file mode 100644
index 0000000..c8d2f2e
Binary files /dev/null and b/__pycache__/nodes.cpython-312.pyc differ
diff --git a/__pycache__/normal_to_height.cpython-312.pyc b/__pycache__/normal_to_height.cpython-312.pyc
new file mode 100644
index 0000000..bd10a0d
Binary files /dev/null and b/__pycache__/normal_to_height.cpython-312.pyc differ
diff --git a/chord b/chord
new file mode 160000
index 0000000..3ce8358
--- /dev/null
+++ b/chord
@@ -0,0 +1 @@
+Subproject commit 3ce835852e3fff59993c7e3ec7f35487ca08ea7b
diff --git a/examples/chord_example.json b/examples/chord_example.json
new file mode 100644
index 0000000..c0c2b4f
--- /dev/null
+++ b/examples/chord_example.json
@@ -0,0 +1,412 @@
+{
+ "id": "39d23e25-903c-4f5f-88e1-3bbadf18b7df",
+ "revision": 0,
+ "last_node_id": 25,
+ "last_link_id": 38,
+ "nodes": [
+ {
+ "id": 25,
+ "type": "ChordMaterialEstimation",
+ "pos": [
+ 669.1222107554908,
+ 312.46323813855713
+ ],
+ "size": [
+ 270,
+ 118
+ ],
+ "flags": {},
+ "order": 2,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "model",
+ "type": "MODEL",
+ "link": 32
+ },
+ {
+ "name": "image",
+ "type": "IMAGE",
+ "link": 33
+ }
+ ],
+ "outputs": [
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 34
+ ]
+ },
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 35,
+ 38
+ ]
+ },
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 36
+ ]
+ },
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 37
+ ]
+ }
+ ],
+ "properties": {
+ "aux_id": "ubisoft/ComfyUI-Chord",
+ "ver": "cfadf742e1838ded9990d85e43650d55a96e7391",
+ "Node name for S&R": "ChordMaterialEstimation"
+ },
+ "widgets_values": []
+ },
+ {
+ "id": 14,
+ "type": "PreviewImage",
+ "pos": [
+ 1035.9285017533348,
+ 125.65524984290931
+ ],
+ "size": [
+ 140,
+ 246
+ ],
+ "flags": {},
+ "order": 3,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "images",
+ "type": "IMAGE",
+ "link": 34
+ }
+ ],
+ "outputs": [],
+ "title": "Basecolor",
+ "properties": {
+ "cnr_id": "comfy-core",
+ "ver": "0.3.70",
+ "Node name for S&R": "PreviewImage"
+ },
+ "widgets_values": []
+ },
+ {
+ "id": 16,
+ "type": "PreviewImage",
+ "pos": [
+ 1183.84326933871,
+ 126.43252808866265
+ ],
+ "size": [
+ 140,
+ 246
+ ],
+ "flags": {},
+ "order": 5,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "images",
+ "type": "IMAGE",
+ "link": 38
+ }
+ ],
+ "outputs": [],
+ "title": "Normal",
+ "properties": {
+ "cnr_id": "comfy-core",
+ "ver": "0.3.70",
+ "Node name for S&R": "PreviewImage"
+ },
+ "widgets_values": []
+ },
+ {
+ "id": 24,
+ "type": "PreviewImage",
+ "pos": [
+ 1035.5098787324794,
+ 407.8680029468901
+ ],
+ "size": [
+ 140,
+ 246
+ ],
+ "flags": {},
+ "order": 7,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "images",
+ "type": "IMAGE",
+ "link": 37
+ }
+ ],
+ "outputs": [],
+ "title": "Roughness",
+ "properties": {
+ "cnr_id": "comfy-core",
+ "ver": "0.3.70",
+ "Node name for S&R": "PreviewImage"
+ },
+ "widgets_values": []
+ },
+ {
+ "id": 15,
+ "type": "PreviewImage",
+ "pos": [
+ 1180,
+ 410
+ ],
+ "size": [
+ 140,
+ 246.00000000000006
+ ],
+ "flags": {},
+ "order": 6,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "images",
+ "type": "IMAGE",
+ "link": 36
+ }
+ ],
+ "outputs": [],
+ "title": "Metalness",
+ "properties": {
+ "cnr_id": "comfy-core",
+ "ver": "0.3.70",
+ "Node name for S&R": "PreviewImage"
+ },
+ "widgets_values": []
+ },
+ {
+ "id": 20,
+ "type": "ChordNormalToHeight",
+ "pos": [
+ 1333.731988297213,
+ 253.143266404242
+ ],
+ "size": [
+ 270,
+ 58
+ ],
+ "flags": {},
+ "order": 4,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "normal",
+ "type": "IMAGE",
+ "link": 35
+ }
+ ],
+ "outputs": [
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 25
+ ]
+ }
+ ],
+ "properties": {
+ "aux_id": "ubisoft/ComfyUI-Chord",
+ "ver": "cfadf742e1838ded9990d85e43650d55a96e7391",
+ "Node name for S&R": "ChordNormalToHeight"
+ },
+ "widgets_values": []
+ },
+ {
+ "id": 21,
+ "type": "PreviewImage",
+ "pos": [
+ 1455.691637213612,
+ 354.32688389208334
+ ],
+ "size": [
+ 140,
+ 246.00000000000006
+ ],
+ "flags": {},
+ "order": 8,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "images",
+ "type": "IMAGE",
+ "link": 25
+ }
+ ],
+ "outputs": [],
+ "title": "Height",
+ "properties": {
+ "cnr_id": "comfy-core",
+ "ver": "0.3.70",
+ "Node name for S&R": "PreviewImage"
+ },
+ "widgets_values": []
+ },
+ {
+ "id": 11,
+ "type": "ChordLoadModel",
+ "pos": [
+ 374.58235480468545,
+ 141.5289281758196
+ ],
+ "size": [
+ 270,
+ 58
+ ],
+ "flags": {},
+ "order": 0,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [
+ {
+ "name": "MODEL",
+ "type": "MODEL",
+ "links": [
+ 32
+ ]
+ }
+ ],
+ "properties": {
+ "aux_id": "ubisoft/ComfyUI-Chord",
+ "ver": "cfadf742e1838ded9990d85e43650d55a96e7391",
+ "Node name for S&R": "ChordLoadModel"
+ },
+ "widgets_values": [
+ [
+ "chord_v1.ckpt"
+ ]
+ ]
+ },
+ {
+ "id": 13,
+ "type": "LoadImage",
+ "pos": [
+ 370,
+ 330
+ ],
+ "size": [
+ 270,
+ 314
+ ],
+ "flags": {},
+ "order": 1,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 33
+ ]
+ },
+ {
+ "name": "MASK",
+ "type": "MASK",
+ "links": null
+ }
+ ],
+ "properties": {
+ "cnr_id": "comfy-core",
+ "ver": "0.3.70",
+ "Node name for S&R": "LoadImage"
+ },
+ "widgets_values": [
+ "ref_render.png",
+ "image"
+ ]
+ }
+ ],
+ "links": [
+ [
+ 25,
+ 20,
+ 0,
+ 21,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 32,
+ 11,
+ 0,
+ 25,
+ 0,
+ "MODEL"
+ ],
+ [
+ 33,
+ 13,
+ 0,
+ 25,
+ 1,
+ "IMAGE"
+ ],
+ [
+ 34,
+ 25,
+ 0,
+ 14,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 35,
+ 25,
+ 1,
+ 20,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 36,
+ 25,
+ 2,
+ 15,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 37,
+ 25,
+ 3,
+ 24,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 38,
+ 25,
+ 1,
+ 16,
+ 0,
+ "IMAGE"
+ ]
+ ],
+ "groups": [],
+ "config": {},
+ "extra": {
+ "ds": {
+ "scale": 1.1102851430715899,
+ "offset": [
+ -148.2607809733155,
+ 51.621946976768974
+ ]
+ },
+ "frontendVersion": "1.28.8"
+ },
+ "version": 0.4
+}
\ No newline at end of file
diff --git a/examples/chord_example.png b/examples/chord_example.png
new file mode 100644
index 0000000..7da5c6c
Binary files /dev/null and b/examples/chord_example.png differ
diff --git a/nodes.py b/nodes.py
new file mode 100644
index 0000000..0d59c46
--- /dev/null
+++ b/nodes.py
@@ -0,0 +1,135 @@
+import os
+import torch
+from omegaconf import OmegaConf
+from torchvision.transforms import v2
+from .normal_to_height import normal_to_height
+
+# Modules from ComfyUI
+import folder_paths
+from comfy_api.latest import ComfyExtension, io
+
+# Modules from ComfyUI-Chord
+from chord import ChordModel
+from chord.io import read_image, save_maps
+
+def apply_padding(model, mode):
+ for layer in [layer for _, layer in model.named_modules() if isinstance(layer, torch.nn.Conv2d)]:
+ layer.padding_mode = mode
+
+def apply_circular_padding(model):
+ if hasattr(model, 'sd'):
+ apply_padding(model.sd.vae, 'circular')
+ apply_padding(model.sd.unet, 'circular')
+ else:
+ apply_padding(model, 'circular')
+
+class ChordLoadModel(io.ComfyNode):
+ """Node to load Chord Model"""
+
+ @classmethod
+ def define_schema(cls) -> io.Schema:
+ return io.Schema(
+ node_id="ChordLoadModel",
+ display_name="Chord - Load Model",
+ category="Chord",
+ description="A node to load a Chord model.",
+ inputs=[
+ io.Combo.Input(
+ "ckpt_name",
+ options=([x for x in folder_paths.get_filename_list("checkpoints") if x.endswith("ckpt")], {})
+ ),
+ ],
+ outputs=[io.Model.Output("model")],
+ )
+
+ @classmethod
+ def execute(cls, ckpt_name) -> io.NodeOutput:
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ if type(ckpt_name) is list:
+ ckpt_name = ckpt_name[0]
+ ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
+ config = OmegaConf.load(os.path.join(os.path.dirname(__file__), "chord/config/chord.yaml"))
+ model = ChordModel(config)
+ ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+ try:
+ model.load_state_dict(ckpt["state_dict"])
+ except RuntimeError as e:
+ raise RuntimeError('Failed to load model, check if the checkpoint file is correct.\n{}'.format(repr(e)))
+ model.to(device)
+ model.eval()
+ return io.NodeOutput(model)
+
+class ChordMaterialEstimation(io.ComfyNode):
+ """Chord Material Estimation Node"""
+
+ @classmethod
+ def define_schema(cls) -> io.Schema:
+ return io.Schema(
+ node_id="ChordMaterialEstimation",
+ display_name="Chord - Material Estimation",
+ category="Chord",
+ description="A node to estimate material maps from a texture image using the Chord model.",
+ inputs=[
+ io.Model.Input(
+ "model",
+ tooltip="The Chord model used to estimate material."
+ ),
+ io.Image.Input(
+ "image",
+ ),
+ ],
+ outputs=[
+ io.Image.Output("basecolor"),
+ io.Image.Output("normal"),
+ io.Image.Output("roughness"),
+ io.Image.Output("metalness"),
+ ]
+ )
+
+ @classmethod
+ def execute(cls, model, image) -> io.NodeOutput:
+ device = next(model.parameters()).device
+ apply_circular_padding(model)
+ image = image.permute(0,3,1,2).to(device)
+ ori_h, ori_w = image.shape[-2:]
+ x = v2.Resize(size=(1024, 1024), antialias=True)(image)
+ with torch.no_grad() as no_grad, torch.autocast(device_type=device.type) as amp:
+ output = model(x)
+ for key in output.keys():
+ output[key] = v2.Resize(size=(ori_h, ori_w), antialias=True)(output[key])
+ if output[key].ndim == 4:
+ output[key] = output[key].permute(0,2,3,1)
+ return io.NodeOutput(output['basecolor'], output['normal'], output['roughness'], output['metalness'])
+
+class ChordNormalToHeight(io.ComfyNode):
+ """Integrate normal map to height map using Poisson solver with overlapping subregions."""
+
+ @classmethod
+ def define_schema(cls) -> io.Schema:
+ return io.Schema(
+ node_id="ChordNormalToHeight",
+ display_name="Chord - Normal to Height",
+ category="Chord",
+ description="A node to convert normal map to height map using Poisson solver with overlapping subregions.",
+ inputs=[
+ io.Image.Input(
+ "normal",
+ ),
+ ],
+ outputs=[
+ io.Image.Output("height"),
+ ],
+ )
+
+ @classmethod
+ def execute(cls, normal) -> io.NodeOutput:
+ normal = normal.permute(0,3,1,2)
+ height_var_threshold = 5e-4
+ ori_h, ori_w = normal.shape[-2:]
+ x = v2.Resize(size=(1024, 1024), antialias=True)(normal)
+ height = normal_to_height(x)[None, None].squeeze(1)
+ if height.var() < height_var_threshold and height.var() > 0:
+ height = normal_to_height(x, skip_normalize_normal=True)[None, None].squeeze(1)
+
+ height = v2.Resize(size=(ori_h, ori_w), antialias=True)(height)
+ return io.NodeOutput(height)
\ No newline at end of file
diff --git a/normal_to_height.py b/normal_to_height.py
new file mode 100644
index 0000000..861c7e9
--- /dev/null
+++ b/normal_to_height.py
@@ -0,0 +1,175 @@
+import torch
+import torch.nn.functional as F
+from torch.signal.windows import hann
+import torch.fft as fft_module
+
+def compute_divergence(fx: torch.tensor, fy: torch.tensor) -> torch.tensor:
+ div_x = F.pad(fx[:, :, 1:] - fx[:, :, :-1], (0, 1, 0, 0), mode='constant')
+ div_y = F.pad(fy[:, 1:, :] - fy[:, :-1, :], (0, 0, 0, 1), mode='constant')
+ return div_x + div_y
+
+def solve_poisson_fft(div: torch.tensor, h: int, w: int) -> torch.tensor:
+ fft_div = fft_module.fft2(div)
+
+ kx = fft_module.fftfreq(2*w, device=div.device) * 2 * torch.pi
+ ky = fft_module.fftfreq(2*h, device=div.device) * 2 * torch.pi
+ kx, ky = torch.meshgrid(kx, ky, indexing='xy')
+ epsilon = 1e-9
+ denom = 4 - 2 * torch.cos(kx) - 2 * torch.cos(ky)
+ denom = torch.where(torch.abs(denom) > epsilon, denom, epsilon)
+
+ height_map_full = torch.real(fft_module.ifft2(fft_div / denom))
+ return torch.nan_to_num(height_map_full[:, :h, :w])
+
+def apply_window_function(gradient: torch.tensor) -> torch.tensor:
+ hann_window = hann(gradient.shape[-2], device=gradient.device)[:, None] * hann(gradient.shape[-1], device=gradient.device)[None, :]
+ return gradient * hann_window
+
+def compute_height(normal_map: torch.tensor, epsilon: float = 1e-8) -> torch.tensor:
+ h, w = normal_map.shape[-2:]
+ nz = normal_map[:, 2]
+ nz_safe = torch.where(torch.abs(nz) > epsilon, nz, epsilon)
+ fx = normal_map[:, 0] / nz_safe
+ fy = normal_map[:, 1] / nz_safe
+
+ fx = apply_window_function(fx)
+ fy = apply_window_function(fy)
+
+ div = compute_divergence(fx, fy)
+ div = F.pad(div, (0, w, 0, h), mode='constant')
+ height_map = solve_poisson_fft(div, h, w)
+ return height_map - torch.mean(height_map)
+
+def extend_normal_map(normal_map: torch.tensor, region_size: int) -> torch.tensor:
+ larger_normal_map = F.pad(normal_map, (region_size, region_size, region_size, region_size), mode='circular')
+ return larger_normal_map
+
+def define_subregions(
+ h: int,
+ w: int,
+ min_region_size: int=128,
+ overlap_factor: float = 0.5
+) -> list:
+ step_size = int(min_region_size - min_region_size * overlap_factor)
+ if step_size <= 0:
+ step_size = min_region_size
+ overlap_size = int(min_region_size * overlap_factor)
+ subregions = []
+ for y in range(0, h, step_size):
+ for x in range(0, w, step_size):
+ y_end = min(y + min_region_size + overlap_size, h)
+ x_end = min(x + min_region_size + overlap_size, w)
+ subregions.append((y, y_end, x, x_end))
+ return subregions
+
+def cosine_smoothing(x: torch.tensor) -> torch.tensor:
+ return 0.5 * (1 - torch.cos(torch.pi * x))
+
+def create_subregions(
+ normal_map: torch.tensor,
+ subregions: list,
+) -> list:
+ results = []
+ for region in subregions:
+ y, y_end, x, x_end = region
+ sub_map = normal_map[:, y:y_end, x:x_end]
+ results.append(sub_map)
+ return results
+
+def cosine_blending_update_weight(
+ sub_weight_map: torch.tensor,
+ subregion,
+ shape,
+ region_size: int = 128,
+) -> None:
+ h, w = sub_weight_map.shape[-2:]
+ d = sub_weight_map.device
+ y_start, y_end, x_start, x_end = subregion
+ if y_start > 0:
+ overlap = min(region_size, h)
+ y_smooth = cosine_smoothing(torch.linspace(0, 1, overlap, device=d))[:, None]
+ sub_weight_map[:overlap, :] *= y_smooth
+ if y_end < shape[0]:
+ overlap = min(region_size, h)
+ y_smooth = cosine_smoothing(torch.linspace(1, 0, overlap, device=d))[:, None]
+ sub_weight_map[-overlap:, :] *= y_smooth
+ if x_start > 0:
+ overlap = min(region_size, w)
+ x_smooth = cosine_smoothing(torch.linspace(0, 1, overlap, device=d))
+ sub_weight_map[:, :overlap] *= x_smooth
+ if x_end < shape[1]:
+ overlap = min(region_size, w)
+ x_smooth = cosine_smoothing(torch.linspace(1, 0, overlap, device=d))
+ sub_weight_map[:, -overlap:] *= x_smooth
+
+def combine_sub_height_maps(
+ height_maps: list,
+ subregions: list,
+ h: int,
+ w: int,
+ epsilon: float = 1e-8,
+) -> torch.tensor:
+ height_map = torch.zeros((h, w), device=height_maps[0].device)
+ weight_map = torch.zeros((h, w), device=height_maps[0].device)
+
+ for subregion, sub_height_map in zip(subregions, height_maps):
+ y_start, y_end, x_start, x_end = subregion
+ y_slice = slice(y_start, y_end)
+ x_slice = slice(x_start, x_end)
+
+ sub_weight_map = torch.ones_like(sub_height_map)
+ cosine_blending_update_weight(
+ sub_weight_map,
+ subregion,
+ (h, w)
+ )
+
+ height_map[y_slice, x_slice] += sub_height_map * sub_weight_map
+ weight_map[y_slice, x_slice] += sub_weight_map
+
+ return height_map / (weight_map + epsilon)
+
+def crop_height_map(height_map: torch.tensor, original_shape: tuple, region_size: int=128) -> torch.tensor:
+ original_height, original_width = original_shape
+ return height_map[
+ region_size:region_size + original_height,
+ region_size:region_size + original_width
+ ]
+
+def normalize_height_map(data: torch.tensor, eps: float = 1e-8) -> torch.tensor:
+ return (data - data.min()) / (data.max() - data.min() + eps)
+
+def normal_to_height(
+ normal_map: torch.tensor,
+ subdivisions: int = 16,
+ min_region_size: int = 128,
+ skip_normalize_normal: bool = False,
+) -> torch.tensor:
+ '''
+ Args:
+ normal_map: torch.tensor(1, 3, H, W)
+ subdivisions: int, subdivision level at each edge
+ min_region_size: int, minimal region size
+ skip_normalize_normal: bool, if skip normalization of input normal map
+ '''
+ if normal_map.dim() == 4:
+ try: assert normal_map.shape[0] == 1
+ except: breakpoint()
+ normal_map = normal_map.squeeze(0)
+ h, w = normal_map.shape[-2:]
+ if not skip_normalize_normal:
+ normal_map = F.normalize(normal_map * 2.0 - 1.0, dim=0)
+ region_size = min(
+ max(min(h, w) // subdivisions, min_region_size),
+ min(h, w)
+ )
+ larger_normal_map = extend_normal_map(normal_map, region_size)
+ lh, lw = larger_normal_map.shape[-2:]
+ subregions = define_subregions(lh, lw, region_size)
+ map_batch = create_subregions(larger_normal_map, subregions)
+ sub_height_maps = []
+ for m in map_batch:
+ sub_height_maps.append(compute_height(m.unsqueeze(0)).squeeze(0))
+ height_combined = combine_sub_height_maps(sub_height_maps, subregions, lh, lw)
+ height_cropped = crop_height_map(height_combined, (h, w), region_size)
+ return normalize_height_map(height_cropped)
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..ce2754b
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,14 @@
+[project]
+name = "comfyui-chord"
+description = "PBR material estimation custom nodes for paper: Chord: Chain of Rendering Decomposition for PBR Material Estimation from Generated Texture Images"
+version = "v1.0"
+license = { text = "Ubisoft Machine Learning License (Research-Only - Copyleft)" }
+dependencies = []
+
+[project.urls]
+Repository = "https://github.com/ubisoft/ComfyUI-Chord"
+
+[tool.comfy]
+PublisherId = "Ubisoft-LaForge"
+DisplayName = "ComfyUI-Chord"
+Icon = "https://www.ubisoft.com/favicon.ico"
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..385a6b1
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,7 @@
+huggingface_hub[hf_xet]
+diffusers
+transformers
+tokenizers
+typer
+omegaconf
+imageio
\ No newline at end of file