InstancedBatchedSkinnedMesh

InstancedBatchedSkinnedMesh.js

// from https://x.com/Cody_J_Bennett/status/1818025947565584873
// demo https://jsfiddle.net/cbenn/Las0poyu

import * as THREE from "three";
import Stats from "three/addons/libs/stats.module.js";
import { OrbitControls } from "three/addons/controls/OrbitControls.js";
import { GLTFLoader } from "three/addons/loaders/GLTFLoader.js";
import { mergeVertices } from "three/addons/utils/BufferGeometryUtils.js";

THREE.ShaderChunk.skinning_pars_vertex =
    THREE.ShaderChunk.skinning_pars_vertex +
    /* glsl */ `
      #ifdef USE_BATCHED_SKINNING

          attribute vec4 skinIndex;
          attribute vec4 skinWeight;

          uniform highp usampler2D batchingKeyframeTexture;
          uniform highp sampler2D boneTexture;

          float getBatchedKeyframe( const in float batchId ) {

              int size = textureSize( batchingKeyframeTexture, 0 ).x;
              int j = int ( batchId );
              int x = j % size;
              int y = j / size;
              return float( texelFetch( batchingKeyframeTexture, ivec2( x, y ), 0 ).r );

          }

          mat4 getBatchedBoneMatrix( const in float i ) {

              float batchId = getIndirectIndex( gl_DrawID );
              float batchKeyframe = getBatchedKeyframe( batchId );

              int size = textureSize( boneTexture, 0 ).x;
              int j = int( batchKeyframe + i ) * 4;
              int x = j % size;
              int y = j / size;
              vec4 v1 = texelFetch( boneTexture, ivec2( x, y ), 0 );
              vec4 v2 = texelFetch( boneTexture, ivec2( x + 1, y ), 0 );
              vec4 v3 = texelFetch( boneTexture, ivec2( x + 2, y ), 0 );
              vec4 v4 = texelFetch( boneTexture, ivec2( x + 3, y ), 0 );

              return mat4( v1, v2, v3, v4 );

          }

      #endif
  `;

THREE.ShaderChunk.skinning_vertex =
    THREE.ShaderChunk.skinning_vertex +
    /* glsl */ `
      #ifdef USE_BATCHED_SKINNING

          vec4 skinVertex = vec4( transformed, 1.0 );

          mat4 boneMatX = getBatchedBoneMatrix( skinIndex.x );
          mat4 boneMatY = getBatchedBoneMatrix( skinIndex.y );
          mat4 boneMatZ = getBatchedBoneMatrix( skinIndex.z );
          mat4 boneMatW = getBatchedBoneMatrix( skinIndex.w );

          vec4 skinned = vec4( 0.0 );
          skinned += boneMatX * skinVertex * skinWeight.x;
          skinned += boneMatY * skinVertex * skinWeight.y;
          skinned += boneMatZ * skinVertex * skinWeight.z;
          skinned += boneMatW * skinVertex * skinWeight.w;

          transformed = skinned.xyz;

      #endif
  `;

const _offsetMatrix = new THREE.Matrix4();

export class InstancedBatchedSkinnedMesh extends THREE.BatchedMesh {
    skeletons = [];
    clips = [];
    animationIds = [];
    offsets = [];
    times;
    fps = 60;

    boneTexture = null;

    constructor(maxInstanceCount, maxVertexCount, maxIndexCount, material) {
        super(maxInstanceCount, maxVertexCount, maxIndexCount, material);

        this.animationIds = Array(maxInstanceCount).fill(-1);
        this.offsets = Array(maxInstanceCount).fill(0);
        this.times = Array(maxInstanceCount).fill(0);

        let size = Math.sqrt(maxInstanceCount);
        size = Math.ceil(size);

        this.batchingKeyframeTexture = new THREE.DataTexture(
            new Uint32Array(size * size),
            size,
            size,
            THREE.RedIntegerFormat,
            THREE.UnsignedIntType
        );
        this.batchingKeyframeTexture.needsUpdate = true;

        this.material.onBeforeCompile = (shader) => {
            if (this.boneTexture === null) this.computeBoneTexture();

            shader.defines ??= {};
            shader.defines.USE_BATCHED_SKINNING = "";
            shader.uniforms.batchingKeyframeTexture = {
                value: this.batchingKeyframeTexture,
            };
            shader.uniforms.boneTexture = { value: this.boneTexture };
        };
    }

    addAnimation(skeleton, clip) {
        clip.optimize();

        this.skeletons.push(skeleton);
        this.clips.push(clip);

        return this.skeletons.length - 1;
    }

    setAnimationAt(instanceId, animationId) {
        this.animationIds[instanceId] = animationId;
    }

    computeBoneTexture() {
        let offset = 0;
        for (let i = 0; i < this.skeletons.length; i++) {
            const skeleton = this.skeletons[i];
            const clip = this.clips[i];

            const steps = Math.ceil(clip.duration * this.fps);
            this.offsets[i] = offset;
            offset += skeleton.bones.length * steps;
        }

        let size = Math.sqrt(offset * 4);
        size = Math.ceil(size / 4) * 4;
        size = Math.max(size, 4);

        const boneMatrices = new Float32Array(size * size * 4);
        this.boneTexture = new THREE.DataTexture(
            boneMatrices,
            size,
            size,
            THREE.RGBAFormat,
            THREE.FloatType
        );
        this.boneTexture.needsUpdate = true;

        for (let i = 0; i < this.skeletons.length; i++) {
            const skeleton = this.skeletons[i];
            const clip = this.clips[i];

            const steps = Math.ceil(clip.duration * this.fps);
            const offset = this.offsets[i];

            let root = skeleton.bones[0];
            while (root.parent !== null && root.parent instanceof THREE.Bone) {
                root = root.parent;
            }
            const mixer = new THREE.AnimationMixer(root);

            const action = mixer.clipAction(clip);
            action.play();

            for (let j = 0; j < steps; j++) {
                mixer.update(1 / this.fps);
                root.updateMatrixWorld(true);

                for (let k = 0; k < skeleton.bones.length; k++) {
                    const matrix = skeleton.bones[k].matrixWorld;
                    _offsetMatrix.multiplyMatrices(
                        matrix,
                        skeleton.boneInverses[k]
                    );
                    _offsetMatrix.toArray(
                        boneMatrices,
                        (offset + (j * skeleton.bones.length + k)) * 16
                    );
                }
            }
        }
    }

    update(delta) {
        for (let i = 0; i < this.maxInstanceCount; i++) {
            const animationId = this.animationIds[i];
            if (animationId === -1) continue;

            const skeleton = this.skeletons[animationId];
            const clip = this.clips[animationId];

            const steps = Math.ceil(clip.duration * this.fps);
            const offset = this.offsets[animationId];

            this.times[i] += delta;
            this.times[i] = THREE.MathUtils.clamp(
                this.times[i] -
                    Math.floor(this.times[i] / clip.duration) * clip.duration,
                0,
                clip.duration
            );

            const frame = Math.floor(this.times[i] * this.fps) % steps;

            this.batchingKeyframeTexture.image.data[i] =
                offset + frame * skeleton.bones.length;
        }
        this.batchingKeyframeTexture.needsUpdate = true;
    }
}

const renderer = new THREE.WebGLRenderer({ alpha: true });
document.body.appendChild(renderer.domElement);

const camera = new THREE.PerspectiveCamera(35);
camera.position.set(-5, 4, 5);

const controls = new OrbitControls(camera, renderer.domElement);
controls.target.set(0, 2, 0);
controls.enableDamping = true;

const scene = new THREE.Scene();

const loader = new GLTFLoader();
const gltf = await loader.loadAsync(
    "https://raw.githack.com/KhronosGroup/glTF-Sample-Assets/main/Models/Fox/glTF-Binary/Fox.glb"
);

let fox;
gltf.scene.traverse((node) => {
    if (node instanceof THREE.SkinnedMesh) {
        fox = node;
    }
});

const columns = 25;
const rows = 25;

const geometry = mergeVertices(fox.geometry);

const batchedMesh = new InstancedBatchedSkinnedMesh(
    columns * rows,
    geometry.attributes.position.count,
    geometry.index.count
);
const geometryId = batchedMesh.addGeometry(geometry);
for (const animation of gltf.animations)
    batchedMesh.addAnimation(fox.skeleton, animation);

const matrix = new THREE.Matrix4();
matrix.makeScale(0.01, 0.01, 0.01);

const color = new THREE.Color();

for (let x = 0; x < columns; x++) {
    for (let y = 0; y < rows; y++) {
        const index = y * columns + x;
        matrix.setPosition(x * 1.2, 0, -y * 1.2);

        const instanceId = batchedMesh.addInstance(geometryId);
        batchedMesh.setMatrixAt(instanceId, matrix);
        batchedMesh.setColorAt(
            instanceId,
            color.setHex(Math.random() * 0xffffff)
        );
        batchedMesh.setAnimationAt(
            instanceId,
            ((gltf.animations.length - 1) * Math.random()) | 0
        );

        batchedMesh.times[index] = x;
    }
}
scene.add(batchedMesh);

onresize = () => {
    renderer.setSize(innerWidth, innerHeight);
    renderer.setPixelRatio(devicePixelRatio);
    camera.aspect = innerWidth / innerHeight;
    camera.updateProjectionMatrix();
};
onresize();

const stats = new Stats();
document.body.appendChild(stats.dom);

let prev = performance.now();

renderer.setAnimationLoop((time) => {
    const delta = (time - prev) / 1000;
    prev = time;

    controls.update(delta);
    batchedMesh.update(delta);
    renderer.render(scene, camera);
    stats.update();
});