yum/3ner

A toon shader for Unity's BIRP.

git clone https://git.yummers.dev/yum/3ner

yumCloth bugfixes; burley tiling scaling fix69f2735

master
branchesmasterall refs →
9.0 KiB282 linesraw
1#!/usr/bin/env python3
2
3# To run with uv:
4# uv run -w OpenEXR -w numba ./make_dfg_lut.py
5
6import argparse
7import math
8import numpy as np
9import OpenEXR
10import Imath
11import numba
12import random
13import concurrent.futures
14import os
15from functools import partial
16
17
18@numba.njit(cache=True)
19def rcp(a):
20    return 1.0 / a
21
22
23@numba.njit(cache=True)
24def lerp(a, b, t):
25    return a + (b - a) * t
26
27
28@numba.njit(cache=True)
29def saturate(a):
30    if a < 0.0: return 0.0
31    if a > 1.0: return 1.0
32    return a
33
34
35# Standard BRDF components.
36@numba.njit(cache=True)
37def F_Schlick(LoH, f0, f90=1.0):
38    term = 1.0 - LoH
39    term2 = term * term
40    term5 = term2 * term2 * term
41    return f0 + (f90 - f0) * term5
42
43
44@numba.njit(cache=True)
45def D_GGX(roughness, NoH):
46    r2 = roughness * roughness
47    NoH2 = NoH * NoH
48    NoH4 = NoH2 * NoH2
49    k = rcp(NoH2) - 1.0
50    r2_plus_k = r2 + k
51    denom = NoH4 * r2_plus_k * r2_plus_k
52    return r2 / (denom + 1e-6)
53
54
55@numba.njit(cache=True)
56def G_GGXSmith(roughness, NoL, NoV):
57    denom = 2.0 * lerp(2.0 * NoL * NoV, NoL + NoV, roughness)
58    return rcp(denom + 1e-6)
59
60
61# Cloth BRDF components.
62@numba.njit(cache=True)
63def D_Cloth(roughness, NoH):
64    if roughness < 1e-4: return 0.0
65    r_rcp = rcp(roughness)
66    sin2H = 1.0 - NoH * NoH
67    return (2.0 + r_rcp) * pow(sin2H, r_rcp * 0.5) / (2.0 * math.pi)
68
69
70@numba.njit(cache=True)
71def G_Cloth_L(x, a, b, c, d, e):
72    return a / (1.0 + b * pow(x, c)) + d * x + e
73
74
75@numba.njit(cache=True)
76def Lambda_Cloth_Raw(roughness, cos_theta):
77    a0, a1 = 25.3245, 21.5473
78    b0, b1 = 3.32435, 3.82987
79    c0, c1 = 0.16801, 0.19823
80    d0, d1 = -1.27393, -1.97760
81    e0, e1 = -4.85967, -4.32054
82
83    one_minus_r = 1.0 - roughness
84    interp = one_minus_r * one_minus_r
85    rough_weight = 1.0 - interp
86
87    lambda_val = 0.0
88    if cos_theta < 0.5:
89        L0 = G_Cloth_L(cos_theta, a0, b0, c0, d0, e0)
90        L1 = G_Cloth_L(cos_theta, a1, b1, c1, d1, e1)
91        L = lerp(L0, L1, rough_weight)
92        lambda_val = math.exp(L)
93    else:
94        L_05_0 = G_Cloth_L(0.5, a0, b0, c0, d0, e0)
95        L_05_1 = G_Cloth_L(0.5, a1, b1, c1, d1, e1)
96        L_05 = lerp(L_05_0, L_05_1, rough_weight)
97
98        one_minus_cos = 1.0 - cos_theta
99        L_c_0 = G_Cloth_L(one_minus_cos, a0, b0, c0, d0, e0)
100        L_c_1 = G_Cloth_L(one_minus_cos, a1, b1, c1, d1, e1)
101        L_c = lerp(L_c_0, L_c_1, rough_weight)
102
103        lambda_val = math.exp(2.0 * L_05 - L_c)
104
105    return lambda_val
106
107
108@numba.njit(cache=True)
109def Lambda_Cloth_Softened(roughness, cos_theta):
110    lambda_val = Lambda_Cloth_Raw(roughness, cos_theta)
111    return pow(lambda_val, 1.0 + 2.0 * pow(1.0 - cos_theta, 8.0))
112
113
114@numba.njit(cache=True)
115def V_Cloth_Outgoing(roughness, NoL, NoV):
116    # Height-correlated Smith: G2 / (4 * NoL * NoV)
117    lambda_l = Lambda_Cloth_Softened(roughness, NoL)
118    lambda_v = Lambda_Cloth_Raw(roughness, NoV)
119    return 1.0 / ((1.0 + lambda_l + lambda_v) * 4.0 * NoL * NoV + 1e-6)
120
121
122@numba.njit(cache=True)
123def V_Cloth_Incoming(roughness, NoL, NoV):
124    lambda_l = Lambda_Cloth_Softened(roughness, NoL)
125    lambda_v = Lambda_Cloth_Raw(roughness, NoV)
126    return 1.0 / ((1.0 + lambda_l + lambda_v) * 4.0 * NoL * NoV + 1e-6)
127
128
129@numba.njit(cache=True)
130def integrate_brdf_jitted(roughness, NoV, num_samples):
131    V_x = math.sqrt(1.0 - NoV * NoV)
132    V_y = 0.0
133    V_z = NoV
134
135    # R: GGX scale, G: GGX bias, B: cloth outgoing albedo, A: cloth incoming albedo
136    std_scale, std_bias, cloth_out, cloth_in = 0.0, 0.0, 0.0, 0.0
137
138    for i in range(num_samples):
139        e1, e2 = random.random(), random.random()
140
141        # Importance sample GGX
142        a = roughness
143        a2 = a * a
144
145        phi = 2.0 * math.pi * e1
146        cos_theta = math.sqrt((1.0 - e2) / (1.0 + (a2 - 1.0) * e2))
147        sin_theta = math.sqrt(1.0 - cos_theta * cos_theta)
148
149        H_x = math.cos(phi) * sin_theta
150        H_y = math.sin(phi) * sin_theta
151        H_z = cos_theta
152
153        VoH = H_x * V_x + H_y * V_y + H_z * V_z
154        if VoH <= 0: continue
155
156        L_x = 2.0 * VoH * H_x - V_x
157        L_y = 2.0 * VoH * H_y - V_y
158        L_z = 2.0 * VoH * H_z - V_z
159
160        NoL = saturate(L_z)
161        NoH = saturate(H_z)
162        NoV_proxy = saturate(V_z)
163
164        if NoL > 0:
165            # --- Standard BRDF ---
166            # D cancels between numerator and PDF.
167            G = G_GGXSmith(roughness, NoL, NoV_proxy)
168            Fc_term = pow(1.0 - VoH, 5.0)
169
170            # PDF = D_GGX * NoH / (4 * VoH), so (D * G * NoL) / PDF simplifies to:
171            common_term = (G * NoL * 4.0 * VoH) / max(NoH, 1e-5)
172
173            std_scale += common_term * (1.0 - Fc_term)
174            std_bias += common_term * Fc_term
175
176            # --- Cloth BRDF ---
177            # Same GGX importance samples, reweighted for cloth D and V.
178            if roughness >= 1e-4:
179                D_c = D_Cloth(roughness, NoH)
180                pdf_ggx = D_GGX(roughness, NoH) * NoH / (4.0 * VoH + 1e-6)
181                V_out = V_Cloth_Outgoing(roughness, NoL, NoV_proxy)
182                V_in = V_Cloth_Incoming(roughness, NoV_proxy, NoL)
183                cloth_out += (D_c * V_out * NoL) / (pdf_ggx + 1e-6)
184                cloth_in += (D_c * V_in * NoL) / (pdf_ggx + 1e-6)
185
186    inv_n = 1.0 / num_samples
187    return std_scale * inv_n, std_bias * inv_n, cloth_out * inv_n, cloth_in * inv_n
188
189
190def calculate_pixel(coords, resolution, num_samples):
191    x, y = coords
192    u = (x + 0.5) / resolution
193    v = (y + 0.5) / resolution
194
195    NoV = saturate(u)
196    perceptual_roughness = saturate(v)
197    roughness = max(perceptual_roughness * perceptual_roughness, 1e-4)
198    if NoV < 1e-4: return x, y, 0.0, 0.0, 0.0, 0.0
199
200    std_scale, std_bias, cloth_out, cloth_in = integrate_brdf_jitted(roughness, NoV, num_samples)
201
202    # R: GGX scale, G: GGX bias, B: cloth outgoing albedo, A: cloth incoming albedo
203    return x, y, std_scale, std_bias, cloth_out, cloth_in
204
205
206def generate_exr(resolution, output_filename, num_samples, num_workers):
207    print(f"Generating {resolution}x{resolution} EXR '{output_filename}' (R=GGX scale, G=GGX bias, B=cloth out, A=cloth in) ({num_samples} samples/pixel) using {num_workers} workers.")
208    header = OpenEXR.Header(resolution, resolution)
209    pt = Imath.PixelType(Imath.PixelType.FLOAT)
210    header['channels'] = {
211        'R': Imath.Channel(pt),
212        'G': Imath.Channel(pt),
213        'B': Imath.Channel(pt),
214        'A': Imath.Channel(pt),
215    }
216
217    pixel_data = np.zeros((resolution, resolution, 4), dtype=np.float32)
218
219    coords_to_process = [(x, y) for y in range(resolution) for x in range(resolution)]
220    worker_func = partial(calculate_pixel, resolution=resolution, num_samples=num_samples)
221
222    processed_count = 0
223    total_pixels = len(coords_to_process)
224    print(f"Starting pixel processing...")
225
226    with concurrent.futures.ProcessPoolExecutor(max_workers=num_workers) as executor:
227        futures = {executor.submit(worker_func, coord): coord for coord in coords_to_process}
228
229        for future in concurrent.futures.as_completed(futures):
230            try:
231                x, y, r, g, b, a = future.result()
232                pixel_data[y, x] = (r, g, b, a)
233            except Exception as exc:
234                coord = futures[future]
235                print(f'\nPixel at {coord} generated an exception: {exc}')
236
237            processed_count += 1
238            print(f"  ...processed {processed_count}/{total_pixels} pixels ({processed_count/total_pixels:.1%})", end='\r')
239
240    print(f"\nProcessing complete. Writing to {output_filename}...")
241    try:
242        # Vertically flip to match UV coordinates (0,0 at bottom-left).
243        pixel_data = np.flipud(pixel_data)
244
245        exr_file = OpenEXR.OutputFile(output_filename, header)
246        r_data = pixel_data[:, :, 0].ravel().tobytes()
247        g_data = pixel_data[:, :, 1].ravel().tobytes()
248        b_data = pixel_data[:, :, 2].ravel().tobytes()
249        a_data = pixel_data[:, :, 3].ravel().tobytes()
250        exr_file.writePixels({'R': r_data, 'G': g_data, 'B': b_data, 'A': a_data})
251        exr_file.close()
252        print(f"Successfully generated {output_filename}")
253    except Exception as e:
254        raise RuntimeError(f"Failed to write EXR file '{output_filename}': {e}")
255
256def main():
257    parser = argparse.ArgumentParser(description='Generate packed DFG LUT (R=GGX scale, G=GGX bias, B=cloth out, A=cloth in).')
258    parser.add_argument('-r', '--resolution', type=int, default=512,
259                        help='Resolution of the square EXR image (default: 512)')
260    parser.add_argument('-s', '--samples', type=int, default=8192,
261                        help='Number of samples per pixel for integration (default: 8192)')
262    parser.add_argument('-o', '--output', default='dfg.exr',
263                        help='Output filename (default: dfg.exr)')
264    parser.add_argument('-j', '--workers', type=int, default=os.cpu_count(),
265                        help=f'Number of worker processes (default: {os.cpu_count()})')
266
267    args = parser.parse_args()
268
269    if args.resolution <= 0:
270        print("Error: Resolution must be a positive integer")
271        return 1
272
273    try:
274        generate_exr(args.resolution, args.output, args.samples, args.workers)
275    except Exception as e:
276        print(f"Error: {e}")
277        return 1
278
279    return 0
280
281if __name__ == '__main__':
282    exit(main())