StreamDiffusionView.kt
src/jvmMain/kotlin/afield/media/video/diffusion/StreamDiffusionView.kt · 253 lines
package afield.media.video.diffusion
import afield.graphics.vulkan.Renderer
import afield.graphics.vulkan.Target
import afield.graphics.vulkan.device.CommandBuffer
import afield.graphics.vulkan.targets.RenderToTexture
import org.lwjgl.system.MemoryStack
import org.lwjgl.vulkan.*
import org.lwjgl.vulkan.VK13.*A render target that sends a source RenderToTexture through the streaming-diffusion worker
and returns the result as an ordinary sampled texture, StreamDiffusionView.output.
The frames stay on the GPU throughout: the source is copied into shared buffers, the CUDA worker
writes its result into others, and the view copies that back into output. Construct and close
a view between frames. With renderSource = true the view schedules the source itself; with
false the caller renders the source earlier in the same command buffer. An optional
StreamDiffusionMotionProvider enables display reprojection and history blending.
class StreamDiffusionView
GPU RTT -> optional CUDA worker -> ordinary sampled RTT. Construct/close between frames. With renderSource=true this target owns scheduling of source; do not also register source. With renderSource=false the caller renders source earlier in the same command buffer.
class StreamDiffusionView(
/** The renderer whose graphics queue hands the frames over, and whose cleanup frees them. */
val renderer: Renderer,
/** What is stylized: a single-layer, single-sample, transfer-source RGBA8 or BGRA8 target,
* with both dimensions multiples of 32 for the diffusion backend. */
val source: RenderToTexture,
/** Where the backend is and how it draws. */
val config: StreamDiffusionConfig,
/** The view's key in the renderer's pre-target list; two views need two names. */
val name: String = "stream-diffusion",
/** Whether this view draws [StreamDiffusionView.source] itself. Pass false, and render it earlier in the same
* command buffer, when something else already schedules it. */
val renderSource: Boolean = true,
/** The clock stamped on each captured frame and returned with its result, for a caller whose
* source does not advance in real time. */
val sourceTimeNanos: () -> Long = System::nanoTime,
/** Per-pixel correspondence from the renderer, which is what display reprojection, history
* blending and noise advection need; without one the worker's output is published as it is. */
val motion: StreamDiffusionMotionProvider? = null
) : Target, AutoCloseable {val output: RenderToTexture
The texture to sample: an ordinary sampled target holding the newest stylized frame.
val session: StreamDiffusionSession
The worker process behind this view, with its status, readiness and timings.
private val rawOutput: RenderToTexture
private val temporal: StreamDiffusionTemporalOutput?var displayReprojection
Warp the newest result onto the current frame before showing it, so the picture follows the
geometry while the worker is still behind. Needs StreamDiffusionView.motion.
@Volatile
var displayReprojection = config.displayReprojectionvar historyBlend
How much of the previous result to carry into a new one along the correspondence, 0 to 1.
Steadies the image between results; too much of it smears a fast-moving stroke. Needs StreamDiffusionView.motion.
@Volatile
var historyBlend = config.historyBlend
set(value) { require(value.isFinite() && value in 0.0..1.0); field = value }
private var previousInput: StreamDiffusionFrame? = nullval displayReprojectionRequests: Long get() = temporal?.displayRequests ?: 0L
Display correspondences asked of StreamDiffusionView.motion so far; zero without one.
val historyBlendRequests: Long get() = temporal?.historyRequests ?: 0L
History correspondences asked of StreamDiffusionView.motion so far; zero without one, or at zero blend.
var noiseMotionFrames = 0L; private set
Captured frames that also carried a motion texture for the worker's noise.
var capturedFrames = 0L; private set
Source frames sent to the worker, which is at most StreamDiffusionConfig.fps a second.
var skippedCaptures = 0L; private set
Captures given up because every input slot was still with the worker.
var presentedFrames = 0L; private set
Results copied into StreamDiffusionView.output.
var droppedOutputs = 0L; private set
Results thrown away: from a generation before the last reset, or behind a backlog.
var lastSourceFrameId = -1L; private set
Which source frame the newest presented result came from, or -1 before the first one.
var lastSourceTimeNanos = 0L; private set
The StreamDiffusionView.sourceTimeNanos reading of that frame, so a caller can measure the round trip.
private var frameId = 0L
private var captureDue = 0L
private var presentDue = 0L
private val interval = 1_000_000_000L / config.fps
private var initialized = false
private var hasOutput = false
@Volatile private var closed = falseinit
init {
check(renderer.currentCommandBuffer == null && !renderer.isClosed)
val support = StreamDiffusionSupport.probe(renderer, config)
check(support.available) { support.reason }
if (config.backend == "diffusion") require(source.width % 32 == 0 && source.height % 32 == 0) {
"Diffusion input dimensions must be multiples of 32"
}
require(source.colorLayerCount == 1 && source.colorImageRefs[0]?.imageData?.sampleCount == VK_SAMPLE_COUNT_1_BIT)
require(source.colorFormats[0] in setOf(VK_FORMAT_R8G8B8A8_UNORM, VK_FORMAT_R8G8B8A8_SRGB,
VK_FORMAT_B8G8R8A8_UNORM, VK_FORMAT_B8G8R8A8_SRGB)) { "Diffusion input must be RGBA8/BGRA8" }
require(source.parameters.colorTransferSrc)
if (motion != null) require(source.parameters.colorSampled && source.defaultSampler != null) {
"Diffusion reprojection needs a sampled source RTT with a default sampler"
}
check(!renderer.renderList.preTargets.containsKey(name))
rawOutput = RenderToTexture(renderer) {
width = source.width; height = source.height; format = source.colorFormats[0]
useDepth = false; colorSampled = true; colorTransferDst = true
}
rawOutput.setDefaultSampler(anisotropyEnable = false)
temporal = try { motion?.let { StreamDiffusionTemporalOutput(renderer, source, rawOutput, it) } }
catch (e: Throwable) { cleanupRawOutput(); throw e }
output = temporal?.output ?: rawOutput
try { session = StreamDiffusionSession(renderer, source.width, source.height, source.colorFormats[0], config, motion != null) }
catch (e: Throwable) { cleanupOutput(); throw e }
renderer.renderList.preTargets[name] = this
}override fun render(command: CommandBuffer, index: Int)
One graphics frame of the exchange, called by the renderer rather than by box code: present a finished result if one is due, then capture the source again if the interval has passed and the worker has a free slot. Until the first result arrives, StreamDiffusionView.output carries a copy of the source, so there is a picture to draw from the moment the view is made.
body — 101 lines, from line 123
override fun render(command: CommandBuffer, index: Int) {
if (closed) return
if (renderSource) source.render(command, index)
val current = StreamDiffusionFrame(frameId++, sourceTimeNanos(), session.generation)
motion?.onFrame(current)
var arrived: StreamDiffusionFrame? = null
// Show an actual source frame while the backend loads, using a GPU image copy.
if (!hasOutput) {
source.recordColorTransferSource(command) { from ->
rawOutput.recordColorTransferDestination(command) { to ->
MemoryStack.stackPush().use { s ->
val r = VkImageCopy.calloc(1, s)
.srcSubresource { it.aspectMask(VK_IMAGE_ASPECT_COLOR_BIT).mipLevel(0).baseArrayLayer(0).layerCount(1) }
.dstSubresource { it.aspectMask(VK_IMAGE_ASPECT_COLOR_BIT).mipLevel(0).baseArrayLayer(0).layerCount(1) }
.extent { it.width(source.width).height(source.height).depth(1) }
vkCmdCopyImage(command.vkCommandBuffer, from, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
to, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, r)
}
}
}
}
if (!session.isClosed && !initialized) {
session.outputs.forEach { it.barrier(command, acquire = false, access = 0) }
initialized = true
}
val now = System.nanoTime()
while (true) {
if (session.isClosed) break
val result = session.results.peek() ?: break
val stale = result.generation != session.generation || session.results.size > 4
if (!stale && now < presentDue) break
session.results.poll()
val slot = session.outputs[result.slot]
renderer.externalGraphicsSync(command, slot.semaphore, result.value, wait = true)
slot.barrier(command, acquire = true, access = VK_ACCESS_2_TRANSFER_READ_BIT)
if (!stale) {
rawOutput.recordColorTransferDestination(command) { image ->
MemoryStack.stackPush().use { s ->
vkCmdCopyBufferToImage(command.vkCommandBuffer, slot.buffer, image,
VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, slot.copyRegion(s, source.width, source.height))
}
}
hasOutput = true; presentedFrames++
lastSourceFrameId = result.frameId; lastSourceTimeNanos = result.timeNanos
arrived = StreamDiffusionFrame(result.frameId, result.timeNanos, result.generation)
presentDue = if (presentDue == 0L) now + interval else presentDue + ((now - presentDue) / interval + 1) * interval
} else {
droppedOutputs++
}
slot.barrier(command, acquire = false, access = VK_ACCESS_2_TRANSFER_READ_BIT)
renderer.externalGraphicsSync(command, slot.semaphore, result.value + 1, wait = false)
renderer.afterCurrentGraphicsFrameComplete { session.outputFree(result.slot, result.value + 1) }
if (!stale) break
}
temporal?.let {
it.displayReprojection = displayReprojection; it.blend = historyBlend
it.resolve(command, index, current, arrived)
}
if (session.isClosed || !session.ready || now < captureDue) return
captureDue = if (captureDue == 0L) now + interval else captureDue + ((now - captureDue) / interval + 1) * interval
val id = current.id
val inputIndex = session.reserveInput()
if (inputIndex < 0) { skippedCaptures++; return }
val slot = session.inputs[inputIndex]
if (slot.used) {
renderer.externalGraphicsSync(command, slot.semaphore, slot.value + 1, wait = true)
slot.barrier(command, acquire = true, access = VK_ACCESS_2_TRANSFER_WRITE_BIT)
}
slot.value += if (slot.used) 2 else 1
slot.used = true
source.recordColorTransferSource(command) { image ->
MemoryStack.stackPush().use { s ->
vkCmdCopyImageToBuffer(command.vkCommandBuffer, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
slot.buffer, slot.copyRegion(s, source.width, source.height))
}
}
val previous = previousInput?.takeIf { it.generation == current.generation }
val mapping = if (motion != null && previous != null)
StreamDiffusionTemporalOutput.checkedMotion(motion.record(command,
StreamDiffusionMotionRequest(current, previous, StreamDiffusionMotionUse.NOISE)), source) else null
if (mapping != null) {
mapping.recordColorTransferSource(command) { image ->
MemoryStack.stackPush().use { stack ->
val region = slot.copyRegion(stack, source.width, source.height)
.bufferOffset(session.motionOffset)
vkCmdCopyImageToBuffer(command.vkCommandBuffer, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
slot.buffer, region)
}
}
noiseMotionFrames++
}
previousInput = current
val motionReference = if (mapping != null) previous!!.id else -1L
slot.barrier(command, acquire = false, access = VK_ACCESS_2_TRANSFER_WRITE_BIT)
renderer.externalGraphicsSync(command, slot.semaphore, slot.value, wait = false)
val value = slot.value
val generation = current.generation
val time = current.timeNanos
renderer.afterCurrentGraphicsFrameComplete { session.frame(inputIndex, value, id, time, generation, motionReference) }
capturedFrames++
}fun reset(prompt: String = config.prompt, strength: Double = config.strength, seed: Int = config.seed, noiseAdvection: Double = config.noiseAdvection)
Change artistic controls; preserves the last valid image during causal warm-up.
fun reset(prompt: String = config.prompt, strength: Double = config.strength, seed: Int = config.seed, noiseAdvection: Double = config.noiseAdvection) =
session.reset(prompt, strength, seed, noiseAdvection)
private fun cleanupOutput() {
temporal?.close()
cleanupRawOutput()
}
private fun cleanupRawOutput() {
renderer.cleanup += Runnable {
rawOutput.imageViewArray.forEach { it.imageView.cleanup() }
rawOutput.colorImageRefs.forEach { it?.cleanup() }
rawOutput.defaultSampler?.cleanup()
}
}override fun close()
Unregisters the view, stops the worker and releases its targets. Call it between graphics frames: there is a command buffer open during one, and closing then is refused.
override fun close() {
check(renderer.currentCommandBuffer == null) { "Close diffusion between graphics frames" }
if (closed) return
closed = true
renderer.renderList.preTargets.remove(name, this)
session.close()
cleanupOutput()
}
}