[RNE Rewrite] feat(tti): drop 8da4w, add SDXS CoreML fp16 + MLX int4

Remove the XNNPACK 8da4w SDXS variant (no speed win over fp32 since both
are conv-bound, and noticeably poorer quality on this quant-sensitive
model). Add iOS backends: CoreML fp16 (quality-preserving) and MLX int4.

- models.textToImage.SDXS_512_DREAMSHAPER: XNNPACK_FP32 (default),
  COREML_FP16, MLX_INT4
- demo screen exposes the CoreML/MLX options (iOS-only)

Refs #1244, #939

Author: msluszniak

apps/computer-vision/app/textToImage/index.tsx

@@ -22,8 +22,14 @@ const MODEL_OPTIONS: ModelOption[] = [
     value: models.textToImage.SDXS_512_DREAMSHAPER.XNNPACK_FP32,
   },
   {
-    label: 'SDXS-512-DreamShaper (XNNPACK 8DA4W)',
-    value: models.textToImage.SDXS_512_DREAMSHAPER.XNNPACK_8DA4W,
+    label: 'SDXS-512-DreamShaper (CoreML FP16)',
+    value: models.textToImage.SDXS_512_DREAMSHAPER.COREML_FP16,
+    disabled: Platform.OS !== 'ios',
+  },
+  {
+    label: 'SDXS-512-DreamShaper (MLX INT4)',
+    value: models.textToImage.SDXS_512_DREAMSHAPER.MLX_INT4,
+    disabled: Platform.OS !== 'ios',
   },
 ];
 
@@ -36,9 +42,13 @@ function TextToImageContent() {
   const [latency, setLatency] = useState<number | null>(null);
   const [error, setError] = useState<string | null>(null);
 
-  const { isReady, downloadProgress, error: loadError, generate, generateWorklet } = useTextToImage(
-    selectedModel
-  );
+  const {
+    isReady,
+    downloadProgress,
+    error: loadError,
+    generate,
+    generateWorklet,
+  } = useTextToImage(selectedModel);
 
   const runGenerate = async (sync: boolean) => {
     if (!generate || !generateWorklet || !prompt.trim()) return;

packages/react-native-executorch/src/extensions/cv/tasks/sdxsTextToImage.ts

@@ -38,10 +38,14 @@ export type SdxsOptions = {
   readonly initNoiseSigma: number;
   /** The single training timestep fed to the UNet for the distilled step. */
   readonly timestep: number;
-  /** Cumulative alpha product at {@link timestep} (from the scheduler schedule). */
-  readonly alphaCumprod: number;
-  /** What the UNet output represents, which fixes the scheduler update formula. */
-  readonly predictionType: 'epsilon' | 'v_prediction' | 'sample';
+  /**
+   * Scheduler-step coefficients. For the distilled single step the DEIS update
+   * is exactly linear in the latents and the UNet output:
+   * `clean = sampleCoeff * latents + noiseCoeff * modelOutput`. Both are pinned
+   * from the reference scheduler during export.
+   */
+  readonly sampleCoeff: number;
+  readonly noiseCoeff: number;
   /** Per-channel or scalar scale applied to the decoder output when mapping to `[0..255]`. */
   readonly outAlpha: number | number[];
   /** Per-channel or scalar bias applied to the decoder output when mapping to `[0..255]`. */
@@ -101,23 +105,17 @@ function seededGaussian(size: number, seed: number): Float32Array {
 }
 
 // Single-step scheduler update: turns the UNet output into the clean latent that
-// is fed to the decoder. For a distilled 1-step model the update to t=0 reduces
-// to recovering x0 from the model prediction.
+// is fed to the decoder. For the distilled single step the DEIS update is exactly
+// linear in the latents and the model output.
 function toCleanLatents(
   latents: Float32Array,
   modelOutput: Float32Array,
   opts: SdxsOptions
 ): Float32Array {
   'worklet';
-  if (opts.predictionType === 'sample') return modelOutput;
-  const sqrtAlpha = Math.sqrt(opts.alphaCumprod);
-  const sqrtBeta = Math.sqrt(1 - opts.alphaCumprod);
   const out = new Float32Array(latents.length);
   for (let i = 0; i < latents.length; i++) {
-    out[i] =
-      opts.predictionType === 'v_prediction'
-        ? sqrtAlpha * latents[i]! - sqrtBeta * modelOutput[i]!
-        : (latents[i]! - sqrtBeta * modelOutput[i]!) / sqrtAlpha;
+    out[i] = opts.sampleCoeff * latents[i]! + opts.noiseCoeff * modelOutput[i]!;
   }
   return out;
 }

packages/react-native-executorch/src/models.ts

@@ -594,18 +594,19 @@ const CLIP_VIT_BASE_PATCH32_IMAGE_XNNPACK_INT8: ImageEmbedderModel = {
 // =============================================================================
 // Text to Image
 // =============================================================================
-// NOTE: the numeric fields below (initNoiseSigma, timestep, alphaCumprod,
-// predictionType, outAlpha/outBeta) are pinned during the `.pte` export by
-// numerically matching the reference diffusers output; the values here are the
-// SDXS defaults and are updated once the export is finalized.
+// The numeric fields below are pinned from the SDXS reference pipeline
+// (DEISMultistepScheduler, single step at t=999, epsilon prediction, TAESD
+// decoder). The single-step scheduler update is exactly linear in the latents
+// and the UNet output, so it is captured as (sampleCoeff, noiseCoeff). The
+// exported `decode` method emits RGB in [0,1], hence outAlpha=255, outBeta=0.
 const SDXS_512_DREAMSHAPER_OPTS = {
   imageSize: 512,
   latentChannels: 4,
   numInferenceSteps: 1,
   initNoiseSigma: 1.0,
   timestep: 999,
-  alphaCumprod: 0.00466,
-  predictionType: 'epsilon' as const,
+  sampleCoeff: 14.642591,
+  noiseCoeff: -14.579279,
   outAlpha: 255.0,
   outBeta: 0.0,
 };
@@ -615,8 +616,13 @@ const SDXS_512_DREAMSHAPER_XNNPACK_FP32: SdxsModel = {
   tokenizerPath: SDXS_512_DREAMSHAPER_TOKENIZER,
   opts: SDXS_512_DREAMSHAPER_OPTS,
 };
-const SDXS_512_DREAMSHAPER_XNNPACK_8DA4W: SdxsModel = {
-  modelPath: `${BASE_URL}-sdxs-512-dreamshaper/${NEXT_VERSION_TAG}/xnnpack/sdxs_512_dreamshaper_xnnpack_8da4w.pte`,
+const SDXS_512_DREAMSHAPER_COREML_FP16: SdxsModel = {
+  modelPath: `${BASE_URL}-sdxs-512-dreamshaper/${NEXT_VERSION_TAG}/coreml/sdxs_512_dreamshaper_coreml_fp16.pte`,
+  tokenizerPath: SDXS_512_DREAMSHAPER_TOKENIZER,
+  opts: SDXS_512_DREAMSHAPER_OPTS,
+};
+const SDXS_512_DREAMSHAPER_MLX_INT4: SdxsModel = {
+  modelPath: `${BASE_URL}-sdxs-512-dreamshaper/${NEXT_VERSION_TAG}/mlx/sdxs_512_dreamshaper_mlx_int4.pte`,
   tokenizerPath: SDXS_512_DREAMSHAPER_TOKENIZER,
   opts: SDXS_512_DREAMSHAPER_OPTS,
 };
@@ -880,7 +886,8 @@ export const models = {
     SDXS_512_DREAMSHAPER: {
       ...SDXS_512_DREAMSHAPER_XNNPACK_FP32,
       XNNPACK_FP32: SDXS_512_DREAMSHAPER_XNNPACK_FP32,
-      XNNPACK_8DA4W: SDXS_512_DREAMSHAPER_XNNPACK_8DA4W,
+      COREML_FP16: SDXS_512_DREAMSHAPER_COREML_FP16,
+      MLX_INT4: SDXS_512_DREAMSHAPER_MLX_INT4,
     },
   },
 };