Upgrade to LiteRT-LM runtime for Gemma 3n E4B support

- Add com.google.ai.edge.litertlm:litertlm-android:0.9.0-alpha05 dependency - New LiteRTModel backend: loads .litertlm files via Engine API with GPU backend (CPU fallback if GPU init fails) - OnDeviceModel.create() priority: LiteRT-LM → MediaPipe → Gemini Nano - ModelDownloader: updated to Gemma 3n E4B (4.9 GB) and E4B Web (4.3 GB) from google/gemma-3n-E4B-it-litert-lm (HuggingFace, license required) - Download uses Bearer token auth; token saved to SharedPreferences Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
Fix downloads: HuggingFace auth + Gemma 3 1B IT models
2026-02-28 22:49:22 +01:00 · 2026-02-28 22:36:05 +01:00 · 2026-02-28 22:25:49 +01:00
7 changed files with 299 additions and 45 deletions
@@ -51,9 +51,12 @@ dependencies {
    // ML Kit GenAI — Gemini Nano via AICore (recommended for Pixel 10)
    implementation("com.google.mlkit:genai-prompt:1.0.0-beta1")

-    // MediaPipe LLM Inference — for custom models (Gemma, etc.)
+    // MediaPipe LLM Inference — legacy fallback for .task/.bin models
    implementation("com.google.mediapipe:tasks-genai:0.10.24")

+    // LiteRT-LM — primary backend for Gemma 3n .litertlm models
+    implementation("com.google.ai.edge.litertlm:litertlm-android:0.9.0-alpha05")
+
    // Embedded HTTP server
    implementation("org.nanohttpd:nanohttpd:2.3.1")

@@ -0,0 +1,137 @@
+package com.pixel10.ai.inference
+
+import android.content.Context
+import android.util.Log
+import com.google.ai.edge.litertlm.Backend
+import com.google.ai.edge.litertlm.Engine
+import com.google.ai.edge.litertlm.EngineConfig
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.flow.catch
+import kotlinx.coroutines.withContext
+import java.io.File
+
+/**
+ * LiteRT-LM backend for Gemma 3n models (.litertlm format).
+ *
+ * This replaces MediaPipe for the newer Gemma 3n E4B/E2B models which use
+ * the LiteRT-LM runtime. Runs fully on-device using the Tensor G5 GPU.
+ *
+ * Model files must be placed in the app's files directory (see [ModelDownloader]).
+ */
+class LiteRTModel private constructor(
+    private val engine: Engine,
+    private val modelName: String
+) : OnDeviceModel {
+
+    override val backendName = "LiteRT-LM ($modelName)"
+
+    @Volatile
+    override var isReady: Boolean = true
+        private set
+
+    override suspend fun generate(
+        prompt: String,
+        maxTokens: Int,
+        temperature: Float
+    ): String = withContext(Dispatchers.Default) {
+        val conversation = engine.createConversation()
+        try {
+            conversation.sendMessage(prompt).toString()
+        } catch (e: Exception) {
+            Log.e(TAG, "LiteRT inference error", e)
+            throw OnDeviceModel.InferenceException("Generation failed: ${e.message}", e)
+        } finally {
+            conversation.close()
+        }
+    }
+
+    override suspend fun generateStreaming(
+        prompt: String,
+        onToken: (String) -> Unit
+    ): String = withContext(Dispatchers.Default) {
+        val conversation = engine.createConversation()
+        val sb = StringBuilder()
+        try {
+            conversation.sendMessageAsync(prompt)
+                .catch { e ->
+                    throw OnDeviceModel.InferenceException("Streaming failed: ${e.message}", e)
+                }
+                .collect { message ->
+                    val token = message.toString()
+                    sb.append(token)
+                    onToken(token)
+                }
+        } finally {
+            conversation.close()
+        }
+        sb.toString()
+    }
+
+    override fun close() {
+        isReady = false
+        engine.close()
+    }
+
+    companion object {
+        private const val TAG = "LiteRTModel"
+
+        private val MODEL_EXTENSIONS = listOf("litertlm")
+
+        suspend fun create(context: Context): LiteRTModel = withContext(Dispatchers.IO) {
+            val modelPath = findModelPath(context)
+                ?: throw OnDeviceModel.InferenceException(
+                    "No LiteRT-LM model file found.\n" +
+                    "Download a .litertlm model via the app or place one in:\n" +
+                    "  ${context.filesDir.absolutePath}/"
+                )
+
+            val modelName = File(modelPath).name
+            Log.i(TAG, "Loading LiteRT-LM model: $modelPath")
+
+            try {
+                val config = EngineConfig(
+                    modelPath = modelPath,
+                    backend = Backend.GPU
+                )
+                val engine = Engine(config)
+                withContext(Dispatchers.Default) {
+                    engine.initialize()
+                }
+                Log.i(TAG, "LiteRT-LM model loaded: $modelName")
+                LiteRTModel(engine, modelName)
+            } catch (gpuError: Exception) {
+                Log.w(TAG, "GPU backend failed, trying CPU: ${gpuError.message}")
+                try {
+                    val config = EngineConfig(
+                        modelPath = modelPath,
+                        backend = Backend.CPU
+                    )
+                    val engine = Engine(config)
+                    withContext(Dispatchers.Default) {
+                        engine.initialize()
+                    }
+                    Log.i(TAG, "LiteRT-LM model loaded on CPU: $modelName")
+                    LiteRTModel(engine, modelName)
+                } catch (e: Exception) {
+                    throw OnDeviceModel.InferenceException(
+                        "Failed to load LiteRT-LM model from $modelPath: ${e.message}", e
+                    )
+                }
+            }
+        }
+
+        private fun findModelPath(context: Context): String? {
+            val searchDirs = listOfNotNull(
+                context.filesDir,
+                File(context.filesDir, "models"),
+                context.getExternalFilesDir(null)
+            )
+            for (dir in searchDirs) {
+                if (!dir.exists()) continue
+                dir.listFiles()?.firstOrNull { it.extension in MODEL_EXTENSIONS }
+                    ?.let { return it.absolutePath }
+            }
+            return null
+        }
+    }
+}
@@ -10,23 +10,50 @@ import java.net.HttpURLConnection
 import java.net.URL

 /**
- * Downloads a MediaPipe-compatible Gemma model for background-safe inference.
+ * Downloads a MediaPipe-compatible model for background-safe inference.
 *
- * Gemini Nano (ML Kit) blocks inference when the app is backgrounded (ErrorCode 30).
- * MediaPipe with a local model file has no such restriction — it runs entirely in
- * the app process using the Tensor G5 GPU via OpenCL/Vulkan.
+ * All models are hosted on HuggingFace and require a free API token.
+ * Get one at: https://huggingface.co/settings/tokens
 *
- * The downloaded model is stored in the app's private files directory and
- * survives app restarts. Only needs to be downloaded once (~1.3 GB).
+ * Models use the MediaPipe `.task` format, compatible with [MediaPipeModel].
+ * Gemma 3n E4B/E2B (`.litertlm` format) requires a runtime upgrade — coming later.
 */
 object ModelDownloader {

    private const val TAG = "ModelDownloader"
+    private const val HF_BASE = "https://huggingface.co"

-    const val MODEL_FILENAME = "gemma-2b-it-gpu-int4.bin"
-    private const val MODEL_URL =
-        "https://storage.googleapis.com/mediapipe-models/llm_inference/" +
-        "gemma-2b-it-gpu-int4/float16/1/gemma-2b-it-gpu-int4.bin"
+    /** Available model specs downloadable from HuggingFace (requires token + license acceptance). */
+    enum class ModelSpec(
+        val displayName: String,
+        val filename: String,
+        val repo: String,
+        val sizeMb: Int,
+        val description: String
+    ) {
+        /**
+         * Gemma 3n E4B INT4 — best quality, Tensor G5 optimised, background-safe.
+         * Accept license at: https://huggingface.co/google/gemma-3n-E4B-it-litert-lm
+         */
+        GEMMA_3N_E4B(
+            displayName = "Gemma 3n E4B",
+            filename = "gemma-3n-E4B-it-int4.litertlm",
+            repo = "google/gemma-3n-E4B-it-litert-lm",
+            sizeMb = 4920,
+            description = "Best quality — Tensor G5 optimised (~4.9 GB)"
+        ),
+        /**
+         * Gemma 3n E4B Web INT4 — smaller variant, slightly lower quality.
+         * Same license as above.
+         */
+        GEMMA_3N_E4B_WEB(
+            displayName = "Gemma 3n E4B (Web)",
+            filename = "gemma-3n-E4B-it-int4-Web.litertlm",
+            repo = "google/gemma-3n-E4B-it-litert-lm",
+            sizeMb = 4280,
+            description = "Slightly smaller variant (~4.3 GB)"
+        )
+    }

    data class Progress(
        val downloadedBytes: Long,
@@ -34,34 +61,58 @@ object ModelDownloader {
        val percent: Int = if (totalBytes > 0) (downloadedBytes * 100 / totalBytes).toInt() else 0
    )

+    /** Returns true if any supported model is present in the app's files directory. */
    fun isModelPresent(context: Context): Boolean =
-        modelFile(context).let { it.exists() && it.length() > 1_000_000L }
+        ModelSpec.values().any { modelFile(context, it).let { f -> f.exists() && f.length() > 1_000_000L } }

-    fun modelFile(context: Context): File = File(context.filesDir, MODEL_FILENAME)
+    /** Returns the installed [ModelSpec], or null if no model is present. */
+    fun installedSpec(context: Context): ModelSpec? =
+        ModelSpec.values().firstOrNull { modelFile(context, it).let { f -> f.exists() && f.length() > 1_000_000L } }
+
+    fun modelFile(context: Context, spec: ModelSpec): File =
+        File(context.filesDir, spec.filename)
+
+    /** Returns the file of the installed model, or E4B path as default. */
+    fun modelFile(context: Context): File =
+        installedSpec(context)?.let { modelFile(context, it) }
+            ?: modelFile(context, ModelSpec.GEMMA_3N_E4B)

    /**
-     * Download the model, reporting progress via [onProgress].
+     * Download [spec] from HuggingFace, using [hfToken] for authentication.
     * Supports resume — if a partial file exists, continues from where it left off.
+     *
+     * Get a free token at https://huggingface.co/settings/tokens
     */
    suspend fun download(
        context: Context,
+        spec: ModelSpec = ModelSpec.GEMMA_3N_E4B,
+        hfToken: String,
        onProgress: (Progress) -> Unit
    ) = withContext(Dispatchers.IO) {
-        val dest = modelFile(context)
+        if (hfToken.isBlank()) throw OnDeviceModel.InferenceException(
+            "HuggingFace token required.\nGet a free token at huggingface.co/settings/tokens"
+        )
+
+        val dest = modelFile(context, spec)
        val alreadyDownloaded = if (dest.exists()) dest.length() else 0L
+        val downloadUrl = "$HF_BASE/${spec.repo}/resolve/main/${spec.filename}"

-        Log.i(TAG, "Download starting (already have $alreadyDownloaded bytes)")
+        Log.i(TAG, "Download starting ${spec.displayName} from $downloadUrl (already have $alreadyDownloaded bytes)")

-        val conn = URL(MODEL_URL).openConnection() as HttpURLConnection
+        val conn = URL(downloadUrl).openConnection() as HttpURLConnection
        try {
            conn.connectTimeout = 30_000
            conn.readTimeout   = 60_000
+            conn.setRequestProperty("Authorization", "Bearer $hfToken")
            if (alreadyDownloaded > 0) {
                conn.setRequestProperty("Range", "bytes=$alreadyDownloaded-")
            }
            conn.connect()

            val code = conn.responseCode
+            if (code == 401 || code == 403) throw OnDeviceModel.InferenceException(
+                "Authentication failed (HTTP $code).\nCheck your HuggingFace token."
+            )
            val resuming = code == HttpURLConnection.HTTP_PARTIAL   // 206
            if (code != HttpURLConnection.HTTP_OK && !resuming) {
                throw OnDeviceModel.InferenceException("Download failed: HTTP $code")
@@ -90,7 +141,7 @@ object ModelDownloader {
    }

    fun deleteModel(context: Context) {
-        modelFile(context).delete()
-        Log.i(TAG, "Model deleted")
+        ModelSpec.values().forEach { modelFile(context, it).delete() }
+        Log.i(TAG, "All models deleted")
    }
 }
@@ -112,9 +112,19 @@ interface OnDeviceModel {
         * Tap "Download Model" in the app UI to get the MediaPipe model automatically.
         */
        suspend fun create(context: Context): OnDeviceModel = withContext(Dispatchers.IO) {
-            // MediaPipe first — background-safe, GPU-accelerated via Tensor G5
+            // LiteRT-LM first — Gemma 3n .litertlm format, GPU-accelerated, background-safe
            try {
-                Log.i(TAG, "Attempting MediaPipe LLM with local model...")
+                Log.i(TAG, "Attempting LiteRT-LM with local .litertlm model...")
+                val litert = LiteRTModel.create(context)
+                Log.i(TAG, "LiteRT-LM model ready: ${litert.backendName}")
+                return@withContext litert
+            } catch (e: Exception) {
+                Log.w(TAG, "LiteRT-LM not available: ${e.message}")
+            }
+
+            // MediaPipe fallback — .task/.bin format, background-safe
+            try {
+                Log.i(TAG, "Attempting MediaPipe LLM with local .task model...")
                val mediapipe = MediaPipeModel.create(context)
                Log.i(TAG, "MediaPipe model ready: ${mediapipe.backendName}")
                return@withContext mediapipe
@@ -122,7 +132,7 @@ interface OnDeviceModel {
                Log.w(TAG, "MediaPipe not available: ${e.message}")
            }

-            // Gemini Nano fallback — only works when app is in foreground
+            // Gemini Nano last resort — foreground only
            try {
                Log.i(TAG, "Attempting Gemini Nano via ML Kit (foreground only)...")
                val nano = GeminiNanoModel.create(context)
@@ -134,11 +144,10 @@ interface OnDeviceModel {

            throw InferenceException(
                "No model loaded yet.\n\n" +
-                "Tap 'Download Model' in the app to download Gemma 2B (~1.3 GB).\n" +
+                "Tap 'Download Model' in the app to download Gemma 3n E4B.\n" +
                "Once downloaded the server works fully in the background.\n\n" +
-                "Or place a compatible model file in:\n" +
-                "  ${context.filesDir.absolutePath}/\n" +
-                "  Supported: gemma-2b-it-gpu-int4.bin, gemma-3n-E2B.task, etc."
+                "Or place a .litertlm file in:\n" +
+                "  ${context.filesDir.absolutePath}/"
            )
        }
    }
@@ -5,6 +5,7 @@ import android.content.ComponentName
 import android.content.Context
 import android.content.Intent
 import android.content.ServiceConnection
+import android.content.SharedPreferences
 import android.content.pm.PackageManager
 import android.graphics.drawable.GradientDrawable
 import android.net.wifi.WifiManager
@@ -19,6 +20,7 @@ import androidx.lifecycle.lifecycleScope
 import com.pixel10.ai.R
 import com.pixel10.ai.databinding.ActivityMainBinding
 import com.pixel10.ai.inference.ModelDownloader
+import com.pixel10.ai.inference.ModelDownloader.ModelSpec
 import com.pixel10.ai.server.ApiServerService
 import kotlinx.coroutines.launch
 import java.text.SimpleDateFormat
@@ -28,6 +30,7 @@ import java.util.Locale
 class MainActivity : AppCompatActivity() {

    private lateinit var binding: ActivityMainBinding
+    private lateinit var prefs: SharedPreferences
    private var service: ApiServerService? = null
    private var bound = false
    private var downloading = false
@@ -67,14 +70,23 @@ class MainActivity : AppCompatActivity() {
        binding = ActivityMainBinding.inflate(layoutInflater)
        setContentView(binding.root)

+        prefs = getSharedPreferences("pixel10_prefs", MODE_PRIVATE)
        requestNotificationPermission()

+        // Restore saved HF token
+        binding.etHfToken.setText(prefs.getString("hf_token", ""))
+
        binding.btnToggle.setOnClickListener {
            if (service?.isRunning == true) stopServer() else startServer()
        }

        binding.btnDownloadModel.setOnClickListener {
-            startModelDownload()
+            saveHfToken()
+            startModelDownload(ModelSpec.GEMMA_3N_E4B)
+        }
+        binding.btnDownloadGemma3Q8.setOnClickListener {
+            saveHfToken()
+            startModelDownload(ModelSpec.GEMMA_3N_E4B_WEB)
        }

        updateModelCard()
@@ -121,34 +133,43 @@ class MainActivity : AppCompatActivity() {
        }
    }

-    private fun startModelDownload() {
+    private fun saveHfToken() {
+        val token = binding.etHfToken.text.toString().trim()
+        prefs.edit().putString("hf_token", token).apply()
+    }
+
+    private fun startModelDownload(spec: ModelSpec) {
        if (downloading) return
+        val token = binding.etHfToken.text.toString().trim()
+        if (token.isBlank()) {
+            binding.tvModelDownloadStatus.text = "Enter your HuggingFace token first"
+            return
+        }
        downloading = true
-        binding.btnDownloadModel.isEnabled = false
-        binding.btnDownloadModel.text = "Downloading…"
+        setDownloadButtonsEnabled(false)
        binding.progressDownload.visibility = View.VISIBLE
-        binding.tvModelDownloadStatus.text = "Starting download…"
+        binding.tvModelDownloadStatus.text = "Starting download: ${spec.displayName}…"

        lifecycleScope.launch {
            try {
-                ModelDownloader.download(this@MainActivity) { progress ->
+                ModelDownloader.download(this@MainActivity, spec, token) { progress ->
                    runOnUiThread {
                        binding.progressDownload.progress = progress.percent
                        val mb = progress.downloadedBytes / 1_048_576
                        val total = progress.totalBytes / 1_048_576
-                        binding.tvModelDownloadStatus.text = "Downloading… ${mb}MB / ${total}MB (${progress.percent}%)"
+                        binding.tvModelDownloadStatus.text =
+                            "${spec.displayName}: ${mb}MB / ${total}MB (${progress.percent}%)"
                    }
                }
                runOnUiThread {
                    downloading = false
                    updateModelCard()
-                    appendLog("Model downloaded — background inference enabled")
+                    appendLog("${spec.displayName} downloaded — background inference enabled")
                }
            } catch (e: Exception) {
                runOnUiThread {
                    downloading = false
-                    binding.btnDownloadModel.isEnabled = true
-                    binding.btnDownloadModel.text = getString(R.string.btn_download_model)
+                    setDownloadButtonsEnabled(true)
                    binding.progressDownload.visibility = View.GONE
                    binding.tvModelDownloadStatus.text = "Download failed: ${e.message}"
                    appendLog("Download error: ${e.message}")
@@ -157,17 +178,25 @@ class MainActivity : AppCompatActivity() {
        }
    }

+    private fun setDownloadButtonsEnabled(enabled: Boolean) {
+        binding.btnDownloadModel.isEnabled = enabled
+        binding.btnDownloadGemma3Q8.isEnabled = enabled
+    }
+
    private fun updateModelCard() {
-        val present = ModelDownloader.isModelPresent(this)
-        if (present) {
-            binding.tvModelDownloadStatus.text = getString(R.string.model_downloaded)
+        val spec = ModelDownloader.installedSpec(this)
+        if (spec != null) {
+            binding.tvModelDownloadStatus.text = getString(R.string.model_downloaded, spec.displayName)
+            binding.etHfToken.visibility = View.GONE
            binding.btnDownloadModel.visibility = View.GONE
+            binding.btnDownloadGemma3Q8.visibility = View.GONE
            binding.progressDownload.visibility = View.GONE
        } else {
            binding.tvModelDownloadStatus.text = getString(R.string.model_not_downloaded)
+            binding.etHfToken.visibility = View.VISIBLE
            binding.btnDownloadModel.visibility = View.VISIBLE
-            binding.btnDownloadModel.isEnabled = true
-            binding.btnDownloadModel.text = getString(R.string.btn_download_model)
+            binding.btnDownloadGemma3Q8.visibility = View.VISIBLE
+            setDownloadButtonsEnabled(true)
            binding.progressDownload.visibility = View.GONE
        }
    }
@@ -132,6 +132,19 @@
                android:textColor="@color/log_text"
                android:textSize="13sp" />

+            <com.google.android.material.textfield.TextInputEditText
+                android:id="@+id/etHfToken"
+                android:layout_width="match_parent"
+                android:layout_height="48dp"
+                android:layout_marginTop="8dp"
+                android:hint="@string/hf_token_hint"
+                android:inputType="textPassword"
+                android:textColor="@color/on_surface"
+                android:textColorHint="@color/log_text"
+                android:textSize="13sp"
+                android:fontFamily="monospace"
+                android:backgroundTint="@color/primary" />
+
            <ProgressBar
                android:id="@+id/progressDownload"
                style="@android:style/Widget.ProgressBar.Horizontal"
@@ -147,7 +160,17 @@
                android:layout_width="wrap_content"
                android:layout_height="wrap_content"
                android:layout_marginTop="8dp"
-                android:text="@string/btn_download_model"
+                android:text="@string/btn_download_gemma3_q4"
+                android:textSize="13sp"
+                app:cornerRadius="8dp" />
+
+            <com.google.android.material.button.MaterialButton
+                android:id="@+id/btnDownloadGemma3Q8"
+                style="@style/Widget.MaterialComponents.Button.OutlinedButton"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:layout_marginTop="4dp"
+                android:text="@string/btn_download_gemma3_q8"
                android:textSize="13sp"
                app:cornerRadius="8dp" />
        </LinearLayout>
@@ -18,9 +18,11 @@
    <string name="model_not_loaded">Model: not loaded</string>

    <!-- Controls -->
-    <string name="btn_download_model">Download Model (~1.3 GB)</string>
-    <string name="model_downloaded">✓ Gemma 2B ready — background inference enabled</string>
-    <string name="model_not_downloaded">No local model. Download to enable background inference.</string>
+    <string name="hf_token_hint">HuggingFace token (huggingface.co/settings/tokens)</string>
+    <string name="btn_download_gemma3_q4">⭐ Gemma 3n E4B — Best quality (~4.9 GB)</string>
+    <string name="btn_download_gemma3_q8">Gemma 3n E4B Web — Smaller (~4.3 GB)</string>
+    <string name="model_downloaded">✓ %s ready — background inference enabled</string>
+    <string name="model_not_downloaded">No local model. Enter HuggingFace token and download.</string>
    <string name="btn_start">Start Server</string>
    <string name="btn_stop">Stop Server</string>
    <string name="port_label">Port:</string>