notes.md

# Notes / lab log

Running log of what Claude Opus 4.8 in Claude Code figured out. Newest at the bottom of each section.

## Environment
- macOS (darwin arm64), Apple Silicon. No CUDA → torch CPU/MPS build.
- System python 3.9.6; using `uv` to manage an isolated env.
- git-lfs 3.7.1 present. Weights cloned to `/tmp/Moebius/Moebius-weights`
  (pretrained, ft_celebahq, ft_ffhq, ft_places2 — each ~450MB fp32 `.bin`).
- Code repo at `/tmp/Moebius/Moebius`.

## Code map (what matters for the port)
- Entry: `infer/infer_moebius.py` → `infer/utils.py:build_pipeline`.
- Pipeline: `removal/v1_2/pipeline.py` (`RemovalSDXLPipeline_BatchMode`).
- Model wrapper: `removal/v1_2/removal_model.py` (`RemovalModel` = embedding + diff UNet).
- Core helpers: `utils_infer.py` (`encode_clean_latents`, `predict_noise`).
- UNet impl: `model_lib/nets/unet_lambda_prune_lite.py` (+ lambda layers under
  `model_lib/nets/layers/λ/vanillaλ.py`).
- Config: `config/model_cfg/moebius.yaml`.

## Key findings
- The CUDA/Triton `fla` dependency is ONLY imported in `model_lib/nets/layers/gla/gla.py`
  (the GLA teacher variant). Moebius's student UNet (lambda-DWConv) does not need it —
  must avoid importing `unet_gla` to keep the graph clean for export.
- "Prompt" conditioning is a plain `nn.Embedding(20, 3072)`. CFG uses fixed ids:
  cond=[0..9], uncond=[10..19]. So encoder_hidden_states is a constant per branch →
  can be precomputed and baked into the ONNX UNet as a constant, OR passed as input.
- 9-channel UNet input = cat([noisy_latents(4), resized_mask(1), masked_latents(4)], dim=1).
- CFG batches uncond+cond into one forward (batch dim ×2), then splits.
- einsum is used in the λ layers (linear attention). Supported in ONNX; need to check
  ORT-Web WebGPU coverage.

## Phase 1 results (reference inference — DONE)
- Got the real pipeline running end-to-end on CPU (macOS, torch 2.7.1).
- Patches needed to load student on CPU/mac:
  - `model_lib/__init__.py`: wrapped teacher `unet_gla` import in try/except (needs `fla`).
  - Don't import `utils_train` (drags in orjson/library); `build_vae` is just
    `AutoencoderKL.from_pretrained(vae_dir)`.
- **VAE scaling_factor = 0.13025** (NOT the usual SD 0.18215!). Custom VAE.
  block_out_channels = [128,256,512,512] → vae_scale_factor 8. This MUST be hardcoded
  correctly in the JS port or colors/contrast will be wrong.
- removal_model params = 226.04M confirmed. load_state_dict: all keys matched.
- Perf: ~8.9 s/step on CPU (×19 steps + CFG ×2 = 38 UNet passes ≈ 2:48 total). WebGPU
  expected far faster. Confirms CPU/WASM is unusable; WebGPU is the whole game.
- Output saved to reference_out/reference_result.png — plausible inpaint. Mask convention:
  white(255) → 1 → region to inpaint (zeroed in masked_image); black → keep.
- num_inference_steps=20 with strength=0.99 → DDIM uses 19 steps (drops first).

## Parity strategy
- Won't try to reproduce torch RNG in JS. For PyTorch↔ONNX parity: dump identical input
  tensors and compare outputs. For the web app: generate noise with a seedable JS RNG;
  diffusion is robust to the particular noise draw, so visual results will be valid even
  if not bit-identical to the torch reference.
- DDIM `scale_model_input` is identity → skip in TS. Need to reproduce DDIM alphas/betas
  (scaled_linear, beta 0.00085→0.012, 1000 steps) and the DDIM step update in JS.

## Architecture: spatial size is FIXED (important!)
- Self-attn (attn1): MQSλ with `r=15` → local-context path (Conv3d pos_conv). Spatially
  dynamic, fine at any size.
- Cross-attn (attn2): MQCλ with NO `r` → global path → `rel_pos_emb` is an
  `nn.Parameter(n*n, m, dim_k, dim_u)` where n = per-block sample_size, m = 10. This is
  TIED to the trained spatial resolution. Different spatial size → wrong/oob indexing.
- ⇒ Export at STATIC 512×512 image (64×64 latent). Web app resizes user input to 512×512,
  inpaints, resizes result back + pastes. Square only. This is the benchmark resolution.

## ONNX export plan
- Three graphs, spatial static, batch dynamic where cheap:
  - vae_encoder: (B,3,512,512) → moments (B,8,64,64); JS uses mean=moments[:,:4]*sf.
  - unet (RemovalModel): (B,9,64,64), timesteps(B,), input_ids(B,10) → noise(B,4,64,64).
    Embedding (nn.Embedding 20×3072) stays IN the graph (cheap gather). CFG batches B=2.
  - vae_decoder: (B,4,64,64) → (B,3,512,512).
- scaling_factor = 0.13025 applied in JS (encode: latent*sf; decode: latent/sf).

## Phase 2 results (ONNX export — DONE)
- torch.onnx.export (legacy tracer, opset 18) traced all 3 graphs cleanly. No op-coverage
  failures. The einsum/lambda/Conv3d ops all exported.
- Parity vs PyTorch (CPU EP): decoder 5.7e-5, unet 3.6e-6, encoder mean ch ~2e-2.
- FULL pipeline parity test (python/onnx_pipeline.py): reimplemented DDIM+CFG+9ch+scaling
  in numpy on the ONNX sessions, vs torch models with identical noise:
    final latents max|Δ| 0.149, decoded image mean|Δ| 0.0022, max 0.090 → visually identical.
  This validates the ENTIRE orchestration I'll port to TS.
- numpy DDIM vs diffusers DDIMScheduler: step max|Δ| 5e-7, timesteps identical. ✓

## DDIM constants for JS (validated)
- betas = linspace(sqrt(0.00085), sqrt(0.012), 1000)^2 ; alphas_cumprod = cumprod(1-betas)
- timesteps(20 steps) = [950,900,...,50,0]; strength 0.99 ⇒ drop first ⇒ [900,...,0] (19).
- ddim_step (eta=0, clip_sample=False):
    pred_x0 = (sample - sqrt(1-ac_t)*eps) / sqrt(ac_t)
    prev    = sqrt(ac_prev)*pred_x0 + sqrt(1-ac_prev)*eps
    ac_prev = alphas_cumprod[prev_t], or final_alpha_cumprod=1.0 when prev_t<0 (last step).
- noise_offset 0.0357: noise += 0.0357 * randn(B,4,1,1). (optional; small)

## Web pipeline recipe (numpy → TS)
1. resize image+mask to 512×512 (mask NEAREST, binarize ≥128).
2. img→[-1,1] CHW; masked = img*(1-mask).
3. encode img & masked → moments; take mean[:4] * 0.13025.
4. mask→64×64 NEAREST, 1ch.
5. latents = randn(1,4,64,64) [+ noise_offset].
6. loop t in timesteps: nine=cat([latents,mask64,maskedLat]); batch×2; unet;
   cfg = u + g*(c-u); latents = ddim_step.
7. decode(latents/0.13025); (x+1)/2; clip; → image.
8. paste: out*blur(mask) + (1-blur(mask))*orig.

## Phase 3 (web app) — in progress
- Vite + TS + onnxruntime-web (1.27.0). Default `onnxruntime-web/webgpu` import resolves
  to the self-contained `ort.webgpu.bundle.min.mjs`.
- Models served LOCALLY: web/public/models -> ../models symlink, at /models/*.onnx.
  Total ~1.24GB fetched over localhost (no internet). ORT runtime served from
  web/ort-dist at /ort/* via a custom static middleware (see vite.config.ts).
- BUG FIXED: ORT glue .mjs must NOT be in /public (Vite tries to module-transform it).
  Fix = serve /ort/* as raw static files via configureServer middleware.
- COOP/COEP headers set (needed for threaded WASM / SharedArrayBuffer).

## WebGPU op coverage (confirmed from ORT source at /tmp/Moebius/onnxruntime)
- js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts registers: Einsum ✓, Conv ✓
  (conv.ts has computeConv3DInfo / createConv3DNaiveProgramInfo → Conv3d pos_conv works,
  naive kernel so possibly slow), InstanceNormalization ✓, MatMul/Gemm ✓, Softmax ✓,
  Reduce* ✓, Transpose/Concat/Gather/Pad/Resize/Where ✓.
- GroupNorm: not registered by that name, BUT torch.onnx exports nn.GroupNorm as a
  Reshape→InstanceNormalization→Reshape→Mul→Add decomposition → covered. (VAE decoder
  CPU-EP parity was 5.7e-5, so the graph is decomposed, not a single GroupNorm op.)
- ⇒ No expected silent CPU fallback for the heavy ops. Confirm empirically in console.

## Verification without a GPU browser (sandbox can't drive Chrome — user's live Chrome
## holds the playwright profile)
- web/test/fixture/*.bin: dumped inputs + reference final latents from the validated
  numpy/ONNX pipeline, to check the TS port (ddim.ts + 9ch assembly) in Node.

## TODO / unknowns
- fp16 export to ~450MB UNet for real deployment (VAE fp16 unstable: decoder Cast issue;
  keep VAE fp32). Quality risk in λ layers — validate before shipping fp16.
- Confirm in-browser: WebGPU selected, no CPU fallback, end-to-end correctness + timing.