Publish verified LiteRT text-to-image pipeline and CPU examples

ac49327 verified 13 days ago

3.09 kB

package com.iris3b

import java.io.Closeable
import java.io.File
import java.io.RandomAccessFile
import java.nio.ByteOrder
import java.nio.ShortBuffer
import java.nio.channels.FileChannel

/**
 * The iris-3b prompt layout and the host-side embedding lookup.
 *
 * Layout (upstream `iris3b/text/qwen3_vl.py`, `_run`): the system-prompt prefix (34 ids), the
 * caption cut to 295 ids, the assistant-turn suffix (5 ids), right-padded with `<|endoftext|>` to
 * 334; the attention mask is 1 on prefix + caption + suffix. The suffix is placed after the cut, so
 * a long caption never loses the assistant marker.
 *
 * The embedding table is the checkpoint's bf16 `embed_tokens` as raw little-endian bytes,
 * memory-mapped; bf16 -> fp32 is a 16-bit shift, so the lookup is exact.
 */
class PromptTokens(assetsDir: File) : Closeable {

    private val tokenizer = QwenTokenizer(assetsDir)
    private val prefix = readIds(File(assetsDir, "prompt_prefix_ids.txt"))
    private val suffix = readIds(File(assetsDir, "prompt_suffix_ids.txt"))
    private val embeddingFile = RandomAccessFile(File(assetsDir, EMBEDDING_FILE), "r")
    private val embeddings: ShortBuffer = embeddingFile.channel
        .map(FileChannel.MapMode.READ_ONLY, 0, embeddingFile.length())
        .order(ByteOrder.LITTLE_ENDIAN)
        .asShortBuffer()

    /** Token ids and attention mask, both [SEQ_LEN], for one caption. */
    fun encode(caption: String): Tokens {
        val captionIds = tokenizer.encode(caption)
        val budget = WINDOW - suffix.size
        val n = minOf(captionIds.size, budget)
        val ids = IntArray(SEQ_LEN) { PAD_TOKEN_ID }
        prefix.copyInto(ids, 0)
        for (i in 0 until n) {
            ids[prefix.size + i] = captionIds[i]
        }
        suffix.copyInto(ids, prefix.size + n)
        val realCount = prefix.size + n + suffix.size
        return Tokens(ids, realCount, captionIds.size)
    }

    /** inputs_embeds [1, SEQ_LEN, HIDDEN] as fp32, gathered from the bf16 table. */
    fun embed(tokens: Tokens): FloatArray {
        val out = FloatArray(SEQ_LEN * HIDDEN)
        for (position in 0 until SEQ_LEN) {
            val base = tokens.ids[position] * HIDDEN
            for (channel in 0 until HIDDEN) {
                out[position * HIDDEN + channel] = Float.fromBits(embeddings.get(base + channel).toInt() shl 16)
            }
        }
        return out
    }

    override fun close() {
        embeddingFile.close()
    }

    /** Padded ids, how many are real (attention mask = 1), and the untruncated caption length. */
    class Tokens(val ids: IntArray, val realCount: Int, val captionTokens: Int) {
        fun mask(): IntArray = IntArray(SEQ_LEN) { if (it < realCount) 1 else 0 }
    }

    companion object {
        const val SEQ_LEN = 334
        const val WINDOW = 300
        const val HIDDEN = 2560
        const val PAD_TOKEN_ID = 151643
        const val EMBEDDING_FILE = "embed_bf16.bin"

        private fun readIds(file: File): IntArray =
            file.readText().trim().split(',').map { it.trim().toInt() }.toIntArray()
    }
}

	package com.iris3b

	import java.io.Closeable
	import java.io.File
	import java.io.RandomAccessFile
	import java.nio.ByteOrder
	import java.nio.ShortBuffer
	import java.nio.channels.FileChannel

	/**
	* The iris-3b prompt layout and the host-side embedding lookup.
	*
	* Layout (upstream `iris3b/text/qwen3_vl.py`, `_run`): the system-prompt prefix (34 ids), the
	* caption cut to 295 ids, the assistant-turn suffix (5 ids), right-padded with `<\|endoftext\|>` to
	* 334; the attention mask is 1 on prefix + caption + suffix. The suffix is placed after the cut, so
	* a long caption never loses the assistant marker.
	*
	* The embedding table is the checkpoint's bf16 `embed_tokens` as raw little-endian bytes,
	* memory-mapped; bf16 -> fp32 is a 16-bit shift, so the lookup is exact.
	*/
	class PromptTokens(assetsDir: File) : Closeable {

	private val tokenizer = QwenTokenizer(assetsDir)
	private val prefix = readIds(File(assetsDir, "prompt_prefix_ids.txt"))
	private val suffix = readIds(File(assetsDir, "prompt_suffix_ids.txt"))
	private val embeddingFile = RandomAccessFile(File(assetsDir, EMBEDDING_FILE), "r")
	private val embeddings: ShortBuffer = embeddingFile.channel
	.map(FileChannel.MapMode.READ_ONLY, 0, embeddingFile.length())
	.order(ByteOrder.LITTLE_ENDIAN)
	.asShortBuffer()

	/** Token ids and attention mask, both [SEQ_LEN], for one caption. */
	fun encode(caption: String): Tokens {
	val captionIds = tokenizer.encode(caption)
	val budget = WINDOW - suffix.size
	val n = minOf(captionIds.size, budget)
	val ids = IntArray(SEQ_LEN) { PAD_TOKEN_ID }
	prefix.copyInto(ids, 0)
	for (i in 0 until n) {
	ids[prefix.size + i] = captionIds[i]
	}
	suffix.copyInto(ids, prefix.size + n)
	val realCount = prefix.size + n + suffix.size
	return Tokens(ids, realCount, captionIds.size)
	}

	/** inputs_embeds [1, SEQ_LEN, HIDDEN] as fp32, gathered from the bf16 table. */
	fun embed(tokens: Tokens): FloatArray {
	val out = FloatArray(SEQ_LEN * HIDDEN)
	for (position in 0 until SEQ_LEN) {
	val base = tokens.ids[position] * HIDDEN
	for (channel in 0 until HIDDEN) {
	out[position * HIDDEN + channel] = Float.fromBits(embeddings.get(base + channel).toInt() shl 16)
	}
	}
	return out
	}

	override fun close() {
	embeddingFile.close()
	}

	/** Padded ids, how many are real (attention mask = 1), and the untruncated caption length. */
	class Tokens(val ids: IntArray, val realCount: Int, val captionTokens: Int) {
	fun mask(): IntArray = IntArray(SEQ_LEN) { if (it < realCount) 1 else 0 }
	}

	companion object {
	const val SEQ_LEN = 334
	const val WINDOW = 300
	const val HIDDEN = 2560
	const val PAD_TOKEN_ID = 151643
	const val EMBEDDING_FILE = "embed_bf16.bin"

	private fun readIds(file: File): IntArray =
	file.readText().trim().split(',').map { it.trim().toInt() }.toIntArray()
	}
	}