OpenDlssNrStudy.kt
src/jvmMain/kotlin/afield/graphics/render/neural/OpenDlssNrStudy.kt · 207 lines
An editable source image fed through the native neural-rendering (NR) pass, with a graded comparison.
OpenDlssNrStudy renders a source, either a fragment shader, a loaded image or a small geometry scene, into a texture, runs it through OpenDlssNrPass when OpenDlssNrSupport reports the native library ready, and displays a graded comparison. Without the library it still shows the source. OpenDlssNrStudy.reload swaps shaders live and OpenDlssNrStudy.saveFrame captures a frame.
package afield.graphics.render.neural
import afield.graphics.vulkan.Renderer
import afield.graphics.vulkan.Target
import afield.graphics.vulkan.device.CommandBuffer
import afield.document.context.*
import afield.graphics.render.*
import afield.graphics.scene.*
import afield.graphics.shader.*
import afield.graphics.vulkan.targets.RenderToTexture
import afield.graphics.texture.LoadedTexture
import afield.graphics.texture.TextureSampler
import afield.graphics.texture.HasImageViewArray
import afield.graphics.texture.PreparesSampledImages
import afield.runtime.lifecycle.RunScope
import afield.math.linalg.Vector4f
import org.lwjgl.vulkan.VK13.*
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.nio.file.Path
import java.util.concurrent.CompletableFutureclass OpenDlssNrStudy
Editable source -> native NR -> graded comparison. Source mode works without the optional library. The current box run owns the study: stopping or rerunning it closes the old resources. The renderer also closes the study before destroying its device. The comparison and capture use the very same targets as the live Starting Point.
class OpenDlssNrStudy(
val renderer: Renderer,
val size: Int = 768,
val name: String = "nr-study",
imagePath: String? = null,
val geometry: Boolean = false,
vertex: String = if(geometry)resource("cubes.vert") else defaultVertex,
fragment: String = if(geometry)resource("cubes.frag") else defaultFragment,
attach: Boolean = true
) : Target, AutoCloseable {init
init { require(!geometry || imagePath==null) { "The geometry study needs shaders rather than an input image" } }val availability = OpenDlssNrSupport.availability(renderer)
private val image = imagePath?.let { LoadedTexture.fromFile(it, renderer, VK_FORMAT_R8G8B8A8_UNORM) }val source
val source = RenderToTexture(renderer) {
width=size; height=size; useDepth=geometry; format=VK_FORMAT_R16G16B16A16_SFLOAT
colorTransferSrc=true; debugName="$name-source"
}
private val imageProvider=image?.let { texture -> object : HasImageViewArray,PreparesSampledImages {
override val imageViewArray=listOf(texture)
override fun prepareSampledImagesFor(command: CommandBuffer)=texture.prepareSampledImagesFor(command)
} }
private val sampler=TextureSampler(renderer,1,false,VK_FILTER_LINEAR)
private fun buildSourceModule(vertex: String,fragment: String) = Module(parameters(renderer) {
this.name="$name-source"; +source; depthTest=geometry; depthWrite=geometry; blend=false
textureResolver={ if(it.name=="image" && imageProvider!=null)imageProvider to sampler else null }
this.vertex=vertex
this.fragment=if(image==null)fragment else """
#version 460
layout(location=0) in vec2 uv;
layout(location=0) out vec4 color;
layout(binding=0) uniform sampler2D image;
layout(push_constant) uniform Study { vec4 frame; } pc;
void main() {
vec3 c=texture(image,uv).rgb;
c=mix(c/12.92,pow((c+.055)/1.055,vec3(2.4)),step(vec3(.04045),c));
color=vec4(c+pc.frame.w,1);
}
""".trimIndent()
})
private var sourceModule = buildSourceModule(vertex,fragment)
private var sourceMesh = triangle(sourceModule, if(geometry)3+39*36 else 3)A rebuilt source (new stage interface) waiting for the next frame boundary.
@Volatile
private var pendingSource: Pair<Module,Mesh>?=nullval pass = if(availability.ready)OpenDlssNrPass(renderer,size,size) else null
val output
val output = RenderToTexture(renderer) {
width=size*3; height=size+64; useDepth=false; format=VK_FORMAT_R16G16B16A16_SFLOAT
colorTransferSrc=true; debugName="$name-comparison"
}
private val compareModule = Module(parameters(renderer) {
this.name="$name-comparison"; +output; depthTest=false; depthWrite=false; blend=false
this.vertex=defaultVertex; this.fragment=resource("compare.frag")
textureResolver={ when(it.name) { "source" -> source to sampler; "neural" -> (pass?.output ?: source) to sampler; else -> null } }
})
private val compareMesh = triangle(compareModule)
private var screen: ModulePointer?=null
private var frame=0L
private var closed=false
private val ownerScope=RunScope.current()Closes this study when its renderer tears its device down, so the boxes using it need not stop first.
private val closeWithRenderer=Runnable { close() }val isClosed: Boolean get()=closed
Whether close has run, by the box or by the renderer going away.
var time=0f
var animate=true
var recording=false
The editable capture box reserves scene time while recording.
var controls=OpenDlssNrControls(style=1,tone=0f,structure=10f)
var grade=Vector4f(.65f,1.1f,1.4f,1f)
init
init {
check(sourceModule.valid && compareModule.valid) { "NR study shader compilation failed" }
source.modules["source"]=sourceModule.pointer()
output.modules["compare"]=compareModule.pointer()
renderer.beforeDeviceCleanup+=closeWithRenderer
if(attach)attach()
ownerScope?.own(this)
}
private fun triangle(module: Module, vertices: Int=3): Mesh = Mesh(renderer,module.pointer()).also {
it.resize(vertices,0,module.pointer()); it.push.setMembers(module.fi.pushConstants.single().members)
module.meshes["triangle"]=it
}fun attach(): OpenDlssNrStudy
fun attach(): OpenDlssNrStudy {
check(!closed)
renderer.renderList.preTargets[name]=this
screen=output.blitToScreen(name=name)
return this
}fun reload(vertex: String=if(geometry)resource("cubes.vert") else defaultVertex,
A shader edit replaces the source at a frame boundary. An edit that keeps the stage interface swaps the pipeline in place; one that changes it (adding or removing a varying, a push-constant member, a binding) rebuilds the source module and mesh. Invalid edits retain the current source.
fun reload(vertex: String=if(geometry)resource("cubes.vert") else defaultVertex,
fragment: String=if(geometry)resource("cubes.frag") else defaultFragment): Result<Unit> = runCatching {
check(!closed && image==null) { "Shader reload requires the procedural source" }
val errors=mutableListOf<String>()
val v=Introspection(vertex,ShaderKind.VERTEX_SHADER,errorHandler={message,_ -> errors+=message})
val f=Introspection(fragment,ShaderKind.FRAGMENT_SHADER,errorHandler={message,_ -> errors+=message})
check(v.valid && f.valid) {errors.joinToString("\n")}
val current=pendingSource?.first ?: sourceModule
if(graphicsShaderAbiCheck(current.vi,current.fi,v,f).compatible) {
current.installPipelineUpdate(current.buildPipelineUpdate(v,f))
} else {
val module=buildSourceModule(vertex,fragment)
if(!module.valid) { renderer.cleanup+=Runnable { module.cleanup() }; error("NR study shader compilation failed") }
val mesh=triangle(module,if(geometry)3+39*36 else 3)
pendingSource?.let { (old,_) -> renderer.cleanup+=Runnable { old.cleanup() } }
pendingSource=module to mesh
}
pass?.reset("source shader edit")
}val reloadPending: Boolean get()=pendingSource!=null
Whether the last interface-changing reload has reached the GPU.
fun saveFrame(path: Path): CompletableFuture<Path>
Queue the next comparison frame for a lossless PNG after its GPU fence completes.
fun saveFrame(path: Path): CompletableFuture<Path> {
check(!closed)
return output.readPixelsAsync().thenApplyAsync { bytes ->
val buffer=ByteBuffer.wrap(bytes).order(ByteOrder.nativeOrder())
val pixels=FloatArray(output.width*output.height*4) { java.lang.Float.float16ToFloat(buffer.getShort(it*2)) }
OpenDlssNrDiagnostics.png(pixels,output.width,output.height,path)
path
}
}override fun render(command: CommandBuffer,index: Int)
override fun render(command: CommandBuffer,index: Int) {
if(closed)return
pendingSource?.let { (module,mesh) ->
pendingSource=null
val old=sourceModule
sourceModule=module; sourceMesh=mesh
source.modules["source"]=module.pointer()
renderer.cleanup+=Runnable { old.cleanup() }
}
if(geometry) {
sourceMesh.push["sizeJitter"]=Vector4f(size.toFloat(),size.toFloat(),0f,0f)
sourceMesh.push["times"]=Vector4f(time,time,0f,0f)
sourceMesh.push["settings"]=Vector4f(0f,1f,0f,0f)
} else sourceMesh.push["frame"]=Vector4f(time,size.toFloat(),size.toFloat(),0f)
source.render(command,index)
pass?.let {
it.controls=controls
it.evaluate(command,OpenDlssNrInputs(source,OpenDlssNrColorEncoding.LINEAR_LDR),OpenDlssNrFrame(frame))
}
compareMesh.push["grade"]=grade
compareMesh.push["dimensions"]=Vector4f(output.width.toFloat(),output.height.toFloat(),64f,if(pass==null)0f else 1f)
output.render(command,index)
if(animate)time+=1f/30f
frame++
}override fun close()
override fun close() {
if(closed)return
closed=true
ownerScope?.disown(this)
renderer.beforeDeviceCleanup.remove(closeWithRenderer)
renderer.renderList.preTargets.remove(name,this)
screen?.let { renderer.renderList.modules.remove(name,it) }
pass?.close()
renderer.cleanup+=Runnable {
screen?.current?.cleanup(); sourceModule.cleanup(); pendingSource?.first?.cleanup(); compareModule.cleanup()
for(target in listOf(source,output)) {
target.imageViewArray.forEach {it.imageView.cleanup()}
target.colorImageRefs.forEach {it?.cleanup()}
target.depthImageView?.cleanup(); target.depthImage?.cleanup()
target.defaultSampler?.cleanup()
}
image?.cleanup(); sampler.cleanup()
}
}companion object
companion object {
private fun resource(file: String)=OpenDlssNrStudy::class.java.getResource("/afield/four/neuralrendering/$file")!!.readText()