feature/onnx/com.shifthackz.aisdv1.feature.onnx.ai.unet/UNet

UNet

internal class UNet(val deviceNNAPIFlagProvider: DeviceNNAPIFlagProvider, val ortEnvironmentProvider: OrtEnvironmentProvider, val fileProviderDescriptor: FileProviderDescriptor, val localModelIdProvider: LocalModelIdProvider, val preferenceManager: PreferenceManager)(source)

Coordinates UNet behavior in the SDAI ONNX local diffusion feature layer.

Author

Dmitriy Moroz

Constructors

UNet

constructor(deviceNNAPIFlagProvider: DeviceNNAPIFlagProvider, ortEnvironmentProvider: OrtEnvironmentProvider, fileProviderDescriptor: FileProviderDescriptor, localModelIdProvider: LocalModelIdProvider, preferenceManager: PreferenceManager)

Types

Callback

interface Callback

Defines the Callback contract for the SDAI ONNX local diffusion feature layer.

Properties

callback

private var callback: UNet.Callback?

Exposes the callback value used by the SDAI ONNX local diffusion feature layer.

decoder

private var decoder: VaeDecoder?

Exposes the decoder value used by the SDAI ONNX local diffusion feature layer.

deviceNNAPIFlagProvider

private val deviceNNAPIFlagProvider: DeviceNNAPIFlagProvider

Exposes the deviceNNAPIFlagProvider value used by the SDAI ONNX local diffusion feature layer.

fileProviderDescriptor

private val fileProviderDescriptor: FileProviderDescriptor

Exposes the fileProviderDescriptor value used by the SDAI ONNX local diffusion feature layer.

height

private var height: Int

Exposes the height value used by the SDAI ONNX local diffusion feature layer.

localModelIdProvider

private val localModelIdProvider: LocalModelIdProvider

Exposes the localModelIdProvider value used by the SDAI ONNX local diffusion feature layer.

ortEnvironmentProvider

private val ortEnvironmentProvider: OrtEnvironmentProvider

Exposes the ortEnvironmentProvider value used by the SDAI ONNX local diffusion feature layer.

preferenceManager

private val preferenceManager: PreferenceManager

Exposes the preferenceManager value used by the SDAI ONNX local diffusion feature layer.

random

private val random: Random

Exposes the random value used by the SDAI ONNX local diffusion feature layer.

session

private var session: OrtSession?

Exposes the session value used by the SDAI ONNX local diffusion feature layer.

width

private var width: Int

Exposes the width value used by the SDAI ONNX local diffusion feature layer.

Functions

fun close()

Executes the close step in the SDAI ONNX local diffusion feature layer.

createUNetModelInput

private fun createUNetModelInput(encoderHiddenStates: OnnxTensor, sample: OnnxTensor, timeStep: OnnxTensor): Map<String, OnnxTensor>

Creates the SDAI value produced by createUNetModelInput.

decode

fun decode(latents: LocalDiffusionTensor<*>): Bitmap

Executes the decode step in the SDAI ONNX local diffusion feature layer.

generateLatentSample

private fun generateLatentSample(batchSize: Int, height: Int, width: Int, seed: Long, initNoiseSigma: Float): LocalDiffusionTensor<*>

Executes the generateLatentSample step in the SDAI ONNX local diffusion feature layer.

inference

fun inference(seedNum: Long, numInferenceSteps: Int, textEmbeddings: OnnxTensor, guidanceScale: Double, batchSize: Int, width: Int, height: Int)

Executes the inference step in the SDAI ONNX local diffusion feature layer.

initialize

fun initialize()

Executes the initialize step in the SDAI ONNX local diffusion feature layer.

performGuidance

private fun performGuidance(noisePrediction: Array3D<FloatArray>, noisePredictionText: Array3D<FloatArray>, guidanceScale: Double)

Executes the performGuidance step in the SDAI ONNX local diffusion feature layer.

setCallback

fun setCallback(callback: UNet.Callback?)

Executes the setCallback step in the SDAI ONNX local diffusion feature layer.