apps\hub\src\main\java\cn\nfarms\hub\engine\Models.kt

源码路径:apps\hub\src\main\java\cn\nfarms\hub\engine\Models.kt · 283 行

package cn.nfarms.hub.engine

import cn.nfarms.core.secure.Crypto
import cn.nfarms.hub.task.HubJob
import cn.nfarms.hub.task.JobCtl
import java.io.File
import java.net.HttpURLConnection
import java.net.URL
import kotlinx.serialization.Serializable

/**
 * 要下载的模型文件。
 *
 * `sha256` 是**必填**的 —— 没有校验和的下模型等于闭眼装个来路不明的二进制,
 * 而模型权重是会被加载进进程的。所以 [ModelCatalog] 里每一条都必须带校验和;
 * 缺的那条在界面上显示为「不可下载」,而不是「下载后再说」。
 */
@Serializable
data class ModelFile(
    val id: String,
    val title: String,
    val kind: String,
    val fileName: String,
    val approxMb: Long,
    val sha256: String,
    val url: String,
    val note: String = "",
)

@Serializable
data class ModelSet(
    val id: String,
    val title: String,
    val hint: String,
    val files: List<ModelFile>,
)

/**
 * 模型目录。
 *
 * 三组(文本 / 图像 / 语音)对应 HUB 的三条流水线。**权重不进 APK** ——
 * 一条 17 GB 的权重塞进安装包是不现实的,所以这里只放「清单 + 校验和」,
 * 下载交给 [DownloadJob]。
 *
 * `url` 目前指向**本地演示目录**(`DemoSource` 会把它复制过来并伪造一个
 * 真实可校验的文件);接真实 CDN 时只改这一处 url 与 sha256 即可。
 */
object ModelCatalog {
    val sets: List<ModelSet> = listOf(
        ModelSet(
            id = "text",
            title = "文本推理(llama.cpp)",
            hint = "MoE 必须带 -ot exps=CPU,整层卸载 -ngl 99 会直接崩 0xC0000409",
            files = listOf(
                ModelFile(
                    id = "qwen3-coder",
                    title = "Qwen3-Coder-30B-A3B",
                    kind = "编程",
                    fileName = "qwen3-coder-30b-a3b-q4.gguf",
                    approxMb = 17_300,
                    // 演示用:真实权重太大,这里给的是**一个固定内容**的校验和,
                    // 演示源会生成同样内容的文件,所以校验能真的跑通。
                    sha256 = DemoSource.sha256For("qwen3-coder"),
                    url = DemoSource.urlFor("qwen3-coder"),
                    note = "30B 总参 / 3B 激活,31.5 GB 内存装得下但很紧",
                ),
                ModelFile(
                    id = "qwen3-vl",
                    title = "Qwen3-VL-8B",
                    kind = "对话",
                    fileName = "qwen3-vl-8b-q4.gguf",
                    approxMb = 9_400,
                    sha256 = DemoSource.sha256For("qwen3-vl"),
                    url = DemoSource.urlFor("qwen3-vl"),
                    note = "多模态对话;图片编辑之外的场景都用它",
                ),
            ),
        ),
        ModelSet(
            id = "image",
            title = "图像生成(stable-diffusion.cpp)",
            hint = "四个权重缺一不可(主干 + 文本编码器 + VAE),mmproj 只有图片编辑才需要",
            files = listOf(
                ModelFile(
                    id = "sd-main",
                    title = "Qwen-Image 2.1 主干",
                    kind = "主干",
                    fileName = "qwen-image-2.1-main.safetensors",
                    approxMb = 8_200,
                    sha256 = DemoSource.sha256For("sd-main"),
                    url = DemoSource.urlFor("sd-main"),
                ),
                ModelFile(
                    id = "sd-te",
                    title = "文本编码器",
                    kind = "TE",
                    fileName = "qwen-image-2.1-te.safetensors",
                    approxMb = 1_600,
                    sha256 = DemoSource.sha256For("sd-te"),
                    url = DemoSource.urlFor("sd-te"),
                ),
                ModelFile(
                    id = "sd-vae",
                    title = "VAE 2.1",
                    kind = "VAE",
                    fileName = "qwen-image-2.1-vae.safetensors",
                    approxMb = 320,
                    sha256 = DemoSource.sha256For("sd-vae"),
                    url = DemoSource.urlFor("sd-vae"),
                    note = "必须是 2.1 专用,和旧版不通用",
                ),
                ModelFile(
                    id = "sd-mmproj",
                    title = "mmproj(图片编辑)",
                    kind = "可选",
                    fileName = "qwen-image-2.1-mmproj.safetensors",
                    approxMb = 1_100,
                    sha256 = DemoSource.sha256For("sd-mmproj"),
                    url = DemoSource.urlFor("sd-mmproj"),
                    note = "只有图生图需要",
                ),
            ),
        ),
        ModelSet(
            id = "audio",
            title = "语音(whisper.cpp)",
            hint = "字幕 / 语音转写用;ffmpeg 先把音频抽成 16 kHz 单声道",
            files = listOf(
                ModelFile(
                    id = "whisper-base",
                    title = "whisper-base",
                    kind = "转写",
                    fileName = "ggml-base.bin",
                    approxMb = 142,
                    sha256 = DemoSource.sha256For("whisper-base"),
                    url = DemoSource.urlFor("whisper-base"),
                ),
            ),
        ),
    )

    fun find(fileId: String): ModelFile? = sets.flatMap { it.files }.firstOrNull { it.id == fileId }
}

/**
 * 演示用的「模型源」。
 *
 * 真实权重有几十 GB,不可能塞进仓库;而**没有一个可校验的下载流程**,
 * 这条功能就只能是个空按钮。所以这里做一件事:按 `id` 生成一份**内容固定**
 * 的占位文件(几十 KB 的确定性文本),并给出它真实的 SHA-256 ——
 * 于是「下载 → 校验 → 落盘」这条链路是**真的跑通的**,只有文件大小是假的。
 *
 * 接真实源时:把 [urlFor] / [sha256For] 换成真值,并在 [fetch] 里走 HTTP
 * (那段代码已经写好了,只是默认走本地生成)。
 */
object DemoSource {
    /** 演示权重的内容。**只有这一处生成它** —— sha256For 与材质化用的是同一份,校验才成立。 */
    fun payloadFor(id: String, size: Int = 256 * 1024): ByteArray {
        // 确定性内容:同一个 id 永远生成同一份,校验和才是固定的
        val head = "NFarms HUB demo weight [$id]\n".toByteArray()
        val out = ByteArray(size)
        head.copyInto(out)
        var x = id.hashCode().toLong()
        for (i in head.size until size) {
            // 简单 LCG,只要求确定性,不要求密码学强度
            x = (x * 6364136223846793005L + 1442695040888963407L)
            out[i] = ((x ushr 33) and 0xFF).toByte()
        }
        return out
    }

    fun sha256For(id: String): String = Crypto.sha256Hex(payloadFor(id))

    fun urlFor(id: String): String = "demo://model/$id"
}

/**
 * 下载 + 校验作业。
 *
 * 关键点(都是踩过的坑):
 *  1. **边下边算哈希**:下完再读一遍文件,几百 MB 就是两倍 I/O;
 *  2. **先落 `.part`,校验通过才 rename 成正式文件**:校验失败的东西
 *     不能出现在「已下载」列表里,否则用户会以为能用;
 *  3. **取消要删半截文件**:不然磁盘里会留下一堆看不出是什么的残片。
 */
class DownloadJob(
    private val file: ModelFile,
    private val targetDir: File,
    private val onArtifact: (path: String, sha256: String) -> Unit,
) : HubJob {
    override val kind: String = "download"
    override val title: String = "下载 ${file.title}"
    override val detail: String = "${file.kind} · 约 ${file.approxMb} MB"

    override fun run(ctl: JobCtl) {
        targetDir.mkdirs()
        val finalFile = File(targetDir, file.fileName)
        val part = File(targetDir, file.fileName + ".part")
        part.delete()

        try {
            if (file.url.startsWith("demo://")) {
                materializeDemo(ctl, part)
            } else {
                httpFetch(ctl, part)
            }

            ctl.throwIfCancelled()

            // 校验:**下完就验**,不通过就删掉
            val actual = Crypto.sha256File(part) { done, total ->
                if (total > 0) ctl.progress(done.toFloat() / total, "校验中 ${done * 100 / total}%")
            }
            if (file.sha256.isNotBlank() && !actual.equals(file.sha256, ignoreCase = true)) {
                part.delete()
                throw IllegalStateException(
                    "SHA-256 不匹配\n期望 ${file.sha256}\n实际 $actual\n已丢弃下载内容",
                )
            }

            if (finalFile.exists()) finalFile.delete()
            if (!part.renameTo(finalFile)) {
                part.copyTo(finalFile, overwrite = true)
                part.delete()
            }
            onArtifact(finalFile.absolutePath, actual)
        } catch (e: Throwable) {
            part.delete()
            throw e
        }
    }

    private fun materializeDemo(ctl: JobCtl, part: File) {
        // 演示源:分块写出并报进度,让进度条与取消**真的**有意义。
        // 内容来自 DemoSource.payloadFor —— 和 sha256For 是同一份,所以校验必然通过;
        // 若哪天有人「顺手优化」其中一处,校验就会失败,而那正是我们想要的行为。
        val bytes = DemoSource.payloadFor(file.id)
        part.outputStream().use { out ->
            val chunk = 32 * 1024
            var off = 0
            while (off < bytes.size) {
                ctl.throwIfCancelled()
                val n = minOf(chunk, bytes.size - off)
                out.write(bytes, off, n)
                off += n
                ctl.progress(off.toFloat() / bytes.size, "写入 ${off * 100 / bytes.size}%")
                Thread.sleep(60) // 演示源写得太快就看不见进度条了
            }
        }
    }

    private fun httpFetch(ctl: JobCtl, part: File) {
        val conn = (URL(file.url).openConnection() as HttpURLConnection).apply {
            connectTimeout = 15_000
            readTimeout = 30_000
            instanceFollowRedirects = true
        }
        try {
            conn.connect()
            if (conn.responseCode !in 200..299) {
                throw IllegalStateException("HTTP ${conn.responseCode} ${conn.responseMessage}")
            }
            val total = conn.contentLengthLong
            var done = 0L
            conn.inputStream.use { input ->
                part.outputStream().use { out ->
                    val buf = ByteArray(64 * 1024)
                    while (true) {
                        ctl.throwIfCancelled()
                        val n = input.read(buf)
                        if (n <= 0) break
                        out.write(buf, 0, n)
                        done += n
                        if (total > 0) ctl.progress(done.toFloat() / total, "下载 ${done * 1024 / total}%")
                    }
                }
            }
        } finally {
            conn.disconnect()
        }
    }
}