@needle-tools/materialx 1.0.0 → 1.0.1-next.19d0723

package/src/materialx.material.ts

@@ -0,0 +1,220 @@
+import { Camera, DoubleSide, FrontSide, GLSL3, Matrix3, Matrix4, Mesh, Object3D, ShaderMaterial, Texture, Vector3 } from "three";
+import { debug } from "./utils.js";
+import { MaterialXEnvironment, state } from "./materialx.js";
+import { getUniformValues, Loaders } from "./materialx.helper.js";
+import { Context } from "@needle-tools/engine";
+
+
+// Add helper matrices for uniform updates (similar to MaterialX example)
+const identityMatrix = new Matrix4();
+const normalMat = new Matrix3();
+const viewProjMat = new Matrix4();
+const worldViewPos = new Vector3();
+
+declare type MaterialXMaterialInitParameters = {
+    name: string,
+    shader: any,
+    loaders: Loaders,
+    transparent?: boolean,
+}
+
+export class MaterialXMaterial extends ShaderMaterial {
+
+    // copy(source: MaterialXMaterial): this {
+    //     super.copy(source);
+    //     this.name = source.name;
+    //     this.uniforms = { ...source.uniforms }; // Shallow copy of uniforms
+    //     this.envMapIntensity = source.envMapIntensity;
+    //     this.envMap = source.envMap;
+    //     this.updateUniforms = source.updateUniforms; // Copy the update function
+    //     return this;
+    // }
+
+    constructor(init?: MaterialXMaterialInitParameters) {
+
+        // 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_tangent\b/g, 'tangent');
+        vertexShader = vertexShader.replace(/\bi_color_0\b/g, 'color');
+
+        // Patch fragmentShader
+        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_tangent\b/g, 'tangent');
+        fragmentShader = fragmentShader.replace(/\bi_color_0\b/g, 'color');
+
+        // 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+vec3\s+uv;/g, '');
+        vertexShader = vertexShader.replace(/in\s+vec3\s+uv1;/g, '');
+        vertexShader = vertexShader.replace(/in\s+vec4\s+tangent;/g, '');
+        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
+        vertexShader = vertexShader.replace(/texcoord_0 = uv;/g, 'texcoord_0 = vec3(uv, 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 searchPath = ""; // Could be derived from the asset path if needed
+        const flipV = false; // Set based on your geometry requirements
+        const isTransparent = init.transparent ?? false;
+        super({
+            name: init.name,
+            uniforms: {
+                ...getUniformValues(init.shader.getStage('vertex'), init.loaders, searchPath, flipV),
+                ...getUniformValues(init.shader.getStage('pixel'), init.loaders, searchPath, flipV),
+            },
+            vertexShader: vertexShader,
+            fragmentShader: fragmentShader,
+            glslVersion: GLSL3,
+            transparent: isTransparent,
+            side: FrontSide,
+            depthTest: true,
+            depthWrite: !isTransparent,
+        });
+
+
+
+        Object.assign(this.uniforms, {
+            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: [] }, // Array of light data
+        });
+
+        if (debug) {
+            // Get lighting and environment data from MaterialX environment
+            console.group("[MaterialX]: ", name);
+            console.log("Vertex shader length:", vertexShader.length, vertexShader);
+            console.log("Fragment shader length:", fragmentShader.length, fragmentShader);
+            console.groupEnd();
+        }
+
+    }
+
+
+    envMapIntensity: number = 1.0; // Default intensity for environment map
+    envMap: Texture | null = null; // Environment map texture, can be set externally
+    updateUniforms = (context: Context, environment: MaterialXEnvironment, object: Object3D, camera: Camera) => {
+
+        const uniforms = this.uniforms;
+
+        // TODO remove. Not sure why this is needed, but without it
+        // we currently get some "swimming" where matrices are not up to date.
+        camera.updateMatrixWorld(true);
+
+        // Update standard transformation matrices
+        if (uniforms.u_worldMatrix) {
+            if (!uniforms.u_worldMatrix.value?.isMatrix4) uniforms.u_worldMatrix.value = new Matrix4();
+            uniforms.u_worldMatrix.value = object.matrixWorld;
+        }
+
+        if (uniforms.u_viewProjectionMatrix) {
+            if (!uniforms.u_viewProjectionMatrix.value?.isMatrix4) uniforms.u_viewProjectionMatrix.value = new Matrix4();
+            uniforms.u_viewProjectionMatrix.value.multiplyMatrices(camera.projectionMatrix, camera.matrixWorldInverse);
+        }
+
+        if (uniforms.u_viewPosition) {
+            if (!uniforms.u_viewPosition.value?.isVector3) uniforms.u_viewPosition.value = new Vector3();
+            uniforms.u_viewPosition.value.copy(camera.getWorldPosition(worldViewPos));
+        }
+
+        if (uniforms.u_worldInverseTransposeMatrix) {
+            if (!uniforms.u_worldInverseTransposeMatrix.value?.isMatrix4) uniforms.u_worldInverseTransposeMatrix.value = new Matrix4();
+            uniforms.u_worldInverseTransposeMatrix.value.setFromMatrix3(normalMat.getNormalMatrix(object.matrixWorld));
+        }
+
+        // Update time uniforms
+        if (uniforms.u_time) {
+            uniforms.u_time.value = context.time.time;
+        }
+        if (uniforms.u_frame) {
+            uniforms.u_frame.value = context.time.frame;
+        }
+
+        // Update light uniforms
+        this.updateEnvironmentUniforms(environment);
+
+        this.uniformsNeedUpdate = true;
+    }
+
+    private updateEnvironmentUniforms = (environment: MaterialXEnvironment) => {
+
+        // 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 (this.uniforms.u_numActiveLightSources && lightCount >= 0) {
+            this.uniforms.u_numActiveLightSources.value = lightCount;
+        }
+
+        // Update light data
+        if (lightData) {
+            this.uniforms.u_lightData.value = lightData;
+        }
+
+        // Update environment uniforms
+        if (this.uniforms.u_envMatrix) {
+            this.uniforms.u_envMatrix.value = identityMatrix;
+        }
+        if (this.uniforms.u_envRadiance) {
+            this.uniforms.u_envRadiance.value = textures.radianceTexture || null;
+        }
+        if (this.uniforms.u_envRadianceMips) {
+            this.uniforms.u_envRadianceMips.value = Math.trunc(Math.log2(Math.max(textures.radianceTexture?.source.data.width ?? 0, textures.radianceTexture?.source.data.height ?? 0))) + 1;
+        }
+        if (this.uniforms.u_envIrradiance) {
+            this.uniforms.u_envIrradiance.value = textures.irradianceTexture;
+        }
+
+        this.uniformsNeedUpdate = true;
+    }
+}

package/src/materialx.ts

@@ -1,32 +1,15 @@
-import { Context, delay, isDevEnvironment, ObjectUtils } from "@needle-tools/engine";
+import { Context, delay, isDevEnvironment, ObjectUtils, GameObject, onBeforeRender } from "@needle-tools/engine";
+import type { MaterialX as MX } from "./materialx.types.js";
 import MaterialX from "../bin/JsMaterialXGenShader.js";
 import { debug } from "./utils.js";
 import { renderPMREMToEquirect } from "./textureHelper.js";
-import { Light, MeshBasicMaterial, Object3D, PMREMGenerator } from "three";
-import { registerLights } from "./helper.js";
-
-// Configure MaterialX with the correct path for its data files
-const materialXConfig = {
-    locateFile: (path: string, scriptDirectory: string) => {
-        if (debug) console.debug("MaterialX locateFile called:", { path, scriptDirectory });
-
-        // Return the correct path for MaterialX data files
-        if (path.endsWith('.data') || path.endsWith('.wasm')) {
-            // For Vite dev server, we need to use the correct module path
-            const correctPath = new URL(`../bin/${path}`, import.meta.url).href;
-            if (debug) console.log("Resolved path:", correctPath);
-            return correctPath;
-        }
-        return scriptDirectory + path;
-    },
-    // Add buffer allocation to handle the data file properly
-    wasmBinary: null,
-    wasmMemory: null
-};
-
-// Global MaterialX module instance - initialized lazily
+import { Light, Material, MeshBasicMaterial, Object3D, PMREMGenerator, Texture } from "three";
+import { registerLights, getLightData } from "./materialx.helper.js";
+import type { MaterialXMaterial } from "./materialx.material.js";
+
+
 export const state = new class {
-    materialXModule: any = null;
+    materialXModule: MX.MODULE | null = null;
     materialXGenerator: any = null;
     materialXGenContext: any = null;
     materialXStdLib: any = null;
@@ -40,19 +23,45 @@ export const state = new class {
     }
 }
 
-
-// Initialize MaterialX WASM module lazily
-export async function initializeMaterialX(): Promise<void> {
+/** Initialize the MaterialX module. Must be awaited before trying to create materials */
+export async function ready(): Promise<void> {
     if (state.materialXInitPromise) {
         return state.materialXInitPromise;
     }
     return state.materialXInitPromise = (async () => {
         if (state.materialXModule) return; // Already initialized
-        if (debug) console.log("Initializing MaterialX WASM module...");
+        if (debug) console.log("[MaterialX] Initializing WASM module...");
         try {
-            const module = await MaterialX(materialXConfig);
-            if (debug) console.log("MaterialXLoader module loaded", module);
-            state.materialXModule = module;
+
+            const urls: Array<string> = await Promise.all([
+                /** @ts-ignore */
+                import(`../bin/JsMaterialXCore.wasm?url`).then(m => m.default || m),
+                /** @ts-ignore */
+                import(`../bin/JsMaterialXGenShader.wasm?url`).then(m => m.default || m),
+                /** @ts-ignore */
+                import(`../bin/JsMaterialXGenShader.data.txt?url`).then(m => m.default || m),
+            ]);
+            const [JsMaterialXCore, JsMaterialXGenShader, JsMaterialXGenShader_data] = urls;
+
+            const module = await MaterialX({
+                locateFile: (path: string, scriptDirectory: string) => {
+                    if (debug) console.debug("[MaterialX]  locateFile called:", { path, scriptDirectory });
+
+                    if (path.includes("JsMaterialXCore.wasm")) {
+                        return JsMaterialXCore; // Use the URL for the core WASM file
+                    }
+                    else if (path.includes("JsMaterialXGenShader.wasm")) {
+                        return JsMaterialXGenShader; // Use the URL for the shader WASM file
+                    }
+                    else if (path.includes("JsMaterialXGenShader.data")) {
+                        return JsMaterialXGenShader_data; // Use the URL for the shader data file
+                    }
+
+                    return scriptDirectory + path;
+                },
+            });
+            if (debug) console.log("[MaterialX] module loaded", module);
+            state.materialXModule = module as MX.MODULE
 
             // Initialize shader generator and context
             state.materialXGenerator = module.EsslShaderGenerator.create();
@@ -63,158 +72,190 @@ export async function initializeMaterialX(): Promise<void> {
             state.materialXStdLib = module.loadStandardLibraries(state.materialXGenContext);
             tempDoc.setDataLibrary(state.materialXStdLib);
 
-            // Initialize basic lighting with default light rig
-            const defaultLightRigXml = `<?xml version="1.0"?>
-<materialx version="1.39">
-  <!-- Default directional light -->
-  <directional_light name="default_light" type="lightshader">
-    <input name="direction" type="vector3" value="0.0, -1.0, -0.5" />
-    <input name="color" type="color3" value="1.0, 1.0, 1.0" />
-    <input name="intensity" type="float" value="1.0" />
-  </directional_light>
-</materialx>`;
+            // TODO ShaderInterfaceType.SHADER_INTERFACE_REDUCED would be better, but doesn't actually seem to be supported in the MaterialX javascript bindings
+            state.materialXGenContext.getOptions().shaderInterfaceType = state.materialXModule.ShaderInterfaceType.SHADER_INTERFACE_COMPLETE;
+
+            // SPECULAR_ENVIRONMENT_NONE: Do not use specular environment maps.
+            // SPECULAR_ENVIRONMENT_FIS: Use Filtered Importance Sampling for specular environment/indirect lighting.
+            // SPECULAR_ENVIRONMENT_PREFILTER: Use pre-filtered environment maps for specular environment/indirect lighting.
+            state.materialXGenContext.getOptions().hwSpecularEnvironmentMethod = state.materialXModule.HwSpecularEnvironmentMethod.SPECULAR_ENVIRONMENT_FIS;
+
+            // TRANSMISSION_REFRACTION: Use a refraction approximation for transmission rendering.
+            // TRANSMISSION_OPACITY: Use opacity for transmission rendering.
+            // state.materialXGenContext.getOptions().hwTransmissionRenderMethod = state.materialXModule.HwTransmissionRenderMethod.TRANSMISSION_REFRACTION;
+
+            // Turned off because we're doing color space conversion the three.js way
+            state.materialXGenContext.getOptions().hwSrgbEncodeOutput = false;
+
+            // Enables the generation of a prefiltered environment map.
+            // TODO Would be great to use but requires setting more uniforms (like u_envPrefilterMip).
+            // When set to true, the u_envRadiance map is expected to be a prefiltered environment map.
+            // state.materialXGenContext.getOptions().hwWriteEnvPrefilter = true;
+
+            // Set a reasonable default for max active lights
+            state.materialXGenContext.getOptions().hwMaxActiveLightSources = 4;
 
             // This prewarms the shader generation context to have all light types
-            await registerLights(state.materialXModule, [], state.materialXGenContext);
+            await registerLights(state.materialXModule, state.materialXGenContext);
 
-            if (debug) console.log("MaterialX generator initialized successfully");
+            if (debug) console.log("[MaterialX]  generator initialized successfully");
         } catch (error) {
-            console.error("Failed to load MaterialX module:", error);
+            console.error("[MaterialX] Failed to load MaterialX module:", error);
             throw error;
         }
     })();
 }
 
+type EnvironmentTextureSet = {
+    radianceTexture: Texture | null;
+    irradianceTexture: Texture | null;
+}
+
 // MaterialX Environment Manager - handles lighting and environment setup
 export class MaterialXEnvironment {
-    private lights: any[] = [];
-    private lightData: any = null;
-    private radianceTexture: any = null;
-    private irradianceTexture: any = null;
-    private context: Context | null = null;
-    private initialized: boolean = false;
+    private _context: Context | null = null;
+    private _lights: Array<Light> = [];
+    private _lightData: any = null;
+    private _lightCount: number = 0;
+    private _initializePromise: Promise<boolean> | null = null;
+
+    private _unsubscribehook: (() => void) | null = null;
 
     constructor() {
-        if (debug) console.log("MaterialX Environment created");
+        if (debug) console.log("[MaterialX] Environment created");
     }
 
-    setContext(context: Context) {
-        this.context = context;
+    // Initialize with Needle Engine context
+    async initialize(context: Context): Promise<boolean> {
+        if (this._initializePromise) {
+            return this._initializePromise;
+        }
+        return this._initializePromise = this._initialize(context);
     }
 
-    /*
-    // Initialize MaterialX lighting system based on the reference implementation
-    async initializeLighting(lightRigXml: string, renderer?: any, radianceTexture?: any, irradianceTexture?: any): Promise<void> {
-        if (!materialXModule || !materialXGenContext) {
-            console.warn("MaterialX module not initialized, skipping lighting setup");
-            return;
+    get lightData() { return this._lightData; }
+    get lightCount() { return this._lightCount || 0; }
+    getTextures(material: MaterialXMaterial) {
+        if (material.envMap) {
+            // If the material has its own envMap, we don't use the irradiance texture
+            return this._getTextures(material.envMap);
         }
-
-        registerLights(materialXModule, this.lights, materialXGenContext);
+        return this._getTextures(this._context?.scene.environment);
     }
-    */
 
-    // Initialize with Needle Engine context
-    async initializeFromContext(): Promise<void> {
-        if (!this.context) {
-            console.warn("No Needle context available for MaterialX environment initialization");
-            return;
-        }
+    private _pmremGenerator: PMREMGenerator | null = null;
+    private readonly _texturesCache: Map<Texture | null, EnvironmentTextureSet> = new Map();
 
-        // Prevent multiple initializations
-        if (this.initialized) {
-            if (debug) console.log("MaterialX environment already initialized, skipping");
-            return;
-        }
+    private _initialize: (context: Context) => Promise<boolean> = async (context: Context) => {
 
-        // Clean up previous textures if they exist
-        if (this.radianceTexture) {
-            if (debug) console.log("Disposing previous radiance texture");
-            this.radianceTexture.dispose();
-            this.radianceTexture = null;
-        }
-        if (this.irradianceTexture) {
-            if (debug) console.log("Disposing previous irradiance texture");
-            this.irradianceTexture.dispose();
-            this.irradianceTexture = null;
-        }
+        this._context = context;
+        this._pmremGenerator = new PMREMGenerator(context.renderer);
 
-        // Get renderer from context
-        const renderer = this.context.renderer;
+        this._unsubscribehook?.();
+        this._unsubscribehook = onBeforeRender(() => {
+            this.updateLighting(false);
+            this._getTextures(context.scene.environment);
+        })
 
         // TODO remove this delay; we should wait for the scene lighting to be ready
         // and then update the uniforms
-        let envMap = this.context.scene.environment;
-        while (!envMap) {
-            await delay(200);
-            envMap = this.context.scene.environment;
-        }
-        var pmrem = new PMREMGenerator(renderer);
-        const target = pmrem.fromEquirectangular(envMap);
-
-        const radianceRenderTarget = renderPMREMToEquirect(renderer, target.texture, 0.0, 1024, 512, target.height);
-        const irradianceRenderTarget = renderPMREMToEquirect(renderer, target.texture, 1.0, 32, 16, target.height);
-
-        this.radianceTexture = radianceRenderTarget.texture;
-        this.irradianceTexture = irradianceRenderTarget.texture;
-
-        // Clean up PMREM generator and its render target
-        target.dispose();
-        pmrem.dispose();
-
-        if (debug) {
-            console.log({ radiance: this.radianceTexture, irradiance: this.irradianceTexture });
-            // Show both of them on cubes in the scene
-            const unlitMat = new MeshBasicMaterial();
-            const radianceMat = unlitMat.clone();
-            radianceMat.map = this.radianceTexture;
-            const radianceCube = ObjectUtils.createPrimitive("Cube", { material: radianceMat });
-            const irradianceMat = unlitMat.clone();
-            irradianceMat.map = this.irradianceTexture;
-            const irradianceCube = ObjectUtils.createPrimitive("Cube", { material: irradianceMat });
-            this.context.scene.add(radianceCube);
-            this.context.scene.add(irradianceCube);
-            radianceCube.position.set(2, 0, 0);
-            radianceCube.scale.y = 0.00001;
-            irradianceCube.position.set(-2, 0, 0);
-            irradianceCube.scale.y = 0.00001;
-            // await this.initializeLighting(defaultLightRigXml, renderer);
-            console.log("MaterialX environment initialized from Needle context", this, this.context.scene);
+        while (!context.scene.environment) {
+            await delay(5);
         }
+        this._getTextures(context.scene.environment);
 
-        // Find lights in scene
-        let lights = new Array<Light>();
-        this.context.scene.traverse((object: Object3D) => {
-            if ((object as Light).isLight) lights.push(object as Light);
-        });
+        // if (debug) {
+        //     console.log({ radiance: this._radianceTexture, irradiance: this._irradianceTexture });
+        //     // Show both of them on cubes in the scene
+        //     const unlitMat = new MeshBasicMaterial();
+        //     unlitMat.side = 2;
+        //     const radianceMat = unlitMat.clone();
+        //     radianceMat.map = this._radianceTexture;
+        //     const radianceCube = ObjectUtils.createPrimitive("Quad", { material: radianceMat });
+        //     const irradianceMat = unlitMat.clone();
+        //     irradianceMat.map = this._irradianceTexture;
+        //     const irradianceCube = ObjectUtils.createPrimitive("Quad", { material: irradianceMat });
+        //     context.scene.add(radianceCube);
+        //     context.scene.add(irradianceCube);
+        //     radianceCube.position.set(2, 0, 0);
+        //     irradianceCube.position.set(-2, 0, 0);
+        //     console.log("[MaterialX] environment initialized from Needle context", this, this._context.scene);
+        // }
 
-        this.lightData = await registerLights(state.materialXModule, lights, state.materialXGenContext);
+        this.updateLighting(true);
 
         // Mark as initialized
-        this.initialized = true;
+        return true;
     }
 
-    // getLights() { return this.lights; }
-    getLightData() { return this.lightData; }
-    getRadianceTexture() { return this.radianceTexture; }
-    getIrradianceTexture() { return this.irradianceTexture; }
-
-    setRadianceTexture(texture: any) { this.radianceTexture = texture; }
-    setIrradianceTexture(texture: any) { this.irradianceTexture = texture; }
-
     // Reset the environment to allow re-initialization
     reset() {
-        if (debug) console.log("Resetting MaterialX environment");
-        if (this.radianceTexture) {
-            this.radianceTexture.dispose();
-            this.radianceTexture = null;
+        if (debug) console.log("[MaterialX] Resetting environment");
+        this._initializePromise = null;
+        this._lights = [];
+        this._lightData = null;
+        this._lightCount = 0;
+        this._pmremGenerator?.dispose();
+        this._pmremGenerator = null;
+        for(const textureSet of this._texturesCache.values()) {
+            textureSet.radianceTexture?.dispose();
+            textureSet.irradianceTexture?.dispose();
+        }
+        this._texturesCache.clear();
+
+        this._unsubscribehook?.();
+        this._unsubscribehook = null;
+    }
+
+    private _getTextures(texture: Texture | null | undefined): {
+        radianceTexture: Texture | null,
+        irradianceTexture: Texture | null
+    } {
+
+        let res: EnvironmentTextureSet | undefined = this._texturesCache.get(texture || null);
+        if (res) {
+            return res;
         }
-        if (this.irradianceTexture) {
-            this.irradianceTexture.dispose();
-            this.irradianceTexture = null;
+
+        if (this._context && this._pmremGenerator && texture) {
+            if (debug) console.log("[MaterialX] Generating environment textures", texture.name);
+
+            const target = this._pmremGenerator.fromEquirectangular(texture);
+            const radianceRenderTarget = renderPMREMToEquirect(this._context.renderer, target.texture, 0.0, 1024, 512, target.height);
+            const irradianceRenderTarget = renderPMREMToEquirect(this._context.renderer, target.texture, 1.0, 32, 16, target.height);
+            target.dispose();
+            res = {
+                radianceTexture: radianceRenderTarget.texture,
+                irradianceTexture: irradianceRenderTarget.texture
+            }
+        }
+        else {
+            res = {
+                radianceTexture: null,
+                irradianceTexture: null
+            }
+        }
+        this._texturesCache.set(texture || null, res);
+        return res;
+    }
+
+    private updateLighting = (collectLights: boolean = false) => {
+        if (!this._context) return;
+
+        // Find lights in scene
+        if (collectLights) {
+            const lights = new Array<Light>();
+            this._context.scene.traverse((object: Object3D) => {
+                if ((object as Light).isLight && GameObject.isActiveInHierarchy(object))
+                    lights.push(object as Light);
+            });
+            this._lights = lights;
+        }
+
+        if (state.materialXGenContext) {
+            const { lightData, lightCount } = getLightData(this._lights, state.materialXGenContext);
+            this._lightData = lightData;
+            this._lightCount = lightCount;
         }
-        this.initialized = false;
-        // this.lights = [];
-        this.lightData = null;
     }
 }

package/src/materialx.types.d.ts

@@ -0,0 +1,50 @@
+
+
+
+export namespace MaterialX {
+
+    export type MODULE = {
+        ShaderInterfaceType: any;
+        HwSpecularEnvironmentMethod: any;
+        HwShaderGenerator: {
+            bindLightShader(def: any, id: number, genContext: GenContext): void;
+            unbindLightShaders(context: any): void;
+        };
+        createDocument(): Document;
+        readFromXmlString(doc: Document, xml: string, unknown: string): void;
+        loadStandardLibraries(genContext: GenContext): StandardLibrary;
+        isTransparentSurface(renderableElement: any, target: string): boolean;
+    }
+
+
+    export type GenContext = {
+    }
+
+    export type StandardLibrary = {
+
+    }
+    
+    // https://github.com/AcademySoftwareFoundation/MaterialX/blob/b74787db6544283dc32afc8085ebc93cabe937cb/source/MaterialXGenShader/ShaderStage.h#L56
+    export type ShaderStage = {
+        getUniformBlocks(): Record<string, any>;
+    }
+
+    export type Document = {
+        setDataLibrary(lib: StandardLibrary): void;
+        importLibrary(lib: Document): void;
+
+        getNodes(): Node[];
+    }
+
+    export type Node = {
+        getType(): string;
+    }
+
+    export type Matrix = {
+        numRows(): number;
+        numColumns(): number;
+        get size(): number;
+        getItem(row: number, col: number): number;
+    }
+
+}