From 9c592c30c1b4c87be75c5fdea25c9d8c2fb25592 Mon Sep 17 00:00:00 2001 From: ashish-aesthisia Date: Tue, 11 Aug 2026 06:53:47 +0000 Subject: [PATCH 1/4] feat(minimax-h3): portable LoRA keys, pruned and fp8 checkpoints, chunked fuse --- README.md | 551 +++++++++++------- TRAINING.md | 8 +- core/CLAUDE.md | 16 + core/src/inline_core/models/keymap.py | 31 +- core/src/inline_core/models/lora.py | 98 +++- .../src/inline_core/models/minimaxh3/adaln.py | 102 ++++ core/src/inline_core/models/minimaxh3/keys.py | 34 +- core/src/inline_core/models/minimaxh3/load.py | 137 ++++- .../inline_core/models/minimaxh3/lora_keys.py | 235 ++++++++ .../inline_core/models/minimaxh3/pipeline.py | 9 + .../models/minimaxh3/requirements.py | 96 ++- core/src/inline_core/training/arch.py | 12 + core/src/inline_core/training/h3.py | 19 + core/src/inline_core/training/trainer.py | 31 +- core/tests/test_lora.py | 43 ++ core/tests/test_minimaxh3_load.py | 153 +++++ core/tests/test_minimaxh3_lora_keys.py | 249 ++++++++ core/tests/test_minimaxh3_nodes.py | 163 +++++- core/tests/test_webui_install.py | 3 + core/uv.lock | 8 +- 20 files changed, 1696 insertions(+), 302 deletions(-) create mode 100644 core/src/inline_core/models/minimaxh3/lora_keys.py create mode 100644 core/tests/test_minimaxh3_lora_keys.py diff --git a/README.md b/README.md index fc36e15..c24198d 100644 --- a/README.md +++ b/README.md @@ -2,52 +2,55 @@

AI filmmaking on a node canvas

-

A free and open-source app for AI filmmaking on a single node canvas. Train your own LoRAs and generate locally on your own GPU, with hosted models when you want them. Every render is kept as a versioned take.

+

Inline Studio is a free, open-source app for AI filmmakers. Generate locally on your own GPU and train your own LoRAs on the same node canvas, with the built-in Inline Core engine and hosted fal models. Build a whole visual pipeline from moodboard to final cut, and every render is kept as a versioned, non-destructive take.

License: GPLv3 Python 3.11+ - Latest release + Latest release Join our Discord

![Inline Studio node canvas showing a generative AI film pipeline with frames, takes, and connectors](https://raw.githubusercontent.com/inlineresearch/Inline-Studio/main/screenshots/screenshot-dashboard-2.png) -[**New here? Start with the getting started guide →**](https://inlinestudio.art/getting-started) +[**New here? Check out our getting started guide →**](https://inlinestudio.art/getting-started) -## Supported models +**Contents:** [What is Inline Studio?](#what-is-inline-studio) · [Get Started](#get-started) · +[Features](#features) · [LoRA training](#lora-training) · [How it works](#how-it-works) · +[Two ways to generate](#two-ways-to-generate) · [Inline Core engine](#inline-core-generation-engine) +([Krea 2](#krea-2), [FLUX.2](#flux2), [MiniMax H3](#minimax-h3), [ControlNet](#controlnet)) · +[API Nodes](#api-nodes) · [FAQ](#faq) · [Contributing](#contributing) - +## What is Inline Studio? -| Model | Train | Generate | 16GB card | -| ----------------------------------------------------------------- | ----- | -------- | --------------- | -| [FLUX.2](https://bfl.ai/blog/flux-2) (klein Base 4B) | yes | yes | ~8.6GB | -| [Krea 2](https://www.krea.ai/) (RAW, 4-bit) | yes | yes | ~11.9GB | -| Z-Image Turbo | yes | yes | ~13.4GB | -| [MiniMax H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) (video) | yes | yes | ~20.6GB, slowly | -| Hosted models (API Nodes) | no | yes | no GPU needed | +Inline Studio is a free, open-source app for **AI filmmaking on a node canvas**, powered by the built-in **Inline Core** engine (local diffusion models) and hosted [fal](https://fal.ai) models. It gives AI filmmakers a free-form canvas to build a whole visual pipeline, from moodboard to final cut. -VRAM is the training peak at 512px. Training is cheaper than generating, and a LoRA trained at 512 -applies at any generation resolution. Full per-card matrix and timings: -[Benchmark results](TRAINING.md#benchmark-results). +- **Non-destructive by default** - every render is kept as a versioned take; generating again adds one, nothing is overwritten. +- **Local diffusion generation engine** - the built-in Inline Core engine runs popular diffusion models locally, on your own GPU, from a single model file, no external server. Currently supported: **Z-Image Turbo**, **Krea 2** (RAW + Turbo), **FLUX.2**, and **MiniMax H3** for video with sound. +- **Train LoRAs locally** - the Trainer canvas fine-tunes Z-Image, Krea 2, FLUX.2 or MiniMax H3 on your own images, on your own GPU. H3 also trains on short video clips, so a LoRA can learn motion and not just look. With a 4-bit base, Krea 2 trains at 512px inside about 12GB, so a 16GB card can train a LoRA for a 26GB model. See [LoRA training](#lora-training). +- **Hosted models via API Nodes** - reach for closed models with no GPU and no setup for instant creative range; see [API Nodes](#api-nodes). +- **Mix both in the same film** - Inline Studio handles everything around the render: exploring options, keeping what works, and shaping a repeatable process you can iterate on and share. -## Install +It runs as a **single process on one port**: the Inline Core engine (Python) serves the web UI _and_ does the generation: `python core/main.py` and open the browser. No desktop install, no separate backend. -You need [Python 3.11+](https://python.org). The web UI ships as a Python package, so there is no -Node step. `--install --extra all` installs everything: the engine, the model runtime, the trainer -and the UI. +**Who it's for:** AI filmmakers, motion artists, and generative creators who want to make AI short films and longer cuts without losing every good version along the way. + +## Get Started + +The built web UI ships as a Python package, so all you need is [Python 3.11+](https://python.org), no Node. **`--install --extra all` is the single command that installs everything** - the engine, the local model runtime, the LoRA trainer, and the UI. On an NVIDIA machine it reads the GPU's compute capability and pulls the CUDA build of PyTorch that has kernels for it, RTX 50-series included. + +[**No GPU? Deploy Inline Studio on RunPod →**](https://console.runpod.io/deploy?template=c0qkyaypuv&ref=hs2l4qhc) **macOS / Linux:** ```bash git clone https://github.com/inlineresearch/Inline-Studio.git cd Inline-Studio/core -./webui.sh --install --extra all -./webui.sh # http://127.0.0.1:8848 +./webui.sh --install --extra all # one command: installs everything +./webui.sh # then run, on http://127.0.0.1:8848 ``` -**Windows** (use `webui.bat`; `webui.sh` is a bash script and will not run in PowerShell): +**Windows** (use `webui.bat` - `webui.sh` is a bash script and won't run in PowerShell; you can also double-click it): ```powershell git clone https://github.com/inlineresearch/Inline-Studio.git @@ -55,295 +58,419 @@ cd Inline-Studio\core .\webui.bat --install --extra all .\webui.bat -rem If the CUDA build is wrong for your card, name the index yourself: +rem If the CUDA build turns out wrong for your card, name the index yourself: .\webui.bat --install --extra all --torch-index cu130 ``` -On NVIDIA, `--install` reads your GPU's compute capability and pulls the matching CUDA build of -PyTorch, RTX 50-series included. Everything lands in `core/.venv`, which Inline Studio owns; an -environment already activated in your shell is never touched. Re-running `--install` is safe. +That's it: `--install` sets up the environment and installs everything once, then `webui.sh` / `webui.bat` runs the app on one port. See **[Command-line options](#command-line-options)** for every flag (`--listen`, `--port`, `--lowvram`, `--multi-gpu`, …). + +Everything lands in `core/.venv`, which Inline Studio owns. If you already have another virtualenv or conda env activated in that shell (a ComfyUI one, say), it is left completely untouched - `--install` says so and carries on. Re-running `--install` is safe: an existing `core/.venv` is reused, so adding an extra later is just another `--install --extra NAME`. -Prefer pip? `pip install -r requirements.txt` from the repo root installs the whole app from PyPI, -then run `inline-studio`. +Prefer pip over the launcher? `pip install -r requirements.txt` (from the repo root) installs the whole app - engine, UI, model runtime, and trainer - from PyPI; then run `inline-studio`. + +### Hardware support
-Hardware support, RTX 50-series, AMD ROCm, Apple Silicon +GPU, CPU, Apple Silicon, and ROCm setup -Honest status, what has actually been run versus what has a code path nobody has verified: +Honest status - what's actually been run, versus what has a code path but no one has verified: -| Hardware | Status | Extra steps | -| ----------------------- | ------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------- | -| **NVIDIA, Linux** | **Tested**, Z-Image Turbo 1024² on a T4 (16GB); Krea 2 1024² and LoRA training on an L40S (48GB) | None | -| **NVIDIA, Windows** | Supported | PyPI's default torch is CPU-only on Windows, so `--install` picks the CUDA build for your card | -| **Apple Silicon (MPS)** | Code path exists, **untested** | None. int8 does not apply on MPS, so a model too big for unified memory will not fit | -| **AMD (ROCm), Linux** | **Untested**, reports welcome | Needs a ROCm build of PyTorch, see below | -| **CPU only** | Works, very slow | `./webui.sh --cpu` | +| Hardware | Status | Extra steps | +| ----------------------- | --------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **NVIDIA, Linux** | **Tested** - Z-Image Turbo 1024² on a T4 (16 GB); Krea 2 1024² and LoRA training on an L40S (48 GB) | None. `webui.sh --install` picks the CUDA build automatically. | +| **NVIDIA, Windows** | Supported, needs one step | Run `.\webui.bat --install` (the Windows launcher). PyPI's default `torch` is **CPU-only on Windows**, so `--install` reads your GPU's compute capability and pulls the matching CUDA build: `cu130` for RTX 50-series (Blackwell), `cu126` for everything older. Override it with `--torch-index` - see [RTX 50-series](#rtx-50-series-blackwell) below. | +| **Apple Silicon (MPS)** | Code path exists, **untested** | None. int8 quantisation doesn't apply on MPS, so a model too big for unified memory won't fit. | +| **AMD (ROCm), Linux** | **Untested** - reports welcome | Needs a ROCm build of PyTorch - see [AMD (ROCm) setup](#amd-rocm-setup) below. | +| **CPU only** | Works, very slow | `./webui.sh --cpu` (Windows: `.\webui.bat --cpu`) | #### RTX 50-series (Blackwell) -RTX 50-series cards are compute capability `sm_120`, and no wheel built for CUDA 12.4 or 12.6 has -kernels for them. `--install` reads the capability off the driver and picks `cu130`, so a plain -`.\webui.bat --install --extra all` is all you need. +RTX 50-series cards (5060/5070/5080/5090 and the RTX PRO Blackwell line) are compute capability **sm_120**, and no PyTorch wheel built for CUDA 12.4 or 12.6 has kernels for them. `--install` handles this: it reads the compute capability off the driver and picks `cu130`, so a plain `.\webui.bat --install --extra all` is all you need. -**Old driver?** CUDA 13 needs driver R580 or newer. If yours predates it, `--install` picks `cu128` -and says so: cu128 still has `sm_120` but is frozen at torch 2.11 and will never update, so update -the driver when you can. +**Old driver?** CUDA 13 needs driver R580 or newer. If yours predates it, `--install` picks `cu128` for you and says so: cu128 still has `sm_120` but is **frozen at torch 2.11** and will never update, so updating the driver and re-running `--install` is worth doing when you can. -`--torch-index` takes a short name (`cu130`, `cu128`, `cu126`), a full index URL, or `cpu`. Naming it -explicitly also **replaces** an already-installed torch, which a plain re-run will not do, so you -rarely need `--recreate`. `INLINE_TORCH_INDEX` does the same thing. +To name an index yourself: -Not sure what you have? `.\webui.bat --print-torch-index` prints what the driver reported and which -index would be used, and installs nothing. That one line is what to paste into a bug report. +```powershell +.\webui.bat --install --extra all --torch-index cu130 -#### AMD (ROCm) +rem Or set it once for the shell, same effect +set INLINE_TORCH_INDEX=cu130 +``` + +`--torch-index` takes a short name (`cu130`, `cu128`, `cu126`), a full index URL, or `cpu` to force the CPU-only build. `webui.sh` takes the same flag. Naming it explicitly also **replaces** an already-installed torch, which a plain re-run will not do, so you rarely need `--recreate`. + +Not sure what you have? `.\webui.bat --print-torch-index` prints what the driver reported and which index would be used, and installs nothing. Paste that into a bug report. If the installed build turns out to have no kernels for your card, Core also says so by name at startup rather than leaving you with PyTorch's own `sm_120 is not compatible` warning. -Nobody has verified this yet, so treat it as a starting point. Install normally **first**, then -replace PyTorch, so nothing can overwrite your ROCm build afterwards: +#### AMD (ROCm) setup + +Nobody has verified Inline Studio on AMD yet, so treat this as a starting point rather than a supported path. Install everything normally **first**, then replace PyTorch with the ROCm build - doing it in this order means nothing can quietly overwrite your ROCm torch afterwards: ```bash cd core -./webui.sh --install --extra runtime +./webui.sh --install --extra runtime # engine + runtime (pulls the default PyPI torch) -# Pick the index matching YOUR ROCm version: https://pytorch.org/get-started/locally/ +# Replace torch with the ROCm build. Pick the index that matches YOUR ROCm version - +# check https://pytorch.org/get-started/locally/ (rocm6.2 shown here as an example). +# --python pins the install to Inline Studio's venv, whatever is activated in your shell. uv pip install --python .venv/bin/python --force-reinstall \ --index-url https://download.pytorch.org/whl/rocm6.2 torch -# hip should print a version, not None +# Verify you actually got a ROCm build (hip should print a version, not None): .venv/bin/python -c "import torch; print(torch.cuda.is_available(), torch.version.hip)" ``` -Do not run `uv sync` or pass `--recreate` afterwards; both put the PyPI torch back over your ROCm -build. The dtype heuristics key off NVIDIA compute capability, which is meaningless on RDNA and -CDNA, so [open an issue](https://github.com/inlineresearch/Inline-Studio/issues) either way. +Then run `./webui.sh` as usual. + +Three gotchas: + +- **Don't run `uv sync` afterwards** - it re-resolves the environment against the lockfile and will pull the PyPI torch back over your ROCm build. Use `uv pip install --python .venv/bin/python` for follow-up installs. The same applies to a hand-picked CUDA index. +- **Don't pass `--recreate`** - it rebuilds `.venv` from scratch and your ROCm torch goes with it. A plain `--install` re-run reuses the venv and is safe. +- ROCm presents itself through `torch.cuda`, so the engine will treat it as a CUDA device and may largely work. But the dtype heuristics key off **NVIDIA** compute capability (`< 8.0` → fp16), which is meaningless on RDNA/CDNA, and the int8 (torchao) path is unverified on ROCm. If it works - or doesn't - [open an issue](https://github.com/inlineresearch/Inline-Studio/issues); that's the fastest way to get AMD properly supported. + +**Known limits, so you can judge before installing:** + +- **Local model coverage is Z-Image Turbo, Krea 2 and FLUX.2** today. SDXL and others are planned; hosted models via [API Nodes](#api-nodes) need no GPU at all. +- **Krea 2 is a 12.9B model and needs a big card to generate.** The bf16 checkpoint is 26 GB on disk, and generation peaks around 36 GB at 1024 with guidance on, so a 40 GB+ GPU is the practical floor for inference. **Training is cheaper than generating**, because the 4-bit base path puts Krea 2 LoRA training at 512 inside 12 GB - see [Benchmark results](TRAINING.md#benchmark-results). Z-Image remains the low-VRAM path for generation. +- **1024² with Guidance (CFG) above 0 needs more than 16 GB.** CFG runs the prompt and negative prompt together, doubling the denoise. Z-Image Turbo is distilled to run CFG-free - at Guidance 0, 1024² fits in ~11.5 GB. FLUX.2 klein 4B peaks near 17.9 GB at bf16, so 24 GB holds it resident and a 16 GB card runs it quantized instead. + +
+ +### From source (for UI development) + +
+Build the UI and run the engine locally + +To hack on the web UI you need [Node.js](https://nodejs.org) 20.11+ as well, and you serve a local SPA build: + +```bash +git clone https://github.com/inlineresearch/Inline-Studio.git && cd Inline-Studio + +# 1. Build the web UI +npm install +npm run build:spa # -> dist-web/ + +# 2. Set up + run the engine, serving your local build +cd core +uv sync --extra server --extra runtime # server + the local model runtime (torch/diffusers) +uv run python main.py --front-end-root ../dist-web +``` + +`uv sync` here manages `core/.venv` as a project environment - it is exact, so it removes anything not in the lockfile (including the `inline-studio-frontend` package a previous `--install` may have added, which doesn't matter when you're serving `--front-end-root ../dist-web`). -#### Generation VRAM, so you can judge before downloading +Then open **http://127.0.0.1:8848**. Add your [fal.ai API key](https://fal.ai/dashboard/keys) in Settings for hosted models, and set up local generation as in [Two ways to generate](#two-ways-to-generate). The canvas and planning work without any models. -Krea 2 is 26GB on disk and generation peaks near 36GB at 1024, so a 40GB card is the practical floor -for **inference**. Training is far cheaper, see the table above. Z-Image Turbo is the low-VRAM path -for generation: it is distilled to run CFG-free, so 1024² fits in about 11.5GB. +**Hot-reload:** run the engine as above, then in another terminal `npm run dev:web` (Vite serves the UI with HMR and proxies API calls to Core).
+### Command-line options + +The friendly launcher (in `core/`) maps flags onto the engine's `INLINE_*` environment knobs: `webui.sh` on macOS/Linux, `webui.bat` on Windows. `core/main.py` takes the same flags when you run the engine directly. `./webui.sh --help` (or `.\webui.bat --help`) lists them all. +
-Command-line options - -`webui.sh` (macOS/Linux) and `webui.bat` (Windows) map friendly flags onto the engine's `INLINE_*` -environment variables. `core/main.py` takes the same flags. Run `--help` for the full list. - -| Flag | Env var | What it does | -| --------------------- | ------------------------ | ------------------------------------------------------------ | -| `--listen` | `INLINE_HOST=0.0.0.0` | Bind all interfaces so other machines can reach it | -| `--port N` | `INLINE_PORT` | Port to serve on (default 8848) | -| `--models-dir PATH` | `INLINE_MODELS_DIR` | Where weights are scanned from (default `./models`) | -| `--data-dir PATH` | `INLINE_DATA_DIR` | Where runs and takes are written | -| `--lowvram` | `INLINE_PROFILE=lowvram` | Tight-VRAM profile (tiling, slicing, int8) | -| `--cpu` | `INLINE_PROFILE=cpu` | Force CPU generation | -| `--vram-budget GB` | `INLINE_VRAM_BUDGET_GB` | Treat the GPU as having GB of usable VRAM | -| `--multi-gpu [SPEC]` | `INLINE_PARALLEL` | Split one image's denoise across GPUs; auto with 2+ GPUs | -| `--torch-index WHICH` | `INLINE_TORCH_INDEX` | With `--install`, override the PyTorch wheel index | -| `--print-torch-index` | n/a | Print the GPU probe and chosen index, then exit | -| `--extra NAME` | n/a | Add an install extra: `runtime`, `server`, `training`, `all` | -| `--recreate` | n/a | Rebuild `.venv` from scratch | -| `--dev` / `--rebuild` | n/a | Live-reload dev loop / force a fresh SPA build | - -**From source (UI development):** build the SPA with `npm ci && npm run build:spa`, then serve it -with `cd core && uv run python main.py --front-end-root ../dist-web`. Or `./webui.sh --dev` for -Vite HMR on `:5173`. +Show all command-line flags + +| `webui.sh` / `main.py` flag | Env var | What it does | +| ---------------------------------- | ------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `--listen` | `INLINE_HOST=0.0.0.0` | Bind all interfaces so other machines can reach it | +| `--host ADDR` | `INLINE_HOST` | Bind a specific address (default `127.0.0.1`) | +| `--port N` | `INLINE_PORT` | Port to serve on (default `8848`) | +| `--models-dir PATH` | `INLINE_MODELS_DIR` | Where model weights are scanned from (default `./models`) | +| `--data-dir PATH` | `INLINE_DATA_DIR` | Where runs + takes are written (default `./.inline`) | +| `--lowvram` | `INLINE_PROFILE=lowvram` | Tight-VRAM profile (VAE tiling/slicing, attention slicing) | +| `--cpu` | `INLINE_PROFILE=cpu` | Force CPU generation | +| `--profile NAME` | `INLINE_PROFILE` | Set the profile explicitly: `gpu-max` \| `lowvram` \| `cpu` | +| `--vram-budget GB` | `INLINE_VRAM_BUDGET_GB` | Treat the GPU as having GB of usable VRAM | +| `--multi-gpu [SPEC]` | `INLINE_PARALLEL` | Split one image's denoise across GPUs (e.g. `pipefusion=2`); auto with 2+ GPUs | +| `--front-end-root DIR` _(main.py)_ | `INLINE_FRONTEND_ROOT` | Serve a local SPA build instead of the installed UI package (dev) | +| `--rebuild` _(webui.sh)_ | n/a | Force a fresh SPA build (`npm run build:spa`) from source and serve it on the one port; use after UI changes when not running `--dev`. Needs the repo checkout + Node/npm | +| `--torch-index WHICH` | `INLINE_TORCH_INDEX` | With `--install`, override the PyTorch wheel index picked from your GPU's compute capability. A short name (`cu130`, `cu128`, `cu126`), a full index URL, or `cpu`. Naming it also replaces an already-installed torch | +| `--print-torch-index` | n/a | Print what the GPU probe read and which index would be used, then exit without installing. The one line to paste into a bug report |
-## Train a LoRA +`webui.sh` also has `--install` / `--extra NAME` to set up the venv, plus `--torch-index WHICH` (`INLINE_TORCH_INDEX`) to override the PyTorch wheel index picked from your GPU's compute capability, `--recreate` (rebuild `.venv` from scratch) and `--use-active-env` (install into / run from the environment activated in your shell instead of `.venv`). New to Inline Studio? The [Getting Started guide](https://inlinestudio.art/getting-started) walks you through your first render. + +## Features -Train on your own images, or on short video clips, on your own GPU with no cloud step. The -**Trainer** tab is a second canvas: wire the nodes, press Start, watch it run. The finished -`.safetensors` lands in `models/loras/`, where the LoRA loader node picks it up, so you can generate -with it in the Studio tab straight away. +- Free-form node canvas +- Versioned, non-destructive takes +- Chain frames into a generative pipeline +- Video editing on the canvas +- Local generation, built in +- Multi-reference composition +- Train your own LoRAs locally +- API Nodes for hosted models +- Community extensions +- Free & open source (GPL-3.0) + +[**Follow our Animated Short Film with LTX 2.3 and GPT Image Generation tutorial →**](https://inlinestudio.art/projects/circuit-race) + +## LoRA training + +Train a LoRA on your own images without leaving the app, on your own GPU, with no cloud step. The **Trainer** tab is a second canvas: wire up the nodes, press Start, and watch it run. When the run finishes, the `.safetensors` lands in `models/loras/`, where the LoRA loader node picks it up automatically, so you can generate with it over in the Studio tab straight away. ![Inline Studio Trainer tab showing the LoRA training node graph with a dataset, live logs, and a loss curve](https://raw.githubusercontent.com/inlineresearch/Inline-Studio/main/screenshots/lora-trainer.png) +Five nodes, wired left to right, with the hyperparameters behind an Adjust button so the node face stays a status surface: + ``` [ Load Dataset ] --> [ Caption ] --> [ Train LoRA ] --> [ Graph ] | +--> Resources (VRAM monitor) ``` -Hyperparameters sit behind an Adjust button, so the node face stays a status surface. MiniMax H3 -trains on stills for look and style, or on clips to learn motion as well; one dataset can hold both. +### Does my card fit? + +Peak VRAM at 512px, 12 steps, rank 16, batch 1, gradient checkpointing on: + +| Architecture | 512px peak | 16GB card | +| ------------------------- | ---------- | ----------- | +| FLUX.2 (klein Base 4B) | ~8.6GB | yes | +| Krea 2 (4-bit base) | ~11.9GB | yes | +| Z-Image | ~13.4GB | yes | +| MiniMax H3 (4-bit, video) | ~20.6GB | yes, slowly | + +Training is cheaper than generating, and a LoRA trained at 512 applies at any generation resolution. Full per-card matrix, both resolutions and the timings: [Benchmark results](TRAINING.md#benchmark-results). + +If you installed with `--extra all` from [Get Started](#get-started), the trainer is ready. Otherwise: + +```bash +cd core +./webui.sh --install --extra training # Windows: .\webui.bat --install --extra training +``` + +For a worked example, see [`inlineresearch/skin-lora-krea-2-raw`](https://huggingface.co/inlineresearch/skin-lora-krea-2-raw), a photorealistic skin LoRA trained here on the Krea 2 RAW base from the 26 image and caption pairs published as [`inlineresearch/krea2-skin-lora`](https://huggingface.co/datasets/inlineresearch/krea2-skin-lora). + +**[TRAINING.md](TRAINING.md) is the full reference:** [which base to train on](TRAINING.md#architecture-and-base-model-modes) · [measured benchmarks](TRAINING.md#benchmark-results) · [datasets and outputs](TRAINING.md#datasets-and-outputs) · [stop and resume](TRAINING.md#stop-and-resume) · [trigger words](TRAINING.md#trigger-words) · [base precision](TRAINING.md#base-precision) + +## How it works + +Generating a single frame is the easy part. The work that makes an AI film is what comes after: exploring options, keeping what's good, and shaping a repeatable process out of it. Inline Studio is the layer where that happens, organised around one model: + +### Export the whole pipeline, not just the final render -Already installed with `--extra all`? The trainer is ready. Otherwise -`./webui.sh --install --extra training`. +From the home screen, **Export** zips a project into one archive. Import it on the other side and you get everything back: the inputs (every imported asset), the outputs (all the generated takes), and the graph that turned one into the other. Whoever opens it can re-run the pipeline exactly and keep iterating. -**[TRAINING.md](TRAINING.md) is the full reference:** -[which base to train on](TRAINING.md#architecture-and-base-model-modes) · -[benchmarks](TRAINING.md#benchmark-results) · -[training on clips](TRAINING.md#training-on-clips) · -[datasets and outputs](TRAINING.md#datasets-and-outputs) · -[stop and resume](TRAINING.md#stop-and-resume) · -[trigger words](TRAINING.md#trigger-words) +## Two ways to generate -A worked example: [`skin-lora-krea-2-raw`](https://huggingface.co/inlineresearch/skin-lora-krea-2-raw), -trained here on Krea 2 RAW from the 26 pairs published as -[`krea2-skin-lora`](https://huggingface.co/datasets/inlineresearch/krea2-skin-lora). +Pick whatever fits the shot, and mix both in one film. However you render, the frame keeps its full take history, so you never lose a good version. -## Generate +| How you render | What it's like | What you need | +| --------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **Local GPU: Inline Core** _(built in)_ | Drop a **Z-Image Turbo**, **Krea 2**, **FLUX.2**, or **MiniMax H3** node, wire a prompt, hit Run: one node, no loader/sampler wiring. A single `.safetensors` is usually all you bring (a prequantized build can be a diffusers folder); the engine pairs it with a VAE + text-encoder and downloads nothing behind your back. Two or more GPUs? It can split one image's denoise across them (xDiT). | Runs locally on your own GPU. No account, no external server. Low-VRAM friendly: it auto-fits the model to your card (streaming weights, int8, then NF4) so a model too big for full precision still runs, with no flags. | +| **Hosted: API Nodes** | Add a Generate node and pick a model: hosted, closed models across image, video, and audio. No GPU, instant range. See [API Nodes](#api-nodes) for the model list and providers. | A provider key (currently [fal](https://fal.ai/dashboard/keys)); it stays on your machine, and you pay per render (each node estimates the price first). | -Drop a model node, wire a prompt, hit Run. One node, no loader or sampler wiring. Either put a -`.safetensors` in `core/models/diffusion_models/`, or use the node's model popup to download the -diffusion model, VAE and text encoder with visible progress. Nothing is fetched behind your back. +For local generation, either drop a `.safetensors` into `core/models/diffusion_models/`, or add a model node and use its **model popup** (a blinking hint shows up when something's missing) to download the diffusion model, VAE, and text-encoder into `core/models/`, with visible progress. The canvas and planning work with no models at all. See [Krea 2](#krea-2) and [FLUX.2](#flux2) for those models' files and VRAM. + +## Inline Core generation engine + +Inline Core is a from-scratch generation engine for local rendering. It keeps the open node-graph model (a typed DAG of nodes and edges → immutable "takes"), and Inline Studio drives it as a single process. ![Z-Image Turbo generating locally on the Inline Core engine](https://raw.githubusercontent.com/inlineresearch/Inline-Studio/main/screenshots/zit.png) -- **Z-Image Turbo** is the low-VRAM starting point, distilled to run CFG-free. -- **[Krea 2](https://www.krea.ai/)** is a 12.9B MMDiT in two halves: train on **RAW**, generate with - **Turbo**. A LoRA trained on RAW applies to Turbo unchanged. -- **[FLUX.2](https://bfl.ai/blog/flux-2)** is natively multi-reference: wire several images and the - prompt addresses them by position. One node covers klein 4B and 9B, their Base builds, and dev. -- **[MiniMax H3](https://huggingface.co/MiniMaxAI/MiniMax-H3)** generates video **and its - soundtrack** in one pass, as four nodes (text, image, first and last frame, reference). 24fps, 5 to - 15 seconds. See the [open weights guide](https://inlinestudio.art/minimax-h3-open-weights). -- **ControlNet** steers a local render with a pose, depth or edge map. **Control Space** is a 3D pose - editor in a node, so you can build the skeleton rather than find a reference photo. +- **One process, one port** - Inline Studio is a **web SPA** (React) served by Inline Core (a headless Python engine, in `core/`). `core/main.py` runs Core, which serves the built UI and is the app's backend. +- **Core owns the backend** - the browser reaches it over a small typed RPC/WebSocket contract; Core owns the project database, the filesystem, generation, and the ffmpeg timeline. No Electron, no separate Node server, nothing external to stand up. +- **Typed graph, checked before it runs** - named params and type-checked edges, so a bad graph is rejected at submit rather than dying part-way through a denoise. +- **Immutable takes** - regenerating adds a take; nothing is ever overwritten. The take history is the point. +- **Durable runs** - a run survives a restart, and progress streams over a WebSocket. +- **Graph decoupled from GPU work** - the graph is the unit of caching; a batched sampler is the unit of batching, grouping compatible jobs across requests. +- **A single device policy owns all placement** - device, dtype, offload, and attention, so the same graph runs on a 4090, a 6 GB laptop, pure CPU, or split across several GPUs without touching the graph. +- **Bring your own models, no hidden downloads** - a drop-in `models/` layout feeds a typed catalog and versioned node descriptors; nothing is fetched behind your back. -
-Model files: what goes where +### Krea 2 + +[Krea 2](https://www.krea.ai/) is a 12.9B single-stream MMDiT, released as two checkpoints that work together: **RAW** is the undistilled base you fine-tune, **Turbo** is an 8-step distilled checkpoint you generate with. A LoRA trained on RAW applies to Turbo unchanged, which is the workflow both nodes are built around. -Only the **bf16** builds load. The `fp8_scaled`, `int8_convrot`, `mxfp8`, `nvfp4` and `pruned` files -in the community repos carry ComfyUI-specific scale tensors that only ComfyUI reads; the node says so -rather than failing deep in a load. Memory saving is the device policy's job instead. +Both nodes read the ComfyUI-style files from [`Comfy-Org/Krea-2`](https://huggingface.co/Comfy-Org/Krea-2): ``` core/models/ - diffusion_models/ krea2_turbo_bf16.safetensors <- Krea 2 Turbo (generate) - krea2_raw_bf16.safetensors <- Krea 2 RAW (train) - flux-2-klein-4b.safetensors <- FLUX.2 default, Apache 2.0 - flux-2-klein-base-4b.safetensors <- FLUX.2 base build, for training - minimax_h3_fl2va_bf16.safetensors <- H3 text, image, first/last frame - minimax_h3_ref2va_bf16.safetensors <- H3 reference node - text_encoders/ qwen3vl_4b_bf16.safetensors <- Krea 2 - qwen_3_4b.safetensors <- FLUX.2 klein 4B, shared with Z-Image - MiniMax-H3-text-encoder/ <- Qwen3-VL-32B, a folder - MiniMax-H3-processor/ + diffusion_models/ krea2_turbo_bf16.safetensors <- for the Krea 2 Turbo node + krea2_raw_bf16.safetensors <- for the Krea 2 RAW node (and training) + text_encoders/ qwen3vl_4b_bf16.safetensors vae/ qwen_image_vae_diffusers.safetensors - flux2-vae.safetensors - minimax_h3_video_vae_fp16.safetensors + loras/ krea2_retroanime.safetensors <- the official style LoRAs, optional +``` + +Two things are worth knowing before you download 26 GB twice: + +- **Only the `bf16` builds load.** The `fp8_scaled`, `int8_convrot`, `mxfp8` and `nvfp4` files in that repo carry ComfyUI-specific scale tensors that only ComfyUI can read, and the node says so rather than failing deep in a load. Memory saving is the device policy's job instead. +- **The VAE is the diffusers-format one**, fetched from [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image). ComfyUI's `qwen_image_vae.safetensors` holds the same weights in a different module layout that diffusers cannot read. The node's model popup downloads the right file for you. + +Nothing here needs a Hugging Face token: every repo involved is public, and Krea's own gated repos are never touched. + +### FLUX.2 + +[FLUX.2](https://bfl.ai/blog/flux-2) from Black Forest Labs is the first family here that is natively **multi-reference**: reference images ride in the denoiser's token sequence, so "the character from image 1 wearing the jacket from image 2" is a first-class capability rather than a workaround. + +One node, **FLUX.2**, covers the whole family. Pick a checkpoint in the node's Adjust sidebar and it identifies itself: klein 4B, klein 9B, either Base build, the KV variant, or dev. Steps and guidance default to "from model", so switching a distilled checkpoint for its Base build moves 4 steps at guidance 1.0 to 50 at guidance 4.0 without touching a setting. + +- **Reference images** - wire one image to edit it, or several to compose from them. The node numbers them on its face, and the prompt addresses them by position. There is no denoise-strength slider because FLUX.2 has no img2img: a single reference _is_ the edit. +- **klein 4B is Apache 2.0** and the recommended starting point. It renders 1024px in four steps. At bf16 it is 16.1 GB of weights peaking around 17.9 GB, so a 24 GB card holds it resident; on a 16 GB card the fit ladder drops it to int8 or fp16 and it still runs. Its text encoder is the same Qwen3-4B file Z-Image already uses, so if you have run Z-Image you have most of it. +- **dev is 32B.** The fp8 build the model popup offers wants around 32 GB. Bring a prequantized NF4 folder yourself and it fits a 24 GB card instead, because the prompt is encoded first and the text encoder freed before the transformer loads. + +Files come from the ungated ComfyUI repacks. The model popup fetches klein 4B, its text encoder and the VAE in one click, and offers the rest of the family as optional extras: + +``` +core/models/ + diffusion_models/ flux-2-klein-4b.safetensors <- the default, Apache 2.0 + flux-2-klein-base-4b.safetensors <- the base build, for LoRA training + flux-2-klein-9b-int8-ConvRot-comfyui.safetensors <- klein 9B, int8, ~12 GB + flux2_dev_fp8mixed.safetensors <- dev, fp8, ~32 GB + text_encoders/ qwen_3_4b.safetensors <- the 4B builds, shared with Z-Image + qwen_3_8b.safetensors <- klein 9B + mistral_3_small_flux2_fp8.safetensors <- dev + vae/ flux2-vae.safetensors +``` + +Each klein size needs its own text encoder, and dev uses Mistral-3 rather than Qwen3. + +For dev on a 24 GB card, take the ungated [`diffusers/FLUX.2-dev-bnb-4bit`](https://huggingface.co/diffusers/FLUX.2-dev-bnb-4bit) instead: an 18.1 GB NF4 transformer beside a 15.4 GB NF4 encoder, far cheaper than the fp8 single file. Clone it into `diffusion_models/` as a folder. The popup does not fetch it, and a diffusers folder is a valid checkpoint anywhere a single file is. + +Worth knowing: + +- **Prompts are prose, not tags.** FLUX.2 wants natural language, and keyword stuffing works against it. Word order carries weight. +- **Only the Base klein checkpoints take a negative prompt.** The distilled builds run no classifier-free guidance and dev is guidance-distilled, so a negative prompt is logged and ignored there rather than silently pretending to apply. +- **dev and every 9B build are non-commercial.** klein 4B, its Base build, and the VAE are Apache 2.0. + +### MiniMax H3 + +[MiniMax H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) is the first video model in Inline Core, and the first model here that generates a soundtrack rather than a silent clip. One transformer denoises the video and its 32 kHz stereo audio in a single pass, so a take is one MP4 with sound in it, ready to drop straight into the timeline. + +Four nodes, because the inputs genuinely differ and a node should say what it takes: + +- **Text → Video** for a shot from nothing but a prompt. +- **Image → Video** to bring a still you already like into motion. +- **First and Last Frame** to pin the opening frame, the closing frame, or both, and let the model fill in between. +- **Reference → Video** for consistency across shots. Wire up to nine images, three video clips and three audio clips, and address them by position in the prompt. Wiring order is the numbering you see on the node. + +Output is 24 fps between 5 and 15 seconds at a 768 pixel short edge. Duration snaps to the frame grid the video decoder works in, so asking for 14.9 seconds renders 14.4 rather than failing. There is no guidance slider and no negative prompt, because the released checkpoints are guidance-distilled and neither exists. + +``` +core/models/ + diffusion_models/ minimax_h3_fl2va_bf16.safetensors <- text, image, and first/last frame + minimax_h3_ref2va_bf16.safetensors <- the reference node + text_encoders/ MiniMax-H3-text-encoder/ <- Qwen3-VL-32B, a folder + MiniMax-H3-processor/ <- tokenizer and processor + vae/ minimax_h3_video_vae_fp16.safetensors minimax_h3_audio_vae_fp32.safetensors - loras/ your trained adapters land here - controlnet/ ControlNet and control-LoRA files ``` -Krea 2's VAE is the **diffusers-format** one from [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image); -ComfyUI's `qwen_image_vae.safetensors` holds the same weights in a layout diffusers cannot read. -Every repo involved is public, so no Hugging Face token is needed. +Worth knowing before you start a download this size: -**MiniMax H3 is big:** 144GB for the first three nodes, 210GB with the reference node. Measured on a -45GB card, a 10 second clip at 960x544 takes about 7.2 minutes, peaking at 38.9GB VRAM and 46.7GB of -system RAM, so plan on 64GB of RAM. Canvas size is the biggest speed lever: 960x544 renders about -2.3x faster per step than 1344x768. +- **This is a big model.** 144 GB for the first three nodes, 210 GB with the reference node. A 33B transformer runs beside a 32B conditioner, but not at the same time: the prompt is encoded first, the conditioner then steps off the card, and the denoiser takes it for the whole denoise. Two things make that fit. The modulation weights, 40% of the transformer, are factorised at load and take it from 66.3 GB to 40.3 GB. The video VAE stays resident when the card has room, rather than streaming leaf by leaf. Measured on a 45 GB card: **a 10 second clip at 960x544 takes about 7.2 minutes, peaking at 38.9 GB VRAM and 46.7 GB of system RAM that cannot be reclaimed**, so plan on 64 GB of RAM. Smaller cards stream more and are slower. If that is out of reach, the same model is available as an API node with no setup at all. +- **Canvas is the biggest speed lever.** 960x544 renders about 2.3x faster per step than the trained 1344x768, and the difference is far larger than any other setting. +- **The bf16, `pruned` and `pruned_fp8_scaled` builds load.** The pruned builds ship the modulation branch already reduced and no timestep path, which Inline reads directly; `pruned_fp8_scaled` is 21.0 GB against 66.3 GB for the same model. The `int8_convrot` files still do not load, because their weights are stored rotated and only ComfyUI can undo that. The picker lists what it cannot read with the reason. +- **A smaller file is a smaller download, not a smaller model.** All of them occupy the same memory once loaded, so the choice is bandwidth and disk, not VRAM. **Training needs the bf16 build**: a pruned one has no timestep path to derive the modulation basis from, and it would save nothing anyway, since the base trains at 4-bit whichever file it starts from. +- **LoRAs work.** Every H3 node has a LoRA input, and adapters are fused into each block as it streams, before the factorisation and the quantisation. You can train one in the Trainer tab: see [LoRA training](#lora-training). An H3 LoRA is trained on stills and applies to video, so it carries look and style rather than motion. -**FLUX.2 dev on a 24GB card:** take the ungated -[`diffusers/FLUX.2-dev-bnb-4bit`](https://huggingface.co/diffusers/FLUX.2-dev-bnb-4bit) folder rather -than the fp8 single file. A diffusers folder is a valid checkpoint anywhere a single file is. +### ControlNet -
+Steer a local render with a pose, depth, or edge map. Wire a control map into a gen node's **Control** input and pick a ControlNet in the node's Adjust sidebar. -
-Hosted models (API Nodes) +- **Control Space** - a 3D pose editor in a node. Pose one or more characters, frame a camera, and render the scene as an OpenPose skeleton or a depth map. No reference photo needed. +- **Apply ControlNet** - turn any image into a control map (OpenPose, Depth-Anything V2, MiDaS depth, or Canny edges). Detector weights download once on first use. +- **Z-Image Turbo** - full ControlNet via the Fun Union model. Use the distilled `-2602-8steps` build; the plain one is blurry at 8 steps. +- **Krea 2** - depth control via the [`Patil/Krea-2-depth-controlnet`](https://huggingface.co/Patil/Krea-2-depth-controlnet) control-LoRA. +- **FLUX.2** - two routes. On any variant, a control map wired into **Control** is used as a reference image, which steers loosely and costs nothing extra. On **dev**, pick the [`alibaba-pai/FLUX.2-dev-Fun-Controlnet-Union`](https://huggingface.co/alibaba-pai/FLUX.2-dev-Fun-Controlnet-Union) model for tight structural adherence: one union model covering canny, depth, pose, HED, MLSD, scribble and gray, with no mode to select. It needs real headroom, rendering at 512px on a 24 GB card. +- **Control strength** - dial how hard the map is followed, per node. -Add a Generate node and pick a model: hosted, closed models across image, video and audio, with no -GPU and no setup. Bring your own provider key; it stays on your machine and you pay the provider per -render, with each node estimating the price first. +Drop ControlNet files into `core/models/controlnet/`, or use the node's model popup to download them. -The initial provider is **[fal](https://fal.ai)**: FLUX.2, FLUX.2 Edit, GPT Image 2, Nano Banana, -Seedance, MiniMax H3, LTX, Sonilo and more. Add your [key](https://fal.ai/dashboard/keys) in -Settings. MiniMax H3 is on the canvas both ways, as an API node and as local nodes with no -per-render cost. +### Community extensions -Local and hosted mix freely in one film, and either way the frame keeps its full take history. +Install community-built nodes straight from a GitHub repo, from the Extensions dialog or a repo URL. -
+- **One-click install**, with a live stepper showing download, security review, dependency resolution, and activation. +- **Every install is reviewed.** Code that could replace Inline's PyTorch, hide a payload, or run at install time is blocked outright; subprocesses, sockets, and unknown network hosts need your explicit approval. +- **Extensions can't break your setup.** Their dependencies install into their own folder and can never touch the shared torch/diffusers runtime, and genuine conflicts fail at install with both versions named. +- **Nodes appear on the canvas immediately**, with their own params, model downloads, take history, and Run control. No restart, and no frontend code from the author. +- **Toggle any node on or off**, roll back to a previous version, or uninstall, and see when an update is available. +- **Publish by tagging.** Authors list once in the registry; after that a new tag reaches users with no further PR. + +Browse the [**extension registry**](https://github.com/inlineresearch/Inline-Registry), or copy the +[**extension guide**](https://github.com/inlineresearch/Inline-Studio-Extension-Guide) to build your +own: four working nodes, declared model downloads, and a full authoring reference.
Multi-GPU: split one image across GPUs -With two or more GPUs, Inline Core can cut a single image's latency by running its denoise loop -collectively across them. This is not one image per GPU; it is one image whose sampling is shared, so -a single render finishes faster. +Got two or more GPUs? Inline Core can cut a single image's latency by running its **denoise loop** (the expensive, iterative sampling step) collectively across them. This is not "one image per GPU" (independent renders); it's **one image whose sampling is shared by all the GPUs**, so a single render finishes faster. -Built on [xDiT](https://github.com/xdit-project/xDiT) in an isolated worker group, one process per -GPU. The split method follows the interconnect Core detects: PipeFusion over PCIe, Ulysses with -NVLink. Turn it on with `./webui.sh --multi-gpu` after `uv pip install -e ".[parallel]"`. +It's done with [xDiT](https://github.com/xdit-project/xDiT) (`xfuser`), which parallelizes diffusion-transformer inference in an isolated worker group (one process per GPU via `torchrun`, over local IPC). The HTTP server, database, and graph stay single-process; only the denoise distributes, and it sits behind a sampler seam so single-GPU/CPU runs pay no overhead. The split method is chosen from the interconnect Core detects: **PipeFusion** (default, works over plain PCIe) or **Ulysses** (sequence-parallel attention, used when NVLink is present). Turn it on with `./webui.sh --multi-gpu` after `uv pip install -e ".[parallel]"`.
-## Features +For the full engineering story (the graph/sampler/device-policy design, the node vocabularies, and the xDiT worker group), see **[core/README.md](core/README.md)** and **[core/CLAUDE.md](core/CLAUDE.md)**. -- Free-form node canvas with versioned, non-destructive takes -- Train your own styles and consistent characters, locally -- Local generation built in, with the Inline Core engine -- Multi-reference composition, and video with sound -- Every locally generated image embeds the graph that made it: drop the file back on the canvas to - rebuild the pipeline -- Video Director, Trim Video and Trim Audio nodes -- Export and import the whole project as one archive -- Community extensions, and API Nodes for hosted models +## API Nodes -[**Follow the Animated Short Film tutorial →**](https://inlinestudio.art/projects/circuit-race) +**API Nodes** bring hosted, closed models onto the same canvas: no GPU, no setup, instant creative range. Add a Generate node, pick a model, and bring your own provider key (it stays on your machine); you pay the provider per render, and each node estimates the price before you run. -## How it works +The initial provider is **[fal](https://fal.ai)**, with models across image, video, and audio: **FLUX.2**, **FLUX.2 Edit**, **GPT Image 2**, **Nano Banana**, **Seedance**, **MiniMax H3**, **LTX**, **Sonilo**, and many more. Add your [fal.ai key](https://fal.ai/dashboard/keys) in Settings to use them. More providers will follow behind the same API Node surface. + +However you render, the frame keeps its full, non-destructive take history, so you can mix API Nodes and local generation in the same film without ever losing a good version. + +### MiniMax H3 + +H3 (Hailuo 03) is on the canvas as three API nodes, all at 2K, 5 to 15 seconds. The open weights also run locally on the Inline Core engine, as [four nodes with no per-render cost](#minimax-h3) if you have the hardware for them. -Generating one frame is the easy part. The work that makes a film is what comes after: exploring -options, keeping what is good, and shaping a repeatable process out of it. +- **Text → Video** for a shot from nothing but a prompt. +- **Image → Video** for a still you already like. It has two image dots: wire a start frame on its own, or add an end frame and H3 interpolates between the two. +- **Reference → Video** for consistency across shots. Wire up to nine images, plus reference video and audio, and address them by position in the prompt: "Image 1 is the lead, Image 2 is her dog". Wiring order is the numbering you see on the node. -A **frame** is a slot with a history of **takes**, never a single file. Generating again adds a take -and nothing is overwritten. **Export** zips a project into one archive, inputs, outputs and the graph -that turned one into the other, so whoever opens it can re-run the pipeline exactly. +Video costs $0.26 per second at 2K, so the node's price badge reads about $1.30 for a five second clip. Reference images past the first five and any reference video cost extra on top of that. ![Inline Studio dashboard with recent AI film projects](https://raw.githubusercontent.com/inlineresearch/Inline-Studio/main/screenshots/screenshot-dashboard.png) -It runs as a single process on one port: the Inline Core engine (Python) serves the web UI and does -the generation. No desktop install, no separate backend. For the engineering story see -[core/README.md](core/README.md) and [core/CLAUDE.md](core/CLAUDE.md). +## FAQ -## Extensions +### Is Inline Studio free? -Install community-built nodes from a GitHub repo, from the Extensions dialog or a repo URL. Every -install is security-reviewed, dependencies are isolated from the shared torch runtime, and nodes -appear on the canvas immediately with no restart. +Yes. Inline Studio is free and open source under the [GPL-3.0 license](LICENSE). There's no paid tier to use the app. -Browse the [registry](https://github.com/inlineresearch/Inline-Registry), or copy the -[extension guide](https://github.com/inlineresearch/Inline-Studio-Extension-Guide) to build your own. +### Do I need a GPU? -## FAQ +Only for **local** generation. The built-in Inline Core engine renders on the GPU of whatever machine runs it (you can also run it on a remote GPU box and open the UI from your laptop). Hosted **fal** models need no GPU at all, and the canvas + planning work with no GPU either. -**Is Inline Studio free?** Yes, free and open source under GPL-3.0. Local generation and training -cost nothing to run. Hosted models are billed by the provider. +### Can I train a LoRA locally? -**Do I need a GPU?** Not for the canvas, planning, editing or hosted models. Local generation and -LoRA training need one; see the table at the top. +Yes, that is what the [Trainer tab](#lora-training) is for, and it runs entirely on your own GPU with no cloud step. It trains LoRAs for Z-Image, Krea 2, FLUX.2, and MiniMax H3. Training is cheaper than generating: FLUX.2 klein Base trains at 512px in about 8.6GB, Z-Image in about 13GB, and Krea 2 with a 4-bit base in about 12GB, so a 16GB card handles all three. See [Benchmark results](TRAINING.md#benchmark-results) for the measured table. -**Can I train a LoRA locally?** Yes, for all four local models, on your own GPU. See -[TRAINING.md](TRAINING.md). +### What models can I run? -**What models can I run?** Locally: Z-Image Turbo, FLUX.2, Krea 2 and MiniMax H3. Hosted: the fal -catalogue, with more providers to follow. +See [Two ways to generate](#two-ways-to-generate): Z-Image, [Krea 2](#krea-2), [FLUX.2](#flux2), or [MiniMax H3](#minimax-h3) locally on your own GPU, or hosted fal models. You can train a LoRA locally for any of the four as well, see [LoRA training](#lora-training). Adding a new local model is a Core change (a model runner), no UI release. ## Contributing -Issues, ideas and pull requests are all welcome. Start with [CONTRIBUTING.md](CONTRIBUTING.md) for -setup and the checks to run; [CLAUDE.md](CLAUDE.md) is the deeper engineering guide. By taking part -you agree to our [Code of Conduct](CODE_OF_CONDUCT.md). +Inline Studio is early and moving fast, any issues, ideas, and pull requests are all welcome. Start with [CONTRIBUTING.md](CONTRIBUTING.md) for setup, the checks to run, and how to open a PR. [CLAUDE.md](CLAUDE.md) is the deeper engineering guide: the architecture, the data model, and the conventions to follow. By taking part you agree to our [Code of Conduct](CODE_OF_CONDUCT.md). + +Want to help by using it for real? Try the [creator task](task.md): build a short 20-second AI film in Inline Studio and send us your feedback. ## Credits -- [**xDiT**](https://github.com/xdit-project/xDiT) for the PipeFusion and Ulysses parallelism behind the multi-GPU denoise. -- [**ai-toolkit**](https://github.com/ostris/ai-toolkit) by ostris, for the approach to training on a step-distilled model, and the [Z-Image](https://huggingface.co/ostris/zimage_turbo_training_adapter) and [Krea 2](https://huggingface.co/ostris/krea2_turbo_training_adapter) training adapters. -- [**diffusers**](https://github.com/huggingface/diffusers) for the Krea 2 and MiniMax H3 reference implementations. -- [**Krea AI**](https://www.krea.ai/) for Krea 2, under the [Krea AI Community License](https://www.krea.ai/krea-2-licensing). -- [**Black Forest Labs**](https://bfl.ai/blog/flux-2) for FLUX.2: klein 4B, its Base build and the VAE are Apache 2.0; dev and the 9B builds are non-commercial. -- [**MiniMax**](https://huggingface.co/MiniMaxAI/MiniMax-H3) for MiniMax H3, under the MiniMax H3 Community License. +Inline Core's multi-GPU denoise builds on [**xDiT**](https://github.com/xdit-project/xDiT)'s PipeFusion and Ulysses parallelism. + +The LoRA trainer's approach to training on a step-distilled model follows [**ai-toolkit**](https://github.com/ostris/ai-toolkit) by ostris, and the Turbo modes use his training adapters for [Z-Image](https://huggingface.co/ostris/zimage_turbo_training_adapter) and [Krea 2](https://huggingface.co/ostris/krea2_turbo_training_adapter). + +Krea 2 support follows the reference implementations in [**diffusers**](https://github.com/huggingface/diffusers) (`Krea2Pipeline` and the Krea 2 DreamBooth LoRA example). Krea 2 is released by [Krea AI](https://www.krea.ai/) under the [Krea AI Community License](https://www.krea.ai/krea-2-licensing); the weights are the user's to obtain and use under that license. + +[**FLUX.2**](https://bfl.ai/blog/flux-2) is released by Black Forest Labs. klein 4B, its Base build, and the VAE are Apache 2.0; dev and the 9B builds carry non-commercial terms. The weights are the user's to obtain and use under those. + +[**MiniMax H3**](https://huggingface.co/MiniMaxAI/MiniMax-H3) is released by MiniMax under the MiniMax H3 Community License; the weights are the user's to obtain and use under that. Support is built on the [**diffusers**](https://github.com/huggingface/diffusers) integration from its MiniMax-H3 pull request, vendored with provenance until it lands upstream. + +## Help shape Inline Studio + +Are you an AI filmmaker who wants to help us make this better? We run a **paid trial feedback program**: use Inline Studio on real work, tell us what helps and what gets in your way, and get paid for your time. + +Come say hi on our [Discord](https://discord.gg/cSUS88VdY9) and reach out, we'll get you set up. + +[![Join our Discord](https://img.shields.io/badge/Discord-Join%20the%20community-5865F2?logo=discord&logoColor=white&style=for-the-badge)](https://discord.gg/cSUS88VdY9) ## License -[GPL-3.0](LICENSE). Model weights are yours to obtain and carry their own licences, which the GPL -does not change. +Copyright (C) 2026 Inline Studio. Licensed under the [GNU General Public License v3.0](LICENSE): you may use, study, share and modify it, and any work you distribute that builds on it must also be GPL-3.0. + +The models you run carry their own licenses, which the GPL does not change: Krea 2 is under the [Krea AI Community License](https://www.krea.ai/krea-2-licensing), Z-Image under Tongyi's terms, and FLUX.2 is split - klein 4B, its Base build, and the VAE are Apache 2.0, while dev and every 9B build are non-commercial. You bring your own weights and use them under those. diff --git a/TRAINING.md b/TRAINING.md index 791eb46..ce9d913 100644 --- a/TRAINING.md +++ b/TRAINING.md @@ -2,9 +2,9 @@ The full reference for Inline Studio's **Trainer**: which base to train on, what it costs in VRAM on a real card, and every setting that shapes the result. For the short version and a screenshot of -the canvas, see [Train a LoRA in the README](README.md#train-a-lora). +the canvas, see [LoRA training in the README](README.md#lora-training). -Inline Studio trains LoRAs for **Z-Image Turbo**, **FLUX.2**, **Krea 2** and **MiniMax H3** on your own +Inline Studio trains LoRAs for **Z-Image**, **Krea 2**, **FLUX.2** and **MiniMax H3** on your own GPU, with no cloud step and nothing uploaded. Training is cheaper than generating: a 16GB card trains all three image models at 512px, and a LoRA trained at 512 applies at any generation resolution. @@ -56,7 +56,7 @@ The Trainer's Adjust panel picks the **architecture** first (Z-Image, Krea 2, FL **MiniMax H3** is the video model, and it trains on **still images**: -- **FL2VA** is the only base, and it is undistilled, so there is no adapter and nothing to drift. Put `minimax_h3_fl2va_bf16.safetensors` in `models/diffusion_models/`, train on stills, then wire the LoRA into any of the four H3 nodes. It loads on the Reference to Video node too, which uses a different checkpoint file: the two partitions are the same architecture. +- **FL2VA** is the only base, and it is undistilled, so there is no adapter and nothing to drift. Put `minimax_h3_fl2va_bf16.safetensors` in `models/diffusion_models/`, train on stills, then wire the LoRA into any of the four H3 nodes. **It has to be the bf16 file.** The smaller `pruned` and `pruned_fp8_scaled` builds generate but cannot train: they ship no timestep path for the modulation basis to be derived from, and they would save nothing anyway, because the base trains at 4-bit whichever file it starts from. The trainer says so rather than failing part way in. It loads on the Reference to Video node too, which uses a different checkpoint file: the two partitions are the same architecture. - **Stills or short clips.** Drop images and it learns appearance: look, style, character, lighting. Drop video and it learns motion too. Sound is never learned either way, because the audio rows are empty. See [Training on clips](#training-on-clips). - **The base is 4-bit, always.** H3 is 40GB after the AdaLN factorisation and 11.7GB after quantisation, so full precision is refused rather than offered and then failing. There is no base-precision control for H3 for the same reason. - **A 24GB card is comfortable and a 16GB card works, slowly.** The run encodes latents and captions in two passes that never overlap, because H3's fp32 video VAE and its 32B conditioner cannot be resident together. On a card that holds the conditioner it peaks at 20.6GB; on one that does not, the conditioner runs on the CPU and the peak drops to 12.7GB while a step goes from 0.6s to 16s. Either way there is about seven minutes of startup, and 64GB of system RAM for the smaller card. See [Benchmark results](#benchmark-results) for the split. The download is about 124GB before any of that. @@ -104,7 +104,7 @@ audio rows, so an adapter changes what a clip looks like and never what it sound ## Install -If you installed with `--extra all` from [Install](README.md#install), the trainer is already set up - nothing more to do. To add it to a leaner install, its dependencies (PEFT, 8-bit Adam, the captioner) sit behind the `training` extra: +If you installed with `--extra all` from [Get Started](README.md#get-started), the trainer is already set up - nothing more to do. To add it to a leaner install, its dependencies (PEFT, 8-bit Adam, the captioner) sit behind the `training` extra: ```bash cd core diff --git a/core/CLAUDE.md b/core/CLAUDE.md index 4fc326f..edcd141 100644 --- a/core/CLAUDE.md +++ b/core/CLAUDE.md @@ -163,6 +163,17 @@ between nodes and are never takes. → sequential offload → wont-fit. NF4 is what makes a 32B model viable on a 24 GB card. Note int8 forces bf16 (torchao's weight-only int8 silently no-ops under fp16) while NF4 does not, so a Turing card keeps its fp16 tensor cores under NF4. +- **A pre-reduced checkpoint must not be re-reduced, structurally or numerically.** MiniMax H3's + `pruned` builds ship the AdaLN branch already factorised to rank 8 and drop the timestep path + entirely, so re-running our factorisation multiplies a `[96768, 8]` projection by a full-width + basis. `minimaxh3/pipeline.py` turns `factorise_adaln` off for those, the same way the rule below + turns quantization off for a prequantized file. Both are the same rule: the source is already in + the target form. +- **Size a checkpoint by what it becomes, not by what it weighs.** A pruned file has already lost + its AdaLN branch and an fp8 file stores half the bytes it will occupy once dequantised, so scaling + the on-disk number under-sizes both, by up to 3x. `minimaxh3.requirements.resident_bytes` counts + from the header. Under-sizing is the dangerous direction: the fit ladder then promises a machine + that dies to a host-RAM OOM kill instead of raising. - **Prequantized checkpoints must not be re-quantized.** A checkpoint that ships already quantized (`flux2/variants.is_prequantized`) has an on-disk size that already _is_ its resident size, so the ladder's assumption that quantization halves it does not hold, and handing diffusers a second, @@ -230,6 +241,11 @@ real codec that moves tensors lives with the model runner. Don't scatter it. - **Bring-your-own models.** Nothing is downloaded by the engine. The catalog scans; the user places files. A model picker is a `SELECT` param with `options_from=""`. +- **Adapter strength is not a quality metric, and a threshold on it is a false-positive machine.** + Measured against real bases, published LoRAs that work well span `|B@A| / |W|` from 0.017% + (a style LoRA) to 1.2% (a restoration LoRA), so "this adapter looks weak" is not a finding. What + predicts a LoRA doing nothing is whether its per-weight change clears one quantization step. Warn + on that, and only when the base is actually quantized. - **Verify image models by rendering.** The FLUX.2 work shipped five bugs past a green test suite, and every one produced a _wrong image rather than an error_: a mis-keyed checkpoint, a vision-language encoder loaded in place of a text one, an unnormalized latent, and a control context diff --git a/core/src/inline_core/models/keymap.py b/core/src/inline_core/models/keymap.py index b363056..7d4d279 100644 --- a/core/src/inline_core/models/keymap.py +++ b/core/src/inline_core/models/keymap.py @@ -262,7 +262,7 @@ def transform( raise ComponentError(f"{key} has {tensor.shape[0]} rows, not divisible into {parts} parts.") if verify_layout: assert_layout(tensor, action, key=key) - source = _deinterleave(tensor, parts, action.head_dim) if ( + source = deinterleave_rows(tensor, parts, action.head_dim) if ( action.layout is RowLayout.INTERLEAVED ) else tensor block = source.shape[0] // parts @@ -270,16 +270,31 @@ def transform( yield target, source[index * block : (index + 1) * block] -def _deinterleave(tensor: Any, parts: int, head_dim: int) -> Any: - """``[p0_h0; p1_h0; p2_h0][p0_h1; …]`` to ``[p0_all; p1_all; p2_all]``. - - ``transpose`` is not the same call in torch and numpy - torch swaps two axes, numpy wants a full - permutation - so the swap is spelled per backend rather than duck-typed. - """ +def deinterleave_rows(tensor: Any, parts: int, head_dim: int) -> Any: + """``[p0_h0; p1_h0; p2_h0][p0_h1; …]`` to ``[p0_all; p1_all; p2_all]``.""" if head_dim < 1: raise ComponentError("De-interleaving needs the head dimension the parts are grouped by.") heads = tensor.shape[0] // (parts * head_dim) - reshaped = tensor.reshape(heads, parts, head_dim, *tensor.shape[1:]) + return _swap01(tensor, (heads, parts, head_dim)) + + +def interleave_rows(tensor: Any, parts: int, head_dim: int) -> Any: + """``[p0_all; p1_all; p2_all]`` back to per-head groups: the inverse of ``deinterleave_rows``. + + Needed to *write* a checkpoint or adapter in a publisher's interleaved layout, where the load + path only ever reads one.""" + if head_dim < 1: + raise ComponentError("Interleaving needs the head dimension the parts are grouped by.") + heads = tensor.shape[0] // (parts * head_dim) + return _swap01(tensor, (parts, heads, head_dim)) + + +def _swap01(tensor: Any, shape: tuple[int, int, int]) -> Any: + """Reshape to ``shape`` plus the trailing dims, exchange the first two, flatten back. + + ``transpose`` is not the same call in torch and numpy - torch swaps two axes, numpy wants a full + permutation - so the swap is spelled per backend rather than duck-typed.""" + reshaped = tensor.reshape(*shape, *tensor.shape[1:]) if _is_torch(tensor): moved = reshaped.transpose(0, 1).contiguous() else: diff --git a/core/src/inline_core/models/lora.py b/core/src/inline_core/models/lora.py index a6327d8..5d52391 100644 --- a/core/src/inline_core/models/lora.py +++ b/core/src/inline_core/models/lora.py @@ -29,6 +29,10 @@ #: while ostris' training adapter uses the reference names. Alias = Callable[[str], str | None] +#: Rewrites a whole adapter before it is matched, for an arch whose checkpoint keys need more than a +#: rename - MiniMax H3 ships attention fused, so three of our modules are one of theirs. +Translate = Callable[[dict[str, Any]], dict[str, Any]] + _DOWN = ("lora_down.weight", "lora_A.weight", "lora_A.default.weight") _UP = ("lora_up.weight", "lora_B.weight", "lora_B.default.weight") # Prefixes checkpoints put in front of the module path; stripped when matching against the model. @@ -39,12 +43,22 @@ LoraPlan = dict[str, list[tuple[Any, Any, float]]] -def fuse_loras(model: Any, loras: tuple[LoraRef, ...], alias: Alias | None = None) -> None: +def fuse_loras( + model: Any, + loras: tuple[LoraRef, ...], + alias: Alias | None = None, + translate: Translate | None = None, +) -> None: """Merge each LoRA into ``model``'s weights in order. No-op for an empty stack.""" - apply_plan(model, plan_loras(model, loras, alias)) + apply_plan(model, plan_loras(model, loras, alias, translate)) -def plan_loras(model: Any, loras: tuple[LoraRef, ...], alias: Alias | None = None) -> LoraPlan: +def plan_loras( + model: Any, + loras: tuple[LoraRef, ...], + alias: Alias | None = None, + translate: Translate | None = None, +) -> LoraPlan: """Resolve every LoRA against ``model``'s module names, without touching any weights. Split from the fusing so a streaming loader can validate the whole stack **before** reading a @@ -53,7 +67,7 @@ def plan_loras(model: Any, loras: tuple[LoraRef, ...], alias: Alias | None = Non plan: LoraPlan = {} names = _linear_module_names(model) for lora in loras: - _plan_one(plan, names, lora.file, lora.strength, alias) + _plan_one(plan, names, lora.file, lora.strength, alias, translate) return plan @@ -74,7 +88,12 @@ def apply_plan(module: Any, plan: LoraPlan, prefix: str = "") -> None: def _plan_one( - plan: LoraPlan, names: dict[str, None], path: str, strength: float, alias: Alias | None + plan: LoraPlan, + names: dict[str, None], + path: str, + strength: float, + alias: Alias | None, + translate: Translate | None = None, ) -> None: from safetensors.torch import load_file @@ -83,6 +102,8 @@ def _plan_one( except Exception as exc: # noqa: BLE001 raise ComponentError(f"Could not read LoRA {path!r}: {exc}") from exc + if translate is not None: + state = translate(state) pairs, alphas = _group(state) if not pairs: raise ComponentError(f"LoRA {path!r} contains no recognisable lora_down/lora_up pairs.") @@ -104,27 +125,39 @@ def _plan_one( ) +#: Most fp32 delta held at once. The product is computed a slice of output rows at a time because +#: the whole of it is enormous: one MiniMax H3 block's six Linears come to 1.5GB and the model to +#: 80GB, which is host RAM during a staged load and took a 60GB box down three times. +_DELTA_CHUNK_BYTES = 64 * 1024 * 1024 + + def _add_delta(weight: Any, up: Any, down: Any, scale: float) -> None: """Fuse ``scale * (up @ down)`` into ``weight`` in place. - Computed on the weight's own device: a big LoRA (Krea 2's are ~260 modules on a 12.9B model) - materializes tens of GB of fp32 deltas, and doing that on the CPU costs ~20s of maths plus the - transfer where the GPU takes ~2s. Falls back to the CPU if the device runs out of memory, so a - tight card still fuses, just slowly.""" + Computed on the weight's own device: doing it on the CPU costs ~20s of maths plus the transfer + where the GPU takes ~2s. Falls back to the CPU if the device runs out of memory, so a tight card + still fuses, just slowly.""" import torch try: - weight.add_(_delta(up, down, weight, weight.device) * scale) + _accumulate(weight, up, down, scale, weight.device) except torch.cuda.OutOfMemoryError: - weight.add_((_delta(up, down, weight, "cpu") * scale).to(weight.device, weight.dtype)) + _accumulate(weight, up, down, scale, "cpu") -def _delta(up: Any, down: Any, weight: Any, device: Any) -> Any: - """``up @ down`` on ``device``, shaped and typed for the target weight. Conv LoRAs flatten the - spatial dims.""" +def _accumulate(weight: Any, up: Any, down: Any, scale: float, device: Any) -> None: + """Add the product into ``weight`` in row slices. Conv LoRAs flatten the spatial dims.""" dtype = _fuse_dtype(up) - delta = up.to(device, dtype=dtype).flatten(1) @ down.to(device, dtype=dtype).flatten(1) - return delta.reshape(weight.shape).to(weight.dtype) + rows = up.to(device, dtype=dtype).flatten(1) + cols = down.to(device, dtype=dtype).flatten(1) + if not weight.is_contiguous(): # a non-contiguous target cannot be written through a view + weight.add_((rows @ cols).reshape(weight.shape).to(weight.dtype) * scale) + return + target = weight.view(rows.shape[0], -1) + step = max(1, _DELTA_CHUNK_BYTES // max(1, cols.shape[1] * 4)) + for start in range(0, rows.shape[0], step): + stop = start + step + target[start:stop].add_((rows[start:stop] @ cols).to(weight.dtype) * scale) def _fuse_dtype(tensor: Any) -> Any: @@ -141,20 +174,29 @@ def _alpha_scale(alpha: Any, rank: int) -> float: return float(alpha.item() if hasattr(alpha, "item") else alpha) / float(rank) +def split_key(key: str) -> tuple[str, str] | None: + """``…to_q.lora_A.weight`` to ``("…to_q", "down")``. None for anything that is not a LoRA key. + + Public so a per-arch key translator groups by exactly the suffixes the fuser recognises: a + convention known to one and not the other would drop tensors silently.""" + for suffix in _DOWN: + if key.endswith("." + suffix): + return key[: -len(suffix) - 1], "down" + for suffix in _UP: + if key.endswith("." + suffix): + return key[: -len(suffix) - 1], "up" + if key.endswith(".alpha"): + return key[: -len(".alpha")], "alpha" + return None + + def _group(state: dict[str, Any]) -> tuple[dict[str, tuple[Any, Any]], dict[str, Any]]: - downs: dict[str, Any] = {} - ups: dict[str, Any] = {} - alphas: dict[str, Any] = {} + parts: dict[str, dict[str, Any]] = {} for key, value in state.items(): - for suffix in _DOWN: - if key.endswith("." + suffix): - downs[key[: -len(suffix) - 1]] = value - for suffix in _UP: - if key.endswith("." + suffix): - ups[key[: -len(suffix) - 1]] = value - if key.endswith(".alpha"): - alphas[key[: -len(".alpha")]] = value - return {k: (downs[k], ups[k]) for k in downs if k in ups}, alphas + if (split := split_key(key)) is not None: + parts.setdefault(split[0], {})[split[1]] = value + pairs = {k: (v["down"], v["up"]) for k, v in parts.items() if "down" in v and "up" in v} + return pairs, {k: v["alpha"] for k, v in parts.items() if "alpha" in v} def _linear_module_names(model: Any) -> dict[str, None]: diff --git a/core/src/inline_core/models/minimaxh3/adaln.py b/core/src/inline_core/models/minimaxh3/adaln.py index e3ad59a..8a53a27 100644 --- a/core/src/inline_core/models/minimaxh3/adaln.py +++ b/core/src/inline_core/models/minimaxh3/adaln.py @@ -41,6 +41,11 @@ import torch from torch import nn +from .vendor.transformer_minimax_h3 import ( + MiniMaxH3AdaLayerNormModulation, + MiniMaxH3AdaLayerNormOut, +) + logger = logging.getLogger("inline_core.minimaxh3") #: Rank kept. Five directions carry the energy; eight matches the published build and leaves slack. @@ -144,3 +149,100 @@ def nbytes(module: Any) -> int: len(model.transformer_blocks), saved / 1e9, ) return saved + + +# --- the published pruned builds ------------------------------------------------------------ +# +# MiniMax ship `pruned` checkpoints that do this same rank-8 reduction ahead of time, and go one +# step further: the timestep path itself is gone. There is no `time_embedder` in the file at all, +# only `adaln_t_table [1025, 8]`, holding `silu(temb)` already projected into their basis at 1025 +# points across t in [0, 1]. So the branch cannot be rebuilt as a basis applied to a `silu(temb)` we +# compute; the table has to be read directly. +# +# Off-grid timesteps are interpolated. Measured against the full bf16 weights at 24 random t, the +# table reaches 1.636e-4 relative with linear interpolation and 1.874e-4 taking the nearest row, +# where one bf16 ulp of the reference is 1.307e-3. The grid is dense enough that interpolating costs +# nothing and removes the sampler constraint a lookup would otherwise impose. + +#: Rows in the published table. Checked, not assumed: a build on a different grid must not be read +#: as though it were on this one. +TABLE_ROWS = 1025 + + +class TableEmbedder(nn.Module): + """Stands in for ``time_proj`` + ``time_embedder``, returning the pruned build's rank-8 row. + + ``linear_1`` exists because the port reads ``time_embedder.linear_1.weight.dtype`` to cast its + input; it carries the table's dtype and nothing else, which keeps ``vendor/`` verbatim. + """ + + #: Declared so the buffer reads as a tensor; ``register_buffer`` alone types as ``Module``. + table: torch.Tensor + + def __init__(self, table: torch.Tensor) -> None: + super().__init__() + self.register_buffer("table", table, persistent=True) + self.linear_1 = nn.Linear(1, 1, bias=False, dtype=table.dtype) + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + rows = self.table.shape[0] + position = timestep.to(self.table.dtype).flatten() * (rows - 1) + low = position.floor().long().clamp(0, rows - 2) + frac = (position - low).unsqueeze(1) + return self.table[low] * (1 - frac) + self.table[low + 1] * frac + + +class TabulatedModulation(MiniMaxH3AdaLayerNormModulation): + """``adaln_proj`` reading a table row. It already holds ``silu(temb)``, so no activation.""" + + def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]: + out = self.linear(temb.to(self.linear.weight.dtype)).view(-1, 6 * self.hidden_size) + return out.chunk(6, dim=-1) + + +class TabulatedNormOut(MiniMaxH3AdaLayerNormOut): + """``norm_out`` reading a table row, otherwise the port's own forward.""" + + def forward( + self, hidden_states: torch.Tensor, temb: torch.Tensor, timestep_indices: torch.Tensor + ) -> torch.Tensor: + shift, scale = self.linear(temb.to(self.linear.weight.dtype)).chunk(2, dim=-1) + hidden_states = self.norm(hidden_states) + return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + + +@torch.no_grad() +def tabulate(model: Any, table: torch.Tensor) -> None: + """Rebuild the timestep path around a pruned build's table, on the meta device before streaming. + + Every module the table feeds changes shape, so this has to happen before any weight is placed: + the rank-8 ``adaln_proj`` in the file would otherwise be assigned into a ``[96768, 2688]`` slot. + """ + if table.ndim != 2 or table.shape[0] != TABLE_ROWS: + raise ValueError( + f"adaln_t_table is {tuple(table.shape)}, not [{TABLE_ROWS}, rank]. This build is on a " + "different timestep grid from the one measured, and reading it as if it were not would " + "shift the modulation at every step while still rendering." + ) + rank = int(table.shape[1]) + model.time_proj = nn.Identity() + model.time_embedder = TableEmbedder(table) + for block in model.transformer_blocks: + block.adaln_proj = _retyped(block.adaln_proj, TabulatedModulation, rank) + model.norm_out = _retyped(model.norm_out, TabulatedNormOut, rank) + + +def _retyped(module: Any, cls: type, rank: int) -> Any: + """The same module with its projection narrowed to ``rank`` inputs, still on meta.""" + replacement = module + replacement.__class__ = cls + old = module.linear + with torch.device("meta"): + replacement.linear = nn.Linear(rank, old.out_features, bias=old.bias is not None) + return replacement + + +#: Parameters ``tabulate`` creates that no checkpoint fills. +TABULATED_SELF_COMPUTED = ("time_embedder.linear_1.weight", "time_embedder.table") diff --git a/core/src/inline_core/models/minimaxh3/keys.py b/core/src/inline_core/models/minimaxh3/keys.py index 79a75fd..f4a4e3d 100644 --- a/core/src/inline_core/models/minimaxh3/keys.py +++ b/core/src/inline_core/models/minimaxh3/keys.py @@ -23,7 +23,7 @@ from __future__ import annotations -from ..keymap import AssertEqual, KeyPlan, Rename, RowLayout, Split, SwapHalves +from ..keymap import AssertEqual, Drop, KeyPlan, Rename, RowLayout, Split, SwapHalves #: Bumping this invalidates every prepared artifact built by the old plan. PLAN_VERSION = "minimax-h3.keys.1" @@ -81,9 +81,17 @@ def build_plan( num_blocks: int = NUM_BLOCKS, num_refiner_blocks: int = NUM_REFINER_BLOCKS, head_dim: int = HEAD_DIM, + pruned: bool = False, + sidecars: tuple[str, ...] = (), ) -> KeyPlan: """The plan for a publisher's layout. ``source`` selects how the fused QKV rows are arranged. + ``pruned`` is the published rank-8 build: it ships ``adaln_t_table`` in place of the whole + timestep path, so the two ``time_embedder`` projections are simply not in the file. + + ``sidecars`` are the quantisation tensors an fp8 build carries beside each weight. They are + consumed while streaming and dropped here, so the coverage check still accounts for every key. + The counts are arguments so a round-trip test can exercise the same code at a size that fits in memory; the defaults are the released geometry. """ @@ -95,7 +103,16 @@ def build_plan( ) from None actions: dict[str, object] = {} - for stem, target in _TOP_LEVEL.items(): + top_level = { + stem: target + for stem, target in _TOP_LEVEL.items() + if not (pruned and stem.startswith("time_embedder.")) + } + if pruned: + actions["adaln_t_table"] = Drop("read before streaming, to rebuild the timestep path") + for key in sidecars: + actions[key] = Drop("a quantisation scale, applied to its weight while streaming") + for stem, target in top_level.items(): for suffix in ("weight", "bias"): actions[f"{stem}.{suffix}"] = Rename(f"{target}.{suffix}") for key, target in _WEIGHT_ONLY.items(): @@ -126,9 +143,16 @@ def build_plan( continue actions[f"{src}.{stem}"] = Rename(f"{dst}.{target}") - return KeyPlan(version=f"{PLAN_VERSION}+{source}", actions=actions) # type: ignore[arg-type] + suffix = source + ("+pruned" if pruned else "") + ("+fp8" if sidecars else "") + return KeyPlan(version=f"{PLAN_VERSION}+{suffix}", actions=actions) # type: ignore[arg-type] -def self_computed_targets() -> set[str]: +def self_computed_targets(*, pruned: bool = False) -> set[str]: """Targets the port builds itself, which ``check_coverage`` must not demand be filled.""" - return {"rope.inv_freq"} + if not pruned: + return {"rope.inv_freq"} + return {"rope.inv_freq", *_PRUNED_SELF_COMPUTED} + + +#: What ``adaln.tabulate`` creates in place of the timestep path. No checkpoint fills these. +_PRUNED_SELF_COMPUTED = ("time_embedder.table", "time_embedder.linear_1.weight") diff --git a/core/src/inline_core/models/minimaxh3/load.py b/core/src/inline_core/models/minimaxh3/load.py index fde51c3..cd53c25 100644 --- a/core/src/inline_core/models/minimaxh3/load.py +++ b/core/src/inline_core/models/minimaxh3/load.py @@ -13,6 +13,7 @@ import json import logging +import re from collections.abc import Iterator from pathlib import Path from typing import Any @@ -32,6 +33,7 @@ row_stats, transform, ) +from . import adaln, lora_keys from . import keys as h3keys from .vendor import MiniMaxH3Transformer3DModel @@ -40,6 +42,19 @@ #: The tensor the layout is measured on. Block 0's fused QKV is present in every published build. _PROBE_KEY = "blocks.0.attn.qkv_proj.weight" +#: Only the pruned builds carry it, and it stands in for the entire timestep path. +_TABLE_KEY = "adaln_t_table" + +#: What an fp8 build writes beside each quantised weight. ``weight_scale`` is the scalar the weight +#: multiplies back by; ``input_scale`` is for an fp8 matmul we do not do, and ``comfy_quant`` names +#: the format, which ``requirements.py`` has already checked by the time a load starts. +_SIDECAR_SUFFIXES = (".weight_scale", ".input_scale", ".comfy_quant") +_SCALE_SUFFIX = ".weight_scale" + +#: The rank-8 projections a pruned build ships. 75 MB across the model, so they stay in float32: +#: rounding them to bf16 would cost more accuracy than the whole factorisation does. +_ADALN_LINEAR = re.compile(r"(?:adaln_proj|norm_out)\.linear\.(?:weight|bias)$") + #: Source config name -> the vendored port's constructor argument. The port's defaults already match #: the released checkpoints, but a future build may not, so the file wins over the default. _CONFIG_MAP = { @@ -110,6 +125,7 @@ def load_transformer( layout: RowLayout | None = None, shrink: Any = None, loras: tuple[Any, ...] = (), + quantised: bool = False, ) -> MiniMaxH3Transformer3DModel: """Build the port and stream ``path`` into it through the key plan. @@ -128,9 +144,12 @@ def load_transformer( # Resolved against module names on the meta model, so a LoRA trained for another architecture # is refused before the 62 GB read rather than a block into it. - lora_plan = lora_module.plan_loras(model, loras) if loras else {} + lora_plan = ( + lora_module.plan_loras(model, loras, translate=lora_keys.adapt) if loras else {} + ) fused: set[str] = set() - shrink = _fusing_shrink(lora_plan, fused, shrink) if lora_plan else shrink + strength: list[float] = [] + shrink = _fusing_shrink(lora_plan, fused, shrink, strength) if lora_plan else shrink with safe_open(str(path), framework="pt") as handle: source_keys = list(handle.keys()) # noqa: SIM118 - safe_open has no __contains__ @@ -149,10 +168,30 @@ def load_transformer( handle.get_tensor(_PROBE_KEY), head_dim=geometry["head_dim"] ) logger.info("MiniMax H3 checkpoint %s: QKV rows are %s", path.name, measured.value) - plan = h3keys.build_plan(_source_for(measured), **geometry) - _check_plan(plan, source_keys, model) + # The pruned builds ship no timestep path at all, so the modules the table feeds have to be + # rebuilt at their reduced width before a single weight is placed into them. + sidecars = tuple(k for k in source_keys if k.endswith(_SIDECAR_SUFFIXES)) + scales = { + f"{k[: -len(_SCALE_SUFFIX)]}.weight": handle.get_tensor(k) + for k in source_keys + if k.endswith(_SCALE_SUFFIX) + } + if scales: + logger.info( + "MiniMax H3 %s is an fp8 build: dequantising %d weights on the way in", + path.name, len(scales), + ) + pruned = _TABLE_KEY in source_keys + if pruned: + adaln.tabulate(model, handle.get_tensor(_TABLE_KEY)) + logger.info("MiniMax H3 %s is a pruned build: reading its AdaLN table", path.name) + plan = h3keys.build_plan( + _source_for(measured), pruned=pruned, sidecars=sidecars, **geometry + ) + _check_plan(plan, source_keys, model, pruned=pruned) filled = _stream_into( - model, handle, plan, dtype=dtype, device=device, shrink=shrink + model, handle, plan, dtype=dtype, device=device, shrink=shrink, pruned=pruned, + scales=scales, ) if lora_plan: @@ -165,12 +204,22 @@ def load_transformer( len(lora_plan), ", ".join(f"{Path(ref.file).name}@{ref.strength:g}" for ref in loras), ) - _assert_nothing_left_on_meta(model, filled) + _report_strength(strength, quantised) + _assert_nothing_left_on_meta(model, filled, pruned=pruned) model.eval() return model -def _fusing_shrink(plan: Any, fused: set[str], inner: Any) -> Any: +#: A fused delta smaller than this fraction of one quantisation step is rounded away rather than +#: applied. Measured on a real adapter: 0.008 of a step, which flipped 0.92% of int8 codes and +#: delivered none of its intended output change. Adapter strength itself is not the test - published +#: LoRAs that work well measure anywhere from 0.017% to 1.2% of the weight norm. +_LOST_TO_QUANTISATION = 0.25 + + +def _fusing_shrink( + plan: Any, fused: set[str], inner: Any, strength: list[float] | None = None +) -> Any: """Fuse a block's share of the LoRA stack the moment it lands, then hand off to ``shrink``. The only window that works: after the stream the weights exist, before ``shrink`` they are @@ -181,24 +230,74 @@ def shrink(model: Any, prefix: str) -> None: module = model for part in prefix.split("."): module = module[int(part)] if part.isdigit() else getattr(module, part) - fused.update(_fuse_subtree(module, plan, f"{prefix}.")) + fused.update(_fuse_subtree(module, plan, f"{prefix}.", strength)) if inner is not None: inner(model, prefix) return shrink -def _fuse_subtree(module: Any, plan: Any, prefix: str) -> set[str]: +def _fuse_subtree( + module: Any, plan: Any, prefix: str, strength: list[float] | None = None +) -> set[str]: """Apply the plan's share for one subtree, reporting which of its targets were covered.""" hit = { path for name, _child in module.named_modules() if (path := f"{prefix}{name}" if prefix else name) in plan } + if strength is not None and not strength and hit: + _measure_strength(module, plan, prefix, hit, strength) lora_module.apply_plan(module, plan, prefix) return hit +def _measure_strength( + module: Any, plan: Any, prefix: str, hit: set[str], out: list[float] +) -> None: + """``|B@A| / |W|`` on the first adapted layer, read before the fuse and before quantisation. + + An adapter that trained but barely moved renders exactly like one that never loaded, and until + this number was in the log the only way to tell them apart was to render twice and compare.""" + path = sorted(hit)[0] + child = dict(module.named_modules()).get(path[len(prefix) :]) + weight = getattr(child, "weight", None) + if weight is None: + return + reference = weight.detach().to(torch.float32) + delta = None + for down, up, scale in plan[path]: + step = up.to(reference.device, torch.float32) @ down.to(reference.device, torch.float32) + delta = step * scale if delta is None else delta + step * scale + norm = float(reference.norm()) + if delta is None or not norm: + return + delta = delta.reshape(reference.shape) + # int8 is symmetric per output row, so one step is the row's largest magnitude over 127. + quant_step = float((reference.abs().amax(dim=1) / 127.0).median()) if reference.ndim == 2 else 0 + out.append(float(delta.norm()) / norm) + out.append(float(delta.abs().mean()) / quant_step if quant_step else 0.0) + + +def _report_strength(strength: list[float], quantised: bool) -> None: + """Say what the adapter did, and warn only when quantisation is about to discard it.""" + if len(strength) < 2: + return + ratio, per_step = strength + logger.info( + "MiniMax H3: the LoRA moves that layer by %.3f%% of its weight norm (%.2f of an int8 step)", + ratio * 100, per_step, + ) + if quantised and per_step < _LOST_TO_QUANTISATION: + logger.warning( + "MiniMax H3: this base is quantised and the LoRA changes each weight by only %.2f of " + "one int8 step, so rounding discards nearly all of it and the render will look " + "unadapted. This is not a weak LoRA, it is the quantisation. Raising the LoRA strength " + "does not fix it. A card that holds the base unquantised does.", + per_step, + ) + + def _finish_fuse(model: Any, plan: Any, fused: set[str]) -> None: """Fuse what the block callback never saw, then prove nothing was missed. @@ -225,11 +324,13 @@ def _source_for(layout: RowLayout) -> str: raise ComponentError(f"No key plan for a {layout.value} checkpoint.") -def _check_plan(plan: Any, source_keys: list[str], model: Any) -> None: +def _check_plan(plan: Any, source_keys: list[str], model: Any, *, pruned: bool = False) -> None: from ..keymap import check_coverage targets = set(dict(model.named_parameters()) | dict(model.named_buffers())) - check_coverage(plan, source_keys, sorted(targets - h3keys.self_computed_targets())) + check_coverage( + plan, source_keys, sorted(targets - h3keys.self_computed_targets(pruned=pruned)) + ) def _stream_into( @@ -240,6 +341,8 @@ def _stream_into( dtype: torch.dtype, device: str, shrink: Any = None, + pruned: bool = False, + scales: dict[str, torch.Tensor] | None = None, ) -> set[str]: """Place every tensor, optionally shrinking each transformer block as soon as it is complete. @@ -260,12 +363,16 @@ def _stream_into( _assign(model, key, shipped.to(device=device)) filled.add(key) continue - for target, value in transform(key, handle.get_tensor(key), action): + source = handle.get_tensor(key) + if scales and (scale := scales.get(key)) is not None: + source = source.to(torch.float32) * scale.to(torch.float32) + for target, value in transform(key, source, action): block = _block_prefix(target) if shrink is not None and pending is not None and block != pending: shrink(model, pending) pending = block if shrink is not None else None - _assign(model, target, value.to(dtype=dtype, device=device)) + placed = torch.float32 if pruned and _ADALN_LINEAR.search(target) else dtype + _assign(model, target, value.to(dtype=placed, device=device)) filled.add(target) if shrink is not None and pending is not None: shrink(model, pending) @@ -308,14 +415,14 @@ def _assign(model: Any, key: str, tensor: torch.Tensor) -> None: raise ComponentError(f"The model has no parameter or buffer named {key}.") -def _assert_nothing_left_on_meta(model: Any, filled: set[str]) -> None: +def _assert_nothing_left_on_meta(model: Any, filled: set[str], *, pruned: bool = False) -> None: """A parameter still on the meta device was never assigned, which the coverage check should have caught. Belt and braces, because the failure downstream is an inscrutable meta-tensor error deep in a forward pass.""" stranded = sorted( name for name, tensor in (dict(model.named_parameters()) | dict(model.named_buffers())).items() - if tensor.is_meta and name not in h3keys.self_computed_targets() + if tensor.is_meta and name not in h3keys.self_computed_targets(pruned=pruned) ) if stranded: raise ComponentError( diff --git a/core/src/inline_core/models/minimaxh3/lora_keys.py b/core/src/inline_core/models/minimaxh3/lora_keys.py new file mode 100644 index 0000000..32e55c4 --- /dev/null +++ b/core/src/inline_core/models/minimaxh3/lora_keys.py @@ -0,0 +1,235 @@ +"""Translate MiniMax H3 LoRAs between our diffusers module names and the reference's fused ones. + +An adapter trained here attaches to the diffusers port's modules +(``transformer_blocks.0.attn.to_q``), while every other tool keys off the released checkpoint's own +names (``blocks.0.attn.qkv_proj``). The two disagree by the same transforms ``keys.py`` already +declares for the base weights, so this is the LoRA-shaped view of that one map rather than a second +copy of it that can drift. + +Two facts make the round trip exact: + +* Every transform acts on **output rows**, so it rewrites ``lora_B`` and leaves ``lora_A`` alone. + The fused QKV is the exception, because three modules here are one there. +* Fusing q, k and v cannot keep rank ``r``: each has its own ``A``. Stacking the three ``A`` + matrices and making ``B`` block-diagonal gives the identical delta at rank ``3r``, and the alpha + triples with it so ``alpha / rank`` is unchanged. Missing that last part divides the adapter by + three, silently. + +The unresolvable half is stated on ``import_reference``: fused row order differs by publisher, and +an adapter carries no base weights to measure it against. +""" + +from __future__ import annotations + +from typing import Any + +from ...errors import ComponentError +from ..keymap import RowLayout, deinterleave_rows, interleave_rows +from ..lora import split_key +from .keys import HEAD_DIM, SOURCE_LAYOUTS + +#: Reference stem -> ours, for the block Linears an adapter can attach to. The norms in +#: ``keys._BLOCK_RENAMES`` are absent on purpose: they are not Linear and never carry a LoRA. +_RENAMES = {"attn.out_proj": "attn.to_out.0", "mlp.fc2": "ff.net.2"} + +#: Same, but the two halves are exchanged, for the ``SwiGLU`` reason ``keys.py`` gives. +_SWAPPED = {"mlp.fc1": "ff.net.0.proj"} + +#: The fused stem and the three it becomes, in row order. +_QKV = "attn.qkv_proj" +_QKV_PARTS = ("attn.to_q", "attn.to_k", "attn.to_v") + +#: Linears outside any block. ``keys.py`` maps more; these are the only ones +#: ``training/arch._MINIMAX_H3_TARGETS`` lets an adapter reach. +_TOP_LEVEL = {"condition_proj": "context_embedder"} + +#: Block prefixes, longest first so the refiner is not eaten by the bare ``blocks.`` rule. +_PREFIXES = ( + ("token_refiner.blocks.", "token_refiner.refiner_blocks."), + ("blocks.", "transformer_blocks."), +) + +#: Stripped from an incoming key. ComfyUI writes the first, PEFT the third. +_WRAPPERS = ("diffusion_model.", "transformer.", "base_model.model.") + +_UNSWAP = {v: k for k, v in _SWAPPED.items()} +_UNRENAME = {v: k for k, v in _RENAMES.items()} +_UNTOP = {v: k for k, v in _TOP_LEVEL.items()} + + +def export_reference(state: dict[str, Any], *, target: str = "comfy-org") -> dict[str, Any]: + """Our adapter in the reference's fused key names, so other tools can load it.""" + layout = _layout(target) + grouped = _group(state) + out: dict[str, Any] = {} + for stem in sorted(grouped): + parsed = _parse(stem, reference=False) + if parsed is None: + if (name := _UNTOP.get(stem)) is not None: + _emit(out, f"diffusion_model.{name}", grouped[stem]) + continue + prefix, index, tail = parsed + block = f"{prefix[0]}{index}." + if tail in _QKV_PARTS: + # Only q drives the fuse; k and v come with it and are skipped when their turn comes. + if tail == _QKV_PARTS[0]: + fused = _fuse_qkv(f"{prefix[1]}{index}.", grouped, layout) + _emit(out, f"diffusion_model.{block}{_QKV}", fused) + elif (name := _UNSWAP.get(tail)) is not None: + _emit(out, f"diffusion_model.{block}{name}", _with_swapped_up(grouped[stem])) + elif (name := _UNRENAME.get(tail)) is not None: + _emit(out, f"diffusion_model.{block}{name}", grouped[stem]) + if not out: + raise ComponentError("Nothing in this adapter maps onto MiniMax H3's key names.") + return out + + +def import_reference(state: dict[str, Any], *, source: str = "comfy-org") -> dict[str, Any]: + """A third-party H3 adapter rewritten onto the diffusers port's module names. + + ``source`` picks the fused QKV row order. It cannot be measured here: + ``keymap.detect_row_layout`` works off a checkpoint's real weights and an adapter has none, so a + wrong guess renders a plausible wrong video rather than failing. The default is the layout + published H3 LoRAs are overwhelmingly trained against; the caller names the other explicitly. + """ + layout = _layout(source) + grouped = _group(state) + out: dict[str, Any] = {} + for stem in sorted(grouped): + parsed = _parse(stem, reference=True) + if parsed is None: + if (name := _TOP_LEVEL.get(stem)) is not None: + _emit(out, name, grouped[stem]) + continue + prefix, index, tail = parsed + block = f"{prefix[1]}{index}." + if tail == _QKV: + for part, tensors in _split_qkv(block, grouped[stem], layout).items(): + _emit(out, part, tensors) + elif (name := _SWAPPED.get(tail)) is not None: + _emit(out, f"{block}{name}", _with_swapped_up(grouped[stem])) + elif (name := _RENAMES.get(tail)) is not None: + _emit(out, f"{block}{name}", grouped[stem]) + if not out: + raise ComponentError( + "No MiniMax H3 layers found in this adapter; it was trained for a different model." + ) + return out + + +def adapt(state: dict[str, Any], *, source: str = "comfy-org") -> dict[str, Any]: + """Bring any H3 adapter onto our module names: translate a reference-keyed one, pass ours on.""" + if not is_reference(state): + return state + return import_reference(state, source=source) + + +def is_reference(state: dict[str, Any]) -> bool: + """Whether this adapter is keyed to the released checkpoint rather than the diffusers port. + + Checked on the prefix, not on ``qkv_proj``: an adapter that only touched the feed-forward has no + fused tensor to give it away.""" + for key in state: + if (split := split_key(key)) is None: + continue + stem = split[0] + for wrapper in _WRAPPERS: + stem = stem.removeprefix(wrapper) + if stem in _TOP_LEVEL or any(stem.startswith(pair[0]) for pair in _PREFIXES): + return True + return False + + +def _layout(name: str) -> RowLayout: + try: + return SOURCE_LAYOUTS[name] + except KeyError: + raise ComponentError( + f"Unknown MiniMax H3 LoRA layout {name!r}; expected one of {sorted(SOURCE_LAYOUTS)}." + ) from None + + +def _parse(stem: str, *, reference: bool) -> tuple[tuple[str, str], str, str] | None: + """``(prefix pair, block index, module tail)``, or None when the stem is not inside a block.""" + for pair in _PREFIXES: + prefix = pair[0] if reference else pair[1] + if stem.startswith(prefix): + index, _, tail = stem[len(prefix) :].partition(".") + return pair, index, tail + return None + + +def _group(state: dict[str, Any]) -> dict[str, dict[str, Any]]: + """LoRA tensors bucketed by module stem, with the wrapper prefixes stripped.""" + grouped: dict[str, dict[str, Any]] = {} + for key, value in state.items(): + if (split := split_key(key)) is None: + continue + stem, role = split + for wrapper in _WRAPPERS: + if stem.startswith(wrapper): + stem = stem[len(wrapper) :] + grouped.setdefault(stem, {})[role] = value + return grouped + + +def _emit(out: dict[str, Any], name: str, tensors: dict[str, Any]) -> None: + for role, suffix in (("down", "lora_A.weight"), ("up", "lora_B.weight"), ("alpha", "alpha")): + if role in tensors: + out[f"{name}.{suffix}" if role != "alpha" else f"{name}.alpha"] = tensors[role] + + +def _with_swapped_up(tensors: dict[str, Any]) -> dict[str, Any]: + return {**tensors, "up": _swap_halves(tensors["up"])} + + +def _fuse_qkv( + block: str, grouped: dict[str, dict[str, Any]], layout: RowLayout +) -> dict[str, Any]: + """``[Bq@Aq; Bk@Ak; Bv@Av]`` as one rank-``3r`` adapter: stacked ``A``, block-diagonal ``B``.""" + import torch + + parts = [grouped.get(f"{block}{part}") for part in _QKV_PARTS] + present = [p for p in parts if p is not None] + if len(present) != len(_QKV_PARTS): + raise ComponentError( + f"{block}attn needs q, k and v adapted together to fuse into one qkv_proj, but " + f"{len(present)} of 3 are in this adapter." + ) + up: Any = torch.block_diag(*(p["up"] for p in present)) + if layout is RowLayout.INTERLEAVED: + up = interleave_rows(up, len(_QKV_PARTS), HEAD_DIM) + tensors: dict[str, Any] = {"down": torch.cat([p["down"] for p in present], dim=0), "up": up} + if (alpha := present[0].get("alpha")) is not None: + # Rank tripled, so alpha must too, or alpha/rank quietly divides the adapter by three. + tensors["alpha"] = alpha * len(_QKV_PARTS) + return tensors + + +def _split_qkv( + block: str, tensors: dict[str, Any], layout: RowLayout +) -> dict[str, dict[str, Any]]: + """The fused adapter back to three, each keeping the shared ``lora_A``.""" + up = tensors["up"] + parts = len(_QKV_PARTS) + if up.shape[0] % parts: + raise ComponentError(f"{block}{_QKV} has {up.shape[0]} rows, not a multiple of {parts}.") + if layout is RowLayout.INTERLEAVED: + up = deinterleave_rows(up, parts, HEAD_DIM) + rows = up.shape[0] // parts + out: dict[str, dict[str, Any]] = {} + for index, part in enumerate(_QKV_PARTS): + split = {"down": tensors["down"], "up": up[index * rows : (index + 1) * rows]} + if (alpha := tensors.get("alpha")) is not None: + split["alpha"] = alpha + out[f"{block}{part}"] = split + return out + + +def _swap_halves(tensor: Any) -> Any: + import torch + + half = tensor.shape[0] // 2 + if half * 2 != tensor.shape[0]: + raise ComponentError("A gated FFN adapter must have an even number of output rows.") + return torch.cat([tensor[half:], tensor[:half]], dim=0) diff --git a/core/src/inline_core/models/minimaxh3/pipeline.py b/core/src/inline_core/models/minimaxh3/pipeline.py index 7a90436..9019260 100644 --- a/core/src/inline_core/models/minimaxh3/pipeline.py +++ b/core/src/inline_core/models/minimaxh3/pipeline.py @@ -116,6 +116,14 @@ def load_pipeline( ) audio_vae = _require(reqs.resolve("vae", reqs.AUDIO_VAE_FILE), "the audio VAE") + # A pruned build has already had this done to it, and re-running the transform would multiply a + # rank-8 projection by a full-width basis. Same shape of rule as never re-quantising a + # prequantized checkpoint, and the same reason: the source is already in the target form. + if reqs.inspect_file(transformer_path).pruned: + if factorise_adaln: + logger.info("%s ships its AdaLN reduced; skipping ours.", transformer_path.name) + factorise_adaln = False + # Hand the policy the on-disk sizes so it fits dtype and quantisation to THIS card, then refuse # an impossible load up front. Without this it falls back to coarse VRAM buckets and tries to # place a 62 GB transformer resident, which no consumer card can hold. Residency is staged, so @@ -317,6 +325,7 @@ def _build( # `render_staged` moves it across once the conditioner is parked. device="cpu" if (recipe.denoiser_offload or staged) else device, loras=loras, + quantised=recipe.quantizes, shrink=_block_shrinker( recipe, transformer_path if factorise_adaln else None, diff --git a/core/src/inline_core/models/minimaxh3/requirements.py b/core/src/inline_core/models/minimaxh3/requirements.py index e9dd1df..e99a8d4 100644 --- a/core/src/inline_core/models/minimaxh3/requirements.py +++ b/core/src/inline_core/models/minimaxh3/requirements.py @@ -13,9 +13,11 @@ import json import struct +from collections.abc import Iterator from dataclasses import dataclass from functools import lru_cache from pathlib import Path +from typing import Any, cast from ...config import models_dir from ..requirements import ModelComponent @@ -26,6 +28,9 @@ MINIMAX_REPO = "MiniMaxAI/MiniMax-H3" FL2VA_FILE = "minimax_h3_fl2va_bf16.safetensors" +#: A third the download for the same model. Generation only: the trainer needs the timestep path a +#: pruned build does not ship, and it saves nothing in VRAM because the base is quantised anyway. +FL2VA_FP8_FILE = "minimax_h3_fl2va_pruned_fp8_scaled.safetensors" REF2VA_FILE = "minimax_h3_ref2va_bf16.safetensors" TEXT_ENCODER_DIR = "FL2VA/text_encoder" VIDEO_VAE_FILE = "minimax_h3_video_vae_fp16.safetensors" @@ -36,8 +41,13 @@ _PROBE_SHAPE = [21504, 5376] #: Only in the pruned builds, whose AdaLN branch is a rank-8 lookup, not a projection. _PRUNED_MARKER = "adaln_t_table" -#: ComfyUI's own quantisation, which carries scale tensors nothing else can read. +#: ComfyUI's own quantisation, which carries scale tensors alongside the weights. _COMFY_QUANT_SUFFIX = ".comfy_quant" +#: The quantised dtypes the published builds use, and whether the loader can read one. +#: fp8 is a scalar scale per weight and nothing else, so it dequantises exactly. int8 is only +#: published with ComfyUI's rotation applied, which is a transform we cannot invert. +_FP8_DTYPE = "F8_E4M3" +_INT8_DTYPE = "I8" @dataclass(frozen=True) @@ -48,26 +58,29 @@ class Candidate: is_h3: bool pruned: bool = False comfy_quantised: bool = False + #: "", "float8_e4m3fn" or "int8". Read from the weight dtypes, not from the filename. + quantisation: str = "" @property def usable(self) -> bool: - return self.is_h3 and not self.pruned and not self.comfy_quantised + return self.is_h3 and self.quantisation in ("", "float8_e4m3fn") @property def reason(self) -> str: - """Why an H3 file cannot be loaded, for the picker to show instead of hiding it.""" + """Why an H3 file cannot be loaded, for the picker to show instead of hiding it. + + Empty for anything loadable, so a caller can treat a reason as proof of refusal.""" + if self.usable: + return "" if not self.is_h3: return "not a MiniMax H3 transformer" - if self.comfy_quantised: - return ( - "a ComfyUI int8 build: it carries comfy_quant scale tensors that only ComfyUI " - "reads. Use the bf16 file; memory saving is the device policy's job here." - ) - if self.pruned: + if self.quantisation in ("int8", "unknown"): return ( - "a pruned build: its AdaLN branch is a rank-8 lookup table rather than a " - "projection, which is a different graph from the one this node runs." + "a ComfyUI int8 build: its weights are stored rotated (convrot), which is a " + "transform only ComfyUI can invert. The fp8 build is the same size and loads." ) + if self.quantisation: + return f"quantised as {self.quantisation}, which this node cannot read" return "" @@ -90,6 +103,13 @@ def read_header(path: Path) -> dict[str, object] | None: return header +def _entries(header: dict[str, object]) -> Iterator[tuple[str, dict[str, Any]]]: + """The tensor records in a header, skipping anything that is not one.""" + for name, info in header.items(): + if isinstance(info, dict): + yield name, cast("dict[str, Any]", info) + + def inspect_file(path: Path) -> Candidate: """Classify a checkpoint from its header alone.""" return _inspect_cached(str(path), *_stamp(path)) @@ -114,14 +134,28 @@ def _inspect_cached(path_str: str, mtime: int, size: int) -> Candidate: is_h3 = isinstance(probe, dict) and list(probe.get("shape", [])) == _PROBE_SHAPE if not is_h3: return Candidate(path, is_h3=False) + dtypes = {str(info.get("dtype")) for _, info in _entries(header)} return Candidate( path, is_h3=True, pruned=any(_PRUNED_MARKER in key for key in header), comfy_quantised=any(key.endswith(_COMFY_QUANT_SUFFIX) for key in header), + quantisation=_quantisation(dtypes, any(k.endswith(_COMFY_QUANT_SUFFIX) for k in header)), ) +def _quantisation(dtypes: set[str], has_sidecars: bool) -> str: + """What a build is quantised as, refusing anything unrecognised rather than guessing. + + A ``comfy_quant`` sidecar in a format we do not know is named ``unknown`` and refused: reading + quantised weights with the wrong recipe renders a plausible wrong video, not an error.""" + if _FP8_DTYPE in dtypes: + return "float8_e4m3fn" + if _INT8_DTYPE in dtypes: + return "int8" + return "unknown" if has_sidecars else "" + + def usable_transformers() -> list[Path]: """Every H3 transformer in ``diffusion_models/`` this node can actually load.""" root = models_dir() / "diffusion_models" @@ -228,6 +262,9 @@ def components(partition: str = "fl2va") -> list[ModelComponent]: "MiniMax-H3-processor", MINIMAX_REPO, "FL2VA/processor"), _file("h3-ref2va", "Ref2VA transformer (66.3 GB)", "diffusion_models", REF2VA_FILE, COMFY_REPO, f"diffusion_models/{REF2VA_FILE}", optional=not ref2va_required), + _file("h3-fl2va-fp8", "FL2VA transformer, fp8 (21.0 GB, generation only)", + "diffusion_models", FL2VA_FP8_FILE, + COMFY_REPO, f"diffusion_models/{FL2VA_FP8_FILE}", optional=True), ] return entries @@ -259,6 +296,32 @@ def _folder( #: under half that. ADALN_SHARE = 0.392 +#: What the model weighs once loaded, always bf16 whatever the file holds. +_RESIDENT_BYTES_PER_PARAM = 2 + + +def resident_bytes(path: Path) -> int: + """What ``path`` will occupy once placed, counted from its own header. + + Not the file size. A pruned build has already had its AdaLN branch reduced, and an fp8 build + stores half a byte-per-param of what it will occupy once dequantised, so scaling the on-disk + number would under-size both. Under-sizing is the dangerous direction: the fit ladder would + promise a machine that then dies to a host-RAM OOM kill rather than raising. + """ + header = read_header(path) + if header is None: + return 0 + total = 0 + for _, info in _entries(header): + shape = info.get("shape") + if not isinstance(shape, list): + continue + count = 1 + for dim in cast("list[Any]", shape): + count *= int(dim) + total += count * _RESIDENT_BYTES_PER_PARAM + return total + def footprint_bytes( partition: str = "fl2va", @@ -284,8 +347,15 @@ def size(path: Path | None) -> int: encoder_bytes = sum(f.stat().st_size for f in encoder.rglob("*") if f.is_file()) if ( encoder.is_dir() ) else 0 - diffusion = size(transformer if transformer is not None else resolve_transformer(partition)) - if factorised: + chosen = transformer if transformer is not None else resolve_transformer(partition) + diffusion = size(chosen) + if chosen is not None and (candidate := inspect_file(chosen)).is_h3: + # Counted from the header, so a pruned or fp8 build is sized by what it becomes rather than + # by what it weighs on disk. + diffusion = resident_bytes(chosen) + if factorised and not candidate.pruned: + diffusion = int(diffusion * (1 - ADALN_SHARE)) + elif factorised: diffusion = int(diffusion * (1 - ADALN_SHARE)) video = video_vae if video_vae is not None else resolve("vae", VIDEO_VAE_FILE) return { diff --git a/core/src/inline_core/training/arch.py b/core/src/inline_core/training/arch.py index b38b06f..39f2683 100644 --- a/core/src/inline_core/training/arch.py +++ b/core/src/inline_core/training/arch.py @@ -127,6 +127,9 @@ class TrainingArch: target: Callable[[Any, Any], Any] #: (transformer, noisy, timestep, item) -> the prediction, same shape as the clean latent. forward: Callable[..., Any] + #: Rewrites the finished adapter into the published checkpoint's key names, for an arch whose + #: names differ from the port's. Without it the LoRA only ever loads back into Inline. + export_keys: Callable[[dict[str, Any]], dict[str, Any]] | None = None # --- Z-Image ------------------------------------------------------------------------------------- @@ -254,6 +257,14 @@ def _flux2_forward(transformer: Any, noisy: Any, timestep: Any, item: dict[str, # --- MiniMax H3 ----------------------------------------------------------------------------------- + +def _h3_export_keys(state: dict[str, Any]) -> dict[str, Any]: + """H3 ships attention fused, so three of the port's modules are one of the checkpoint's.""" + from ..models.minimaxh3.lora_keys import export_reference + + return export_reference(state) + + #: H3's (t, h, w) patch. A still is one latent frame, so only the spatial half ever bites. _H3_PATCH = (1, 2, 2) @@ -321,6 +332,7 @@ def _h3_forward(transformer: Any, noisy: Any, timestep: Any, item: dict[str, Any ), MINIMAX_H3: TrainingArch( key=MINIMAX_H3, + export_keys=_h3_export_keys, target_modules=_MINIMAX_H3_TARGETS, # Same shift expression as Z-Image, at the scheduler's video shift of 12.0. sigma=_zimage_sigma, diff --git a/core/src/inline_core/training/h3.py b/core/src/inline_core/training/h3.py index 0cdf5ca..55c22a2 100644 --- a/core/src/inline_core/training/h3.py +++ b/core/src/inline_core/training/h3.py @@ -404,6 +404,7 @@ def load_base(models_dir: str, device: str, dtype: Any, quant: Any) -> Any: from . import models path = Path(models._require(Path(models_dir), "minimax-h3", "diffusion_models")) + _refuse_pruned(path) # Derived before the stream: the callback needs it while the first block lands, and # `time_embedder.*` sorts after `blocks.*`. Only two tensors are read, about 60 MB of 62 GB. basis = _adaln_basis(path) @@ -416,6 +417,24 @@ def load_base(models_dir: str, device: str, dtype: Any, quant: Any) -> Any: return model +def _refuse_pruned(path: Path) -> None: + """Training needs the full build. A pruned one has no timestep path to derive the basis from. + + Worth saying plainly rather than letting it fail: the smaller builds save download size, not + VRAM, because the base is quantised to 4-bit either way. Nobody gains a card by using one. + """ + from ..models.minimaxh3 import requirements as reqs + + if not reqs.inspect_file(path).pruned: + return + raise RuntimeError( + f"{path.name} is a pruned MiniMax H3 build. It ships the AdaLN branch already reduced and " + "no timestep path at all, which training needs to derive the basis, so it can generate but " + "not train. Use minimax_h3_fl2va_bf16.safetensors. It will not cost you any VRAM: the base " + "trains at 4-bit whichever file you start from." + ) + + def _shrinker(basis: Any, quant: Any, device: str, dtype: Any) -> Any: """Shrink one transformer block as soon as its weights land, and place it.""" from ..device.policy import Quantization diff --git a/core/src/inline_core/training/trainer.py b/core/src/inline_core/training/trainer.py index 9eeb51c..359da9f 100644 --- a/core/src/inline_core/training/trainer.py +++ b/core/src/inline_core/training/trainer.py @@ -94,17 +94,30 @@ def _resume_step(ckpt_dir: Path) -> int: return 0 -def _save_lora(transformer: Any, output_path: str) -> None: - """Write the PEFT adapter as safetensors. Its ``base_model.model...lora_A/lora_B`` keys are read - directly by the loader's fuser (``models/lora.py`` strips the ``base_model.model.`` prefix).""" +def _save_lora( + transformer: Any, output_path: str, *, alpha: int, arch: archs.TrainingArch +) -> None: + """Write the finished adapter as safetensors, in the keys other tools read. + + An arch with ``export_keys`` is written in its published checkpoint's names rather than the + diffusers port's, because a LoRA that only loads back into the app that made it is not much of a + deliverable. Our own loader translates it back on the way in. + + The ``.alpha`` written beside each pair is not decoration. PEFT trains with a scale of + ``alpha / rank`` and saves the factors raw, so an adapter without it fuses at 1.0: correct only + while ``alpha == rank``, which is the default and is why this went unnoticed.""" import torch from peft import get_peft_model_state_dict from safetensors.torch import save_file - state = { + state: dict[str, Any] = { k: v.detach().to("cpu", dtype=torch.float32).contiguous() for k, v in get_peft_model_state_dict(transformer).items() } + for key in [k for k in state if k.endswith(".lora_A.weight")]: + state[f"{key[: -len('.lora_A.weight')]}.alpha"] = torch.tensor(float(alpha)) + if arch.export_keys is not None: + state = arch.export_keys(state) Path(output_path).parent.mkdir(parents=True, exist_ok=True) save_file(state, output_path) @@ -230,10 +243,11 @@ def train(manifest: dict[str, Any]) -> str | None: # flow, so it looks fine, but it is not the path peft tests and merging would be wrong. if quant is Quantization.NF4: transformer.is_loaded_in_4bit = True + lora_alpha = int(hp.get("alpha") or hp["rank"]) transformer.add_adapter( LoraConfig( r=int(hp["rank"]), - lora_alpha=int(hp.get("alpha") or hp["rank"]), + lora_alpha=lora_alpha, lora_dropout=0.0, target_modules=archs.target_modules(arch, str(hp.get("loraScope") or "full")), ) @@ -299,5 +313,10 @@ def train(manifest: dict[str, Any]) -> str | None: accelerator.wait_for_everyone() if accelerator.is_main_process: - _save_lora(accelerator.unwrap_model(transformer), manifest["outputPath"]) + _save_lora( + accelerator.unwrap_model(transformer), + manifest["outputPath"], + alpha=lora_alpha, + arch=arch, + ) return manifest["outputPath"] diff --git a/core/tests/test_lora.py b/core/tests/test_lora.py index 35b9a40..ec1d740 100644 --- a/core/tests/test_lora.py +++ b/core/tests/test_lora.py @@ -10,6 +10,7 @@ from inline_core.graph.loader_runners import LoadLoraRunner, LoraRef from inline_core.graph.registry import build_default_registry from inline_core.graph.schema import Node +from inline_core.models import lora from inline_core.models.loaders import _device_key, _dtype_key, lora_cache_key torch = pytest.importorskip("torch") @@ -216,3 +217,45 @@ def test_fuse_of_an_empty_stack_is_a_noop() -> None: fuse_loras(model, ()) assert torch.equal(model.proj.weight, before) + + +# --- fusing without materialising the whole delta ------------------------------------------------ + + +def _unchunked(weight, up, down, scale): # type: ignore[no-untyped-def] + """What the fuse used to do: one full fp32 product, cast, scaled, added.""" + product = up.to(torch.float32).flatten(1) @ down.to(torch.float32).flatten(1) + return weight + product.reshape(weight.shape).to(weight.dtype) * scale + + +@pytest.mark.parametrize("shape,rank", [((512, 256), 8), ((97, 33), 4), ((64, 8, 3, 3), 4)]) +def test_chunked_fuse_matches_the_whole_product_exactly(shape, rank) -> None: # type: ignore[no-untyped-def] + """The chunking is a memory change only, so the arithmetic must be bit-identical, including for + a shape that does not divide evenly and for a conv LoRA that flattens its spatial dims.""" + torch.manual_seed(0) + fan_in = 1 + for dim in shape[1:]: + fan_in *= dim + weight = torch.randn(*shape, dtype=torch.bfloat16) + up = torch.randn(shape[0], rank, dtype=torch.bfloat16) + down = torch.randn(rank, fan_in, dtype=torch.bfloat16) + + want = _unchunked(weight.clone(), up, down, 1.7) + got = weight.clone() + lora._add_delta(got, up, down, 1.7) + assert torch.equal(got, want) + + +def test_the_delta_is_never_materialised_whole(monkeypatch: pytest.MonkeyPatch) -> None: + """One MiniMax H3 block's Linears are 1.5GB of fp32 product and the model 80GB, which is host + RAM during a staged load. A chunk size below one row still has to land on the same answer.""" + monkeypatch.setattr(lora, "_DELTA_CHUNK_BYTES", 1) + torch.manual_seed(1) + weight = torch.randn(300, 128, dtype=torch.bfloat16) + up = torch.randn(300, 8, dtype=torch.bfloat16) + down = torch.randn(8, 128, dtype=torch.bfloat16) + + want = _unchunked(weight.clone(), up, down, 0.5) + got = weight.clone() + lora._add_delta(got, up, down, 0.5) + assert torch.equal(got, want) diff --git a/core/tests/test_minimaxh3_load.py b/core/tests/test_minimaxh3_load.py index 59ec68c..6e8b80e 100644 --- a/core/tests/test_minimaxh3_load.py +++ b/core/tests/test_minimaxh3_load.py @@ -16,6 +16,8 @@ torch = pytest.importorskip("torch") pytest.importorskip("safetensors") +from safetensors.torch import load_file # noqa: E402 + from inline_core.errors import ComponentError # noqa: E402 from inline_core.models.keymap import ( # noqa: E402 AssertEqual, @@ -24,7 +26,9 @@ Split, SwapHalves, ) +from inline_core.models.minimaxh3 import adaln # noqa: E402 from inline_core.models.minimaxh3 import keys as h3keys # noqa: E402 +from inline_core.models.minimaxh3 import load as h3_load # noqa: E402 from inline_core.models.minimaxh3.load import ( # noqa: E402 detect_source_layout, expected_inv_freq, @@ -255,3 +259,152 @@ def test_detect_source_layout_reads_a_real_shaped_tensor() -> None: ) assert detect_source_layout(contiguous, head_dim=head_dim) is RowLayout.CONTIGUOUS assert detect_source_layout(interleaved, head_dim=head_dim) is RowLayout.INTERLEAVED + + +# --- the published pruned builds ------------------------------------------------------------ + + +def _pruned_source(model: Any, state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """A pruned checkpoint for ``model``: no timestep path, a table, rank-reduced projections. + + Built the way MiniMax's own build is. ``V`` is an orthonormal basis for ``silu(temb)``, the + table holds ``silu(temb) @ V`` sampled across t, and each projection becomes ``W @ V``, so that + ``(silu(temb) @ V) @ (W @ V).T`` is the modulation the unpruned model computes. + """ + plan = h3keys.build_plan("comfy-org", pruned=True, **TINY_PLAN) + source = _invert(plan, state, RowLayout.CONTIGUOUS) + grid = torch.linspace(0, 1, adaln.TABLE_ROWS) + with torch.no_grad(): + activated = torch.nn.functional.silu(model.time_embedder(model.time_proj(grid))).float() + basis = torch.linalg.svd(activated, full_matrices=False)[2].T.contiguous() + source[_TABLE] = (activated @ basis).contiguous() + for key in [k for k in source if k.endswith("adaln_proj.linear.weight")]: + source[key] = (source[key].float() @ basis).contiguous() + return source + + +_TABLE = "adaln_t_table" + + +def test_a_pruned_build_reproduces_the_unpruned_modulation(tmp_path: Path) -> None: + """The whole point: the same modulation, from a file with no timestep path in it.""" + model = _tiny_model() + state = {k: v.detach().clone() for k, v in model.state_dict().items()} + path = _write(tmp_path / "h3_pruned.safetensors", _pruned_source(model, state)) + + loaded = load_transformer(path, dtype=torch.float32).eval() + steps = torch.tensor([0.03, 0.271, 0.5, 0.8125, 0.99]) + with torch.no_grad(): + want = model.transformer_blocks[0].adaln_proj( + model.time_embedder(model.time_proj(steps)) + ) + got = loaded.transformer_blocks[0].adaln_proj( + loaded.time_embedder(loaded.time_proj(steps)) + ) + for a, b in zip(want, got, strict=True): + assert torch.allclose(a, b, atol=2e-4), (a - b).abs().max() + + +def test_a_pruned_build_replaces_the_timestep_path(tmp_path: Path) -> None: + model = _tiny_model() + state = {k: v.detach().clone() for k, v in model.state_dict().items()} + path = _write(tmp_path / "h3_pruned2.safetensors", _pruned_source(model, state)) + loaded = load_transformer(path, dtype=torch.float32) + assert isinstance(loaded.time_embedder, adaln.TableEmbedder) + assert isinstance(loaded.transformer_blocks[0].adaln_proj, adaln.TabulatedModulation) + assert isinstance(loaded.norm_out, adaln.TabulatedNormOut) + # The rank-8 projections stay in float32; rounding them costs more than the factorisation does. + assert loaded.transformer_blocks[0].adaln_proj.linear.weight.dtype is torch.float32 + + +def test_an_unpruned_build_is_untouched_by_any_of_this(reference) -> None: # type: ignore[no-untyped-def] + path, _ = reference(RowLayout.CONTIGUOUS) + loaded = load_transformer(path, dtype=torch.float32) + assert not isinstance(loaded.time_embedder, adaln.TableEmbedder) + assert not isinstance(loaded.norm_out, adaln.TabulatedNormOut) + + +def test_a_table_on_a_different_grid_is_refused(tmp_path: Path) -> None: + """A build sampled at another resolution would shift the modulation at every step and render.""" + model = _tiny_model() + state = {k: v.detach().clone() for k, v in model.state_dict().items()} + source = _pruned_source(model, state) + source[_TABLE] = source[_TABLE][::2].contiguous() + path = _write(tmp_path / "h3_offgrid.safetensors", source) + with pytest.raises(ValueError, match="different timestep grid"): + load_transformer(path, dtype=torch.float32) + + +def test_the_pruned_plan_is_versioned_apart(tmp_path: Path) -> None: + """Prepared artifacts are keyed on the plan version; the two must not collide.""" + plain = h3keys.build_plan("comfy-org", **TINY_PLAN).version + pruned = h3keys.build_plan("comfy-org", pruned=True, **TINY_PLAN).version + assert plain != pruned and "pruned" in pruned + + +def test_an_fp8_weight_is_dequantised_by_its_scale(reference, tmp_path: Path) -> None: # type: ignore[no-untyped-def] + """Values chosen to be exact in fp8, so this measures the scale and nothing else. + + Silently ignoring ``weight_scale`` would leave every quantised weight off by a constant factor, + which renders a washed-out video rather than raising.""" + path, state = reference(RowLayout.CONTIGUOUS) + tensors = dict(load_file(str(path))) + target = "blocks.0.mlp.fc2.weight" + shape = tensors[target].shape + codes = torch.tensor([-2.0, -0.5, 0.5, 2.0]).repeat(shape.numel() // 4).reshape(shape) + tensors[target] = codes.to(torch.float8_e4m3fn) + tensors["blocks.0.mlp.fc2.weight_scale"] = torch.tensor(0.25) + tensors["blocks.0.mlp.fc2.input_scale"] = torch.tensor(1.0) + tensors["blocks.0.mlp.fc2.comfy_quant"] = torch.zeros(27, dtype=torch.uint8) + quantised = _write(tmp_path / "h3_fp8.safetensors", tensors) + + loaded = load_transformer(quantised, dtype=torch.float32) + assert torch.equal(loaded.transformer_blocks[0].ff.net[2].weight, codes * 0.25) + + +def test_the_fp8_plan_is_versioned_apart() -> None: + plain = h3keys.build_plan("comfy-org", **TINY_PLAN).version + fp8 = h3keys.build_plan( + "comfy-org", sidecars=("blocks.0.mlp.fc2.weight_scale",), **TINY_PLAN + ).version + assert plain != fp8 and "fp8" in fp8 + + +# --- what a fused LoRA is worth once the base is quantised ------------------------------------ + + +def test_a_lora_lost_to_quantisation_is_warned_about(caplog) -> None: # type: ignore[no-untyped-def] + """The measured case: 0.008 of an int8 step reaches the weights, so rounding discards it.""" + import logging + + with caplog.at_level(logging.WARNING, logger="inline_core.minimaxh3"): + h3_load._report_strength([0.00038, 0.008], quantised=True) + assert "quantisation" in caplog.text + assert "not a weak LoRA" in caplog.text + + +def test_no_warning_when_the_base_is_not_quantised(caplog) -> None: # type: ignore[no-untyped-def] + """The same adapter applies fine in full precision, so warning there would be noise.""" + import logging + + with caplog.at_level(logging.WARNING, logger="inline_core.minimaxh3"): + h3_load._report_strength([0.00038, 0.008], quantised=False) + assert caplog.text == "" + + +def test_no_warning_when_the_delta_clears_a_quantisation_step(caplog) -> None: # type: ignore[no-untyped-def] + import logging + + with caplog.at_level(logging.WARNING, logger="inline_core.minimaxh3"): + h3_load._report_strength([0.02, 1.5], quantised=True) + assert caplog.text == "" + + +def test_adapter_strength_alone_does_not_trigger_the_warning(caplog) -> None: # type: ignore[no-untyped-def] + """Published LoRAs that work well measure 0.017% to 1.2% of the weight norm, so a threshold on + strength alone would fire on almost all of them.""" + import logging + + with caplog.at_level(logging.WARNING, logger="inline_core.minimaxh3"): + h3_load._report_strength([0.00017, 0.9], quantised=True) + assert caplog.text == "" diff --git a/core/tests/test_minimaxh3_lora_keys.py b/core/tests/test_minimaxh3_lora_keys.py new file mode 100644 index 0000000..224f112 --- /dev/null +++ b/core/tests/test_minimaxh3_lora_keys.py @@ -0,0 +1,249 @@ +"""MiniMax H3 LoRA key translation, both directions. + +Every transform here fails silently when it is wrong - a mis-split QKV or a backwards gated FFN +still loads and still renders - so these tests check the *delta each module receives*, not the key +names. A rename test would pass on all of them. +""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") + +from inline_core.errors import ComponentError # noqa: E402 +from inline_core.models.minimaxh3.lora_keys import ( # noqa: E402 + adapt, + export_reference, + import_reference, + is_reference, +) + +RANK = 4 +HEADS = 2 +HEAD_DIM = 128 +PART = HEADS * HEAD_DIM # rows one of q, k or v contributes +FFN = 16 # rows of the gated FFN's fused pair, half gate and half value +WIDTH = 8 + + +def _pair(rows: int, seed: int) -> dict[str, torch.Tensor]: + g = torch.Generator().manual_seed(seed) + return { + "down": torch.randn(RANK, WIDTH, generator=g), + "up": torch.randn(rows, RANK, generator=g), + } + + +def _ours() -> dict[str, torch.Tensor]: + """A one-block adapter over every module H3 training targets.""" + modules = { + "transformer_blocks.0.attn.to_q": PART, + "transformer_blocks.0.attn.to_k": PART, + "transformer_blocks.0.attn.to_v": PART, + "transformer_blocks.0.attn.to_out.0": PART, + "transformer_blocks.0.ff.net.0.proj": FFN, + "transformer_blocks.0.ff.net.2": FFN, + "token_refiner.refiner_blocks.1.attn.to_q": PART, + "token_refiner.refiner_blocks.1.attn.to_k": PART, + "token_refiner.refiner_blocks.1.attn.to_v": PART, + "context_embedder": WIDTH, + } + state: dict[str, torch.Tensor] = {} + for seed, (name, rows) in enumerate(modules.items()): + for role, suffix in (("down", "lora_A"), ("up", "lora_B")): + state[f"base_model.model.{name}.{suffix}.weight"] = _pair(rows, seed)[role] + return state + + +def _deltas(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """``B @ A`` per module stem, which is the only thing the model ever sees.""" + out: dict[str, torch.Tensor] = {} + for key, value in state.items(): + if key.endswith("lora_B.weight"): + stem = key[: -len(".lora_B.weight")] + out[stem] = value @ state[f"{stem}.lora_A.weight"] + return out + + +def test_export_uses_the_reference_key_names() -> None: + keys = set(export_reference(_ours())) + assert "diffusion_model.blocks.0.attn.qkv_proj.lora_A.weight" in keys + assert "diffusion_model.blocks.0.mlp.fc1.lora_B.weight" in keys + assert "diffusion_model.blocks.0.mlp.fc2.lora_B.weight" in keys + assert "diffusion_model.blocks.0.attn.out_proj.lora_B.weight" in keys + assert "diffusion_model.token_refiner.blocks.1.attn.qkv_proj.lora_B.weight" in keys + assert "diffusion_model.condition_proj.lora_B.weight" in keys + # The split names must be gone, or a tool that matches loosely applies the LoRA twice. + assert not [k for k in keys if "to_q" in k or "transformer_blocks" in k] + + +def test_fused_qkv_carries_the_same_delta_as_the_three_it_replaces() -> None: + ours = _ours() + exported = _deltas(export_reference(ours)) + mine = _deltas(ours) + stacked = torch.cat( + [mine[f"base_model.model.transformer_blocks.0.attn.to_{p}"] for p in "qkv"], dim=0 + ) + fused = exported["diffusion_model.blocks.0.attn.qkv_proj"] + assert fused.shape == (3 * PART, WIDTH) + assert torch.allclose(fused, stacked, atol=1e-6) + + +def test_round_trip_returns_every_original_delta() -> None: + ours = _ours() + back = _deltas(import_reference(export_reference(ours))) + for stem, delta in _deltas(ours).items(): + name = stem.removeprefix("base_model.model.") + assert torch.allclose(back[name], delta, atol=1e-6), name + assert set(back) == {s.removeprefix("base_model.model.") for s in _deltas(ours)} + + +def test_round_trip_survives_the_interleaved_layout() -> None: + ours = _ours() + exported = export_reference(ours, target="minimaxai") + back = _deltas(import_reference(exported, source="minimaxai")) + for stem, delta in _deltas(ours).items(): + assert torch.allclose(back[stem.removeprefix("base_model.model.")], delta, atol=1e-6) + + +def test_the_two_layouts_are_not_the_same_bytes() -> None: + """If they were, the layout argument would be decoration and a wrong guess would be harmless.""" + key = "diffusion_model.blocks.0.attn.qkv_proj.lora_B.weight" + ours = _ours() + assert not torch.equal( + export_reference(ours, target="comfy-org")[key], + export_reference(ours, target="minimaxai")[key], + ) + + +def test_reading_a_layout_the_wrong_way_round_corrupts_the_delta() -> None: + """The failure this defaults around: no error, just the wrong weights.""" + ours = _ours() + wrong = _deltas(import_reference(export_reference(ours, target="minimaxai"))) + right = _deltas(ours)["base_model.model.transformer_blocks.0.attn.to_q"] + assert not torch.allclose(wrong["transformer_blocks.0.attn.to_q"], right, atol=1e-6) + + +def test_gated_ffn_halves_are_exchanged_and_exchanged_back() -> None: + ours = _ours() + mine = _deltas(ours)["base_model.model.transformer_blocks.0.ff.net.0.proj"] + theirs = _deltas(export_reference(ours))["diffusion_model.blocks.0.mlp.fc1"] + half = FFN // 2 + assert torch.allclose(theirs[:half], mine[half:], atol=1e-6) + assert torch.allclose(theirs[half:], mine[:half], atol=1e-6) + + +def test_alpha_triples_with_the_rank_so_the_scale_is_unchanged() -> None: + ours = _ours() + for part in "qkv": + ours[f"base_model.model.transformer_blocks.0.attn.to_{part}.alpha"] = torch.tensor(2.0) + exported = export_reference(ours) + alpha = exported["diffusion_model.blocks.0.attn.qkv_proj.alpha"] + rank = exported["diffusion_model.blocks.0.attn.qkv_proj.lora_A.weight"].shape[0] + assert rank == 3 * RANK + assert float(alpha) / rank == pytest.approx(2.0 / RANK) + + +def test_the_effective_scale_survives_the_split_back() -> None: + """``alpha / rank`` is what the fuser multiplies by, so that ratio is the invariant.""" + ours = _ours() + for part in "qkv": + ours[f"base_model.model.transformer_blocks.0.attn.to_{part}.alpha"] = torch.tensor(2.0) + back = import_reference(export_reference(ours)) + alpha = float(back["transformer_blocks.0.attn.to_q.alpha"]) + rank = back["transformer_blocks.0.attn.to_q.lora_A.weight"].shape[0] + assert alpha / rank == pytest.approx(2.0 / RANK) + + +def test_a_third_party_rank_r_adapter_splits_to_rank_r() -> None: + """The common import: a fused rank-r adapter must not be inflated on the way in.""" + theirs = { + "diffusion_model.blocks.0.attn.qkv_proj.lora_down.weight": torch.randn(RANK, WIDTH), + "diffusion_model.blocks.0.attn.qkv_proj.lora_up.weight": torch.randn(3 * PART, RANK), + "diffusion_model.blocks.0.attn.qkv_proj.alpha": torch.tensor(2.0), + } + out = import_reference(theirs) + assert out["transformer_blocks.0.attn.to_q.lora_A.weight"].shape[0] == RANK + assert float(out["transformer_blocks.0.attn.to_q.alpha"]) == pytest.approx(2.0) + + +def test_comfy_lora_down_up_naming_is_read() -> None: + theirs = { + "diffusion_model.blocks.0.attn.qkv_proj.lora_down.weight": torch.randn(RANK, WIDTH), + "diffusion_model.blocks.0.attn.qkv_proj.lora_up.weight": torch.randn(3 * PART, RANK), + } + out = import_reference(theirs) + assert set(out) == { + f"transformer_blocks.0.attn.to_{p}.lora_{s}.weight" for p in "qkv" for s in "AB" + } + assert torch.equal( + out["transformer_blocks.0.attn.to_k.lora_B.weight"], + theirs["diffusion_model.blocks.0.attn.qkv_proj.lora_up.weight"][PART : 2 * PART], + ) + + +def test_the_three_split_parts_share_one_lora_a() -> None: + theirs = { + "diffusion_model.blocks.0.attn.qkv_proj.lora_down.weight": torch.randn(RANK, WIDTH), + "diffusion_model.blocks.0.attn.qkv_proj.lora_up.weight": torch.randn(3 * PART, RANK), + } + out = import_reference(theirs) + first = out["transformer_blocks.0.attn.to_q.lora_A.weight"] + for part in "kv": + assert torch.equal(out[f"transformer_blocks.0.attn.to_{part}.lora_A.weight"], first) + + +def test_exporting_a_partial_attention_is_refused() -> None: + ours = _ours() + for suffix in ("lora_A", "lora_B"): + del ours[f"base_model.model.transformer_blocks.0.attn.to_v.{suffix}.weight"] + with pytest.raises(ComponentError, match="2 of 3"): + export_reference(ours) + + +def test_an_adapter_for_another_model_is_refused() -> None: + with pytest.raises(ComponentError, match="different model"): + import_reference( + {"diffusion_model.double_blocks.0.img_attn.qkv.lora_down.weight": torch.zeros(4, 8)} + ) + + +def test_an_unknown_layout_names_the_ones_that_exist() -> None: + with pytest.raises(ComponentError, match="comfy-org"): + import_reference(_ours(), source="nonsense") + + +def test_adapt_passes_our_own_adapter_through_untouched() -> None: + """Old files, and every fresh non-H3 one, must not be rewritten on the way in.""" + ours = _ours() + assert adapt(ours) is ours + assert not is_reference(ours) + + +def test_adapt_translates_a_reference_keyed_adapter() -> None: + theirs = { + "diffusion_model.blocks.0.attn.qkv_proj.lora_down.weight": torch.randn(RANK, WIDTH), + "diffusion_model.blocks.0.attn.qkv_proj.lora_up.weight": torch.randn(3 * PART, RANK), + } + assert is_reference(theirs) + assert "transformer_blocks.0.attn.to_q.lora_A.weight" in adapt(theirs) + + +def test_a_feed_forward_only_adapter_is_still_recognised() -> None: + """Detection is on the block prefix: an adapter that never touched attention has no fused + tensor to give it away.""" + assert is_reference( + { + "diffusion_model.blocks.3.mlp.fc2.lora_down.weight": torch.randn(RANK, WIDTH), + "diffusion_model.blocks.3.mlp.fc2.lora_up.weight": torch.randn(FFN, RANK), + } + ) + + +def test_a_round_trip_through_save_and_load_keeps_the_delta() -> None: + """What actually ships: train here, write the portable file, read it back here.""" + ours = _ours() + back = _deltas(adapt(export_reference(ours))) + for stem, delta in _deltas(ours).items(): + assert torch.allclose(back[stem.removeprefix("base_model.model.")], delta, atol=1e-6) diff --git a/core/tests/test_minimaxh3_nodes.py b/core/tests/test_minimaxh3_nodes.py index fe273e9..92450bd 100644 --- a/core/tests/test_minimaxh3_nodes.py +++ b/core/tests/test_minimaxh3_nodes.py @@ -218,9 +218,11 @@ def test_the_seed_is_resolved_to_a_concrete_value() -> None: # --- recognising checkpoints ---------------------------------------------------------------------- -def _fake_checkpoint(path: Path, keys: dict[str, list[int]]) -> Path: +def _fake_checkpoint( + path: Path, keys: dict[str, list[int]], dtypes: dict[str, str] | None = None +) -> Path: header = { - name: {"dtype": "BF16", "shape": shape, "data_offsets": [0, 0]} + name: {"dtype": (dtypes or {}).get(name, "BF16"), "shape": shape, "data_offsets": [0, 0]} for name, shape in keys.items() } blob = json.dumps(header).encode() @@ -254,14 +256,58 @@ def test_another_architecture_is_not_offered(models_root: Path) -> None: assert not reqs.inspect_file(path).is_h3 -def test_the_pruned_build_is_rejected_with_a_reason(models_root: Path) -> None: +def test_a_pruned_bf16_build_is_accepted(models_root: Path) -> None: + """40.2 GB against 66.3 GB, and the loader reads its AdaLN table directly.""" path = _fake_checkpoint( - models_root / "diffusion_models" / "minimax_h3_fl2va_pruned_int8_convrot.safetensors", + models_root / "diffusion_models" / "minimax_h3_fl2va_pruned_bf16.safetensors", {**_H3_PROBE, "adaln_t_table": [1025, 8]}, ) candidate = reqs.inspect_file(path) + assert candidate.is_h3 and candidate.pruned and candidate.usable + assert candidate.reason == "" + + +def test_a_pruned_fp8_build_is_accepted(models_root: Path) -> None: + """21.0 GB. A scalar scale per weight and no rotation, so it dequantises exactly.""" + path = _fake_checkpoint( + models_root / "diffusion_models" / "minimax_h3_fl2va_pruned_fp8_scaled.safetensors", + { + **_H3_PROBE, + "adaln_t_table": [1025, 8], + "blocks.0.attn.qkv_proj.comfy_quant": [27], + "blocks.0.attn.qkv_proj.weight_scale": [], + }, + dtypes={"blocks.0.attn.qkv_proj.weight": "F8_E4M3"}, + ) + candidate = reqs.inspect_file(path) + assert candidate.usable and candidate.quantisation == "float8_e4m3fn" + # A reason on a loadable file reads as a refusal to anything that shows it. + assert candidate.reason == "" + + +def test_an_int8_convrot_build_is_still_rejected(models_root: Path) -> None: + """The rotation is the part we cannot invert, and inverting it wrongly still renders.""" + path = _fake_checkpoint( + models_root / "diffusion_models" / "minimax_h3_fl2va_pruned_int8_convrot.safetensors", + { + **_H3_PROBE, + "adaln_t_table": [1025, 8], + "blocks.0.attn.qkv_proj.comfy_quant": [70], + }, + dtypes={"blocks.0.attn.qkv_proj.weight": "I8"}, + ) + candidate = reqs.inspect_file(path) assert candidate.is_h3 and not candidate.usable - assert "rank-8 lookup table" in candidate.reason + assert "convrot" in candidate.reason + + +def test_an_unrecognised_quantisation_is_refused_rather_than_guessed(models_root: Path) -> None: + """A comfy_quant format we do not know would render a plausible wrong video, not raise.""" + path = _fake_checkpoint( + models_root / "diffusion_models" / "minimax_h3_fl2va_something_new.safetensors", + {**_H3_PROBE, "blocks.0.attn.qkv_proj.comfy_quant": [40]}, + ) + assert not reqs.inspect_file(path).usable def test_the_comfy_int8_build_is_rejected_with_a_reason(models_root: Path) -> None: @@ -282,8 +328,11 @@ def test_the_picker_offers_the_usable_file_and_explains_the_rest(models_root: Pa ) provider = MiniMaxH3Provider("fl2va") - assert provider.catalog_options("diffusion_models") == ["good.safetensors"] - assert [r["file"] for r in provider.rejected()] == ["pruned.safetensors"] + assert provider.catalog_options("diffusion_models") == [ + "good.safetensors", + "pruned.safetensors", + ] + assert provider.rejected() == [] assert provider.catalog_options("loras") is None # not ours to filter @@ -319,7 +368,7 @@ def test_a_picked_transformer_is_what_gets_sized(models_root: Path) -> None: assert reqs.footprint_bytes("fl2va", factorised=False)["diffusion_bytes"] == 0 raw = reqs.footprint_bytes("fl2va", factorised=False, transformer=picked)["diffusion_bytes"] - assert raw == picked.stat().st_size > 0 + assert raw == reqs.resident_bytes(picked) > 0 def test_the_factorised_share_comes_off_whichever_file_was_picked(models_root: Path) -> None: @@ -380,3 +429,101 @@ def test_each_blockset_gets_its_denoiser_under_the_name_it_declares() -> None: assert _denoiser_name(MiniMaxH3Blocks()) == "transformer" assert _denoiser_name(MiniMaxH3Ref2VABlocks()) == "transformer_ref" + + +def test_a_pruned_build_is_sized_by_what_it_becomes_not_its_file_size(models_root: Path) -> None: + """Under-sizing is the dangerous direction: the fit ladder would promise a machine that then + dies to a host-RAM OOM kill instead of raising. A pruned build has already lost its AdaLN + branch, so taking the usual 39 percent off it a second time under-counts by that much again.""" + path = _fake_checkpoint( + models_root / "diffusion_models" / reqs.FL2VA_FILE, + {**_H3_PROBE, "adaln_t_table": [1025, 8]}, + ) + resident = reqs.resident_bytes(path) + assert resident == (21504 * 5376 + 1025 * 8) * 2 + sizes = reqs.footprint_bytes("fl2va", transformer=path) + assert sizes["diffusion_bytes"] == resident + + +def test_an_unpruned_build_still_has_the_adaln_share_taken_off(models_root: Path) -> None: + path = _fake_checkpoint(models_root / "diffusion_models" / reqs.FL2VA_FILE, _H3_PROBE) + sizes = reqs.footprint_bytes("fl2va", transformer=path) + assert sizes["diffusion_bytes"] == int(reqs.resident_bytes(path) * (1 - reqs.ADALN_SHARE)) + + +def test_an_fp8_build_is_sized_at_its_dequantised_weight(models_root: Path) -> None: + """The file is half the size it will occupy, and sizing from disk would halve the estimate.""" + path = _fake_checkpoint( + models_root / "diffusion_models" / reqs.FL2VA_FILE, + {**_H3_PROBE, "adaln_t_table": [1025, 8]}, + dtypes={"blocks.0.attn.qkv_proj.weight": "F8_E4M3"}, + ) + expected = (21504 * 5376 + 1025 * 8) * 2 + assert reqs.footprint_bytes("fl2va", transformer=path)["diffusion_bytes"] == expected + + +def test_a_pruned_build_is_not_factorised_again(models_root: Path, monkeypatch) -> None: # type: ignore[no-untyped-def] + """Re-running the transform would multiply a rank-8 projection by a full-width basis. The load + fails long before that, at the first shape mismatch, so this asserts the flag rather than the + crash.""" + pytest.importorskip("torch") + from inline_core.models.minimaxh3 import pipeline as pl + + path = _fake_checkpoint( + models_root / "diffusion_models" / reqs.FL2VA_FILE, + {**_H3_PROBE, "adaln_t_table": [1025, 8]}, + ) + for folder in ("MiniMax-H3-text-encoder", "MiniMax-H3-processor"): + (models_root / "text_encoders" / folder).mkdir(parents=True, exist_ok=True) + for name in (reqs.VIDEO_VAE_FILE, reqs.AUDIO_VAE_FILE): + (models_root / "vae").mkdir(parents=True, exist_ok=True) + (models_root / "vae" / name).write_bytes(b"") + seen: dict[str, object] = {} + + def stop(*_a: object, **kw: object) -> None: + seen.update(kw) + raise RuntimeError("stop here") + + monkeypatch.setattr(pl.reqs, "footprint_bytes", stop) + with pytest.raises(RuntimeError, match="stop here"): + pl.load_pipeline(_NullPolicy(), params={"model": path.name}, partition="fl2va") + assert seen["factorised"] is False + + +class _NullPolicy: + """Enough policy for the resolve-and-size prologue, which is all this reaches.""" + + def set_footprint(self, *_a: object) -> None: ... + def fit_estimate(self) -> None: return None + + +def test_the_fp8_build_is_offered_as_an_optional_download(models_root: Path) -> None: + """A third the download for the same model, so it belongs in the popup. Optional, because the + trainer cannot use it and the bf16 file stays the one a full install needs.""" + entries = {c.id: c for c in reqs.components("fl2va")} + fp8 = entries["h3-fl2va-fp8"] + assert fp8.optional and fp8.filename == reqs.FL2VA_FP8_FILE + assert "generation only" in fp8.label + assert not entries["h3-fl2va"].optional + + +def test_training_refuses_a_pruned_build_by_name(models_root: Path) -> None: + """It would otherwise fail reading a timestep tensor the file does not contain, which reads + like a corrupt download rather than the wrong build.""" + pytest.importorskip("torch") + from inline_core.training import h3 as train_h3 + + path = _fake_checkpoint( + models_root / "diffusion_models" / reqs.FL2VA_FP8_FILE, + {**_H3_PROBE, "adaln_t_table": [1025, 8]}, + ) + with pytest.raises(RuntimeError, match="pruned MiniMax H3 build"): + train_h3._refuse_pruned(path) + + +def test_training_accepts_the_full_build(models_root: Path) -> None: + pytest.importorskip("torch") + from inline_core.training import h3 as train_h3 + + path = _fake_checkpoint(models_root / "diffusion_models" / reqs.FL2VA_FILE, _H3_PROBE) + train_h3._refuse_pruned(path) # does not raise diff --git a/core/tests/test_webui_install.py b/core/tests/test_webui_install.py index ab8c7b5..519d315 100644 --- a/core/tests/test_webui_install.py +++ b/core/tests/test_webui_install.py @@ -110,6 +110,9 @@ def sandbox(tmp_path: Path) -> Sandbox: stubs = tmp_path / "stubs" stubs.mkdir() + # Shadows the real driver so the suite means the same thing on a GPU box as on a bare one. + # Without it, /usr/bin/nvidia-smi leaks in and every no-GPU case silently tests the opposite. + _stub(stubs / "nvidia-smi", "exit 1\n") uv_log = tmp_path / "uv.log" _stub(stubs / "uv", f'printf "%s\\n" "$*" >> "{uv_log}"\nexit 0\n') diff --git a/core/uv.lock b/core/uv.lock index 0099078..0e35cbb 100644 --- a/core/uv.lock +++ b/core/uv.lock @@ -601,7 +601,7 @@ wheels = [ [[package]] name = "inline-core" -version = "1.2.63" +version = "1.2.66" source = { editable = "." } dependencies = [ { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, @@ -612,6 +612,7 @@ dependencies = [ [package.optional-dependencies] all = [ { name = "accelerate" }, + { name = "av" }, { name = "bitsandbytes", marker = "sys_platform != 'darwin'" }, { name = "controlnet-aux" }, { name = "diffusers" }, @@ -678,6 +679,7 @@ training = [ requires-dist = [ { name = "accelerate", marker = "extra == 'all'", specifier = ">=0.30" }, { name = "accelerate", marker = "extra == 'runtime'", specifier = ">=0.30" }, + { name = "av", marker = "extra == 'all'", specifier = ">=12" }, { name = "av", marker = "extra == 'runtime'", specifier = ">=12" }, { name = "bitsandbytes", marker = "sys_platform != 'darwin' and extra == 'all'", specifier = ">=0.43" }, { name = "bitsandbytes", marker = "sys_platform != 'darwin' and extra == 'training'", specifier = ">=0.43" }, @@ -690,8 +692,8 @@ requires-dist = [ { name = "fastapi", marker = "extra == 'all'", specifier = ">=0.110" }, { name = "fastapi", marker = "extra == 'server'", specifier = ">=0.110" }, { name = "httpx", marker = "extra == 'dev'", specifier = ">=0.27" }, - { name = "huggingface-hub", marker = "extra == 'all'", specifier = ">=0.23" }, - { name = "huggingface-hub", marker = "extra == 'runtime'", specifier = ">=0.23" }, + { name = "huggingface-hub", marker = "extra == 'all'", specifier = ">=0.32" }, + { name = "huggingface-hub", marker = "extra == 'runtime'", specifier = ">=0.32" }, { name = "imageio-ffmpeg", marker = "extra == 'all'", specifier = ">=0.4" }, { name = "imageio-ffmpeg", marker = "extra == 'server'", specifier = ">=0.4" }, { name = "numpy", specifier = ">=1.26" }, From ac4a6213e00a6e4185d6d65941ac39e35620b2e0 Mon Sep 17 00:00:00 2001 From: ashish-aesthisia Date: Tue, 11 Aug 2026 09:45:26 +0000 Subject: [PATCH 2/4] fix(minimax-h3): restore the run queue after a refresh, drop the fp32 AdaLN promotion --- TODO | 5 ++ core/src/inline_core/models/minimaxh3/load.py | 36 +++++--------- core/src/inline_core/studio/fal.py | 26 ++++++---- core/src/inline_core/studio/generation.py | 18 +++++++ core/src/inline_core/studio/handlers.py | 10 ++++ core/tests/test_minimaxh3_load.py | 47 +++++++++---------- core/tests/test_studio_fal.py | 16 +++++++ core/tests/test_studio_generation.py | 40 ++++++++++++++++ src/renderer/store/generationStore.ts | 22 +++++++++ .../views/Moodboard/MoodboardPanel.tsx | 4 ++ src/shared/ipc.ts | 8 ++++ src/shared/types.ts | 11 +++++ 12 files changed, 184 insertions(+), 59 deletions(-) diff --git a/TODO b/TODO index e69de29..d6c16b9 100644 --- a/TODO +++ b/TODO @@ -0,0 +1,5 @@ +- queue management +- past runs, cancel active runs +- models side panel +- templates +- keep runs active between project switch diff --git a/core/src/inline_core/models/minimaxh3/load.py b/core/src/inline_core/models/minimaxh3/load.py index cd53c25..306d1b8 100644 --- a/core/src/inline_core/models/minimaxh3/load.py +++ b/core/src/inline_core/models/minimaxh3/load.py @@ -13,7 +13,6 @@ import json import logging -import re from collections.abc import Iterator from pathlib import Path from typing import Any @@ -51,10 +50,6 @@ _SIDECAR_SUFFIXES = (".weight_scale", ".input_scale", ".comfy_quant") _SCALE_SUFFIX = ".weight_scale" -#: The rank-8 projections a pruned build ships. 75 MB across the model, so they stay in float32: -#: rounding them to bf16 would cost more accuracy than the whole factorisation does. -_ADALN_LINEAR = re.compile(r"(?:adaln_proj|norm_out)\.linear\.(?:weight|bias)$") - #: Source config name -> the vendored port's constructor argument. The port's defaults already match #: the released checkpoints, but a future build may not, so the file wins over the default. _CONFIG_MAP = { @@ -210,13 +205,6 @@ def load_transformer( return model -#: A fused delta smaller than this fraction of one quantisation step is rounded away rather than -#: applied. Measured on a real adapter: 0.008 of a step, which flipped 0.92% of int8 codes and -#: delivered none of its intended output change. Adapter strength itself is not the test - published -#: LoRAs that work well measure anywhere from 0.017% to 1.2% of the weight norm. -_LOST_TO_QUANTISATION = 0.25 - - def _fusing_shrink( plan: Any, fused: set[str], inner: Any, strength: list[float] | None = None ) -> Any: @@ -280,22 +268,21 @@ def _measure_strength( def _report_strength(strength: list[float], quantised: bool) -> None: - """Say what the adapter did, and warn only when quantisation is about to discard it.""" + """Report what the adapter actually did to the weights. + + Reported and not judged. A threshold on strength is a false-positive machine, since published + LoRAs that work well span 0.017% to 1.2% of the weight norm, and the obvious second theory is + wrong too: fusing into a quantised base preserves the delta along its intended direction at + about 100%, measured, so a small step ratio is not evidence the adapter will be invisible. + """ if len(strength) < 2: return ratio, per_step = strength logger.info( - "MiniMax H3: the LoRA moves that layer by %.3f%% of its weight norm (%.2f of an int8 step)", - ratio * 100, per_step, + "MiniMax H3: the LoRA moves that layer by %.3f%% of its weight norm (%.2f of an int8 " + "step%s)", + ratio * 100, per_step, "" if quantised else ", this base is not quantised", ) - if quantised and per_step < _LOST_TO_QUANTISATION: - logger.warning( - "MiniMax H3: this base is quantised and the LoRA changes each weight by only %.2f of " - "one int8 step, so rounding discards nearly all of it and the render will look " - "unadapted. This is not a weak LoRA, it is the quantisation. Raising the LoRA strength " - "does not fix it. A card that holds the base unquantised does.", - per_step, - ) def _finish_fuse(model: Any, plan: Any, fused: set[str]) -> None: @@ -371,8 +358,7 @@ def _stream_into( if shrink is not None and pending is not None and block != pending: shrink(model, pending) pending = block if shrink is not None else None - placed = torch.float32 if pruned and _ADALN_LINEAR.search(target) else dtype - _assign(model, target, value.to(dtype=placed, device=device)) + _assign(model, target, value.to(dtype=dtype, device=device)) filled.add(target) if shrink is not None and pending is not None: shrink(model, pending) diff --git a/core/src/inline_core/studio/fal.py b/core/src/inline_core/studio/fal.py index 42f02ec..ecce0a5 100644 --- a/core/src/inline_core/studio/fal.py +++ b/core/src/inline_core/studio/fal.py @@ -268,6 +268,8 @@ def __init__(self, store: Any, events: Any) -> None: self._store = store self._events = events self._active: dict[str, bool] = {} + #: Last progress per frame, so a reloaded page can rebuild its queue. + self._last: dict[str, tuple[float, str | None]] = {} def run(self, frame_id: str, request: dict[str, Any]) -> None: key = self._store.fal_key() @@ -283,6 +285,20 @@ def run(self, frame_id: str, request: dict[str, Any]) -> None: def cancel(self, frame_id: str | None = None) -> None: for fid in [frame_id] if frame_id else list(self._active.keys()): self._active.pop(fid, None) + self._last.pop(fid, None) + + def active(self) -> list[dict[str, Any]]: + """The fal runs still in flight, for a client that has lost its own copy of the queue.""" + from .generation import active_entry + + return [active_entry(f, self._last.get(f)) for f in self._active] + + def _progress(self, frame_id: str, fraction: float, status: str | None) -> None: + self._last[frame_id] = (fraction, status) + self._events.broadcast( + "events:generationProgress", + {"frameId": frame_id, "fraction": fraction, "status": status}, + ) async def _run(self, frame_id: str, request: dict[str, Any], key: str) -> None: import httpx @@ -292,10 +308,7 @@ async def _run(self, frame_id: str, request: dict[str, Any], key: str) -> None: output_kind = request.get("outputKind") or "image" headers = {"Authorization": f"Key {key}"} try: - self._events.broadcast( - "events:generationProgress", - {"frameId": frame_id, "fraction": 0.05, "status": "Queued"}, - ) + self._progress(frame_id, 0.05, "Queued") async with httpx.AsyncClient(timeout=600) as client: sub = await client.post(f"{_QUEUE_BASE}/{endpoint}", headers=headers, json=body) sub.raise_for_status() @@ -310,10 +323,7 @@ async def _run(self, frame_id: str, request: dict[str, Any], key: str) -> None: res.raise_for_status() status = res.json() fraction, label = _progress_from_status(status) - self._events.broadcast( - "events:generationProgress", - {"frameId": frame_id, "fraction": fraction, "status": label}, - ) + self._progress(frame_id, fraction, label) state = status.get("status") if state == "COMPLETED": break diff --git a/core/src/inline_core/studio/generation.py b/core/src/inline_core/studio/generation.py index 5214b07..876b68f 100644 --- a/core/src/inline_core/studio/generation.py +++ b/core/src/inline_core/studio/generation.py @@ -50,6 +50,9 @@ def __init__(self, store: Any, manager: Any, events: Any, registry: Any = None) self._events = events self._registry = registry self._active: dict[str, str] = {} # canvas item id -> run id + # Last progress per item, so a reloaded page can rebuild its queue. Progress is broadcast + # and forgotten otherwise, and a run mid-model-load emits nothing for minutes. + self._last: dict[str, tuple[float, str | None]] = {} def _is_list_port(self, node_type: str, port_id: str) -> bool: """Whether a port accepts several wires. Only the registry knows, and the canvas needs it to @@ -80,6 +83,7 @@ def cancel(self, item_id: str | None = None) -> None: ids = [item_id] if item_id else list(self._active.keys()) for iid in ids: run_id = self._active.pop(iid, None) + self._last.pop(iid, None) if run_id: self._manager.cancel(run_id) @@ -137,8 +141,14 @@ async def _drain(self, item_id: str, record: Any) -> None: finally: record.subscribers.discard(queue) self._active.pop(item_id, None) + self._last.pop(item_id, None) + + def active(self) -> list[dict[str, Any]]: + """The runs still in flight, for a client that has lost its own copy of the queue.""" + return [active_entry(i, self._last.get(i)) for i in self._active] def _progress(self, item_id: str, fraction: float, status: str | None) -> None: + self._last[item_id] = (fraction, status) self._events.broadcast( "events:generationProgress", {"frameId": item_id, "fraction": fraction, "status": status}, @@ -205,3 +215,11 @@ def _is_primary_output(self, item_id: str, kind: str) -> bool: return True output_kind = self._registry.get(node_type).output_kind return output_kind is None or _kind_str(output_kind) == kind + + +def active_entry(frame_id: str, last: tuple[float, str | None] | None) -> dict[str, Any]: + """One in-flight run, shaped like a progress event so a client reuses the same reducer. + + Public because the fal runner reports the same shape into the same merged queue.""" + fraction, status = last if last is not None else (None, None) + return {"frameId": frame_id, "fraction": fraction, "status": status} diff --git a/core/src/inline_core/studio/handlers.py b/core/src/inline_core/studio/handlers.py index 0be876d..77db498 100644 --- a/core/src/inline_core/studio/handlers.py +++ b/core/src/inline_core/studio/handlers.py @@ -257,6 +257,16 @@ def cancel_generation(frame_id: str | None = None) -> None: fal_generation.cancel(frame_id) reg("generation:cancel", cancel_generation) + + def active_generations() -> list[dict[str, Any]]: + """What is still running, so a reloaded page rebuilds its queue instead of losing it.""" + out: list[dict[str, Any]] = [] + for source in (generation, fal_generation): + if source is not None: + out.extend(source.active()) + return out + + reg("generation:active", active_generations) reg("generation:resumePending", lambda: None) # --- LoRA training (dataset CRUD + the training run subprocess) ------------------------------ diff --git a/core/tests/test_minimaxh3_load.py b/core/tests/test_minimaxh3_load.py index 6e8b80e..af3a216 100644 --- a/core/tests/test_minimaxh3_load.py +++ b/core/tests/test_minimaxh3_load.py @@ -309,12 +309,15 @@ def test_a_pruned_build_replaces_the_timestep_path(tmp_path: Path) -> None: model = _tiny_model() state = {k: v.detach().clone() for k, v in model.state_dict().items()} path = _write(tmp_path / "h3_pruned2.safetensors", _pruned_source(model, state)) - loaded = load_transformer(path, dtype=torch.float32) + loaded = load_transformer(path, dtype=torch.bfloat16) assert isinstance(loaded.time_embedder, adaln.TableEmbedder) assert isinstance(loaded.transformer_blocks[0].adaln_proj, adaln.TabulatedModulation) assert isinstance(loaded.norm_out, adaln.TabulatedNormOut) - # The rank-8 projections stay in float32; rounding them costs more than the factorisation does. - assert loaded.transformer_blocks[0].adaln_proj.linear.weight.dtype is torch.float32 + # The compute dtype like every other weight, and not float32. Float32 here promoted the hidden + # states through the modulation multiply, which doubled activation memory and dropped rms_norm + # off its fused kernel for the whole denoise. + for module in (loaded.transformer_blocks[0].adaln_proj, loaded.norm_out): + assert module.linear.weight.dtype is torch.bfloat16 def test_an_unpruned_build_is_untouched_by_any_of_this(reference) -> None: # type: ignore[no-untyped-def] @@ -373,38 +376,30 @@ def test_the_fp8_plan_is_versioned_apart() -> None: # --- what a fused LoRA is worth once the base is quantised ------------------------------------ -def test_a_lora_lost_to_quantisation_is_warned_about(caplog) -> None: # type: ignore[no-untyped-def] - """The measured case: 0.008 of an int8 step reaches the weights, so rounding discards it.""" - import logging - - with caplog.at_level(logging.WARNING, logger="inline_core.minimaxh3"): - h3_load._report_strength([0.00038, 0.008], quantised=True) - assert "quantisation" in caplog.text - assert "not a weak LoRA" in caplog.text - - -def test_no_warning_when_the_base_is_not_quantised(caplog) -> None: # type: ignore[no-untyped-def] - """The same adapter applies fine in full precision, so warning there would be noise.""" +def test_the_fuse_reports_both_numbers(caplog) -> None: # type: ignore[no-untyped-def] + """Strength alone says nothing, so the step ratio is reported beside it.""" import logging - with caplog.at_level(logging.WARNING, logger="inline_core.minimaxh3"): - h3_load._report_strength([0.00038, 0.008], quantised=False) - assert caplog.text == "" + with caplog.at_level(logging.INFO, logger="inline_core.minimaxh3"): + h3_load._report_strength([0.00818, 0.15], quantised=True) + assert "0.818%" in caplog.text + assert "0.15 of an int8 step" in caplog.text -def test_no_warning_when_the_delta_clears_a_quantisation_step(caplog) -> None: # type: ignore[no-untyped-def] +def test_a_small_step_ratio_is_not_warned_about(caplog) -> None: # type: ignore[no-untyped-def] + """It was, and the warning was wrong. Fusing into a quantised base preserves the delta along + its intended direction at about 100%, measured, so a small ratio is not evidence of anything. + Published LoRAs that work well span 0.017% to 1.2% of the weight norm.""" import logging with caplog.at_level(logging.WARNING, logger="inline_core.minimaxh3"): - h3_load._report_strength([0.02, 1.5], quantised=True) + h3_load._report_strength([0.00038, 0.008], quantised=True) assert caplog.text == "" -def test_adapter_strength_alone_does_not_trigger_the_warning(caplog) -> None: # type: ignore[no-untyped-def] - """Published LoRAs that work well measure 0.017% to 1.2% of the weight norm, so a threshold on - strength alone would fire on almost all of them.""" +def test_an_unquantised_base_says_so(caplog) -> None: # type: ignore[no-untyped-def] import logging - with caplog.at_level(logging.WARNING, logger="inline_core.minimaxh3"): - h3_load._report_strength([0.00017, 0.9], quantised=True) - assert caplog.text == "" + with caplog.at_level(logging.INFO, logger="inline_core.minimaxh3"): + h3_load._report_strength([0.02, 1.5], quantised=False) + assert "not quantised" in caplog.text diff --git a/core/tests/test_studio_fal.py b/core/tests/test_studio_fal.py index 46bfd2e..2a41d07 100644 --- a/core/tests/test_studio_fal.py +++ b/core/tests/test_studio_fal.py @@ -234,3 +234,19 @@ def test_the_same_asset_can_feed_two_ports(tmp_path) -> None: assert first["id"] != second["id"] # Re-adding the same (asset, handle) pair still dedups. assert fr.add_input(conn, gen["frameId"], "a1", "image")["id"] == first["id"] + + +def test_fal_active_reports_runs_in_flight() -> None: + """Same contract as the Core side, since the handler merges both into one queue.""" + + class _Events: + def broadcast(self, channel: str, payload: dict) -> None: ... + + gen = fal.FalGeneration(store=None, events=_Events()) + gen._active["frame-1"] = True + gen._progress("frame-1", 0.25, "Queued") + + assert gen.active() == [{"frameId": "frame-1", "fraction": 0.25, "status": "Queued"}] + + gen.cancel("frame-1") + assert gen.active() == [] diff --git a/core/tests/test_studio_generation.py b/core/tests/test_studio_generation.py index 45343de..cbb7eea 100644 --- a/core/tests/test_studio_generation.py +++ b/core/tests/test_studio_generation.py @@ -274,3 +274,43 @@ def test_single_take_image_node_is_unaffected_by_the_gate(tmp_path) -> None: output = mb.get_item(store.conn(), z["id"])["data"]["core"]["output"] assert output["kind"] == "image" and output["takeId"] == "tk1" + + +# --- surviving a page refresh -------------------------------------------------------------------- + + +class _Events: + def __init__(self) -> None: + self.sent: list[tuple[str, dict]] = [] + + def broadcast(self, channel: str, payload: dict) -> None: + self.sent.append((channel, payload)) + + +def test_active_reports_runs_in_flight_with_their_last_progress() -> None: + """A refresh throws away the client's queue while the run carries on. Without this the UI shows + an empty queue against a GPU that is still working.""" + gen = CoreGeneration(store=None, manager=None, events=_Events()) + gen._active["item-1"] = "run-1" + gen._progress("item-1", 0.4, "Sampling") + + assert gen.active() == [{"frameId": "item-1", "fraction": 0.4, "status": "Sampling"}] + + +def test_a_run_that_has_not_reported_yet_still_appears() -> None: + """H3 emits nothing for minutes while the base loads, which is exactly when someone refreshes, + so an entry with no progress yet has to come back rather than be omitted.""" + gen = CoreGeneration(store=None, manager=None, events=_Events()) + gen._active["item-2"] = "run-2" + + assert gen.active() == [{"frameId": "item-2", "fraction": None, "status": None}] + + +def test_a_finished_run_is_not_reported() -> None: + gen = CoreGeneration(store=None, manager=None, events=_Events()) + gen._active["item-3"] = "run-3" + gen._progress("item-3", 0.9, "Decoding") + gen._active.pop("item-3") + gen._last.pop("item-3") + + assert gen.active() == [] diff --git a/src/renderer/store/generationStore.ts b/src/renderer/store/generationStore.ts index c2e7753..19d7fae 100644 --- a/src/renderer/store/generationStore.ts +++ b/src/renderer/store/generationStore.ts @@ -34,6 +34,8 @@ interface GenerationState { cancel: (frameId?: string) => Promise /** Ask main to re-poll + finish any generations left in flight from a previous session. */ resumePending: () => Promise + /** Rebuild the queue from what Core still has running, after a page refresh threw ours away. */ + hydrateActive: () => Promise /** Persist a fal frame's param values (optimistically updates the frame store). */ setParams: (frameId: string, params: Record) => Promise /** Switch a fal frame to a different model (resets params + output kind). */ @@ -171,6 +173,26 @@ export const useGenerationStore = create((set) => ({ } }, + hydrateActive: async () => { + try { + const res = await studio().generation.active() + if (!res.ok) return + set((s) => { + const busy = { ...s.busyByFrame } + const progress = { ...s.progressByFrame } + const status = { ...s.statusByFrame } + for (const run of res.value) { + busy[run.frameId] = true + progress[run.frameId] = run.fraction + status[run.frameId] = run.status + } + return { busyByFrame: busy, progressByFrame: progress, statusByFrame: status } + }) + } catch { + // A backend that predates the channel simply has no queue to restore. + } + }, + setParams: async (frameId, params) => { // Optimistic: reflect the edit on the frame immediately so the node re-renders. useFrameStore.setState((s) => ({ diff --git a/src/renderer/views/Moodboard/MoodboardPanel.tsx b/src/renderer/views/Moodboard/MoodboardPanel.tsx index 106e6a7..cafb7ac 100644 --- a/src/renderer/views/Moodboard/MoodboardPanel.tsx +++ b/src/renderer/views/Moodboard/MoodboardPanel.tsx @@ -418,6 +418,10 @@ function Board(): React.JSX.Element { // Finish any generations that were still running when the app last closed (they keep // running server-side); their progress/completion arrives through the events below. void gen.resumePending() + // A page refresh throws away this tab's copy of the queue while Core keeps working, and a run + // inside a long model load emits nothing for minutes, so waiting for the next event is not + // enough. Ask Core what is still in flight. + void gen.hydrateActive() const unsubs = [ studio().events.onGenerationProgress((e) => { gen.setBusy(e.frameId, true) diff --git a/src/shared/ipc.ts b/src/shared/ipc.ts index 2764483..82c06d3 100644 --- a/src/shared/ipc.ts +++ b/src/shared/ipc.ts @@ -7,6 +7,7 @@ * preload. The renderer imports this type; the main process implements it. */ import type { + ActiveGeneration, Project, RecentProject, Asset, @@ -123,6 +124,8 @@ export const IpcChannels = { runWorkflow: 'generation:runWorkflow', /** Abort the in-flight generation run (optionally just one frame's). */ cancel: 'generation:cancel', + /** What is still running, so a reloaded page can rebuild its queue. */ + active: 'generation:active', /** Re-poll + finish any runs that were in flight when the app last closed. */ resumePending: 'generation:resumePending', }, @@ -431,6 +434,11 @@ export interface InlineStudioApi { runWorkflow(itemId: string): Promise> /** Abort the in-flight run - a specific frame's, or all when no id is given. */ cancel(frameId?: string): Promise> + /** + * The runs Core still has in flight. A page refresh throws away the renderer's copy of the + * queue while the backend keeps working, so this is what rebuilds it on mount. + */ + active(): Promise> /** Re-poll + finish any generations that were in flight when the app last closed. */ resumePending(): Promise> } diff --git a/src/shared/types.ts b/src/shared/types.ts index e4bd8d7..8d5e87e 100644 --- a/src/shared/types.ts +++ b/src/shared/types.ts @@ -79,6 +79,17 @@ export interface FrameInput { } /** Every ComfyUI render of a frame becomes an immutable Take. */ +/** + * A generation Core still has in flight. Shaped like the progress event so the renderer can feed + * it through the same reducer when rebuilding its queue after a page refresh. + */ +export interface ActiveGeneration { + frameId: string + /** 0..1, or null when the run has not reported yet (a long model load reports nothing). */ + fraction: number | null + status?: string +} + export interface Take { id: string frameId: string From de4d8164c9307a6a4ea86ac8e12335b7f2360590 Mon Sep 17 00:00:00 2001 From: ashish-aesthisia Date: Tue, 11 Aug 2026 09:48:15 +0000 Subject: [PATCH 3/4] update TODO --- TODO | 1 + 1 file changed, 1 insertion(+) diff --git a/TODO b/TODO index d6c16b9..1889ed3 100644 --- a/TODO +++ b/TODO @@ -3,3 +3,4 @@ - models side panel - templates - keep runs active between project switch +- provide gen node in trainer tab for quick test with different lora strenths From f0963d01efaff3b92ff48148170396dfb9f52a27 Mon Sep 17 00:00:00 2001 From: ashish-aesthisia Date: Tue, 11 Aug 2026 09:57:52 +0000 Subject: [PATCH 4/4] fix(minimax-h3): report denoise progress, restore the run queue after a refresh --- TODO | 1 + .../inline_core/models/minimaxh3/runner.py | 14 ++++ .../inline_core/models/pipeline_runtime.py | 71 +++++++++++++++++++ core/tests/test_minimaxh3_nodes.py | 60 ++++++++++++++++ 4 files changed, 146 insertions(+) diff --git a/TODO b/TODO index 1889ed3..c03ecb0 100644 --- a/TODO +++ b/TODO @@ -4,3 +4,4 @@ - templates - keep runs active between project switch - provide gen node in trainer tab for quick test with different lora strenths +- Run Graph drop down, download workflow json \ No newline at end of file diff --git a/core/src/inline_core/models/minimaxh3/runner.py b/core/src/inline_core/models/minimaxh3/runner.py index 69dbae9..2e5c136 100644 --- a/core/src/inline_core/models/minimaxh3/runner.py +++ b/core/src/inline_core/models/minimaxh3/runner.py @@ -325,6 +325,20 @@ def run(self, node: Node, inputs: dict[str, list[Any]], ctx: ExecutionContext) - call = call_kwargs(request, self._variant, inputs) call["generator"] = torch.Generator(device="cpu").manual_seed(request.seed) + + def on_step(done: int, total: int) -> None: + if ctx.cancel.cancelled: + raise CancelledError("Run cancelled.") + ctx.emitter.emit( + rt.progress_event( + ctx, node, Phase.SAMPLE, done / max(total, 1), + step=done, step_count=total, status=f"Step {done}/{total}", + ) + ) + + # Without this the run reports nothing between loading and saving, and a denoise that takes + # minutes reads as a stuck model load. + rt.attach_step_progress(pipe, on_step) started = time.perf_counter() try: state = render_staged(pipe, self._policy.placement('denoiser').device, **call) diff --git a/core/src/inline_core/models/pipeline_runtime.py b/core/src/inline_core/models/pipeline_runtime.py index fb5efae..f1ba07e 100644 --- a/core/src/inline_core/models/pipeline_runtime.py +++ b/core/src/inline_core/models/pipeline_runtime.py @@ -649,3 +649,74 @@ def try_call(fn: Any, *args: Any, **kwargs: Any) -> None: "try_call", "wont_fit_message", ] + + +class _StepReporter: + """Stands in for the progress bar a modular denoise loop drives, and reports each step onward. + + A modular blockset has no ``callback_on_step_end``: its loop calls ``self.progress_bar`` and + then ``.update()`` per step. That bar is therefore the only per-step hook, and using it keeps + ``vendor/`` verbatim. The real bar is still driven, so the terminal output is unchanged. + """ + + def __init__(self, inner: Any, total: int, on_step: Any) -> None: + self._inner = inner + self._total = total + self._on_step = on_step + self._done = 0 + + def __enter__(self) -> _StepReporter: + self._inner.__enter__() + return self + + def __exit__(self, *exc: Any) -> Any: + return self._inner.__exit__(*exc) + + def update(self, n: int = 1) -> None: + self._inner.update(n) + self._done += n + self._on_step(self._done, self._total) + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + +def attach_step_progress(pipe: Any, on_step: Any) -> bool: + """Report every denoising step of ``pipe``. Returns whether a loop was found to hook. + + Without this a long denoise emits nothing between "loading" and "saving", and the UI shows the + load phase for the whole render, which reads as a hang. + """ + found = False + for blocks in _blocksets(pipe): + loop = _denoise_loop(blocks) + if loop is None: + continue + found = True + original = loop.progress_bar + + def progress_bar( + total: Any = None, _original: Any = original, **kw: Any + ) -> _StepReporter: + return _StepReporter(_original(total=total, **kw), int(total or 0), on_step) + + loop.progress_bar = progress_bar + return found + + +def _blocksets(pipe: Any) -> list[Any]: + """Every blockset that might run: a staged pipeline keeps the denoise in its second half.""" + phases = getattr(pipe, "_inline_phases", None) + targets = list(phases) if phases else [pipe] + return [t.blocks for t in targets if getattr(t, "blocks", None) is not None] + + +def _denoise_loop(blocks: Any) -> Any: + """The loop block, found by the ``loop_step`` that defines one, not by name.""" + if hasattr(blocks, "loop_step"): + return blocks + for child in getattr(blocks, "sub_blocks", {}).values(): + found = _denoise_loop(child) + if found is not None: + return found + return None diff --git a/core/tests/test_minimaxh3_nodes.py b/core/tests/test_minimaxh3_nodes.py index 92450bd..98d606d 100644 --- a/core/tests/test_minimaxh3_nodes.py +++ b/core/tests/test_minimaxh3_nodes.py @@ -527,3 +527,63 @@ def test_training_accepts_the_full_build(models_root: Path) -> None: path = _fake_checkpoint(models_root / "diffusion_models" / reqs.FL2VA_FILE, _H3_PROBE) train_h3._refuse_pruned(path) # does not raise + + +# --- per-step progress --------------------------------------------------------------------------- + + +class _FakeBar: + def __init__(self) -> None: + self.updates = 0 + + def __enter__(self) -> _FakeBar: + return self + + def __exit__(self, *_exc: object) -> bool: + return False + + def update(self, n: int = 1) -> None: + self.updates += n + + +class _FakeLoop: + """Stands in for the vendored denoise block: it has a loop_step and drives a progress bar.""" + + def __init__(self) -> None: + self.bar = _FakeBar() + self.progress_bar = lambda total=None, **_kw: self.bar + + def loop_step(self) -> None: ... + + +def test_step_progress_hooks_the_denoise_loop() -> None: + """The modular loop has no callback_on_step_end, so the progress bar is the only per-step hook. + Without it a long denoise emits nothing and the UI shows 'loading model' for the whole render.""" + from inline_core.models import pipeline_runtime as rt + + loop = _FakeLoop() + + class _Blocks: + sub_blocks = {"denoise": loop} + + class _Pipe: + blocks = _Blocks() + + seen: list[tuple[int, int]] = [] + assert rt.attach_step_progress(_Pipe(), lambda done, total: seen.append((done, total))) + + with loop.progress_bar(total=3) as bar: + for _ in range(3): + bar.update() + + assert seen == [(1, 3), (2, 3), (3, 3)] + assert loop.bar.updates == 3 # the real bar is still driven + + +def test_step_progress_reports_when_there_is_no_loop_to_hook() -> None: + class _Pipe: + blocks = None + + from inline_core.models import pipeline_runtime as rt + + assert rt.attach_step_progress(_Pipe(), lambda *_a: None) is False