@needle-tools/materialx 1.0.0-alpha

package/src/materialx.ts

@@ -0,0 +1,220 @@
+import { Context, delay, isDevEnvironment, ObjectUtils } from "@needle-tools/engine";
+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
+export const state = new class {
+    materialXModule: any = null;
+    materialXGenerator: any = null;
+    materialXGenContext: any = null;
+    materialXStdLib: any = null;
+    materialXInitPromise: Promise<void> | null = null;
+
+    // Global MaterialX environment instance
+    private _materialXEnvironment: MaterialXEnvironment | null = null;
+    get materialXEnvironment() {
+        this._materialXEnvironment ??= new MaterialXEnvironment();
+        return this._materialXEnvironment;
+    }
+}
+
+
+// Initialize MaterialX WASM module lazily
+export async function initializeMaterialX(): 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...");
+        try {
+            const module = await MaterialX(materialXConfig);
+            if (debug) console.log("MaterialXLoader module loaded", module);
+            state.materialXModule = module;
+
+            // Initialize shader generator and context
+            state.materialXGenerator = module.EsslShaderGenerator.create();
+            state.materialXGenContext = new module.GenContext(state.materialXGenerator);
+
+            // Load standard libraries
+            const tempDoc = module.createDocument();
+            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>`;
+
+            // This prewarms the shader generation context to have all light types
+            await registerLights(state.materialXModule, [], state.materialXGenContext);
+
+            if (debug) console.log("MaterialX generator initialized successfully");
+        } catch (error) {
+            console.error("Failed to load MaterialX module:", error);
+            throw error;
+        }
+    })();
+}
+
+// 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;
+
+    constructor() {
+        if (debug) console.log("MaterialX Environment created");
+    }
+
+    setContext(context: Context) {
+        this.context = 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;
+        }
+
+        registerLights(materialXModule, this.lights, materialXGenContext);
+    }
+    */
+
+    // Initialize with Needle Engine context
+    async initializeFromContext(): Promise<void> {
+        if (!this.context) {
+            console.warn("No Needle context available for MaterialX environment initialization");
+            return;
+        }
+
+        // Prevent multiple initializations
+        if (this.initialized) {
+            if (debug) console.log("MaterialX environment already initialized, skipping");
+            return;
+        }
+
+        // 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;
+        }
+
+        // Get renderer from context
+        const renderer = this.context.renderer;
+
+        // 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);
+        }
+
+        // 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);
+        });
+
+        this.lightData = await registerLights(state.materialXModule, lights, state.materialXGenContext);
+
+        // Mark as initialized
+        this.initialized = 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 (this.irradianceTexture) {
+            this.irradianceTexture.dispose();
+            this.irradianceTexture = null;
+        }
+        this.initialized = false;
+        // this.lights = [];
+        this.lightData = null;
+    }
+}

package/src/textureHelper.ts

@@ -0,0 +1,170 @@
+import { WebGLRenderer, Scene, WebGLRenderTarget, PlaneGeometry, OrthographicCamera, ShaderMaterial, RGBAFormat, FloatType, LinearFilter, Mesh, EquirectangularReflectionMapping, RepeatWrapping, LinearMipMapLinearFilter, Texture, WebGLUtils } from 'three';
+import { getParam } from '@needle-tools/engine';
+
+const debug = getParam("debugmaterialx");
+
+/**
+ * Renders a PMREM environment map to an equirectangular texture with specified roughness
+ * @param {WebGLRenderer} renderer - Three.js WebGL renderer
+ * @param {Texture} pmremTexture - PMREM texture (2D CubeUV layout) to convert
+ * @param {number} roughness - Roughness value (0.0 to 1.0)
+ * @param {number} width - Output texture width (default: 1024)
+ * @param {number} height - Output texture height (default: 512)
+ * @param {number} renderTargetHeight - Original render target height (optional, for proper PMREM parameter calculation)
+ * @returns {WebGLRenderTarget} Render target containing the equirectangular texture
+ * @example // Creating an equirectangular texture from a PMREM environment map at a certain roughness level:
+ * const pmremRenderTarget = pmremGenerator.fromEquirectangular(envMap);
+ * const equirectRenderTarget = await renderPMREMToEquirect(renderer, pmremRenderTarget.texture, 0.5, 2048, 1024, pmremRenderTarget.height);
+    
+    // Use the rendered equirectangular texture
+    const equirectTexture = equirectRenderTarget.texture;
+    
+    // Apply to your material or save/export
+    someMaterial.map = equirectTexture;
+    
+    // Don't forget to dispose when done
+    // equirectRenderTarget.dispose();
+ */
+export function renderPMREMToEquirect(renderer: WebGLRenderer, pmremTexture: Texture, roughness = 0.0, width = 1024, height = 512, renderTargetHeight?: number) {
+    // TODO Validate inputs
+    // console.log(renderer, pmremTexture);
+
+    // Calculate PMREM parameters
+    // For PMREM CubeUV layout, we need the cube face size to calculate proper parameters
+    // Use renderTargetHeight if provided, otherwise try to derive from texture
+    let imageHeight;
+    if (renderTargetHeight) {
+        imageHeight = renderTargetHeight;
+    } else if (pmremTexture.image) {
+        imageHeight = pmremTexture.image.height / 4; // Fallback: assume CubeUV layout height / 4
+    } else {
+        imageHeight = 256; // Final fallback
+    }
+    
+    const maxMip = Math.log2(imageHeight) - 2;
+    const cubeUVHeight = imageHeight;
+    const cubeUVWidth = 3 * Math.max(Math.pow(2, maxMip), 7 * 16);
+
+    // Create render target for equirectangular output
+    const renderTarget = new WebGLRenderTarget(width, height, {
+        format: RGBAFormat,
+        type: FloatType,
+        minFilter: LinearMipMapLinearFilter,
+        magFilter: LinearFilter,
+        generateMipmaps: true,
+        wrapS: RepeatWrapping,
+        anisotropy: renderer.capabilities.getMaxAnisotropy(),
+    });
+
+    // Create fullscreen quad geometry and camera
+    const geometry = new PlaneGeometry(2, 2);
+    const camera = new OrthographicCamera(-1, 1, 1, -1, 0, 1);
+
+    // Create shader material for PMREM to equirectangular conversion
+    const material = new ShaderMaterial({
+        defines: {
+            USE_ENVMAP: '',
+            ENVMAP_TYPE_CUBE_UV: '',
+            CUBEUV_TEXEL_WIDTH: 1.0 / cubeUVWidth,
+            CUBEUV_TEXEL_HEIGHT: 1.0 / cubeUVHeight,
+            CUBEUV_MAX_MIP: (maxMip + 0) + '.0',
+        },
+        uniforms: {
+            envMap: { value: pmremTexture },
+            roughness: { value: roughness }
+        },
+        vertexShader: `
+            varying vec2 vUv;
+            
+            void main() {
+                vUv = uv;
+                gl_Position = vec4(position.xy, 0.0, 1.0);
+            }
+        `,
+        fragmentShader: `
+            uniform sampler2D envMap;
+            uniform float roughness;
+            varying vec2 vUv;
+            
+            #include <common>
+            #include <cube_uv_reflection_fragment>
+            
+            void main() {
+                // Convert UV coordinates to equirectangular direction
+                vec2 uv = vUv;
+                
+                // Map UV (0,1) to spherical coordinates
+                // Longitude: -π to π, Latitude: 0 to π
+                float phi = uv.x * 2.0 * PI - PI; // Longitude (-π to π)
+                float theta = uv.y * PI;          // Latitude (0 to π)
+                // Rotate 90° around Y
+                phi -= PI / 2.0; // Adjust to match Three.js convention
+
+                // Convert spherical to cartesian coordinates
+                vec3 direction = vec3(
+                    sin(theta) * cos(phi),  // x
+                    cos(theta),             // y  
+                    sin(theta) * sin(phi)   // z
+                );
+                
+                // Sample the PMREM cube texture using the direction and roughness
+                #ifdef ENVMAP_TYPE_CUBE_UV
+                    vec4 envColor = textureCubeUV(envMap, direction, roughness);
+                #else
+                    vec4 envColor = vec4(1.0, 0.0, 1.0, 1.0); // Magenta fallback
+                #endif
+
+                gl_FragColor = vec4(envColor.rgb, 1.0);
+            }
+        `
+    });
+
+    // Create temporary scene and mesh for rendering
+    const tempScene = new Scene();
+    const mesh = new Mesh(geometry, material);
+    tempScene.add(mesh);
+
+    // Store current renderer state
+    const currentRenderTarget = renderer.getRenderTarget();
+    const currentAutoClear = renderer.autoClear;
+    const currentXrEnabled = renderer.xr.enabled;
+    const currentShadowMapEnabled = renderer.shadowMap.enabled;
+    
+    renderTarget.texture.generateMipmaps = true;
+    
+    try {
+        // Disable XR and shadow mapping during our render to avoid interference
+        renderer.xr.enabled = false;
+        renderer.shadowMap.enabled = false;
+        
+        // Render to our target
+        renderer.autoClear = true;
+        renderer.setRenderTarget(renderTarget);
+        renderer.clear(); // Explicitly clear the render target
+        renderer.render(tempScene, camera);
+    } finally {
+        // Restore renderer state completely
+        renderer.setRenderTarget(currentRenderTarget);
+        renderer.autoClear = currentAutoClear;
+        renderer.xr.enabled = currentXrEnabled;
+        renderer.shadowMap.enabled = currentShadowMapEnabled;
+        
+        // Clean up temporary objects
+        geometry.dispose();
+        material.dispose();
+        tempScene.remove(mesh);
+    }
+
+    renderTarget.texture.name = 'PMREM_Equirectangular_Texture_' + roughness.toFixed(2);
+    renderTarget.texture.mapping = EquirectangularReflectionMapping;
+
+    // Log mipmap infos
+    if (debug) console.log('PMREM to Equirect Render Target:', {
+        width: renderTarget.width,
+        height: renderTarget.height,
+        mipmaps: renderTarget.texture.mipmaps?.length,
+        roughness: roughness,
+    });
+
+    return renderTarget;
+}

package/src/utils.ts

@@ -0,0 +1,58 @@
+import { Context, getParam } from "@needle-tools/engine";
+import { Mesh } from "three";
+
+export const debug = getParam("debugmaterialx");
+
+
+
+
+
+/**
+ * =====================================
+ * Unused
+ */
+
+// Patch WebGL2 methods for debugging purposes
+
+const patchWebGL2 = () => {
+    const getUniformLocation = WebGL2RenderingContext.prototype.getUniformLocation;
+    const programAndNameToUniformLocation = new WeakMap<WebGLUniformLocation, { program: WebGLProgram, name: string }>();
+    WebGL2RenderingContext.prototype.getUniformLocation = function (program: WebGLProgram, name: string) {
+        const location = getUniformLocation.call(this, program, name);
+        if (location) {
+            programAndNameToUniformLocation.set(location, { program, name });
+        }
+        return location;
+    };
+
+    const uniform4fv = WebGL2RenderingContext.prototype.uniform4fv;
+    WebGL2RenderingContext.prototype.uniform4fv = function (location: WebGLUniformLocation | null, v: Float32Array | number[]) {
+        if (location) {
+            const uniformName = programAndNameToUniformLocation.get(location);
+            if (true) console.log("Calling uniform4fv", { location, v, name: uniformName?.name });
+        }
+        return uniform4fv.call(this, location, v);
+    };
+};
+// patchWebGL2();
+
+
+// TODO doesn't actually reset yet...
+function resetShaders(ctx: Context) {
+    const scene = ctx.scene;
+    const gl = ctx.renderer;
+    console.log(gl.properties, gl.info)
+    scene.traverse(object => {
+        if ((object as Mesh).isMesh) {
+            const mesh = object as Mesh;
+            if (Array.isArray(mesh.material)) {
+                mesh.material.forEach(mat => gl.properties.remove(mat));
+            }
+            else {
+                gl.properties.remove(mesh.material);
+            }
+        }
+    })
+    if (gl.info?.programs)
+        gl.info.programs.length = 0;
+}

package/tsconfig.json

@@ -0,0 +1,20 @@
+{
+  "compilerOptions": {
+    "target": "ESNext",
+    "useDefineForClassFields": true,
+    "module": "ESNext",
+    "lib": ["ESNext", "DOM"],
+    "moduleResolution": "Node",
+    "strict": true,
+    "sourceMap": true,
+    "resolveJsonModule": true,
+    "esModuleInterop": true,
+    "noEmit": true,
+    "noUnusedLocals": false,
+    "noUnusedParameters": true,
+    "noImplicitReturns": true,
+    "noImplicitAny": false,
+    "experimentalDecorators": true
+  },
+  "include": ["."]
+}

package/workspace.code-workspace

@@ -0,0 +1,18 @@
+{
+  "folders": [
+    {
+      "path": "."
+    },
+    {
+      "name": "Needle",
+      "path": "./node_modules/@needle-tools"
+    }
+  ],
+  "settings": {
+    "files.exclude": {
+      "**/.git": true,
+      "**/.DS_Store": true,
+      "**/*.meta": true
+    }
+  }
+}