@needle-tools/materialx
Version:
Web runtime support to load and display MaterialX materials in Needle Engine and three.js via the MaterialX WebAssembly library. glTF files containing the `NEEDLE_materials_mtlx` extension can be loaded with this package. There is also experimental suppor
361 lines (310 loc) • 17 kB
JavaScript
import { BufferGeometry, Camera, FrontSide, GLSL3, Group, Matrix3, Matrix4, Object3D, Scene, ShaderMaterial, Texture, Vector3, WebGLRenderer } from "three";
import { debug, getFrame, getTime } from "./utils.js";
import { MaterialXEnvironment } from "./materialx.js";
import { generateMaterialPropertiesForUniforms, getUniformValues } from "./materialx.helper.js";
import { cloneUniforms, cloneUniformsGroups } from "three/src/renderers/shaders/UniformsUtils.js";
// Add helper matrices for uniform updates (similar to MaterialX example)
const normalMat = new Matrix3();
const worldViewPos = new Vector3();
/**
* @typedef {Object} MaterialXMaterialInitParameters
* @property {string} name
* @property {string | null} [shaderName] - Optional name of the shader
* @property {any} shader
* @property {import('./materialx.helper.js').Callbacks} loaders
* @property {import('./materialx.js').MaterialXContext} context
* @property {import('three').MaterialParameters} [parameters] - Optional parameters
* @property {boolean} [debug] - Debug flag
*/
/**
* @typedef {"highp" | "mediump" | "lowp"} Precision
*/
// @dont-generate-component
export class MaterialXMaterial extends ShaderMaterial {
/** The original name of the shader
* @type {string | null} */
shaderName = null;
/**
* @param {MaterialXMaterial} source
* @returns {this}
*/
copy(source) {
super.copy(source);
this._context = source._context;
this._shader = source._shader;
this.uniforms = cloneUniforms(source.uniforms);
this.uniformsGroups = cloneUniformsGroups(source.uniformsGroups);
this.envMapIntensity = source.envMapIntensity;
this.envMap = source.envMap;
generateMaterialPropertiesForUniforms(this, this._shader.getStage('pixel'));
generateMaterialPropertiesForUniforms(this, this._shader.getStage('vertex'));
this.needsUpdate = true;
return this;
}
/** @type {import('./materialx.js').MaterialXContext | null} */
_context = null;
/** @type {any} */
_shader = null;
/** @type {boolean} */
_needsTangents = false;
/**
* @param {MaterialXMaterialInitParameters} [init]
*/
constructor(init) {
// TODO: we need to properly copy the uniforms and other properties from the source material
if (!init) {
super();
return;
}
// Get vertex and fragment shader source, and remove #version directive for newer js.
// It's added by three.js glslVersion.
let vertexShader = init.shader.getSourceCode("vertex");
let fragmentShader = init.shader.getSourceCode("pixel");
vertexShader = vertexShader.replace(/^#version.*$/gm, '').trim();
fragmentShader = fragmentShader.replace(/^#version.*$/gm, '').trim();
// MaterialX uses different attribute names than js defaults,
// so we patch the MaterialX shaders to match the js standard names.
// Otherwise, we'd have to modify the mesh attributes (see original MaterialX for reference).
// Patch vertexShader
vertexShader = vertexShader.replace(/\bi_position\b/g, 'position');
vertexShader = vertexShader.replace(/\bi_normal\b/g, 'normal');
vertexShader = vertexShader.replace(/\bi_texcoord_0\b/g, 'uv');
vertexShader = vertexShader.replace(/\bi_texcoord_1\b/g, 'uv1');
vertexShader = vertexShader.replace(/\bi_texcoord_2\b/g, 'uv2');
vertexShader = vertexShader.replace(/\bi_texcoord_3\b/g, 'uv3');
vertexShader = vertexShader.replace(/\bi_tangent\b/g, 'tangent');
vertexShader = vertexShader.replace(/\bi_color_0\b/g, 'color');
// TODO: do we need to add depthbuffer fragments? https://discourse.threejs.org/t/shadermaterial-render-order-with-logarithmicdepthbuffer-is-wrong/49221/4
// Add logdepthbuf_pars_vertex at the beginning of the vertex shader before main()
// vertexShader = vertexShader.replace(/void\s+main\s*\(\s*\)\s*{/, `#include <logdepthbuf_pars_vertex>\nvoid main() {`);
// // Add logdepthbuf_vertex to vertex shader if not present at end of main()
// vertexShader = vertexShader.replace(/void\s+main\s*\(\s*\)\s*{/, `void main() {\n #include <logdepthbuf_vertex>\n`);
// Patch fragmentShader
const precision = init.parameters?.precision || "highp";
vertexShader = vertexShader.replace(/precision mediump float;/g, `precision ${precision} float;`);
vertexShader = vertexShader.replace(/#define M_FLOAT_EPS 1e-8/g, precision === "highp" ? `#define M_FLOAT_EPS 1e-8` : `#define M_FLOAT_EPS 1e-3`);
fragmentShader = fragmentShader.replace(/precision mediump float;/g, `precision ${precision} float;`);
fragmentShader = fragmentShader.replace(/#define M_FLOAT_EPS 1e-8/g, precision === "highp" ? `#define M_FLOAT_EPS 1e-8` : `#define M_FLOAT_EPS 1e-3`);
fragmentShader = fragmentShader.replace(/\bi_position\b/g, 'position');
fragmentShader = fragmentShader.replace(/\bi_normal\b/g, 'normal');
fragmentShader = fragmentShader.replace(/\bi_texcoord_0\b/g, 'uv');
fragmentShader = fragmentShader.replace(/\bi_texcoord_1\b/g, 'uv1');
fragmentShader = fragmentShader.replace(/\bi_texcoord_2\b/g, 'uv2');
fragmentShader = fragmentShader.replace(/\bi_texcoord_3\b/g, 'uv3');
fragmentShader = fragmentShader.replace(/\bi_tangent\b/g, 'tangent');
fragmentShader = fragmentShader.replace(/\bi_color_0\b/g, 'color');
// Capture some vertex shader properties
const uv_is_vec2 = vertexShader.includes('in vec2 uv;'); // check if uv is vec2; e.g. https://matlib.gpuopen.com/main/materials/all?material=da6ec531-f5c1-4790-ac14-8a5c51d0314e
const uv1_is_vec2 = vertexShader.includes('in vec2 uv1;');
const uv2_is_vec2 = vertexShader.includes('in vec2 uv2;');
const uv3_is_vec2 = vertexShader.includes('in vec2 uv3;');
// Remove `in vec3 position;` and so on since they're already declared by ShaderMaterial
vertexShader = vertexShader.replace(/in\s+vec3\s+position;/g, '');
vertexShader = vertexShader.replace(/in\s+vec3\s+normal;/g, '');
vertexShader = vertexShader.replace(/in\s+vec2\s+uv;/g, '');
vertexShader = vertexShader.replace(/in\s+vec3\s+uv;/g, '');
var hasUv1 = vertexShader.includes('in vec3 uv1;');
vertexShader = vertexShader.replace(/in\s+vec3\s+uv1;/g, '');
var hasUv2 = vertexShader.includes('in vec3 uv2;');
vertexShader = vertexShader.replace(/in\s+vec3\s+uv2;/g, '');
var hasUv3 = vertexShader.includes('in vec3 uv3;');
vertexShader = vertexShader.replace(/in\s+vec3\s+uv3;/g, '');
var hasTangent = vertexShader.includes('in vec4 tangent;');
vertexShader = vertexShader.replace(/in\s+vec4\s+tangent;/g, '');
var hasColor = vertexShader.includes('in vec4 color;');
vertexShader = vertexShader.replace(/in\s+vec4\s+color;/g, '');
// Patch uv 2-component to 3-component (`texcoord_0 = uv;` needs to be replaced with `texcoord_0 = vec3(uv, 0.0);`)
// TODO what if we actually have a 3-component UV? Not sure what three.js does then
if (!uv_is_vec2) vertexShader = vertexShader.replace(/texcoord_0 = uv;/g, 'texcoord_0 = vec3(uv, 0.0);');
if (!uv1_is_vec2) vertexShader = vertexShader.replace(/texcoord_1 = uv1;/g, 'texcoord_1 = vec3(uv1, 0.0);');
if (!uv2_is_vec2) vertexShader = vertexShader.replace(/texcoord_2 = uv2;/g, 'texcoord_2 = vec3(uv2, 0.0);');
if (!uv3_is_vec2) vertexShader = vertexShader.replace(/texcoord_3 = uv3;/g, 'texcoord_3 = vec3(uv3, 0.0);');
// Patch units – seems MaterialX uses different units and we end up with wrong light values?
// result.direction = light.position - position;
fragmentShader = fragmentShader.replace(
/result\.direction\s*=\s*light\.position\s*-\s*position;/g,
'result.direction = (light.position - position) * 10.0 / 1.0;');
// Add tonemapping and colorspace handling
// Replace `out vec4 out1;` with `out vec4 gl_FragColor;`
fragmentShader = fragmentShader.replace(
/out\s+vec4\s+out1;/,
'layout(location = 0) out vec4 pc_fragColor;\n#define gl_FragColor pc_fragColor');
// Replace `out1 = vec4(<CAPTURE>)` with `gl_FragColor = vec4(<CAPTURE>)` and tonemapping/colorspace handling
fragmentShader = fragmentShader.replace(/^\s*out1\s*=\s*vec4\((.*)\);/gm, `
gl_FragColor = vec4($1);
#include <tonemapping_fragment>
#include <colorspace_fragment>`);
const defines = {};
if (hasUv1) defines['USE_UV1'] = '';
if (hasUv2) defines['USE_UV2'] = '';
if (hasUv3) defines['USE_UV3'] = '';
if (hasTangent) defines['USE_TANGENT'] = '';
if (hasColor) defines['USE_COLOR'] = '';
const searchPath = ""; // Could be derived from the asset path if needed
const isTransparent = init.parameters?.transparent ?? false;
super({
name: init.name,
uniforms: {},
vertexShader: vertexShader,
fragmentShader: fragmentShader,
glslVersion: GLSL3,
depthTest: true,
depthWrite: !isTransparent,
defines: defines,
...init.parameters, // Spread any additional parameters passed to the material
});
this.shaderName = init.shaderName || null;
this._context = init.context;
this._shader = init.shader;
this._needsTangents = vertexShader.includes('in vec4 tangent;') || vertexShader.includes('in vec3 tangent;');
Object.assign(this.uniforms, {
...getUniformValues(init.shader.getStage('vertex'), init.loaders, searchPath),
...getUniformValues(init.shader.getStage('pixel'), init.loaders, searchPath),
u_worldMatrix: { value: new Matrix4() },
u_viewProjectionMatrix: { value: new Matrix4() },
u_viewPosition: { value: new Vector3() },
u_worldInverseTransposeMatrix: { value: new Matrix4() },
// u_shadowMatrix: { value: new Matrix4() },
// u_shadowMap: { value: null, type: 't' }, // Shadow map
u_envMatrix: { value: new Matrix4() },
u_envRadiance: { value: null, type: 't' },
u_envRadianceMips: { value: 8, type: 'i' },
// TODO we need to figure out how we can set a PMREM here... doing many texture samples is prohibitively expensive
u_envRadianceSamples: { value: 8, type: 'i' },
u_envIrradiance: { value: null, type: 't' },
u_refractionEnv: { value: true },
u_numActiveLightSources: { value: 0 },
u_lightData: { value: [], needsUpdate: false }, // Array of light data. We need to set needsUpdate to false until we actually update it
});
generateMaterialPropertiesForUniforms(this, init.shader.getStage('pixel'));
generateMaterialPropertiesForUniforms(this, init.shader.getStage('vertex'));
if (debug || init.debug) {
// Get lighting and environment data from MaterialX environment
console.group("[MaterialX]: ", this.name);
console.log(`Vertex shader length: ${vertexShader.length}\n`, vertexShader);
console.log(`Fragment shader length: ${fragmentShader.length}\n`, fragmentShader);
console.groupEnd();
}
}
/** @type {boolean} */
_missingTangentsWarned = false;
/**
* @param {WebGLRenderer} renderer
* @param {Scene} _scene
* @param {Camera} camera
* @param {BufferGeometry} geometry
* @param {Object3D} object
* @param {Group} _group
*/
onBeforeRender(renderer, _scene, camera, geometry, object, _group) {
if (this._needsTangents && !geometry.attributes.tangent) {
if (!this._missingTangentsWarned) {
this._missingTangentsWarned = true;
console.warn(`[MaterialX] Tangents are required for this material (${this.name}) but not present in the geometry.`);
// TODO: can we compute tangents here?
}
}
const time = this._context?.getTime?.() || getTime();
const frame = this._context?.getFrame?.() || getFrame();
const env = MaterialXEnvironment.get(_scene);
if (env) {
env.update(frame, _scene, renderer);
this.updateUniforms(env, renderer, object, camera, time, frame);
}
}
/** @type {number} */
envMapIntensity = 1.0; // Default intensity for environment map
/** @type {Texture | null} */
envMap = null; // Environment map texture, can be set externally
/**
* @param {MaterialXEnvironment} environment
* @param {WebGLRenderer} _renderer
* @param {Object3D} object
* @param {Camera} camera
* @param {number} [time]
* @param {number} [frame]
*/
updateUniforms = (environment, _renderer, object, camera, time, frame) => {
const uniforms = this.uniforms;
// Update standard transformation matrices
if (uniforms.u_worldMatrix) {
uniforms.u_worldMatrix.value = object.matrixWorld;
// @ts-ignore
uniforms.u_worldMatrix.needsUpdate = true;
}
if (uniforms.u_viewProjectionMatrix) {
uniforms.u_viewProjectionMatrix.value.multiplyMatrices(camera.projectionMatrix, camera.matrixWorldInverse);
// @ts-ignore
uniforms.u_viewProjectionMatrix.needsUpdate = true;
}
if (uniforms.u_worldInverseTransposeMatrix) {
uniforms.u_worldInverseTransposeMatrix.value.setFromMatrix3(normalMat.getNormalMatrix(object.matrixWorld));
// @ts-ignore
uniforms.u_worldInverseTransposeMatrix.needsUpdate = true;
}
if (uniforms.u_viewPosition) {
uniforms.u_viewPosition.value.copy(camera.getWorldPosition(worldViewPos));
// @ts-ignore
uniforms.u_viewPosition.needsUpdate = true;
}
// if (uniforms.u_shadowMap) {
// const light = environment.lights?.[2] || null;
// uniforms.u_shadowMatrix.value = light?.shadow?.matrix.clone().premultiply(object.matrixWorld.clone()).invert();
// uniforms.u_shadowMap.value = light.shadow?.map || null;
// uniforms.u_shadowMap.needsUpdate = true;
// console.log("[MaterialX] Renderer shadow map updated", light);
// }
// Update time uniforms
if (uniforms.u_time) {
if (time === undefined) time = getTime();
uniforms.u_time.value = time;
}
if (uniforms.u_frame) {
if (frame === undefined) frame = getFrame();
uniforms.u_frame.value = frame;
}
this.uniformsNeedUpdate = true;
// Update light uniforms
this.updateEnvironmentUniforms(environment);
}
/**
* @private
* @param {MaterialXEnvironment} environment
*/
updateEnvironmentUniforms = (environment) => {
const uniforms = this.uniforms;
// Get lighting data from environment
const lightData = environment.lightData || null;
const lightCount = environment.lightCount || 0;
const textures = environment.getTextures(this) || null;
// Update light count
if (uniforms.u_numActiveLightSources && lightCount >= 0) {
uniforms.u_numActiveLightSources.value = lightCount;
}
// Update light data
if (lightData?.length) {
uniforms.u_lightData.value = lightData;
if ("needsUpdate" in uniforms.u_lightData && uniforms.u_lightData.needsUpdate === false) {
if (debug) console.debug(`[MaterialX] LightData assigned (${this.name}, ${this.uuid})`, lightData);
uniforms.u_lightData.needsUpdate = undefined;
}
}
// Update environment uniforms
if (uniforms.u_envRadiance) {
const prev = uniforms.u_envRadiance.value;
uniforms.u_envRadiance.value = textures.radianceTexture;
// @ts-ignore
if (prev != textures.radianceTexture) uniforms.u_envRadiance.needsUpdate = true;
}
if (uniforms.u_envRadianceMips) {
uniforms.u_envRadianceMips.value = Math.trunc(Math.log2(Math.max(textures.radianceTexture?.source.data.width ?? 0, textures.radianceTexture?.source.data.height ?? 0))) + 1;
}
if (uniforms.u_envIrradiance) {
const prev = uniforms.u_envIrradiance.value;
uniforms.u_envIrradiance.value = textures.irradianceTexture;
// @ts-ignore
if (prev != textures.irradianceTexture) uniforms.u_envIrradiance.needsUpdate = true;
}
this.uniformsNeedUpdate = true;
}
}