rayzee 11.4.0 → 11.6.0

package/src/Stages/LocalExposure.js

@@ -495,7 +495,8 @@ export class LocalExposure extends RenderStage {
 
 	/**
 	 * The gain for a tone-mapping pass outside three.js (`PackedToneMapper`'s `gain`), or null before the
-	 * first build. `cpu()` reads the grid back for the CPU tone map.
+	 * first build. `cpu()` reads the grid back for the CPU tone map; `copyTo( device )` copies it to another
+	 * GPU device (the neural models' own).
 	 */
 	toneGain() {
 
@@ -508,10 +509,12 @@ export class LocalExposure extends RenderStage {
 		this._paramsBuffer ??= backend.device.createBuffer( { label: 'local-exposure-params', size: 64, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST } );
 		backend.device.queue.writeBuffer( this._paramsBuffer, 0, this._params );
 		const params = this._params.slice();
+		const read = () => Promise.all( [ this.renderer.getArrayBufferAsync( this._grid.value ), this.renderer.getArrayBufferAsync( this._blur.value ) ] );
 
 		return {
 			key: 'local-exposure',
 			wgsl: LOCAL_EXPOSURE_WGSL,
+			device: backend.device,
 			entries: [
 				{ binding: 0, resource: { buffer: grid } },
 				{ binding: 1, resource: { buffer: blur } },
@@ -519,11 +522,31 @@ export class LocalExposure extends RenderStage {
 			],
 			cpu: async () => {
 
-				const [ g, b ] = await Promise.all( [ this.renderer.getArrayBufferAsync( this._grid.value ), this.renderer.getArrayBufferAsync( this._blur.value ) ] );
+				const [ g, b ] = await read();
 				const gridData = new Float32Array( g ), blurData = new Float32Array( b );
 				return ( u, v, lum, logExposure ) => localExposureGain( params, gridData, blurData, u, v, lum, logExposure );
 
 			},
+			copyTo: async device => {
+
+				const [ g, b ] = await read();
+				const buffers = [[ g, GPUBufferUsage.STORAGE ], [ b, GPUBufferUsage.STORAGE ], [ params, GPUBufferUsage.UNIFORM ]].map( ( [ data, usage ] ) => {
+
+					const buffer = device.createBuffer( { label: 'local-exposure-copy', size: data.byteLength, usage: usage | GPUBufferUsage.COPY_DST } );
+					device.queue.writeBuffer( buffer, 0, data );
+					return buffer;
+
+				} );
+
+				return {
+					key: 'local-exposure',
+					wgsl: LOCAL_EXPOSURE_WGSL,
+					device,
+					entries: buffers.map( ( buffer, binding ) => ( { binding, resource: { buffer } } ) ),
+					destroy: () => buffers.forEach( buffer => buffer.destroy() ),
+				};
+
+			},
 		};
 
 	}

package/src/Stages/PathTracer.js

@@ -28,6 +28,7 @@ import {
 } from '../TSL/SortGlobalKernels.js';
 import { ENGINE_DEFAULTS } from '../EngineDefaults.js';
 import { MAX_STORAGE_TEXTURE_SIZE } from '../Processor/StorageTexturePool.js';
+import { readPixels } from '../Processor/TextureReadback.js';
 import { ISSUE_CODES } from '../EngineIssues.js';
 import { createLogger, fmt } from '../utils/Logger.js';
 import { deviceMemoryGB } from '../Processor/HostMemory.js';
@@ -758,16 +759,7 @@ export class PathTracer extends PathTracerStage {
 		const height = pool.renderHeight;
 		const pixels = width * height;
 
-		const readAttachment = async index => {
-
-			const padded = await this.renderer.readRenderTargetPixelsAsync( pool.readTarget, 0, 0, width, height, index );
-			const stride = Math.ceil( width * 16 / 256 ) * 64;
-			if ( stride === width * 4 ) return new Float32Array( padded.buffer, padded.byteOffset, pixels * 4 ).slice();
-			const tight = new Float32Array( pixels * 4 );
-			for ( let y = 0; y < height; y ++ ) tight.set( padded.subarray( y * stride, y * stride + width * 4 ), y * width * 4 );
-			return tight;
-
-		};
+		const readAttachment = index => readPixels( this.renderer, pool.readTarget, width, height, index );
 
 		const readBuffer = async ( attr, Type ) => attr ? new Type( await this.renderer.getArrayBufferAsync( attr, null, 0, pixels * 4 ) ) : null;
 		const aux = this._auxGBufferEnabled && this._auxSamples > 0;

package/src/managers/DenoisingManager.js

@@ -137,6 +137,11 @@ export class DenoisingManager extends EventDispatcher {
 		this._retouch = null;
 		this._superResModule = null;
 		this._neuralPostBusy = false;
+		// The finished render's post passes: its guard, the run in flight, a queued rerun, and a cancel.
+		this._postIsStillComplete = null;
+		this._postRun = null;
+		this._postRedo = null;
+		this._postCancelled = false;
 
 		// The tier the finished image uses. The loaded tier is not always this one: while the
 		// image is still accumulating we run a cheaper model (see previewQuality).
@@ -192,6 +197,7 @@ export class DenoisingManager extends EventDispatcher {
 		this._denoiserEndHandler = null;
 		this._denoiserTileHandler = null;
 		this._denoiserErrorHandler = null;
+		this._denoiserLoadHandler = null;
 		this._upscalerResChangedHandler = null;
 		this._upscalerStartHandler = null;
 		this._upscalerProgressHandler = null;
@@ -372,6 +378,12 @@ export class DenoisingManager extends EventDispatcher {
 		this.denoiser.addEventListener( 'tileProgress', this._denoiserTileHandler );
 		this.denoiser.addEventListener( 'error', this._denoiserErrorHandler );
 
+		// 'error' too: a failed load never sends 'loaded'. The first load starts inside the constructor, before these listen.
+		this._denoiserLoadHandler = e =>
+			this.dispatchEvent( { type: EngineEvents.DENOISER_LOADING, loading: e.type === 'loading' } );
+		for ( const type of [ 'loading', 'loaded', 'error' ] ) this.denoiser.addEventListener( type, this._denoiserLoadHandler );
+		if ( this.denoiser.state?.isLoading ) this.dispatchEvent( { type: EngineEvents.DENOISER_LOADING, loading: true } );
+
 	}
 
 	/**
@@ -1148,6 +1160,7 @@ export class DenoisingManager extends EventDispatcher {
 	/** Turns the neural-rendering (detail) pass on or off, and sets its appearance controls. */
 	setNeuralRendering( enabled, settings = null ) {
 
+		const changed = this.neuralRendering !== !! enabled;
 		this.neuralRendering = !! enabled;
 		if ( settings ) this.neuralRenderingSettings = { ...this.neuralRenderingSettings, ...settings };
 		if ( ! this.neuralRendering ) {
@@ -1157,6 +1170,8 @@ export class DenoisingManager extends EventDispatcher {
 
 		}
 
+		if ( changed ) this._redoPostPasses();
+
 	}
 
 	_releaseNeuralRetouch() {
@@ -1167,14 +1182,8 @@ export class DenoisingManager extends EventDispatcher {
 	}
 
 	/**
-	 * The neural post-chain over the picture the closing denoise just produced: super resolution
-	 * first, then the detail pass, either or both.
-	 *
-	 * Order is forced by the models. Super resolution hands back light that can be fed onward; the
-	 * detail pass writes an `rgba8unorm` texture with no `COPY_SRC`, so nothing can read its result
-	 * and it has to be last — and it presents in its own display space rather than the engine's tone
-	 * curve. Which shape super resolution returns therefore depends on what follows it: packed
-	 * halves when the detail pass is next, display bytes when it is the end of the chain.
+	 * The post-chain over the picture the closing denoise just produced: the detail pass at render size,
+	 * then the chosen upscaler (neural super resolution or the ONNX one), either or both.
 	 *
 	 * Deliberately fire-and-forget: the render is already finished and a failure must not break it.
 	 */
@@ -1182,6 +1191,7 @@ export class DenoisingManager extends EventDispatcher {
 
 		const wantSR = this.upscaler?.enabled && this.upscalerBackend === 'neural';
 		const wantNR = this.neuralRendering;
+		const wantESRGAN = this.upscaler?.enabled && ! wantSR;
 		if ( ! wantSR && ! wantNR ) return;
 
 		// A second completion can land while the first pass is still in flight — the models are
@@ -1230,7 +1240,12 @@ export class DenoisingManager extends EventDispatcher {
 
 			}
 
-			if ( ! wantSR && ! runNR ) return;
+			if ( ! wantSR && ! runNR ) {
+
+				if ( wantESRGAN && ! this._postCancelled ) await this.upscaler.start();
+				return;
+
+			}
 
 			// Must be the effective value, not the settings one: AutoExposure overwrites
 			// `renderer.toneMappingExposure` every frame and never touches settings, so reading
@@ -1291,6 +1306,15 @@ export class DenoisingManager extends EventDispatcher {
 
 			}
 
+			if ( this._postRedo || this._postCancelled ) return;
+
+			if ( wantESRGAN ) {
+
+				await this.upscaler.start( new ImageData( image.rgba8, image.width, image.height ) );
+				return;
+
+			}
+
 			sr.presentRGBA8( this.upscalerCanvas, image.rgba8, image.width, image.height );
 
 			if ( ! isStillComplete() ) {
@@ -1360,10 +1384,61 @@ export class DenoisingManager extends EventDispatcher {
 
 	}
 
+	/** The neural chain or the ONNX upscaler over the finished render, as the switches ask. */
+	_runPostPasses( isStillComplete ) {
+
+		this._postIsStillComplete = isStillComplete;
+		this._postCancelled = false;
+
+		// The neural chain owns the overlay canvas when it runs, and hands its result to the ONNX upscaler.
+		if ( this.neuralRendering || ( this.upscaler?.enabled && this.upscalerBackend === 'neural' ) ) {
+
+			this._postRun = this._runNeuralPost( isStillComplete );
+
+		} else {
+
+			this._postRun = this.upscaler?.enabled ? this.upscaler.start() : null;
+
+		}
+
+	}
+
+	/** Runs the post passes again on a finished render whose switches changed, once the run in flight ends. */
+	_redoPostPasses() {
+
+		const isStillComplete = this._postIsStillComplete;
+		// No chain yet: the closing denoise is still running and the chain reads the switches when it starts.
+		if ( ! isStillComplete?.() || this._postRedo ) return;
+
+		this.upscaler?.abort();
+		this._postCancelled = false;
+		this._postRedo = Promise.resolve( this._postRun ).catch( () => {} ).then( () => {
+
+			this._postRedo = null;
+			if ( this._postCancelled || this._postIsStillComplete !== isStillComplete || ! isStillComplete() ) return;
+			this._restoreRenderDisplay();
+			this._runPostPasses( isStillComplete );
+
+		} );
+
+	}
+
+	/** Stops the post passes on the finished render, which stays on screen as rendered. The switches stay as they are. */
+	cancelPostPasses() {
+
+		this._postCancelled = true;
+		this.upscaler?.abort();
+		this._restoreRenderDisplay();
+		// A retouch in flight cannot be stopped, only discarded, and ends later.
+		this.dispatchEvent( { type: EngineEvents.UPSCALING_END } );
+
+	}
+
 	onRenderComplete( { isStillComplete } ) {
 
 		// Remove any stale completion-chain listener from a previous render cycle
 		this._cleanupCompletionListener();
+		this._postIsStillComplete = null;
 
 		// What closes the render, if anything. A full pass at the chosen tier when that switch is
 		// on; otherwise, if OIDN owns the live view, one last cheap refresh — the cadence's last
@@ -1383,16 +1458,7 @@ export class DenoisingManager extends EventDispatcher {
 			// denoise. Ordered after `_denoiserEndHandler`, which publishes the picture.
 			if ( closing ) this._onDisplayRefresh?.();
 
-			// One owner of the overlay canvas, so the neural chain and the ONNX upscaler are
-			// exclusive. The neural chain also runs on its own when only its detail pass is on.
-			if ( this.neuralRendering || ( this.upscaler?.enabled && this.upscalerBackend === 'neural' ) ) {
-
-				this._runNeuralPost( isStillComplete );
-				return;
-
-			}
-
-			if ( this.upscaler?.enabled ) this.upscaler.start();
+			this._runPostPasses( isStillComplete );
 
 		};
 
@@ -1467,6 +1533,7 @@ export class DenoisingManager extends EventDispatcher {
 
 		// Remove stale completion-chain listener before aborting
 		this._cleanupCompletionListener();
+		this._postIsStillComplete = null;
 
 		// The 2D canvas is the upscaler's alone: a reset means its enlarged result no longer
 		// describes anything, so the render canvas comes back.
@@ -1509,6 +1576,7 @@ export class DenoisingManager extends EventDispatcher {
 
 		// Remove pending completion-chain listener
 		this._cleanupCompletionListener();
+		this._postIsStillComplete = null;
 		// Owns a second GPUDevice plus 3 MB of weights, so it must not outlive the manager.
 		this._releaseNeuralUpscaler();
 		this._releaseNeuralRetouch();
@@ -1523,6 +1591,13 @@ export class DenoisingManager extends EventDispatcher {
 			if ( this._denoiserEndHandler ) this.denoiser.removeEventListener( 'end', this._denoiserEndHandler );
 			if ( this._denoiserTileHandler ) this.denoiser.removeEventListener( 'tileProgress', this._denoiserTileHandler );
 			if ( this._denoiserErrorHandler ) this.denoiser.removeEventListener( 'error', this._denoiserErrorHandler );
+			if ( this._denoiserLoadHandler ) {
+
+				for ( const type of [ 'loading', 'loaded', 'error' ] ) this.denoiser.removeEventListener( type, this._denoiserLoadHandler );
+				if ( this.denoiser.state?.isLoading ) this.dispatchEvent( { type: EngineEvents.DENOISER_LOADING, loading: false } );
+
+			}
+
 			this.denoiser.dispose();
 			this.denoiser = null;
 
@@ -1547,6 +1622,7 @@ export class DenoisingManager extends EventDispatcher {
 		this._denoiserEndHandler = null;
 		this._denoiserTileHandler = null;
 		this._denoiserErrorHandler = null;
+		this._denoiserLoadHandler = null;
 		this._upscalerResChangedHandler = null;
 		this._upscalerStartHandler = null;
 		this._upscalerProgressHandler = null;
@@ -1737,10 +1813,12 @@ export class DenoisingManager extends EventDispatcher {
 	/** Enables or disables the AI upscaler. */
 	setUpscalerEnabled( enabled ) {
 
+		const changed = !! this.upscaler && this.upscaler.enabled !== !! enabled;
 		if ( this.upscaler ) this.upscaler.enabled = enabled;
 		else if ( enabled ) this._recordNoOverlayCanvas( 'upscaler' );
 		// The enlarged picture describes a setting that is no longer on.
 		if ( ! enabled ) this._restoreRenderDisplay();
+		if ( changed ) this._redoPostPasses();
 
 	}