diff --git a/README.md b/README.md index f8778fe..cec42a6 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,8 @@ A free, open-source app for AI generation where your characters stay the same. B | Model | Train | Generate | Trains on a 16GB card | | ------------------------------------------------------------------------ | ----- | -------- | ----------------------- | -| [FLUX.2](https://bfl.ai/blog/flux-2) (klein Base 4B) | yes | yes | yes, ~8.6GB | +| [FLUX.2](https://bfl.ai/blog/flux-2) (klein Base 4B / 9B) | yes | yes | yes, ~8.6GB | +| [FLUX.1](https://bfl.ai/blog/flux-1) dev (4-bit) | yes | yes | yes, ~10.4GB | | [Krea 2](https://www.krea.ai/) (RAW, 4-bit) | yes | yes | yes, ~11.9GB | | Z-Image Turbo | yes | yes | yes, ~13.4GB | | [MiniMax H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) (video, sound) | yes | yes | yes but slowly, ~12.7GB | @@ -241,6 +242,11 @@ trains on stills for look and style, or on clips to learn motion as well, and on both. LTX-2.5 trains on clips, and can also learn a transform between a reference clip and a target one, which upstream calls an IC-LoRA. +The image models train on stills. **FLUX.1** trains on dev - it is guidance-distilled rather than +step-distilled, so it is its own training base and needs no adapter, and 4-bit puts it on a 12GB +card at 512px. **FLUX.2** trains on a klein Base build, 4B or 9B, and the adapter still loads on the +distilled build afterwards. **Krea 2** and **Z-Image** are covered in TRAINING.md. + Already installed with `--extra all`? The trainer is ready. Otherwise `./webui.sh --install --extra training`. @@ -297,16 +303,22 @@ core/models/ 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 + flux-2-klein-base-9b.safetensors <- FLUX.2 9B base, for training (gated) + flux1-dev.safetensors <- FLUX.1 dev, generate and train minimax_h3_fl2va_bf16.safetensors <- H3 text, image, first/last frame minimax_h3_ref2va_bf16.safetensors <- H3 reference node ltx-2.5-22b-distilled-transformer-bf16.safetensors <- LTX fast mode ltx-2.5-22b-dev-transformer-bf16.safetensors <- LTX quality mode, and training text_encoders/ qwen3vl_4b_bf16.safetensors <- Krea 2 qwen_3_4b.safetensors <- FLUX.2 klein 4B, shared with Z-Image + qwen_3_8b.safetensors <- FLUX.2 klein 9B + t5xxl_fp16.safetensors <- FLUX.1 sequence encoder + clip_l.safetensors <- FLUX.1 pooled encoder MiniMax-H3-text-encoder/ <- Qwen3-VL-32B, a folder MiniMax-H3-processor/ gemma4-12b-with-proj-ltx-2.5-bf16.safetensors <- LTX vae/ qwen_image_vae_diffusers.safetensors + ae.safetensors <- FLUX.1, the same file Z-Image uses flux2-vae.safetensors minimax_h3_video_vae_fp16.safetensors minimax_h3_audio_vae_fp32.safetensors @@ -409,10 +421,11 @@ cost nothing to run. Hosted models are billed by the provider. **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. -**Can I train a LoRA locally?** Yes, for all five local models, on your own GPU. See +**Can I train a LoRA locally?** Yes, for all six local models, on your own GPU. See [TRAINING.md](TRAINING.md). -**What models can I run?** Locally: Z-Image Turbo, FLUX.2, Krea 2, MiniMax H3 and LTX-2.5. Hosted: +**What models can I run?** Locally: Z-Image Turbo, FLUX.1, FLUX.2, Krea 2, MiniMax H3 and +LTX-2.5. Hosted: the fal catalogue, with more providers to follow. Adding a new local model is a Core change, not a UI release. @@ -436,7 +449,7 @@ on [Discord](https://discord.gg/cSUS88VdY9), or try the [creator task](task.md). - [**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. +- [**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. And for [FLUX.1](https://bfl.ai/blog/flux-1) dev, which is non-commercial - a LoRA trained on it inherits that. - [**MiniMax**](https://huggingface.co/MiniMaxAI/MiniMax-H3) for MiniMax H3, under the MiniMax H3 Community License. - [**Lightricks**](https://huggingface.co/Lightricks/LTX-2.5) for LTX-2.5, under the LTX-2 Community License, and for the [paired dataset pipeline](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-trainer/docs/dataset-preparation.md) the control LoRA trainer follows. diff --git a/TRAINING.md b/TRAINING.md index 39d9f84..f40c472 100644 --- a/TRAINING.md +++ b/TRAINING.md @@ -39,16 +39,29 @@ A dataset's trigger word is prepended to every caption during training, so the m ## Architecture and base model modes -The Trainer's Adjust panel picks the **architecture** first (Z-Image, Krea 2, FLUX.2, MiniMax H3, or LTX-2.5), then a base within it. Training directly on a step-distilled checkpoint breaks the distillation down (turbo drift), so each architecture offers a way around that. +The Trainer's Adjust panel picks the **architecture** first (Z-Image, Krea 2, FLUX.1, FLUX.2, MiniMax H3, or LTX-2.5), then a base within it. Training directly on a step-distilled checkpoint breaks the distillation down (turbo drift), so each architecture offers a way around that. **Krea 2** avoids the problem outright, which is why it is the recommended path: - **Krea 2 RAW** trains on the undistilled base. Nothing to fuse, nothing to drift. Put `krea2_raw_bf16.safetensors` in `models/diffusion_models/`, train, then generate with the **Krea 2 Turbo** node - the LoRA carries over unchanged. - **Krea 2 Turbo + training adapter** exists for people who only hold Turbo. Put [ostris/krea2_turbo_training_adapter](https://huggingface.co/ostris/krea2_turbo_training_adapter) in `models/loras/`, or point `INLINE_KREA2_TRAIN_ADAPTER` at it. -**FLUX.2** works like Krea 2, with no adapter to download: +**FLUX.2** works like Krea 2, with no adapter to download. Both bases are undistilled builds, and the trainer refuses a distilled checkpoint rather than letting a run produce a bad adapter hours later: -- **FLUX.2 Base** is the only option, and the trainer refuses a distilled checkpoint rather than letting a run produce a bad adapter hours later. Put `flux-2-klein-base-4b.safetensors` in `models/diffusion_models/`, train, then generate with the distilled **klein 4B** checkpoint. The LoRA carries over unchanged. +- **klein Base 4B** is the recommended path and the default. Apache-2.0 and ungated. Put `flux-2-klein-base-4b.safetensors` in `models/diffusion_models/`, train, then generate with the distilled **klein 4B** checkpoint. The LoRA carries over unchanged. +- **klein Base 9B** is the larger tier, and wants about 22GB of VRAM and 64GB of system RAM. Put `flux-2-klein-base-9b.safetensors` in `models/diffusion_models/` alongside the **Qwen3-8B** text encoder in `models/text_encoders/` - 9B does not use the 4B encoder, and the model popup fetches the right one once the Base setting is on 9B. Unlike the 4B line this checkpoint is gated and non-commercial, so you accept BFL's licence on Hugging Face first. + +Both bases can sit in `models/diffusion_models/` at once; the **Base** setting picks which one a run trains, rather than whichever sorts first. + +**FLUX.1** trains on **dev**, and needs no adapter either - but for a different reason to FLUX.2: + +- **dev is guidance-distilled, not step-distilled.** A LoRA trains _through_ guidance distillation by pinning the guidance embedding to 1 for every step, which is what every reference trainer does. There is nothing to de-distill, so there is no adapter and no base/turbo choice. +- **It conditions on two text encoders** and needs both: T5-XXL for the sequence, CLIP-L for the pooled vector. Put `flux1-dev.safetensors` in `models/diffusion_models/`, `t5xxl_fp16.safetensors` and `clip_l.safetensors` in `models/text_encoders/`, and `ae.safetensors` in `models/vae/` - that VAE is the file Z-Image already uses, so an existing install has it. +- **schnell is refused**, and by content rather than by name: it is the one FLUX.1 build with no guidance embedder at all, so the Trainer spots it whatever the file is called. It is step-distilled and collapses the way a distilled FLUX.2 does. +- **The Fill, Canny and Depth builds are refused too.** They take a mask or a stacked hint in extra input channels that a plain image dataset cannot fill, and the mismatch would otherwise surface as a shape error some way into the run. Kontext is left out for a different reason: it would train, but it learns an _edit_ between a pair of images, and a dataset of single images teaches it nothing it is used for. +- **A 24GB bf16 base beside a 10GB T5 encoder**, so most cards train it in 4-bit. Base precision on Auto picks that for you. The two encoders are loaded, used to cache latents and captions, then freed before the transformer loads, so the peak is one half or the other rather than the sum. +- **4-bit takes it from a 24GB card to a 12GB one, for about 10 percent a step.** Measured on an L40S: 24.9GB at 512 and 26.5GB at 1024 in bf16, against 10.4GB and 11.6GB in 4-bit, at 0.82s and 2.23s a step against bf16's 0.66s and 2.08s. A 500-step run at 512 is about six minutes. +- **dev's weights are non-commercial**, whichever mirror they come from, and a LoRA trained on them is a derivative that inherits it. The model popup says so on the row. **MiniMax H3** is the video model, and it trains on **still images**: @@ -204,11 +217,19 @@ well as its own `0001.ref.mp4`, so a downloaded set can be trained on directly. | FLUX.2 | Base (klein 4B) | 512 | **4-bit** | 8.6GB | not measured | | FLUX.2 | Base (klein 4B) | 1024 | bf16 | 9.9GB | not measured | | FLUX.2 | Base (klein 4B) | 1024 | **4-bit** | 9.9GB | not measured | +| FLUX.2 | Base (klein 9B) | 512 | bf16 | 19.2GB | not measured | +| FLUX.2 | Base (klein 9B) | 512 | **4-bit** | **15.7GB** | not measured | +| FLUX.2 | Base (klein 9B) | 1024 | bf16 | 20.8GB | not measured | +| FLUX.2 | Base (klein 9B) | 1024 | **4-bit** | **16.9GB** | not measured | | MiniMax H3 | FL2VA | 512 | **4-bit** | **20.6GB** | **12.7GB** | | MiniMax H3 | FL2VA | 768 | **4-bit** | **20.6GB** | not measured | | MiniMax H3 | FL2VA | 1024 | **4-bit** | **20.6GB** | not measured | | MiniMax H3 | FL2VA, clips | 512 | **4-bit** | **20.4GB** | not measured | | LTX-2.5 | dev, clips | 512 | bf16 | **42.0GB** | not supported | +| FLUX.1 | dev | 512 | bf16 | 24.9GB | not measured | +| FLUX.1 | dev | 512 | **4-bit** | **10.4GB** | not measured | +| FLUX.1 | dev | 1024 | bf16 | 26.5GB | not measured | +| FLUX.1 | dev | 1024 | **4-bit** | **11.6GB** | not measured | **H3 peaks lower on a 16GB card than it does on a 48GB one.** The run has three phases that never overlap, and on a big card the tallest is the caption pass: @@ -257,16 +278,25 @@ against the same 13. Precache is keyed and reused, so a second run over the same dataset, resolution and clip length skips most of it. Prefer long runs, and expect a short one to be dominated by setup. -**FLUX.2 is the cheapest of the three to train, and 4-bit does nothing for it.** Both precisions peak at the same number, because klein's base is 7.4GB against a 7.5GB Qwen3-4B text encoder, so the caption and latent caching pass at the start of the run costs more than the training does. Dropping the frozen base to 4-bit shrinks a part of the run that was never the high-water mark, and the step gets slower for nothing. Leave base precision on Auto for FLUX.2, which is what it already picks. The rows above are klein Base 4B, the only checkpoint the trainer accepts for this architecture. +**FLUX.1's 4-bit peak depends on where the base is quantised.** bitsandbytes quantises on the move +to CUDA, so streaming the base straight to the card materialises all 24GB of it first and the run +peaks there whatever precision you asked for. Loading to host RAM and quantising on the way +across drops the peak from 23.8GB to 6.2GB for the same 70 seconds, which is the difference between +FLUX.1 needing a 24GB card and running on a 12GB one. + +**FLUX.2 is the cheapest of the three to train, and 4-bit does nothing for it.** Both precisions peak at the same number, because klein's base is 7.4GB against a 7.5GB Qwen3-4B text encoder, so the caption and latent caching pass at the start of the run costs more than the training does. Dropping the frozen base to 4-bit shrinks a part of the run that was never the high-water mark, and the step gets slower for nothing. Leave base precision on Auto for FLUX.2, which is what it already picks. + +**klein Base 9B measures the same way, one tier up.** 19.2GB at 512 and 20.8GB at 1024 in bf16, against 15.7GB and 16.9GB in 4-bit, at roughly 0.53s and 1.95s a step. 4-bit drops the resident base from 18.2GB to 4.7GB but the peak only to ~16GB, because the ceiling is the Qwen3-8B caching pass rather than the transformer - the same reason 4-bit does little for 4B. So a 24GB card runs it in bf16 and 4-bit buys a card that is short of that, not a faster run. BFL's own floor is 22GB of VRAM and 64GB of system RAM, which matches. -Which card fits what (24GB and 32GB are interpolated, not measured, as are the FLUX.2 columns on 16GB: those peaks were measured on an L40S and leave room on a smaller card, but no 16GB run has been done): +Which card fits what. The 16GB row is measured on a T4 for Z-Image, Krea 2 and MiniMax H3; **every other cell below 48GB is interpolated**, including all of FLUX.1 and FLUX.2. A peak that leaves room on a 48GB card should also fit a smaller one, but no run has been done on one. -| Card | Z-Image 512 | Z-Image 1024 | Krea 2 512 | Krea 2 1024 | FLUX.2 512 | FLUX.2 1024 | MiniMax H3 | LTX-2.5 512 | -| ---- | ----------- | ------------ | ---------- | ----------- | ---------- | ----------- | ----------- | ----------- | -| 16GB | yes | no | yes, 4-bit | no | yes | yes | yes, slowly | no | -| 24GB | yes | yes | yes | no | yes | yes | yes | no | -| 32GB | yes | yes | yes | yes, 4-bit | yes | yes | yes | no | -| 48GB | yes | yes | yes | 4-bit only | yes | yes | yes | yes | +| Card | Z-Image 512 | Z-Image 1024 | Krea 2 512 | Krea 2 1024 | FLUX.1 512 | FLUX.1 1024 | FLUX.2 512 | FLUX.2 1024 | MiniMax H3 | LTX-2.5 512 | +| ---- | ----------- | ------------ | ---------- | ----------- | ---------- | ----------- | ---------- | ----------- | ----------- | ----------- | +| 12GB | no | no | no | no | yes, 4-bit | no | yes | yes | no | no | +| 16GB | yes | no | yes, 4-bit | no | yes, 4-bit | yes, 4-bit | yes | yes | yes, slowly | no | +| 24GB | yes | yes | yes | no | yes, 4-bit | yes, 4-bit | yes | yes | yes | no | +| 32GB | yes | yes | yes | yes, 4-bit | yes | yes | yes | yes | yes | no | +| 48GB | yes | yes | yes | 4-bit only | yes | yes | yes | yes | yes | yes | H3 has one column because resolution barely moves it. The 16GB entry is measured on a T4 with 64GB of RAM, where the conditioner spills to the CPU: it fits in 12.7GB of VRAM but costs 16.2s a step and a 19 minute caption pass. The 24GB entry is interpolated from the 20.6GB peak, not measured on a 24GB card. A 16GB card with only 16GB of RAM is refused up front. @@ -319,7 +349,7 @@ Krea 2's base is 26GB at bf16, which is what makes it expensive to fine-tune. Th - **Full precision (bf16)** forces the unquantized base. - **4-bit (NF4)** forces the quantized base. -The setting appears for Krea 2 and FLUX.2, but it only pays off on Krea 2. Z-Image has no 4-bit path and does not need one: it trains in about 15GB at 1024, so bf16 already fits the cards people have. FLUX.2 has the path and gains nothing from it, because klein 4B is smaller than its own text encoder and the peak sits in the caching pass either way, so Auto leaves it at bf16. See [Benchmark results](#benchmark-results). +The setting appears for Krea 2, FLUX.1 and FLUX.2, but it only pays off on Krea 2 and FLUX.1. Z-Image has no 4-bit path and does not need one: it trains in about 15GB at 1024, so bf16 already fits the cards people have. FLUX.2 has the path and gains nothing from it, because klein 4B is smaller than its own text encoder and the peak sits in the caching pass either way, so Auto leaves it at bf16. FLUX.1 does need it: a 24GB base is the largest here after Krea 2's, and 4-bit takes the peak from 24.9GB to 10.4GB at 512. Auto sizes the base plus its activations against your card, so anything under 32GB gets 4-bit. See [Benchmark results](#benchmark-results). To keep the peak down, the VAE and text encoder are loaded first, used to cache latents and captions, then freed before the transformer loads, so the two never stack. Which half then owns the peak depends on the model: for Z-Image and Krea 2 it is the transformer, for FLUX.2 klein it is the caching pass. If you do hit an out-of-memory error, lower the training resolution before changing anything else. diff --git a/core/pyproject.toml b/core/pyproject.toml index c25f3ea..bd48507 100644 --- a/core/pyproject.toml +++ b/core/pyproject.toml @@ -1,7 +1,7 @@ [project] # PyPI name; the import package is `inline_core` (src/inline_core). name = "openchar-core" -version = "1.3.19" +version = "1.3.20" description = "The generation engine behind Inline Studio." readme = "README.md" license = "GPL-3.0-or-later" diff --git a/core/src/inline_core/device/memory.py b/core/src/inline_core/device/memory.py index ca49137..a0dc7d7 100644 --- a/core/src/inline_core/device/memory.py +++ b/core/src/inline_core/device/memory.py @@ -52,6 +52,24 @@ # int8 still cannot fit; it is CUDA-only and, like int8, never combined with CPU offload. _NF4_FACTOR = 0.28 +#: The ladder's own rungs, so a caller sizing a load reads them here rather than restating them. +_RESIDENT_FACTORS = { + Quantization.NONE: 1.0, + Quantization.INT8: _INT8_FACTOR, + Quantization.NF4: _NF4_FACTOR, +} + + +def resident_factor(quant: Quantization) -> float: + """What fraction of a checkpoint's on-disk bf16 weight bytes survive a quantization rung.""" + return _RESIDENT_FACTORS.get(quant, 1.0) + + +def activation_headroom_gb() -> float: + """VRAM the fit ladder reserves beyond the weights, for activations and allocator slack.""" + return _ACTIVATION_HEADROOM_GB + + def _system_ram_gb() -> float | None: try: diff --git a/core/src/inline_core/models/flux1/__init__.py b/core/src/inline_core/models/flux1/__init__.py new file mode 100644 index 0000000..199093b --- /dev/null +++ b/core/src/inline_core/models/flux1/__init__.py @@ -0,0 +1 @@ +"""FLUX.1: the diffusers ``FluxPipeline`` family, as a generation node and a LoRA training base.""" diff --git a/core/src/inline_core/models/flux1/provider.py b/core/src/inline_core/models/flux1/provider.py new file mode 100644 index 0000000..eded57d --- /dev/null +++ b/core/src/inline_core/models/flux1/provider.py @@ -0,0 +1,87 @@ +"""FLUX.1's answer to "what do I need on disk" - the model popup's data source for the node.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from ...config import models_dir +from ..requirements import ModelComponent +from .requirements import ( + flux1_checkpoints, + flux1_encoders, + flux1_requirements, + footprint_bytes, + resolve_clip, + resolve_diffusion, + resolve_text_encoder, + resolve_vae, + resolved_variant, +) + + +class Flux1Provider: + """Requirements + fit estimate for the FLUX.1 node. + + One provider covers every build: dev, schnell, Kontext, Krea and the Fill/Control checkpoints + share the same VAE and the same two encoders, so only the transformer row changes. + """ + + def components(self, params: dict[str, object] | None = None) -> list[ModelComponent]: + return flux1_requirements(params) + + def download_target(self, component: ModelComponent) -> Path: + return models_dir() / component.category + + def resolved(self) -> dict[str, str]: + """What the node would load right now, so its pickers open on the real files rather than + on "auto". Names are relative to their category folder, matching the dropdown values.""" + picks = { + "model": resolve_diffusion(None), + "vae": resolve_vae(None), + "text_encoder": resolve_text_encoder(None), + "clip": resolve_clip(None), + } + out = {key: path.name for key, path in picks.items() if path is not None} + variant = resolved_variant(None) + if variant is not None: + out["variant"] = variant.key + return out + + def catalog_options(self, category: str) -> list[str] | None: + """Only the files this node can actually load. The categories are shared with Z-Image, + Krea 2 and FLUX.2, so an unfiltered list offers checkpoints that would fail on load.""" + if category == "diffusion_models": + return [p.name for p in flux1_checkpoints()] + if category == "text_encoders": + return [p.name for p in flux1_encoders()] + return None + + def estimate(self, policy: Any) -> dict[str, Any] | None: + """Whether the installed checkpoint fits this machine, and how, so the popup warns before a + load. Pure ``stat`` plus a live VRAM/RAM probe; None when it cannot be sized.""" + if policy is None: + return None + try: + from ...device.policy import ModelFootprint + except ImportError: + return None + footprint = ModelFootprint( + **footprint_bytes( + resolve_diffusion(None), resolve_vae(None), resolve_text_encoder(None), + resolve_clip(None), + ) + ) + fit = policy.estimate_fit(footprint) # pure - never mutates the shared policy + if fit is None: + return None + soft = not fit.fits or fit.plan in ("int8", "nf4", "offload") + return { + "plan": fit.plan, + "fits": fit.fits, + "requiredVramMb": int(fit.required_vram_gb * 1024), + "totalVramMb": int(fit.total_vram_gb * 1024) if fit.total_vram_gb else None, + "freeVramMb": policy.free_vram_mb(), + "freeRamMb": policy.free_ram_mb(), + "warning": fit.note if soft else None, + } diff --git a/core/src/inline_core/models/flux1/requirements.py b/core/src/inline_core/models/flux1/requirements.py new file mode 100644 index 0000000..b6650c8 --- /dev/null +++ b/core/src/inline_core/models/flux1/requirements.py @@ -0,0 +1,346 @@ +"""What FLUX.1 needs on disk, and whether it is there - the data behind the node's model popup. + +**No hidden downloads.** A component counts as present only because the user placed the file under +``models/`` or fetched it through the popup. Nothing here ever reaches the network. + +Torch-free (pure filesystem + safetensors headers) so the popup works on an install with no ML +stack. Files are matched by reading each candidate's header rather than trusting its name, which is +what lets a Z-Image, Krea 2, FLUX.1 and FLUX.2 checkpoint share ``diffusion_models/`` safely - and +what keeps FLUX.1's two encoders out of FLUX.2's encoder slot, which they would otherwise fit. + +**Licence.** dev's weights are non-commercial whichever mirror they come from, and a LoRA trained on +them is a derivative that inherits it. The row labels say so; that is the only place a user sees it. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +from ...config import models_dir +from ..catalog import resolve_picked +from ..requirements import ModelComponent +from . import variants as V + +__all__ = [ + "CLIP_FILE", + "DIFFUSION_FILE", + "TEXT_ENCODER_FILE", + "VAE_FILE", + "download_target", + "flux1_checkpoints", + "flux1_encoders", + "footprint_bytes", + "flux1_requirements", + "resolve_clip", + "resolve_diffusion", + "resolve_text_encoder", + "resolve_vae", + "resolved_variant", +] + +#: ComfyUI's repackaged single files: ungated, one consolidated ``.safetensors`` per component. +#: BFL's own repos are gated for every FLUX.1 build - schnell included - which breaks the popup. +#: Note the transformer sits at the repo root here, unlike the FLUX.2 mirror's ``split_files/``. +DIFFUSION_REPO = "Comfy-Org/flux1-dev" +DIFFUSION_FILE = "flux1-dev.safetensors" +#: Both encoders come from one ungated Apache-2.0 repo. +ENCODER_REPO = "comfyanonymous/flux_text_encoders" +TEXT_ENCODER_FILE = "t5xxl_fp16.safetensors" +CLIP_FILE = "clip_l.safetensors" +#: FLUX.1's VAE is ``ae.safetensors``, and Z-Image ships the *same weights* under the same name - +#: identical tensor names, shapes and values, and the same 0.3611/0.1159 scale and shift. So a +#: Z-Image install already has this file, and pointing at its repo means no second 335MB download. +VAE_REPO = "Comfy-Org/z_image" +VAE_FILE = "ae.safetensors" + +_WEIGHT_SUFFIXES = (".safetensors", ".sft", ".gguf") + +#: Offered as suggestions; none of them block a run. +_EXTRAS: tuple[tuple[str, str, str, str, str], ...] = ( + ( + "diffusion_fp8", + "FLUX.1 dev fp8 (12 GB, non-commercial)", + "diffusion_models", + "Kijai/flux-fp8", + "flux1-dev-fp8.safetensors", + ), + ( + "text_encoder_fp8", + "T5-XXL fp8 scaled (5 GB, for a smaller card)", + "text_encoders", + ENCODER_REPO, + "t5xxl_fp8_e4m3fn_scaled.safetensors", + ), +) + + +# --- filesystem resolution ----------------------------------------------------------------------- + + +def _category(name: str) -> Path: + return models_dir() / name + + +def _weight_files(category: str) -> list[Path]: + """Every candidate in a category: consolidated single files, plus diffusers-format folders.""" + root = _category(category) + if not root.is_dir(): + return [] + return sorted( + p + for p in root.iterdir() + if (p.is_file() and p.suffix.lower() in _WEIGHT_SUFFIXES) + or (p.is_dir() and (p / "config.json").is_file()) + ) + + +#: Header reads are cheap but the popup opens often, so identification is memoized on +#: (path, size, mtime) - a replaced file re-identifies, an untouched one does not. +_IDENTIFIED: dict[tuple[str, int, int], V.Flux1Variant | None] = {} + + +def _identify(path: Path) -> V.Flux1Variant | None: + try: + target = path / "config.json" if path.is_dir() else path + stat = target.stat() + except OSError: + return None + key = (str(path), stat.st_size, int(stat.st_mtime)) + if key not in _IDENTIFIED: + _IDENTIFIED[key] = V.detect(path) + return _IDENTIFIED[key] + + +def flux1_checkpoints() -> list[Path]: + """Every installed file that identifies as a FLUX.1 transformer.""" + return [p for p in _weight_files("diffusion_models") if _identify(p) is not None] + + +def _encoder_kind(path: Path) -> str | None: + """``"t5xxl"`` | ``"clip-l"`` | None, from the header alone. + + By content, never by name: both files live in ``text_encoders/`` beside Qwen3 and Mistral-3, and + T5's 4096-wide embedding is exactly the width FLUX.2 klein 9B looks for. + """ + if path.is_dir() or path.suffix.lower() not in (".safetensors", ".sft"): + return None + try: + from ..checkpoint import CheckpointReader + + keys = CheckpointReader(path).shapes() + except Exception: # noqa: BLE001 - an unreadable file simply does not match + return None + if any(k.startswith("encoder.block.") for k in keys) and "shared.weight" in keys: + return "t5xxl" + if any(k.startswith("text_model.") for k in keys): + return "clip-l" + return None + + +def _is_flux1_vae(path: Path) -> bool: + """Whether this is FLUX.1's LDM-style VAE, from the header alone. + + ``vae/`` is shared, and a name check is not enough: FLUX.2's file is called ``flux2-vae`` and + would match any "flux" fallback, while being a different VAE entirely. It is told apart by the + running batch-norm buffers it carries **in place of** FLUX.1's scalar scale and shift. + """ + if path.is_dir() or path.suffix.lower() not in (".safetensors", ".sft"): + return False + try: + from ..checkpoint import CheckpointReader + + keys = CheckpointReader(path).shapes() + except Exception: # noqa: BLE001 - an unreadable file simply does not match + return False + if any(k.startswith("bn.") or ".bn." in k for k in keys): + return False # FLUX.2's VAE + return "decoder.conv_out.weight" in keys and "encoder.conv_out.weight" in keys + + +def resolve_diffusion(params: dict[str, object] | None = None) -> Path | None: + """The FLUX.1 checkpoint to load: an explicit pick, ``INLINE_FLUX1_MODEL``, else the first file + in ``diffusion_models/`` that identifies as FLUX.1.""" + env = os.environ.get("INLINE_FLUX1_MODEL", "").strip() + if env: + path = Path(env) + return path if path.exists() else None + chosen = (params or {}).get("model") + if str(chosen or "").strip(): + return resolve_picked("diffusion_models", chosen) + return next(iter(flux1_checkpoints()), None) + + +def resolved_variant(params: dict[str, object] | None = None) -> V.Flux1Variant | None: + """Which variant the node will run: the explicit ``variant`` param, else what the resolved + checkpoint identifies as.""" + forced = V.get(str((params or {}).get("variant") or "")) + if forced is not None: + return forced + diffusion = resolve_diffusion(params) + return _identify(diffusion) if diffusion is not None else None + + +def resolve_vae(params: dict[str, object] | None = None) -> Path | None: + """The FLUX.1 VAE, which is also Z-Image's - the same file under the same name. + + ``vae/`` is shared, so the canonical name is matched first and a FLUX-named repack only after. + Krea 2's Qwen-Image VAE and FLUX.2's carry their own names and are never reached.""" + env = os.environ.get("INLINE_FLUX1_VAE", "").strip() + if env: + path = Path(env) + return path if path.exists() else None + chosen = (params or {}).get("vae") + if str(chosen or "").strip(): + return resolve_picked("vae", chosen) + exact = _category("vae") / VAE_FILE + if exact.is_file() and _is_flux1_vae(exact): + return exact + return next((p for p in _weight_files("vae") if _is_flux1_vae(p)), None) + + +def resolve_text_encoder(params: dict[str, object] | None = None) -> Path | None: + """T5-XXL, identified by its own key layout rather than by name or width.""" + env = os.environ.get("INLINE_FLUX1_TEXT_ENCODER", "").strip() + if env: + path = Path(env) + return path if path.exists() else None + chosen = (params or {}).get("text_encoder") + if str(chosen or "").strip(): + return resolve_picked("text_encoders", chosen) + return next((p for p in _weight_files("text_encoders") if _encoder_kind(p) == "t5xxl"), None) + + +def resolve_clip(params: dict[str, object] | None = None) -> Path | None: + """CLIP-L, the pooled half of FLUX.1's conditioning.""" + env = os.environ.get("INLINE_FLUX1_CLIP", "").strip() + if env: + path = Path(env) + return path if path.exists() else None + chosen = (params or {}).get("clip") + if str(chosen or "").strip(): + return resolve_picked("text_encoders", chosen) + return next((p for p in _weight_files("text_encoders") if _encoder_kind(p) == "clip-l"), None) + + +def flux1_encoders() -> list[Path]: + """Both of FLUX.1's encoders, so the node's picker offers neither Qwen3 nor Mistral-3.""" + return [p for p in _weight_files("text_encoders") if _encoder_kind(p) is not None] + + +# --- memory footprint ---------------------------------------------------------------------------- + + +def _file_bytes(path: object) -> int: + text = str(path or "").strip() + if not text: + return 0 + try: + p = Path(text) + if p.is_dir(): + return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) + return p.stat().st_size if p.is_file() else 0 + except OSError: + return 0 + + +def footprint_bytes( + diffusion: object = None, + vae: object = None, + text_encoder: object = None, + clip: object = None, +) -> dict[str, int]: + """On-disk sizes keyed to match ``ModelFootprint``, for the device policy's fit estimate. + + CLIP-L is folded into the text-encoder total rather than dropped: it is only 246 MB, but both + encoders are resident together while the prompt is encoded. + """ + return { + "diffusion_bytes": _file_bytes(diffusion), + "text_encoder_bytes": _file_bytes(text_encoder) + _file_bytes(clip), + "vae_bytes": _file_bytes(vae), + "controlnet_bytes": 0, + } + + +# --- the requirements view (the popup's data) ----------------------------------------------------- + + +def _component( + *, id: str, label: str, category: str, filename: str, present: bool, repo: str, repo_file: str, + optional: bool = False, +) -> ModelComponent: + return ModelComponent( + id=id, + label=label, + category=category, + present=present, + filename=filename, + repo=repo, + repo_file=repo_file, + optional=optional, + ) + + +def flux1_requirements(params: dict[str, object] | None = None) -> list[ModelComponent]: + """The popup's rows: the four required components, then the optional extras.""" + variant = resolved_variant(params) + label = f" ({variant.label})" if variant else "" + required = [ + _component( + id="diffusion", + # The licence rides on the label because this is the only place a user reads it. + label=f"Diffusion model{label} - non-commercial licence", + category="diffusion_models", + filename=DIFFUSION_FILE, + present=resolve_diffusion(params) is not None, + repo=DIFFUSION_REPO, + repo_file=DIFFUSION_FILE, + ), + _component( + id="text_encoder", + label="Text encoder (T5-XXL)", + category="text_encoders", + filename=TEXT_ENCODER_FILE, + present=resolve_text_encoder(params) is not None, + repo=ENCODER_REPO, + repo_file=TEXT_ENCODER_FILE, + ), + _component( + id="clip", + label="CLIP-L (pooled conditioning)", + category="text_encoders", + filename=CLIP_FILE, + present=resolve_clip(params) is not None, + repo=ENCODER_REPO, + repo_file=CLIP_FILE, + ), + _component( + id="vae", + label="VAE", + category="vae", + filename=VAE_FILE, + present=resolve_vae(params) is not None, + repo=VAE_REPO, + repo_file=VAE_FILE, + ), + ] + extras = [ + _component( + id=extra_id, + label=extra_label, + category=category, + filename=Path(repo_file).name, + present=(_category(category) / Path(repo_file).name).is_file(), + repo=repo, + repo_file=repo_file, + optional=True, + ) + for extra_id, extra_label, category, repo, repo_file in _EXTRAS + ] + return required + extras + + +def download_target(component: ModelComponent) -> Path: + """Where the component's file lands: its category folder, flat, under the models root.""" + return _category(component.category) diff --git a/core/src/inline_core/models/flux1/runner.py b/core/src/inline_core/models/flux1/runner.py new file mode 100644 index 0000000..f45ba84 --- /dev/null +++ b/core/src/inline_core/models/flux1/runner.py @@ -0,0 +1,477 @@ +"""FLUX.1 runner: prompt (+ optional image) -> one rendered take. + +A single generation node, ``black-forest-labs/flux-1``, backed by diffusers' ``FluxPipeline`` and +``FluxImg2ImgPipeline``. Placement, the pipeline cache, prompt pre-encoding and the OOM messages all +come from ``models/pipeline_runtime.py``; this module holds only what is FLUX.1 specific. + +Two things separate it from its siblings. Conditioning is **two encoders** - T5-XXL for the sequence +and CLIP-L for the pooled vector - so both are staged, parked and detached together. And dev is +guidance-distilled, so ``guidance`` is an embedding the transformer consumes rather than a CFG pass: +there is no negative prompt on this node at all. + +torch + diffusers are imported at module top on purpose: an absent ``runtime`` extra makes this +import raise, and ``server.bootstrap`` skips the model so the engine still boots. +""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Callable +from dataclasses import replace +from typing import Any + +import torch +from diffusers import FluxImg2ImgPipeline + +from ...device.policy import DevicePolicy, ModelFootprint, Profile, Quantization +from ...errors import CancelledError, ComponentError +from ...graph.descriptor import NodeDescriptor, ParamField, Port, Widget +from ...graph.loader_runners import LoraRef +from ...graph.runners import NodeResult, NodeRunner +from ...graph.schema import Node, PortKind +from ...media import MediaKind +from ...runtime.context import ExecutionContext +from ...runtime.progress import Phase +from ...runtime.store import TakeStore +from .. import loaders +from .. import pipeline_runtime as rt +from ..sampling import SamplingFamily, apply_sampling, sampling_param_fields +from . import requirements as reqs +from . import variants as V + +# Every model this node needs comes from files under models/ (see `requirements.py`). Nothing is +# ever downloaded here: every load runs local_files_only=True. +_ARCH = "flux1" +_LABEL = "FLUX.1" +#: Both encoders, in the order the pipeline names them. Parked and detached as a set. +_ENCODERS = ("text_encoder", "text_encoder_2") + +#: Sentinels meaning "take the checkpoint's own default", so switching schnell for dev in the +#: dropdown moves 4 steps / guidance 0 to 28 / 3.5 without the user editing either field. +_AUTO_STEPS = 0 +_AUTO_GUIDANCE = -1.0 + +logger = logging.getLogger("inline_core.flux1") + + +FLUX1 = NodeDescriptor( + type="black-forest-labs/flux-1", + title="FLUX.1", + category="Generate", + icon="wand", + output_kind=MediaKind.IMAGE, + inputs=( + Port("prompt", "Prompt", PortKind.TEXT, required=True), + # Optional component handles from load/* subnodes - wire one to override the dropdown. + Port("model", "Diffusion model", PortKind.MODEL, required=False), + Port("vae", "VAE", PortKind.VAE, required=False), + Port("text_encoder", "Text encoder (T5)", PortKind.TEXT_ENCODER, required=False), + Port("clip", "CLIP-L", PortKind.TEXT_ENCODER, required=False), + Port("lora", "LoRA", PortKind.LORA, required=False), + Port("image", "Image (img2img)", PortKind.IMAGE, required=False), + ), + outputs=(Port("image", "Image", PortKind.IMAGE),), + params=( + # No negative prompt: dev is guidance-distilled and FluxPipeline has no negative path. + ParamField("width", "Width", Widget.NUMBER, 1024, min=256, max=2048, step=64), + ParamField("height", "Height", Widget.NUMBER, 1024, min=256, max=2048, step=64), + ParamField("steps", "Steps (0 = from model)", Widget.NUMBER, _AUTO_STEPS, min=0, max=100, + step=1), + ParamField("guidance", "Guidance (-1 = from model)", Widget.NUMBER, _AUTO_GUIDANCE, + min=-1.0, max=30.0, step=0.5), + *sampling_param_fields(SamplingFamily.FLOW_MATCH), + ParamField( + "strength", "Denoise strength", Widget.NUMBER, 0.6, min=0.0, max=1.0, step=0.05, + advanced=True, + ), + ParamField("seed", "Seed (-1 = random)", Widget.SEED, -1), + # Advanced: pick a specific file per component. "" = auto. + ParamField("model", "Diffusion model", Widget.SELECT, "", + options_from="diffusion_models", advanced=True), + ParamField("text_encoder", "Text encoder (T5-XXL)", Widget.SELECT, "", + options_from="text_encoders", advanced=True), + ParamField("clip", "CLIP-L", Widget.SELECT, "", + options_from="text_encoders", advanced=True), + ParamField("vae", "VAE", Widget.SELECT, "", options_from="vae", advanced=True), + ), +) + + +def register_flux1(registry: Any, store: TakeStore, policy: DevicePolicy) -> None: + """Register the FLUX.1 node and its runner. Called best-effort by server.bootstrap.""" + registry.register(FLUX1, Flux1Runner(store, policy)) + + +def _snap(value: int) -> int: + """Down onto a multiple of 16: the VAE's 8x downscale then the pipeline's 2x2 latent fold.""" + return max(256, (int(value) // 16) * 16) + + +def _needs_staged_encode( + diffusion: str, vae: str, text: str, clip: str, quant: Quantization, policy: DevicePolicy +) -> bool: + """Whether the text encoders and the transformer are too big to be resident together. + + Unlike FLUX.2, on-disk size is *not* the right measure here: dev ships as a plain bf16 file, so + what it weighs is not what it occupies once the ladder picks a quantization. The estimate is + scaled by the planned rung instead, and the VAE is left unscaled because it is never quantized. + + Note the transformer still peaks at its full bf16 size inside ``from_single_file`` - bnb only + quantizes on the move to CUDA - so staging buys the encoders' room, not the transformer's. + """ + from ...device.memory import activation_headroom_gb, resident_factor + + budget_mb = policy.vram_budget_mb() + if not budget_mb: + return False # CPU or an unmeasurable device takes the normal path + sizes = reqs.footprint_bytes(diffusion, vae, text, clip) + gb = 1024**3 + factor = resident_factor(quant) + transformer = sizes["diffusion_bytes"] * factor / gb + encoders = sizes["text_encoder_bytes"] * factor / gb + vae_gb = sizes["vae_bytes"] / gb + budget = budget_mb / 1024 - activation_headroom_gb() + return transformer + encoders + vae_gb > budget and transformer + vae_gb <= budget + + +def _resolve_steps(params: dict[str, Any], variant: V.Flux1Variant | None) -> int: + picked = int(params.get("steps") or _AUTO_STEPS) + return max(1, picked) if picked > 0 else (variant.steps if variant else 28) + + +def _resolve_guidance(params: dict[str, Any], variant: V.Flux1Variant | None) -> float: + picked = float(params.get("guidance", _AUTO_GUIDANCE)) + return picked if picked >= 0 else (variant.guidance if variant else 3.5) + + +class Flux1Runner(NodeRunner): + produces_takes = True + + def __init__(self, store: TakeStore, policy: DevicePolicy) -> None: + self._store = store + self._policy = policy + + def run(self, node: Node, inputs: dict[str, list[Any]], ctx: ExecutionContext) -> NodeResult: + prompt = rt.first_str(inputs.get("prompt")) + if not prompt: + raise ComponentError("FLUX.1 needs a prompt.") + params = {**FLUX1.defaults(), **node.params} + width, height = _snap(int(params["width"])), _snap(int(params["height"])) + seed = rt.resolve_seed(params.get("seed")) + sampler, scheduler = str(params["sampler"]), str(params["scheduler"]) + image_ref = rt.first(inputs.get("image")) + img2img = image_ref is not None + + # Wired component handles from load/* subnodes override the dropdowns. + model_ref = rt.component_ref(inputs, "model", "diffusion", _LABEL) + vae_ref = rt.component_ref(inputs, "vae", "vae", _LABEL) + # Both encoders are "text_encoder" handles - there is no separate CLIP kind - so which + # requirement a wired one satisfies comes from the port it arrived on, not from its kind. + te_ref = rt.component_ref(inputs, "text_encoder", "text_encoder", _LABEL) + clip_ref = rt.component_ref(inputs, "clip", "text_encoder", _LABEL) + loras = rt.lora_stack(inputs, _LABEL) + by_row = { + "diffusion": model_ref, "vae": vae_ref, "text_encoder": te_ref, "clip": clip_ref, + } + wired = {row for row, ref in by_row.items() if ref is not None} + + # No hidden downloads: a required component that is neither wired nor on disk fails fast. + missing = [ + c.label + for c in reqs.flux1_requirements(params) + if not c.present and not c.optional and c.id not in wired + ] + if missing: + raise ComponentError( + "FLUX.1 models missing: " + + ", ".join(missing) + + ". Download them from the node's model popup (the hint on the node)." + ) + + source = model_ref.file if model_ref else rt.path_or_none(reqs.resolve_diffusion(params)) + if not source: # defensive: the missing-check above already covers this + raise ComponentError("FLUX.1 diffusion model not found in diffusion_models/.") + variant = V.detect(source) + config = V.config_for(source) + if config is None: + raise ComponentError(f"{source} is not a FLUX.1 checkpoint.") + steps = _resolve_steps(params, variant) + guidance = _resolve_guidance(params, variant) + + vae_file = vae_ref.file if vae_ref else rt.path_or_none(reqs.resolve_vae(params)) + te_file = te_ref.file if te_ref else rt.path_or_none(reqs.resolve_text_encoder(params)) + clip_file = clip_ref.file if clip_ref else rt.path_or_none(reqs.resolve_clip(params)) + + # Size-aware placement: hand the policy the on-disk sizes so it fits dtype/quant/offload to + # THIS GPU, then refuse an impossible load up front rather than OOM-killing the server. + self._policy.set_footprint( + ModelFootprint(**reqs.footprint_bytes(source, vae_file, te_file, clip_file)) + ) + fit = self._policy.fit_estimate() + if fit is not None and not fit.fits: + raise ComponentError(rt.wont_fit_message(fit)) + + # A checkpoint that ships already quantized loads as-is: re-quantizing it is a hard error, + # and its on-disk weights are already the resident ones. + quant = Quantization.NONE if V.is_prequantized(source) else self._policy.quantization() + staged = _needs_staged_encode(source, vae_file, te_file, clip_file, quant, self._policy) + logger.info( + "FLUX.1 (%s) run: %dx%d, %d steps, guidance=%.1f, img2img=%s | %s", + variant.label if variant else "unknown", width, height, steps, guidance, img2img, + rt.device_report(self._policy), + ) + rt.reset_peak_vram() + rt.raise_if_cancelled(ctx) # bail before a 24GB load if already cancelled + ctx.emitter.emit(rt.progress_event(ctx, node, Phase.LOADING, 0.0, status="Loading model…")) + try: + pipe = _load_pipeline( + self._policy, + img2img=img2img, + source=source, + config=config, + vae=vae_file, + text=te_file, + clip=clip_file, + quant=quant, + loras=loras, + staged=staged, + cancel_check=lambda: rt.raise_if_cancelled(ctx), + ) + except CancelledError: + rt.free_vram() # a cancelled load must return whatever VRAM it placed + raise + except torch.cuda.OutOfMemoryError as error: + rt.free_vram() + raise ComponentError(_oom(width, height)) from error + except MemoryError as error: + rt.free_vram() + raise ComponentError(_oom(width, height, host=True)) from error + + placement = self._policy.placement("denoiser") + on_cpu = placement.offload or self._policy.profile is Profile.CPU + gen_device = "cpu" if on_cpu else str(placement.device) + generator = torch.Generator(device=gen_device).manual_seed(seed) + + def on_step_end(_pipe: Any, step: int, _t: Any, kwargs: dict[str, Any]) -> dict[str, Any]: + if ctx.cancel.cancelled: + raise CancelledError("Run cancelled.") + done = step + 1 + ctx.emitter.emit( + rt.progress_event( + ctx, node, Phase.SAMPLE, done / steps, + step=done, step_count=steps, status=f"Step {done}/{steps}", + ) + ) + return kwargs + + call: dict[str, Any] = dict( + height=height, + width=width, + num_inference_steps=steps, + guidance_scale=guidance, + generator=generator, + output_type="pil", + callback_on_step_end=on_step_end, + ) + call.update(_prompt_kwargs(pipe, self._policy, prompt=prompt)) + if staged and getattr(pipe, "transformer", None) is None: + if "prompt_embeds" not in call: + raise ComponentError( + "FLUX.1 needs its prompt encoded before the transformer loads, but encoding " + "did not produce embeddings. Free some VRAM and retry." + ) + ctx.emitter.emit( + rt.progress_event(ctx, node, Phase.LOADING, 0.5, status="Loading transformer…") + ) + loaders.attach_flux1_transformer( + pipe, + arch=_ARCH, + diffusion_file=source, + config=config, + vae_file=vae_file, + dtype=rt.torch_dtype(placement), + quant=quant, + device=None if placement.offload else str(placement.device), + loras=loras, + ) + rt.configure_pipeline(pipe, self._policy) + rt.capture_base_scheduler_config(pipe) + if img2img: + call["image"] = rt.load_image(image_ref, _LABEL) + call["strength"] = float(params.get("strength", 0.6)) + + base_config = getattr(pipe, "_inline_base_scheduler_config", None) + if base_config is not None: + sigmas = apply_sampling( + pipe, base_config, SamplingFamily.FLOW_MATCH, sampler, scheduler, steps + ) + if sigmas is not None: + call["sigmas"] = sigmas + + logger.info( + "FLUX.1 sampling %d steps on %s (sampler=%s, scheduler=%s)…", + steps, gen_device, sampler, scheduler, + ) + rt.raise_if_cancelled(ctx) # cancelled during load? don't start the denoise + sample_start = time.perf_counter() + try: + with rt.text_encoder_detached(pipe, "prompt_embeds" in call, _ENCODERS): + image = pipe(**call).images[0] + except CancelledError: + rt.free_vram() # release partial-denoise activations so the next run isn't starved + raise + except torch.cuda.OutOfMemoryError as error: + rt.free_vram() + raise ComponentError(_oom(width, height)) from error + except MemoryError as error: + rt.free_vram() + raise ComponentError(_oom(width, height, host=True)) from error + elapsed = time.perf_counter() - sample_start + peak_gb = rt.peak_vram_gb() + logger.info( + "FLUX.1 sampled %dx%d in %.1fs (%.2fs/step)%s | %s", + width, height, elapsed, elapsed / steps, + f", peak VRAM {peak_gb:.1f}GB" if peak_gb else "", rt.device_report(self._policy), + ) + rt.free_vram() # return fragmented free blocks to the driver (keeps the model resident) + + save_status = "Saving…" + (f" (peak VRAM {peak_gb:.1f}GB)" if peak_gb else "") + ctx.emitter.emit(rt.progress_event(ctx, node, Phase.SAVE, 1.0, status=save_status)) + take = self._store.save( + ctx.run_id, + node.id, + image, + { + "model": source, + "variant": variant.key if variant else "", + "prompt": prompt, + "width": width, + "height": height, + "steps": steps, + "guidance": guidance, + "sampler": sampler, + "scheduler": scheduler, + "seed": seed, + **({"strength": call["strength"]} if img2img else {}), + **( + {"loras": [{"file": lo.file, "strength": lo.strength} for lo in loras]} + if loras + else {} + ), + }, + ) + return NodeResult(outputs={"image": take}, takes=(take,)) + + +def _prompt_kwargs(pipe: Any, policy: DevicePolicy, *, prompt: str) -> dict[str, Any]: + """Precomputed embeddings (encoders then parked on the CPU), or the raw prompt as a fallback. + + Both encoders are staged together: FLUX.1's ``encode_prompt`` returns the T5 sequence *and* the + CLIP pooled vector, and the transformer needs both. + """ + + def raw() -> dict[str, Any]: + return {"prompt": prompt} + + def encode(device: str) -> dict[str, Any]: + # encode_prompt called directly is not wrapped in the pipeline's @torch.no_grad (only + # __call__ is); the caller supplies it. + prompt_embeds, pooled, _ids = pipe.encode_prompt( + prompt=prompt, prompt_2=None, device=torch.device(device), max_sequence_length=512 + ) + return { + "prompt_embeds": rt.embeds_to(prompt_embeds, device), + "pooled_prompt_embeds": rt.embeds_to(pooled, device), + } + + return rt.encoded_prompt_kwargs( + pipe, policy, encode=encode, fallback=raw, encoders=_ENCODERS + ) + + +def _oom(width: int, height: int, *, host: bool = False) -> str: + where = "System RAM" if host else "VRAM" + return ( + f"{where} ran out generating {width}x{height} with FLUX.1. Its base is 24GB at bf16 beside " + "a 10GB T5 encoder, so try a smaller size, or let the fit ladder quantize by leaving the " + "memory profile on auto." + ) + + +def _load_pipeline( + policy: DevicePolicy, + *, + img2img: bool, + source: str, + config: dict[str, Any], + vae: str, + text: str, + clip: str, + quant: Quantization = Quantization.NONE, + loras: tuple[LoraRef, ...] = (), + staged: bool = False, + cancel_check: Callable[[], None] | None = None, +) -> Any: + # CLIP rides in the key's `controlnet` slot, which is the only field for a second weight file: + # `component_files` keeps it alive and `evict_stale` compares it, which is all it has to do. + key = rt.PipelineKey( + arch=_ARCH, + diffusion=source, + vae=vae, + text_encoder=text, + variant="i2i" if img2img else "t2i", + quant=quant.value, + loras=loaders.lora_cache_key(loras), + controlnet=clip, + ) + with rt.PIPELINES.lock: + cached = rt.PIPELINES.get(key) + if cached is not None: + logger.info( + "Pipeline cache hit (%s, img2img=%s) - reusing loaded weights", source, img2img + ) + return cached + if cancel_check is not None: + cancel_check() # bail before the disk read if the run was cancelled while queued + started = time.perf_counter() + logger.info( + "Loading FLUX.1 pipeline: source=%s, img2img=%s, staged=%s | %s", + source, img2img, staged, rt.device_report(policy), + ) + # Free any *other* model still resident before loading this one, so switching checkpoints + # doesn't stack VRAM. Keeps this key's own components, including a t2i base reused below. + rt.PIPELINES.evict_stale(key) + placement = policy.placement("denoiser") + # Resident placement streams weights straight to the GPU; the offload path loads to CPU so + # accelerate can install its hooks before placing. + device = None if placement.offload else str(placement.device) + if staged: + logger.info("FLUX.1: encoding before the transformer loads - they do not fit together") + pipe = loaders.assemble_flux1_encoders( + arch=_ARCH, img2img=img2img, vae_file=vae, text_encoder_file=text, clip_file=clip, + dtype=rt.torch_dtype(placement), quant=quant, + vae_dtype=rt.torch_dtype(policy.placement("vae")), device=device, + ) + rt.capture_base_scheduler_config(pipe) + # Deliberately not cached: it holds no transformer, so a later run would reuse a + # half-built pipeline. The caller attaches one and caches nothing either. + return pipe + base = rt.PIPELINES.get(replace(key, variant="t2i")) + if img2img and base is not None: + pipe = FluxImg2ImgPipeline.from_pipe(base) + logger.info( + "Built img2img pipeline from cached base in %.1fs", time.perf_counter() - started + ) + else: + pipe = loaders.assemble_flux1_pipeline( + arch=_ARCH, img2img=img2img, diffusion_file=source, config=config, vae_file=vae, + text_encoder_file=text, clip_file=clip, dtype=rt.torch_dtype(placement), + quant=quant, vae_dtype=rt.torch_dtype(policy.placement("vae")), device=device, + loras=loras, cancel_check=cancel_check, + ) + rt.configure_pipeline(pipe, policy) + rt.capture_base_scheduler_config(pipe) + rt.PIPELINES.put(key, pipe) + logger.info("FLUX.1 pipeline ready in %.1fs", time.perf_counter() - started) + return pipe diff --git a/core/src/inline_core/models/flux1/variants.py b/core/src/inline_core/models/flux1/variants.py new file mode 100644 index 0000000..7c01893 --- /dev/null +++ b/core/src/inline_core/models/flux1/variants.py @@ -0,0 +1,363 @@ +"""Which FLUX.1 checkpoint a file is, read from its own tensor shapes. + +``diffusion_models/`` is shared across every architecture, so a checkpoint is identified by content +and never by filename. The geometry is reconstructed from the header rather than bundled per build, +which is what lets one node load dev, schnell, Kontext and the Fill/Control builds - and a future +one - with no code change. + +Torch-free and header-only: nothing here reads tensor data or imports torch, so the model popup +works on an install with no ML stack. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import Path + +__all__ = [ + "VARIANTS", + "Flux1Variant", + "config_for", + "derive_transformer_config", + "detect", + "folder_config", + "get", + "is_prequantized", + "quantization_of", + "single_file_blocker", + "trainable", +] + +#: The one number no tensor shape reveals. FLUX.1 splits its rotary embedding 16/56/56 across the +#: three position axes, which sums to the 128 head dim the shapes do give. +_FIXED_GEOMETRY: dict[str, object] = { + "patch_size": 1, + "out_channels": None, + "axes_dims_rope": [16, 56, 56], +} + +#: Every FLUX.1 build diffusers can load through ``FluxTransformer2DModel``. The three that share a +#: geometry are told apart by name; everything else falls out of the shapes (see ``detect``). +_INNER_DIM = 3072 + + +@dataclass(frozen=True) +class Flux1Variant: + """One FLUX.1 checkpoint: how to build it, and what it wants at sampling time.""" + + key: str + label: str + #: Which diffusers pipeline family to build: "t2i" | "kontext" | "fill" | "control". + pipeline: str + #: **Step**-distilled. schnell alone is; dev is *guidance*-distilled, which is a different thing + #: and is trained through by pinning guidance at 1 rather than around. + distilled: bool + #: 64 for a plain latent, 128 for a Control build's stacked hint, 384 for Fill's mask channels. + in_channels: int + guidance_embeds: bool + #: The loader arch key. One asset bundle serves the family: same VAE, same CLIP-L, same T5-XXL. + arch: str + steps: int + guidance: float + + +VARIANTS: tuple[Flux1Variant, ...] = ( + Flux1Variant( + key="dev", label="dev", pipeline="t2i", distilled=False, in_channels=64, + guidance_embeds=True, arch="flux1", steps=28, guidance=3.5, + ), + Flux1Variant( + key="schnell", label="schnell", pipeline="t2i", distilled=True, in_channels=64, + guidance_embeds=False, arch="flux1", steps=4, guidance=0.0, + ), + Flux1Variant( + key="krea-dev", label="Krea dev", pipeline="t2i", distilled=False, in_channels=64, + guidance_embeds=True, arch="flux1", steps=28, guidance=4.5, + ), + Flux1Variant( + key="kontext-dev", label="Kontext dev", pipeline="kontext", distilled=False, in_channels=64, + guidance_embeds=True, arch="flux1", steps=28, guidance=2.5, + ), + Flux1Variant( + key="fill-dev", label="Fill dev", pipeline="fill", distilled=False, in_channels=384, + guidance_embeds=True, arch="flux1", steps=50, guidance=30.0, + ), + Flux1Variant( + key="canny-dev", label="Canny dev", pipeline="control", distilled=False, in_channels=128, + guidance_embeds=True, arch="flux1", steps=50, guidance=30.0, + ), + Flux1Variant( + key="depth-dev", label="Depth dev", pipeline="control", distilled=False, in_channels=128, + guidance_embeds=True, arch="flux1", steps=30, guidance=10.0, + ), +) + +_BY_KEY = {v.key: v for v in VARIANTS} + + +def get(key: str | None) -> Flux1Variant | None: + return _BY_KEY.get((key or "").strip()) + + +def trainable(variant: Flux1Variant) -> bool: + """Whether a LoRA run can train against this build. + + Three exclusions, all so a run fails now rather than twenty minutes into a precache. schnell is + step-distilled and collapses the same way a distilled FLUX.2 does. Fill and Control take a mask + or a stacked hint in their extra input channels, which the dataset exporter does not produce - + their shapes only disagree once the first batch reaches the transformer. Kontext would train: + it is 64-channel and undistilled, and conditions through the token sequence rather than the + channels. But it learns an *edit* between a pair, and a dataset of single images teaches it + nothing it is used for, so it is left out until paired datasets exist. + """ + return not variant.distilled and variant.in_channels == 64 and variant.pipeline == "t2i" + + +# --- identifying a checkpoint -------------------------------------------------------------------- + +#: Prefixes ComfyUI-style repacks put in front of the keys. Stripped before matching so a Comfy +#: single file and a diffusers export identify the same way. +_PREFIXES = ("model.diffusion_model.", "diffusion_model.", "model.") + +_BLOCK_RE = re.compile(r"^transformer_blocks\.(\d+)\.") +_SINGLE_BLOCK_RE = re.compile(r"^single_transformer_blocks\.(\d+)\.") +_WEIGHT_SUFFIXES = (".safetensors", ".sft") + +#: The shipped checkpoints use BFL's key layout, not diffusers'. diffusers converts at load time +#: (``convert_flux_transformer_checkpoint_to_diffusers``), but identification runs on the raw +#: header, so the handful of keys we read are renamed here first. Only the identifying keys are +#: mapped - this is not a checkpoint converter, and it must not become one. +_BFL_RENAMES: tuple[tuple[re.Pattern[str], str], ...] = ( + (re.compile(r"^double_blocks\."), "transformer_blocks."), + (re.compile(r"^single_blocks\."), "single_transformer_blocks."), + (re.compile(r"^txt_in\.weight$"), "context_embedder.weight"), + (re.compile(r"^img_in\.weight$"), "x_embedder.weight"), + ( + re.compile(r"^vector_in\.in_layer\.weight$"), + "time_text_embed.text_embedder.linear_1.weight", + ), + ( + re.compile(r"^guidance_in\.in_layer\.weight$"), + "time_text_embed.guidance_embedder.linear_1.weight", + ), + (re.compile(r"(img_)?attn\.norm\.query_norm\.scale$"), "attn.norm_q.weight"), + (re.compile(r"(?<=\.)norm\.query_norm\.scale$"), "attn.norm_q.weight"), +) + + +def _strip(key: str) -> str: + for prefix in _PREFIXES: + if key.startswith(prefix): + key = key[len(prefix) :] + break + for pattern, replacement in _BFL_RENAMES: + key = pattern.sub(replacement, key) + return key + + +def derive_transformer_config(shapes: dict[str, list[int]]) -> dict[str, object] | None: + """Reconstruct a ``FluxTransformer2DModel`` config from a checkpoint's tensor shapes. + + Returns None when the file is not a FLUX.1 transformer. ``text_embedder`` is what separates + FLUX.1 from FLUX.2: only FLUX.1 conditions on a pooled CLIP vector, so FLUX.2's checkpoints + have no such key and cannot be claimed here. + """ + keys = {_strip(key): shape for key, shape in shapes.items()} + context = keys.get("context_embedder.weight") + x_embed = keys.get("x_embedder.weight") + pooled = keys.get("time_text_embed.text_embedder.linear_1.weight") + if not context or not x_embed or not pooled: + return None + if len(context) != 2 or len(x_embed) != 2 or len(pooled) != 2: + return None + + inner_dim, joint_attention_dim = context[0], context[1] + # diffusers' single-file converter hardcodes inner_dim 3072 and mlp_ratio 4.0 to split each + # single block's fused linear1, so a wider checkpoint would load mis-split rather than raise. + if inner_dim != _INNER_DIM: + return None + head_dim = next( + (shape[0] for key, shape in keys.items() if key.endswith("attn.norm_q.weight") and shape), + 0, + ) + if not head_dim or inner_dim % head_dim: + return None + + layers = {int(m.group(1)) for key in keys if (m := _BLOCK_RE.match(key))} + single_layers = {int(m.group(1)) for key in keys if (m := _SINGLE_BLOCK_RE.match(key))} + if not layers or not single_layers: + return None + + return { + **_FIXED_GEOMETRY, + "attention_head_dim": head_dim, + "guidance_embeds": any( + "guidance_embedder" in key or key.startswith("guidance_in.") for key in keys + ), + "in_channels": x_embed[1], + "joint_attention_dim": joint_attention_dim, + "num_attention_heads": inner_dim // head_dim, + "num_layers": len(layers), + "num_single_layers": len(single_layers), + "pooled_projection_dim": pooled[1], + } + + +def folder_config(path: str | Path) -> dict[str, object] | None: + """The transformer config of a diffusers-format checkpoint **folder**, or None. + + A folder ships its shards beside a ``config.json``, so there is no header to derive geometry + from - and no need. Recognised by its ``_class_name``, so a folder belonging to another + architecture is left alone. + """ + folder = Path(path) + marker = folder / "config.json" + if not folder.is_dir() or not marker.is_file(): + return None + try: + import json + + config = json.loads(marker.read_text()) + except (OSError, ValueError): + return None + if not isinstance(config, dict): + return None + if config.get("_class_name") != "FluxTransformer2DModel": + return None + return {k: v for k, v in config.items() if not k.startswith("_")} + + +#: Safetensors spells fp8 as ``F8_E4M3`` / ``F8_E5M2``. Never match on a ``.scale`` suffix alone: a +#: plain FLUX.1 checkpoint carries 152 RMSNorm ``norm.query_norm.scale`` weights that are not +#: quantization scales at all. Only a scale riding on a *fused qkv* is a repack artifact, and it is +#: the one diffusers' converter chokes on. +_FP8_PREFIX = "F8_" +_INT_DTYPES = frozenset({"I8", "U8"}) +_QKV_SCALE_SUFFIXES = ("qkv.scale", "qkv.weight_scale", "qkv.scale_weight") + + +def quantization_of(path: str | Path) -> str | None: + """The quantization a checkpoint already carries, or None if it is plain weights. + + A folder declares it in ``config.json``. A single file does not, so it is sniffed from the + header: an fp8 dtype, or the per-tensor scale tensors every int8/fp8 repack ships alongside. + """ + target = Path(path) + if target.is_dir(): + try: + import json + + config = json.loads((target / "config.json").read_text()) + except (OSError, ValueError): + return None + if not isinstance(config, dict): + return None + declared = config.get("quantization_config") + if not isinstance(declared, dict): + return None + return str(declared.get("quant_method", "quantized")) + + if not target.is_file() or target.suffix.lower() not in _WEIGHT_SUFFIXES: + return None + try: + from ..checkpoint import CheckpointReader + + dtypes = CheckpointReader(target).dtypes() + except Exception: # noqa: BLE001 - unreadable means "not something we can classify" + return None + weights = {d for key, d in dtypes.items() if key.endswith(".weight")} + if any(d.startswith(_FP8_PREFIX) for d in weights): + return "fp8" + if weights & _INT_DTYPES: + return "int8" + if any(key.endswith(_QKV_SCALE_SUFFIXES) for key in dtypes): + return "quantized" + return None + + +def is_prequantized(path: str | Path) -> bool: + """Whether a checkpoint carries its own quantization (an NF4 folder, an fp8 single file). + + Such a checkpoint must not be quantized again: its on-disk size already is its resident size, + and handing diffusers a second, different quantization config is a hard error. + """ + return quantization_of(path) is not None + + +def single_file_blocker(path: str | Path) -> str | None: + """Why a failed single-file load probably failed, or None if this is not the known cause. + + Diagnostic only, never a gate. The converter maps a ``.scale`` key to a weight and chunks it + into q/k/v, so a **per-tensor** (0-dim) scale on a fused qkv dies deep inside diffusers without + naming the file. + """ + target = Path(path) + if not target.is_file() or target.suffix.lower() not in _WEIGHT_SUFFIXES: + return None + try: + from ..checkpoint import CheckpointReader + + shapes = CheckpointReader(target).shapes() + except Exception: # noqa: BLE001 - unreadable is not our call to make here + return None + scalar = [k for k, s in shapes.items() if k.endswith(_QKV_SCALE_SUFFIXES) and not s] + if not scalar: + return None + return f"it carries a single scale value on each of {len(scalar)} fused qkv tensors" + + +def config_for(path: str | Path) -> dict[str, object] | None: + """The transformer geometry for a checkpoint, whether it is a single file or a folder.""" + folder = folder_config(path) + if folder is not None: + return folder + shapes = _shapes_of(path) + return derive_transformer_config(shapes) if shapes else None + + +def _shapes_of(path: str | Path) -> dict[str, list[int]] | None: + file = Path(path) + if not file.is_file() or file.suffix.lower() not in _WEIGHT_SUFFIXES: + return None + try: + from ..checkpoint import CheckpointReader + + return CheckpointReader(file).shapes() + except Exception: # noqa: BLE001 - an unreadable or foreign file is simply "not FLUX.1" + return None + + +def _name_flags(name: str) -> tuple[bool, bool]: + """(is_kontext, is_krea) read from a filename. + + dev, Kontext dev and Krea dev are byte-identical geometry - same widths, same block counts, + same guidance embedder - so no shape tells them apart and the name is the only signal left. + Everything else here is identified by content. + """ + padded = "-" + re.sub(r"[^a-z0-9]+", "-", name.lower()) + "-" + return "-kontext-" in padded, "-krea-" in padded + + +def detect(path: str | Path, shapes: dict[str, list[int]] | None = None) -> Flux1Variant | None: + """Which FLUX.1 variant a checkpoint is, or None if it is not one. + + Handles both a single ``.safetensors`` and a diffusers folder. Pass ``shapes`` when the header + has already been read, so it is not read twice. Never reads tensor data, never imports torch. + """ + file = Path(path) + config = derive_transformer_config(shapes) if shapes is not None else config_for(file) + if config is None: + return None + family = [v for v in VARIANTS if v.in_channels == config["in_channels"]] + if not family: + return None + # No guidance embedder means schnell, whatever the file is called. + if not config["guidance_embeds"]: + return next((v for v in family if v.distilled), None) + family = [v for v in family if not v.distilled] + is_kontext, is_krea = _name_flags(file.name) + if len(family) > 1: + wanted = "kontext-dev" if is_kontext else "krea-dev" if is_krea else "dev" + return next((v for v in family if v.key == wanted), family[0]) + return family[0] if family else None diff --git a/core/src/inline_core/models/flux2/requirements.py b/core/src/inline_core/models/flux2/requirements.py index 16e04f2..243a02c 100644 --- a/core/src/inline_core/models/flux2/requirements.py +++ b/core/src/inline_core/models/flux2/requirements.py @@ -76,6 +76,14 @@ "black-forest-labs/FLUX.2-klein-9B", "flux-2-klein-9b.safetensors", ), + ( + "diffusion_klein_9b_base", + "Klein 9B Base (for LoRA training)", + "diffusion_models", + # Gated too, for the same reason as the distilled 9B above: no ungated mirror has it. + "black-forest-labs/FLUX.2-klein-base-9B", + "flux-2-klein-base-9b.safetensors", + ), ( "text_encoder_qwen3_8b", "Qwen3-8B text encoder (for Klein 9B)", @@ -159,6 +167,13 @@ def _identify_gguf(path: Path) -> V.Flux2Variant | None: return V.get(f"klein-{size}-base" if is_base else f"klein-{size}") +#: Key fragments that rule a checkpoint out as a plain text-only Qwen3 encoder, each for a model +#: that really does sit in ``text_encoders/`` beside one: Qwen3-VL (Krea 2), then T5-XXL and CLIP-L +#: (FLUX.1's two). T5 is the dangerous one - its ``encoder.embed_tokens.weight`` is 32128 x 4096, +#: and 4096 is exactly the width klein 9B matches on. +_NOT_A_QWEN3_ENCODER = (".visual.", ".vision_", ".language_model.", "encoder.block.", "text_model.") + + def _encoder_width(path: Path) -> int | None: """The hidden width of a **plain** text-only Qwen3 checkpoint, or None if it is anything else. @@ -166,7 +181,8 @@ def _encoder_width(path: Path) -> int | None: and which sits in the same ``text_encoders/`` folder) share an embedding matrix of exactly 151936 x 2560, so matching on width alone silently loaded the vision-language model into FLUX.2's text-only encoder and rendered structured noise. A multimodal checkpoint carries a - vision tower and nests its text stack under ``language_model``; both are rejected here. + vision tower and nests its text stack under ``language_model``; both are rejected here, as are + the two encoders FLUX.1 brings to the same folder (see ``_NOT_A_QWEN3_ENCODER``). """ if path.is_dir(): return _folder_encoder_width(path) @@ -178,7 +194,7 @@ def _encoder_width(path: Path) -> int | None: keys = CheckpointReader(path).shapes() except Exception: # noqa: BLE001 - an unreadable file simply does not match return None - if any(".visual." in k or ".vision_" in k or ".language_model." in k for k in keys): + if any(fragment in k for k in keys for fragment in _NOT_A_QWEN3_ENCODER): return None for key, shape in keys.items(): if key.endswith("embed_tokens.weight") and len(shape) == 2: diff --git a/core/src/inline_core/models/loaders.py b/core/src/inline_core/models/loaders.py index c0f5370..ba7f3b8 100644 --- a/core/src/inline_core/models/loaders.py +++ b/core/src/inline_core/models/loaders.py @@ -181,6 +181,33 @@ class ArchSpec: ), ) +#: FLUX.1's own repos are gated - dev *and* schnell - so the small assets come from the ungated +#: Apache-2.0 Flex.1-alpha, which is built on FLUX.1's VAE and both of its text encoders (verified: +#: scaling_factor 0.3611, shift_factor 0.1159, CLIPTextModel 768, T5EncoderModel d_model 4096). +#: Freepik/flux.1-lite-8B carries the same layout if that repo ever disappears. One spec serves the +#: whole family: dev, schnell, Kontext, Krea and the Fill/Control builds share every one of these. +#: +#: The transformer config is NOT fetched - it is derived from the picked checkpoint (see +#: flux1/variants.py). Neither is the scheduler: Flex's is schnell-flavoured and would silently +#: give dev the wrong shift, so FLUX1_SCHEDULER_CONFIG is a code constant below. +_FLUX1 = ArchSpec( + key="flux1", + assets_repo="ostris/Flex.1-alpha", + asset_files=( + AssetFile("vae/config.json"), + AssetFile("text_encoder/config.json"), + AssetFile("text_encoder_2/config.json"), + # CLIP's slow tokenizer builds from vocab + merges, so no tokenizer.json is needed here. + AssetFile("tokenizer/vocab.json"), + AssetFile("tokenizer/merges.txt"), + AssetFile("tokenizer/tokenizer_config.json"), + AssetFile("tokenizer/special_tokens_map.json"), + AssetFile("tokenizer_2/tokenizer.json"), + AssetFile("tokenizer_2/tokenizer_config.json"), + AssetFile("tokenizer_2/special_tokens_map.json"), + ), +) + #: MiniMax H3's conditioner. The repo lays its encoder out under `FL2VA/text_encoder/`, so each #: file is re-homed to the `text_encoder/` subfolder the staged dir loads from. _MINIMAX_H3 = ArchSpec( @@ -207,6 +234,7 @@ class ArchSpec: _FLUX2_KLEIN_4B.key: _FLUX2_KLEIN_4B, _FLUX2_KLEIN_9B.key: _FLUX2_KLEIN_9B, _FLUX2_DEV.key: _FLUX2_DEV, + _FLUX1.key: _FLUX1, _MINIMAX_H3.key: _MINIMAX_H3, } @@ -314,7 +342,7 @@ def _link_or_copy(src: Path, dst: Path) -> None: shutil.copy2(src, dst) -def _staged_encoder_dir(arch: str, file: str) -> Path: +def _staged_encoder_dir(arch: str, file: str, subfolder: str = "text_encoder") -> Path: """A tiny engine-owned dir transformers can load the text encoder from as a normal model: the bundled config next to the user's weights file linked in as ``model.safetensors``. @@ -324,16 +352,26 @@ def _staged_encoder_dir(arch: str, file: str) -> Path: stays ≈ one tensor. Idempotent via a ``.complete`` marker; keyed by the weights path so a different file stages afresh.""" root = ensure_assets(arch) - te_config = root / "text_encoder" - digest = hashlib.sha1(str(file).encode()).hexdigest()[:16] + te_config = root / subfolder + # The subfolder joins the key only when it is not the default, so an already-staged encoder is + # not re-staged - that fallback is a full copy of several GB where symlinks are unavailable. + keyed = str(file) if subfolder == "text_encoder" else f"{subfolder}\x00{file}" + digest = hashlib.sha1(keyed.encode()).hexdigest()[:16] stage = assets_root(arch) / "te_stage" / digest marker = stage / ".complete" - if marker.is_file(): + # The marker alone is not enough: these are symlinks, and a models dir that moved leaves them + # dangling inside a stage that still says it is complete. `exists` follows the link. + weights = stage / "model.safetensors" + if marker.is_file() and weights.exists(): return stage with _ASSETS_LOCK: - if marker.is_file(): + if marker.is_file() and weights.exists(): return stage stage.mkdir(parents=True, exist_ok=True) + # A broken link is not overwritten by _link_or_copy, which treats any entry as done. + for stale in stage.iterdir(): + if stale.is_symlink() and not stale.exists(): + stale.unlink() for name in ("config.json", "generation_config.json"): src = te_config / name if src.is_file(): @@ -736,6 +774,102 @@ def load_scheduler(arch: str) -> Any: return FlowMatchEulerDiscreteScheduler() +# --- FLUX.1 ------------------------------------------------------------------------------------ + +#: dev's own scheduler settings, as a code constant rather than a fetched asset: the ungated repo +#: the rest of FLUX.1's assets come from is a schnell derivative whose scheduler does no dynamic +#: shifting, and taking it would silently give dev the wrong shift. +FLUX1_SCHEDULER_CONFIG = { + "base_image_seq_len": 256, + "max_image_seq_len": 4096, + "base_shift": 0.5, + "max_shift": 1.15, + "num_train_timesteps": 1000, + "shift": 3.0, + "use_dynamic_shifting": True, +} + + +def load_flux1_scheduler() -> Any: + """FLUX.1's flow-match scheduler, from the constant above (config-only - never downloads).""" + from diffusers import FlowMatchEulerDiscreteScheduler + + return FlowMatchEulerDiscreteScheduler.from_config(FLUX1_SCHEDULER_CONFIG) + + +def load_flux1_text_encoders( + arch: str, + clip_file: str, + t5_file: str, + dtype: Any, + quant: Quantization = Quantization.NONE, + device: str | None = None, +) -> tuple[Any, Any, Any, Any]: + """``(clip, clip_tokenizer, t5, t5_tokenizer)`` - FLUX.1 conditions on two encoders. + + Both load from a staging directory rather than a pre-loaded state dict, for the reason + ``_staged_encoder_dir`` gives: T5-XXL is ~10 GB and a state dict would materialize all of it in + host RAM. The published single files are already exactly what transformers expects - CLIP's keys + are all ``text_model.*``, T5's are ``shared`` plus ``encoder.block.*`` with no decoder - so + there is no key conversion here. + + ``quant`` reaches T5 alone. CLIP-L is 246 MB, so quantizing it saves nothing measurable and only + risks the pooled vector the timestep embedder mixes in. Each half is cached under its own key, + so replacing one file does not reload the other. + """ + from transformers import CLIPTokenizer, T5TokenizerFast + + root = ensure_assets(arch) + + def build_clip() -> Any: + from transformers import CLIPTextModel + + return CLIPTextModel.from_pretrained( + str(_staged_encoder_dir(arch, clip_file)), + torch_dtype=dtype, + low_cpu_mem_usage=True, + device_map={"": device} if device else None, + local_files_only=True, + ) + + def build_t5() -> Any: + from transformers import T5EncoderModel + + # An fp8 repack (t5xxl_fp8_e4m3fn_scaled) already is its resident size; quantizing on top + # is a hard error, and transformers drops its scales as unexpected keys. + carried = checkpoint.prequantized_kind(t5_file) + if carried: + logger.info( + "%s is a %s text encoder; loading it as-is rather than quantizing on top.", + Path(t5_file).name, carried, + ) + return T5EncoderModel.from_pretrained( + str(_staged_encoder_dir(arch, t5_file, "text_encoder_2")), + torch_dtype=dtype, + low_cpu_mem_usage=True, + device_map={"": device} if device else None, + local_files_only=True, + quantization_config=_quant_config( + Quantization.NONE if carried else quant, framework="transformers" + ), + ) + + clip = _cached( + (arch, "clip", clip_file, _dtype_key(dtype), Quantization.NONE.value, _device_key(device)), + build_clip, + ) + t5 = _cached( + (arch, "t5", t5_file, _dtype_key(dtype), quant.value, _device_key(device)), build_t5 + ) + # Explicit classes, not AutoTokenizer: FluxPipeline's own type hints name exactly these two. + return ( + clip, + CLIPTokenizer.from_pretrained(str(root / "tokenizer"), local_files_only=True), + t5, + T5TokenizerFast.from_pretrained(str(root / "tokenizer_2"), local_files_only=True), + ) + + # --- Krea 2 ------------------------------------------------------------------------------------ #: The reference "single_mmdit_large_wide" geometry, shared by RAW and Turbo. These are also @@ -1075,14 +1209,18 @@ def assemble_zimage_pipeline( # --- FLUX.2 -------------------------------------------------------------------------------------- -def _flux2_config_dir(arch: str, file: str, config: dict[str, Any]) -> Path: +def _transformer_config_dir( + arch: str, file: str, config: dict[str, Any], class_name: str +) -> Path: """A tiny staging dir holding the transformer config derived from ``file``, so diffusers' ``from_single_file`` has a local config to read. The config is derived from the checkpoint's own tensor shapes rather than fetched, which is what lets one node load klein 4B, klein 9B, dev and any later build. Keyed by the weights path, so a different checkpoint stages its own.""" - digest = hashlib.sha1(f"{file}:{sorted(config.items())!r}".encode()).hexdigest()[:16] + digest = hashlib.sha1( + f"{class_name}:{file}:{sorted(config.items())!r}".encode() + ).hexdigest()[:16] stage = assets_root(arch) / "transformer_config" / digest marker = stage / "config.json" if marker.is_file(): @@ -1093,7 +1231,7 @@ def _flux2_config_dir(arch: str, file: str, config: dict[str, Any]) -> Path: import json stage.mkdir(parents=True, exist_ok=True) - payload = {"_class_name": "Flux2Transformer2DModel", **config} + payload = {"_class_name": class_name, **config} marker.write_text(json.dumps(payload, indent=2)) return stage @@ -1118,6 +1256,203 @@ def _gguf_config(dtype: Any) -> Any: ) from error +def load_flux1_transformer( + arch: str, + file: str, + config: dict[str, Any], + dtype: Any, + quant: Quantization = Quantization.NONE, + device: str | None = None, + loras: tuple[LoraRef, ...] = (), +) -> Any: + """The FLUX.1 transformer from a single ``.safetensors``, a ``.gguf``, or a diffusers folder. + + ``config`` is the geometry derived from the checkpoint (see ``flux1/variants.py``), so dev, + schnell, Kontext and the Fill/Control builds all load through here without a per-build config. + """ + + def build() -> Any: + from diffusers import FluxTransformer2DModel + + if Path(file).is_dir(): + # A folder carries its own config and, when prequantized, already-reduced shards; + # from_pretrained streams them rather than materializing anything at full size. + return FluxTransformer2DModel.from_pretrained( + file, + torch_dtype=dtype, + low_cpu_mem_usage=True, + local_files_only=True, + **({"device_map": {"": device}} if device else {}), + ) + + root = _transformer_config_dir(arch, file, config, "FluxTransformer2DModel") + kwargs: dict[str, Any] = {} + if _is_gguf(file): + kwargs["quantization_config"] = _gguf_config(dtype) + # NF4 loads to the CPU on purpose: bitsandbytes only quantizes on the move to CUDA, so + # streaming straight to the card materializes the whole bf16 base first. Measured on dev: + # 23.81GB peak that way against 6.23GB this way, for the same 70s. + load_device = None if quant is Quantization.NF4 else device + try: + model = FluxTransformer2DModel.from_single_file( + file, + config=str(root), + torch_dtype=dtype, + low_cpu_mem_usage=True, + device=load_device, + local_files_only=True, + **kwargs, + ) + except Exception as error: + # Every checkpoint gets attempted - a user's own file is never refused on suspicion. + # Explain afterwards, because diffusers' converter fails deep inside itself without + # naming the file (see flux1/variants.single_file_blocker). + from .flux1 import variants as flux1_variants + + blocked = flux1_variants.single_file_blocker(file) + if blocked is None: + raise + raise ComponentError( + f"{Path(file).name} did not load: {blocked}, and diffusers splits those as if they " + "were weights. A full-precision or .gguf build of the same model will work, and " + "OpenChar quantizes on load to fit your card either way." + ) from error + if _is_gguf(file): + if loras: + raise ComponentError( + "LoRAs cannot be fused into a .gguf checkpoint. Use a .safetensors model, or " + "remove the LoRA." + ) + return model + # Fuse before quantizing: the fuse adds a full-precision delta into each weight, which a + # quantized weight is not a plain tensor to accept. Same ordering as the FLUX.2 path. + if loras: + from .lora import fuse_loras + + fuse_loras(model, loras) + if quant is Quantization.NF4: + _swap_to_4bit(model) + if device: + model.to(device) # bitsandbytes quantizes each weight during this move + else: + # from_single_file ignores quantization_config, so int8 is applied after the load. + _quantize_in_place(model, quant) + return model + + key = ( + arch, "diffusion", file, _dtype_key(dtype), quant.value, _device_key(device), + *lora_cache_key(loras), + ) + return _cached(key, build) + + +def assemble_flux1_pipeline( + *, + arch: str, + img2img: bool, + diffusion_file: str, + config: dict[str, Any], + vae_file: str, + text_encoder_file: str, + clip_file: str, + dtype: Any, + quant: Quantization = Quantization.NONE, + vae_dtype: Any = None, + device: str | None = None, + loras: tuple[LoraRef, ...] = (), + cancel_check: Callable[[], None] | None = None, +) -> Any: + """A whole FLUX.1 pipeline from the user's single files. + + Encoders first and the transformer last, so peak memory is one component rather than the sum - + the transformer is by far the largest and nothing else is still being read when it lands. + """ + from diffusers import FluxImg2ImgPipeline, FluxPipeline + + if cancel_check is not None: + cancel_check() + vae = load_vae(arch, vae_file, dtype if vae_dtype is None else vae_dtype, device=device) + _release_transient() + clip, clip_tok, t5, t5_tok = load_flux1_text_encoders( + arch, clip_file, text_encoder_file, dtype, quant, device=device + ) + _release_transient() + if cancel_check is not None: + cancel_check() + transformer = load_flux1_transformer( + arch, diffusion_file, config, dtype, quant, device=device, loras=loras + ) + _release_transient() + cls = FluxImg2ImgPipeline if img2img else FluxPipeline + return cls( + scheduler=load_flux1_scheduler(), vae=vae, text_encoder=clip, tokenizer=clip_tok, + text_encoder_2=t5, tokenizer_2=t5_tok, transformer=transformer, + ) + + +def assemble_flux1_encoders( + *, + arch: str, + img2img: bool, + vae_file: str, + text_encoder_file: str, + clip_file: str, + dtype: Any, + quant: Quantization = Quantization.NONE, + vae_dtype: Any = None, + device: str | None = None, +) -> Any: + """A **transformer-less** FLUX.1 pipeline: VAE + both text encoders + scheduler only. + + dev is 24 GB of transformer against a 10 GB T5, so on a card that cannot hold both the prompt is + encoded through this first, the encoders are freed, and only then does the transformer load. + Same trick the trainer uses to precache before its base loads. + """ + from diffusers import FluxImg2ImgPipeline, FluxPipeline + + vae = load_vae(arch, vae_file, dtype if vae_dtype is None else vae_dtype, device=device) + _release_transient() + clip, clip_tok, t5, t5_tok = load_flux1_text_encoders( + arch, clip_file, text_encoder_file, dtype, quant, device=device + ) + _release_transient() + cls = FluxImg2ImgPipeline if img2img else FluxPipeline + return cls( + scheduler=load_flux1_scheduler(), vae=vae, text_encoder=clip, tokenizer=clip_tok, + text_encoder_2=t5, tokenizer_2=t5_tok, transformer=None, + ) + + +def attach_flux1_transformer( + pipe: Any, + *, + arch: str, + diffusion_file: str, + config: dict[str, Any], + vae_file: str, + dtype: Any, + quant: Quantization = Quantization.NONE, + device: str | None = None, + loras: tuple[LoraRef, ...] = (), +) -> Any: + """Free **both** text encoders, then load the transformer into the VRAM they were holding. + + The VAE is kept: it is small and the decode still needs it. Dropped rather than moved to the + CPU, because a host that needs this staging cannot take a 10 GB encoder in RAM either. + """ + pipe.text_encoder = None + pipe.tokenizer = None + pipe.text_encoder_2 = None + pipe.tokenizer_2 = None + unload_components(keep_files={vae_file}) + _release_transient() + pipe.transformer = load_flux1_transformer( + arch, diffusion_file, config, dtype, quant, device=device, loras=loras + ) + _release_transient() + return pipe + + def load_flux2_transformer( arch: str, file: str, @@ -1148,17 +1483,21 @@ def build() -> Any: **({"device_map": {"": device}} if device else {}), ) - root = _flux2_config_dir(arch, file, config) + root = _transformer_config_dir(arch, file, config, "Flux2Transformer2DModel") kwargs: dict[str, Any] = {} if _is_gguf(file): kwargs["quantization_config"] = _gguf_config(dtype) + # NF4 loads to the CPU on purpose: bitsandbytes only quantizes on the move to CUDA, so + # streaming straight to the card materializes the whole bf16 base first. Measured on + # klein 4B: 7.77GB peak that way against 2.16GB this way, and 85s against 22s. + load_device = None if quant is Quantization.NF4 else device try: model = Flux2Transformer2DModel.from_single_file( file, config=str(root), torch_dtype=dtype, low_cpu_mem_usage=True, - device=device, + device=load_device, local_files_only=True, **kwargs, ) diff --git a/core/src/inline_core/models/pipeline_runtime.py b/core/src/inline_core/models/pipeline_runtime.py index 3b5d24e..cfaab3b 100644 --- a/core/src/inline_core/models/pipeline_runtime.py +++ b/core/src/inline_core/models/pipeline_runtime.py @@ -286,22 +286,29 @@ def supports_prompt_embeds(pipe: Any) -> bool: @contextmanager -def text_encoder_detached(pipe: Any, active: bool) -> Iterator[None]: - """Temporarily remove the text encoder from the pipeline for the denoise, then restore it. +def text_encoder_detached( + pipe: Any, active: bool, encoders: tuple[str, ...] = ("text_encoder",) +) -> Iterator[None]: + """Temporarily remove the text encoder(s) from the pipeline for the denoise, then restore them. diffusers infers the execution device from *some* registered module and iterates a set, so with the encoder parked on the CPU the pick is non-deterministic and can build latents on the CPU while the generator is on CUDA. Detaching leaves only CUDA modules. No-op on the raw path. + + ``encoders`` names every attribute to detach, because a pipeline can have more than one and + leaving either behind reinstates the non-determinism this exists to remove. """ if not active: yield return - saved = getattr(pipe, "text_encoder", None) - pipe.text_encoder = None + saved = {name: getattr(pipe, name, None) for name in encoders} + for name in encoders: + setattr(pipe, name, None) try: yield finally: - pipe.text_encoder = saved + for name, value in saved.items(): + setattr(pipe, name, value) def split_blocks(blocks: Any, *, through: str) -> tuple[Any, Any]: @@ -363,6 +370,7 @@ def encoded_prompt_kwargs( *, encode: Callable[[str], dict[str, Any]], fallback: Callable[[], dict[str, Any]], + encoders: tuple[str, ...] = ("text_encoder",), ) -> dict[str, Any]: """Pre-encode the prompt on the GPU, park the text encoder on the CPU, and return the pipeline call kwargs, so the encoder's GB go to the denoise instead of idling on the card. @@ -371,27 +379,35 @@ def encoded_prompt_kwargs( dequantizes per-op, while a CPU encode would dequantize the whole encoder into host RAM, and without no_grad the full activation graph is retained. Parking is a plain tensor copy. Any failure falls back to the raw-prompt path, so this optimization can never break a run.""" - text_encoder = getattr(pipe, "text_encoder", None) - if not is_resident(policy) or text_encoder is None or not supports_prompt_embeds(pipe): + # Every named encoder, in order. FLUX.1 has two, and parking only the first would move its + # 246MB CLIP while leaving a 9.8GB T5 on the card - the opposite of the point. + staged = [(name, getattr(pipe, name, None)) for name in encoders] + staged = [(name, module) for name, module in staged if module is not None] + if not is_resident(policy) or not staged or not supports_prompt_embeds(pipe): return fallback() + # The park decision is taken on the largest encoder, which is the one that crowds the denoise. + text_encoder = max((m for _n, m in staged), key=module_bytes) device = str(policy.placement("denoiser").device) try: logger.info("Encoding prompt on %s (no_grad) | host RAM %.1fGB", device, host_ram_gb()) - text_encoder.to(device) + for _name, module in staged: + module.to(device) with torch.no_grad(): kwargs = encode(device) park = should_park_encoder(policy, text_encoder) if park: - # torchao's .to() round-trip on quantized weights is unreliable, so only the encoder - # moves. - text_encoder.to("cpu") + # torchao's .to() round-trip on quantized weights is unreliable, so only the encoders + # move. + for _name, module in staged: + module.to("cpu") free_vram() except Exception as error: # noqa: BLE001 - an optimization must never break generation logger.warning( "Text-encoder GPU encode failed (%s); denoising with the encoder resident.", error ) - try_call(text_encoder.to, device) + for _name, module in staged: + try_call(module.to, device) return fallback() if park: diff --git a/core/src/inline_core/models/trainingreqs.py b/core/src/inline_core/models/trainingreqs.py index 47de80c..f961e22 100644 --- a/core/src/inline_core/models/trainingreqs.py +++ b/core/src/inline_core/models/trainingreqs.py @@ -26,6 +26,19 @@ def _hyperparams(params: dict[str, Any] | None) -> dict[str, Any]: return raw if isinstance(raw, dict) else {} +#: FLUX.2 base mode -> (variant, the optional rows a training run promotes, the required rows they +#: stand in for). Mirrors ``training/models._FLUX2_BASES``; ``raw`` stays 4B so saved runs resolve +#: to the checkpoint they were trained against. +_FLUX2_TRAINING_BASES: dict[str, tuple[str, tuple[str, ...], tuple[str, ...]]] = { + "raw": ("klein-4b-base", ("diffusion_klein_4b_base",), ("diffusion",)), + "raw_9b": ( + "klein-9b-base", + ("diffusion_klein_9b_base", "text_encoder_qwen3_8b"), + ("diffusion", "text_encoder"), + ), +} + + def base_components(arch: str, base_mode: str) -> list[ModelComponent]: """The required components for one architecture's training base, newest-arch-first by name.""" if arch == "krea2": @@ -38,15 +51,27 @@ def base_components(arch: str, base_mode: str) -> list[ModelComponent]: from .zimage import requirements as reqs return _required(reqs.zimage_requirements()) + _adapter(arch, base_mode) + if arch == "flux1": + from .flux1 import requirements as reqs + + # No row swap, unlike FLUX.2 below: dev is guidance-distilled rather than step-distilled, so + # the checkpoint the popup already lists as required *is* the training base. + return _required(reqs.flux1_requirements()) if arch == "flux2": from .flux2 import requirements as reqs # The distilled build is what the generation node wants and what the trainer refuses, so - # the Base checkpoint FLUX.2 lists as an optional extra is the required one here. - rows = reqs.flux2_requirements() - base = next((c for c in rows if c.id == "diffusion_klein_4b_base"), None) - keep = [c for c in _required(rows) if c.id != "diffusion"] - return ([replace(base, optional=False)] if base else _required(rows)) + keep + # the Base checkpoint the popup lists as an optional extra is the required one here. 9B + # brings its own encoder: the required row points at the 4B one whatever the base. + variant, promote, replaced = _FLUX2_TRAINING_BASES.get( + base_mode, _FLUX2_TRAINING_BASES["raw"] + ) + rows = reqs.flux2_requirements({"variant": variant}) + by_id = {c.id: c for c in rows} + promoted = [replace(by_id[row], optional=False) for row in promote if row in by_id] + if not promoted: + return _required(rows) + return promoted + [c for c in _required(rows) if c.id not in replaced] if arch == "ltx-2-5": from .ltx25 import requirements as reqs diff --git a/core/src/inline_core/server/bootstrap.py b/core/src/inline_core/server/bootstrap.py index e3bd6a9..29f7588 100644 --- a/core/src/inline_core/server/bootstrap.py +++ b/core/src/inline_core/server/bootstrap.py @@ -97,6 +97,15 @@ def _register_builtins( registered.append(FLUX2.type) except ImportError: pass + try: + from ..models.flux1.provider import Flux1Provider + from ..models.flux1.runner import FLUX1, register_flux1 + + register_flux1(registry, store, policy) + requirements.register(FLUX1.type, Flux1Provider()) + registered.append(FLUX1.type) + except ImportError: + pass try: from ..models.minimaxh3.provider import MiniMaxH3Provider from ..models.minimaxh3.runner import VARIANTS, register_minimax_h3 diff --git a/core/src/inline_core/training/arch.py b/core/src/inline_core/training/arch.py index 35c06c8..35ad080 100644 --- a/core/src/inline_core/training/arch.py +++ b/core/src/inline_core/training/arch.py @@ -7,7 +7,7 @@ All of them are rectified flow, but with **opposite conventions**, which is exactly the kind of detail a test should pin: Z-Image and MiniMax H3 predict ``clean - noise`` at timestep -``1 - sigma``, while Krea 2, FLUX.2 and LTX-2.5 predict ``noise - clean`` at timestep ``sigma``. +``1 - sigma``, while Krea 2, FLUX.1, FLUX.2 and LTX-2.5 predict ``noise - clean`` at ``sigma``. LTX-2.5 is the one that looks like a third convention and is not. Its published model is an ``X0Model`` returning a denoised latent, but that is a weightless wrapper over a velocity model, and @@ -23,6 +23,7 @@ Z_IMAGE = "z-image" KREA2 = "krea2" +FLUX1 = "flux1" FLUX2 = "flux2" MINIMAX_H3 = "minimax-h3" LTX25 = "ltx-2-5" @@ -60,6 +61,37 @@ ] +#: FLUX.1: every Linear in the MMDiT stack, confirmed against FluxTransformer2DModel's own +#: named_modules (19 double + 38 single blocks on dev). The single blocks are ``pre_only``, so they +#: carry no ``attn.to_out`` at all - the ModuleList/Linear suffix clash FLUX.2 has below does not +#: recur here, and ``to_out.0`` reaches the double blocks alone. +#: +#: ``proj_out`` matches **twice**, deliberately: PEFT matches by module-name suffix, so it reaches +#: the 38 single blocks' output projections and, by exact name, the model's own final ``proj_out``. +#: Both are plain Linears, so nothing breaks, and the extra one is a 3072x64 tail. Dropping it to +#: avoid that would also drop the single-block projections, which every published FLUX.1 LoRA +#: carries. The AdaLN modulation Linears (``norm1.linear``, ``norm.linear``, ``norm_out.linear``) +#: stay out, for the reason ``_MINIMAX_H3_TARGETS`` gives about ``adaln_proj``. +_FLUX1_TARGETS = [ + "to_q", + "to_k", + "to_v", + "to_out.0", + "add_q_proj", + "add_k_proj", + "add_v_proj", + "to_add_out", + "ff.net.0.proj", + "ff.net.2", + "ff_context.net.0.proj", + "ff_context.net.2", + "proj_mlp", + "proj_out", + "x_embedder", + "context_embedder", +] + + #: FLUX.2: every Linear in the MMDiT stack, confirmed against Flux2Transformer2DModel's own #: named_parameters (5 double + 20 single blocks on klein 4B, 169 tensors, no biases anywhere). #: Double blocks carry separate image/text streams (``add_*`` is the text side); single blocks fuse @@ -251,6 +283,66 @@ def _krea2_forward(transformer: Any, noisy: Any, timestep: Any, item: dict[str, +# --- FLUX.1 --------------------------------------------------------------------------------------- + + +#: What the guidance embedder is trained at. dev is guidance-distilled, so the embedder is part of +#: the model rather than a sampler setting, and training it at the 3.5 dev generates with teaches +#: the LoRA to change how guidance itself behaves. Every reference trainer pins it to 1. +_FLUX1_TRAIN_GUIDANCE = 1.0 + + +def _flux1_sigma(device: Any, shift: float) -> Any: + import torch + + # Logit-normal, matching the reference trainers. FLUX.1's shift is resolution dependent and + # computed at inference (``calculate_shift``), so it is not baked into the training + # distribution - the same reasoning as Krea 2 and FLUX.2. + del shift + return torch.sigmoid(torch.randn((), device=device)) + + +def _flux1_forward(transformer: Any, noisy: Any, timestep: Any, item: dict[str, Any]) -> Any: + """One prediction from FluxTransformer2DModel, mirroring FluxPipeline's denoise call. + + The precached latent is the VAE's own 16-channel H/8 output, so packing is the pipeline's 2x2 + fold rather than FLUX.2's plain flatten. ``timestep`` is the raw sigma: the pipeline passes + ``timestep / 1000`` where its own timesteps are ``sigma * 1000``. + + Conditioning is two tensors, which no other arch here has: T5-XXL's sequence and CLIP-L's pooled + vector. Guidance is passed unconditionally, which is safe only because ``_flux1_base_file`` + refuses schnell - the model picks a 2- or 3-argument embedder off its own ``guidance_embeds``, + so a mismatch either way is a TypeError, and ``arch.forward`` cannot read that config reliably + because under DDP it is handed the wrapper rather than the model. The model scales by 1000 + itself, so this is the raw value. + """ + import torch + from diffusers import FluxPipeline as P + + channels, height, width = noisy.shape + packed = P._pack_latents(noisy.unsqueeze(0), 1, channels, height, width) + img_ids = P._prepare_latent_image_ids(1, height // 2, width // 2, noisy.device, noisy.dtype) + embed = item["embed"] + embed = embed.unsqueeze(0) if embed.dim() == 2 else embed + pooled = item["pooled"] + pooled = pooled.unsqueeze(0) if pooled.dim() == 1 else pooled + # 2-D on purpose: a 3-D txt_ids is deprecated and silently indexed back down to this. + txt_ids = torch.zeros(embed.shape[1], 3, device=noisy.device, dtype=noisy.dtype) + + out = transformer( + hidden_states=packed, + encoder_hidden_states=embed, + pooled_projections=pooled, + timestep=timestep.reshape(1), + img_ids=img_ids, + txt_ids=txt_ids, + guidance=torch.full((1,), _FLUX1_TRAIN_GUIDANCE, device=noisy.device, dtype=torch.float32), + return_dict=False, + )[0] + # Pixel dimensions and the VAE's 8x factor, exactly what the pipeline passes. + return P._unpack_latents(out, height * 8, width * 8, 8)[0] + + # --- FLUX.2 --------------------------------------------------------------------------------------- @@ -442,6 +534,17 @@ def _ltx25_forward(transformer: Any, noisy: Any, timestep: Any, item: dict[str, target=lambda clean, noise: noise - clean, forward=_krea2_forward, ), + FLUX1: TrainingArch( + key=FLUX1, + target_modules=_FLUX1_TARGETS, + sigma=_flux1_sigma, + # Rectified flow, Krea 2's convention: x_t = (1 - sigma) * clean + sigma * noise, so + # d x_t / d sigma is noise - clean, and the model is called at the sigma itself. Verified + # against FluxPipeline, whose scheduler timesteps are sigma * 1000 and which divides back. + timestep=lambda sigma: sigma, + target=lambda clean, noise: noise - clean, + forward=_flux1_forward, + ), FLUX2: TrainingArch( key=FLUX2, target_modules=_FLUX2_TARGETS, diff --git a/core/src/inline_core/training/cache.py b/core/src/inline_core/training/cache.py index 6ef89e7..687950c 100644 --- a/core/src/inline_core/training/cache.py +++ b/core/src/inline_core/training/cache.py @@ -30,6 +30,7 @@ def build( dtype: Any, resolution: int, *, + base_mode: str = "raw", flip: bool = False, dropout: float = 0.0, clip_frames: int = 1, @@ -54,6 +55,13 @@ def build( # A motion run caches a different shape of item, so it must not read a clip run's latents. "training_mode": training_mode, } + # Only for a family whose members encode differently: klein 4B and 9B caption through + # different-width encoders, so one's cache is the wrong shape for the other. Added rather than + # always present because every key here changes the fingerprint, and an arch whose base cannot + # vary would re-encode a valid cache for nothing - 19 minutes of it, on H3. + base = models.loader_arch(arch, models_dir, base_mode) + if base != arch: + settings["base"] = base say = on_status or (lambda _text: None) key = None if cache_dir: @@ -67,8 +75,8 @@ def build( items, unconditional, shift = _encode( dataset_dir, models_dir, arch, device, dtype, resolution, - flip=flip, dropout=dropout, clip_frames=clip_frames, clip_window=clip_window, - training_mode=training_mode, on_status=on_status, + base_mode=base_mode, flip=flip, dropout=dropout, clip_frames=clip_frames, + clip_window=clip_window, training_mode=training_mode, on_status=on_status, ) if cache_dir and key is not None: store.save(Path(cache_dir), key, items, unconditional, shift) @@ -83,6 +91,7 @@ def _encode( dtype: Any, resolution: int, *, + base_mode: str = "raw", flip: bool = False, dropout: float = 0.0, clip_frames: int = 1, @@ -108,7 +117,9 @@ def _encode( ) return items, unconditional, _VIDEO_SHIFTS[archs.LTX25] - encoders = models.load_encoders(models_dir, arch, device, dtype) + # Before the encoders, not after: this is the cheap check and they are the expensive load. + ds.check_usable(dataset_dir, arch) + encoders = models.load_encoders(models_dir, arch, device, dtype, base_mode) items = ds.precache(dataset_dir, encoders, arch, device, dtype, resolution, flip=flip) unconditional = ds.precache_empty(encoders, arch, device) if dropout > 0 else None shift = float(encoders.scheduler.config.get("shift", 1.0) or 1.0) diff --git a/core/src/inline_core/training/dataset.py b/core/src/inline_core/training/dataset.py index 55e3db9..7ee7ac2 100644 --- a/core/src/inline_core/training/dataset.py +++ b/core/src/inline_core/training/dataset.py @@ -151,6 +151,30 @@ def _to_tensor(image: Any, resolution: int, flip: bool = False) -> Any: return torch.from_numpy(arr).permute(2, 0, 1) # CHW +def check_usable(dataset_dir: str, arch: str) -> None: + """Fail before the encoders load when the dataset holds nothing this arch can train on. + + Worth its own pass because the alternative is a ten-gigabyte text encoder loading first and the + run then reporting an empty dataset - true, but not the reason, and not what to do about it. + """ + from . import arch as archs + + root = Path(dataset_dir) + if not root.is_dir(): + raise RuntimeError(f"The dataset folder {dataset_dir} does not exist.") + images = len(_pairs(root, _IMAGE_SUFFIXES)) + clips = len(_pairs(root, _VIDEO_SUFFIXES)) + if images or (clips and archs.get(arch).clip is not None): + return + if clips: + raise RuntimeError( + f"This dataset holds {clips} video clip{'s' if clips != 1 else ''} and no images, but " + f"{arch} trains on still images. Train it on a dataset of images, or pick MiniMax H3 " + "or LTX-2.5, which train on clips." + ) + raise RuntimeError(f"The dataset at {dataset_dir} has no images or clips in it.") + + def precache( dataset_dir: str, components: Any, @@ -263,6 +287,40 @@ def _krea2_caption(components: Any, caption: str, device: str) -> dict[str, Any] # --- FLUX.2 ------------------------------------------------------------------------------------- +def _flux1_latent(vae: Any, pixels: Any) -> Any: + """FLUX.1 shares Z-Image's scalar VAE normalization, and adds one constraint on the grid. + + ``_pack_latents`` folds the H/8 latent 2x2, and ``_unpack_latents`` floors on the way back, so + an odd latent grid silently loses a row rather than raising. Refused here, at the first step of + the precache, instead of training against latents that are quietly the wrong size. + """ + height, width = pixels.shape[-2], pixels.shape[-1] + if (height // 8) % 2 or (width // 8) % 2: + raise RuntimeError( + f"FLUX.1 trains at a resolution that is a multiple of 16, not {height}x{width}: the " + "latent grid is folded 2x2 after the VAE's 8x downscale." + ) + return _zimage_latent(vae, pixels) + + +def _flux1_caption(components: Any, caption: str, device: str) -> dict[str, Any]: + """Caption -> conditioning, straight from ``FluxPipeline.encode_prompt``. + + Two tensors, not one, and both are needed: T5-XXL gives the (512, 4096) sequence and CLIP-L the + (768,) pooled vector the timestep embedder mixes in. Routed through the pipeline because the + details are easy to get subtly wrong and impossible to notice - CLIP truncates at 77 tokens + while T5 pads to 512, and the pooled vector is the EOS position rather than a mean. + """ + import torch + + with torch.no_grad(): + embeds, pooled, _ids = components.pipeline.encode_prompt( + prompt=caption or "", prompt_2=None, device=torch.device(device), + max_sequence_length=512, + ) + return {"embed": embeds[0], "pooled": pooled[0]} + + def _flux2_latent(vae: Any, pixels: Any) -> Any: """Pixels -> the latent FLUX.2 actually denoises, matching ``_encode_vae_image``. @@ -304,6 +362,9 @@ def _flux2_caption(components: Any, caption: str, device: str) -> dict[str, Any] #: a training latent is byte-for-byte what generation would produce from the same image. _ENCODERS = { archs.KREA2: (_krea2_latent, _krea2_caption), + # FLUX.1 normalizes with the same scalar (latent - shift) * scale Z-Image does, not FLUX.2's + # running batch-norm statistics, so the latent encoder is shared rather than copied. + archs.FLUX1: (_flux1_latent, _flux1_caption), archs.FLUX2: (_flux2_latent, _flux2_caption), } diff --git a/core/src/inline_core/training/models.py b/core/src/inline_core/training/models.py index 37c9c37..8b7ee4c 100644 --- a/core/src/inline_core/training/models.py +++ b/core/src/inline_core/training/models.py @@ -32,6 +32,14 @@ "text_encoders": "INLINE_KREA2_TEXT_ENCODER", "adapter": "INLINE_KREA2_TRAIN_ADAPTER", }, + # No adapter entry: dev is *guidance*-distilled, not step-distilled, so it is itself the + # training base and there is nothing to de-distill first. + archs.FLUX1: { + "diffusion_models": "INLINE_FLUX1_MODEL", + "vae": "INLINE_FLUX1_VAE", + "text_encoders": "INLINE_FLUX1_TEXT_ENCODER", + "clip": "INLINE_FLUX1_CLIP", + }, archs.FLUX2: { "diffusion_models": "INLINE_FLUX2_MODEL", "vae": "INLINE_FLUX2_VAE", @@ -66,9 +74,18 @@ class Encoders: #: Krea 2 only: a transformer-less Krea2Pipeline, so caption encoding goes through diffusers' #: own prompt template and 12-layer tap rather than a copy that could drift from inference. pipeline: Any = None + #: FLUX.1 only: CLIP-L beside T5-XXL. Held here as well as on the pipeline so ``free_encoders`` + #: drops all ~10GB by name rather than relying on the pipeline being the last reference. + text_encoder_2: Any = None + tokenizer_2: Any = None + +#: A component whose name is not its folder. FLUX.1's CLIP-L is a second text encoder, so it lives +#: in ``text_encoders/`` beside T5 rather than in a folder of its own. +_CATEGORY_FOLDER = {"clip": "text_encoders"} -def _require(root: Path, arch: str, category: str) -> str: + +def _require(root: Path, arch: str, category: str, variant: Any = None) -> str: """The weight file for a component, resolved by the arch's own requirements module. Never "the first file in the folder": every architecture shares ``vae/`` and ``text_encoders/``, @@ -77,17 +94,21 @@ def _require(root: Path, arch: str, category: str) -> str: env = os.environ.get(_ENV[arch][category]) if env: return env - picked = _resolve(arch, category) + picked = _resolve(arch, category, variant) if picked: return str(picked) + folder = _CATEGORY_FOLDER.get(category, category) raise RuntimeError( - f"No {arch} {category} weight found under {root / category}. Add it there " + f"No {arch} {category} weight found under {root / folder}. Add it there " f"(or set {_ENV[arch][category]})." ) -def _resolve(arch: str, category: str) -> Any: - """The arch's own answer for a component, so training and generation pick the same file.""" +def _resolve(arch: str, category: str, variant: Any = None) -> Any: + """The arch's own answer for a component, so training and generation pick the same file. + + ``variant`` pins a family whose members need different components, so every lookup answers for + the checkpoint this run trains rather than rescanning and landing elsewhere.""" if arch == archs.MINIMAX_H3: from ..models.minimaxh3 import requirements as h3_reqs @@ -102,11 +123,25 @@ def _resolve(arch: str, category: str) -> Any: if arch == archs.FLUX2: from ..models.flux2 import requirements as flux2_reqs + # klein 4B and 9B want encoders of different widths, so the picked variant is forwarded + # rather than re-derived: resolve_text_encoder would otherwise scan for itself. + params = {"variant": variant.key} if variant is not None else None if category == "vae": - return flux2_reqs.resolve_vae(None) + return flux2_reqs.resolve_vae(params) if category == "text_encoders": - return flux2_reqs.resolve_text_encoder(None) - return flux2_reqs.resolve_diffusion(None) + return flux2_reqs.resolve_text_encoder(params) + return flux2_reqs.resolve_diffusion(params) + + if arch == archs.FLUX1: + from ..models.flux1 import requirements as flux1_reqs + + if category == "vae": + return flux1_reqs.resolve_vae(None) + if category == "text_encoders": + return flux1_reqs.resolve_text_encoder(None) + if category == "clip": + return flux1_reqs.resolve_clip(None) + return flux1_reqs.resolve_diffusion(None) if arch == archs.KREA2: from ..models.krea2 import requirements as krea2_reqs @@ -129,18 +164,33 @@ def _resolve(arch: str, category: str) -> Any: -def loader_arch(arch: str, models_dir: str | None = None) -> str: +def flux2_variant(root: Path, base_mode: str) -> Any: + """The FLUX.2 variant this run trains, resolved once from the base checkpoint it picked. + + Everything downstream keys off this rather than scanning again, because two scans over one + folder can disagree: the base file is the first *undistilled* checkpoint while + ``resolve_diffusion`` takes the first FLUX.2 checkpoint at all, so a distilled klein 4B kept + for generation would pair a 9B base with a 4B text encoder. + """ + from ..models.flux2 import variants as flux2_variants + + detected = flux2_variants.detect(_flux2_base_file(root, base_mode)) + return detected if detected is not None else flux2_variants.get(_FLUX2_BASES["raw"]) + + +def loader_arch(arch: str, models_dir: str | None = None, base_mode: str = "raw") -> str: """The ``models/loaders.py`` arch key for a training arch. They are not the same namespace: training says ``flux2`` while the loaders key their config and - tokenizer bundles per variant (``flux2-klein-4b``, ``flux2-dev``), because a 4B and a 9B need - different encoder configs. Resolved from whichever checkpoint is installed. + tokenizer bundles per variant (``flux2-klein-4b``, ``flux2-klein-9b``), because a 4B and a 9B + need different encoder configs. Resolved from the base checkpoint the run trains on. """ if arch != archs.FLUX2: return arch - from ..models.flux2 import requirements as flux2_reqs + from ..config import models_dir as default_models_dir - variant = flux2_reqs.resolved_variant(None) + root = Path(models_dir) if models_dir else default_models_dir() + variant = flux2_variant(root, base_mode) return variant.arch if variant is not None else "flux2-klein-4b" @@ -148,6 +198,12 @@ def _adapter_path(root: Path, arch: str, base_mode: str) -> str | None: """The de-distillation adapter for a turbo base mode, or None when the base is undistilled.""" if base_mode not in ("turbo_adapter",): return None + if arch == archs.FLUX1: + raise RuntimeError( + "FLUX.1 has no de-distillation adapter. dev is guidance-distilled rather than " + "step-distilled, so it is itself the training base - a LoRA trains through that by " + "pinning guidance at 1. Set Base back to FLUX.1 dev." + ) if arch == archs.FLUX2: raise RuntimeError( "FLUX.2 has no de-distillation adapter. Train against a -base- checkpoint instead; " @@ -199,13 +255,18 @@ def compute_dtype() -> Any: return torch.float32 -def load_encoders(models_dir: str, arch: str, device: str, dtype: Any) -> Encoders: +def load_encoders( + models_dir: str, arch: str, device: str, dtype: Any, base_mode: str = "raw" +) -> Encoders: """The VAE + text encoder + tokenizer + scheduler, for the precache pass.""" from ..models import loaders root = Path(models_dir) - vae_file = _require(root, arch, "vae") - encoder_file = _require(root, arch, "text_encoders") + # Resolved once, from the base this run trains: the encoder, its config bundle and the + # transformer all have to describe the same checkpoint (see ``flux2_variant``). + variant = flux2_variant(root, base_mode) if arch == archs.FLUX2 else None + vae_file = _require(root, arch, "vae", variant) + encoder_file = _require(root, arch, "text_encoders", variant) if arch == archs.KREA2: from diffusers import FlowMatchEulerDiscreteScheduler, Krea2Pipeline @@ -223,10 +284,29 @@ def load_encoders(models_dir: str, arch: str, device: str, dtype: Any) -> Encode ) return Encoders(vae, text_encoder, tokenizer, scheduler, pipeline) + if arch == archs.FLUX1: + from diffusers import FluxPipeline + + clip_file = _require(root, arch, "clip") + vae = loaders.load_vae(arch, vae_file, dtype, device=device) + clip, clip_tok, t5, t5_tok = loaders.load_flux1_text_encoders( + arch, clip_file, encoder_file, dtype, device=device + ) + # Transformer-less, so caption encoding goes through diffusers' own two-encoder path: CLIP + # truncates at 77 tokens while T5 pads to 512, and the pooled vector is the EOS position. + pipeline = FluxPipeline( + scheduler=loaders.load_flux1_scheduler(), vae=None, text_encoder=clip, + tokenizer=clip_tok, text_encoder_2=t5, tokenizer_2=t5_tok, transformer=None, + ) + return Encoders( + vae, clip, clip_tok, pipeline.scheduler, pipeline, + text_encoder_2=t5, tokenizer_2=t5_tok, + ) + if arch == archs.FLUX2: from diffusers import Flux2KleinPipeline - larch = loader_arch(arch, models_dir) + larch = variant.arch if variant is not None else "flux2-klein-4b" vae = loaders.load_flux2_vae(larch, vae_file, dtype, device=device) text_encoder, tokenizer = loaders.load_text_encoder( larch, encoder_file, dtype, device=device @@ -259,6 +339,8 @@ def free_encoders(encoders: Encoders) -> None: encoders.vae = None encoders.text_encoder = None encoders.tokenizer = None + encoders.text_encoder_2 = None + encoders.tokenizer_2 = None encoders.pipeline = None loaders.unload_components(keep_files=set()) # the transformer isn't loaded yet - drop it all gc.collect() @@ -268,12 +350,27 @@ def free_encoders(encoders: Encoders) -> None: #: Architectures whose loader can build a 4-bit base. Z-Image is not here because it does not need #: to be: it trains in ~15GB at 1024, so bf16 already fits the cards people have. -_QUANTIZABLE = {archs.KREA2, archs.FLUX2} +_QUANTIZABLE = {archs.KREA2, archs.FLUX1, archs.FLUX2} #: Peak activation cost per image token, measured at rank 16, batch 1, gradient checkpointing on. #: Krea 2's wider blocks cost roughly 7x Z-Image's per token, and activations - not weights - are -#: what decides whether 1024 fits. -_ACTIVATION_MB_PER_TOKEN = {archs.KREA2: 5.2, archs.Z_IMAGE: 0.8, archs.FLUX2: 1.6} +#: what decides whether 1024 fits. FLUX.2 is keyed per variant, not per arch: a 4B and a 9B are the +#: same arch at 3072x25 and 4096x36 blocks, which is nearly twice the activation per token. +#: +#: flux1 is measured on an L40S at rank 16, batch 1: 24.90GB at 512 and 26.51GB at 1024 against a +#: 23.8GB bf16 base, which is a slope of 0.52MB per token. Rounded up, because under-estimating +#: promises bf16 to a card that then OOMs while over-estimating only reaches for NF4 sooner. +#: +#: flux2-klein-9b is measured the same way: 19.15GB at 512 and 20.80GB at 1024 against an 18.2GB +#: base, a slope of 0.55MB per token. Its 4B sibling's 1.6 is left as it was rather than re-derived +#: from a run this change did not make. +_ACTIVATION_MB_PER_TOKEN = { + archs.KREA2: 5.2, + archs.Z_IMAGE: 0.8, + "flux2-klein-4b": 1.6, + "flux2-klein-9b": 0.6, + archs.FLUX1: 0.6, +} #: Room left for the adapter, its 8-bit optimizer state and allocator slack. _MARGIN_BYTES = 2 * 1024**3 @@ -333,7 +430,9 @@ def resolve_quant( base = _base_size(models_dir, arch, base_mode) if not base: return Quantization.NONE # unmeasurable, so do not guess at the user's expense - needed = base + _activation_bytes(arch, resolution) + _MARGIN_BYTES + needed = base + _activation_bytes( + _activation_key(arch, models_dir, base_mode), resolution + ) + _MARGIN_BYTES fits = torch.cuda.get_device_properties(0).total_memory >= needed return Quantization.NONE if fits else Quantization.NF4 @@ -368,15 +467,30 @@ def resolve_offload( base = _base_size(models_dir, arch, base_mode) if not base: return False # unmeasurable, so do not pay the offload cost on a guess - needed = base + _activation_bytes(arch, resolution) + _MARGIN_BYTES + needed = base + _activation_bytes( + _activation_key(arch, models_dir, base_mode), resolution + ) + _MARGIN_BYTES return torch.cuda.get_device_properties(0).total_memory < needed -def _activation_bytes(arch: str, resolution: int) -> int: +def _activation_key(arch: str, models_dir: str, base_mode: str) -> str: + """Which ``_ACTIVATION_MB_PER_TOKEN`` row this run reads: FLUX.2's is per variant.""" + if arch != archs.FLUX2: + return arch + try: + return loader_arch(arch, models_dir, base_mode) + except RuntimeError: + return arch # unresolvable, so fall through to the conservative default below + + +def _activation_bytes(key: str, resolution: int) -> int: """Estimated peak activation memory. Image tokens are the VAE's 8x downscale then 2x2 - patching, and cost is linear in them - attention is memory-efficient, so there is no square.""" + patching, and cost is linear in them - attention is memory-efficient, so there is no square. + + An unknown key takes Krea 2's number, the largest here: over-estimating costs a slower 4-bit + run, under-estimating costs an OOM partway through one.""" tokens = (max(resolution, 1) // 16) ** 2 - per_token = _ACTIVATION_MB_PER_TOKEN.get(arch, 5.2) + per_token = _ACTIVATION_MB_PER_TOKEN.get(key, 5.2) return int(tokens * per_token * 1024**2) @@ -477,6 +591,16 @@ def load_transformer( loras: tuple[LoraRef, ...] = (LoraRef(file=adapter, strength=1.0),) if adapter else () diffusion = _base_file(root, arch, base_mode) + if arch == archs.FLUX1: + from ..models.checkpoint import CheckpointReader + from ..models.flux1 import variants as flux1_variants + + config = flux1_variants.derive_transformer_config(CheckpointReader(diffusion).shapes()) + if config is None: + raise RuntimeError(f"{Path(diffusion).name} is not a FLUX.1 checkpoint.") + return loaders.load_flux1_transformer( + arch, diffusion, config, dtype, quant, device=device, loras=loras, + ) if arch == archs.FLUX2: from ..models.checkpoint import CheckpointReader from ..models.flux2 import variants as flux2_variants @@ -485,7 +609,7 @@ def load_transformer( if config is None: raise RuntimeError(f"{Path(diffusion).name} is not a FLUX.2 checkpoint.") return loaders.load_flux2_transformer( - loader_arch(arch, models_dir), diffusion, config, dtype, quant, + loader_arch(arch, models_dir, base_mode), diffusion, config, dtype, quant, device=device, loras=loras, ) if arch == archs.KREA2: @@ -510,8 +634,10 @@ def _base_file(root: Path, arch: str, base_mode: str) -> str: """The base checkpoint this run trains against.""" if arch == archs.LTX25: return _ltx25_base_file() + if arch == archs.FLUX1: + return _flux1_base_file(root) if arch == archs.FLUX2: - return _flux2_base_file(root) + return _flux2_base_file(root, base_mode) if arch != archs.KREA2: return _require(root, arch, "diffusion_models") @@ -528,25 +654,84 @@ def _base_file(root: Path, arch: str, base_mode: str) -> str: return str(diffusion) -def _flux2_base_file(root: Path) -> str: - """The **undistilled** FLUX.2 checkpoint to train against. +def _flux1_base_file(root: Path) -> str: + """The FLUX.1 checkpoint to train against. + + dev is the training base. It is *guidance*-distilled, which a LoRA trains through by pinning + guidance at 1 rather than around; schnell is the *step*-distilled build and collapses the same + way a distilled FLUX.2 does. schnell is spotted by content, not by name: it is the one FLUX.1 + build with no guidance embedder at all. + + Fill and Control resolve as undistilled too, but their extra input channels carry a mask or a + stacked hint the dataset exporter does not produce, so they are refused here rather than as a + shape error twenty minutes into a precache. + """ + from ..models.flux1 import variants as flux1_variants + + override = os.environ.get(_ENV[archs.FLUX1]["diffusion_models"]) + if override: + return override + folder = root / "diffusion_models" + candidates = sorted(folder.iterdir()) if folder.is_dir() else [] + detected = [(p, flux1_variants.detect(p)) for p in candidates if p.is_file()] + trainable = [p for p, v in detected if v is not None and flux1_variants.trainable(v)] + if trainable: + return str(trainable[0]) + found = [(p, v) for p, v in detected if v is not None] + if found: + name, variant = found[0][0].name, found[0][1] + why = ( + "the step-distilled schnell build, which trains badly" + if variant.distilled + else f"the {variant.label} build, which needs paired data this app does not export" + ) + raise RuntimeError( + f"{name} is {why}. Download FLUX.1 dev from the node's model popup and train on that." + ) + raise RuntimeError( + f"No FLUX.1 checkpoint found under {folder}. Download one from the node's model popup " + f"(or set {_ENV[archs.FLUX1]['diffusion_models']})." + ) + + +#: FLUX.2 base mode -> the variant it trains. ``raw`` stays klein 4B so runs saved before 9B was +#: offered keep resolving to the checkpoint they were trained against. +_FLUX2_BASES = {"raw": "klein-4b-base", "raw_9b": "klein-9b-base"} + + +def _flux2_base_file(root: Path, base_mode: str) -> str: + """The **undistilled** FLUX.2 checkpoint to train against, for the picked base mode. Training on a step-distilled build is the documented cause of the collapse reports: BFL and musubi-tuner both say to train on ``-base-`` and load the adapter onto the distilled model afterwards, which is also faster and usually better. So a distilled checkpoint is refused here rather than silently producing a bad LoRA hours later. + + The mode picks the size as well, because both bases can be installed at once and sorted order + would otherwise always hand back 4B. """ from ..models.flux2 import variants as flux2_variants override = os.environ.get(_ENV[archs.FLUX2]["diffusion_models"]) if override: return override + wanted = _FLUX2_BASES.get(base_mode) + if wanted is None: + raise RuntimeError(f"FLUX.2 has no {base_mode!r} base mode.") folder = root / "diffusion_models" candidates = sorted(folder.iterdir()) if folder.is_dir() else [] detected = [(p, flux2_variants.detect(p)) for p in candidates if p.is_file()] - base = [p for p, v in detected if v is not None and not v.distilled] + base = [p for p, v in detected if v is not None and v.key == wanted] if base: return str(base[0]) + label = flux2_variants.get(wanted) + other = [(p, v) for p, v in detected if v is not None and not v.distilled] + if other: + raise RuntimeError( + f"No FLUX.2 {label.label if label else wanted} checkpoint found under {folder} - the " + f"undistilled build there is {other[0][1].label}. Download the one you picked from the " + "FLUX.2 node's model popup, or switch the Base setting to match what you have." + ) distilled = [(p, v) for p, v in detected if v is not None] if distilled: name, variant = distilled[0][0].name, distilled[0][1] diff --git a/core/src/inline_core/training/trainer.py b/core/src/inline_core/training/trainer.py index 6ef812c..ace2549 100644 --- a/core/src/inline_core/training/trainer.py +++ b/core/src/inline_core/training/trainer.py @@ -202,7 +202,7 @@ def _activation_offload(enabled: bool) -> Any: #: The cached-item keys that carry activations and take the compute dtype. Everything else moves #: unchanged: a bool mask would become weights, index tensors would stop addressing anything, and #: H3's float64 rotary grid would lose its mantissa. None of it raises. -_ACTIVATION_KEYS = frozenset({"latent", "embed", "audio"}) +_ACTIVATION_KEYS = frozenset({"latent", "embed", "pooled", "audio"}) def _to_device(item: dict[str, Any], device: Any, dtype: Any) -> dict[str, Any]: @@ -249,6 +249,7 @@ def train(manifest: dict[str, Any]) -> str | None: # that reaches the UI: this subprocess installs no logging handler. data, unconditional, shift = cache.build( manifest["datasetDir"], manifest["modelsDir"], arch.key, str(device), dtype, resolution, + base_mode=manifest["baseMode"], flip=bool(hp.get("flipAugment")), dropout=dropout, clip_frames=archs.clip_frames(arch, hp.get("clipSeconds")), clip_window=str(hp.get("clipWindow") or "start"), diff --git a/core/tests/test_flux1_resolve.py b/core/tests/test_flux1_resolve.py new file mode 100644 index 0000000..d675b38 --- /dev/null +++ b/core/tests/test_flux1_resolve.py @@ -0,0 +1,143 @@ +"""Which files the FLUX.1 node picks off disk, and what the popup reports. + +``text_encoders/`` is the crowded one: FLUX.1 puts T5-XXL and CLIP-L in the same folder that already +holds Qwen3 for Z-Image and FLUX.2 klein, and Mistral-3 for FLUX.2 dev. T5's embedding is 4096 wide, +which is exactly what klein 9B looks for, so identification is by key layout and never by width. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest + +from inline_core.models.flux1 import requirements as reqs +from inline_core.models.flux1 import variants as V +from tests.test_flux1_variants import BFL_DEV_HEADER +from tests.test_flux2_resolve import _write_header_only + +#: The identifying subset of each real encoder's header. +T5_HEADER = { + "shared.weight": [32128, 4096], + "encoder.embed_tokens.weight": [32128, 4096], + "encoder.block.0.layer.0.SelfAttention.q.weight": [4096, 4096], +} +CLIP_HEADER = { + "text_model.embeddings.token_embedding.weight": [49408, 768], + "text_model.encoder.layers.0.self_attn.q_proj.weight": [768, 768], +} +#: FLUX.1's VAE, which Z-Image ships as the same file. FLUX.2's carries running batch-norm buffers +#: in place of the scalar scale and shift, which is what tells the two apart. +VAE_HEADER = {"encoder.conv_out.weight": [32, 512, 3, 3], "decoder.conv_out.weight": [3, 128, 3, 3]} +FLUX2_VAE_HEADER = {**VAE_HEADER, "bn.running_mean": [128], "bn.running_var": [128]} + + +@pytest.fixture +def models(tmp_path: Path, monkeypatch: Any) -> Path: + root = tmp_path / "models" + monkeypatch.setenv("INLINE_MODELS_DIR", str(root)) + for category in ("diffusion_models", "vae", "text_encoders"): + (root / category).mkdir(parents=True) + reqs._IDENTIFIED.clear() # keyed by path, and tmp_path is reused + return root + + +def test_resolve_skips_checkpoints_from_other_architectures(models: Path) -> None: + _write_header_only(models / "diffusion_models" / "z_image_bf16.safetensors", {"foo": [4, 4]}) + assert reqs.resolve_diffusion() is None + + flux = _write_header_only( + models / "diffusion_models" / "flux1-dev.safetensors", BFL_DEV_HEADER + ) + assert reqs.resolve_diffusion() == flux + assert reqs.resolved_variant() is V.get("dev") + + +def test_the_two_encoders_are_told_apart_by_their_key_layout(models: Path) -> None: + # Named to sort the wrong way round, so an order-based pick would swap them. + t5 = _write_header_only(models / "text_encoders" / "zzz_t5.safetensors", T5_HEADER) + clip = _write_header_only(models / "text_encoders" / "aaa_clip.safetensors", CLIP_HEADER) + assert reqs.resolve_text_encoder() == t5 + assert reqs.resolve_clip() == clip + + +def test_a_qwen3_encoder_is_never_offered_to_flux1(models: Path) -> None: + # The other direction of the shared-folder problem: Z-Image and klein's encoder is not ours. + _write_header_only( + models / "text_encoders" / "qwen_3_4b.safetensors", + {"model.embed_tokens.weight": [151936, 2560]}, + ) + assert reqs.resolve_text_encoder() is None + assert reqs.resolve_clip() is None + + +def test_a_flux1_t5_encoder_is_not_offered_to_flux2(models: Path) -> None: + """The regression this guards, from the other side: T5's ``encoder.embed_tokens.weight`` is + 32128 x 4096, and 4096 is exactly the width FLUX.2 klein 9B matches on.""" + from inline_core.models.flux2 import requirements as flux2_reqs + from tests.test_flux2_variants import KLEIN_9B + from tests.test_flux2_variants import _shapes as _flux2_shapes + + flux2_reqs._IDENTIFIED.clear() + _write_header_only(models / "text_encoders" / "aaa_t5xxl.safetensors", T5_HEADER) + _write_header_only(models / "text_encoders" / "aab_clip_l.safetensors", CLIP_HEADER) + _write_header_only( + models / "diffusion_models" / "flux-2-klein-9b.safetensors", _flux2_shapes(KLEIN_9B) + ) + assert flux2_reqs.resolve_text_encoder() is None, "FLUX.1's encoders are not klein 9B's" + + +def test_the_popup_lists_four_required_components_then_the_extras(models: Path) -> None: + components = reqs.flux1_requirements() + required = [c for c in components if not c.optional] + assert [c.id for c in required] == ["diffusion", "text_encoder", "clip", "vae"] + assert not any(c.present for c in required), "an empty models dir has nothing" + # The licence is on the row, because the popup is the only place a user reads it. + assert "non-commercial" in next(c.label for c in required if c.id == "diffusion") + assert {c.id for c in components if c.optional} == {"diffusion_fp8", "text_encoder_fp8"} + + _write_header_only(models / "diffusion_models" / "flux1-dev.safetensors", BFL_DEV_HEADER) + _write_header_only(models / "text_encoders" / "t5xxl_fp16.safetensors", T5_HEADER) + _write_header_only(models / "text_encoders" / "clip_l.safetensors", CLIP_HEADER) + _write_header_only(models / "vae" / "ae.safetensors", VAE_HEADER) + reqs._IDENTIFIED.clear() + assert all(c.present for c in reqs.flux1_requirements() if not c.optional) + + +def test_an_env_override_wins_over_the_scan(models: Path, monkeypatch: Any) -> None: + picked = _write_header_only(models / "elsewhere.safetensors", BFL_DEV_HEADER) + monkeypatch.setenv("INLINE_FLUX1_MODEL", str(picked)) + assert reqs.resolve_diffusion() == picked + + +def test_the_node_offers_only_files_it_can_load(models: Path) -> None: + from inline_core.models.flux1.provider import Flux1Provider + + _write_header_only(models / "diffusion_models" / "flux1-dev.safetensors", BFL_DEV_HEADER) + _write_header_only(models / "diffusion_models" / "z_image.safetensors", {"foo": [4, 4]}) + _write_header_only(models / "text_encoders" / "t5xxl_fp16.safetensors", T5_HEADER) + _write_header_only(models / "text_encoders" / "clip_l.safetensors", CLIP_HEADER) + _write_header_only( + models / "text_encoders" / "qwen_3_4b.safetensors", + {"model.embed_tokens.weight": [151936, 2560]}, + ) + reqs._IDENTIFIED.clear() + provider = Flux1Provider() + assert provider.catalog_options("diffusion_models") == ["flux1-dev.safetensors"] + assert sorted(provider.catalog_options("text_encoders") or []) == [ + "clip_l.safetensors", + "t5xxl_fp16.safetensors", + ] + assert provider.resolved()["variant"] == "dev" + + +def test_the_vae_is_shared_with_z_image_but_never_taken_from_flux2(models: Path) -> None: + """FLUX.1's VAE is ``ae.safetensors``, and Z-Image ships the same weights under the same name - + identical tensor names, shapes and values. FLUX.2's is a different VAE that would match any + "flux" name check, so the two are told apart by the batch-norm buffers FLUX.2 carries.""" + _write_header_only(models / "vae" / "flux2-vae.safetensors", FLUX2_VAE_HEADER) + assert reqs.resolve_vae() is None, "FLUX.2's VAE is not FLUX.1's" + + shared = _write_header_only(models / "vae" / "ae.safetensors", VAE_HEADER) + assert reqs.resolve_vae() == shared diff --git a/core/tests/test_flux1_training.py b/core/tests/test_flux1_training.py new file mode 100644 index 0000000..4d328e9 --- /dev/null +++ b/core/tests/test_flux1_training.py @@ -0,0 +1,268 @@ +"""The FLUX.1 LoRA training arch: what it adapts, what it predicts, and which base it demands. + +Two rules decide whether a run is worth anything, and neither raises when broken - they produce a +plausible-but-wrong adapter hours later. Guidance is trained at 1, not at the 3.5 dev generates +with, because dev is guidance-distilled and the embedder is part of the model. And conditioning is +two tensors, not one: T5's sequence and CLIP's pooled vector, both at the compute dtype. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from inline_core.training import arch as archs +from tests.test_flux1_variants import BFL_DEV_HEADER, DEV +from tests.test_flux2_resolve import _write_header_only + +models = pytest.importorskip("inline_core.training.models") + +#: A miniature FLUX.1 with the real topology - both block types, the pre_only single-block +#: attention, the guidance embedder - so the forward and adapter run without a 12B checkpoint. +#: ``axes_dims_rope`` must sum to ``attention_head_dim``; ``in_channels`` stays 64 because that is +#: what a 16-channel latent becomes once the pipeline folds it 2x2. +_TINY = { + "attention_head_dim": 32, + "axes_dims_rope": [8, 12, 12], + "guidance_embeds": True, + "in_channels": 64, + "joint_attention_dim": 192, + "num_attention_heads": 4, + "num_layers": 2, + "num_single_layers": 2, + "out_channels": None, + "patch_size": 1, + "pooled_projection_dim": 64, +} +_EMBED, _POOLED = 192, 64 + + +def _tiny_model(): + torch = pytest.importorskip("torch") + diffusers = pytest.importorskip("diffusers") + return diffusers.FluxTransformer2DModel(**_TINY).to(torch.float32).eval() + + +def test_flux1_is_a_registered_training_arch() -> None: + a = archs.get("flux1") + assert a.key == archs.FLUX1 + # Rectified flow, Krea 2's convention: x_t = (1-s)*clean + s*noise, so d/ds is noise - clean. + assert a.target(clean=2.0, noise=5.0) == 3.0 + assert a.timestep(0.25) == 0.25 + assert a.clip is None, "FLUX.1 trains on stills" + + +def _linears() -> set[str]: + torch = pytest.importorskip("torch") + diffusers = pytest.importorskip("diffusers") + # On meta, so the real 19 + 38 topology costs nothing and the counts below are the real ones. + with torch.device("meta"): + model = diffusers.FluxTransformer2DModel(**DEV) + return {n for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)} + + +def _resolved(target: str, linears: set[str]) -> set[str]: + """The modules PEFT would attach to for one target: an exact name, or a dotted suffix.""" + return {n for n in linears if n == target or n.endswith("." + target)} + + +def test_every_target_matches_the_real_model() -> None: + linears = _linears() + for target in archs.get("flux1").target_modules: + assert _resolved(target, linears), f"{target} matches nothing" + + +def test_to_out_reaches_the_double_blocks_alone() -> None: + """FLUX.1's single blocks are ``pre_only``, so they carry no ``attn.to_out`` - the ModuleList + versus Linear suffix clash that forces FLUX.2 to drop its single-block output projection does + not arise here.""" + linears = _linears() + hit = _resolved("to_out.0", linears) + assert len(hit) == 19 + assert not any("single_transformer_blocks" in n for n in hit) + + +def test_proj_out_reaches_the_single_blocks_and_the_models_own_tail() -> None: + """The double match is a decision, not an accident: PEFT matches by suffix, so ``proj_out`` + takes the 38 single-block projections *and* the model's own final ``proj_out``. Both are plain + Linears, so nothing breaks. Pinned here so a later change to it is deliberate.""" + hit = _resolved("proj_out", _linears()) + assert len(hit) == 39 + assert "proj_out" in hit + assert sum(1 for n in hit if n.startswith("single_transformer_blocks.")) == 38 + + +def test_the_adaln_modulation_linears_are_left_alone() -> None: + # Adapting the modulation projections fights the checkpoint the way H3's adaln_proj does. + targets = archs.get("flux1").target_modules + assert not {"norm1.linear", "norm1_context.linear", "norm.linear", "norm_out.linear"} & set( + targets + ) + + +def test_attention_scope_narrows_to_the_projections() -> None: + narrowed = archs.target_modules(archs.get("flux1"), "attention") + assert {"to_q", "to_k", "to_v", "to_out.0"}.issubset(set(narrowed)) + assert not {"proj_mlp", "proj_out", "ff.net.2", "x_embedder"} & set(narrowed) + assert set(narrowed) < set(archs.get("flux1").target_modules) + + +def test_one_training_step_produces_a_prediction_shaped_like_its_target() -> None: + torch = pytest.importorskip("torch") + a = archs.get("flux1") + torch.manual_seed(0) + model = _tiny_model() + + clean = torch.randn(16, 32, 32) # a 16-channel H/8 latent, as the VAE produces + noise = torch.randn_like(clean) + sigma = a.sigma("cpu", 3.0) + noisy = (1 - sigma) * clean + sigma * noise + item = {"embed": torch.randn(77, _EMBED), "pooled": torch.randn(_POOLED)} + pred = a.forward(model, noisy, a.timestep(sigma), item) + + assert pred.shape == a.target(clean, noise).shape + assert torch.isfinite(pred).all() + + +def test_the_forward_passes_guidance_of_one_and_the_pooled_projections() -> None: + """Neither is visible in a shape test, and both are wrong in ways that only show up as a bad + adapter: dev generates at 3.5, and training the embedder there teaches the LoRA to change how + guidance behaves rather than what the images look like.""" + torch = pytest.importorskip("torch") + a = archs.get("flux1") + # No model here, so these are the real widths: T5's 4096 sequence and CLIP's 768 pooled vector. + seen: dict[str, object] = {} + + def capture(**kwargs: object) -> tuple[object]: + seen.update(kwargs) + return (torch.zeros(1, 256, 64),) + + a.forward( + capture, torch.randn(16, 32, 32), torch.tensor(0.5), + {"embed": torch.randn(77, 4096), "pooled": torch.randn(768)}, + ) + guidance = seen["guidance"] + assert tuple(guidance.shape) == (1,) + assert float(guidance[0]) == 1.0, "dev is guidance-distilled; training pins guidance at 1" + assert tuple(seen["pooled_projections"].shape) == (1, 768) + # 2-D on purpose: a 3-D txt_ids is deprecated and silently indexed back down. + assert seen["txt_ids"].dim() == 2 + + +def test_a_two_tensor_conditioning_round_trips_through_the_precache_store(tmp_path: Path) -> None: + """FLUX.1 is the first arch whose conditioning is two tensors, so the on-disk cache has to carry + both. Nothing in the store is arch-aware, which is what makes this worth pinning rather than + assuming.""" + torch = pytest.importorskip("torch") + from inline_core.training import precache_store as ps + + items = [{"latent": torch.randn(4, 8, 8), "embed": torch.randn(77, 16), + "pooled": torch.randn(32)}] + uncond = {"embed": torch.randn(77, 16), "pooled": torch.randn(32)} + ps.save(tmp_path, "k", items, uncond, 3.0) + loaded, loaded_uncond, shift = ps.load(tmp_path, "k") + + assert shift == 3.0 + assert torch.equal(loaded[0]["pooled"], items[0]["pooled"]) + assert torch.equal(loaded[0]["embed"], items[0]["embed"]) + assert loaded_uncond is not None + assert torch.equal(loaded_uncond["pooled"], uncond["pooled"]) + + +def test_the_pooled_vector_takes_the_compute_dtype() -> None: + """Left out of the activation set it would still reach the device, but stay fp32 while the + model runs bf16 - a dtype error at best, a silent upcast under autocast at worst.""" + torch = pytest.importorskip("torch") + from inline_core.training import trainer + + assert "pooled" in trainer._ACTIVATION_KEYS + item = {"latent": torch.randn(2, 2), "embed": torch.randn(2, 2), "pooled": torch.randn(4)} + moved = trainer._to_device(item, torch.device("cpu"), torch.float16) + assert moved["pooled"].dtype is torch.float16 + + +def test_a_lora_attaches_to_both_block_types_and_receives_gradient() -> None: + torch = pytest.importorskip("torch") + peft = pytest.importorskip("peft") + a = archs.get("flux1") + torch.manual_seed(0) + model = _tiny_model() + model.requires_grad_(False) + model.add_adapter( + peft.LoraConfig(r=4, lora_alpha=4, target_modules=a.target_modules, init_lora_weights=False) + ) + trainable = [(n, p) for n, p in model.named_parameters() if p.requires_grad] + assert trainable and all("lora" in n for n, _ in trainable) + + clean = torch.randn(16, 32, 32) + noise = torch.randn_like(clean) + sigma = a.sigma("cpu", 3.0) + noisy = (1 - sigma) * clean + sigma * noise + pred = a.forward( + model, noisy, a.timestep(sigma), + {"embed": torch.randn(77, _EMBED), "pooled": torch.randn(_POOLED)}, + ) + torch.nn.functional.mse_loss(pred.float(), a.target(clean, noise).float()).backward() + + got = [n for n, p in trainable if p.grad is not None and p.grad.abs().sum() > 0] + assert got, "no adapter parameter received gradient" + assert any("transformer_blocks." in n and "single" not in n for n in got) + assert any("single_transformer_blocks." in n for n in got) + + +# --- which checkpoint a run trains against ------------------------------------------------------- + + +@pytest.fixture +def models_root(tmp_path: Path, monkeypatch) -> Path: + root = tmp_path / "models" + (root / "diffusion_models").mkdir(parents=True) + monkeypatch.setenv("INLINE_MODELS_DIR", str(root)) + from inline_core.models.flux1 import requirements as reqs + + reqs._IDENTIFIED.clear() # keyed by path, and tmp_path is reused across tests + return root + + +def _put(root: Path, name: str, header: dict | None = None) -> Path: + return _write_header_only(root / "diffusion_models" / name, header or BFL_DEV_HEADER) + + +def test_dev_is_the_training_base(models_root: Path) -> None: + _put(models_root, "flux1-dev.safetensors") + assert models._base_file(models_root, "flux1", "raw").endswith("flux1-dev.safetensors") + + +def test_schnell_is_refused_with_a_pointer_to_dev(models_root: Path) -> None: + # Identified by content, not by name: schnell is the one build with no guidance embedder. + schnell = {k: v for k, v in BFL_DEV_HEADER.items() if not k.startswith("guidance_in.")} + _put(models_root, "flux1-schnell.safetensors", schnell) + with pytest.raises(RuntimeError, match="schnell"): + models._base_file(models_root, "flux1", "raw") + + +def test_an_empty_models_dir_says_where_to_get_a_checkpoint(models_root: Path) -> None: + with pytest.raises(RuntimeError, match="model popup"): + models._base_file(models_root, "flux1", "raw") + + +def test_flux1_has_no_de_distillation_adapter(models_root: Path) -> None: + """dev is *guidance*-distilled, not step-distilled, so it is itself the training base - unlike + Z-Image and Krea 2 there is nothing to fuse first.""" + with pytest.raises(RuntimeError, match="guidance-distilled"): + models._adapter_path(models_root, "flux1", "turbo_adapter") + assert models._adapter_path(models_root, "flux1", "raw") is None + + +def test_flux1_can_train_in_4bit() -> None: + # A 24GB bf16 base needs the NF4 rung to reach the cards people have. + assert archs.FLUX1 in models._QUANTIZABLE + + +def test_the_four_components_are_all_required_for_a_run() -> None: + from inline_core.models import trainingreqs + + ids = [c.id for c in trainingreqs.base_components("flux1", "raw")] + # No row swap, unlike FLUX.2: the popup's default checkpoint already is the training base. + assert ids == ["diffusion", "text_encoder", "clip", "vae"] diff --git a/core/tests/test_flux1_variants.py b/core/tests/test_flux1_variants.py new file mode 100644 index 0000000..b5c157d --- /dev/null +++ b/core/tests/test_flux1_variants.py @@ -0,0 +1,153 @@ +"""Which FLUX.1 checkpoint a file is, derived from its own tensor shapes. + +``diffusion_models/`` is shared with Z-Image, Krea 2 and FLUX.2, so identification is by content. +The values pinned here are read from the real ``Comfy-Org/flux1-dev/flux1-dev.safetensors`` header: +780 BF16 tensors, ``txt_in [3072, 4096]``, ``img_in [3072, 64]``, ``vector_in.in_layer +[3072, 768]``, ``guidance_in.in_layer [3072, 256]``, a ``[128]`` QK-norm vector, and 19+38 blocks. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from inline_core.models.flux1 import variants as V + +DEV: dict[str, object] = { + "attention_head_dim": 128, + "axes_dims_rope": [16, 56, 56], + "guidance_embeds": True, + "in_channels": 64, + "joint_attention_dim": 4096, + "num_attention_heads": 24, + "num_layers": 19, + "num_single_layers": 38, + "out_channels": None, + "patch_size": 1, + "pooled_projection_dim": 768, +} +SCHNELL = {**DEV, "guidance_embeds": False} +FILL = {**DEV, "in_channels": 384} +CANNY = {**DEV, "in_channels": 128} + + +def _shapes(config: dict[str, object]) -> dict[str, list[int]]: + """Real parameter shapes for a config, straight from diffusers on the meta device.""" + torch = pytest.importorskip("torch") + diffusers = pytest.importorskip("diffusers") + with torch.device("meta"): + model = diffusers.FluxTransformer2DModel(**config) + shapes = {name: list(p.shape) for name, p in model.named_parameters()} + shapes.update({name: list(b.shape) for name, b in model.named_buffers()}) + return shapes + + +#: The identifying subset of the real checkpoint, in BFL's key layout rather than diffusers'. +BFL_DEV_HEADER: dict[str, list[int]] = { + "txt_in.weight": [3072, 4096], + "img_in.weight": [3072, 64], + "vector_in.in_layer.weight": [3072, 768], + "guidance_in.in_layer.weight": [3072, 256], + "final_layer.linear.weight": [64, 3072], + "single_blocks.0.linear1.weight": [21504, 3072], + "single_blocks.0.linear2.weight": [3072, 15360], + **{f"double_blocks.{i}.img_attn.norm.query_norm.scale": [128] for i in range(19)}, + **{f"single_blocks.{i}.norm.query_norm.scale": [128] for i in range(38)}, +} + + +@pytest.mark.parametrize( + ("name", "config"), [("dev", DEV), ("schnell", SCHNELL), ("fill", FILL), ("canny", CANNY)] +) +def test_config_round_trips_from_tensor_shapes(name: str, config: dict[str, object]) -> None: + assert V.derive_transformer_config(_shapes(config)) == config, name + + +def test_the_shipped_bfl_checkpoint_derives_the_dev_config() -> None: + # The published file is in BFL's key layout; diffusers renames at load, but identification runs + # on the raw header, so the renames have to happen here first. + assert V.derive_transformer_config(BFL_DEV_HEADER) == DEV + assert V.detect("flux1-dev.safetensors", BFL_DEV_HEADER) is V.get("dev") + + +def test_comfy_style_key_prefixes_are_stripped() -> None: + prefixed = {f"model.diffusion_model.{k}": v for k, v in BFL_DEV_HEADER.items()} + assert V.derive_transformer_config(prefixed) == DEV + + +def test_schnell_is_identified_by_its_missing_guidance_embedder() -> None: + """schnell is the one step-distilled build, and no filename is needed to spot it: it is the + only FLUX.1 checkpoint with no guidance embedder at all.""" + schnell = {k: v for k, v in BFL_DEV_HEADER.items() if not k.startswith("guidance_in.")} + assert V.detect("some-name.safetensors", schnell) is V.get("schnell") + assert not V.trainable(V.get("schnell")) + + +def test_the_three_identical_builds_are_told_apart_by_name_and_nothing_else() -> None: + # dev, Kontext dev and Krea dev share every shape, so the name is the only signal left. + assert V.detect("flux1-dev.safetensors", BFL_DEV_HEADER) is V.get("dev") + assert V.detect("flux1-kontext-dev.safetensors", BFL_DEV_HEADER) is V.get("kontext-dev") + assert V.detect("flux1-krea-dev.safetensors", BFL_DEV_HEADER) is V.get("krea-dev") + + +def test_the_control_and_fill_builds_are_told_apart_by_their_input_channels() -> None: + assert V.detect("flux1-fill-dev.safetensors", _shapes(FILL)) is V.get("fill-dev") + assert V.detect("flux1-canny-dev.safetensors", _shapes(CANNY)) is V.get("canny-dev") + + +def test_only_a_plain_undistilled_build_trains() -> None: + """Fill and Control resolve as undistilled, but their extra input channels carry a mask or a + stacked hint the dataset exporter does not produce - a shape error twenty minutes into a + precache rather than a refusal up front. Kontext would train, and is left out for a different + reason: it learns an edit between a pair, so single images teach it nothing it is used for.""" + assert V.trainable(V.get("dev")) and V.trainable(V.get("krea-dev")) + assert not V.trainable(V.get("fill-dev")) + assert not V.trainable(V.get("canny-dev")) + assert not V.trainable(V.get("schnell")) + assert not V.trainable(V.get("kontext-dev")) + + +def test_a_wider_checkpoint_is_refused_rather_than_loaded_mis_split() -> None: + """diffusers' single-file converter hardcodes inner_dim 3072 and mlp_ratio 4.0 to split each + single block's fused linear1, so a wider build would load silently mis-split - a wrong image, + not an error. Refused at identification instead.""" + wide = {k: ([6144, v[1]] if k in ("txt_in.weight", "img_in.weight") else v) + for k, v in BFL_DEV_HEADER.items()} + wide["vector_in.in_layer.weight"] = [6144, 768] + assert V.derive_transformer_config(wide) is None + + +def test_a_checkpoint_without_pooled_conditioning_is_not_flux1() -> None: + # Only FLUX.1 conditions on a pooled CLIP vector; this is what keeps FLUX.2 out. + no_pooled = {k: v for k, v in BFL_DEV_HEADER.items() if k != "vector_in.in_layer.weight"} + assert V.derive_transformer_config(no_pooled) is None + + +def test_flux1_and_flux2_never_claim_each_other() -> None: + """Both directions. The FLUX.2 half passes because 4096 matches no FLUX.2 joint width, which a + future variant row could change - so it is pinned rather than left to luck.""" + from inline_core.models.flux2 import variants as V2 + from tests.test_flux2_variants import DEV as FLUX2_DEV + from tests.test_flux2_variants import KLEIN_4B, KLEIN_9B + from tests.test_flux2_variants import _shapes as _flux2_shapes + + assert V2.detect("flux1-dev.safetensors", BFL_DEV_HEADER) is None + for config in (KLEIN_4B, KLEIN_9B, FLUX2_DEV): + assert V.derive_transformer_config(_flux2_shapes(config)) is None + + +def test_rmsnorm_scales_do_not_make_a_plain_checkpoint_look_quantized(tmp_path: Path) -> None: + # A plain FLUX.1 file carries 152 `norm.query_norm.scale` weights that are not quant scales. + from tests.test_flux2_resolve import _write_header_only + + plain = _write_header_only(tmp_path / "flux1-dev.safetensors", BFL_DEV_HEADER) + assert V.quantization_of(plain) is None + assert not V.is_prequantized(plain) + assert V.single_file_blocker(plain) is None + + +def test_every_variant_names_a_loader_arch_and_a_pipeline() -> None: + assert {v.arch for v in V.VARIANTS} == {"flux1"} + assert {v.pipeline for v in V.VARIANTS} <= {"t2i", "kontext", "fill", "control"} + assert len({v.key for v in V.VARIANTS}) == len(V.VARIANTS) diff --git a/core/tests/test_flux2_resolve.py b/core/tests/test_flux2_resolve.py index 33c1a4a..3f06337 100644 --- a/core/tests/test_flux2_resolve.py +++ b/core/tests/test_flux2_resolve.py @@ -104,6 +104,41 @@ def test_a_vision_language_encoder_is_never_picked(models: Path) -> None: assert reqs.resolve_text_encoder() == good +def test_flux1_encoders_are_never_offered_to_flux2(models: Path) -> None: + """The bug this guards: T5-XXL's ``encoder.embed_tokens.weight`` is 32128 x 4096, and 4096 is + exactly the width klein 9B matches on - so FLUX.1's sequence encoder, dropped in the same + folder, was offered to klein 9B as its Qwen3-8B. CLIP-L is ruled out the same way.""" + # Both named to sort first, so a width-only match would reach them before the real encoder. + t5 = models / "text_encoders" / "aaa_t5xxl_fp16.safetensors" + _write_header_only( + t5, + { + "shared.weight": [32128, 4096], + "encoder.embed_tokens.weight": [32128, 4096], + "encoder.block.0.layer.0.SelfAttention.q.weight": [4096, 4096], + }, + ) + assert reqs._encoder_width(t5) is None + + clip = models / "text_encoders" / "aab_clip_l.safetensors" + _write_header_only( + clip, + { + "text_model.embeddings.token_embedding.weight": [49408, 768], + "text_model.encoder.layers.0.self_attn.q_proj.weight": [768, 768], + }, + ) + assert reqs._encoder_width(clip) is None + + _write_header_only( + models / "diffusion_models" / "flux-2-klein-9b.safetensors", _shapes(KLEIN_9B) + ) + assert reqs.resolve_text_encoder() is None, "FLUX.1's encoders are not klein 9B's" + + good = _encoder(models / "text_encoders" / "qwen_3_8b.safetensors", 4096) + assert reqs.resolve_text_encoder() == good + + def test_klein_4b_reuses_the_z_image_encoder_file(models: Path) -> None: # klein 4B's encoder is stock Qwen3-4B - the same file Z-Image downloads - so a user who has # run Z-Image already has it and must not be asked to fetch it again. @@ -143,6 +178,7 @@ def test_popup_blocks_on_the_required_three_then_lists_the_family(models: Path) assert {c.id for c in components if c.optional} == { "diffusion_klein_4b_base", "diffusion_klein_9b", + "diffusion_klein_9b_base", "text_encoder_qwen3_8b", "diffusion_dev", "text_encoder_mistral", diff --git a/core/tests/test_flux2_training.py b/core/tests/test_flux2_training.py index 378a46a..31b1c68 100644 --- a/core/tests/test_flux2_training.py +++ b/core/tests/test_flux2_training.py @@ -14,7 +14,7 @@ from inline_core.training import arch as archs from tests.test_flux2_resolve import _write_header_only -from tests.test_flux2_variants import KLEIN_4B, _shapes +from tests.test_flux2_variants import KLEIN_4B, KLEIN_9B, _shapes models = pytest.importorskip("inline_core.training.models") @@ -197,3 +197,78 @@ def test_a_lora_adapter_attaches_to_every_target_and_receives_gradient() -> None # Both block types must learn: the single blocks hold most of FLUX.2's parameters. assert any("transformer_blocks." in n and "single" not in n for n in got_grad) assert any("single_transformer_blocks." in n for n in got_grad) + + +def _installed(root: Path, *files: tuple[str, dict]) -> None: + """Write header-only checkpoints into ``diffusion_models/``.""" + for name, config in files: + _write_header_only(root / "diffusion_models" / name, _shapes(config)) + + +@pytest.fixture +def models_root(tmp_path: Path, monkeypatch) -> Path: + root = tmp_path / "models" + (root / "diffusion_models").mkdir(parents=True) + monkeypatch.setenv("INLINE_MODELS_DIR", str(root)) + from inline_core.models.flux2 import requirements as reqs + + reqs._IDENTIFIED.clear() # keyed by path, and tmp_path is reused across tests + return root + + +def test_the_base_mode_picks_the_size_not_only_the_undistilled_build(models_root: Path) -> None: + # Both bases installed at once: sorted order alone would always hand back 4B, because + # "flux-2-klein-base-4b" sorts before "flux-2-klein-base-9b". + _installed( + models_root, + ("flux-2-klein-base-4b.safetensors", KLEIN_4B), + ("flux-2-klein-base-9b.safetensors", KLEIN_9B), + ) + assert models._base_file(models_root, "flux2", "raw").endswith("base-4b.safetensors") + assert models._base_file(models_root, "flux2", "raw_9b").endswith("base-9b.safetensors") + + +def test_the_encoder_bundle_follows_the_base_rather_than_a_second_scan(models_root: Path) -> None: + """The bug this guards: the base file was the first *undistilled* checkpoint while the loader + arch came from a separate scan for the first FLUX.2 checkpoint at all. A distilled klein 4B + kept for generation sorts first, so a 9B run paired a 9B transformer with a 4B text encoder.""" + _installed( + models_root, + ("flux-2-klein-4b.safetensors", KLEIN_4B), # distilled, sorts first, generation-only + ("flux-2-klein-base-9b.safetensors", KLEIN_9B), + ) + assert models._base_file(models_root, "flux2", "raw_9b").endswith("base-9b.safetensors") + assert models.loader_arch("flux2", str(models_root), "raw_9b") == "flux2-klein-9b" + assert models.flux2_variant(models_root, "raw_9b").key == "klein-9b-base" + + +def test_a_9b_run_is_sized_against_its_own_row(models_root: Path) -> None: + """The sizing is keyed per variant rather than per arch, so a 9B run cannot read 4B's number. + + Only the routing is asserted, not which row is larger: measured on an L40S the two slopes are + 0.43 and 0.55 MB per token, so the rows are close, and 4B's constant sits well above its own + measurement. A magnitude comparison would be pinning that padding, not the behaviour.""" + _installed(models_root, ("flux-2-klein-base-9b.safetensors", KLEIN_9B)) + assert models._activation_key("flux2", str(models_root), "raw_9b") == "flux2-klein-9b" + assert models._activation_key("flux2", str(models_root), "raw") != "flux2-klein-9b" + # Both rows exist, so neither run silently falls through to the Krea 2 default. + assert {"flux2-klein-4b", "flux2-klein-9b"} <= set(models._ACTIVATION_MB_PER_TOKEN) + + +def test_asking_for_a_base_that_is_not_installed_names_what_is(models_root: Path) -> None: + _installed(models_root, ("flux-2-klein-base-4b.safetensors", KLEIN_4B)) + with pytest.raises(RuntimeError, match="Klein 4B Base"): + models._base_file(models_root, "flux2", "raw_9b") + + +def test_the_9b_base_and_its_own_encoder_are_both_required(models_root: Path) -> None: + from inline_core.models import trainingreqs + + ids = [c.id for c in trainingreqs.base_components("flux2", "raw_9b")] + assert "diffusion_klein_9b_base" in ids + # The required `text_encoder` row points at the 4B encoder whatever the base, so 9B swaps it. + assert "text_encoder_qwen3_8b" in ids + assert "text_encoder" not in ids and "diffusion" not in ids + + four = [c.id for c in trainingreqs.base_components("flux2", "raw")] + assert "diffusion_klein_4b_base" in four and "text_encoder" in four diff --git a/core/tests/test_flux2_variants.py b/core/tests/test_flux2_variants.py index ff74903..e0b3a5c 100644 --- a/core/tests/test_flux2_variants.py +++ b/core/tests/test_flux2_variants.py @@ -32,8 +32,9 @@ "rope_theta": 2000, "timestep_guidance_channels": 256, } -KLEIN_9B = {**KLEIN_4B, "joint_attention_dim": 12288, "num_attention_heads": 32, "num_layers": 6, - "num_single_layers": 30} +#: Block counts read from the published checkpoint's own header, not guessed: 8 + 24, not 6 + 30. +KLEIN_9B = {**KLEIN_4B, "joint_attention_dim": 12288, "num_attention_heads": 32, "num_layers": 8, + "num_single_layers": 24} DEV = {**KLEIN_4B, "joint_attention_dim": 15360, "num_attention_heads": 48, "num_layers": 8, "num_single_layers": 48, "guidance_embeds": True} diff --git a/core/tests/test_studio_rpc.py b/core/tests/test_studio_rpc.py index ede461d..ced0956 100644 --- a/core/tests/test_studio_rpc.py +++ b/core/tests/test_studio_rpc.py @@ -61,6 +61,7 @@ def test_full_project_and_canvas_flow(client) -> None: visible = {m["type"] for m in models["models"] if not m.get("hidden")} assert visible == { "alibaba/z-image-turbo", + "black-forest-labs/flux-1", "black-forest-labs/flux-2", "krea/krea-2-turbo", "krea/krea-2-raw", diff --git a/core/tests/test_training_dataset.py b/core/tests/test_training_dataset.py index 8235c8e..19736f4 100644 --- a/core/tests/test_training_dataset.py +++ b/core/tests/test_training_dataset.py @@ -139,3 +139,30 @@ def test_an_unpaired_item_is_a_clip_item_not_an_error(tmp_path) -> None: triples = media_triples(tmp_path) assert len(triples) == 1 assert triples[0].reference is None + + +def test_a_clip_dataset_on_a_stills_arch_says_so_before_the_encoders_load(tmp_path) -> None: + """The real report this came from: a 173-clip H3 dataset trained as FLUX.1 reported only "The + exported dataset is empty" - true, but not the reason, and only after a 10GB encoder had + loaded. The check names what is there and what will train it.""" + from inline_core.training import dataset as ds + + for i in range(3): + (tmp_path / f"{i:04d}.mp4").write_bytes(b"") + + with pytest.raises(RuntimeError, match="3 video clips and no images"): + ds.check_usable(str(tmp_path), "flux1") + # The video archs take the same folder without complaint. + ds.check_usable(str(tmp_path), "minimax-h3") + ds.check_usable(str(tmp_path), "ltx-2-5") + + # One image alongside is enough for a stills arch; the clips are simply skipped. + (tmp_path / "0000.png").write_bytes(b"") + ds.check_usable(str(tmp_path), "flux1") + + +def test_an_empty_dataset_folder_is_named(tmp_path) -> None: + from inline_core.training import dataset as ds + + with pytest.raises(RuntimeError, match="no images or clips"): + ds.check_usable(str(tmp_path), "flux1") diff --git a/package-lock.json b/package-lock.json index acedc10..95cbd98 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "inline-studio", - "version": "1.3.19", + "version": "1.3.20", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "inline-studio", - "version": "1.3.19", + "version": "1.3.20", "license": "GPL-3.0-or-later", "dependencies": { "@react-three/drei": "^10.7.7", diff --git a/package.json b/package.json index aa20c45..4e8d664 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "inline-studio", - "version": "1.3.19", + "version": "1.3.20", "description": "AI filmmaking on a node canvas. Generate locally on your own GPU and train your own LoRAs on the same canvas, with the built-in Inline Core engine and hosted models. Every render is kept as a versioned take.", "keywords": [ "ai-filmmaking", diff --git a/packages/frontend/pyproject.toml b/packages/frontend/pyproject.toml index d3ad1ac..e9e0b33 100644 --- a/packages/frontend/pyproject.toml +++ b/packages/frontend/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "openchar-frontend" -version = "1.3.19" +version = "1.3.20" description = "Prebuilt Inline Studio web UI (SPA), served by Inline Core. Mirrors comfyui-frontend-package." requires-python = ">=3.9" readme = "README.md" diff --git a/src/renderer/assets/h3-character.jpg b/src/renderer/assets/h3-character.jpg new file mode 100644 index 0000000..e7f454f Binary files /dev/null and b/src/renderer/assets/h3-character.jpg differ diff --git a/src/renderer/assets/race.png b/src/renderer/assets/race.png deleted file mode 100644 index 3566e77..0000000 Binary files a/src/renderer/assets/race.png and /dev/null differ diff --git a/src/renderer/components/Logo.tsx b/src/renderer/components/Logo.tsx index 27f7a98..14db8ff 100644 --- a/src/renderer/components/Logo.tsx +++ b/src/renderer/components/Logo.tsx @@ -1,11 +1,11 @@ import logoUrl from '../assets/logo.svg' -/** The Inline Studio logo mark. Sized via `size` (px). */ +/** The OpenChar logo mark. Sized via `size` (px). */ export function Logo({ size = 28 }: { size?: number }): React.JSX.Element { return ( Inline Studio

Try a demo

- {currentVersion && Inline Studio v{currentVersion}} + {currentVersion && OpenChar v{currentVersion}} {updateAvailable && ( / {project.name} diff --git a/src/shared/clipGrid.ts b/src/shared/clipGrid.ts index f11bfbc..72a5a88 100644 --- a/src/shared/clipGrid.ts +++ b/src/shared/clipGrid.ts @@ -29,6 +29,7 @@ export interface ClipGrid { export const CLIP_GRIDS: Record = { 'z-image': null, krea2: null, + flux1: null, flux2: null, 'minimax-h3': { fps: 24, grid: 17, offset: 5 }, 'ltx-2-5': { fps: 24, grid: 8, offset: 1 }, diff --git a/src/shared/types.ts b/src/shared/types.ts index e05316f..dea0173 100644 --- a/src/shared/types.ts +++ b/src/shared/types.ts @@ -645,7 +645,7 @@ export interface ModelDownloadErrorEvent { * The model family a LoRA is trained for. `minimax-h3` and `ltx-2-5` are the video models: they * train on stills or clips, and the adapter applies to every node of that family at generation time. */ -export type TrainingArch = 'z-image' | 'krea2' | 'flux2' | 'minimax-h3' | 'ltx-2-5' +export type TrainingArch = 'z-image' | 'krea2' | 'flux1' | 'flux2' | 'minimax-h3' | 'ltx-2-5' /** * Which shape an LTX-2.5 run trains in. `clip` learns a look and how it moves, from single clips. @@ -678,9 +678,11 @@ export interface DatasetRepoPreview { /** * Which base checkpoint a run trains against. `raw` is Krea 2's undistilled base (the recommended * path - the LoRA then applies to Turbo at generation time); the turbo modes train against a - * distilled checkpoint and need a training adapter to avoid turbo drift. + * distilled checkpoint and need a training adapter to avoid turbo drift. `raw_9b` is FLUX.2 klein + * Base 9B - a size rather than a distillation choice, and `raw` stays 4B so saved runs resolve to + * the checkpoint they were trained against. */ -export type TrainingBaseMode = 'turbo_adapter' | 'deturbo' | 'raw' +export type TrainingBaseMode = 'turbo_adapter' | 'deturbo' | 'raw' | 'raw_9b' /** * Precision of the frozen base during training. The LoRA itself is always full precision, so `nf4`