Skip to content

Commit

Permalink
NodeMaterial: Introduce .geometryNode and jelly example (#29551)
Browse files Browse the repository at this point in the history 
* fix `compute()` during `render()` call

* NodeMaterial: Introduce `.geometryNode`

* update compute geometry example

* rename jellyNode -> jelly

* cleanup
  • Loading branch information
@sunag
sunag authored Oct 5, 2024
1 parent ea86688 commit 5f96ed4
Show file tree
Hide file tree
Showing 4 changed files with 148 additions and 45 deletions.
Binary file modified examples/screenshots/webgpu_compute_geometry.jpg
View file
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
179 changes: 134 additions & 45 deletions examples/webgpu_compute_geometry.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,83 +25,139 @@
<script type="module">

import * as THREE from 'three';
import { vec3, cos, sin, mat3, storage, Fn, instanceIndex, timerLocal } from 'three/tsl';
import { vec3, vec4, storage, Fn, If, uniform, instanceIndex, objectWorldMatrix, color, screenUV, attribute } from 'three/tsl';

import { GLTFLoader } from 'three/addons/loaders/GLTFLoader.js';

import { OrbitControls } from 'three/addons/controls/OrbitControls.js';

import { GUI } from 'three/addons/libs/lil-gui.module.min.js';

import Stats from 'three/addons/libs/stats.module.js';

let camera, scene, renderer;
let computeUpdate;
let raycaster, pointer;
let stats;

const pointerPosition = uniform( vec4( 0 ) );
const elasticity = uniform( .4 ); // elasticity ( how "strong" the spring is )
const damping = uniform( .94 ); // damping factor ( energy loss )
const brushSize = uniform( .25 );
const brushStrength = uniform( .22 );

init();

function init() {
const jelly = Fn( ( { renderer, geometry, object } ) => {

camera = new THREE.PerspectiveCamera( 50, window.innerWidth / window.innerHeight, 0.1, 10 );
camera.position.set( 0, 0, 1 );
const count = geometry.attributes.position.count;

scene = new THREE.Scene();
scene.background = new THREE.Color( 0x333333 );
const speedBufferAttribute = new THREE.StorageBufferAttribute( count, 4 );

new GLTFLoader().load( 'models/gltf/LeePerrySmith/LeePerrySmith.glb', function ( gltf ) {
// replace geometry attributes for storage buffer attributes

const mesh = gltf.scene.children[ 0 ];
mesh.scale.setScalar( .1 );
mesh.material = new THREE.MeshNormalMaterial();
scene.add( mesh );
const positionBaseAttribute = geometry.attributes.position;
const positionStorageBufferAttribute = new THREE.StorageBufferAttribute( count, 4 );

geometry.setAttribute( 'storagePosition', positionStorageBufferAttribute );

//
// compute ( jelly )

const positionBaseAttribute = mesh.geometry.attributes.position;
const normalBaseAttribute = mesh.geometry.attributes.normal;
const positionAttribute = storage( positionBaseAttribute, 'vec3', count ).toReadOnly();
const positionStorageAttribute = storage( positionStorageBufferAttribute, 'vec4', count );

// replace geometry attributes for storage buffer attributes
const speedAttribute = storage( speedBufferAttribute, 'vec4', count );

const positionStorageBufferAttribute = new THREE.StorageBufferAttribute( positionBaseAttribute.count, 4 );
const normalStorageBufferAttribute = new THREE.StorageBufferAttribute( normalBaseAttribute.count, 4 );
// vectors

mesh.geometry.setAttribute( 'position', positionStorageBufferAttribute );
mesh.geometry.setAttribute( 'normal', normalStorageBufferAttribute );
const basePosition = vec3( positionAttribute.element( instanceIndex ) );
const currentPosition = positionStorageAttribute.element( instanceIndex );
const currentSpeed = speedAttribute.element( instanceIndex );

// compute shader
//

const computeFn = Fn( () => {
const computeInit = Fn( () => {

const positionAttribute = storage( positionBaseAttribute, 'vec3', positionBaseAttribute.count ).toReadOnly();
const normalAttribute = storage( normalBaseAttribute, 'vec3', normalBaseAttribute.count ).toReadOnly();
// copy position to storage

const positionStorageAttribute = storage( positionStorageBufferAttribute, 'vec4', positionStorageBufferAttribute.count );
const normalStorageAttribute = storage( normalStorageBufferAttribute, 'vec4', normalStorageBufferAttribute.count );
currentPosition.assign( basePosition );

const time = timerLocal( 1 );
const scale = 0.3;
} )().compute( count );

//
//

const position = vec3( positionAttribute.element( instanceIndex ) );
const normal = vec3( normalAttribute.element( instanceIndex ) );
const computeUpdate = Fn( () => {

const theta = sin( time.add( position.y ) ).mul( scale );
// pinch

const c = cos( theta );
const s = sin( theta );
If( pointerPosition.w.equal( 1 ), () => {

const m = mat3(
c, 0, s,
0, 1, 0,
s.negate(), 0, c
);
const worldPosition = objectWorldMatrix( object ).mul( currentPosition.xyz );

const transformed = position.mul( m );
const transformedNormal = normal.mul( m );
const dist = worldPosition.distance( pointerPosition.xyz );
const direction = pointerPosition.xyz.sub( worldPosition ).normalize();

positionStorageAttribute.element( instanceIndex ).assign( transformed );
normalStorageAttribute.element( instanceIndex ).assign( transformedNormal );
const power = brushSize.sub( dist ).max( 0 ).mul( brushStrength );

currentPosition.xyz.addAssign( direction.mul( power ) );

} );

computeUpdate = computeFn().compute( positionBaseAttribute.count );
// update

const distance = basePosition.distance( currentPosition );
const force = elasticity.mul( distance ).mul( basePosition.sub( currentPosition ) );

currentSpeed.addAssign( force );
currentSpeed.mulAssign( damping );

currentPosition.addAssign( currentSpeed );

} )().compute( count );

// initialize the storage buffer with the base position

computeUpdate.onInit( () => renderer.compute( computeInit ) );

//

return computeUpdate;

} );

function init() {

camera = new THREE.PerspectiveCamera( 50, window.innerWidth / window.innerHeight, 0.1, 10 );
camera.position.set( 0, 0, 1 );

scene = new THREE.Scene();

raycaster = new THREE.Raycaster();
pointer = new THREE.Vector2();

// background

const bgColor = screenUV.y.mix( color( 0x9f87f7 ), color( 0xf2cdcd ) );
const bgVignet = screenUV.distance( .5 ).remapClamp( 0.3, .8 ).oneMinus();
const bgIntensity = 4;

scene.backgroundNode = bgColor.mul( bgVignet.mul( color( 0xa78ff6 ).mul( bgIntensity ) ) );

// model

new GLTFLoader().load( 'models/gltf/LeePerrySmith/LeePerrySmith.glb', function ( gltf ) {

// create jelly effect material

const material = new THREE.MeshNormalNodeMaterial();
material.geometryNode = jelly();
material.positionNode = attribute( 'storagePosition' );

// apply the material to the mesh

const mesh = gltf.scene.children[ 0 ];
mesh.scale.setScalar( .1 );
mesh.material = material;
scene.add( mesh );

} );

Expand All @@ -117,7 +173,40 @@
controls.minDistance = .7;
controls.maxDistance = 2;

const gui = new GUI();
gui.add( elasticity, 'value', 0, .5 ).name( 'elasticity' );
gui.add( damping, 'value', .9, .98 ).name( 'damping' );
gui.add( brushSize, 'value', .1, .5 ).name( 'brush size' );
gui.add( brushStrength, 'value', .1, .3 ).name( 'brush strength' );

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

window.addEventListener( 'resize', onWindowResize );
window.addEventListener( 'pointermove', onPointerMove );

}

function onPointerMove( event ) {

pointer.set( ( event.clientX / window.innerWidth ) * 2 - 1, - ( event.clientY / window.innerHeight ) * 2 + 1 );

raycaster.setFromCamera( pointer, camera );

const intersects = raycaster.intersectObject( scene );

if ( intersects.length > 0 ) {

const intersect = intersects[ 0 ];

pointerPosition.value.copy( intersect.point );
pointerPosition.value.w = 1; // enable

} else {

pointerPosition.value.w = 0; // disable

}

}

Expand All @@ -132,7 +221,7 @@

async function animate() {

if ( computeUpdate ) renderer.compute( computeUpdate );
stats.update();

renderer.render( scene, camera );

Expand Down
8 changes: 8 additions & 0 deletions src/materials/nodes/NodeMaterial.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ class NodeMaterial extends Material {
this.alphaTestNode = null;

this.positionNode = null;
this.geometryNode = null;

this.depthNode = null;
this.shadowNode = null;
Expand Down Expand Up @@ -98,6 +99,12 @@ class NodeMaterial extends Material {

builder.stack.outputNode = this.vertexNode || this.setupPosition( builder );

if ( this.geometryNode !== null ) {

builder.stack.outputNode = builder.stack.outputNode.bypass( this.geometryNode );

}

builder.addFlow( 'vertex', builder.removeStack() );

// < FRAGMENT STAGE >
Expand Down Expand Up @@ -634,6 +641,7 @@ class NodeMaterial extends Material {
this.alphaTestNode = source.alphaTestNode;

this.positionNode = source.positionNode;
this.geometryNode = source.geometryNode;

this.depthNode = source.depthNode;
this.shadowNode = source.shadowNode;
Expand Down
6 changes: 6 additions & 0 deletions src/renderers/webgl-fallback/WebGLBackend.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -605,6 +605,12 @@ class WebGLBackend extends Backend {

this.prepareTimestampBuffer( computeGroup );

if ( this._currentContext ) {

this._setFramebuffer( this._currentContext );

}

}

draw( renderObject/*, info*/ ) {
Expand Down

0 comments on commit 5f96ed4

Please sign in to comment.