Skip to content

Commit 893add1

Browse files
authored
TAAUNode: Adding new node for TAA with Upsampling. (#33368)
1 parent 291280b commit 893add1

6 files changed

Lines changed: 1324 additions & 2 deletions

File tree

examples/files.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -493,6 +493,7 @@
493493
"webgpu_tsl_vfx_tornado",
494494
"webgpu_tsl_wood",
495495
"webgpu_upscaling_fsr1",
496+
"webgpu_upscaling_taau",
496497
"webgpu_video_frame",
497498
"webgpu_video_panorama",
498499
"webgpu_volume_caustics",
Lines changed: 283 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,283 @@
1+
import { HalfFloatType, RenderTarget, Vector2, NodeMaterial, RendererUtils, QuadMesh, TempNode, NodeUpdateType } from 'three/webgpu';
2+
import { Fn, float, vec3, vec4, ivec2, int, uv, floor, abs, max, min, exp2, nodeObject, passTexture, textureSize, textureLoad, convertToTexture } from 'three/tsl';
3+
4+
const _quadMesh = /*@__PURE__*/ new QuadMesh();
5+
const _size = /*@__PURE__*/ new Vector2();
6+
7+
let _rendererState;
8+
9+
/**
10+
* Post processing node for contrast-adaptive sharpening (RCAS).
11+
*
12+
* Reference: {@link https://gpuopen.com/fidelityfx-superresolution/}.
13+
*
14+
* @augments TempNode
15+
* @three_import import { sharpen } from 'three/addons/tsl/display/SharpenNode.js';
16+
*/
17+
class SharpenNode extends TempNode {
18+
19+
static get type() {
20+
21+
return 'SharpenNode';
22+
23+
}
24+
25+
/**
26+
* Constructs a new sharpen node.
27+
*
28+
* @param {TextureNode} textureNode - The texture node that represents the input of the effect.
29+
* @param {Node<float>} [sharpness=0.2] - Sharpening strength. 0 = maximum sharpening, 2 = no sharpening.
30+
* @param {Node<bool>} [denoise=false] - Whether to attenuate sharpening in noisy areas.
31+
*/
32+
constructor( textureNode, sharpness = 0.2, denoise = false ) {
33+
34+
super( 'vec4' );
35+
36+
/**
37+
* This flag can be used for type testing.
38+
*
39+
* @type {boolean}
40+
* @readonly
41+
* @default true
42+
*/
43+
this.isSharpenNode = true;
44+
45+
/**
46+
* The texture node that represents the input of the effect.
47+
*
48+
* @type {TextureNode}
49+
*/
50+
this.textureNode = textureNode;
51+
52+
/**
53+
* Sharpening strength. 0 = maximum, 2 = none.
54+
*
55+
* @type {Node<float>}
56+
* @default 0.2
57+
*/
58+
this.sharpness = nodeObject( sharpness );
59+
60+
/**
61+
* Whether to attenuate sharpening in noisy areas.
62+
*
63+
* @type {Node<bool>}
64+
* @default false
65+
*/
66+
this.denoise = nodeObject( denoise );
67+
68+
/**
69+
* The render target for the sharpening pass.
70+
*
71+
* @private
72+
* @type {RenderTarget}
73+
*/
74+
this._renderTarget = new RenderTarget( 1, 1, { depthBuffer: false, type: HalfFloatType } );
75+
this._renderTarget.texture.name = 'SharpenNode.output';
76+
77+
/**
78+
* The result of the effect as a texture node.
79+
*
80+
* @private
81+
* @type {PassTextureNode}
82+
*/
83+
this._textureNode = passTexture( this, this._renderTarget.texture );
84+
85+
/**
86+
* The material for the sharpening pass.
87+
*
88+
* @private
89+
* @type {?NodeMaterial}
90+
*/
91+
this._material = null;
92+
93+
/**
94+
* The `updateBeforeType` is set to `NodeUpdateType.FRAME` since the node renders
95+
* its effect once per frame in `updateBefore()`.
96+
*
97+
* @type {string}
98+
* @default 'frame'
99+
*/
100+
this.updateBeforeType = NodeUpdateType.FRAME;
101+
102+
}
103+
104+
/**
105+
* Sets the output size of the effect.
106+
*
107+
* @param {number} width - The width in pixels.
108+
* @param {number} height - The height in pixels.
109+
*/
110+
setSize( width, height ) {
111+
112+
this._renderTarget.setSize( width, height );
113+
114+
}
115+
116+
/**
117+
* This method is used to render the effect once per frame.
118+
*
119+
* @param {NodeFrame} frame - The current node frame.
120+
*/
121+
updateBefore( frame ) {
122+
123+
const { renderer } = frame;
124+
125+
_rendererState = RendererUtils.resetRendererState( renderer, _rendererState );
126+
127+
//
128+
129+
renderer.getDrawingBufferSize( _size );
130+
this.setSize( _size.x, _size.y );
131+
132+
renderer.setRenderTarget( this._renderTarget );
133+
134+
_quadMesh.material = this._material;
135+
_quadMesh.name = 'Sharpen [ RCAS ]';
136+
_quadMesh.render( renderer );
137+
138+
//
139+
140+
RendererUtils.restoreRendererState( renderer, _rendererState );
141+
142+
}
143+
144+
/**
145+
* Returns the result of the effect as a texture node.
146+
*
147+
* @return {PassTextureNode} A texture node that represents the result of the effect.
148+
*/
149+
getTextureNode() {
150+
151+
return this._textureNode;
152+
153+
}
154+
155+
/**
156+
* This method is used to setup the effect's TSL code.
157+
*
158+
* @param {NodeBuilder} builder - The current node builder.
159+
* @return {PassTextureNode}
160+
*/
161+
setup( builder ) {
162+
163+
const textureNode = this.textureNode;
164+
const inputTex = textureNode.value;
165+
166+
// RCAS: Robust Contrast-Adaptive Sharpening.
167+
//
168+
// Ported from AMD FidelityFX FSR 1 (ffx_fsr1.h). Uses a 5-tap
169+
// cross pattern (center + up/down/left/right) to compute a
170+
// per-pixel sharpening weight that is automatically limited by
171+
// local contrast to avoid ringing. An optional noise-attenuation
172+
// factor reduces sharpening in noisy areas.
173+
174+
const rcas = Fn( () => {
175+
176+
const targetUV = uv();
177+
const texSize = textureSize( textureLoad( inputTex ) );
178+
179+
const p = ivec2( int( floor( targetUV.x.mul( texSize.x ) ) ), int( floor( targetUV.y.mul( texSize.y ) ) ) ).toConst();
180+
181+
const e = textureLoad( inputTex, p );
182+
const b = textureLoad( inputTex, p.add( ivec2( 0, - 1 ) ) );
183+
const d = textureLoad( inputTex, p.add( ivec2( - 1, 0 ) ) );
184+
const f = textureLoad( inputTex, p.add( ivec2( 1, 0 ) ) );
185+
const h = textureLoad( inputTex, p.add( ivec2( 0, 1 ) ) );
186+
187+
// Approximate luminance (luma times 2).
188+
189+
const luma = ( s ) => s.g.add( s.b.add( s.r ).mul( 0.5 ) );
190+
191+
const bL = luma( b );
192+
const dL = luma( d );
193+
const eL = luma( e );
194+
const fL = luma( f );
195+
const hL = luma( h );
196+
197+
// Sharpening amount from user parameter.
198+
199+
const con = exp2( this.sharpness.negate() ).toConst();
200+
201+
// Min and max of ring.
202+
203+
const mn4 = min( min( b.rgb, d.rgb ), min( f.rgb, h.rgb ) ).toConst();
204+
const mx4 = max( max( b.rgb, d.rgb ), max( f.rgb, h.rgb ) ).toConst();
205+
206+
// Compute adaptive lobe weight.
207+
// Limiters based on how much sharpening the local contrast can tolerate.
208+
209+
const RCAS_LIMIT = float( 0.25 - 1.0 / 16.0 ).toConst();
210+
211+
const hitMin = min( mn4, e.rgb ).div( mx4.mul( 4.0 ) ).toConst();
212+
const hitMax = vec3( 1.0 ).sub( max( mx4, e.rgb ) ).div( mn4.mul( 4.0 ).sub( 4.0 ) ).toConst();
213+
const lobeRGB = max( hitMin.negate(), hitMax ).toConst();
214+
215+
const lobe = max(
216+
RCAS_LIMIT.negate(),
217+
min( max( lobeRGB.r, max( lobeRGB.g, lobeRGB.b ) ), float( 0.0 ) )
218+
).mul( con ).toConst();
219+
220+
// Noise attenuation.
221+
222+
const nz = bL.add( dL ).add( fL ).add( hL ).mul( 0.25 ).sub( eL ).toConst();
223+
const nzRange = max( max( bL, dL ), max( eL, max( fL, hL ) ) ).sub( min( min( bL, dL ), min( eL, min( fL, hL ) ) ) ).toConst();
224+
const nzFactor = float( 1.0 ).sub( abs( nz ).div( max( nzRange, float( 1.0 / 65536.0 ) ) ).saturate().mul( 0.5 ) ).toConst();
225+
226+
const effectiveLobe = this.denoise.equal( true ).select( lobe.mul( nzFactor ), lobe ).toConst();
227+
228+
// Resolve: weighted blend of cross neighbors and center.
229+
230+
const result = b.rgb.add( d.rgb ).add( f.rgb ).add( h.rgb ).mul( effectiveLobe ).add( e.rgb )
231+
.div( effectiveLobe.mul( 4.0 ).add( 1.0 ) ).toConst();
232+
233+
return vec4( result, e.a );
234+
235+
} );
236+
237+
//
238+
239+
const context = builder.getSharedContext();
240+
241+
const material = this._material || ( this._material = new NodeMaterial() );
242+
material.fragmentNode = rcas().context( context );
243+
material.name = 'Sharpen_RCAS';
244+
material.needsUpdate = true;
245+
246+
//
247+
248+
const properties = builder.getNodeProperties( this );
249+
properties.textureNode = textureNode;
250+
251+
//
252+
253+
return this._textureNode;
254+
255+
}
256+
257+
/**
258+
* Frees internal resources. This method should be called
259+
* when the effect is no longer required.
260+
*/
261+
dispose() {
262+
263+
this._renderTarget.dispose();
264+
265+
if ( this._material !== null ) this._material.dispose();
266+
267+
}
268+
269+
}
270+
271+
export default SharpenNode;
272+
273+
/**
274+
* TSL function for creating a sharpen node for post processing.
275+
*
276+
* @tsl
277+
* @function
278+
* @param {Node<vec4>} node - The node that represents the input of the effect.
279+
* @param {(number|Node<float>)} [sharpness=0.2] - Sharpening strength. 0 = maximum, 2 = none.
280+
* @param {(boolean|Node<bool>)} [denoise=false] - Whether to attenuate sharpening in noisy areas.
281+
* @returns {SharpenNode}
282+
*/
283+
export const sharpen = ( node, sharpness, denoise ) => new SharpenNode( convertToTexture( node ), sharpness, denoise );

0 commit comments

Comments
 (0)