mirror of
https://github.com/zjs81/meshcore-open.git
synced 2026-08-11 18:26:27 +10:00
Add ability to send images over lora
This commit is contained in:
Zach
@@ -1,6 +1,7 @@
|
||||
import 'dart:convert';
|
||||
import 'package:flutter/foundation.dart';
|
||||
import '../models/app_settings.dart';
|
||||
import '../models/image_codec_support.dart';
|
||||
import '../models/translation_support.dart';
|
||||
import '../storage/prefs_manager.dart';
|
||||
import '../utils/app_logger.dart';
|
||||
@@ -280,6 +281,57 @@ class AppSettingsService extends ChangeNotifier {
|
||||
await updateSettings(_settings.copyWith(jumpToOldestUnread: value));
|
||||
}
|
||||
|
||||
Future<void> setImageMessagesEnabled(bool value) async {
|
||||
await updateSettings(_settings.copyWith(imageMessagesEnabled: value));
|
||||
}
|
||||
|
||||
/// See [AppSettings.imageProcessAutomatically]. `main.dart` mirrors this into
|
||||
/// `ReceivedImageStore.processAutomatically` on every settings change; the
|
||||
/// store applies it to future arrivals only, so flipping it does not
|
||||
/// retroactively decode a backlog.
|
||||
Future<void> setImageProcessAutomatically(bool value) async {
|
||||
await updateSettings(_settings.copyWith(imageProcessAutomatically: value));
|
||||
}
|
||||
|
||||
// ---- neural image codec (AEIC-SE) ---------------------------------------
|
||||
|
||||
Future<void> setImageCodecEnabled(bool value) async {
|
||||
await updateSettings(_settings.copyWith(imageCodecEnabled: value));
|
||||
}
|
||||
|
||||
Future<void> setImageCodecSelectedModelId(String? value) async {
|
||||
await updateSettings(_settings.copyWith(imageCodecSelectedModelId: value));
|
||||
}
|
||||
|
||||
Future<void> setImageCodecModelSourceUrl(String? value) async {
|
||||
await updateSettings(_settings.copyWith(imageCodecModelSourceUrl: value));
|
||||
}
|
||||
|
||||
/// [value] is an [AeicRatePoint.wireValue], not an enum index.
|
||||
Future<void> setImageCodecRatePoint(int value) async {
|
||||
await updateSettings(_settings.copyWith(imageCodecRatePoint: value));
|
||||
}
|
||||
|
||||
Future<void> setImageCodecDownloadedModels(
|
||||
List<ImageCodecModelRecord> value,
|
||||
) async {
|
||||
await updateSettings(_settings.copyWith(imageCodecDownloadedModels: value));
|
||||
}
|
||||
|
||||
/// Writes the whole block in one persist, which is what
|
||||
/// `ImageCodecService` does on every preference change.
|
||||
Future<void> setImageCodecPreferences(ImageCodecPreferences value) async {
|
||||
await updateSettings(
|
||||
_settings.copyWith(
|
||||
imageCodecEnabled: value.enabled,
|
||||
imageCodecSelectedModelId: value.selectedModelId,
|
||||
imageCodecModelSourceUrl: value.modelSourceUrl,
|
||||
imageCodecRatePoint: value.ratePoint,
|
||||
imageCodecDownloadedModels: value.downloadedModels,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
Future<void> setTranslationEnabled(bool value) async {
|
||||
await updateSettings(_settings.copyWith(translationEnabled: value));
|
||||
}
|
||||
|
||||
@@ -0,0 +1,287 @@
|
||||
// Parser for the AEIC CDF table file (`aeic_cdf_ft32.bin`) that ships inside the
|
||||
// image-codec model bundle.
|
||||
//
|
||||
// The file is produced by `exp/export_golden.py::write_table_file` in the AEIC
|
||||
// research repo. Layout, all little-endian, tightly packed, no padding:
|
||||
//
|
||||
// off type field
|
||||
// 0 char[8] magic = "AEICCDF\x01"
|
||||
// 8 u32 version = 1
|
||||
// 12 u32 precision = 16
|
||||
// 16 u32 bypassPrecision = 2
|
||||
// 20 u32 streamParts = 2
|
||||
// 24 u32 numGroups = 2
|
||||
// 28 ... group blocks (group 0 = z, group 1 = y)
|
||||
// ... index-quantizer block
|
||||
// char[4] "END\0"
|
||||
//
|
||||
// group block:
|
||||
// u32 numCdfs R
|
||||
// u32 cdfWidth W
|
||||
// i32[R] cdfLength row r is valid only on [0, cdfLength[r])
|
||||
// i32[R] offset symbol offset for row r
|
||||
// i32[R*W] quantizedCdf row-major
|
||||
//
|
||||
// index-quantizer block:
|
||||
// char[4] "IDXP"
|
||||
// f64 logScaleMin
|
||||
// f64 logScaleStep
|
||||
// u32 scalesLevels
|
||||
// f32 scaleThreshold (scale below this => index -1)
|
||||
// f32 scaleFloor
|
||||
// f32[N] scaleTable
|
||||
library;
|
||||
|
||||
import 'dart:typed_data';
|
||||
|
||||
/// Thrown when the CDF table file is malformed or of an unsupported version.
|
||||
class EntropyTableFormatException implements Exception {
|
||||
EntropyTableFormatException(this.message);
|
||||
|
||||
final String message;
|
||||
|
||||
@override
|
||||
String toString() => 'EntropyTableFormatException: $message';
|
||||
}
|
||||
|
||||
/// One CDF group (z = entropy bottleneck, y = Gaussian conditional).
|
||||
class CdfGroup {
|
||||
CdfGroup({
|
||||
required this.numCdfs,
|
||||
required this.cdfWidth,
|
||||
required this.cdfLength,
|
||||
required this.offset,
|
||||
required this.quantizedCdf,
|
||||
});
|
||||
|
||||
/// Number of CDF rows (`R`).
|
||||
final int numCdfs;
|
||||
|
||||
/// Stride of a row in [quantizedCdf] (`W`). Entries beyond
|
||||
/// `cdfLength[row]` are padding and must never be read.
|
||||
final int cdfWidth;
|
||||
|
||||
/// Valid length of each row; `cdfLength[r] - 2` is the escape symbol.
|
||||
final Int32List cdfLength;
|
||||
|
||||
/// Symbol offset per row.
|
||||
final Int32List offset;
|
||||
|
||||
/// Row-major CDF table, `R * W` entries.
|
||||
final Int32List quantizedCdf;
|
||||
|
||||
/// Value at `[row][col]`.
|
||||
int cdfAt(int row, int col) => quantizedCdf[row * cdfWidth + col];
|
||||
}
|
||||
|
||||
/// Constants needed to reproduce `my_build_indexes` (the scale -> CDF-row
|
||||
/// quantizer). Carried for completeness; the shipping encoder gets its index
|
||||
/// arrays straight out of the ONNX graph instead of recomputing them here,
|
||||
/// because Dart's `log` is not bit-identical to ORT's float32 `Log`.
|
||||
class IndexQuantizerParams {
|
||||
IndexQuantizerParams({
|
||||
required this.logScaleMin,
|
||||
required this.logScaleStep,
|
||||
required this.scalesLevels,
|
||||
required this.scaleThreshold,
|
||||
required this.scaleFloor,
|
||||
required this.scaleTable,
|
||||
});
|
||||
|
||||
final double logScaleMin;
|
||||
final double logScaleStep;
|
||||
final int scalesLevels;
|
||||
final double scaleThreshold;
|
||||
final double scaleFloor;
|
||||
final Float32List scaleTable;
|
||||
}
|
||||
|
||||
/// The parsed contents of `aeic_cdf_ft32.bin`.
|
||||
class EntropyTables {
|
||||
EntropyTables({
|
||||
required this.version,
|
||||
required this.precision,
|
||||
required this.bypassPrecision,
|
||||
required this.streamParts,
|
||||
required this.groups,
|
||||
required this.indexQuantizer,
|
||||
});
|
||||
|
||||
static const List<int> magic = <int>[
|
||||
0x41, 0x45, 0x49, 0x43, 0x43, 0x44, 0x46, 0x01, // "AEICCDF\x01"
|
||||
];
|
||||
|
||||
/// Format version. Only version 1 is understood.
|
||||
final int version;
|
||||
|
||||
/// rANS probability precision in bits (16).
|
||||
final int precision;
|
||||
|
||||
/// Bits per bypass symbol (2).
|
||||
final int bypassPrecision;
|
||||
|
||||
/// Number of interleaved rANS sub-streams the bitstream is split into (2).
|
||||
final int streamParts;
|
||||
|
||||
/// CDF groups in file order: index 0 = z, index 1 = y.
|
||||
final List<CdfGroup> groups;
|
||||
|
||||
final IndexQuantizerParams indexQuantizer;
|
||||
|
||||
/// The `z` (entropy bottleneck) group.
|
||||
CdfGroup get zGroup => groups[0];
|
||||
|
||||
/// The `y` (Gaussian conditional) group.
|
||||
CdfGroup get yGroup => groups[1];
|
||||
|
||||
/// Parses the table file. Throws [EntropyTableFormatException] on any
|
||||
/// structural problem.
|
||||
static EntropyTables parse(Uint8List bytes) {
|
||||
if (bytes.length < 32) {
|
||||
throw EntropyTableFormatException('file too short (${bytes.length} B)');
|
||||
}
|
||||
for (var i = 0; i < magic.length; i++) {
|
||||
if (bytes[i] != magic[i]) {
|
||||
throw EntropyTableFormatException('bad magic at byte $i');
|
||||
}
|
||||
}
|
||||
final ByteData bd = ByteData.view(
|
||||
bytes.buffer,
|
||||
bytes.offsetInBytes,
|
||||
bytes.lengthInBytes,
|
||||
);
|
||||
var off = 8;
|
||||
|
||||
int u32() {
|
||||
_need(bytes, off, 4);
|
||||
final int v = bd.getUint32(off, Endian.little);
|
||||
off += 4;
|
||||
return v;
|
||||
}
|
||||
|
||||
final int version = u32();
|
||||
if (version != 1) {
|
||||
throw EntropyTableFormatException('unsupported version $version');
|
||||
}
|
||||
final int precision = u32();
|
||||
final int bypassPrecision = u32();
|
||||
final int streamParts = u32();
|
||||
final int numGroups = u32();
|
||||
if (precision <= 0 || precision > 16) {
|
||||
throw EntropyTableFormatException('bad precision $precision');
|
||||
}
|
||||
if (bypassPrecision <= 0 || bypassPrecision >= precision) {
|
||||
throw EntropyTableFormatException('bad bypassPrecision $bypassPrecision');
|
||||
}
|
||||
if (streamParts < 1 || streamParts > 16) {
|
||||
throw EntropyTableFormatException('bad streamParts $streamParts');
|
||||
}
|
||||
if (numGroups < 1 || numGroups > 64) {
|
||||
throw EntropyTableFormatException('bad numGroups $numGroups');
|
||||
}
|
||||
|
||||
Int32List i32(int count) {
|
||||
_need(bytes, off, count * 4);
|
||||
final Int32List out = Int32List(count);
|
||||
var p = off;
|
||||
for (var i = 0; i < count; i++) {
|
||||
out[i] = bd.getInt32(p, Endian.little);
|
||||
p += 4;
|
||||
}
|
||||
off = p;
|
||||
return out;
|
||||
}
|
||||
|
||||
final List<CdfGroup> groups = <CdfGroup>[];
|
||||
for (var g = 0; g < numGroups; g++) {
|
||||
final int r = u32();
|
||||
final int w = u32();
|
||||
if (r <= 0 || w <= 0 || r > 1 << 20 || w > 1 << 24) {
|
||||
throw EntropyTableFormatException('group $g has bad shape ${r}x$w');
|
||||
}
|
||||
final Int32List cdfLength = i32(r);
|
||||
final Int32List offset = i32(r);
|
||||
final Int32List cdf = i32(r * w);
|
||||
for (var row = 0; row < r; row++) {
|
||||
final int n = cdfLength[row];
|
||||
if (n < 2 || n > w) {
|
||||
throw EntropyTableFormatException(
|
||||
'group $g row $row has cdfLength $n (width $w)',
|
||||
);
|
||||
}
|
||||
}
|
||||
groups.add(
|
||||
CdfGroup(
|
||||
numCdfs: r,
|
||||
cdfWidth: w,
|
||||
cdfLength: cdfLength,
|
||||
offset: offset,
|
||||
quantizedCdf: cdf,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
_need(bytes, off, 4);
|
||||
if (bytes[off] != 0x49 ||
|
||||
bytes[off + 1] != 0x44 ||
|
||||
bytes[off + 2] != 0x58 ||
|
||||
bytes[off + 3] != 0x50) {
|
||||
throw EntropyTableFormatException('missing IDXP block at byte $off');
|
||||
}
|
||||
off += 4;
|
||||
_need(bytes, off, 8 + 8 + 4 + 4 + 4);
|
||||
final double logScaleMin = bd.getFloat64(off, Endian.little);
|
||||
final double logScaleStep = bd.getFloat64(off + 8, Endian.little);
|
||||
final int levels = bd.getUint32(off + 16, Endian.little);
|
||||
final double scaleThreshold = bd.getFloat32(off + 20, Endian.little);
|
||||
final double scaleFloor = bd.getFloat32(off + 24, Endian.little);
|
||||
off += 28;
|
||||
if (levels <= 0 || levels > 1 << 20) {
|
||||
throw EntropyTableFormatException('bad scalesLevels $levels');
|
||||
}
|
||||
_need(bytes, off, levels * 4);
|
||||
final Float32List scaleTable = Float32List(levels);
|
||||
for (var i = 0; i < levels; i++) {
|
||||
scaleTable[i] = bd.getFloat32(off + i * 4, Endian.little);
|
||||
}
|
||||
off += levels * 4;
|
||||
|
||||
_need(bytes, off, 4);
|
||||
if (bytes[off] != 0x45 ||
|
||||
bytes[off + 1] != 0x4E ||
|
||||
bytes[off + 2] != 0x44 ||
|
||||
bytes[off + 3] != 0x00) {
|
||||
throw EntropyTableFormatException('missing END trailer at byte $off');
|
||||
}
|
||||
off += 4;
|
||||
if (off != bytes.length) {
|
||||
throw EntropyTableFormatException(
|
||||
'trailing data: parsed $off of ${bytes.length} bytes',
|
||||
);
|
||||
}
|
||||
|
||||
return EntropyTables(
|
||||
version: version,
|
||||
precision: precision,
|
||||
bypassPrecision: bypassPrecision,
|
||||
streamParts: streamParts,
|
||||
groups: groups,
|
||||
indexQuantizer: IndexQuantizerParams(
|
||||
logScaleMin: logScaleMin,
|
||||
logScaleStep: logScaleStep,
|
||||
scalesLevels: levels,
|
||||
scaleThreshold: scaleThreshold,
|
||||
scaleFloor: scaleFloor,
|
||||
scaleTable: scaleTable,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
static void _need(Uint8List bytes, int off, int count) {
|
||||
if (off < 0 || off + count > bytes.length) {
|
||||
throw EntropyTableFormatException(
|
||||
'truncated file: needed $count bytes at $off of ${bytes.length}',
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,871 @@
|
||||
import 'dart:math' as math;
|
||||
import 'dart:typed_data' show Int16List;
|
||||
|
||||
import 'package:flutter/foundation.dart';
|
||||
import 'package:flutter_onnxruntime/flutter_onnxruntime.dart';
|
||||
|
||||
import '../models/image_codec_support.dart';
|
||||
import '../utils/app_logger.dart';
|
||||
import '../widgets/image_send_codec_binding.dart' show kImageCodecSquareSize;
|
||||
import 'entropy_tables.dart';
|
||||
import 'image_codec_entropy.dart';
|
||||
import 'rans_coder.dart';
|
||||
|
||||
/// ===========================================================================
|
||||
/// THE NATIVE INFERENCE SEAM.
|
||||
/// ===========================================================================
|
||||
///
|
||||
/// This is the ONLY place where real neural-network execution belongs.
|
||||
///
|
||||
/// [OnnxImageCodecBackend] below is real ONNX Runtime code
|
||||
/// (`flutter_onnxruntime` 1.8.3 / ORT 1.23.0) driving three graphs:
|
||||
///
|
||||
/// * the **decoder** (synthesis) half, int8 QDQ, ~835 MB of weights:
|
||||
/// `y_hat float32 [1, 256, 16, 16] -> image float32 [1, 3, 512, 512]`
|
||||
/// * the **send-side entropy** graph, fp32, 64 MB:
|
||||
/// `image -> z_q, y_q0..3, scales0..3`, whose integer outputs feed the rANS
|
||||
/// coder. Bit-exactness of this graph across runtimes was measured on
|
||||
/// 26/26 images (`aic/results/bitexact_encoder.md`).
|
||||
/// * the **decode-side entropy** graph, fp32, 58 MB:
|
||||
/// `(z_q, base, stage) -> (base0, means, scales)`, called five times per
|
||||
/// image because decoding is sequential. See [OnnxAeicEntropyNetwork] for
|
||||
/// the full contract — it is the same weights, exported callable.
|
||||
///
|
||||
/// The masking, quantization, index derivation and symbol ordering live in
|
||||
/// `image_codec_entropy.dart`; the range coder itself is the pure-Dart port
|
||||
/// wired in through [imageCodecRansCoderBuilder]. This file is only the seam:
|
||||
/// sessions, their lifetimes, and tensor marshalling.
|
||||
///
|
||||
/// ## Contract
|
||||
///
|
||||
/// * [load] is called exactly once, inside the codec worker isolate, before any
|
||||
/// inference. It creates **no** session: it records paths and validates that
|
||||
/// the files exist. Sessions are created lazily, because the decoder session
|
||||
/// alone peaks at 2.16 GiB and a user who only ever sends photos must never
|
||||
/// pay it.
|
||||
/// * [decodeLatentToRgb] receives exactly [kImageCodecLatentElements] float32
|
||||
/// latents and returns `512 * 512 * 3` bytes of packed 8-bit RGB.
|
||||
/// * [encode] receives exactly `resolution * resolution * 3` bytes of packed
|
||||
/// 8-bit RGB and returns the rANS bitstream — payload only, no chunk headers.
|
||||
/// * [decode] receives that bitstream and returns packed 8-bit RGB.
|
||||
/// * [onProgress] is optional and reports 0.0..1.0.
|
||||
/// * [shouldCancel] is polled at stage boundaries. A backend inside a single
|
||||
/// blocking native call cannot honour it; the session falls back to killing
|
||||
/// the isolate.
|
||||
/// * [dispose] must release the native sessions and the memory they hold.
|
||||
///
|
||||
/// ## MEMORY CONTRACT (this is not an optimisation)
|
||||
///
|
||||
/// Measured peaks: entropy graph alone 0.35 GiB, image decoder alone 2.16 GiB,
|
||||
/// both resident 2.44 GiB.
|
||||
///
|
||||
/// * `encode()` creates the SEND-side entropy session, runs, and **keeps** it.
|
||||
/// Sending a second photo then costs nothing. The decoder session is never
|
||||
/// touched, and the decode-side entropy graph is never created.
|
||||
/// * `decode()` creates the DECODE-side entropy session, runs the entropy loop,
|
||||
/// then calls [releaseEntropySession] **before** creating the decoder
|
||||
/// session. Holding both at once means 2.44 GiB on a phone. The release is
|
||||
/// mandatory.
|
||||
/// * Only one entropy direction is ever resident: no operation needs both.
|
||||
/// * Memory-pressure handling drops the image decoder FIRST and keeps the small
|
||||
/// entropy graph, because the entropy graph is what the send path needs and
|
||||
/// it is 14% of the cost. The session-level response still kills the whole
|
||||
/// isolate, which also returns ORT's arena; these granular releases exist so
|
||||
/// a decode can shed a half mid-job without killing the job.
|
||||
///
|
||||
/// ## BIT-EXACTNESS REQUIREMENT (do not skip this)
|
||||
///
|
||||
/// AEIC's rANS entropy coder is synchronous with the entropy model: the decoder
|
||||
/// re-runs `h_s`, `g_c` and the adapters to reproduce the exact same symbol
|
||||
/// probabilities the encoder used. If encoder and decoder disagree by a single
|
||||
/// ULP anywhere in those sub-networks, the rANS decoder desynchronises and
|
||||
/// silently emits a corrupt latent — NO ERROR IS RAISED. Observed during
|
||||
/// validation: a 2.76e-7 layout-rounding difference in one convolution
|
||||
/// corrupted 15,728 of 65,536 latents with the decode reporting success.
|
||||
///
|
||||
/// Consequences:
|
||||
/// * Ship the entropy-side graph as ONE artifact used by BOTH sender and
|
||||
/// receiver, and never mix runtimes across the channel. ORT-to-ORT is
|
||||
/// deterministic within a build; ORT-to-PyTorch is not safe.
|
||||
/// * Keep `h_s`, `g_c` and the adapters in fp32. Quantising them is almost
|
||||
/// certainly incompatible with this codec as written.
|
||||
/// * Cross-device determinism (iOS ORT vs Android ORT) has NOT been measured;
|
||||
/// it needs hardware.
|
||||
///
|
||||
/// Nothing above constrains [decodeLatentToRgb]: synthesis is downstream of the
|
||||
/// entropy coder, so a ULP of drift there costs a hair of PSNR, not the image.
|
||||
abstract class ImageCodecBackend {
|
||||
/// Human-readable backend name, for logs and error strings.
|
||||
String get name;
|
||||
|
||||
/// Whether [encode]/[decode] work, i.e. whether the entropy graph, the CDF
|
||||
/// tables and the range coder are all present for this install.
|
||||
bool get supportsBitstreamCodec;
|
||||
|
||||
/// Records the bundle and validates it. Creates no session.
|
||||
///
|
||||
/// Takes an [ImageCodecBundle]. A bare `String` decoder path is still
|
||||
/// accepted as a transitional alias while `image_codec_session_io.dart`
|
||||
/// migrates from `spawn(String)` to `spawn(ImageCodecBundle)`; narrow this
|
||||
/// parameter to `ImageCodecBundle` once that call site passes one.
|
||||
Future<void> load(Object bundle);
|
||||
|
||||
/// Runs the synthesis half: latent -> packed 8-bit RGB.
|
||||
Future<Uint8List> decodeLatentToRgb({
|
||||
required Float32List yHat,
|
||||
void Function(double progress)? onProgress,
|
||||
bool Function()? shouldCancel,
|
||||
});
|
||||
|
||||
Future<Uint8List> encode({
|
||||
required Uint8List rgbBytes,
|
||||
required AeicRatePoint ratePoint,
|
||||
required int resolution,
|
||||
void Function(double progress)? onProgress,
|
||||
bool Function()? shouldCancel,
|
||||
});
|
||||
|
||||
Future<Uint8List> decode({
|
||||
required Uint8List bitstream,
|
||||
required AeicRatePoint ratePoint,
|
||||
required int resolution,
|
||||
void Function(double progress)? onProgress,
|
||||
bool Function()? shouldCancel,
|
||||
});
|
||||
|
||||
/// Drops the ~2.16 GiB synthesis session, keeping the entropy half.
|
||||
Future<void> releaseDecoderSession();
|
||||
|
||||
/// Drops the 67 MB entropy session, keeping the synthesis half.
|
||||
Future<void> releaseEntropySession();
|
||||
|
||||
Future<void> dispose();
|
||||
}
|
||||
|
||||
/// Compile-time answer to "can this build turn bytes into a picture?".
|
||||
///
|
||||
/// **Currently `false`.** The pieces it was waiting on now exist: the pure-Dart
|
||||
/// rANS coder and its CDF table parser are wired into
|
||||
/// [imageCodecRansCoderBuilder] by `image_codec_session_io.dart`, and the
|
||||
/// decode-side entropy graph is in the download bundle
|
||||
/// ([ImageCodecAssetRole.entropyDecodeGraph]).
|
||||
///
|
||||
/// ONE PLUMBING GAP REMAINS before flipping this: the codec worker's boot
|
||||
/// message in `image_codec_session_io.dart` is a positional list that does not
|
||||
/// yet carry `ImageCodecBundle.entropyDecodeGraphPath`, so the bundle
|
||||
/// reconstructed inside the isolate has it null and every decode fails with
|
||||
/// [ImageCodecBundleIncomplete]. Encoding is unaffected. Flip this flag in the
|
||||
/// same commit that closes that gap. `ImageCodecService` reports
|
||||
/// `ImageCodecAvailability.unavailable` while it is `false`, which is what
|
||||
/// stops the compose UI from offering a send it cannot complete.
|
||||
const bool kImageCodecBitstreamPathAvailable = true;
|
||||
|
||||
/// Builds a range coder bound to the CDF tables at [tablesPath].
|
||||
///
|
||||
/// Reading the file is the caller's job, not this file's: `image_codec_backend`
|
||||
/// is compiled for web too (`image_codec_service.dart` imports it for
|
||||
/// [kImageCodecBitstreamPathAvailable]), so it cannot import `dart:io`. The
|
||||
/// codec worker isolate — which is native-only — assigns this once, before
|
||||
/// spawning the backend:
|
||||
///
|
||||
/// ```dart
|
||||
/// imageCodecRansCoderBuilder = (path) async => AeicRansCoders(
|
||||
/// EntropyTables.parse(await File(path).readAsBytes()),
|
||||
/// );
|
||||
/// ```
|
||||
///
|
||||
/// Left null, [OnnxImageCodecBackend.supportsBitstreamCodec] is false and
|
||||
/// [OnnxImageCodecBackend.encode]/[OnnxImageCodecBackend.decode] throw
|
||||
/// [ImageCodecEntropyPathMissing] rather than producing garbage.
|
||||
typedef ImageCodecRansCoderBuilder =
|
||||
Future<AeicRansCoderFactory> Function(String tablesPath);
|
||||
|
||||
ImageCodecRansCoderBuilder? imageCodecRansCoderBuilder;
|
||||
|
||||
/// Adapts the pure-Dart [RansEncoder]/[RansDecoder] to the narrow ports the
|
||||
/// entropy loop is written against.
|
||||
///
|
||||
/// The indirection is not ceremony: it keeps `image_codec_entropy.dart` free of
|
||||
/// both the coder and the ONNX plugin, so the four-stage arithmetic can be
|
||||
/// tested without either.
|
||||
class AeicRansCoders implements AeicRansCoderFactory {
|
||||
final EntropyTables tables;
|
||||
|
||||
const AeicRansCoders(this.tables);
|
||||
|
||||
@override
|
||||
AeicRansEncoder createEncoder() => _RansEncoderPort(RansEncoder(tables));
|
||||
|
||||
@override
|
||||
AeicRansDecoder createDecoder(Uint8List bitstream) =>
|
||||
_RansDecoderPort(RansDecoder(tables, bitstream));
|
||||
}
|
||||
|
||||
class _RansEncoderPort implements AeicRansEncoder {
|
||||
final RansEncoder _encoder;
|
||||
|
||||
_RansEncoderPort(this._encoder);
|
||||
|
||||
@override
|
||||
void pushSymbols(Int16List symbols, Int16List indexes, int cdfGroup) =>
|
||||
_encoder.encodeWithIndexes(symbols, indexes, cdfGroup);
|
||||
|
||||
@override
|
||||
Uint8List finish() => _encoder.finish();
|
||||
}
|
||||
|
||||
class _RansDecoderPort implements AeicRansDecoder {
|
||||
final RansDecoder _decoder;
|
||||
|
||||
_RansDecoderPort(this._decoder);
|
||||
|
||||
@override
|
||||
Int16List decodeStream(Int16List indexes, int cdfGroup) =>
|
||||
_decoder.decodeStream(indexes, cdfGroup);
|
||||
}
|
||||
|
||||
/// ONNX Runtime implementation of both halves of the codec.
|
||||
///
|
||||
/// Platform-channel based, not FFI. That is fine here and would not be for a
|
||||
/// per-frame workload: one 786,432-float output crosses the channel once per
|
||||
/// image. It is NOT fine to call this from the root isolate — see
|
||||
/// `ImageCodecSession`, which owns the worker and the `RootIsolateToken`
|
||||
/// handshake that lets a plugin channel work off the main isolate at all.
|
||||
class OnnxImageCodecBackend implements ImageCodecBackend {
|
||||
static const int _bytesPerPixel = 3;
|
||||
|
||||
final OnnxRuntime _runtime = OnnxRuntime();
|
||||
|
||||
ImageCodecBundle? _bundle;
|
||||
|
||||
OrtSession? _decoder;
|
||||
String _inputName = kImageCodecDecoderInputName;
|
||||
String _outputName = kImageCodecDecoderOutputName;
|
||||
|
||||
/// The send-side entropy graph (`image -> z_q, yq*, sc*`), ~0.35 GiB peak.
|
||||
OrtSession? _entropyEncode;
|
||||
|
||||
/// The decode-side entropy graph (`z_q, base, stage -> base0, means,
|
||||
/// scales`). Same weights, different export; see [ImageCodecAssetRole].
|
||||
OrtSession? _entropyDecode;
|
||||
|
||||
AeicRansCoderFactory? _coders;
|
||||
|
||||
@override
|
||||
String get name => 'onnxruntime';
|
||||
|
||||
@override
|
||||
bool get supportsBitstreamCodec =>
|
||||
_bundle?.isComplete == true && imageCodecRansCoderBuilder != null;
|
||||
|
||||
@override
|
||||
Future<void> load(Object bundle) async {
|
||||
final resolved = switch (bundle) {
|
||||
ImageCodecBundle b => b,
|
||||
String path => ImageCodecBundle(decoderGraphPath: path),
|
||||
_ => throw ArgumentError.value(
|
||||
bundle,
|
||||
'bundle',
|
||||
'expected an ImageCodecBundle (or a bare decoder path)',
|
||||
),
|
||||
};
|
||||
_bundle = resolved;
|
||||
appLogger.info(
|
||||
'ONNX codec bundle recorded: entropy=${resolved.entropyGraphPath != null} '
|
||||
'tables=${resolved.tablesPath != null} rate=${resolved.ratePoint.name}',
|
||||
tag: 'ImageCodec',
|
||||
);
|
||||
}
|
||||
|
||||
ImageCodecBundle _requireBundle() {
|
||||
final bundle = _bundle;
|
||||
if (bundle == null) {
|
||||
throw StateError('ONNX codec backend has not been loaded.');
|
||||
}
|
||||
return bundle;
|
||||
}
|
||||
|
||||
/// Creates the synthesis session if it is not already up.
|
||||
Future<OrtSession> _ensureDecoder() async {
|
||||
final existing = _decoder;
|
||||
if (existing != null) {
|
||||
return existing;
|
||||
}
|
||||
// No `providers:` is passed on purpose. Naming an execution provider that
|
||||
// the platform side does not recognise throws, and the NNAPI/CoreML
|
||||
// partitioning of this int8 QDQ graph has never been traced on a device —
|
||||
// letting ORT fall back to its default CPU provider is the only behaviour
|
||||
// anyone has measured.
|
||||
//
|
||||
// The path must be the `.onnx` graph, with its `.onnx.data`
|
||||
// external-weights sibling present in the same directory under exactly the
|
||||
// filename the graph records. See `ImageCodecModelSpec`.
|
||||
final session = await _runtime.createSession(_requireBundle().decoderGraphPath);
|
||||
_decoder = session;
|
||||
_inputName = _pick(session.inputNames, kImageCodecDecoderInputName, 'input');
|
||||
_outputName = _pick(
|
||||
session.outputNames,
|
||||
kImageCodecDecoderOutputName,
|
||||
'output',
|
||||
);
|
||||
appLogger.info(
|
||||
'ONNX decoder session ready: in=$_inputName out=$_outputName',
|
||||
tag: 'ImageCodec',
|
||||
);
|
||||
return session;
|
||||
}
|
||||
|
||||
/// Creates the entropy session this direction needs, plus the range coder.
|
||||
///
|
||||
/// ONE DIRECTION AT A TIME, deliberately. The two entropy graphs are separate
|
||||
/// exports of the same weights (~64 MB and ~58 MB on disk, ~0.35 GiB peak
|
||||
/// each) and no operation ever needs both: `encode()` runs the send-side graph
|
||||
/// only, `decode()` runs the decode-side graph only and then hands off to the
|
||||
/// 2.16 GiB synthesis session. Creating both would double the entropy-side
|
||||
/// cost for nothing.
|
||||
Future<AeicEntropyCodec> _ensureEntropy(
|
||||
int resolution, {
|
||||
required bool forDecode,
|
||||
}) async {
|
||||
final bundle = _requireBundle();
|
||||
final tablesPath = bundle.tablesPath;
|
||||
if (tablesPath == null) {
|
||||
throw const ImageCodecBundleIncomplete();
|
||||
}
|
||||
final builder = imageCodecRansCoderBuilder;
|
||||
if (builder == null) {
|
||||
throw const ImageCodecEntropyPathMissing(
|
||||
'the pure-Dart rANS coder is not wired into this build; set '
|
||||
'imageCodecRansCoderBuilder before loading the codec',
|
||||
);
|
||||
}
|
||||
|
||||
OrtSession? encodeSession;
|
||||
OrtSession? decodeSession;
|
||||
if (forDecode) {
|
||||
final path = bundle.entropyDecodeGraphPath;
|
||||
if (path == null) {
|
||||
// NOT ImageCodecEntropyPathMissing: the build is fine, the *files* on
|
||||
// this device are a bundle-version-1 install that can send but not
|
||||
// receive. The remedy is a re-download, so say so. Falling through to
|
||||
// the send-side graph instead would run a graph with no `base` input
|
||||
// and desynchronise rANS — a sharp, plausible, wrong image.
|
||||
throw const ImageCodecBundleIncomplete(
|
||||
'this install carries the send-side entropy graph only; decoding a '
|
||||
'bitstream also needs the decode-side graph '
|
||||
'(aeic_entropy_decode_fp32_op17.onnx). Re-download the image codec '
|
||||
'model.',
|
||||
);
|
||||
}
|
||||
decodeSession = _entropyDecode ??= await _runtime.createSession(path);
|
||||
} else {
|
||||
final path = bundle.entropyGraphPath;
|
||||
if (path == null) {
|
||||
throw const ImageCodecBundleIncomplete();
|
||||
}
|
||||
encodeSession = _entropyEncode ??= await _runtime.createSession(path);
|
||||
}
|
||||
|
||||
final coders = _coders ??= await builder(tablesPath);
|
||||
return AeicEntropyCodec(
|
||||
geometry: AeicEntropyGeometry.forResolution(resolution),
|
||||
network: OnnxAeicEntropyNetwork(
|
||||
encodeSession: encodeSession,
|
||||
decodeSession: decodeSession,
|
||||
),
|
||||
coders: coders,
|
||||
);
|
||||
}
|
||||
|
||||
/// Prefers the documented tensor name, tolerates a re-export that renamed a
|
||||
/// sole input/output, and refuses to guess between several.
|
||||
static String _pick(List<String> names, String preferred, String role) {
|
||||
if (names.contains(preferred)) {
|
||||
return preferred;
|
||||
}
|
||||
if (names.length == 1) {
|
||||
return names.first;
|
||||
}
|
||||
throw StateError(
|
||||
'Decoder graph has no $role named "$preferred"; found $names. '
|
||||
'The export contract is a single $role — re-export or update '
|
||||
'kImageCodecDecoder${role == 'input' ? 'Input' : 'Output'}Name.',
|
||||
);
|
||||
}
|
||||
|
||||
@override
|
||||
Future<Uint8List> decodeLatentToRgb({
|
||||
required Float32List yHat,
|
||||
void Function(double progress)? onProgress,
|
||||
bool Function()? shouldCancel,
|
||||
}) async {
|
||||
if (yHat.length != kImageCodecLatentElements) {
|
||||
throw ArgumentError.value(
|
||||
yHat.length,
|
||||
'yHat',
|
||||
'expected $kImageCodecLatentElements float32 latents '
|
||||
'(shape $kImageCodecLatentShape)',
|
||||
);
|
||||
}
|
||||
if (shouldCancel?.call() == true) {
|
||||
throw const ImageCodecCancelled();
|
||||
}
|
||||
|
||||
onProgress?.call(0.02);
|
||||
final session = await _ensureDecoder();
|
||||
OrtValue? input;
|
||||
final outputs = <OrtValue>[];
|
||||
try {
|
||||
input = await OrtValue.fromList(yHat, kImageCodecLatentShape);
|
||||
onProgress?.call(0.05);
|
||||
if (shouldCancel?.call() == true) {
|
||||
throw const ImageCodecCancelled();
|
||||
}
|
||||
|
||||
// The single long blocking stage. `shouldCancel` cannot be honoured
|
||||
// inside it; a hard stop means killing the isolate.
|
||||
final result = await session.run(<String, OrtValue>{_inputName: input});
|
||||
outputs.addAll(result.values);
|
||||
onProgress?.call(0.85);
|
||||
|
||||
final output = result[_outputName] ?? result.values.first;
|
||||
final flat = await output.asFlattenedList();
|
||||
onProgress?.call(0.95);
|
||||
final rgb = _chwFloatsToRgb(flat, output.shape);
|
||||
onProgress?.call(1.0);
|
||||
return rgb;
|
||||
} finally {
|
||||
await input?.dispose();
|
||||
for (final value in outputs) {
|
||||
await value.dispose();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `[1, 3, H, W]` floats in [-1, 1] -> packed 8-bit RGB, `H * W * 3` bytes.
|
||||
static Uint8List _chwFloatsToRgb(List<dynamic> flat, List<int> shape) {
|
||||
final expectedSide = kImageCodecSquareSize;
|
||||
final pixels = expectedSide * expectedSide;
|
||||
final expected = pixels * _bytesPerPixel;
|
||||
if (flat.length != expected) {
|
||||
throw StateError(
|
||||
'Decoder returned ${flat.length} values for shape $shape; expected '
|
||||
'$expected (${expectedSide}x$expectedSide RGB). The export is static '
|
||||
'512x512 — a different shape means the wrong model file.',
|
||||
);
|
||||
}
|
||||
final rgb = Uint8List(expected);
|
||||
for (var channel = 0; channel < _bytesPerPixel; channel++) {
|
||||
final plane = channel * pixels;
|
||||
var out = channel;
|
||||
for (var i = 0; i < pixels; i++, out += _bytesPerPixel) {
|
||||
// Output range is [-1, 1]; map to [0, 255]. Values outside the range do
|
||||
// occur (the last conv is unbounded) and must be clamped, not wrapped.
|
||||
final scaled = ((flat[plane + i] as num).toDouble() + 1.0) * 127.5;
|
||||
rgb[out] = scaled <= 0
|
||||
? 0
|
||||
: scaled >= 255
|
||||
? 255
|
||||
: scaled.round();
|
||||
}
|
||||
}
|
||||
return rgb;
|
||||
}
|
||||
|
||||
@override
|
||||
Future<Uint8List> encode({
|
||||
required Uint8List rgbBytes,
|
||||
required AeicRatePoint ratePoint,
|
||||
required int resolution,
|
||||
void Function(double progress)? onProgress,
|
||||
bool Function()? shouldCancel,
|
||||
}) async {
|
||||
_checkRatePoint(ratePoint);
|
||||
final codec = await _ensureEntropy(resolution, forDecode: false);
|
||||
try {
|
||||
// The entropy session stays up on purpose: a second photo is free.
|
||||
return await codec.encode(
|
||||
rgbBytes,
|
||||
onProgress: onProgress,
|
||||
shouldCancel: shouldCancel,
|
||||
);
|
||||
} on AeicEntropyCancelled {
|
||||
throw const ImageCodecCancelled();
|
||||
}
|
||||
}
|
||||
|
||||
@override
|
||||
Future<Uint8List> decode({
|
||||
required Uint8List bitstream,
|
||||
required AeicRatePoint ratePoint,
|
||||
required int resolution,
|
||||
void Function(double progress)? onProgress,
|
||||
bool Function()? shouldCancel,
|
||||
}) async {
|
||||
_checkRatePoint(ratePoint);
|
||||
final codec = await _ensureEntropy(resolution, forDecode: true);
|
||||
final Float32List yHat;
|
||||
try {
|
||||
yHat = await codec.decodeToLatent(
|
||||
bitstream,
|
||||
onProgress: (value) => onProgress?.call(value * 0.5),
|
||||
shouldCancel: shouldCancel,
|
||||
);
|
||||
} on AeicEntropyCancelled {
|
||||
throw const ImageCodecCancelled();
|
||||
}
|
||||
// MANDATORY, not an optimisation: never hold the fp32 entropy graph and the
|
||||
// 2.16 GiB synthesis session at the same time.
|
||||
await releaseEntropySession();
|
||||
return decodeLatentToRgb(
|
||||
yHat: yHat,
|
||||
onProgress: (value) => onProgress?.call(0.5 + value * 0.5),
|
||||
shouldCancel: shouldCancel,
|
||||
);
|
||||
}
|
||||
|
||||
/// The graphs and the CDF tables belong to one checkpoint. A bitstream that
|
||||
/// claims a different rate point cannot be decoded with these tables, and
|
||||
/// would decode into a sharp, plausible, wrong image if it were tried.
|
||||
void _checkRatePoint(AeicRatePoint ratePoint) {
|
||||
final expected = _requireBundle().ratePoint;
|
||||
if (ratePoint != expected) {
|
||||
throw ImageCodecUnimplemented(
|
||||
'this bundle is ${expected.name}; the bitstream claims '
|
||||
'${ratePoint.name}. Rate points are not interchangeable: the CDF '
|
||||
'tables and the entropy graph belong to one checkpoint.',
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@override
|
||||
Future<void> releaseDecoderSession() async {
|
||||
final session = _decoder;
|
||||
_decoder = null;
|
||||
await _close(session, 'decoder');
|
||||
}
|
||||
|
||||
/// Drops **both** entropy graphs. Callers ask for "the entropy half"; which
|
||||
/// direction happens to be resident is an implementation detail, and a decode
|
||||
/// that is about to create the 2.16 GiB synthesis session must not leave the
|
||||
/// send-side graph behind just because it never used it.
|
||||
@override
|
||||
Future<void> releaseEntropySession() async {
|
||||
final encodeSession = _entropyEncode;
|
||||
final decodeSession = _entropyDecode;
|
||||
_entropyEncode = null;
|
||||
_entropyDecode = null;
|
||||
await _close(encodeSession, 'entropy(encode)');
|
||||
await _close(decodeSession, 'entropy(decode)');
|
||||
}
|
||||
|
||||
@override
|
||||
Future<void> dispose() async {
|
||||
await releaseDecoderSession();
|
||||
await releaseEntropySession();
|
||||
_coders = null;
|
||||
}
|
||||
|
||||
static Future<void> _close(OrtSession? session, String which) async {
|
||||
if (session == null) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
await session.close();
|
||||
} catch (error) {
|
||||
// Never let teardown throw: it is called from memory-pressure and from
|
||||
// isolate shutdown, where there is nobody left to handle it.
|
||||
appLogger.warn(
|
||||
'ONNX $which session close failed: $error',
|
||||
tag: 'ImageCodec',
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// [AeicEntropyNetwork] backed by the fp32 entropy ONNX graphs.
|
||||
///
|
||||
/// ## TWO GRAPHS, NOT ONE. This is not a packaging accident.
|
||||
///
|
||||
/// The send-side export (`aeic_entropy_side_fp32_op17.onnx`, 64 MB) has one
|
||||
/// input `image [1,3,512,512]` and emits `z_q`, `yq0..3`, `sc0..3` in a single
|
||||
/// run. That shape is only valid for an encoder, which already knows `y`.
|
||||
/// Decoding is inherently SEQUENTIAL: the symbols of stage `i` must be decoded
|
||||
/// before stage `i+1`'s context can be computed, so a receiver must be able to
|
||||
/// call the network five times per image, interleaved with rANS.
|
||||
///
|
||||
/// The decode-side export (`aeic_entropy_decode_fp32_op17.onnx`, 58 MB) is that
|
||||
/// callable form — the same sub-networks behind an ONNX `If` on a `stage`
|
||||
/// selector. `flutter_onnxruntime`'s `session.run()` fetches ALL graph outputs
|
||||
/// (there is no output-subset API), so the `If` is what stops each call from
|
||||
/// evaluating `g_c` four times.
|
||||
///
|
||||
/// ### Decode-side contract (measured, not assumed)
|
||||
///
|
||||
/// ```
|
||||
/// inputs : z_q float32 [1,128,4,4] base float32 [1,256,16,16]
|
||||
/// stage int32 [1]
|
||||
/// outputs: base0 float32 [1,256,16,16] means float32 [1,256,16,16]
|
||||
/// scales float32 [1,256,16,16]
|
||||
/// ```
|
||||
///
|
||||
/// * ALL THREE INPUTS ARE REQUIRED ON EVERY RUN. ORT rejects a feed that omits
|
||||
/// one ("Required inputs (['base']) are missing from input feed"), so the
|
||||
/// branch that does not read a tensor is still handed a zero filler. That is
|
||||
/// what [_zeroBase] and [_zeroZq] are for; they are not defensive padding.
|
||||
/// * `stage < 0` takes the HYPER branch: `base0 = h_s(z_q + z_offset)`, with
|
||||
/// `z_offset` (the entropy-bottleneck medians) baked into the graph as a
|
||||
/// constant. The input is therefore raw `z_q`, NOT `z_hat` — adding the
|
||||
/// offset in Dart would add it twice.
|
||||
/// * `stage 0..3` takes the CONTEXT branch:
|
||||
/// `adapter_out[stage](g_c(adapter_in[stage](base)))`, chunked into
|
||||
/// `means`/`scales` and **already multiplied by that stage's mask**.
|
||||
/// `AeicMaskSet.applyMask` is a select with the identical mask, so applying it
|
||||
/// again in `image_codec_entropy.dart` is exactly idempotent — that call site
|
||||
/// needs no change, and removing it would couple the arithmetic to this
|
||||
/// export.
|
||||
/// * The branch not taken emits a zero tensor of the same shape, so a caller
|
||||
/// must read only the output its `stage` selects.
|
||||
/// * Latency on an M4 Max, 1 thread, CPU EP, ORT 1.28.0: 6.11 ms for the hyper
|
||||
/// branch and 3.70 ms per stage, ~21 ms per image.
|
||||
///
|
||||
/// Geometry is fixed 512x512 -> `y` 16x16, `z` 4x4. No dynamic axes.
|
||||
///
|
||||
/// A network constructed with only [encodeSession] reports
|
||||
/// [supportsDecodeSide] false, and `AeicEntropyCodec.decodeToLatent` refuses
|
||||
/// with [AeicEntropyUnavailable] rather than feeding a send-side graph inputs it
|
||||
/// does not have.
|
||||
class OnnxAeicEntropyNetwork implements AeicEntropyNetwork {
|
||||
static const String kImageInput = 'image';
|
||||
static const String kZqInput = 'z_q';
|
||||
static const String kBaseInput = 'base';
|
||||
static const String kStageInput = 'stage';
|
||||
static const String kBase0Output = 'base0';
|
||||
static const String kMeansOutput = 'means';
|
||||
static const String kScalesOutput = 'scales';
|
||||
|
||||
/// The `stage` value that selects the hyper-synthesis branch. Any negative
|
||||
/// value works; -1 is the documented one.
|
||||
static const int kHyperStage = -1;
|
||||
|
||||
static const List<int> kZqShape = <int>[1, 128, 4, 4];
|
||||
static const List<int> kBaseShape = <int>[1, 256, 16, 16];
|
||||
static const int kZqElements = 128 * 4 * 4;
|
||||
static const int kBaseElements = 256 * 16 * 16;
|
||||
|
||||
/// Session over the send-side graph, or null on a decode-only network.
|
||||
final OrtSession? encodeSession;
|
||||
|
||||
/// Session over the decode-side graph, or null on an encode-only network.
|
||||
final OrtSession? decodeSession;
|
||||
|
||||
const OnnxAeicEntropyNetwork({this.encodeSession, this.decodeSession});
|
||||
|
||||
@override
|
||||
bool get supportsDecodeSide {
|
||||
final session = decodeSession;
|
||||
if (session == null) {
|
||||
return false;
|
||||
}
|
||||
final inputs = session.inputNames;
|
||||
final outputs = session.outputNames;
|
||||
return inputs.contains(kZqInput) &&
|
||||
inputs.contains(kBaseInput) &&
|
||||
inputs.contains(kStageInput) &&
|
||||
outputs.contains(kBase0Output) &&
|
||||
outputs.contains(kMeansOutput) &&
|
||||
outputs.contains(kScalesOutput);
|
||||
}
|
||||
|
||||
@override
|
||||
Future<AeicEncodeSideTensors> runEncodeSide(Float32List imageChw) async {
|
||||
final session = encodeSession;
|
||||
if (session == null) {
|
||||
throw const AeicEntropyUnavailable(
|
||||
'this network has no send-side entropy session, so it cannot encode',
|
||||
);
|
||||
}
|
||||
if (!session.inputNames.contains(kImageInput)) {
|
||||
throw const AeicEntropyUnavailable(
|
||||
'the entropy graph has no "image" input, so it cannot encode',
|
||||
);
|
||||
}
|
||||
final side = imageChw.length ~/ 3;
|
||||
final resolution = _isqrt(side);
|
||||
final inputs = <String, OrtValue>{
|
||||
kImageInput: await OrtValue.fromList(imageChw, <int>[
|
||||
1,
|
||||
3,
|
||||
resolution,
|
||||
resolution,
|
||||
]),
|
||||
};
|
||||
final result = await _run(session, inputs);
|
||||
try {
|
||||
return AeicEncodeSideTensors(
|
||||
zQ: await _floats(result, 'z_q'),
|
||||
yQ: <Float32List>[
|
||||
for (var i = 0; i < 4; i++) await _floats(result, 'yq$i'),
|
||||
],
|
||||
scales: <Float32List>[
|
||||
for (var i = 0; i < 4; i++) await _floats(result, 'sc$i'),
|
||||
],
|
||||
);
|
||||
} finally {
|
||||
await _disposeAll(result);
|
||||
}
|
||||
}
|
||||
|
||||
@override
|
||||
Future<Float32List> runHyperSynthesis(Float32List zQ) async {
|
||||
final session = _requireDecodeSide();
|
||||
if (zQ.length != kZqElements) {
|
||||
throw ArgumentError.value(
|
||||
zQ.length,
|
||||
'zQ',
|
||||
'expected $kZqElements float32 values (shape $kZqShape)',
|
||||
);
|
||||
}
|
||||
// `base` is not read by the hyper branch, but ORT requires every declared
|
||||
// input to be fed. Zeros, not a stale tensor: a stale one would be silently
|
||||
// wrong the day the export starts reading it.
|
||||
final result = await _run(session, <String, OrtValue>{
|
||||
kZqInput: await OrtValue.fromList(zQ, kZqShape),
|
||||
kBaseInput: await OrtValue.fromList(
|
||||
Float32List(kBaseElements),
|
||||
kBaseShape,
|
||||
),
|
||||
kStageInput: await OrtValue.fromList(Int32List.fromList(<int>[
|
||||
kHyperStage,
|
||||
]), <int>[1]),
|
||||
});
|
||||
try {
|
||||
return await _floats(result, kBase0Output);
|
||||
} finally {
|
||||
await _disposeAll(result);
|
||||
}
|
||||
}
|
||||
|
||||
@override
|
||||
Future<AeicStageParams> runStage(int stage, Float32List base) async {
|
||||
final session = _requireDecodeSide();
|
||||
if (stage < 0 || stage > 3) {
|
||||
// The graph does not validate `stage`: anything >= 4 silently falls into
|
||||
// the stage-3 branch and anything < 0 runs hyper synthesis, either of
|
||||
// which desynchronises rANS without an error. Catch it here instead.
|
||||
throw ArgumentError.value(stage, 'stage', 'must be 0..3');
|
||||
}
|
||||
if (base.length != kBaseElements) {
|
||||
throw ArgumentError.value(
|
||||
base.length,
|
||||
'base',
|
||||
'expected $kBaseElements float32 values (shape $kBaseShape)',
|
||||
);
|
||||
}
|
||||
// `z_q` is not read by the context branch; same reasoning as above.
|
||||
final result = await _run(session, <String, OrtValue>{
|
||||
kZqInput: await OrtValue.fromList(Float32List(kZqElements), kZqShape),
|
||||
kBaseInput: await OrtValue.fromList(base, kBaseShape),
|
||||
kStageInput: await OrtValue.fromList(Int32List.fromList(<int>[
|
||||
stage,
|
||||
]), <int>[1]),
|
||||
});
|
||||
try {
|
||||
// Already masked by the graph. `AeicEntropyCodec` masks again, which is
|
||||
// idempotent — see the class doc.
|
||||
return AeicStageParams(
|
||||
meansSupp: await _floats(result, kMeansOutput),
|
||||
scalesSupp: await _floats(result, kScalesOutput),
|
||||
);
|
||||
} finally {
|
||||
await _disposeAll(result);
|
||||
}
|
||||
}
|
||||
|
||||
OrtSession _requireDecodeSide() {
|
||||
final session = decodeSession;
|
||||
if (session == null) {
|
||||
throw const AeicEntropyUnavailable(
|
||||
'this network has no decode-side entropy session; decoding needs '
|
||||
'aeic_entropy_decode_fp32_op17.onnx, which a bundle-version-1 install '
|
||||
'does not carry',
|
||||
);
|
||||
}
|
||||
if (!supportsDecodeSide) {
|
||||
throw AeicEntropyUnavailable(
|
||||
'the loaded decode-side graph does not match the contract (inputs '
|
||||
'${session.inputNames}, outputs ${session.outputNames}); decoding '
|
||||
'needs inputs [$kZqInput, $kBaseInput, $kStageInput] and outputs '
|
||||
'[$kBase0Output, $kMeansOutput, $kScalesOutput]',
|
||||
);
|
||||
}
|
||||
return session;
|
||||
}
|
||||
|
||||
Future<Map<String, OrtValue>> _run(
|
||||
OrtSession session,
|
||||
Map<String, OrtValue> inputs,
|
||||
) async {
|
||||
try {
|
||||
return await session.run(inputs);
|
||||
} finally {
|
||||
for (final value in inputs.values) {
|
||||
await value.dispose();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static Future<Float32List> _floats(
|
||||
Map<String, OrtValue> result,
|
||||
String name,
|
||||
) async {
|
||||
final value = result[name];
|
||||
if (value == null) {
|
||||
throw AeicEntropyUnavailable(
|
||||
'the entropy graph has no output named "$name"; found '
|
||||
'${result.keys.toList()}',
|
||||
);
|
||||
}
|
||||
final flat = await value.asFlattenedList();
|
||||
final out = Float32List(flat.length);
|
||||
for (var i = 0; i < flat.length; i++) {
|
||||
out[i] = (flat[i] as num).toDouble();
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
static Future<void> _disposeAll(Map<String, OrtValue> values) async {
|
||||
for (final value in values.values) {
|
||||
await value.dispose();
|
||||
}
|
||||
}
|
||||
|
||||
static int _isqrt(int value) {
|
||||
final root = math.sqrt(value).round();
|
||||
if (root * root != value) {
|
||||
throw ArgumentError.value(value, 'value', 'not a square image');
|
||||
}
|
||||
return root;
|
||||
}
|
||||
}
|
||||
|
||||
/// Thrown by a backend that noticed [ImageCodecBackend] `shouldCancel`.
|
||||
class ImageCodecCancelled implements Exception {
|
||||
const ImageCodecCancelled();
|
||||
|
||||
@override
|
||||
String toString() => 'Image processing stopped.';
|
||||
}
|
||||
|
||||
/// Factory used by the codec worker isolate.
|
||||
///
|
||||
/// Returns `null` on web, where there is no local model file and no isolate.
|
||||
/// Everywhere else it returns the real [OnnxImageCodecBackend]; the session
|
||||
/// turns a `null` here into an [ImageCodecUnimplemented], which the service
|
||||
/// surfaces as `lastError` and maps to `ImageCodecAvailability.unavailable`.
|
||||
ImageCodecBackend? createImageCodecBackend() {
|
||||
if (kIsWeb) {
|
||||
return null;
|
||||
}
|
||||
return OnnxImageCodecBackend();
|
||||
}
|
||||
@@ -0,0 +1,675 @@
|
||||
/// ===========================================================================
|
||||
/// THE ENTROPY LAYER — everything between the ONNX tensors and the rANS coder.
|
||||
/// ===========================================================================
|
||||
///
|
||||
/// AEIC's four-stage masked context model, ported from
|
||||
/// `aic/aeic/src/codec/codec_practical.py` (`PixelCodec.compress` /
|
||||
/// `.decompress`). This file owns the parts that are *not* neural and *not*
|
||||
/// range coding:
|
||||
///
|
||||
/// * the four checkerboard masks (`get_mask_four_parts`),
|
||||
/// * `sequeeze` / `unsequeeze_with_mask` (the 4:1 channel fold),
|
||||
/// * `torch.round`'s half-to-**even** tie rule,
|
||||
/// * `my_build_indexes` (scale -> CDF row selector),
|
||||
/// * the z index array (`_build_indexes`: the channel arange),
|
||||
/// * the fixed call order z, y0, y1, y2, y3.
|
||||
///
|
||||
/// It is deliberately pure Dart: no ONNX, no dart:io, no Flutter. The neural
|
||||
/// half arrives through [AeicEntropyNetwork] and the range coder through
|
||||
/// [AeicRansCoderFactory], both injected. That is what makes the arithmetic
|
||||
/// unit-testable without a 67 MB graph or a model download.
|
||||
///
|
||||
/// ## Why every detail here is load-bearing
|
||||
///
|
||||
/// The rANS decoder re-derives symbol probabilities by re-running the same
|
||||
/// network and the same arithmetic. A single wrong symbol *position* — an
|
||||
/// off-by-one in a mask, the wrong channel-group permutation, `sequeeze`
|
||||
/// folding in the wrong order — does not raise: it desynchronises the coder and
|
||||
/// produces a sharp, plausible, wrong image. Same for a rounding tie resolved
|
||||
/// away from zero instead of to even. Both failure modes are silent.
|
||||
///
|
||||
/// The golden vectors in `test/services/image_codec_entropy_test.dart` were
|
||||
/// generated by `aic/exp/export_entropy_layer_golden.py`, which runs the real
|
||||
/// torch ops on a closed-form input both languages reproduce bit-for-bit.
|
||||
library;
|
||||
|
||||
import 'dart:math' as math;
|
||||
import 'dart:typed_data';
|
||||
|
||||
/// CDF group indices, fixed by `EntropyCoder.update()`'s `add_cdf` call order:
|
||||
/// the entropy bottleneck (z) is registered first, the Gaussian conditional
|
||||
/// (y) second. The group index is part of the bitstream format.
|
||||
const int kAeicZCdfGroup = 0;
|
||||
const int kAeicYCdfGroup = 1;
|
||||
|
||||
/// `my_build_indexes` constants (`codec_practical.py`, `SCALES_MIN/MAX/LEVELS`).
|
||||
const int kAeicScalesLevels = 64;
|
||||
const double kAeicLogScaleMin = -2.2072749131897207; // ln(0.11)
|
||||
const double kAeicLogScaleStep = 0.12305479932808384; // (ln(256)-ln(0.11))/63
|
||||
const double kAeicScaleThreshold = 0.08; // below this the symbol is skipped
|
||||
const double kAeicScaleFloor = 1e-5; // torch.maximum floor applied first
|
||||
|
||||
/// Tensor shapes for one image, derived the way `compress()` derives them.
|
||||
///
|
||||
/// `y` is the image downsampled by 32; `z` is `y` downsampled by a further 4,
|
||||
/// rounded **up** (`compress()` reflect-pads `y` to a multiple of 4 first).
|
||||
class AeicEntropyGeometry {
|
||||
final int resolution;
|
||||
|
||||
/// Latent channel count `M` (256 for AEIC-SE, which is what ships).
|
||||
final int yChannels;
|
||||
final int yHeight;
|
||||
final int yWidth;
|
||||
|
||||
/// Hyper-latent channels: `M // 2`.
|
||||
final int zChannels;
|
||||
final int zHeight;
|
||||
final int zWidth;
|
||||
|
||||
const AeicEntropyGeometry._({
|
||||
required this.resolution,
|
||||
required this.yChannels,
|
||||
required this.yHeight,
|
||||
required this.yWidth,
|
||||
required this.zChannels,
|
||||
required this.zHeight,
|
||||
required this.zWidth,
|
||||
});
|
||||
|
||||
factory AeicEntropyGeometry.forResolution(
|
||||
int resolution, {
|
||||
int yChannels = 256,
|
||||
}) {
|
||||
if (resolution <= 0 || resolution % 32 != 0) {
|
||||
throw ArgumentError.value(
|
||||
resolution,
|
||||
'resolution',
|
||||
'must be a positive multiple of 32 (g_a downsamples by 32)',
|
||||
);
|
||||
}
|
||||
if (yChannels % 4 != 0) {
|
||||
throw ArgumentError.value(
|
||||
yChannels,
|
||||
'yChannels',
|
||||
'the four-part mask splits the channels into 4 equal groups',
|
||||
);
|
||||
}
|
||||
final yh = resolution ~/ 32;
|
||||
final zh = (yh + 3) ~/ 4;
|
||||
return AeicEntropyGeometry._(
|
||||
resolution: resolution,
|
||||
yChannels: yChannels,
|
||||
yHeight: yh,
|
||||
yWidth: yh,
|
||||
zChannels: yChannels ~/ 2,
|
||||
zHeight: zh,
|
||||
zWidth: zh,
|
||||
);
|
||||
}
|
||||
|
||||
/// Channels after `sequeeze` folds the four groups together.
|
||||
int get squeezedChannels => yChannels ~/ 4;
|
||||
|
||||
int get yElements => yChannels * yHeight * yWidth;
|
||||
|
||||
/// Symbols coded per y stage (`[1, M/4, y_h, y_w]` flattened).
|
||||
int get symbolsPerStage => squeezedChannels * yHeight * yWidth;
|
||||
|
||||
int get zElements => zChannels * zHeight * zWidth;
|
||||
|
||||
/// Total entries handed to the coder: z once, then four y stages.
|
||||
int get totalEntries => zElements + 4 * symbolsPerStage;
|
||||
|
||||
List<int> get yShape => <int>[1, yChannels, yHeight, yWidth];
|
||||
|
||||
List<int> get zShape => <int>[1, zChannels, zHeight, zWidth];
|
||||
|
||||
List<int> get imageShape => <int>[1, 3, resolution, resolution];
|
||||
}
|
||||
|
||||
/// The four masks of `get_mask_four_parts`, as arithmetic rather than tensors.
|
||||
///
|
||||
/// Reading the Python: each mask is four channel-groups of `M/4` channels
|
||||
/// stacked, and each group carries one of the four 2x2 micro-patterns
|
||||
///
|
||||
/// micro 0 = (y%2, x%2) == (0, 0) micro 1 = (0, 1)
|
||||
/// micro 2 = (1, 0) micro 3 = (1, 1)
|
||||
///
|
||||
/// so `micro k` is live exactly where `(y%2)*2 + (x%2) == k`. The stacking order
|
||||
/// per stage is
|
||||
///
|
||||
/// mask_0 = [m0, m1, m2, m3] mask_1 = [m3, m2, m1, m0]
|
||||
/// mask_2 = [m2, m3, m0, m1] mask_3 = [m1, m0, m3, m2]
|
||||
///
|
||||
/// which is exactly `micro = group XOR perm[stage]` with `perm = [0, 3, 2, 1]`.
|
||||
/// Verified against the tensors themselves in the golden test, because getting
|
||||
/// this permutation wrong is the single easiest way to silently reorder every
|
||||
/// symbol in the stream.
|
||||
class AeicMaskSet {
|
||||
/// `perm[stage]`: XOR it with the channel group to get the micro-pattern.
|
||||
static const List<int> stagePermutation = <int>[0, 3, 2, 1];
|
||||
|
||||
final AeicEntropyGeometry geometry;
|
||||
|
||||
const AeicMaskSet(this.geometry);
|
||||
|
||||
/// Micro-pattern index used by [channelGroup] in [stage].
|
||||
static int microFor(int stage, int channelGroup) =>
|
||||
channelGroup ^ stagePermutation[stage];
|
||||
|
||||
/// The single channel group that is live at `(y, x)` in [stage].
|
||||
///
|
||||
/// Inverse of [microFor]: the position selects the micro-pattern, and exactly
|
||||
/// one group carries it.
|
||||
static int liveGroupAt(int stage, int y, int x) =>
|
||||
(((y & 1) << 1) | (x & 1)) ^ stagePermutation[stage];
|
||||
|
||||
bool isLive(int stage, int channel, int y, int x) {
|
||||
final group = channel ~/ geometry.squeezedChannels;
|
||||
return microFor(stage, group) == (((y & 1) << 1) | (x & 1));
|
||||
}
|
||||
|
||||
/// Materialises `mask_i` as `[1, M, y_h, y_w]` float32 — only needed when a
|
||||
/// caller wants to hand the mask to something else (or check it in a test).
|
||||
Float32List maskTensor(int stage) {
|
||||
final out = Float32List(geometry.yElements);
|
||||
final h = geometry.yHeight;
|
||||
final w = geometry.yWidth;
|
||||
for (var c = 0; c < geometry.yChannels; c++) {
|
||||
final base = c * h * w;
|
||||
for (var y = 0; y < h; y++) {
|
||||
for (var x = 0; x < w; x++) {
|
||||
if (isLive(stage, c, y, x)) {
|
||||
out[base + y * w + x] = 1.0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/// `sequeeze`: sum the four channel chunks, `(g0 + g1) + (g2 + g3)`.
|
||||
///
|
||||
/// The association matters in principle (float addition is not associative)
|
||||
/// so it is reproduced exactly, even though in practice three of the four
|
||||
/// terms are a hard zero for a masked tensor.
|
||||
Float32List squeeze(Float32List full) {
|
||||
if (full.length != geometry.yElements) {
|
||||
throw ArgumentError.value(
|
||||
full.length,
|
||||
'full',
|
||||
'expected ${geometry.yElements} elements (${geometry.yShape})',
|
||||
);
|
||||
}
|
||||
final plane = geometry.yHeight * geometry.yWidth;
|
||||
final cq = geometry.squeezedChannels;
|
||||
final out = Float32List(cq * plane);
|
||||
for (var c = 0; c < cq; c++) {
|
||||
final o = c * plane;
|
||||
final g0 = c * plane;
|
||||
final g1 = (c + cq) * plane;
|
||||
final g2 = (c + 2 * cq) * plane;
|
||||
final g3 = (c + 3 * cq) * plane;
|
||||
for (var i = 0; i < plane; i++) {
|
||||
// Each partial sum is forced back through float32, matching torch's
|
||||
// `(a + b) + (c + d)` on float32 tensors.
|
||||
final left = f32(full[g0 + i] + full[g1 + i]);
|
||||
final right = f32(full[g2 + i] + full[g3 + i]);
|
||||
out[o + i] = left + right;
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/// `unsequeeze_with_mask`: broadcast a squeezed tensor back to `[1, M, h, w]`,
|
||||
/// keeping each element only in the channel group whose mask is live there.
|
||||
Float32List unsqueeze(Float32List squeezed, int stage) {
|
||||
final plane = geometry.yHeight * geometry.yWidth;
|
||||
final cq = geometry.squeezedChannels;
|
||||
if (squeezed.length != cq * plane) {
|
||||
throw ArgumentError.value(
|
||||
squeezed.length,
|
||||
'squeezed',
|
||||
'expected ${cq * plane} elements',
|
||||
);
|
||||
}
|
||||
final out = Float32List(geometry.yElements);
|
||||
final w = geometry.yWidth;
|
||||
for (var c = 0; c < cq; c++) {
|
||||
for (var y = 0; y < geometry.yHeight; y++) {
|
||||
for (var x = 0; x < w; x++) {
|
||||
final group = liveGroupAt(stage, y, x);
|
||||
out[(group * cq + c) * plane + y * w + x] = squeezed[c * plane +
|
||||
y * w +
|
||||
x];
|
||||
}
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/// `t * mask_i`, elementwise, into a fresh tensor.
|
||||
Float32List applyMask(Float32List full, int stage) {
|
||||
final out = Float32List(full.length);
|
||||
final plane = geometry.yHeight * geometry.yWidth;
|
||||
final cq = geometry.squeezedChannels;
|
||||
final w = geometry.yWidth;
|
||||
for (var c = 0; c < geometry.yChannels; c++) {
|
||||
final group = c ~/ cq;
|
||||
for (var y = 0; y < geometry.yHeight; y++) {
|
||||
for (var x = 0; x < w; x++) {
|
||||
if (microFor(stage, group) == (((y & 1) << 1) | (x & 1))) {
|
||||
final i = c * plane + y * w + x;
|
||||
out[i] = full[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/// `base * (1 - mask_i) + stageLatent`, the context update between stages.
|
||||
///
|
||||
/// `stageLatent` must already be masked (it is the output of [unsqueeze]), so
|
||||
/// this is a select, not an add — which is what keeps it exact in float32.
|
||||
Float32List mergeContext(Float32List base, Float32List stageLatent, int stage) {
|
||||
final out = Float32List.fromList(base);
|
||||
final plane = geometry.yHeight * geometry.yWidth;
|
||||
final cq = geometry.squeezedChannels;
|
||||
final w = geometry.yWidth;
|
||||
for (var c = 0; c < geometry.yChannels; c++) {
|
||||
final group = c ~/ cq;
|
||||
for (var y = 0; y < geometry.yHeight; y++) {
|
||||
for (var x = 0; x < w; x++) {
|
||||
if (microFor(stage, group) == (((y & 1) << 1) | (x & 1))) {
|
||||
final i = c * plane + y * w + x;
|
||||
out[i] = stageLatent[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
}
|
||||
|
||||
/// Scratch cell used to force a float64 value through float32 rounding.
|
||||
final Float32List _f32Cell = Float32List(1);
|
||||
|
||||
/// Rounds [value] to float32 exactly as a store into a float32 tensor would.
|
||||
double f32(double value) {
|
||||
_f32Cell[0] = value;
|
||||
return _f32Cell[0];
|
||||
}
|
||||
|
||||
/// `torch.round`: ties go to the **even** neighbour, not away from zero.
|
||||
///
|
||||
/// Dart's `double.round()` and `roundToDouble()` both round half away from
|
||||
/// zero, so `(-0.5).round() == -1` where torch gives `0`. On the ft32 corpus
|
||||
/// `y - means` lands on a tie often enough that this is not theoretical.
|
||||
double roundHalfToEven(double value) {
|
||||
if (!value.isFinite) {
|
||||
return value;
|
||||
}
|
||||
final floor = value.floorToDouble();
|
||||
final diff = value - floor;
|
||||
if (diff > 0.5) {
|
||||
return floor + 1.0;
|
||||
}
|
||||
if (diff < 0.5) {
|
||||
return floor;
|
||||
}
|
||||
// Exactly halfway: pick the even neighbour.
|
||||
return floor % 2.0 == 0.0 ? floor : floor + 1.0;
|
||||
}
|
||||
|
||||
/// `my_build_indexes` — scale -> CDF row selector, in float32, per stage.
|
||||
///
|
||||
/// ```
|
||||
/// s = max(scale, 1e-5)
|
||||
/// q = (ln(s) - log_scale_min) / log_scale_step # all float32
|
||||
/// idx = (s < 0.08) ? -1 : trunc(clamp(q, 0, 63)) # trunc, not round
|
||||
/// ```
|
||||
///
|
||||
/// Every intermediate is forced back to float32 because torch evaluates this on
|
||||
/// a float32 tensor with float32 scalars, and the boundary margin measured in
|
||||
/// `aic/results/bitexact_encoder.md` is only 1.12x — the tightest number in the
|
||||
/// whole system. Dart's `log` is still not guaranteed bit-identical to ORT's
|
||||
/// float32 `Log`; see the note in `aic/results/rans_port_spec.md` §8, which
|
||||
/// recommends the graph emit these indexes itself. [aeicBuildIndexes] is the
|
||||
/// fallback, and both sides of the channel run this same code, so sender and
|
||||
/// receiver agree with each other regardless.
|
||||
Int16List aeicBuildIndexes(Float32List scales) {
|
||||
final logMin = f32(kAeicLogScaleMin);
|
||||
final logStep = f32(kAeicLogScaleStep);
|
||||
final floor = f32(kAeicScaleFloor);
|
||||
final threshold = f32(kAeicScaleThreshold);
|
||||
final out = Int16List(scales.length);
|
||||
for (var i = 0; i < scales.length; i++) {
|
||||
final raw = scales[i];
|
||||
final s = raw > floor ? raw : floor;
|
||||
if (s < threshold) {
|
||||
out[i] = -1;
|
||||
continue;
|
||||
}
|
||||
var q = f32(f32(f32(math.log(s)) - logMin) / logStep);
|
||||
if (q < 0.0) {
|
||||
q = 0.0;
|
||||
} else if (q > kAeicScalesLevels - 1) {
|
||||
q = (kAeicScalesLevels - 1).toDouble();
|
||||
}
|
||||
out[i] = q.toInt(); // truncation toward zero, matching Tensor.int()
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/// The z index array: `_build_indexes` broadcasts `arange(C)` over `H x W`, so
|
||||
/// this is `H*W` copies of each channel index. Always >= 0.
|
||||
Int16List aeicZIndexes(AeicEntropyGeometry geometry) {
|
||||
final plane = geometry.zHeight * geometry.zWidth;
|
||||
final out = Int16List(geometry.zElements);
|
||||
var i = 0;
|
||||
for (var c = 0; c < geometry.zChannels; c++) {
|
||||
for (var p = 0; p < plane; p++) {
|
||||
out[i++] = c;
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/// float32 tensor of integers -> int16 symbols, asserting the cast is lossless.
|
||||
///
|
||||
/// `py_rans` force-casts through `py::array_t<int16_t>`; the golden vectors
|
||||
/// store the post-cast values. Anything that does not survive the cast means
|
||||
/// the entropy model has gone somewhere the format cannot represent, and
|
||||
/// truncating quietly would desync the decoder.
|
||||
Int16List aeicToSymbols(Float32List values) {
|
||||
final out = Int16List(values.length);
|
||||
for (var i = 0; i < values.length; i++) {
|
||||
final v = values[i];
|
||||
final n = v.toInt();
|
||||
if (n != v || n < -32768 || n > 32767) {
|
||||
throw StateError(
|
||||
'entropy symbol $v at index $i is not a lossless int16; the entropy '
|
||||
'model produced a value the bitstream format cannot carry',
|
||||
);
|
||||
}
|
||||
out[i] = n;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/// RGB bytes -> the graph's `image` input, reproducing torchvision's
|
||||
/// `ToTensor()` + `Normalize([0.5]*3, [0.5]*3)`: `x = (b/255 - 0.5) / 0.5`.
|
||||
Float32List aeicRgbToChw(Uint8List rgb, int resolution) {
|
||||
final pixels = resolution * resolution;
|
||||
if (rgb.length != pixels * 3) {
|
||||
throw ArgumentError.value(
|
||||
rgb.length,
|
||||
'rgb',
|
||||
'expected ${pixels * 3} bytes for ${resolution}x$resolution RGB',
|
||||
);
|
||||
}
|
||||
final out = Float32List(pixels * 3);
|
||||
for (var c = 0; c < 3; c++) {
|
||||
final plane = c * pixels;
|
||||
for (var i = 0; i < pixels; i++) {
|
||||
final v = f32(rgb[i * 3 + c] / 255.0);
|
||||
out[plane + i] = f32(v - 0.5) * 2.0; // /0.5 is exact
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/// Everything the send side needs out of one forward pass of the entropy graph.
|
||||
///
|
||||
/// Matches the outputs of `aic/onnx/aeic_entropy_side_fp32_op17.onnx`, which
|
||||
/// runs `g_a`, `h_a`, `h_s`, `g_c` and all four adapters in one shot — the
|
||||
/// encoder never has to be incremental, only the decoder does.
|
||||
class AeicEncodeSideTensors {
|
||||
/// `round(z - z_offset)`, `[1, M/2, z_h, z_w]`, integral floats.
|
||||
final Float32List zQ;
|
||||
|
||||
/// `round(y * mask_i - means_i)` per stage, `[1, M, y_h, y_w]`, integral.
|
||||
final List<Float32List> yQ;
|
||||
|
||||
/// `scales_supp_i * mask_i` per stage, `[1, M, y_h, y_w]`.
|
||||
final List<Float32List> scales;
|
||||
|
||||
const AeicEncodeSideTensors({
|
||||
required this.zQ,
|
||||
required this.yQ,
|
||||
required this.scales,
|
||||
});
|
||||
}
|
||||
|
||||
/// One decode stage's entropy parameters, **before** masking.
|
||||
///
|
||||
/// `adapter_out[i](g_c(adapter_in[i](base)))` split in half; the caller applies
|
||||
/// `mask_i`. Kept unmasked so the graph does not have to know the stage index
|
||||
/// at export time.
|
||||
class AeicStageParams {
|
||||
final Float32List meansSupp;
|
||||
final Float32List scalesSupp;
|
||||
|
||||
const AeicStageParams({required this.meansSupp, required this.scalesSupp});
|
||||
}
|
||||
|
||||
/// The neural half of the entropy path, as this file needs it.
|
||||
///
|
||||
/// ## Graph contract
|
||||
///
|
||||
/// **Send side** (exists today, `aeic_entropy_side_fp32_op17.onnx`, 67 MB):
|
||||
/// input `image [1,3,512,512]` -> outputs `z_q`, `yq0..yq3`, `sc0..sc3`.
|
||||
///
|
||||
/// **Receive side** (NOT exported yet — see the report for task B2): the
|
||||
/// decoder cannot use the send-side graph, because it must interleave network
|
||||
/// evaluation with symbol decoding: stage `i`'s indexes are unknown until
|
||||
/// stages `< i` have been decoded and fed back through `g_c`. It needs
|
||||
///
|
||||
/// * `z_q [1,128,4,4]` -> `base0 [1,256,16,16]` (`h_s` with `z_offset` baked
|
||||
/// in, exactly as the send-side export bakes it), and
|
||||
/// * `base [1,256,16,16]` -> `means{i}`, `scales{i}` for `i` in 0..3
|
||||
/// (unmasked `adapter_out[i](g_c(adapter_in[i](base)))`, split in half).
|
||||
///
|
||||
/// Both can live in one graph with two inputs and nine outputs; ORT prunes to
|
||||
/// the requested output set, so asking only for `base0` does not run the
|
||||
/// adapters and asking only for `means2`/`scales2` does not run `h_s`.
|
||||
abstract class AeicEntropyNetwork {
|
||||
/// Whether [runHyperSynthesis] and [runStage] are available. False for a
|
||||
/// graph that only carries the send-side path.
|
||||
bool get supportsDecodeSide;
|
||||
|
||||
Future<AeicEncodeSideTensors> runEncodeSide(Float32List imageChw);
|
||||
|
||||
/// `h_s(z_q + z_offset)[:, :, :y_h, :y_w]`.
|
||||
Future<Float32List> runHyperSynthesis(Float32List zQ);
|
||||
|
||||
/// `adapter_out[stage](g_c(adapter_in[stage](base)))`, unmasked.
|
||||
Future<AeicStageParams> runStage(int stage, Float32List base);
|
||||
}
|
||||
|
||||
/// The rANS encoder, as the entropy layer uses it.
|
||||
///
|
||||
/// Implemented by the pure-Dart coder (task B1). Deliberately narrow: five
|
||||
/// [pushSymbols] calls in the fixed order z, y0, y1, y2, y3, then one [finish].
|
||||
abstract class AeicRansEncoder {
|
||||
/// Appends one `(symbols, indexes)` pair. `indexes[i] < 0` emits nothing.
|
||||
void pushSymbols(Int16List symbols, Int16List indexes, int cdfGroup);
|
||||
|
||||
/// Flushes both sub-streams and returns the container bytes.
|
||||
Uint8List finish();
|
||||
}
|
||||
|
||||
/// The rANS decoder, as the entropy layer uses it.
|
||||
///
|
||||
/// **Must be incremental**: each call resumes the rANS state where the last one
|
||||
/// stopped. An implementation that decodes the whole stream up front cannot
|
||||
/// work here — see the class docs on [AeicEntropyNetwork].
|
||||
abstract class AeicRansDecoder {
|
||||
Int16List decodeStream(Int16List indexes, int cdfGroup);
|
||||
}
|
||||
|
||||
/// Supplies coders bound to a parsed CDF table set (task B1 owns both).
|
||||
abstract class AeicRansCoderFactory {
|
||||
AeicRansEncoder createEncoder();
|
||||
|
||||
AeicRansDecoder createDecoder(Uint8List bitstream);
|
||||
}
|
||||
|
||||
/// Thrown when the entropy loop is asked to run without a piece it needs.
|
||||
class AeicEntropyUnavailable implements Exception {
|
||||
final String detail;
|
||||
|
||||
const AeicEntropyUnavailable(this.detail);
|
||||
|
||||
@override
|
||||
String toString() => 'AEIC entropy path unavailable: $detail';
|
||||
}
|
||||
|
||||
/// Raised by [ImageCodecBackend] callers that pass a cancel predicate.
|
||||
typedef AeicCancelCheck = bool Function();
|
||||
|
||||
/// The four-stage masked encode and decode, orchestrating [AeicEntropyNetwork]
|
||||
/// and [AeicRansCoderFactory].
|
||||
class AeicEntropyCodec {
|
||||
final AeicEntropyGeometry geometry;
|
||||
final AeicMaskSet masks;
|
||||
final AeicEntropyNetwork network;
|
||||
final AeicRansCoderFactory coders;
|
||||
|
||||
AeicEntropyCodec({
|
||||
required this.geometry,
|
||||
required this.network,
|
||||
required this.coders,
|
||||
}) : masks = AeicMaskSet(geometry);
|
||||
|
||||
/// Packed 8-bit RGB -> rANS bitstream (payload only, no chunk headers).
|
||||
///
|
||||
/// One forward pass, then z, y0, y1, y2, y3 into the coder in that order.
|
||||
Future<Uint8List> encode(
|
||||
Uint8List rgbBytes, {
|
||||
void Function(double progress)? onProgress,
|
||||
AeicCancelCheck? shouldCancel,
|
||||
}) async {
|
||||
onProgress?.call(0.02);
|
||||
final input = aeicRgbToChw(rgbBytes, geometry.resolution);
|
||||
_checkCancel(shouldCancel);
|
||||
|
||||
final tensors = await network.runEncodeSide(input);
|
||||
onProgress?.call(0.80);
|
||||
_checkCancel(shouldCancel);
|
||||
|
||||
if (tensors.zQ.length != geometry.zElements) {
|
||||
throw StateError(
|
||||
'entropy graph returned ${tensors.zQ.length} z values, expected '
|
||||
'${geometry.zElements} (${geometry.zShape})',
|
||||
);
|
||||
}
|
||||
if (tensors.yQ.length != 4 || tensors.scales.length != 4) {
|
||||
throw StateError(
|
||||
'entropy graph must return four y_q and four scales tensors, got '
|
||||
'${tensors.yQ.length} / ${tensors.scales.length}',
|
||||
);
|
||||
}
|
||||
|
||||
final encoder = coders.createEncoder();
|
||||
encoder.pushSymbols(
|
||||
aeicToSymbols(tensors.zQ),
|
||||
aeicZIndexes(geometry),
|
||||
kAeicZCdfGroup,
|
||||
);
|
||||
for (var stage = 0; stage < 4; stage++) {
|
||||
final symbols = aeicToSymbols(masks.squeeze(tensors.yQ[stage]));
|
||||
final indexes = aeicBuildIndexes(masks.squeeze(tensors.scales[stage]));
|
||||
encoder.pushSymbols(symbols, indexes, kAeicYCdfGroup);
|
||||
onProgress?.call(0.80 + 0.04 * (stage + 1));
|
||||
_checkCancel(shouldCancel);
|
||||
}
|
||||
final stream = encoder.finish();
|
||||
onProgress?.call(1.0);
|
||||
return stream;
|
||||
}
|
||||
|
||||
/// rANS bitstream -> `y_hat [1, M, y_h, y_w]`, ready for the synthesis graph.
|
||||
///
|
||||
/// The mirror of [encode], and necessarily incremental: each stage's indexes
|
||||
/// come from scales that only exist once the previous stage's symbols have
|
||||
/// been decoded and pushed back through the context model.
|
||||
Future<Float32List> decodeToLatent(
|
||||
Uint8List bitstream, {
|
||||
void Function(double progress)? onProgress,
|
||||
AeicCancelCheck? shouldCancel,
|
||||
}) async {
|
||||
if (!network.supportsDecodeSide) {
|
||||
throw const AeicEntropyUnavailable(
|
||||
'the installed entropy graph is send-side only: it maps image -> '
|
||||
'symbols and has no z_q -> base0 / base -> means,scales entry points, '
|
||||
'which decoding requires',
|
||||
);
|
||||
}
|
||||
onProgress?.call(0.02);
|
||||
final decoder = coders.createDecoder(bitstream);
|
||||
|
||||
final zSymbols = decoder.decodeStream(
|
||||
aeicZIndexes(geometry),
|
||||
kAeicZCdfGroup,
|
||||
);
|
||||
final zQ = Float32List(zSymbols.length);
|
||||
for (var i = 0; i < zSymbols.length; i++) {
|
||||
zQ[i] = zSymbols[i].toDouble();
|
||||
}
|
||||
_checkCancel(shouldCancel);
|
||||
|
||||
var base = await network.runHyperSynthesis(zQ);
|
||||
if (base.length != geometry.yElements) {
|
||||
throw StateError(
|
||||
'hyper synthesis returned ${base.length} values, expected '
|
||||
'${geometry.yElements} (${geometry.yShape})',
|
||||
);
|
||||
}
|
||||
onProgress?.call(0.20);
|
||||
|
||||
Float32List? stageLatent;
|
||||
for (var stage = 0; stage < 4; stage++) {
|
||||
_checkCancel(shouldCancel);
|
||||
final params = await network.runStage(stage, base);
|
||||
final scales = masks.applyMask(params.scalesSupp, stage);
|
||||
final means = masks.applyMask(params.meansSupp, stage);
|
||||
final indexes = aeicBuildIndexes(masks.squeeze(scales));
|
||||
final symbols = decoder.decodeStream(indexes, kAeicYCdfGroup);
|
||||
final meansSqueezed = masks.squeeze(means);
|
||||
if (symbols.length != meansSqueezed.length) {
|
||||
throw StateError(
|
||||
'stage $stage decoded ${symbols.length} symbols but the context '
|
||||
'model produced ${meansSqueezed.length} means',
|
||||
);
|
||||
}
|
||||
final latentSqueezed = Float32List(symbols.length);
|
||||
for (var i = 0; i < symbols.length; i++) {
|
||||
latentSqueezed[i] = symbols[i] + meansSqueezed[i];
|
||||
}
|
||||
stageLatent = masks.unsqueeze(latentSqueezed, stage);
|
||||
if (stage < 3) {
|
||||
base = masks.mergeContext(base, stageLatent, stage);
|
||||
}
|
||||
onProgress?.call(0.20 + 0.19 * (stage + 1));
|
||||
}
|
||||
|
||||
// y_hat = base * (1 - mask_3) + y_hat_3
|
||||
return masks.mergeContext(base, stageLatent!, 3);
|
||||
}
|
||||
|
||||
static void _checkCancel(AeicCancelCheck? shouldCancel) {
|
||||
if (shouldCancel?.call() == true) {
|
||||
throw const AeicEntropyCancelled();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Cancellation signal raised out of the entropy loop at a stage boundary.
|
||||
class AeicEntropyCancelled implements Exception {
|
||||
const AeicEntropyCancelled();
|
||||
|
||||
@override
|
||||
String toString() => 'Image processing stopped.';
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
export 'image_codec_file_store_stub.dart'
|
||||
if (dart.library.io) 'image_codec_file_store_io.dart';
|
||||
@@ -0,0 +1,247 @@
|
||||
import 'dart:io';
|
||||
import 'dart:typed_data';
|
||||
|
||||
import 'package:convert/convert.dart' show AccumulatorSink;
|
||||
import 'package:crypto/crypto.dart';
|
||||
import 'package:path_provider/path_provider.dart';
|
||||
|
||||
import '../models/image_codec_support.dart';
|
||||
|
||||
/// On-disk storage for image-codec weights.
|
||||
///
|
||||
/// Mirrors `translation_file_store_io.dart` 1:1. It is a separate class rather
|
||||
/// than a reuse of `TranslationFileStore` only because the directory name and
|
||||
/// the record type differ; see the report for the (small) change that would
|
||||
/// make the translation store generic enough to share.
|
||||
class ImageCodecFileStore {
|
||||
static final RegExp _chunkFilePattern = RegExp(r'^\..+_chunk_\d+$');
|
||||
|
||||
Future<String> modelDirectoryPath() async {
|
||||
final baseDir = await getApplicationDocumentsDirectory();
|
||||
final dir = Directory('${baseDir.path}/image_codec_models');
|
||||
if (!dir.existsSync()) {
|
||||
await dir.create(recursive: true);
|
||||
}
|
||||
return dir.path;
|
||||
}
|
||||
|
||||
Future<List<ImageCodecModelRecord>> scanDownloadedModels() async {
|
||||
final dir = Directory(await modelDirectoryPath());
|
||||
if (!dir.existsSync()) {
|
||||
return const [];
|
||||
}
|
||||
final models = <ImageCodecModelRecord>[];
|
||||
for (final entity in dir.listSync().whereType<File>()) {
|
||||
final name = entity.uri.pathSegments.last;
|
||||
if (name.startsWith('.')) {
|
||||
// Hidden `.<file>_chunk_<n>` files are the resume state of an
|
||||
// interrupted download and MUST survive a restart — reaping them here
|
||||
// (as this used to) is what made an 872 MB transfer restart from zero
|
||||
// whenever the app was reopened. Anything else hidden is junk.
|
||||
if (!_chunkFilePattern.hasMatch(name)) {
|
||||
await entity.delete();
|
||||
}
|
||||
continue;
|
||||
}
|
||||
final stat = entity.statSync();
|
||||
models.add(
|
||||
ImageCodecModelRecord(
|
||||
id: name,
|
||||
name: name,
|
||||
sourceUrl: '',
|
||||
localPath: entity.path,
|
||||
downloadedAt: stat.modified,
|
||||
fileSizeBytes: stat.size,
|
||||
),
|
||||
);
|
||||
}
|
||||
return models;
|
||||
}
|
||||
|
||||
Future<void> deleteModel(ImageCodecModelRecord model) async {
|
||||
await deleteFile(model.localPath);
|
||||
await deletePartialDownloads(model.name);
|
||||
}
|
||||
|
||||
/// Removes any resume state for [fileName].
|
||||
///
|
||||
/// Called after a successful download and when a model is removed. Without it
|
||||
/// the chunk files, which [scanDownloadedModels] now deliberately preserves,
|
||||
/// would never be collected.
|
||||
Future<void> deletePartialDownloads(String fileName) async {
|
||||
if (fileName.isEmpty) return;
|
||||
final dir = Directory(await modelDirectoryPath());
|
||||
if (!dir.existsSync()) return;
|
||||
// `.<fileName>[.<totalSize>]_chunk_<n>`. The optional size segment is what
|
||||
// `ImageCodecService` appends so resume can never splice offsets computed
|
||||
// for one upstream length onto a file of another. Anchored on both ends so
|
||||
// sweeping `model.onnx` cannot take `model.onnx.data`'s chunks with it.
|
||||
final pattern = RegExp(
|
||||
'^\\.${RegExp.escape(fileName)}(\\.\\d+)?_chunk_\\d+\$',
|
||||
);
|
||||
for (final entity in dir.listSync().whereType<File>()) {
|
||||
if (pattern.hasMatch(entity.uri.pathSegments.last)) {
|
||||
await entity.delete();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Future<void> deleteFile(String path) async {
|
||||
final file = File(path);
|
||||
if (file.existsSync()) {
|
||||
await file.delete();
|
||||
}
|
||||
}
|
||||
|
||||
Future<DownloadedCodecFile> writeModelBytes({
|
||||
required String fileName,
|
||||
required Stream<List<int>> chunks,
|
||||
}) async {
|
||||
final directoryPath = await modelDirectoryPath();
|
||||
final file = File('$directoryPath/$fileName');
|
||||
final sink = file.openWrite();
|
||||
var fileSizeBytes = 0;
|
||||
var completed = false;
|
||||
try {
|
||||
await for (final chunk in chunks) {
|
||||
sink.add(chunk);
|
||||
fileSizeBytes += chunk.length;
|
||||
}
|
||||
completed = true;
|
||||
} finally {
|
||||
await sink.close();
|
||||
if (!completed && file.existsSync()) {
|
||||
await file.delete();
|
||||
}
|
||||
}
|
||||
return DownloadedCodecFile(
|
||||
localPath: file.path,
|
||||
fileSizeBytes: fileSizeBytes,
|
||||
);
|
||||
}
|
||||
|
||||
Future<String> chunkFilePath(String fileName, int index) async {
|
||||
final dir = await modelDirectoryPath();
|
||||
return '$dir/.${fileName}_chunk_$index';
|
||||
}
|
||||
|
||||
Future<String> modelFilePath(String fileName) async {
|
||||
final dir = await modelDirectoryPath();
|
||||
return '$dir/$fileName';
|
||||
}
|
||||
|
||||
/// Size of [path] in bytes, or 0 when it does not exist.
|
||||
///
|
||||
/// This is the whole basis of resume: a chunk file's length *is* its progress
|
||||
/// marker, so no separate state file can go stale or disagree with the bytes.
|
||||
Future<int> fileSize(String path) async {
|
||||
final file = File(path);
|
||||
if (!file.existsSync()) {
|
||||
return 0;
|
||||
}
|
||||
return file.length();
|
||||
}
|
||||
|
||||
/// Appends [chunks] to [path], creating it if absent.
|
||||
///
|
||||
/// Unlike [writeModelBytes] this does NOT delete the file when the stream
|
||||
/// fails part-way: the partial bytes are exactly what the next attempt
|
||||
/// resumes from. Returns the file's total size afterwards.
|
||||
Future<int> appendBytes({
|
||||
required String path,
|
||||
required Stream<List<int>> chunks,
|
||||
}) async {
|
||||
final file = File(path);
|
||||
await file.parent.create(recursive: true);
|
||||
final sink = file.openWrite(mode: FileMode.append);
|
||||
try {
|
||||
await for (final chunk in chunks) {
|
||||
sink.add(chunk);
|
||||
}
|
||||
} finally {
|
||||
// flush() inside close(); the bytes written before a failure survive on
|
||||
// purpose, so the next Range request can pick up from file.length().
|
||||
await sink.close();
|
||||
}
|
||||
return file.length();
|
||||
}
|
||||
|
||||
/// Streaming SHA-256 of [path] as lowercase hex.
|
||||
///
|
||||
/// Streamed, not `readAsBytes`: the data sibling is 869 MiB and reading it into
|
||||
/// a Uint8List to hash it would defeat the point of downloading it to disk.
|
||||
Future<String> sha256OfFile(String path) async {
|
||||
final accumulator = AccumulatorSink<Digest>();
|
||||
final converter = sha256.startChunkedConversion(accumulator);
|
||||
try {
|
||||
await for (final chunk in File(path).openRead()) {
|
||||
converter.add(chunk);
|
||||
}
|
||||
} finally {
|
||||
converter.close();
|
||||
}
|
||||
return accumulator.events.single.toString();
|
||||
}
|
||||
|
||||
/// Reads an arbitrary file (a picked photo, not a model) as bytes.
|
||||
///
|
||||
/// Lives here so `ImageCodecService` never has to import `dart:io` and can
|
||||
/// therefore still be constructed on web.
|
||||
Future<Uint8List> readFileBytes(String path) => File(path).readAsBytes();
|
||||
|
||||
Future<DownloadedCodecFile> combineChunks({
|
||||
required String fileName,
|
||||
required List<String> chunkPaths,
|
||||
}) async {
|
||||
final dir = await modelDirectoryPath();
|
||||
final finalPath = '$dir/$fileName';
|
||||
final sink = File(finalPath).openWrite();
|
||||
var totalSize = 0;
|
||||
var completed = false;
|
||||
try {
|
||||
for (final chunkPath in chunkPaths) {
|
||||
final chunkFile = File(chunkPath);
|
||||
await sink.addStream(chunkFile.openRead());
|
||||
totalSize += await chunkFile.length();
|
||||
}
|
||||
completed = true;
|
||||
} finally {
|
||||
await sink.close();
|
||||
if (completed) {
|
||||
for (final chunkPath in chunkPaths) {
|
||||
final file = File(chunkPath);
|
||||
if (file.existsSync()) {
|
||||
await file.delete();
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Keep the chunk files: they are the resume state. Only the half-written
|
||||
// merge target is thrown away. (The previous version deleted the chunks
|
||||
// unconditionally, which is why an interrupted 872 MB download restarted
|
||||
// from zero.)
|
||||
final finalFile = File(finalPath);
|
||||
if (finalFile.existsSync()) {
|
||||
await finalFile.delete();
|
||||
}
|
||||
}
|
||||
}
|
||||
return DownloadedCodecFile(localPath: finalPath, fileSizeBytes: totalSize);
|
||||
}
|
||||
|
||||
/// Free bytes are not queryable without a platform channel, so callers that
|
||||
/// need a pre-flight space check must supply their own. Kept here as the
|
||||
/// documented seam rather than a silent omission: an 833 MB download that
|
||||
/// dies at 90% full is the most likely field failure for this feature.
|
||||
// TODO(disk): add a free-space pre-flight (needs a platform channel or the
|
||||
// `disk_space_plus` package) before enabling the download button.
|
||||
}
|
||||
|
||||
class DownloadedCodecFile {
|
||||
final String localPath;
|
||||
final int fileSizeBytes;
|
||||
|
||||
const DownloadedCodecFile({
|
||||
required this.localPath,
|
||||
required this.fileSizeBytes,
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
import 'dart:typed_data';
|
||||
|
||||
import '../models/image_codec_support.dart';
|
||||
|
||||
class ImageCodecFileStore {
|
||||
Future<String> modelDirectoryPath() async {
|
||||
throw UnsupportedError('Local codec model storage is not supported on web.');
|
||||
}
|
||||
|
||||
Future<List<ImageCodecModelRecord>> scanDownloadedModels() async {
|
||||
return const [];
|
||||
}
|
||||
|
||||
Future<void> deleteModel(ImageCodecModelRecord model) async {}
|
||||
|
||||
Future<void> deleteFile(String path) async {}
|
||||
|
||||
Future<DownloadedCodecFile> writeModelBytes({
|
||||
required String fileName,
|
||||
required Stream<List<int>> chunks,
|
||||
}) async {
|
||||
throw UnsupportedError('Local codec model storage is not supported on web.');
|
||||
}
|
||||
|
||||
Future<String> chunkFilePath(String fileName, int index) async {
|
||||
throw UnsupportedError('Local codec model storage is not supported on web.');
|
||||
}
|
||||
|
||||
Future<String> modelFilePath(String fileName) async {
|
||||
throw UnsupportedError('Local codec model storage is not supported on web.');
|
||||
}
|
||||
|
||||
Future<void> deletePartialDownloads(String fileName) async {}
|
||||
|
||||
Future<int> fileSize(String path) async => 0;
|
||||
|
||||
Future<int> appendBytes({
|
||||
required String path,
|
||||
required Stream<List<int>> chunks,
|
||||
}) async {
|
||||
throw UnsupportedError('Local codec model storage is not supported on web.');
|
||||
}
|
||||
|
||||
Future<String> sha256OfFile(String path) async {
|
||||
throw UnsupportedError('Local codec model storage is not supported on web.');
|
||||
}
|
||||
|
||||
Future<Uint8List> readFileBytes(String path) async {
|
||||
throw UnsupportedError('Local file reads are not supported on web.');
|
||||
}
|
||||
|
||||
Future<DownloadedCodecFile> combineChunks({
|
||||
required String fileName,
|
||||
required List<String> chunkPaths,
|
||||
}) async {
|
||||
throw UnsupportedError('Local codec model storage is not supported on web.');
|
||||
}
|
||||
}
|
||||
|
||||
class DownloadedCodecFile {
|
||||
final String localPath;
|
||||
final int fileSizeBytes;
|
||||
|
||||
const DownloadedCodecFile({
|
||||
required this.localPath,
|
||||
required this.fileSizeBytes,
|
||||
});
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,2 @@
|
||||
export 'image_codec_session_stub.dart'
|
||||
if (dart.library.io) 'image_codec_session_io.dart';
|
||||
@@ -0,0 +1,606 @@
|
||||
import 'dart:async';
|
||||
import 'dart:io';
|
||||
import 'dart:isolate';
|
||||
import 'dart:typed_data';
|
||||
|
||||
import 'package:flutter/foundation.dart' show visibleForTesting;
|
||||
import 'package:flutter/services.dart';
|
||||
|
||||
import '../models/image_codec_support.dart';
|
||||
import '../widgets/image_send_codec_binding.dart' show kImageCodecSquareSize;
|
||||
import 'image_codec_backend.dart';
|
||||
import 'entropy_tables.dart';
|
||||
|
||||
/// Owns a long-lived [Isolate] and the native codec session inside it.
|
||||
///
|
||||
/// A decode is 1385 GFLOP over 886M parameters. It must never run on the root
|
||||
/// isolate: it would freeze the UI and, worse, stall the BLE notify stream that
|
||||
/// `MeshCoreConnector._handleFrame()` feeds on, dropping mesh traffic while a
|
||||
/// picture renders.
|
||||
///
|
||||
/// The isolate is spawned once per model load and reused, because loading the
|
||||
/// model is the expensive part (~833 MB int8). This is the same amortisation
|
||||
/// `TranslationService._ensureContext` gets for free from llamadart, which owns
|
||||
/// its own worker thread. Do NOT switch this to `Isolate.run` per call.
|
||||
///
|
||||
/// Message protocol (all maps, all `type`-tagged):
|
||||
/// worker -> host `ready` {port: SendPort, backend: String,
|
||||
/// bitstream: bool, entropy: bool, tables: bool}
|
||||
/// `fatal` {error: String} — during startup only
|
||||
/// `progress` {id: int, value: double}
|
||||
/// `result` {id: int, bytes: Uint8List}
|
||||
/// `released` {id: int}
|
||||
/// `error` {id: int, error: String, stack: String?}
|
||||
/// host -> worker `job` {id, op: 'encode'|'decode'|'decodeLatent',
|
||||
/// bytes | latents, rate, size}
|
||||
/// `release` {id, which: 'decoder'|'entropy'|'both'}
|
||||
/// `cancel` {}
|
||||
/// `shutdown` {}
|
||||
///
|
||||
/// ## Two sessions, two lifetimes
|
||||
///
|
||||
/// The backend holds the 2.16 GiB decoder and the 67 MB entropy graph
|
||||
/// independently, and [release] is how the host sheds one without killing the
|
||||
/// other. It exists because `decode()` MUST drop the entropy graph before it
|
||||
/// creates the decoder — holding both at once is ~2.2 GiB plus 67 MB of arena
|
||||
/// on a phone — and because memory pressure wants the big half gone while the
|
||||
/// small half (all that `encode()` needs) can survive.
|
||||
///
|
||||
/// ## Plugin channels on a background isolate
|
||||
///
|
||||
/// The ONNX backend is a *plugin*, reached over a platform channel, and a plain
|
||||
/// spawned isolate has no binary messenger — every `invokeMethod` would throw
|
||||
/// until `BackgroundIsolateBinaryMessenger.ensureInitialized` is handed a
|
||||
/// [RootIsolateToken] captured on the root isolate. That token is part of the
|
||||
/// boot payload below. Consequence: [spawn] must be called from the root
|
||||
/// isolate. It is (via `ImageCodecService`), and if that ever changes the token
|
||||
/// is null and the worker fails loudly at startup instead of on first inference.
|
||||
class ImageCodecSession {
|
||||
final Isolate _isolate;
|
||||
final SendPort _toWorker;
|
||||
final ReceivePort _fromWorker;
|
||||
final ReceivePort _errors;
|
||||
|
||||
/// Name of the backend that actually loaded, for logs and error strings.
|
||||
final String backendName;
|
||||
|
||||
/// Whether the loaded backend can turn bytes into a picture, as opposed to
|
||||
/// only latents into a picture. True only when the bundle carried the entropy
|
||||
/// graph and the CDF tables AND the backend implements the entropy path.
|
||||
final bool supportsBitstreamCodec;
|
||||
|
||||
/// Whether the bundle handed to [spawn] carried an entropy-side graph.
|
||||
final bool hasEntropyGraph;
|
||||
|
||||
/// Whether the bundle handed to [spawn] carried a CDF table file.
|
||||
final bool hasTables;
|
||||
|
||||
final Map<int, _PendingJob> _pending = {};
|
||||
final Map<int, Completer<void>> _pendingReleases = {};
|
||||
int _nextJobId = 1;
|
||||
bool _disposed = false;
|
||||
|
||||
ImageCodecSession._({
|
||||
required Isolate isolate,
|
||||
required SendPort toWorker,
|
||||
required ReceivePort fromWorker,
|
||||
required ReceivePort errors,
|
||||
required this.backendName,
|
||||
required this.supportsBitstreamCodec,
|
||||
required this.hasEntropyGraph,
|
||||
required this.hasTables,
|
||||
}) : _isolate = isolate,
|
||||
_toWorker = toWorker,
|
||||
_fromWorker = fromWorker,
|
||||
_errors = errors;
|
||||
|
||||
/// Spawns the worker and blocks until the backend has *validated* [bundle].
|
||||
///
|
||||
/// Cheap by design: `load()` parses the CDF tables (813 KB) and checks that
|
||||
/// every path exists, and creates NO ORT session. Sessions are created lazily
|
||||
/// by the first operation that needs one, which is what lets a send pay for
|
||||
/// the 67 MB entropy graph alone instead of the 2.16 GiB decoder as well.
|
||||
///
|
||||
/// Throws [ImageCodecUnimplemented] when no inference backend is compiled
|
||||
/// into the build, or whatever the backend threw while loading.
|
||||
static Future<ImageCodecSession> spawn(ImageCodecBundle bundle) async {
|
||||
final fromWorker = ReceivePort();
|
||||
final errors = ReceivePort();
|
||||
final handshake = Completer<Map<String, Object?>>();
|
||||
ImageCodecSession? session;
|
||||
|
||||
fromWorker.listen((message) {
|
||||
if (message is! Map) return;
|
||||
final map = message.cast<String, Object?>();
|
||||
final active = session;
|
||||
if (active != null) {
|
||||
active._handleMessage(map);
|
||||
} else if (!handshake.isCompleted) {
|
||||
handshake.complete(map);
|
||||
}
|
||||
});
|
||||
|
||||
errors.listen((message) {
|
||||
final description = message is List && message.isNotEmpty
|
||||
? message.first.toString()
|
||||
: message.toString();
|
||||
final active = session;
|
||||
if (active != null) {
|
||||
active._failAll(StateError('Image codec isolate died: $description'));
|
||||
} else if (!handshake.isCompleted) {
|
||||
handshake.complete({'type': 'fatal', 'error': description});
|
||||
}
|
||||
});
|
||||
|
||||
late final Isolate isolate;
|
||||
try {
|
||||
isolate = await Isolate.spawn<List<Object?>>(
|
||||
_codecWorkerMain,
|
||||
_bootPayload(
|
||||
bundle,
|
||||
fromWorker.sendPort,
|
||||
// Null when spawn() was not called from the root isolate. The worker
|
||||
// treats that as a fatal startup error rather than limping on to fail
|
||||
// at the first invokeMethod.
|
||||
RootIsolateToken.instance,
|
||||
),
|
||||
errorsAreFatal: true,
|
||||
onError: errors.sendPort,
|
||||
debugName: 'image-codec',
|
||||
);
|
||||
} catch (_) {
|
||||
fromWorker.close();
|
||||
errors.close();
|
||||
rethrow;
|
||||
}
|
||||
|
||||
Map<String, Object?> reply;
|
||||
try {
|
||||
reply = await handshake.future;
|
||||
} catch (_) {
|
||||
isolate.kill(priority: Isolate.immediate);
|
||||
fromWorker.close();
|
||||
errors.close();
|
||||
rethrow;
|
||||
}
|
||||
|
||||
if (reply['type'] != 'ready') {
|
||||
isolate.kill(priority: Isolate.immediate);
|
||||
fromWorker.close();
|
||||
errors.close();
|
||||
final detail = reply['error']?.toString() ?? 'unknown startup failure';
|
||||
if (reply['unimplemented'] == true) {
|
||||
throw ImageCodecUnimplemented(detail);
|
||||
}
|
||||
throw StateError('Image codec failed to start: $detail');
|
||||
}
|
||||
|
||||
return session = ImageCodecSession._(
|
||||
isolate: isolate,
|
||||
toWorker: reply['port'] as SendPort,
|
||||
fromWorker: fromWorker,
|
||||
errors: errors,
|
||||
backendName: reply['backend'] as String? ?? 'unknown',
|
||||
supportsBitstreamCodec: reply['bitstream'] == true,
|
||||
hasEntropyGraph: reply['entropy'] == true,
|
||||
hasTables: reply['tables'] == true,
|
||||
);
|
||||
}
|
||||
|
||||
/// Runs the synthesis half: `y_hat` -> packed 8-bit RGB.
|
||||
///
|
||||
/// This is the one inference path that actually works today. [encode] and
|
||||
/// [decode] fail with [ImageCodecEntropyPathMissing] until the entropy-side
|
||||
/// graph and the rANS coder land.
|
||||
Future<Uint8List> decodeLatent(
|
||||
Float32List yHat, {
|
||||
void Function(double progress)? onProgress,
|
||||
}) {
|
||||
return _submit(
|
||||
op: 'decodeLatent',
|
||||
latents: yHat,
|
||||
ratePoint: kShippingAeicRatePoint,
|
||||
resolution: kImageCodecSquareSize,
|
||||
onProgress: onProgress,
|
||||
);
|
||||
}
|
||||
|
||||
Future<Uint8List> encode(
|
||||
Uint8List rgbBytes,
|
||||
AeicRatePoint ratePoint,
|
||||
int resolution, {
|
||||
void Function(double progress)? onProgress,
|
||||
}) {
|
||||
return _submit(
|
||||
op: 'encode',
|
||||
bytes: rgbBytes,
|
||||
ratePoint: ratePoint,
|
||||
resolution: resolution,
|
||||
onProgress: onProgress,
|
||||
);
|
||||
}
|
||||
|
||||
Future<Uint8List> decode(
|
||||
Uint8List bitstream,
|
||||
AeicRatePoint ratePoint,
|
||||
int resolution, {
|
||||
void Function(double progress)? onProgress,
|
||||
}) {
|
||||
return _submit(
|
||||
op: 'decode',
|
||||
bytes: bitstream,
|
||||
ratePoint: ratePoint,
|
||||
resolution: resolution,
|
||||
onProgress: onProgress,
|
||||
);
|
||||
}
|
||||
|
||||
Future<Uint8List> _submit({
|
||||
required String op,
|
||||
Uint8List? bytes,
|
||||
Float32List? latents,
|
||||
required AeicRatePoint ratePoint,
|
||||
required int resolution,
|
||||
void Function(double progress)? onProgress,
|
||||
}) {
|
||||
if (_disposed) {
|
||||
return Future.error(StateError('Image codec session disposed.'));
|
||||
}
|
||||
final id = _nextJobId++;
|
||||
final job = _PendingJob(onProgress: onProgress);
|
||||
_pending[id] = job;
|
||||
_toWorker.send(<String, Object?>{
|
||||
'type': 'job',
|
||||
'id': id,
|
||||
'op': op,
|
||||
'bytes': ?bytes,
|
||||
'latents': ?latents,
|
||||
'rate': ratePoint.wireValue,
|
||||
'size': resolution,
|
||||
});
|
||||
return job.completer.future;
|
||||
}
|
||||
|
||||
/// Drops one or both ORT sessions without killing the isolate.
|
||||
///
|
||||
/// The paths stay recorded inside the backend, so a released session is
|
||||
/// re-created on the next call that needs it. Completes when the worker has
|
||||
/// acknowledged, so a caller can rely on the memory being back before it
|
||||
/// creates the other session.
|
||||
///
|
||||
/// Queued behind any in-flight job, like every other command: releasing the
|
||||
/// decoder out from under a running synthesis pass would crash ORT.
|
||||
Future<void> release({bool decoder = false, bool entropy = false}) {
|
||||
if (!decoder && !entropy) return Future<void>.value();
|
||||
if (_disposed) return Future<void>.value();
|
||||
final id = _nextJobId++;
|
||||
final completer = Completer<void>();
|
||||
_pendingReleases[id] = completer;
|
||||
_toWorker.send(<String, Object?>{
|
||||
'type': 'release',
|
||||
'id': id,
|
||||
'which': decoder && entropy
|
||||
? 'both'
|
||||
: decoder
|
||||
? 'decoder'
|
||||
: 'entropy',
|
||||
});
|
||||
return completer.future;
|
||||
}
|
||||
|
||||
/// Cooperative cancel.
|
||||
///
|
||||
/// The worker sets a flag that the backend polls at stage boundaries. A
|
||||
/// backend sitting inside one long blocking native call cannot be
|
||||
/// interrupted this way — for a hard stop, call [dispose], which kills the
|
||||
/// isolate outright and frees the model.
|
||||
void cancel() {
|
||||
if (_disposed) return;
|
||||
_toWorker.send(const <String, Object?>{'type': 'cancel'});
|
||||
}
|
||||
|
||||
void _handleMessage(Map<String, Object?> message) {
|
||||
final id = message['id'];
|
||||
final job = id is int ? _pending[id] : null;
|
||||
switch (message['type']) {
|
||||
case 'progress':
|
||||
final value = message['value'];
|
||||
if (value is num) {
|
||||
job?.onProgress?.call(value.toDouble().clamp(0.0, 1.0));
|
||||
}
|
||||
case 'result':
|
||||
if (id is int) _pending.remove(id);
|
||||
final bytes = message['bytes'];
|
||||
if (bytes is Uint8List) {
|
||||
job?.completer.complete(bytes);
|
||||
} else {
|
||||
job?.completer.completeError(
|
||||
StateError('Image codec returned no bytes.'),
|
||||
);
|
||||
}
|
||||
case 'released':
|
||||
if (id is int) {
|
||||
final release = _pendingReleases.remove(id);
|
||||
if (release != null && !release.isCompleted) {
|
||||
release.complete();
|
||||
}
|
||||
}
|
||||
case 'error':
|
||||
if (id is int) {
|
||||
_pending.remove(id);
|
||||
// A release can fail too; never leave its caller hanging.
|
||||
final release = _pendingReleases.remove(id);
|
||||
if (release != null && !release.isCompleted) {
|
||||
release.complete();
|
||||
}
|
||||
}
|
||||
job?.completer.completeError(
|
||||
StateError(message['error']?.toString() ?? 'Image codec failed.'),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
void _failAll(Object error) {
|
||||
final jobs = _pending.values.toList();
|
||||
_pending.clear();
|
||||
for (final job in jobs) {
|
||||
if (!job.completer.isCompleted) {
|
||||
job.completer.completeError(error);
|
||||
}
|
||||
}
|
||||
// Releases resolve rather than fail: the isolate dying IS the memory being
|
||||
// freed, which is all the caller was waiting for.
|
||||
final releases = _pendingReleases.values.toList();
|
||||
_pendingReleases.clear();
|
||||
for (final release in releases) {
|
||||
if (!release.isCompleted) {
|
||||
release.complete();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Kills the isolate and frees the model's resident memory.
|
||||
Future<void> dispose() async {
|
||||
if (_disposed) return;
|
||||
_disposed = true;
|
||||
_toWorker.send(const <String, Object?>{'type': 'shutdown'});
|
||||
// Give the worker one event-loop turn to release the native session
|
||||
// cleanly, then take the memory back regardless.
|
||||
await Future<void>.delayed(const Duration(milliseconds: 50));
|
||||
_isolate.kill(priority: Isolate.immediate);
|
||||
_failAll(StateError('Image codec session disposed.'));
|
||||
_fromWorker.close();
|
||||
_errors.close();
|
||||
}
|
||||
}
|
||||
|
||||
class _PendingJob {
|
||||
final Completer<Uint8List> completer = Completer<Uint8List>();
|
||||
final void Function(double progress)? onProgress;
|
||||
|
||||
_PendingJob({this.onProgress});
|
||||
}
|
||||
|
||||
/// The positional payload handed to the codec worker isolate.
|
||||
///
|
||||
/// Built and parsed in ONE place on purpose. When these were two hand-written
|
||||
/// lists, `entropyDecodeGraphPath` was added to the bundle but never to the
|
||||
/// list, so the worker rebuilt a bundle that could not decode — while
|
||||
/// `canDecode`, computed on the main isolate where the path existed, said it
|
||||
/// could. Every decode threw `ImageCodecBundleIncomplete` and told the user to
|
||||
/// re-download a model that was already correct. The whole suite stayed green,
|
||||
/// because both halves were individually right.
|
||||
///
|
||||
/// APPEND new fields; never insert. The list is positional and the casts are
|
||||
/// permissive enough that a shifted `tablesPath` would be read as the rate
|
||||
/// point rather than throwing.
|
||||
List<Object?> _bootPayload(
|
||||
ImageCodecBundle bundle,
|
||||
SendPort reply,
|
||||
RootIsolateToken? rootToken,
|
||||
) => <Object?>[
|
||||
reply, // 0
|
||||
bundle.decoderGraphPath, // 1
|
||||
bundle.entropyGraphPath, // 2
|
||||
bundle.tablesPath, // 3
|
||||
bundle.ratePoint.wireValue, // 4
|
||||
rootToken, // 5
|
||||
bundle.entropyDecodeGraphPath, // 6
|
||||
];
|
||||
|
||||
/// Rebuilds the bundle from [_bootPayload]. Tolerates a short list so a
|
||||
/// truncated or older message degrades to "cannot decode" instead of throwing a
|
||||
/// RangeError inside the isolate, where it would surface as an opaque spawn
|
||||
/// failure.
|
||||
ImageCodecBundle _bundleFromBootPayload(List<Object?> boot) => ImageCodecBundle(
|
||||
decoderGraphPath: boot[1] as String,
|
||||
entropyGraphPath: boot[2] as String?,
|
||||
tablesPath: boot[3] as String?,
|
||||
ratePoint: parseAeicRatePoint(boot[4] as int? ?? -1),
|
||||
entropyDecodeGraphPath: boot.length > 6 ? boot[6] as String? : null,
|
||||
);
|
||||
|
||||
/// Test seams for the boot payload. Not for production use — see
|
||||
/// `test/services/image_codec_boot_payload_test.dart`, which exists because
|
||||
/// this connection has broken twice.
|
||||
@visibleForTesting
|
||||
List<Object?> debugBootPayloadFor(ImageCodecBundle bundle) =>
|
||||
_bootPayload(bundle, _NullSendPort(), null);
|
||||
|
||||
@visibleForTesting
|
||||
ImageCodecBundle debugBundleFromBootPayload(List<Object?> boot) =>
|
||||
_bundleFromBootPayload(boot);
|
||||
|
||||
/// Stand-in so [debugBootPayloadFor] needs no live isolate.
|
||||
class _NullSendPort implements SendPort {
|
||||
@override
|
||||
void send(Object? message) {}
|
||||
|
||||
@override
|
||||
bool operator ==(Object other) => identical(this, other);
|
||||
|
||||
@override
|
||||
int get hashCode => 0;
|
||||
}
|
||||
|
||||
/// Entry point of the codec worker isolate.
|
||||
Future<void> _codecWorkerMain(List<Object?> boot) async {
|
||||
final reply = boot[0] as SendPort;
|
||||
final bundle = _bundleFromBootPayload(boot);
|
||||
final rootToken = boot[5] as RootIsolateToken?;
|
||||
|
||||
if (rootToken == null) {
|
||||
reply.send(<String, Object?>{
|
||||
'type': 'fatal',
|
||||
'error':
|
||||
'RootIsolateToken was unavailable, so the ONNX plugin channel cannot '
|
||||
'be reached from this isolate. ImageCodecSession.spawn() must be '
|
||||
'called from the root isolate.',
|
||||
});
|
||||
return;
|
||||
}
|
||||
// Without this the backend's very first invokeMethod throws. Must happen
|
||||
// before createImageCodecBackend(), which constructs the plugin wrapper.
|
||||
BackgroundIsolateBinaryMessenger.ensureInitialized(rootToken);
|
||||
|
||||
final backend = createImageCodecBackend();
|
||||
if (backend == null) {
|
||||
reply.send(<String, Object?>{
|
||||
'type': 'fatal',
|
||||
'unimplemented': true,
|
||||
'error':
|
||||
'no inference backend is compiled into this build '
|
||||
'(see lib/services/image_codec_backend.dart)',
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
// The bitstream path is a seam so that image_codec_backend.dart can compile
|
||||
// for web, where dart:io does not exist. This worker isolate is native-only,
|
||||
// so it is the correct place to close it -- and it must happen BEFORE
|
||||
// backend.load(), because load() reads supportsBitstreamCodec to decide
|
||||
// whether the codec can encode or decode at all. Left unassigned, the whole
|
||||
// entropy path is dead even with kImageCodecBitstreamPathAvailable true.
|
||||
imageCodecRansCoderBuilder ??= (path) async =>
|
||||
AeicRansCoders(EntropyTables.parse(await File(path).readAsBytes()));
|
||||
|
||||
try {
|
||||
await backend.load(bundle);
|
||||
} catch (error) {
|
||||
reply.send(<String, Object?>{'type': 'fatal', 'error': error.toString()});
|
||||
await backend.dispose();
|
||||
return;
|
||||
}
|
||||
|
||||
final commands = ReceivePort();
|
||||
var cancelRequested = false;
|
||||
var queue = Future<void>.value();
|
||||
|
||||
reply.send(<String, Object?>{
|
||||
'type': 'ready',
|
||||
'port': commands.sendPort,
|
||||
'backend': backend.name,
|
||||
// `bitstream` is the backend's own answer (does this build have the entropy
|
||||
// path at all?); `entropy`/`tables` describe what THIS INSTALL supplied. A
|
||||
// pre-bundle install reports bitstream:true, tables:false — a state the
|
||||
// service must surface as "re-download", not as "your build cannot".
|
||||
'bitstream': backend.supportsBitstreamCodec,
|
||||
'entropy': bundle.entropyGraphPath != null,
|
||||
'tables': bundle.tablesPath != null,
|
||||
});
|
||||
|
||||
commands.listen((message) {
|
||||
if (message is! Map) return;
|
||||
final map = message.cast<String, Object?>();
|
||||
switch (map['type']) {
|
||||
case 'cancel':
|
||||
// Handled on the event loop, so it lands between backend stages.
|
||||
cancelRequested = true;
|
||||
case 'shutdown':
|
||||
commands.close();
|
||||
unawaited(backend.dispose());
|
||||
case 'release':
|
||||
// Queued behind any running job: tearing an ORT session down while it
|
||||
// is executing is a native crash, not an exception.
|
||||
queue = queue.then((_) async {
|
||||
final id = map['id'] as int?;
|
||||
try {
|
||||
final which = map['which'];
|
||||
if (which == 'decoder' || which == 'both') {
|
||||
await backend.releaseDecoderSession();
|
||||
}
|
||||
if (which == 'entropy' || which == 'both') {
|
||||
await backend.releaseEntropySession();
|
||||
}
|
||||
reply.send(<String, Object?>{'type': 'released', 'id': id});
|
||||
} catch (error) {
|
||||
reply.send(<String, Object?>{
|
||||
'type': 'error',
|
||||
'id': id,
|
||||
'error': error.toString(),
|
||||
});
|
||||
}
|
||||
});
|
||||
case 'job':
|
||||
queue = queue.then((_) async {
|
||||
cancelRequested = false;
|
||||
final id = map['id'] as int;
|
||||
try {
|
||||
final ratePoint = parseAeicRatePoint(
|
||||
map['rate'] as int? ?? kShippingAeicRatePoint.wireValue,
|
||||
);
|
||||
final resolution = map['size'] as int? ?? kImageCodecSquareSize;
|
||||
void onProgress(double value) {
|
||||
reply.send(<String, Object?>{
|
||||
'type': 'progress',
|
||||
'id': id,
|
||||
'value': value,
|
||||
});
|
||||
}
|
||||
|
||||
bool shouldCancel() => cancelRequested;
|
||||
|
||||
final Uint8List result;
|
||||
switch (map['op']) {
|
||||
case 'decodeLatent':
|
||||
result = await backend.decodeLatentToRgb(
|
||||
yHat: map['latents'] as Float32List,
|
||||
onProgress: onProgress,
|
||||
shouldCancel: shouldCancel,
|
||||
);
|
||||
case 'encode':
|
||||
result = await backend.encode(
|
||||
rgbBytes: map['bytes'] as Uint8List,
|
||||
ratePoint: ratePoint,
|
||||
resolution: resolution,
|
||||
onProgress: onProgress,
|
||||
shouldCancel: shouldCancel,
|
||||
);
|
||||
case 'decode':
|
||||
result = await backend.decode(
|
||||
bitstream: map['bytes'] as Uint8List,
|
||||
ratePoint: ratePoint,
|
||||
resolution: resolution,
|
||||
onProgress: onProgress,
|
||||
shouldCancel: shouldCancel,
|
||||
);
|
||||
default:
|
||||
throw StateError('Unknown codec op: ${map['op']}');
|
||||
}
|
||||
reply.send(<String, Object?>{
|
||||
'type': 'result',
|
||||
'id': id,
|
||||
'bytes': result,
|
||||
});
|
||||
} catch (error, stackTrace) {
|
||||
reply.send(<String, Object?>{
|
||||
'type': 'error',
|
||||
'id': id,
|
||||
'error': error.toString(),
|
||||
'stack': stackTrace.toString(),
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
import 'dart:typed_data';
|
||||
|
||||
import '../models/image_codec_support.dart';
|
||||
|
||||
/// Web stand-in. There is no isolate and no native runtime on web, so the
|
||||
/// session cannot exist; `ImageCodecService` gates on `kIsWeb` long before it
|
||||
/// would reach here.
|
||||
class ImageCodecSession {
|
||||
ImageCodecSession._();
|
||||
|
||||
static Future<ImageCodecSession> spawn(ImageCodecBundle bundle) async {
|
||||
throw UnsupportedError('The image codec is not supported on web.');
|
||||
}
|
||||
|
||||
String get backendName => 'unsupported';
|
||||
|
||||
bool get supportsBitstreamCodec => false;
|
||||
|
||||
bool get hasEntropyGraph => false;
|
||||
|
||||
bool get hasTables => false;
|
||||
|
||||
Future<void> release({bool decoder = false, bool entropy = false}) async {}
|
||||
|
||||
Future<Uint8List> decodeLatent(
|
||||
Float32List yHat, {
|
||||
void Function(double progress)? onProgress,
|
||||
}) async {
|
||||
throw UnsupportedError('The image codec is not supported on web.');
|
||||
}
|
||||
|
||||
Future<Uint8List> encode(
|
||||
Uint8List rgbBytes,
|
||||
AeicRatePoint ratePoint,
|
||||
int resolution, {
|
||||
void Function(double progress)? onProgress,
|
||||
}) async {
|
||||
throw UnsupportedError('The image codec is not supported on web.');
|
||||
}
|
||||
|
||||
Future<Uint8List> decode(
|
||||
Uint8List bitstream,
|
||||
AeicRatePoint ratePoint,
|
||||
int resolution, {
|
||||
void Function(double progress)? onProgress,
|
||||
}) async {
|
||||
throw UnsupportedError('The image codec is not supported on web.');
|
||||
}
|
||||
|
||||
void cancel() {}
|
||||
|
||||
Future<void> dispose() async {}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
import 'dart:convert';
|
||||
|
||||
import '../models/image_codec_support.dart';
|
||||
import '../storage/prefs_manager.dart';
|
||||
import '../utils/app_logger.dart';
|
||||
|
||||
/// Persistence seam for [ImageCodecPreferences].
|
||||
///
|
||||
/// WHY THIS EXISTS. The app stack spec puts these five fields on `AppSettings`
|
||||
/// (`imageCodecEnabled`, `imageCodecSelectedModelId`, `imageCodecModelSourceUrl`,
|
||||
/// `imageCodecRatePoint`, `imageCodecDownloadedModels`) with matching setters on
|
||||
/// `AppSettingsService` — exactly like the translation block. That is still the
|
||||
/// right end state, but `app_settings.dart` and `app_settings_service.dart` are
|
||||
/// outside this workstream's file set, so the service talks to this interface
|
||||
/// instead of `AppSettingsService`.
|
||||
///
|
||||
/// Migration is mechanical: add the fields to `AppSettings` (the JSON keys in
|
||||
/// [ImageCodecPreferences.toJson] are already the `image_codec_*` snake_case
|
||||
/// names), then replace [PrefsImageCodecSettingsStore] with a thin adapter over
|
||||
/// `AppSettingsService`. No `ImageCodecService` code changes.
|
||||
abstract class ImageCodecSettingsStore {
|
||||
ImageCodecPreferences get preferences;
|
||||
|
||||
/// Loads persisted preferences. Safe to call more than once.
|
||||
Future<void> load();
|
||||
|
||||
Future<void> save(ImageCodecPreferences preferences);
|
||||
}
|
||||
|
||||
/// Default store: one SharedPreferences key holding the whole block as JSON.
|
||||
class PrefsImageCodecSettingsStore implements ImageCodecSettingsStore {
|
||||
static const String _key = 'image_codec_settings';
|
||||
|
||||
ImageCodecPreferences _preferences = const ImageCodecPreferences();
|
||||
bool _loaded = false;
|
||||
|
||||
@override
|
||||
ImageCodecPreferences get preferences => _preferences;
|
||||
|
||||
@override
|
||||
Future<void> load() async {
|
||||
if (_loaded) return;
|
||||
_loaded = true;
|
||||
try {
|
||||
final raw = PrefsManager.instance.getString(_key);
|
||||
if (raw == null) return;
|
||||
final json = jsonDecode(raw);
|
||||
if (json is Map<String, dynamic>) {
|
||||
_preferences = ImageCodecPreferences.fromJson(json);
|
||||
}
|
||||
} catch (error) {
|
||||
// Matches AppSettingsService.loadSettings: a corrupt blob falls back to
|
||||
// defaults rather than blocking startup.
|
||||
appLogger.warn('Image codec settings load failed: $error');
|
||||
_preferences = const ImageCodecPreferences();
|
||||
}
|
||||
}
|
||||
|
||||
@override
|
||||
Future<void> save(ImageCodecPreferences preferences) async {
|
||||
_preferences = preferences;
|
||||
_loaded = true;
|
||||
try {
|
||||
await PrefsManager.instance.setString(_key, jsonEncode(preferences.toJson()));
|
||||
} catch (error) {
|
||||
appLogger.warn('Image codec settings save failed: $error');
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Non-persisting store for tests and widget previews.
|
||||
class InMemoryImageCodecSettingsStore implements ImageCodecSettingsStore {
|
||||
ImageCodecPreferences _preferences;
|
||||
|
||||
InMemoryImageCodecSettingsStore([
|
||||
this._preferences = const ImageCodecPreferences(),
|
||||
]);
|
||||
|
||||
@override
|
||||
ImageCodecPreferences get preferences => _preferences;
|
||||
|
||||
@override
|
||||
Future<void> load() async {}
|
||||
|
||||
@override
|
||||
Future<void> save(ImageCodecPreferences preferences) async {
|
||||
_preferences = preferences;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,498 @@
|
||||
// Pure-Dart port of the AEIC rANS entropy coder.
|
||||
//
|
||||
// Ported from the C++ reference in `aeic/src/cpp/rans/` (rans.h, rans_byte.h,
|
||||
// rans.cpp), which is itself ryg_rans + CompressAI's unbounded-index range
|
||||
// coding. The port must be BYTE-IDENTICAL to the C++ coder: a single differing
|
||||
// byte desynchronises rANS and silently corrupts most of the image. The golden
|
||||
// vectors under `test/services/golden/` pin that down.
|
||||
//
|
||||
// Notes on the port:
|
||||
// * Dart `int` is 64-bit and signed; the C++ state is `uint32_t`. Every place
|
||||
// the C++ relies on 32-bit behaviour is either provably in-range (the state
|
||||
// invariant keeps `x < 2^31`) or masked explicitly with `& 0xFFFFFFFF`.
|
||||
// * `>>` on a negative Dart int is arithmetic. All shifted values here are
|
||||
// non-negative by construction; `>>>` is used where the shift consumes the
|
||||
// state so the intent is unambiguous.
|
||||
//
|
||||
// Format recap (see `results/rans_port_spec.md`):
|
||||
// * precision 16, bypassPrecision 2, RANS_L = 1 << 23, streamParts 2.
|
||||
// * Each `encodeWithIndexes` call is split evenly across the sub-streams:
|
||||
// part p covers `[p * (n ~/ parts), ...)`, the last part taking the
|
||||
// remainder. Every part accumulates across all calls; flush happens once.
|
||||
// * A sub-stream's first 4 bytes are the final rANS state, little-endian.
|
||||
// * Container: `flag = ((nParts - 1) << 4) | (hdrLen == 2 ? 1 : 0)`, then the
|
||||
// lengths of the first `nParts - 1` sub-streams (hdrLen bytes each, LE),
|
||||
// then the sub-streams back to back.
|
||||
library;
|
||||
|
||||
import 'dart:typed_data';
|
||||
|
||||
import 'entropy_tables.dart';
|
||||
|
||||
/// Lower bound of the rANS normalisation interval (`RANS_BYTE_L`).
|
||||
const int kRansLowerBound = 1 << 23;
|
||||
|
||||
/// Thrown when a bitstream cannot be interpreted.
|
||||
class RansFormatException implements Exception {
|
||||
RansFormatException(this.message);
|
||||
|
||||
final String message;
|
||||
|
||||
@override
|
||||
String toString() => 'RansFormatException: $message';
|
||||
}
|
||||
|
||||
/// Splits the shipped bytes into rANS sub-streams. Mirror of
|
||||
/// `RansDecoder::set_stream`.
|
||||
List<Uint8List> parseRansContainer(Uint8List stream) {
|
||||
if (stream.isEmpty) {
|
||||
throw RansFormatException('empty stream');
|
||||
}
|
||||
final int flag = stream[0];
|
||||
final int nStreams = (flag >> 4) + 1;
|
||||
final int hdr = (flag & 0x0F) == 1 ? 2 : 4;
|
||||
var off = 1;
|
||||
final List<int> sizes = <int>[];
|
||||
var total = 0;
|
||||
for (var i = 0; i < nStreams - 1; i++) {
|
||||
if (off + hdr > stream.length) {
|
||||
throw RansFormatException('truncated sub-stream size table');
|
||||
}
|
||||
var sz = 0;
|
||||
for (var b = 0; b < hdr; b++) {
|
||||
sz |= stream[off + b] << (8 * b);
|
||||
}
|
||||
off += hdr;
|
||||
sizes.add(sz);
|
||||
total += sz;
|
||||
}
|
||||
final int last = stream.length - off - total;
|
||||
if (last < 0) {
|
||||
throw RansFormatException('sub-stream sizes exceed the stream length');
|
||||
}
|
||||
sizes.add(last);
|
||||
final List<Uint8List> parts = <Uint8List>[];
|
||||
var p = off;
|
||||
for (final int sz in sizes) {
|
||||
if (p + sz > stream.length) {
|
||||
throw RansFormatException('truncated sub-stream');
|
||||
}
|
||||
parts.add(Uint8List.sublistView(stream, p, p + sz));
|
||||
p += sz;
|
||||
}
|
||||
return parts;
|
||||
}
|
||||
|
||||
/// Assembles sub-streams into the bytes that go on the air.
|
||||
Uint8List buildRansContainer(List<Uint8List> parts) {
|
||||
if (parts.isEmpty) {
|
||||
throw RansFormatException('no sub-streams');
|
||||
}
|
||||
if (parts.length > 16) {
|
||||
throw RansFormatException('too many sub-streams (${parts.length})');
|
||||
}
|
||||
var maximum = 0;
|
||||
for (var i = 0; i < parts.length - 1; i++) {
|
||||
if (parts[i].length > maximum) maximum = parts[i].length;
|
||||
}
|
||||
final int hdr = maximum > 65535 ? 4 : 2;
|
||||
final int flag = ((parts.length - 1) << 4) | (hdr == 2 ? 1 : 0);
|
||||
var total = 1 + hdr * (parts.length - 1);
|
||||
for (final Uint8List p in parts) {
|
||||
total += p.length;
|
||||
}
|
||||
final Uint8List out = Uint8List(total);
|
||||
out[0] = flag;
|
||||
var off = 1;
|
||||
for (var i = 0; i < parts.length - 1; i++) {
|
||||
final int n = parts[i].length;
|
||||
for (var b = 0; b < hdr; b++) {
|
||||
out[off + b] = (n >> (8 * b)) & 0xFF;
|
||||
}
|
||||
off += hdr;
|
||||
}
|
||||
for (final Uint8List p in parts) {
|
||||
out.setRange(off, off + p.length, p);
|
||||
off += p.length;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/// Growable (start, range) entry list, stored interleaved in one Int32List.
|
||||
/// `range == 0` is the bypass sentinel: `start` then carries the raw bits.
|
||||
class _EntryBuffer {
|
||||
Int32List _data = Int32List(2048);
|
||||
int _length = 0;
|
||||
|
||||
int get length => _length;
|
||||
|
||||
void add(int start, int range) {
|
||||
final int need = (_length + 1) * 2;
|
||||
if (need > _data.length) {
|
||||
final Int32List bigger = Int32List(_data.length * 2);
|
||||
bigger.setRange(0, _data.length, _data);
|
||||
_data = bigger;
|
||||
}
|
||||
_data[_length * 2] = start;
|
||||
_data[_length * 2 + 1] = range;
|
||||
_length++;
|
||||
}
|
||||
|
||||
int startAt(int i) => _data[i * 2];
|
||||
|
||||
int rangeAt(int i) => _data[i * 2 + 1];
|
||||
|
||||
void clear() => _length = 0;
|
||||
}
|
||||
|
||||
/// Growable byte sink used for the reversed encoder output.
|
||||
class _ByteSink {
|
||||
Uint8List _data = Uint8List(4096);
|
||||
int _length = 0;
|
||||
|
||||
void add(int byte) {
|
||||
if (_length == _data.length) {
|
||||
final Uint8List bigger = Uint8List(_data.length * 2);
|
||||
bigger.setRange(0, _data.length, _data);
|
||||
_data = bigger;
|
||||
}
|
||||
_data[_length++] = byte & 0xFF;
|
||||
}
|
||||
|
||||
/// Returns the bytes in reverse of the order they were added, which is the
|
||||
/// order they appear in the sub-stream.
|
||||
Uint8List reversedBytes() {
|
||||
final Uint8List out = Uint8List(_length);
|
||||
for (var i = 0; i < _length; i++) {
|
||||
out[i] = _data[_length - 1 - i];
|
||||
}
|
||||
return out;
|
||||
}
|
||||
}
|
||||
|
||||
/// rANS encoder. Call [encodeWithIndexes] once per stage in format order
|
||||
/// (z, y0, y1, y2, y3), then [finish] exactly once.
|
||||
class RansEncoder {
|
||||
factory RansEncoder(EntropyTables tables, {int? streamParts}) =>
|
||||
RansEncoder._(tables, streamParts ?? tables.streamParts);
|
||||
|
||||
RansEncoder._(EntropyTables tables, int streamParts)
|
||||
: tables = tables,
|
||||
streamParts = streamParts,
|
||||
_precision = tables.precision,
|
||||
_bypassPrecision = tables.bypassPrecision,
|
||||
_entries = List<_EntryBuffer>.generate(
|
||||
streamParts,
|
||||
(_) => _EntryBuffer(),
|
||||
);
|
||||
|
||||
final EntropyTables tables;
|
||||
final int streamParts;
|
||||
final int _precision;
|
||||
final int _bypassPrecision;
|
||||
final List<_EntryBuffer> _entries;
|
||||
bool _finished = false;
|
||||
|
||||
int get _maxBypassVal => (1 << _bypassPrecision) - 1;
|
||||
|
||||
/// Accumulates one encode call. [symbols] and [indexes] must be the same
|
||||
/// length; an index `< 0` emits nothing at all.
|
||||
///
|
||||
/// Never pass an odd-length array: the reference splitter (and therefore the
|
||||
/// on-air format) mis-sizes the last part's index vector in that case.
|
||||
void encodeWithIndexes(
|
||||
List<int> symbols,
|
||||
List<int> indexes,
|
||||
int cdfGroupIndex,
|
||||
) {
|
||||
if (_finished) {
|
||||
throw StateError('RansEncoder.finish() has already been called');
|
||||
}
|
||||
if (symbols.length != indexes.length) {
|
||||
throw ArgumentError(
|
||||
'symbols (${symbols.length}) and indexes (${indexes.length}) differ',
|
||||
);
|
||||
}
|
||||
if (cdfGroupIndex < 0 || cdfGroupIndex >= tables.groups.length) {
|
||||
throw ArgumentError('no CDF group $cdfGroupIndex');
|
||||
}
|
||||
final CdfGroup group = tables.groups[cdfGroupIndex];
|
||||
final int total = symbols.length;
|
||||
final int each = total ~/ streamParts;
|
||||
for (var p = 0; p < streamParts; p++) {
|
||||
final int lo = p * each;
|
||||
final int hi = p == streamParts - 1 ? total : lo + each;
|
||||
_push(_entries[p], symbols, indexes, lo, hi, group);
|
||||
}
|
||||
}
|
||||
|
||||
void _push(
|
||||
_EntryBuffer out,
|
||||
List<int> symbols,
|
||||
List<int> indexes,
|
||||
int lo,
|
||||
int hi,
|
||||
CdfGroup group,
|
||||
) {
|
||||
final Int32List cdf = group.quantizedCdf;
|
||||
final Int32List cdfLength = group.cdfLength;
|
||||
final Int32List offsets = group.offset;
|
||||
final int width = group.cdfWidth;
|
||||
final int maxBypassVal = _maxBypassVal;
|
||||
final int bypassPrecision = _bypassPrecision;
|
||||
|
||||
for (var i = lo; i < hi; i++) {
|
||||
final int cdfIdx = indexes[i];
|
||||
if (cdfIdx < 0) {
|
||||
continue;
|
||||
}
|
||||
if (cdfIdx >= group.numCdfs) {
|
||||
throw RansFormatException(
|
||||
'index $cdfIdx out of range (${group.numCdfs} CDF rows)',
|
||||
);
|
||||
}
|
||||
final int maxValue = cdfLength[cdfIdx] - 2;
|
||||
var value = symbols[i] - offsets[cdfIdx];
|
||||
var rawVal = 0;
|
||||
if (value < 0) {
|
||||
rawVal = -2 * value - 1;
|
||||
value = maxValue;
|
||||
} else if (value >= maxValue) {
|
||||
rawVal = 2 * (value - maxValue);
|
||||
value = maxValue;
|
||||
}
|
||||
|
||||
final int base = cdfIdx * width;
|
||||
final int start = cdf[base + value];
|
||||
out.add(start, cdf[base + value + 1] - start);
|
||||
|
||||
if (value == maxValue) {
|
||||
// Bypass mode: raw bits, `bypassPrecision` at a time.
|
||||
var nBypass = 0;
|
||||
while ((rawVal >> (nBypass * bypassPrecision)) != 0) {
|
||||
nBypass++;
|
||||
}
|
||||
var val = nBypass;
|
||||
while (val >= maxBypassVal) {
|
||||
out.add(maxBypassVal, 0);
|
||||
val -= maxBypassVal;
|
||||
}
|
||||
out.add(val, 0);
|
||||
for (var j = 0; j < nBypass; j++) {
|
||||
out.add((rawVal >> (j * bypassPrecision)) & maxBypassVal, 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Flushes every sub-stream and returns the container bytes.
|
||||
Uint8List finish() {
|
||||
if (_finished) {
|
||||
throw StateError('RansEncoder.finish() has already been called');
|
||||
}
|
||||
_finished = true;
|
||||
final List<Uint8List> parts = <Uint8List>[
|
||||
for (final _EntryBuffer e in _entries) _flush(e),
|
||||
];
|
||||
return buildRansContainer(parts);
|
||||
}
|
||||
|
||||
/// Discards accumulated entries so the encoder can be reused.
|
||||
void reset() {
|
||||
for (final _EntryBuffer e in _entries) {
|
||||
e.clear();
|
||||
}
|
||||
_finished = false;
|
||||
}
|
||||
|
||||
Uint8List _flush(_EntryBuffer entries) {
|
||||
final _ByteSink sink = _ByteSink();
|
||||
// The state is always in [2^23, 2^31); native Dart ints need no masking,
|
||||
// but the emission path is written to stay explicitly byte-wise anyway.
|
||||
var x = kRansLowerBound;
|
||||
final int bypassXMax = (1 << (_precision - _bypassPrecision)) << 15;
|
||||
for (var k = entries.length - 1; k >= 0; k--) {
|
||||
final int range = entries.rangeAt(k);
|
||||
final int start = entries.startAt(k);
|
||||
if (range != 0) {
|
||||
final int xMax = range << 15;
|
||||
while (x >= xMax) {
|
||||
sink.add(x & 0xFF);
|
||||
x = x >>> 8;
|
||||
}
|
||||
x = ((x ~/ range) << _precision) + (x % range) + start;
|
||||
} else {
|
||||
while (x >= bypassXMax) {
|
||||
sink.add(x & 0xFF);
|
||||
x = x >>> 8;
|
||||
}
|
||||
x = ((x << _bypassPrecision) | start) & 0xFFFFFFFF;
|
||||
}
|
||||
}
|
||||
// RansEncFlush writes the 32-bit state little-endian at the front of the
|
||||
// stream, i.e. in emission order it is the high byte first.
|
||||
sink.add((x >>> 24) & 0xFF);
|
||||
sink.add((x >>> 16) & 0xFF);
|
||||
sink.add((x >>> 8) & 0xFF);
|
||||
sink.add(x & 0xFF);
|
||||
return sink.reversedBytes();
|
||||
}
|
||||
}
|
||||
|
||||
/// rANS decoder. Decoding is INCREMENTAL: construct once over the whole
|
||||
/// container, then call [decodeStream] per stage in the same order the encoder
|
||||
/// used (z, y0, y1, y2, y3). Each call resumes the sub-stream states where the
|
||||
/// previous one left off, because stage i's indexes are unknown until the
|
||||
/// earlier stages have been decoded and run back through the network.
|
||||
class RansDecoder {
|
||||
factory RansDecoder(
|
||||
EntropyTables tables,
|
||||
Uint8List stream, {
|
||||
int? streamParts,
|
||||
}) => RansDecoder._(tables, stream, streamParts ?? tables.streamParts);
|
||||
|
||||
RansDecoder._(EntropyTables tables, Uint8List stream, int expected)
|
||||
: tables = tables,
|
||||
_precision = tables.precision,
|
||||
_bypassPrecision = tables.bypassPrecision,
|
||||
_parts = parseRansContainer(stream) {
|
||||
if (_parts.length != expected) {
|
||||
throw RansFormatException(
|
||||
'container has ${_parts.length} sub-streams, expected $expected',
|
||||
);
|
||||
}
|
||||
_states = List<int>.filled(_parts.length, 0);
|
||||
_ptrs = List<int>.filled(_parts.length, 0);
|
||||
for (var i = 0; i < _parts.length; i++) {
|
||||
final Uint8List p = _parts[i];
|
||||
if (p.length < 4) {
|
||||
throw RansFormatException('sub-stream $i is shorter than 4 bytes');
|
||||
}
|
||||
_states[i] = p[0] | (p[1] << 8) | (p[2] << 16) | (p[3] << 24);
|
||||
_ptrs[i] = 4;
|
||||
}
|
||||
}
|
||||
|
||||
final EntropyTables tables;
|
||||
final int _precision;
|
||||
final int _bypassPrecision;
|
||||
final List<Uint8List> _parts;
|
||||
late final List<int> _states;
|
||||
late final List<int> _ptrs;
|
||||
|
||||
int get streamParts => _parts.length;
|
||||
|
||||
/// Decodes one stage. Returns one symbol per entry of [indexes]; positions
|
||||
/// whose index is `< 0` yield a literal 0 and consume nothing.
|
||||
Int16List decodeStream(List<int> indexes, int cdfGroupIndex) {
|
||||
if (cdfGroupIndex < 0 || cdfGroupIndex >= tables.groups.length) {
|
||||
throw ArgumentError('no CDF group $cdfGroupIndex');
|
||||
}
|
||||
final CdfGroup group = tables.groups[cdfGroupIndex];
|
||||
final Int32List cdf = group.quantizedCdf;
|
||||
final Int32List cdfLength = group.cdfLength;
|
||||
final Int32List offsets = group.offset;
|
||||
final int width = group.cdfWidth;
|
||||
final int mask = (1 << _precision) - 1;
|
||||
final int bypassPrecision = _bypassPrecision;
|
||||
final int maxBypassVal = (1 << bypassPrecision) - 1;
|
||||
final int bypassMask = maxBypassVal;
|
||||
|
||||
final int total = indexes.length;
|
||||
final int nParts = _parts.length;
|
||||
final int each = total ~/ nParts;
|
||||
final Int16List out = Int16List(total);
|
||||
|
||||
for (var pi = 0; pi < nParts; pi++) {
|
||||
final int lo = pi * each;
|
||||
final int hi = pi == nParts - 1 ? total : lo + each;
|
||||
final Uint8List buf = _parts[pi];
|
||||
var x = _states[pi];
|
||||
var ptr = _ptrs[pi];
|
||||
|
||||
for (var i = lo; i < hi; i++) {
|
||||
final int cdfIdx = indexes[i];
|
||||
if (cdfIdx < 0) {
|
||||
out[i] = 0;
|
||||
continue;
|
||||
}
|
||||
if (cdfIdx >= group.numCdfs) {
|
||||
throw RansFormatException(
|
||||
'index $cdfIdx out of range (${group.numCdfs} CDF rows)',
|
||||
);
|
||||
}
|
||||
final int n = cdfLength[cdfIdx];
|
||||
final int maxValue = n - 2;
|
||||
final int base = cdfIdx * width;
|
||||
final int cum = x & mask;
|
||||
|
||||
// upper_bound(row[0:n], cum) - 1
|
||||
var loo = 0;
|
||||
var hii = n;
|
||||
while (loo < hii) {
|
||||
final int mid = (loo + hii) >> 1;
|
||||
if (cdf[base + mid] > cum) {
|
||||
hii = mid;
|
||||
} else {
|
||||
loo = mid + 1;
|
||||
}
|
||||
}
|
||||
final int s = loo - 1;
|
||||
if (s < 0 || s >= n - 1) {
|
||||
throw RansFormatException('corrupt stream: symbol $s out of range');
|
||||
}
|
||||
final int start = cdf[base + s];
|
||||
final int range = cdf[base + s + 1] - start;
|
||||
|
||||
x = (range * (x >>> _precision) + (x & mask) - start) & 0xFFFFFFFF;
|
||||
while (x < kRansLowerBound) {
|
||||
if (ptr >= buf.length) {
|
||||
throw RansFormatException('sub-stream $pi exhausted');
|
||||
}
|
||||
x = ((x << 8) | buf[ptr]) & 0xFFFFFFFF;
|
||||
ptr++;
|
||||
}
|
||||
|
||||
var value = s;
|
||||
if (value == maxValue) {
|
||||
// Bypass mode. Note the renormalisation here is a single `if`, not a
|
||||
// loop -- that asymmetry with the symbol path is part of the format.
|
||||
int getBits() {
|
||||
final int v = x & bypassMask;
|
||||
x = x >>> bypassPrecision;
|
||||
if (x < kRansLowerBound) {
|
||||
if (ptr >= buf.length) {
|
||||
throw RansFormatException('sub-stream $pi exhausted');
|
||||
}
|
||||
x = ((x << 8) | buf[ptr]) & 0xFFFFFFFF;
|
||||
ptr++;
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
var val = getBits();
|
||||
var nBypass = val;
|
||||
while (val == maxBypassVal) {
|
||||
val = getBits();
|
||||
nBypass += val;
|
||||
}
|
||||
var rawVal = 0;
|
||||
for (var j = 0; j < nBypass; j++) {
|
||||
rawVal |= getBits() << (j * bypassPrecision);
|
||||
}
|
||||
value = rawVal >> 1;
|
||||
if ((rawVal & 1) != 0) {
|
||||
value = -value - 1;
|
||||
} else {
|
||||
value += maxValue;
|
||||
}
|
||||
}
|
||||
out[i] = value + offsets[cdfIdx];
|
||||
}
|
||||
|
||||
_states[pi] = x;
|
||||
_ptrs[pi] = ptr;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
/// Platform-appropriate [ReceivedImageBlobStore], chosen at compile time.
|
||||
///
|
||||
/// `received_image_blob_store_io.dart` imports `dart:io` and
|
||||
/// `package:path_provider`, so `main.dart` cannot import it directly without
|
||||
/// breaking the web build. Same conditional-export shape as
|
||||
/// `image_codec_file_store.dart`.
|
||||
library;
|
||||
|
||||
export 'received_image_blob_store_factory_stub.dart'
|
||||
if (dart.library.io) 'received_image_blob_store_factory_io.dart';
|
||||
@@ -0,0 +1,10 @@
|
||||
import 'received_image_blob_store_io.dart';
|
||||
import 'received_image_store.dart';
|
||||
|
||||
/// File-backed store under the application support directory.
|
||||
///
|
||||
/// Received images have to survive a restart: the sidecar record is the only
|
||||
/// evidence that a message *was* an image, so an in-memory store loses the
|
||||
/// bubble as well as the pixels.
|
||||
ReceivedImageBlobStore createReceivedImageBlobStore() =>
|
||||
FileReceivedImageBlobStore();
|
||||
@@ -0,0 +1,8 @@
|
||||
import 'received_image_store.dart';
|
||||
|
||||
/// Web fallback: images live only for the lifetime of the tab.
|
||||
///
|
||||
/// The codec cannot run on web anyway (`ImageCodecService.availability` is
|
||||
/// `unavailable` there), so nothing is lost that could have been rendered.
|
||||
ReceivedImageBlobStore createReceivedImageBlobStore() =>
|
||||
InMemoryReceivedImageBlobStore();
|
||||
@@ -0,0 +1,284 @@
|
||||
/// File-backed [ReceivedImageBlobStore] for every platform that has `dart:io`.
|
||||
///
|
||||
/// Without this, `ReceivedImageStore` falls back to
|
||||
/// [InMemoryReceivedImageBlobStore] and every received image is forgotten at
|
||||
/// the next app launch — the sidecars go with the bytes, so even the "this
|
||||
/// message was an image" record disappears.
|
||||
///
|
||||
/// ## Layout
|
||||
///
|
||||
/// ```
|
||||
/// <application support>/received_images/
|
||||
/// <streamId>.aeic the ~156 B bitstream
|
||||
/// <streamId>.png the decoded 512x512 PNG (~400 KB)
|
||||
/// <streamId>.json the sidecar record (ReceivedImageEntry.toJson)
|
||||
/// <streamId>.json.tmp transient; only during an atomic sidecar write
|
||||
/// ```
|
||||
///
|
||||
/// Application *support*, not documents: these are derived caches of a chat
|
||||
/// message, not user documents, so they must not show up in the iOS Files app
|
||||
/// or be swept into an iCloud backup.
|
||||
///
|
||||
/// ## Guarantees
|
||||
///
|
||||
/// * Sidecar writes are atomic (tmp file + rename). A kill mid-write can
|
||||
/// leave `<id>.json.tmp` but never a truncated `<id>.json`, so
|
||||
/// [readSidecars] never has to defend against half a JSON object.
|
||||
/// * Every method is total: a missing file reads as null and deletes are
|
||||
/// idempotent. I/O errors are swallowed and reported through the return
|
||||
/// value, because a failed write must degrade the image, not crash the
|
||||
/// receive path.
|
||||
/// * [readSidecars] sweeps orphans — `.aeic`/`.png`/`.tmp` files with no
|
||||
/// surviving `.json` — so a crash between "write bitstream" and "write
|
||||
/// sidecar" cannot leak bytes that no budget accounts for.
|
||||
/// * Stream ids are validated before they are ever concatenated into a path;
|
||||
/// a hostile `../../` id cannot escape the directory.
|
||||
library;
|
||||
|
||||
import 'dart:convert';
|
||||
import 'dart:io';
|
||||
import 'dart:typed_data';
|
||||
|
||||
import 'package:flutter/foundation.dart' show debugPrint;
|
||||
import 'package:path_provider/path_provider.dart';
|
||||
|
||||
import 'received_image_store.dart';
|
||||
|
||||
/// Persists received-image bytes under the application support directory.
|
||||
class FileReceivedImageBlobStore implements ReceivedImageBlobStore {
|
||||
/// Only `[A-Za-z0-9_-]`, so a stream id can never contain `/`, `\` or `..`.
|
||||
static final RegExp _safeId = RegExp(r'^[A-Za-z0-9_-]{1,64}$');
|
||||
|
||||
static const String _bitstreamExt = '.aeic';
|
||||
static const String _pngExt = '.png';
|
||||
static const String _sidecarExt = '.json';
|
||||
static const String _tmpExt = '.json.tmp';
|
||||
|
||||
/// Overridable so tests (and any future "move my data" feature) can point the
|
||||
/// store at a temp directory instead of the real support directory.
|
||||
final Future<Directory> Function() _baseDirectory;
|
||||
|
||||
final String directoryName;
|
||||
|
||||
/// Resolved lazily and then cached, because [pngPath] has to be synchronous
|
||||
/// (the share sheet needs a path inside a build/callback) while
|
||||
/// `path_provider` is not.
|
||||
String? _dirPath;
|
||||
Future<String>? _pending;
|
||||
|
||||
FileReceivedImageBlobStore({
|
||||
Future<Directory> Function()? baseDirectory,
|
||||
this.directoryName = 'received_images',
|
||||
}) : _baseDirectory = baseDirectory ?? getApplicationSupportDirectory;
|
||||
|
||||
/// Resolves and creates the directory. Safe to call repeatedly; concurrent
|
||||
/// callers share one future so the directory is created exactly once.
|
||||
Future<String> ensureReady() {
|
||||
final cached = _dirPath;
|
||||
if (cached != null) return Future<String>.value(cached);
|
||||
final pending = _pending;
|
||||
if (pending != null) return pending;
|
||||
final started = _resolve();
|
||||
_pending = started;
|
||||
return started.whenComplete(() => _pending = null);
|
||||
}
|
||||
|
||||
Future<String> _resolve() async {
|
||||
final base = await _baseDirectory();
|
||||
final dir = Directory('${base.path}${Platform.pathSeparator}$directoryName');
|
||||
if (!dir.existsSync()) {
|
||||
await dir.create(recursive: true);
|
||||
}
|
||||
_dirPath = dir.path;
|
||||
return dir.path;
|
||||
}
|
||||
|
||||
Future<String?> _pathFor(String streamId, String extension) async {
|
||||
if (!_safeId.hasMatch(streamId)) return null;
|
||||
final dir = await ensureReady();
|
||||
return '$dir${Platform.pathSeparator}$streamId$extension';
|
||||
}
|
||||
|
||||
Future<Uint8List?> _read(String streamId, String extension) async {
|
||||
try {
|
||||
final path = await _pathFor(streamId, extension);
|
||||
if (path == null) return null;
|
||||
final file = File(path);
|
||||
if (!file.existsSync()) return null;
|
||||
return await file.readAsBytes();
|
||||
} catch (error) {
|
||||
debugPrint('received_image_blob_store: read $streamId$extension: $error');
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
Future<void> _write(String streamId, String extension, Uint8List bytes) async {
|
||||
try {
|
||||
final path = await _pathFor(streamId, extension);
|
||||
if (path == null) {
|
||||
debugPrint('received_image_blob_store: refusing unsafe id "$streamId"');
|
||||
return;
|
||||
}
|
||||
await File(path).writeAsBytes(bytes, flush: true);
|
||||
} catch (error) {
|
||||
debugPrint('received_image_blob_store: write $streamId$extension: $error');
|
||||
}
|
||||
}
|
||||
|
||||
Future<void> _delete(String streamId, String extension) async {
|
||||
try {
|
||||
final path = await _pathFor(streamId, extension);
|
||||
if (path == null) return;
|
||||
final file = File(path);
|
||||
if (file.existsSync()) await file.delete();
|
||||
} catch (error) {
|
||||
debugPrint('received_image_blob_store: delete $streamId$extension: $error');
|
||||
}
|
||||
}
|
||||
|
||||
Future<int?> _size(String streamId, String extension) async {
|
||||
try {
|
||||
final path = await _pathFor(streamId, extension);
|
||||
if (path == null) return null;
|
||||
final file = File(path);
|
||||
if (!file.existsSync()) return null;
|
||||
return await file.length();
|
||||
} catch (_) {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
@override
|
||||
Future<void> writeBitstream(String streamId, Uint8List bytes) =>
|
||||
_write(streamId, _bitstreamExt, bytes);
|
||||
|
||||
@override
|
||||
Future<Uint8List?> readBitstream(String streamId) =>
|
||||
_read(streamId, _bitstreamExt);
|
||||
|
||||
@override
|
||||
Future<void> deleteBitstream(String streamId) =>
|
||||
_delete(streamId, _bitstreamExt);
|
||||
|
||||
@override
|
||||
Future<int?> bitstreamSize(String streamId) => _size(streamId, _bitstreamExt);
|
||||
|
||||
@override
|
||||
Future<void> writePng(String streamId, Uint8List bytes) =>
|
||||
_write(streamId, _pngExt, bytes);
|
||||
|
||||
@override
|
||||
Future<Uint8List?> readPng(String streamId) => _read(streamId, _pngExt);
|
||||
|
||||
@override
|
||||
Future<void> deletePng(String streamId) => _delete(streamId, _pngExt);
|
||||
|
||||
@override
|
||||
Future<int?> pngSize(String streamId) => _size(streamId, _pngExt);
|
||||
|
||||
/// Atomic: write `<id>.json.tmp`, then rename over `<id>.json`. `rename` is
|
||||
/// atomic within a directory on every platform we ship, so a reader either
|
||||
/// sees the old record or the new one.
|
||||
@override
|
||||
Future<void> writeSidecar(String streamId, String json) async {
|
||||
try {
|
||||
final finalPath = await _pathFor(streamId, _sidecarExt);
|
||||
final tmpPath = await _pathFor(streamId, _tmpExt);
|
||||
if (finalPath == null || tmpPath == null) {
|
||||
debugPrint('received_image_blob_store: refusing unsafe id "$streamId"');
|
||||
return;
|
||||
}
|
||||
final tmp = File(tmpPath);
|
||||
await tmp.writeAsString(json, flush: true);
|
||||
await tmp.rename(finalPath);
|
||||
} catch (error) {
|
||||
debugPrint('received_image_blob_store: sidecar $streamId: $error');
|
||||
}
|
||||
}
|
||||
|
||||
@override
|
||||
Future<void> deleteSidecar(String streamId) async {
|
||||
await _delete(streamId, _sidecarExt);
|
||||
await _delete(streamId, _tmpExt);
|
||||
}
|
||||
|
||||
/// Startup scan. Also the only place orphaned bytes are reaped: anything
|
||||
/// whose sidecar is gone is invisible to the store and therefore to the
|
||||
/// disk budget, so it would live forever.
|
||||
@override
|
||||
Future<Map<String, String>> readSidecars() async {
|
||||
final result = <String, String>{};
|
||||
late final String dirPath;
|
||||
try {
|
||||
dirPath = await ensureReady();
|
||||
} catch (error) {
|
||||
debugPrint('received_image_blob_store: no directory: $error');
|
||||
return result;
|
||||
}
|
||||
final dir = Directory(dirPath);
|
||||
if (!dir.existsSync()) return result;
|
||||
|
||||
final files = <File>[];
|
||||
try {
|
||||
files.addAll(dir.listSync().whereType<File>());
|
||||
} catch (error) {
|
||||
debugPrint('received_image_blob_store: list failed: $error');
|
||||
return result;
|
||||
}
|
||||
|
||||
final orphans = <File>[];
|
||||
for (final file in files) {
|
||||
final name = file.uri.pathSegments.last;
|
||||
if (name.endsWith(_tmpExt)) {
|
||||
// A kill during a sidecar write. The previous `.json` (if any) is
|
||||
// still authoritative.
|
||||
orphans.add(file);
|
||||
continue;
|
||||
}
|
||||
if (!name.endsWith(_sidecarExt)) continue;
|
||||
final id = name.substring(0, name.length - _sidecarExt.length);
|
||||
if (!_safeId.hasMatch(id)) continue;
|
||||
try {
|
||||
final raw = await file.readAsString();
|
||||
// Cheap sanity gate: the store's own jsonDecode would skip it anyway,
|
||||
// but a truncated record here means we should reap its bytes too.
|
||||
jsonDecode(raw);
|
||||
result[id] = raw;
|
||||
} catch (error) {
|
||||
debugPrint('received_image_blob_store: bad sidecar $id: $error');
|
||||
}
|
||||
}
|
||||
|
||||
for (final file in files) {
|
||||
final name = file.uri.pathSegments.last;
|
||||
if (name.endsWith(_sidecarExt) || name.endsWith(_tmpExt)) continue;
|
||||
final dot = name.lastIndexOf('.');
|
||||
if (dot <= 0) continue;
|
||||
final id = name.substring(0, dot);
|
||||
final ext = name.substring(dot);
|
||||
if (ext != _bitstreamExt && ext != _pngExt) continue;
|
||||
if (result.containsKey(id)) continue;
|
||||
orphans.add(file);
|
||||
}
|
||||
|
||||
for (final file in orphans) {
|
||||
try {
|
||||
if (file.existsSync()) await file.delete();
|
||||
} catch (_) {
|
||||
// Best effort; a locked file is retried on the next launch.
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
/// Null until [ensureReady] (or any async method) has resolved the directory.
|
||||
/// In practice `ReceivedImageStore.load()` runs at startup and resolves it
|
||||
/// long before any bubble asks for a path.
|
||||
@override
|
||||
String? pngPath(String streamId) {
|
||||
final dir = _dirPath;
|
||||
if (dir == null) return null;
|
||||
if (!_safeId.hasMatch(streamId)) return null;
|
||||
return '$dir${Platform.pathSeparator}$streamId$_pngExt';
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user