diff --git a/commands.md b/commands.md new file mode 100644 index 000000000..5687641ed --- /dev/null +++ b/commands.md @@ -0,0 +1,62 @@ +# Commands + +## Download model + +```sh +hf download prism-ml/bonsai-image-binary-4B-unpacked --local-dir models/bonsai-image-binary-4B-unpacked +``` + +## Merge sharded text encoder (required for both C++ and Python) + +```sh +uv run python3 -c " +from safetensors.torch import save_file +from safetensors import safe_open +import os + +d = 'models/bonsai-image-binary-4B-unpacked/text_encoder' +tensors = {} +for shard in ['model-00001-of-00002.safetensors', 'model-00002-of-00002.safetensors']: + with safe_open(os.path.join(d, shard), framework='pt') as f: + for k in f.keys(): + tensors[k] = f.get_tensor(k) +save_file(tensors, os.path.join(d, 'model.safetensors')) +print('merged:', os.path.getsize(os.path.join(d, 'model.safetensors')) / 1e9, 'GB') +" +``` + +## Build C++ with CUDA + +```sh +git submodule init && git submodule update +cmake -B build -DCMAKE_BUILD_TYPE=Release -DSD_CUDA=ON +cmake --build build -j +``` + +## C++ inference + +```sh +mkdir -p outputs +./build/bin/sd-cli \ + --cfg-scale 1 --width 512 --height 512 --steps 4 --seed 42 \ + -p "a cat sitting on a window sill" \ + -o outputs/cat_cpp.png \ + --diffusion-model models/bonsai-image-binary-4B-unpacked/transformer/diffusion_pytorch_model.safetensors \ + --vae models/bonsai-image-binary-4B-unpacked/vae/diffusion_pytorch_model.safetensors \ + --llm models/bonsai-image-binary-4B-unpacked/text_encoder/model.safetensors +``` + +## Python dependencies + +```sh +uv pip install torch diffusers transformers pillow accelerate +``` + +## Python inference + +```sh +mkdir -p outputs +uv run python3 image-studio/backend_gpu/scripts/inference_bf16.py \ + --prompt "a cat sitting on a window sill" \ + --output outputs/cat_python.png +``` diff --git a/image-studio/.gitignore b/image-studio/.gitignore new file mode 100644 index 000000000..ea0b72d1a --- /dev/null +++ b/image-studio/.gitignore @@ -0,0 +1,21 @@ +.venv/ +.vscode/ +__pycache__/ +*.py[cod] +*.png +.DS_Store +smoke/ +DerivedData/ +*.xcuserstate +*.xccheckout +*.xcscmblueprint +*.moved-aside +*.pbxuser +xcuserdata/ +.swiftpm/ +.build/ +apple/build/ +apple/build-derived*/ +apple/.DerivedData/ +!apple/**/*.png +!frontend/app/*.png diff --git a/image-studio/ATTRIBUTIONS.md b/image-studio/ATTRIBUTIONS.md new file mode 100644 index 000000000..871af88e6 --- /dev/null +++ b/image-studio/ATTRIBUTIONS.md @@ -0,0 +1,14 @@ +# Attributions + +Third-party assets bundled with this project. + +## Bonsai silhouette + +- File: `frontend/public/brand/bonsai-tree.svg` (mirrored to `apple/Bonsai/Resources/bonsai-tree.svg`) +- Source: https://commons.wikimedia.org/wiki/File:Logoprojetbonsai.svg +- Original drawing: **Crazou** (French Wikipedia user) +- SVG conversion: **Mouagip** +- License: **CC BY-SA 3.0** (also available under GFDL 1.2+) +- Modifications: viewBox normalized from 323×346 to 400×500 via a `translate`+`scale` wrapper; the original 8-stop vertical `linearGradient` was replaced with `fill="currentColor"` so the silhouette inherits theme tokens on the web and tint color on iOS. + +Our modified SVG is distributed under **CC BY-SA 3.0** to satisfy the share-alike requirement. A pointer to this entry is included as a comment at the top of each SVG file. diff --git a/image-studio/CONTRIBUTING.md b/image-studio/CONTRIBUTING.md new file mode 100644 index 000000000..86ca7a82a --- /dev/null +++ b/image-studio/CONTRIBUTING.md @@ -0,0 +1,60 @@ +# Contributing + +## Dev setup + +Requires Python 3.13 + [`uv`](https://docs.astral.sh/uv/) + Xcode 16+ (for the apple/ target). + +### mflux + mlx git refs + +`pyproject.toml`'s `[tool.uv.sources]` pins `mflux` and `mlx` to specific revisions under `PrismML-Eng/`. `uv sync` resolves both automatically — no sibling checkout required: + +```sh +uv venv .venv +uv pip install --python .venv/bin/python -e . +``` + +Bump the `mlx` rev or switch `mflux` to a sha pin by editing `[tool.uv.sources]` if you need to lock further. + +Per-component setup lives in [`docs/`](docs/). + +## Running tests + +Backend (FastAPI): +```sh +.venv/bin/python -m pytest backend/tests/ +``` + +Backend GPU (deployed to a CUDA host; tests exercise the loaders + server stubs): +```sh +.venv/bin/python -m pytest backend_gpu/tests/ +``` + +Apple (Mac Catalyst, fastest local loop): +```sh +xcodebuild test -project apple/Bonsai.xcodeproj -scheme Bonsai \ + -destination 'platform=macOS,variant=Mac Catalyst,arch=arm64' +``` + +iPhone parity, performance, and end-to-end tests are gated on a checkpoint payload at `BONSAI_TEST_CHECKPOINT_ROOT` (or the documented fallback). They `XCTSkip` cleanly when the payload is missing. + +## Commit messages + +Short, lower-case, imperative. Match `git log` style. No bodies unless the change is genuinely subtle. No co-author trailers, no AI-tool attribution. + +``` +apple: fused norm+RoPE megakernel + Klein parity tests +backend: pixel-stat diag for checkpoint sweep +docs: add LICENSE Apache 2.0 +``` + +## Pull requests + +- One concern per PR. If a change naturally splits into setup + behavior, split it. +- Don't auto-merge. Wait for review. +- Update [`docs/`](docs/) and [`README.md`](README.md) when behavior or env vars change. + +## Code style + +- Swift: follow the existing pattern in `apple/Bonsai/` — 4-space indent, no semicolons, terse `let` over `var`. Run the test suite under Mac Catalyst before pushing. +- Python: 4-space indent, type-hinted public APIs. The `backend/` and `backend_gpu/` modules are kept readable over clever; prefer explicit imports and short functions. +- TypeScript / Next.js (`frontend/`): `npm run lint` before pushing. diff --git a/image-studio/LICENSE b/image-studio/LICENSE new file mode 100644 index 000000000..302700978 --- /dev/null +++ b/image-studio/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for describing the origin of the Work and + reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Support. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or support. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2026 PrismML + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied. See the License for the specific language governing + permissions and limitations under the License. diff --git a/image-studio/README.md b/image-studio/README.md new file mode 100644 index 000000000..0342b27e1 --- /dev/null +++ b/image-studio/README.md @@ -0,0 +1,66 @@ +# image-studio + +Image generation studio for FLUX.2 Klein on-device. Three components: + +- **`backend/`** — FastAPI server fronting mflux backends on a Mac. +- **`backend_gpu/`** — separate FastAPI server for the GPU arm; deployed to a CUDA host. +- **`frontend/`** — Next.js web client that talks to `backend/`. + +Per-component setup, API contracts, and deployment lives under [`docs/`](docs/). High-level pointers: + +- [`docs/backend.md`](docs/backend.md) — `/generate` API + the backend selectors. +- [`docs/backend_gpu.md`](docs/backend_gpu.md) — gemlite GPU arm + deploy notes. +- [`docs/frontend.md`](docs/frontend.md) — Next.js client. + +Contributor workflow: [`CONTRIBUTING.md`](CONTRIBUTING.md). License: Apache 2.0 (see [`LICENSE`](LICENSE), [`ATTRIBUTIONS.md`](ATTRIBUTIONS.md)). + +## Quickstart (Next.js web client) + +Requires Node 20+ and `npm`. + +```sh +cd frontend +npm install +npm run dev # http://localhost:3000 +``` + +The frontend talks to a `backend/` FastAPI server at `http://localhost:8000` by default (`NEXT_PUBLIC_BACKEND_URL` to override). Start the backend first (next section) or it'll show "GPU unavailable" / "Unknown backend" until reachable. + +Build for prod: `npm run build && npm start`. See [`docs/frontend.md`](docs/frontend.md) for env vars + deploy notes (Vercel-ready). + +## Quickstart (local Mac backend) + +Requires Python 3.13 and [`uv`](https://docs.astral.sh/uv/). + +```sh +uv venv .venv +uv pip install --python .venv/bin/python -e . +.venv/bin/uvicorn backend.server:app --port 8000 +``` + +Set `MFLUX_STUDIO_BAKED_MODEL_PATH` to an absolute path to a FLUX.2 Klein ternary checkpoint root (containing `transformer-packed-mflux/`, `text_encoder/`, `tokenizer/`, `vae/`, `scheduler/`). See [`docs/backend.md`](docs/backend.md) for the checkpoint layout. + +## Backends + +`POST /generate` accepts a `backend` field with three values: + +- `bonsai-ternary-mlx` — Bonsai (ternary) Klein on MLX. +- `bfl-klein-bf16` — BFL FLUX.2 Klein-4B at bf16. +- `bonsai-ternary-gemlite` — remote GPU arm (ternary). + +The `bonsai-ternary-gemlite` arm is served by [`backend_gpu/`](docs/backend_gpu.md) over HTTP and is off by default; it gates on `MFLUX_STUDIO_GPU_HOST` + `MFLUX_STUDIO_GPU_TOKEN` being set. + +Exactly one Mac-side backend is resident at a time. On a backend change the server evicts the transformer+VAE and rebuilds (`Flux2Klein`); expect a one-shot swap cost. Concurrent requests serialize behind an `asyncio.Lock`. The remote `bonsai-ternary-gemlite` arm holds no in-process model so its swap is a label change and concurrent calls fan out over HTTP. + +`GET /backends` returns `{available, default, gpu: {available, reason}}` where `reason` is one of `force_disabled`, `no_gpu_host`, `no_gpu_token`, `healthz_failed:`, or `healthz_unreachable`. Probe result is cached for 30 s. Pass `?force_disable=1` to hide the GPU arm for a single response. + +## Env vars + +- `MFLUX_STUDIO_DEFAULT_BACKEND` — one of the three backend values above (default `bonsai-ternary-mlx`). `/backends` falls through to `bonsai-ternary-mlx` when the configured default is unavailable. +- `MFLUX_STUDIO_BAKED_MODEL_PATH` — absolute dir for the ternary MLX arm. Override to point at your local checkpoint. +- `MFLUX_STUDIO_STOCK_MODEL_PATH` — absolute dir for `bfl-klein-bf16`; when unset mflux resolves its own HF default. +- `MFLUX_STUDIO_GPU_HOST` / `MFLUX_STUDIO_GPU_TOKEN` — base URL + bearer token for the remote `bonsai-ternary-gemlite` arm. Both required to expose it via `/backends`. +- `MFLUX_STUDIO_FORCE_DISABLE_GPU` — `1`/`true` hides the GPU arm regardless of probe. Read once at module load — restart the server to flip it. +- `MFLUX_STUDIO_BACKENDS_PROBE_TTL_SECONDS` — override the 30 s `/backends` cache (mostly tests). + +Deprecated: `MFLUX_STUDIO_PRECISION` (`bf16`) and the request-body `precision` field still work but emit `DeprecationWarning` and will be removed next release. diff --git a/image-studio/backend/__init__.py b/image-studio/backend/__init__.py new file mode 100644 index 000000000..97385808f --- /dev/null +++ b/image-studio/backend/__init__.py @@ -0,0 +1 @@ +"""prism-image-studio backend package.""" diff --git a/image-studio/backend/eviction.py b/image-studio/backend/eviction.py new file mode 100644 index 000000000..2c261ddfb --- /dev/null +++ b/image-studio/backend/eviction.py @@ -0,0 +1,145 @@ +"""Per-component reload for Klein's transformer + VAE. + +Klein's stock `load_transformer_and_vae` rebuilds both unless both are already +resident. That makes VAE-only eviction useless: the next gen reloads the +transformer too. We patch the loader so a call with one component resident and +the other None rebuilds only the missing one. + +The patch also short-circuits the slim packed-mflux checkpoint case (only +`transformer-packed-mflux/`, no bf16 `transformer/`) when +`use_klein_fast_transformer = True`. The stock loader's first move is +`_load_weights(model_path)`, which scans the full Klein layout and crashes +when `transformer/` is missing. For klein-fast + slim we build the VAE from +the small-decoder weights and the transformer from the packed artifact +directly, skipping the failing scan entirely. + +Eviction side stays stock — Klein's `evict_transformer_and_vae` still bundles +both when `evict_transformer=True` is passed to `generate_image`. The pipeline +handles the VAE-only case by setting `model.vae = None` post-decode itself. +""" +from __future__ import annotations + +import gc +from pathlib import Path + +import mlx.core as mx + +from mflux.models.flux2.flux2_initializer import ( + FULL_DECODER_CHANNELS, + SMALL_DECODER_CHANNELS, + Flux2Initializer, +) +from mflux.models.common.weights.loading.loaded_weights import LoadedWeights, MetaData +from mflux.models.common.weights.loading.weight_applier import WeightApplier +from mflux.models.flux2.model.flux2_transformer.transformer import Flux2Transformer +from mflux.models.flux2.model.flux2_vae.vae import Flux2VAE +from mflux.models.flux2.weights.flux2_weight_definition import Flux2KleinWeightDefinition + + +_ORIGINAL_LOAD = Flux2Initializer.load_transformer_and_vae + + +def _is_slim_checkpoint(model_path: str | None) -> bool: + if not model_path: + return False + root = Path(model_path) + return (root / "transformer-packed-mflux").exists() and not (root / "transformer").exists() + + +def _load_slim_klein_fast(model) -> None: + """Build VAE (small-decoder weights) + klein-fast transformer for slim ckpts.""" + if model.vae is None: + decoder_channels = ( + SMALL_DECODER_CHANNELS if model._vae_variant == "small" else FULL_DECODER_CHANNELS + ) + model.vae = Flux2VAE(decoder_block_out_channels=decoder_channels) + # Small-decoder weights are the only VAE source we trust for slim ckpts; + # full VAE would need vae/ on disk, which slim does ship — fall through if so. + if model._vae_variant == "small": + vae_weights = Flux2Initializer._load_small_decoder_weights() + else: + # Slim ckpts include vae/; let the stock per-component loader pull it. + from mflux.models.common.resolution.path_resolution import PathResolution + from mflux.models.common.weights.loading.weight_loader import WeightLoader + + root = PathResolution.resolve( + path=model._model_path, + patterns=Flux2KleinWeightDefinition.get_download_patterns(), + ) + vae_component = next( + c for c in Flux2KleinWeightDefinition.get_components() if c.name == "vae" + ) + vae_weights, _, _ = WeightLoader._load_component(root, vae_component) + loaded = LoadedWeights( + components={"vae": vae_weights}, + meta_data=MetaData(quantization_level=None, mflux_version=None), + ) + WeightApplier.apply_and_quantize( + weights=loaded, + quantize_arg=model._quantize_arg, + weight_definition=Flux2KleinWeightDefinition, + models={"vae": model.vae}, + ) + + if model.transformer is None: + Flux2Initializer._load_klein_fast_transformer_weights( + model, model._model_path, precision=model._klein_fast_precision + ) + Flux2Initializer._apply_lora(model, model._lora_paths_arg, model._lora_scales_arg) + + gc.collect() + mx.clear_cache() + + +def _load_transformer_and_vae_per_component(model) -> None: + tx_resident = model.transformer is not None + vae_resident = model.vae is not None + if tx_resident and vae_resident: + return + + if bool(getattr(model, "_use_klein_fast_transformer", False)) and _is_slim_checkpoint( + getattr(model, "_model_path", None) + ): + _load_slim_klein_fast(model) + return + + if not tx_resident and not vae_resident: + _ORIGINAL_LOAD(model) + return + + weights = Flux2Initializer._load_weights(model._model_path) + + if not vae_resident: + decoder_channels = ( + SMALL_DECODER_CHANNELS if model._vae_variant == "small" else FULL_DECODER_CHANNELS + ) + model.vae = Flux2VAE(decoder_block_out_channels=decoder_channels) + if model._vae_variant == "small": + weights.components["vae"] = Flux2Initializer._load_small_decoder_weights() + WeightApplier.apply_and_quantize( + weights=weights, + quantize_arg=model._quantize_arg, + weight_definition=Flux2KleinWeightDefinition, + models={"vae": model.vae}, + ) + + if not tx_resident: + if model._use_klein_fast_transformer: + Flux2Initializer._load_klein_fast_transformer_weights( + model, model._model_path, precision=model._klein_fast_precision + ) + else: + model.transformer = Flux2Transformer(**model.model_config.transformer_overrides) + WeightApplier.apply_and_quantize( + weights=weights, + quantize_arg=model._quantize_arg, + weight_definition=Flux2KleinWeightDefinition, + models={"transformer": model.transformer}, + ) + Flux2Initializer._apply_lora(model, model._lora_paths_arg, model._lora_scales_arg) + + gc.collect() + mx.clear_cache() + + +Flux2Initializer.load_transformer_and_vae = staticmethod(_load_transformer_and_vae_per_component) diff --git a/image-studio/backend/pipeline.py b/image-studio/backend/pipeline.py new file mode 100644 index 000000000..d7742a768 --- /dev/null +++ b/image-studio/backend/pipeline.py @@ -0,0 +1,524 @@ +from __future__ import annotations + +import gc +import io +import logging +import os +import time +from dataclasses import dataclass +from typing import ClassVar, Literal + +import httpx +import mlx.core as mx +from mflux.models.common.vae.tiling_config import TilingConfig +from mflux.models.flux2.variants.txt2img.flux2_klein import Flux2Klein + +from backend import eviction # noqa: F401 (import installs per-component reload patch) +from backend.text_encoder_4bit import load_te_4bit + +Backend = Literal[ + "bonsai-ternary-mlx", + "bonsai-binary-mlx", + "bonsai-binary-gemlite", + "bonsai-ternary-gemlite", +] +BackendKind = Literal["mlx", "gemlite"] +ModelFamily = Literal["bonsai-binary", "bonsai-ternary"] +TiledMode = Literal["auto", "on", "off"] + +LOCAL_BACKENDS: tuple[Backend, ...] = ( + "bonsai-ternary-mlx", + "bonsai-binary-mlx", +) +REMOTE_BACKENDS: tuple[Backend, ...] = ( + "bonsai-binary-gemlite", + "bonsai-ternary-gemlite", +) +BACKENDS: tuple[Backend, ...] = LOCAL_BACKENDS + REMOTE_BACKENDS + +# Family order mirrors the frontend MODEL_FAMILIES rendering order. +MODEL_FAMILIES: tuple[ModelFamily, ...] = ( + "bonsai-binary", + "bonsai-ternary", +) + +BACKEND_TO_KIND: dict[Backend, BackendKind] = { + "bonsai-ternary-mlx": "mlx", + "bonsai-binary-mlx": "mlx", + "bonsai-binary-gemlite": "gemlite", + "bonsai-ternary-gemlite": "gemlite", +} +BACKEND_TO_FAMILY: dict[Backend, ModelFamily] = { + "bonsai-ternary-mlx": "bonsai-ternary", + "bonsai-binary-mlx": "bonsai-binary", + "bonsai-binary-gemlite": "bonsai-binary", + "bonsai-ternary-gemlite": "bonsai-ternary", +} + +DEFAULT_BAKED_MODEL_PATH = "/tmp/bonsai-checkpoints/v1/" +DEFAULT_BAKED_BINARY_MODEL_PATH: str | None = None +DEFAULT_BACKEND: Backend = "bonsai-ternary-mlx" +DEFAULT_TILED_MODE: TiledMode = "auto" +DEFAULT_EVICT_TEXT_ENCODER = True +DEFAULT_LAZY_COMPONENTS = False +DEFAULT_EVICT_TRANSFORMER = False +DEFAULT_EVICT_VAE = False +DEFAULT_MAX_SEQUENCE_LENGTH = 512 +DEFAULT_BUCKETED_SEQ_LEN = False +DEFAULT_TE_4BIT = True +DEFAULT_GPU_HOST: str | None = None +DEFAULT_GPU_TOKEN: str | None = None +DEFAULT_GPU_TIMEOUT_SECONDS = 600 +DEFAULT_GPU_CONNECT_TIMEOUT_SECONDS = 30 +DEFAULT_SEED = 42 +DEFAULT_STEPS = 4 +DEFAULT_HEIGHT = 512 +DEFAULT_WIDTH = 512 +DEFAULT_GUIDANCE = 1.0 + + +def _auto_tile_threshold() -> int: + return 2 * TilingConfig().vae_decode_tile_size + + +def _normalize_backend_id(raw: str) -> Backend: + if raw in BACKENDS: + return raw # type: ignore[return-value] + raise ValueError(f"Unknown backend {raw!r}; expected one of {BACKENDS}.") + + +def _parse_bool_env(name: str, default: bool) -> bool: + raw = os.getenv(name) + if raw is None: + return default + lowered = raw.strip().lower() + if lowered in {"1", "true", "yes", "on"}: + return True + if lowered in {"0", "false", "no", "off"}: + return False + raise ValueError(f"{name} must be a boolean (true/false/1/0/yes/no/on/off), got {raw!r}.") + + +def _parse_int_env(name: str, default: int, *, minimum: int = 1) -> int: + raw = os.getenv(name) + if raw is None: + return default + try: + value = int(raw) + except ValueError as exc: + raise ValueError(f"{name} must be an integer, got {raw!r}.") from exc + if value < minimum: + raise ValueError(f"{name} must be >= {minimum}, got {value}.") + return value + + +log = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class PipelineConfig: + backend: Backend = DEFAULT_BACKEND + baked_model_path: str = DEFAULT_BAKED_MODEL_PATH + baked_binary_model_path: str | None = DEFAULT_BAKED_BINARY_MODEL_PATH + tiled_mode: TiledMode = DEFAULT_TILED_MODE + evict_text_encoder: bool = DEFAULT_EVICT_TEXT_ENCODER + lazy_components: bool = DEFAULT_LAZY_COMPONENTS + evict_transformer: bool = DEFAULT_EVICT_TRANSFORMER + evict_vae: bool = DEFAULT_EVICT_VAE + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH + bucketed_seq_len: bool = DEFAULT_BUCKETED_SEQ_LEN + te_4bit: bool = DEFAULT_TE_4BIT + gpu_host: str | None = DEFAULT_GPU_HOST + gpu_token: str | None = DEFAULT_GPU_TOKEN + gpu_timeout_seconds: int = DEFAULT_GPU_TIMEOUT_SECONDS + gpu_connect_timeout_seconds: int = DEFAULT_GPU_CONNECT_TIMEOUT_SECONDS + + @staticmethod + def from_env() -> "PipelineConfig": + backend_raw = os.getenv("MFLUX_STUDIO_DEFAULT_BACKEND") + if backend_raw is not None: + backend: Backend = _normalize_backend_id(backend_raw) + else: + backend = DEFAULT_BACKEND + + baked_model_path = os.getenv("MFLUX_STUDIO_BAKED_MODEL_PATH") + if baked_model_path is None: + legacy = os.getenv("MFLUX_STUDIO_MODEL_PATH") + baked_model_path = legacy if legacy is not None else DEFAULT_BAKED_MODEL_PATH + if not os.path.isabs(baked_model_path): + raise ValueError( + "MFLUX_STUDIO_BAKED_MODEL_PATH must be an absolute path, " + f"got {baked_model_path!r}." + ) + baked_binary_model_path = os.getenv( + "MFLUX_STUDIO_BAKED_BINARY_MODEL_PATH", DEFAULT_BAKED_BINARY_MODEL_PATH + ) + if baked_binary_model_path is not None and not os.path.isabs(baked_binary_model_path): + raise ValueError( + "MFLUX_STUDIO_BAKED_BINARY_MODEL_PATH must be an absolute path when set, " + f"got {baked_binary_model_path!r}." + ) + + tiled_mode = os.getenv("MFLUX_STUDIO_TILED_VAE", DEFAULT_TILED_MODE) + if tiled_mode not in {"auto", "on", "off"}: + raise ValueError( + "MFLUX_STUDIO_TILED_VAE must be 'auto', 'on', or 'off', " + f"got {tiled_mode!r}." + ) + evict_text_encoder = _parse_bool_env( + "MFLUX_STUDIO_EVICT_TEXT_ENCODER", DEFAULT_EVICT_TEXT_ENCODER + ) + lazy_components = _parse_bool_env( + "MFLUX_STUDIO_LAZY_COMPONENTS", DEFAULT_LAZY_COMPONENTS + ) + evict_transformer = _parse_bool_env( + "MFLUX_STUDIO_EVICT_TRANSFORMER", DEFAULT_EVICT_TRANSFORMER + ) + evict_vae = _parse_bool_env("MFLUX_STUDIO_EVICT_VAE", DEFAULT_EVICT_VAE) + max_sequence_length = _parse_int_env( + "MFLUX_STUDIO_MAX_SEQUENCE_LENGTH", DEFAULT_MAX_SEQUENCE_LENGTH + ) + bucketed_seq_len = _parse_bool_env( + "MFLUX_STUDIO_BUCKETED_SEQ_LEN", DEFAULT_BUCKETED_SEQ_LEN + ) + te_4bit = _parse_bool_env("MFLUX_STUDIO_TE_4BIT", DEFAULT_TE_4BIT) + gpu_host = os.getenv("MFLUX_STUDIO_GPU_HOST", DEFAULT_GPU_HOST) + gpu_token = os.getenv("MFLUX_STUDIO_GPU_TOKEN", DEFAULT_GPU_TOKEN) + gpu_timeout_seconds = _parse_int_env( + "MFLUX_STUDIO_GPU_TIMEOUT_SECONDS", DEFAULT_GPU_TIMEOUT_SECONDS + ) + gpu_connect_timeout_seconds = _parse_int_env( + "MFLUX_STUDIO_GPU_CONNECT_TIMEOUT_SECONDS", DEFAULT_GPU_CONNECT_TIMEOUT_SECONDS + ) + if backend in REMOTE_BACKENDS: + if gpu_host is None: + raise ValueError( + f"Backend {backend!r} requires MFLUX_STUDIO_GPU_HOST to be set." + ) + if gpu_token is None: + raise ValueError( + f"Backend {backend!r} requires MFLUX_STUDIO_GPU_TOKEN to be set." + ) + return PipelineConfig( + backend=backend, + baked_model_path=baked_model_path, + baked_binary_model_path=baked_binary_model_path, + tiled_mode=tiled_mode, + evict_text_encoder=evict_text_encoder, + lazy_components=lazy_components, + evict_transformer=evict_transformer, + evict_vae=evict_vae, + max_sequence_length=max_sequence_length, + bucketed_seq_len=bucketed_seq_len, + te_4bit=te_4bit, + gpu_host=gpu_host, + gpu_token=gpu_token, + gpu_timeout_seconds=gpu_timeout_seconds, + gpu_connect_timeout_seconds=gpu_connect_timeout_seconds, + ) + + +def _resolve_tiling_config( + *, + request_override: bool | None, + server_default: TiledMode, + height: int, + width: int, +) -> TilingConfig | None: + if request_override is True: + return TilingConfig() + if request_override is False: + return None + if server_default == "on": + return TilingConfig() + if server_default == "off": + return None + return TilingConfig() if max(height, width) >= _auto_tile_threshold() else None + + +def _build_model( + *, + backend: Backend, + model_path: str | None, + config: PipelineConfig, +) -> Flux2Klein: + if backend in REMOTE_BACKENDS: + raise ValueError( + f"Backend {backend!r} is remote-only; build via RemoteGpuPipeline, " + "not FluxPipeline._build_model." + ) + if backend == "bonsai-ternary-mlx": + return Flux2Klein( + model_path=model_path, + use_klein_fast_transformer=True, + klein_fast_precision="2bit", + vae_variant="small", + evict_text_encoder=config.evict_text_encoder, + lazy_components=config.lazy_components, + bucketed_seq_len=config.bucketed_seq_len, + ) + if backend == "bonsai-binary-mlx": + return Flux2Klein( + model_path=model_path, + use_klein_fast_transformer=True, + klein_fast_precision="1bit", + vae_variant="small", + evict_text_encoder=config.evict_text_encoder, + lazy_components=config.lazy_components, + bucketed_seq_len=config.bucketed_seq_len, + ) + raise ValueError(f"Unknown backend {backend!r}; expected one of {BACKENDS}.") + + +def _default_model_path_for(backend: Backend, config: PipelineConfig) -> str | None: + if backend == "bonsai-ternary-mlx": + return config.baked_model_path + if backend == "bonsai-binary-mlx": + return config.baked_binary_model_path + return None + + +class FluxPipeline: + is_remote: ClassVar[bool] = False + + def __init__(self, config: PipelineConfig) -> None: + self.config = config + self._backend: Backend | None = None + self._model_path: str | None = None + self._model: Flux2Klein | None = None + self.last_swap_seconds: float | None = None + self.last_peak_memory_mb: float | None = None + self._load(backend=config.backend, model_path=_default_model_path_for(config.backend, config)) + + @property + def backend(self) -> Backend: + assert self._backend is not None + return self._backend + + @property + def model_path(self) -> str | None: + return self._model_path + + def _load(self, *, backend: Backend, model_path: str | None) -> None: + start = time.perf_counter() + self._model = _build_model(backend=backend, model_path=model_path, config=self.config) + if self.config.te_4bit: + # Replace the freshly loaded bf16 TE with the pre-quantized 4-bit Qwen3 + # from the local model dir (or fall back to mlx-community HF if not + # bundled). The marker routes Klein's cache-miss reload path (patched + # in text_encoder_4bit.py) back to load_te_4bit so eviction stays a win. + # Pass model_path explicitly so load_te_4bit's local-first probe sees + # the bundled text_encoder-mlx-4bit/ subdir — without it, this call + # always defaulted to the HF download path. + self._model._studio_te_4bit = True + overrides = self._model.model_config.text_encoder_overrides + self._model.text_encoder = load_te_4bit(overrides, model_path=model_path) + gc.collect() + mx.clear_cache() + self._backend = backend + self._model_path = model_path + self.last_swap_seconds = time.perf_counter() - start + log.info( + "Loaded backend=%s model_path=%s evict_te=%s lazy=%s evict_tx=%s evict_vae=%s " + "max_seq=%d bucketed=%s te_4bit=%s in %.3fs", + backend, + model_path, + self.config.evict_text_encoder, + self.config.lazy_components, + self.config.evict_transformer, + self.config.evict_vae, + self.config.max_sequence_length, + self.config.bucketed_seq_len, + self.config.te_4bit, + self.last_swap_seconds, + ) + + def ensure_backend(self, *, backend: Backend, model_path: str | None) -> None: + resolved = model_path if model_path is not None else _default_model_path_for(backend, self.config) + if backend == self._backend and resolved == self._model_path: + return + log.info( + "Hot-swapping backend %s(%s) -> %s(%s)", + self._backend, + self._model_path, + backend, + resolved, + ) + self._model = None + gc.collect() + mx.clear_cache() + self._load(backend=backend, model_path=resolved) + + def generate_png( + self, + *, + prompt: str, + seed: int = DEFAULT_SEED, + steps: int = DEFAULT_STEPS, + height: int = DEFAULT_HEIGHT, + width: int = DEFAULT_WIDTH, + guidance: float = DEFAULT_GUIDANCE, + tiled_vae: bool | None = None, + max_sequence_length: int | None = None, + ) -> bytes: + assert self._model is not None + tiling = _resolve_tiling_config( + request_override=tiled_vae, + server_default=self.config.tiled_mode, + height=height, + width=width, + ) + self._model.tiling_config = tiling + log.info( + "generate backend=%s size=%dx%d tiled=%s tile_size=%s", + self._backend, + width, + height, + tiling is not None, + tiling.vae_decode_tile_size if tiling is not None else None, + ) + effective_max_seq = ( + max_sequence_length if max_sequence_length is not None else self.config.max_sequence_length + ) + mx.reset_peak_memory() + generated = self._model.generate_image( + seed=seed, + prompt=prompt, + num_inference_steps=steps, + height=height, + width=width, + guidance=guidance, + max_sequence_length=effective_max_seq, + evict_transformer=self.config.evict_transformer, + ) + # VAE-only eviction (tx_evict=False, vae_evict=True). When tx_evict is on + # Klein has already bundled-evicted VAE via evict_transformer_and_vae. + if self.config.evict_vae and not self.config.evict_transformer: + self._model.vae = None + gc.collect() + mx.clear_cache() + self.last_peak_memory_mb = mx.get_peak_memory() / (1024 * 1024) + output = io.BytesIO() + try: + generated.image.save(output, format="PNG") + return output.getvalue() + finally: + output.close() + del generated + gc.collect() + + +class RemoteGpuPipeline: + """HTTP-proxy pipeline for the GPU arm (gemlite/HQQ on remote CUDA host). + + Mirrors FluxPipeline's surface so the FastAPI handler is byte-identical: + same `ensure_backend(backend, model_path)` + `generate_png(...)` returning + PNG bytes. Reuses one `httpx.Client` for HTTP/1.1 keep-alive across calls. + """ + + is_remote: ClassVar[bool] = True + + def __init__(self, config: PipelineConfig) -> None: + if config.gpu_host is None: + raise ValueError( + "RemoteGpuPipeline requires MFLUX_STUDIO_GPU_HOST." + ) + if config.gpu_token is None: + raise ValueError( + "RemoteGpuPipeline requires MFLUX_STUDIO_GPU_TOKEN." + ) + if config.backend not in REMOTE_BACKENDS: + raise ValueError( + f"RemoteGpuPipeline cannot host local backend {config.backend!r}; " + f"expected one of {REMOTE_BACKENDS}." + ) + self.config = config + self._backend: Backend = config.backend + self._client = httpx.Client( + base_url=config.gpu_host, + headers={"Authorization": f"Bearer {config.gpu_token}"}, + timeout=httpx.Timeout( + connect=float(config.gpu_connect_timeout_seconds), + read=float(config.gpu_timeout_seconds), + write=float(config.gpu_connect_timeout_seconds), + pool=float(config.gpu_connect_timeout_seconds), + ), + ) + self.last_peak_memory_mb: float | None = None + self.last_swap_seconds: float | None = 0.0 + log.info( + "RemoteGpuPipeline ready host=%s backend=%s read_timeout=%ds", + config.gpu_host, + self._backend, + config.gpu_timeout_seconds, + ) + + @property + def backend(self) -> Backend: + return self._backend + + @property + def model_path(self) -> str | None: + return None + + def ensure_backend(self, *, backend: Backend, model_path: str | None) -> None: + # Why: remote arm holds no in-process model, so "swap" is a label change. + if backend not in REMOTE_BACKENDS: + raise ValueError( + f"Server started in remote-GPU mode; backend {backend!r} is local-only. " + f"Restart with MFLUX_STUDIO_DEFAULT_BACKEND set to one of {LOCAL_BACKENDS}." + ) + self._backend = backend + + def generate_png( + self, + *, + prompt: str, + seed: int = DEFAULT_SEED, + steps: int = DEFAULT_STEPS, + height: int = DEFAULT_HEIGHT, + width: int = DEFAULT_WIDTH, + guidance: float = DEFAULT_GUIDANCE, + tiled_vae: bool | None = None, + max_sequence_length: int | None = None, + ) -> bytes: + body: dict[str, object] = { + "prompt": prompt, + "seed": seed, + "steps": steps, + "height": height, + "width": width, + "guidance": guidance, + "backend": self._backend, + } + if tiled_vae is not None: + body["tiled_vae"] = tiled_vae + if max_sequence_length is not None: + body["max_sequence_length"] = max_sequence_length + log.info("remote-gpu generate backend=%s size=%dx%d", self._backend, width, height) + response = self._client.post("/generate", json=body) + if response.status_code != 200: + detail: str + try: + detail = str(response.json().get("detail", response.text)) + except Exception: + detail = response.text + raise RuntimeError( + f"Remote GPU /generate returned {response.status_code}: {detail}" + ) + peak_header = response.headers.get("X-Peak-Memory-MB") + try: + self.last_peak_memory_mb = float(peak_header) if peak_header is not None else None + except ValueError: + self.last_peak_memory_mb = None + return response.content + + def close(self) -> None: + self._client.close() + + +def make_pipeline(config: PipelineConfig) -> "FluxPipeline | RemoteGpuPipeline": + """Factory: pick FluxPipeline (local MLX) or RemoteGpuPipeline (HTTP) by backend prefix.""" + if config.backend in REMOTE_BACKENDS: + return RemoteGpuPipeline(config) + return FluxPipeline(config) diff --git a/image-studio/backend/server.py b/image-studio/backend/server.py new file mode 100644 index 000000000..f1303535b --- /dev/null +++ b/image-studio/backend/server.py @@ -0,0 +1,318 @@ +from __future__ import annotations + +import asyncio +import base64 +import gc +import logging +import os +import time +from contextlib import AsyncExitStack, asynccontextmanager + +import httpx +from fastapi import FastAPI, HTTPException, Response +import mlx.core as mx +from pydantic import BaseModel, Field, model_validator +from pydantic.json_schema import SkipJsonSchema + +from backend.pipeline import ( + BACKEND_TO_FAMILY, + BACKEND_TO_KIND, + BACKENDS, + LOCAL_BACKENDS, + MODEL_FAMILIES, + REMOTE_BACKENDS, + Backend, + BackendKind, + DEFAULT_GUIDANCE, + DEFAULT_HEIGHT, + DEFAULT_SEED, + DEFAULT_STEPS, + DEFAULT_WIDTH, + FluxPipeline, + ModelFamily, + PipelineConfig, + RemoteGpuPipeline, + make_pipeline, +) + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", +) + +log = logging.getLogger(__name__) + + +def _parse_truthy(value: str | None) -> bool: + return value is not None and value.strip().lower() in {"1", "true", "yes", "on"} + + +# Force-disable is read once at module load: server restart is the override path. +# Per-session UI overrides ride on the ?force_disable=1 query param instead. +_FORCE_DISABLE_GPU_AT_LOAD: bool = _parse_truthy(os.getenv("MFLUX_STUDIO_FORCE_DISABLE_GPU")) +_BACKENDS_PROBE_TTL_SECONDS: float = float( + os.getenv("MFLUX_STUDIO_BACKENDS_PROBE_TTL_SECONDS", "30") +) +# Cache key is (effective_force_disable, pipeline_kind). The pipeline_kind is +# fixed per process today, but keying on it keeps the cache correct if a +# future change ever flips it mid-lifetime. +_backends_cache: dict[tuple[bool, BackendKind], tuple[float, dict]] = {} + + +def _clear_backends_cache() -> None: + _backends_cache.clear() + + +def _probe_gpu(host: str, token: str) -> tuple[bool, str | None]: + try: + resp = httpx.get( + f"{host.rstrip('/')}/healthz", + headers={"Authorization": f"Bearer {token}"}, + timeout=2.0, + ) + except httpx.HTTPError: + return False, "healthz_unreachable" + if resp.status_code == 200: + return True, None + return False, f"healthz_failed:{resp.status_code}" + + +def _resolve_backends( + force_disable: bool, pipeline_kind: BackendKind, current_backend: Backend +) -> dict: + """Report the relay's single resident kind + which model families it serves. + + The relay is configured for one kind per process — switching MLX↔gemlite + requires a restart. For the gemlite kind we still probe the remote GPU so + the frontend can surface a clear unhealthy state instead of empty errors. + """ + if pipeline_kind == "gemlite": + gpu_host = os.getenv("MFLUX_STUDIO_GPU_HOST") + gpu_token = os.getenv("MFLUX_STUDIO_GPU_TOKEN") + if force_disable: + healthy, reason = False, "force_disabled" + elif not gpu_host: + healthy, reason = False, "no_gpu_host" + elif not gpu_token: + healthy, reason = False, "no_gpu_token" + else: + healthy, reason = _probe_gpu(gpu_host, gpu_token) + else: + healthy, reason = True, None + + kind_backends = [b for b in BACKENDS if BACKEND_TO_KIND[b] == pipeline_kind] + supported_families: list[ModelFamily] = [ + f for f in MODEL_FAMILIES if any(BACKEND_TO_FAMILY[b] == f for b in kind_backends) + ] + default_family = BACKEND_TO_FAMILY[current_backend] + + return { + "kind": pipeline_kind, + "supported_families": supported_families, + "default_family": default_family, + "healthy": healthy, + "reason": reason, + } + + +def _get_backends_payload( + force_disable_query: bool, pipeline_kind: BackendKind, current_backend: Backend +) -> dict: + effective = _FORCE_DISABLE_GPU_AT_LOAD or force_disable_query + cache_key = (effective, pipeline_kind) + cached = _backends_cache.get(cache_key) + now = time.monotonic() + if cached is not None and (now - cached[0]) < _BACKENDS_PROBE_TTL_SECONDS: + return cached[1] + payload = _resolve_backends( + force_disable=effective, + pipeline_kind=pipeline_kind, + current_backend=current_backend, + ) + _backends_cache[cache_key] = (now, payload) + return payload + + +class GenerateRequest(BaseModel): + prompt: str = Field(min_length=1) + seed: int = DEFAULT_SEED + steps: int = Field(default=DEFAULT_STEPS, ge=1) + guidance: float = Field(default=DEFAULT_GUIDANCE, ge=0.0) + backend: Backend | SkipJsonSchema[None] = Field(default=None) + height: int = Field(default=DEFAULT_HEIGHT, ge=16) + width: int = Field(default=DEFAULT_WIDTH, ge=16) + model_path: str | None = Field(default=None) + tiled_vae: bool | None = Field(default=None) + max_sequence_length: int | None = Field(default=None, ge=1) + + +@asynccontextmanager +async def lifespan(app: FastAPI): + pipeline = make_pipeline(PipelineConfig.from_env()) + app.state.pipeline = pipeline + app.state.swap_lock = asyncio.Lock() + try: + yield + finally: + if isinstance(pipeline, RemoteGpuPipeline): + pipeline.close() + + +app = FastAPI(lifespan=lifespan) + + +@app.get("/backends") +async def get_backends(force_disable: bool = False) -> dict: + pipeline: FluxPipeline | RemoteGpuPipeline = app.state.pipeline + pipeline_kind: BackendKind = "gemlite" if pipeline.is_remote else "mlx" + return _get_backends_payload( + force_disable_query=force_disable, + pipeline_kind=pipeline_kind, + current_backend=pipeline.backend, + ) + + +@app.post( + "/generate", + response_class=Response, + responses={ + 200: { + "content": { + "image/png": { + "schema": { + "type": "string", + "format": "binary", + } + } + }, + "description": "Generated PNG image.", + } + }, +) +async def generate(request: GenerateRequest) -> Response: + pipeline: FluxPipeline | RemoteGpuPipeline = app.state.pipeline + lock: asyncio.Lock = app.state.swap_lock + target_backend: Backend = request.backend if request.backend is not None else pipeline.backend + if target_backend not in BACKENDS: + raise HTTPException(status_code=400, detail=f"Unknown backend {target_backend!r}.") + + async with AsyncExitStack() as stack: + if not pipeline.is_remote: + # Why: lock guards in-process MLX swap; remote arm has no resident model so concurrency is safe. + await stack.enter_async_context(lock) + try: + pipeline.ensure_backend(backend=target_backend, model_path=request.model_path) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + try: + gen_start = time.perf_counter() + image_bytes = pipeline.generate_png( + prompt=request.prompt, + seed=request.seed, + steps=request.steps, + height=request.height, + width=request.width, + guidance=request.guidance, + tiled_vae=request.tiled_vae, + max_sequence_length=request.max_sequence_length, + ) + wall_seconds = time.perf_counter() - gen_start + headers = {"X-Wall-Seconds": f"{wall_seconds:.3f}"} + if pipeline.last_peak_memory_mb is not None: + headers["X-Peak-Memory-MB"] = f"{pipeline.last_peak_memory_mb:.1f}" + return Response( + content=image_bytes, + media_type="image/png", + headers=headers, + ) + finally: + if not pipeline.is_remote: + mx.clear_cache() + gc.collect() + + +class CompareRequest(BaseModel): + prompt: str = Field(min_length=1) + seed: int = DEFAULT_SEED + steps: int = Field(default=DEFAULT_STEPS, ge=1) + guidance: float = Field(default=DEFAULT_GUIDANCE, ge=0.0) + height: int = Field(default=DEFAULT_HEIGHT, ge=16) + width: int = Field(default=DEFAULT_WIDTH, ge=16) + # Why: cross-arm compare is incoherent (one resident pipeline arm at a time); + # default to the three MLX backends so legacy callers behave unchanged. + backends: list[Backend] = Field(default_factory=lambda: list(LOCAL_BACKENDS)) + tiled_vae: bool | None = Field(default=None) + max_sequence_length: int | None = Field(default=None, ge=1) + + @model_validator(mode="after") + def _validate_backends(self) -> "CompareRequest": + if not self.backends: + raise ValueError("backends must contain at least one entry.") + unknown = [b for b in self.backends if b not in BACKENDS] + if unknown: + raise ValueError(f"Unknown backend(s): {unknown}; expected subset of {list(BACKENDS)}.") + if len(set(self.backends)) != len(self.backends): + raise ValueError("backends must not contain duplicates.") + return self + + +@app.post("/generate/compare") +async def generate_compare(request: CompareRequest) -> dict: + pipeline: FluxPipeline | RemoteGpuPipeline = app.state.pipeline + lock: asyncio.Lock = app.state.swap_lock + + results = [] + async with AsyncExitStack() as stack: + if not pipeline.is_remote: + # Why: same as /generate — lock guards in-process MLX swap; remote arm holds no model. + await stack.enter_async_context(lock) + try: + for target_backend in request.backends: + swap_start = time.perf_counter() + try: + pipeline.ensure_backend(backend=target_backend, model_path=None) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + swap_seconds = time.perf_counter() - swap_start + + gen_start = time.perf_counter() + image_bytes = pipeline.generate_png( + prompt=request.prompt, + seed=request.seed, + steps=request.steps, + height=request.height, + width=request.width, + guidance=request.guidance, + tiled_vae=request.tiled_vae, + max_sequence_length=request.max_sequence_length, + ) + wall_seconds = time.perf_counter() - gen_start + + results.append( + { + "backend": target_backend, + "png_b64": base64.b64encode(image_bytes).decode("ascii"), + "wall_seconds": wall_seconds, + "swap_seconds": swap_seconds, + } + ) + if not pipeline.is_remote: + mx.clear_cache() + gc.collect() + finally: + if not pipeline.is_remote: + mx.clear_cache() + gc.collect() + + return {"results": results} + + +__all__ = [ + "app", + "GenerateRequest", + "CompareRequest", + "generate", + "generate_compare", + "get_backends", +] diff --git a/image-studio/backend/tests/__init__.py b/image-studio/backend/tests/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/image-studio/backend/tests/__init__.py @@ -0,0 +1 @@ + diff --git a/image-studio/backend/tests/smoke_checkpoint_compare.py b/image-studio/backend/tests/smoke_checkpoint_compare.py new file mode 100644 index 000000000..d58cb923d --- /dev/null +++ b/image-studio/backend/tests/smoke_checkpoint_compare.py @@ -0,0 +1,97 @@ +"""Side-by-side checkpoint comparison on bonsai-ternary-mlx. + +Same prompt + seed + steps + size against two checkpoint paths. Saves PNGs +and computes pixel statistics: a noise output has approximately uniform +per-channel mean ~127 and high std; a real image has skewed distributions. + +Run: + .venv/bin/python -m backend.tests.smoke_checkpoint_compare +""" +from __future__ import annotations + +import gc +import io +import os +import time + +import mlx.core as mx +import numpy as np +from PIL import Image + +from backend.pipeline import FluxPipeline, PipelineConfig + +PROMPT = ( + "A serene Scandinavian woman in her 40s, soft window light, natural skin texture, " + "shallow depth of field, 50mm portrait lens, neutral linen backdrop." +) +SEED = 42 +STEPS = 4 +H = W = 1024 + +CHECKPOINT_A = os.environ.get("BONSAI_SMOKE_CHECKPOINT_A", "/tmp/bonsai-checkpoints/a") +CHECKPOINT_B = os.environ.get("BONSAI_SMOKE_CHECKPOINT_B", "/tmp/bonsai-checkpoints/b") + +OUT_DIR = os.environ.get("BONSAI_SMOKE_OUT_DIR", "/tmp/bonsai_smoke_diag") + + +def _stats(png_bytes: bytes) -> dict: + img = Image.open(io.BytesIO(png_bytes)).convert("RGB") + arr = np.asarray(img, dtype=np.float32) + return { + "shape": arr.shape, + "mean": tuple(float(arr.mean(axis=(0, 1))[i]) for i in range(3)), + "std": tuple(float(arr.std(axis=(0, 1))[i]) for i in range(3)), + "min": tuple(float(arr.min(axis=(0, 1))[i]) for i in range(3)), + "max": tuple(float(arr.max(axis=(0, 1))[i]) for i in range(3)), + } + + +def _run(label: str, model_path: str) -> None: + print(f"\n=== {label}: {model_path} ===", flush=True) + gc.collect() + mx.clear_cache() + + t0 = time.perf_counter() + pipe = FluxPipeline(PipelineConfig( + backend="bonsai-ternary-mlx", + baked_model_path=model_path, + te_4bit=True, + evict_text_encoder=True, + evict_transformer=False, + evict_vae=False, + )) + init_s = time.perf_counter() - t0 + + t0 = time.perf_counter() + png = pipe.generate_png(prompt=PROMPT, seed=SEED, steps=STEPS, height=H, width=W) + gen_s = time.perf_counter() - t0 + + out = f"{OUT_DIR}/{label}.png" + with open(out, "wb") as f: + f.write(png) + + s = _stats(png) + print( + f"init_s={init_s:.2f} gen_s={gen_s:.2f} -> {out}\n" + f" shape={s['shape']}\n" + f" mean R/G/B = ({s['mean'][0]:.1f}, {s['mean'][1]:.1f}, {s['mean'][2]:.1f})\n" + f" std R/G/B = ({s['std'][0]:.1f}, {s['std'][1]:.1f}, {s['std'][2]:.1f})\n" + f" min R/G/B = ({s['min'][0]:.0f}, {s['min'][1]:.0f}, {s['min'][2]:.0f})\n" + f" max R/G/B = ({s['max'][0]:.0f}, {s['max'][1]:.0f}, {s['max'][2]:.0f})", + flush=True, + ) + + del pipe + gc.collect() + mx.clear_cache() + + +def main() -> None: + import os + os.makedirs(OUT_DIR, exist_ok=True) + _run("checkpoint_a", CHECKPOINT_A) + _run("checkpoint_b", CHECKPOINT_B) + + +if __name__ == "__main__": + main() diff --git a/image-studio/backend/tests/smoke_memory_flatline.py b/image-studio/backend/tests/smoke_memory_flatline.py new file mode 100644 index 000000000..78c0ed632 --- /dev/null +++ b/image-studio/backend/tests/smoke_memory_flatline.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +# Manual diagnostic: +# cd +# .venv/bin/python -m backend.tests.smoke_memory_flatline -v + +import asyncio +import resource +import unittest +import warnings + +import mlx.core as mx +from tqdm import tqdm + +from backend.pipeline import FluxPipeline, PipelineConfig +from backend.server import app +from backend.server import GenerateRequest, generate + +GENERATIONS = 20 +HEIGHT = 1024 +WIDTH = 1024 +DRIFT_TOLERANCE_MIB = 50.0 +MONOTONIC_STEP_TOLERANCE_MIB = 5.0 +PROMPT = ( + "A mossy bonsai tree arranged as a premium studio product photograph, " + "soft daylight, stone pedestal, detailed needles, calm editorial composition." +) + + +def bytes_to_mib(value: int) -> float: + return value / (1024 * 1024) + + +class SmokeMemoryFlatline(unittest.TestCase): + def test_active_memory_flatlines(self) -> None: + warnings.filterwarnings("ignore", message="mx\\.metal\\.clear_cache is deprecated.*") + tqdm.monitor_interval = 0 + post_gen_active_mib: list[float] = [] + app.state.pipeline = FluxPipeline(PipelineConfig.from_env()) + + try: + for generation in range(1, GENERATIONS + 1): + mx.reset_peak_memory() + before_active = mx.get_active_memory() + before_peak = mx.get_peak_memory() + before_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + + response = asyncio.run( + generate( + GenerateRequest( + prompt=PROMPT, + seed=42 + generation, + steps=4, + height=HEIGHT, + width=WIDTH, + ) + ) + ) + self.assertEqual(response.status_code, 200) + + after_active = mx.get_active_memory() + after_peak = mx.get_peak_memory() + after_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + after_active_mib = bytes_to_mib(after_active) + post_gen_active_mib.append(after_active_mib) + + print( + "gen=" + f"{generation:02d} " + f"before_active_mib={bytes_to_mib(before_active):.2f} " + f"before_peak_mib={bytes_to_mib(before_peak):.2f} " + f"post_active_mib={after_active_mib:.2f} " + f"post_peak_mib={bytes_to_mib(after_peak):.2f} " + f"rss={after_rss} " + f"rss_delta={after_rss - before_rss}" + ) + finally: + del app.state.pipeline + + drift_mib = post_gen_active_mib[-1] - post_gen_active_mib[0] + monotonic_rise = ( + all( + current >= previous - MONOTONIC_STEP_TOLERANCE_MIB + for previous, current in zip(post_gen_active_mib, post_gen_active_mib[1:]) + ) + and drift_mib > MONOTONIC_STEP_TOLERANCE_MIB + ) + + self.assertLessEqual( + abs(drift_mib), + DRIFT_TOLERANCE_MIB, + f"post-gen active memory drifted by {drift_mib:.2f} MiB", + ) + self.assertFalse( + monotonic_rise, + f"post-gen active memory kept climbing: {post_gen_active_mib}", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/image-studio/backend/tests/smoke_pareto_sweep.py b/image-studio/backend/tests/smoke_pareto_sweep.py new file mode 100644 index 000000000..ea0a8738c --- /dev/null +++ b/image-studio/backend/tests/smoke_pareto_sweep.py @@ -0,0 +1,134 @@ +"""Pareto sweep: (peak MB, wall seconds) across TE / transformer / VAE eviction knobs. + +Matrix (bonsai-ternary-mlx, 1024², 4 steps, same seed+prompt): + + cfg_id | te_4bit | evict_te | evict_tx | evict_vae | notes + bl_bf16 F T F F bf16 baseline (flag-off) + pin_4bit T F F F pinned 4-bit TE (prior) + ff T T F F current default + ft T T F T VAE-only post-decode eviction + tf T T T F Klein bundles both; reload both next gen + tt T T T T Klein bundles both (equiv to tf) + +Per config: fresh pipeline, one cold gen, one warm gen (same prompt, cache hit). +Records peak_mb and gen_s each. Pipeline destroyed between configs so peak +memory is isolated to that config. + +Run: + .venv/bin/python -m backend.tests.smoke_pareto_sweep +""" +from __future__ import annotations + +import gc +import hashlib +import time +from dataclasses import dataclass + +import mlx.core as mx + +from backend.pipeline import FluxPipeline, PipelineConfig + +PROMPT = ( + "A serene Scandinavian woman in her 40s, soft window light, natural skin texture, " + "shallow depth of field, 50mm portrait lens, neutral linen backdrop." +) +SEED = 42 +STEPS = 4 +H = W = 1024 + + +@dataclass(frozen=True) +class Cfg: + cfg_id: str + te_4bit: bool + evict_te: bool + evict_tx: bool + evict_vae: bool + notes: str + + +CONFIGS = [ + Cfg("bl_bf16", False, True, False, False, "bf16 baseline (flag-off)"), + Cfg("pin_4bit", True, False, False, False, "pinned 4-bit TE"), + Cfg("ff", True, True, False, False, "current default"), + Cfg("ft", True, True, False, True, "VAE-only eviction"), + Cfg("tf", True, True, True, False, "Klein bundles tx+vae"), + Cfg("tt", True, True, True, True, "Klein bundles tx+vae"), +] + + +def _reset_peak() -> None: + gc.collect() + mx.clear_cache() + mx.reset_peak_memory() + + +def _gen(pipe: FluxPipeline) -> tuple[float, float, str]: + _reset_peak() + t0 = time.perf_counter() + png = pipe.generate_png(prompt=PROMPT, seed=SEED, steps=STEPS, height=H, width=W) + gen_s = time.perf_counter() - t0 + peak_mb = mx.get_peak_memory() / (1024 * 1024) + sha = hashlib.sha256(png).hexdigest()[:12] + return gen_s, peak_mb, sha + + +def _run_cfg(cfg: Cfg) -> dict: + print(f"\n=== {cfg.cfg_id}: {cfg.notes} (te4bit={cfg.te_4bit} " + f"evict_te={cfg.evict_te} evict_tx={cfg.evict_tx} evict_vae={cfg.evict_vae}) ===", + flush=True) + gc.collect() + mx.clear_cache() + + t0 = time.perf_counter() + pipe = FluxPipeline(PipelineConfig( + backend="bonsai-ternary-mlx", + te_4bit=cfg.te_4bit, + evict_text_encoder=cfg.evict_te, + evict_transformer=cfg.evict_tx, + evict_vae=cfg.evict_vae, + )) + init_s = time.perf_counter() - t0 + + cold_s, cold_peak, cold_sha = _gen(pipe) + warm_s, warm_peak, warm_sha = _gen(pipe) + byte_identical = (cold_sha == warm_sha) + + del pipe + gc.collect() + mx.clear_cache() + + return { + "cfg": cfg, + "init_s": init_s, + "cold_s": cold_s, + "cold_peak_mb": cold_peak, + "warm_s": warm_s, + "warm_peak_mb": warm_peak, + "byte_identical_cold_warm": byte_identical, + "sha": warm_sha, + } + + +def main() -> None: + results = [_run_cfg(c) for c in CONFIGS] + + print("\n=== SUMMARY ===", flush=True) + print( + f"{'cfg_id':10s} {'te4':>3s} {'evTE':>4s} {'evTX':>4s} {'evVA':>4s} " + f"{'init_s':>7s} {'cold_s':>7s} {'cold_MB':>8s} {'warm_s':>7s} {'warm_MB':>8s} identical", + flush=True, + ) + for r in results: + c = r["cfg"] + print( + f"{c.cfg_id:10s} {int(c.te_4bit):>3d} {int(c.evict_te):>4d} " + f"{int(c.evict_tx):>4d} {int(c.evict_vae):>4d} " + f"{r['init_s']:>7.2f} {r['cold_s']:>7.2f} {r['cold_peak_mb']:>8.1f} " + f"{r['warm_s']:>7.2f} {r['warm_peak_mb']:>8.1f} {r['byte_identical_cold_warm']}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/image-studio/backend/tests/smoke_te_4bit_ab.py b/image-studio/backend/tests/smoke_te_4bit_ab.py new file mode 100644 index 000000000..a0fd1ad41 --- /dev/null +++ b/image-studio/backend/tests/smoke_te_4bit_ab.py @@ -0,0 +1,81 @@ +"""A/B: bf16 TE vs 4bit TE, same seed + prompt, both on bonsai-ternary-mlx backend. + +Run: + .venv/bin/python -m backend.tests.smoke_te_4bit_ab +""" +from __future__ import annotations + +import gc +import hashlib +import os +import time +from pathlib import Path + +import mlx.core as mx + +from backend.pipeline import FluxPipeline, PipelineConfig + +PROMPT = ( + "A serene Scandinavian woman in her 40s, soft window light, natural skin texture, " + "shallow depth of field, 50mm portrait lens, neutral linen backdrop." +) +SEED = 42 +STEPS = 4 +H = W = 1024 +OUT_DIR = Path(os.environ.get("BONSAI_SMOKE_OUT_DIR") or Path(__file__).parent / "_ab_out") + + +def run(label: str, *, te_4bit: bool) -> dict: + print(f"\n=== {label} (te_4bit={te_4bit}) ===", flush=True) + gc.collect() + mx.clear_cache() + mx.reset_peak_memory() + + t0 = time.perf_counter() + pipe = FluxPipeline(PipelineConfig(backend="bonsai-ternary-mlx", te_4bit=te_4bit)) + load_s = time.perf_counter() - t0 + + t1 = time.perf_counter() + png = pipe.generate_png(prompt=PROMPT, seed=SEED, steps=STEPS, height=H, width=W) + gen_s = time.perf_counter() - t1 + + peak_mb = mx.get_peak_memory() / (1024 * 1024) + out_path = OUT_DIR / f"{label}.png" + out_path.write_bytes(png) + digest = hashlib.sha256(png).hexdigest()[:16] + + # Dispose pipeline between runs so peak-memory resets cleanly. + del pipe + gc.collect() + mx.clear_cache() + + return { + "label": label, + "te_4bit": te_4bit, + "load_s": load_s, + "gen_s": gen_s, + "peak_mb": peak_mb, + "png_path": str(out_path), + "sha256_16": digest, + } + + +def main() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + results = [ + run("te_bf16", te_4bit=False), + run("te_4bit", te_4bit=True), + ] + print("\n=== SUMMARY ===", flush=True) + for r in results: + print( + f"{r['label']:10s} load={r['load_s']:6.2f}s gen={r['gen_s']:6.2f}s " + f"peak_mb={r['peak_mb']:7.1f} sha={r['sha256_16']} -> {r['png_path']}", + flush=True, + ) + peak_delta = results[0]["peak_mb"] - results[1]["peak_mb"] + print(f"\npeak_mb delta (bf16 - 4bit) = {peak_delta:+.1f} MB", flush=True) + + +if __name__ == "__main__": + main() diff --git a/image-studio/backend/tests/smoke_te_4bit_eviction.py b/image-studio/backend/tests/smoke_te_4bit_eviction.py new file mode 100644 index 000000000..acf778053 --- /dev/null +++ b/image-studio/backend/tests/smoke_te_4bit_eviction.py @@ -0,0 +1,125 @@ +"""Validate eviction + cache-miss reload with te_4bit=True. + +Sequence: + 1. Cold: generate prompt A. TE loads, encodes, evicts. Peak includes TE. + 2. Warm cache hit: generate prompt A again. TE stays None. Peak should drop. + 3. Cache miss: generate prompt B. Patched reload fires (4-bit path), encodes, evicts. + +Also directly times load_te_4bit() in isolation to report the raw reload wall-time +from cold (after the first run evicts). + +Run: + .venv/bin/python -m backend.tests.smoke_te_4bit_eviction +""" +from __future__ import annotations + +import gc +import hashlib +import os +import time +from pathlib import Path + +import mlx.core as mx + +from backend.pipeline import FluxPipeline, PipelineConfig +from backend.text_encoder_4bit import load_te_4bit + +PROMPT_A = ( + "A serene Scandinavian woman in her 40s, soft window light, natural skin texture, " + "shallow depth of field, 50mm portrait lens, neutral linen backdrop." +) +PROMPT_B = ( + "A rusting steam locomotive abandoned in a coastal fog at dawn, " + "wet cobblestones, muted teal palette, cinematic wide shot, 35mm film grain." +) +SEED = 42 +STEPS = 4 +H = W = 1024 +OUT_DIR = Path(os.environ.get("BONSAI_SMOKE_OUT_DIR") or Path(__file__).parent / "_te_evict_out") + + +def _reset_peak() -> None: + gc.collect() + mx.clear_cache() + mx.reset_peak_memory() + + +def _generate(pipe: FluxPipeline, label: str, prompt: str) -> dict: + _reset_peak() + t0 = time.perf_counter() + png = pipe.generate_png(prompt=prompt, seed=SEED, steps=STEPS, height=H, width=W) + gen_s = time.perf_counter() - t0 + peak_mb = mx.get_peak_memory() / (1024 * 1024) + out_path = OUT_DIR / f"{label}.png" + out_path.write_bytes(png) + digest = hashlib.sha256(png).hexdigest()[:16] + te_is_none = pipe._model.text_encoder is None + return { + "label": label, + "gen_s": gen_s, + "peak_mb": peak_mb, + "png_path": str(out_path), + "sha": digest, + "te_is_none_after": te_is_none, + } + + +def main() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + + print("\n=== Init pipeline (te_4bit=True, evict_text_encoder=True) ===", flush=True) + t0 = time.perf_counter() + pipe = FluxPipeline(PipelineConfig(backend="bonsai-ternary-mlx", te_4bit=True, evict_text_encoder=True)) + init_s = time.perf_counter() - t0 + assert pipe._model._evict_text_encoder is True, "eviction must be on" + assert pipe._model._studio_te_4bit is True, "4-bit marker must be set" + print(f" init_s={init_s:.2f} te_4bit={pipe.config.te_4bit} evict_te={pipe.config.evict_text_encoder}", flush=True) + + results = [] + + print("\n=== Run 1: cold generate prompt A ===", flush=True) + results.append(_generate(pipe, "cold_A", PROMPT_A)) + + print("\n=== Run 2: cache-hit generate prompt A (TE should stay None) ===", flush=True) + assert pipe._model.text_encoder is None, "TE must be evicted after run 1" + results.append(_generate(pipe, "hit_A", PROMPT_A)) + assert results[-1]["te_is_none_after"], "TE should remain None through cache hit" + + print("\n=== Run 3: cache-miss generate prompt B (patched 4-bit reload) ===", flush=True) + assert pipe._model.text_encoder is None, "TE must still be evicted before run 3" + results.append(_generate(pipe, "miss_B", PROMPT_B)) + + print("\n=== Direct timing: load_te_4bit() in isolation ===", flush=True) + _reset_peak() + overrides = pipe._model.model_config.text_encoder_overrides + t0 = time.perf_counter() + te_probe = load_te_4bit(overrides) + mx.eval(te_probe.embed_tokens.weight) + te_load_s = time.perf_counter() - t0 + del te_probe + gc.collect() + mx.clear_cache() + print(f" load_te_4bit cold = {te_load_s:.3f}s", flush=True) + + print("\n=== SUMMARY ===", flush=True) + for r in results: + print( + f"{r['label']:10s} gen={r['gen_s']:6.2f}s peak_mb={r['peak_mb']:7.1f} " + f"te_none_after={r['te_is_none_after']} sha={r['sha']}", + flush=True, + ) + print(f"\nTE 4-bit cold load wall-time (isolated) = {te_load_s:.3f}s", flush=True) + print( + f"peak_mb: cold_A={results[0]['peak_mb']:.0f} " + f"hit_A={results[1]['peak_mb']:.0f} " + f"miss_B={results[2]['peak_mb']:.0f}", + flush=True, + ) + print( + f"Δ hit_A vs cold_A = {results[1]['peak_mb'] - results[0]['peak_mb']:+.1f} MB", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/image-studio/backend/tests/smoke_three_backends.py b/image-studio/backend/tests/smoke_three_backends.py new file mode 100644 index 000000000..d607af89d --- /dev/null +++ b/image-studio/backend/tests/smoke_three_backends.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +# Integration smoke: hit /generate serially with each of the three backends and +# measure hot-swap wall time. Constructs real Flux2Klein models; runtime is +# dominated by transformer reloads + MLX JIT. +# +# .venv/bin/python -m backend.tests.smoke_three_backends -v + +import asyncio +import io +import time +import unittest +import warnings + +from PIL import Image + +from backend.pipeline import BACKENDS, FluxPipeline, PipelineConfig +from backend.server import GenerateRequest, app, generate + +PROMPT = "A mossy bonsai tree in a sunlit studio, editorial composition." +HEIGHT = 512 +WIDTH = 512 +STEPS = 4 +SEED = 42 + + +class SmokeThreeBackends(unittest.TestCase): + def test_generate_each_backend(self) -> None: + warnings.filterwarnings("ignore", message="mx\\.metal\\.clear_cache is deprecated.*") + config = PipelineConfig.from_env() + app.state.pipeline = FluxPipeline(config) + app.state.swap_lock = asyncio.Lock() + timings: dict[str, dict[str, float]] = {} + try: + for backend in BACKENDS: + pre_swap_start = time.perf_counter() + gen_start = pre_swap_start + response = asyncio.run( + generate( + GenerateRequest( + prompt=PROMPT, + seed=SEED, + steps=STEPS, + height=HEIGHT, + width=WIDTH, + backend=backend, + ) + ) + ) + gen_elapsed = time.perf_counter() - gen_start + self.assertEqual(response.status_code, 200) + image = Image.open(io.BytesIO(response.body)) + image.verify() + swap_seconds = app.state.pipeline.last_swap_seconds or 0.0 + timings[backend] = { + "generate_total_s": gen_elapsed, + "swap_s": swap_seconds, + } + print( + f"backend={backend} swap_s={swap_seconds:.3f} " + f"generate_total_s={gen_elapsed:.3f} img={image.size}" + ) + finally: + del app.state.pipeline + del app.state.swap_lock + + print("SMOKE_SUMMARY " + str(timings)) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/image-studio/backend/tests/test_backends_endpoint.py b/image-studio/backend/tests/test_backends_endpoint.py new file mode 100644 index 000000000..4ae9f6c31 --- /dev/null +++ b/image-studio/backend/tests/test_backends_endpoint.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +# Unit tests for GET /backends. _build_model is patched (lifespan still +# instantiates a FluxPipeline) and httpx.get is patched for the GPU probe. +# +# .venv/bin/python -m unittest backend.tests.test_backends_endpoint -v + +import unittest +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +from fastapi.testclient import TestClient + +from backend import server +from backend.server import _clear_backends_cache, app + + +class _FakeModel: + def __init__(self, *, backend: str) -> None: + self.backend = backend + self.tiling_config = None + + +@contextmanager +def _patched_app(env: dict[str, str], force_disable_load: bool = False): + _clear_backends_cache() + # _FakeModel lacks model_config; the te_4bit reload path needs it. Tests + # opt out of 4-bit-TE here so FluxPipeline._load works against the stub. + full_env = {"MFLUX_STUDIO_TE_4BIT": "0", **env} + with patch.dict("os.environ", full_env, clear=False): + with patch.object(server, "_FORCE_DISABLE_GPU_AT_LOAD", force_disable_load): + with patch( + "backend.pipeline._build_model", + side_effect=lambda *, backend, **_kw: _FakeModel(backend=backend), + ): + with TestClient(app) as client: + yield client + + +def _ok_resp(status: int = 200) -> MagicMock: + m = MagicMock() + m.status_code = status + return m + + +_ALL_FAMILIES = ["bonsai-binary", "bonsai-ternary"] + + +class MlxKindTest(unittest.TestCase): + """Relay in MLX mode: GPU env is irrelevant, healthy is unconditionally true.""" + + def test_default_env_yields_mlx_kind(self) -> None: + with _patched_app({}) as client: + r = client.get("/backends") + self.assertEqual(r.status_code, 200) + body = r.json() + self.assertEqual(body["kind"], "mlx") + self.assertEqual(body["supported_families"], _ALL_FAMILIES) + self.assertEqual(body["default_family"], "bonsai-ternary") + self.assertTrue(body["healthy"]) + self.assertIsNone(body["reason"]) + + def test_default_family_follows_env_backend(self) -> None: + env = {"MFLUX_STUDIO_DEFAULT_BACKEND": "bonsai-binary-mlx"} + with _patched_app(env) as client: + body = client.get("/backends").json() + self.assertEqual(body["kind"], "mlx") + self.assertEqual(body["default_family"], "bonsai-binary") + + def test_gpu_env_does_not_trigger_probe(self) -> None: + env = { + "MFLUX_STUDIO_GPU_HOST": "http://localhost:8801", + "MFLUX_STUDIO_GPU_TOKEN": "tok", + } + with _patched_app(env) as client, patch("backend.server.httpx.get") as probe: + body = client.get("/backends").json() + probe.assert_not_called() + self.assertEqual(body["kind"], "mlx") + self.assertTrue(body["healthy"]) + + +class GemliteKindTest(unittest.TestCase): + """Relay in gemlite mode: probes the remote GPU, healthy reflects probe result.""" + + BASE_ENV = { + "MFLUX_STUDIO_DEFAULT_BACKEND": "bonsai-ternary-gemlite", + "MFLUX_STUDIO_GPU_HOST": "http://localhost:8801", + "MFLUX_STUDIO_GPU_TOKEN": "tok", + } + + def test_healthy_probe(self) -> None: + with _patched_app(self.BASE_ENV) as client, patch( + "backend.server.httpx.get", return_value=_ok_resp(200) + ): + body = client.get("/backends").json() + self.assertEqual(body["kind"], "gemlite") + self.assertEqual(body["supported_families"], _ALL_FAMILIES) + self.assertEqual(body["default_family"], "bonsai-ternary") + self.assertTrue(body["healthy"]) + self.assertIsNone(body["reason"]) + + def test_healthz_non_200(self) -> None: + with _patched_app(self.BASE_ENV) as client, patch( + "backend.server.httpx.get", return_value=_ok_resp(503) + ): + body = client.get("/backends").json() + self.assertFalse(body["healthy"]) + self.assertEqual(body["reason"], "healthz_failed:503") + + def test_healthz_unreachable(self) -> None: + import httpx + with _patched_app(self.BASE_ENV) as client, patch( + "backend.server.httpx.get", side_effect=httpx.ConnectError("nope") + ): + body = client.get("/backends").json() + self.assertEqual(body["reason"], "healthz_unreachable") + + def test_force_disable_via_module_load(self) -> None: + with _patched_app(self.BASE_ENV, force_disable_load=True) as client, patch( + "backend.server.httpx.get", return_value=_ok_resp(200) + ) as probe: + body = client.get("/backends").json() + probe.assert_not_called() + self.assertFalse(body["healthy"]) + self.assertEqual(body["reason"], "force_disabled") + + def test_force_disable_via_query(self) -> None: + with _patched_app(self.BASE_ENV) as client, patch( + "backend.server.httpx.get", return_value=_ok_resp(200) + ): + healthy = client.get("/backends").json() + disabled = client.get("/backends?force_disable=1").json() + self.assertTrue(healthy["healthy"]) + self.assertFalse(disabled["healthy"]) + self.assertEqual(disabled["reason"], "force_disabled") + + def test_default_family_respects_env_backend(self) -> None: + env = {**self.BASE_ENV, "MFLUX_STUDIO_DEFAULT_BACKEND": "bonsai-binary-gemlite"} + with _patched_app(env) as client, patch( + "backend.server.httpx.get", return_value=_ok_resp(200) + ): + body = client.get("/backends").json() + self.assertEqual(body["default_family"], "bonsai-binary") + + def test_cache_avoids_reprobe(self) -> None: + with _patched_app(self.BASE_ENV) as client, patch( + "backend.server.httpx.get", return_value=_ok_resp(200) + ) as probe: + client.get("/backends") + client.get("/backends") + client.get("/backends") + self.assertEqual(probe.call_count, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/image-studio/backend/tests/test_ensure_backend.py b/image-studio/backend/tests/test_ensure_backend.py new file mode 100644 index 000000000..9573e74ff --- /dev/null +++ b/image-studio/backend/tests/test_ensure_backend.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +# Unit tests for ensure_backend swap logic. +# No MLX model is constructed — FluxPipeline internals are patched. +# +# .venv/bin/python -m unittest backend.tests.test_ensure_backend -v + +import unittest +from unittest.mock import MagicMock, patch + +from mflux.models.common.vae.tiling_config import TilingConfig + +from backend.pipeline import ( + FluxPipeline, + PipelineConfig, + _default_model_path_for, + _resolve_tiling_config, +) +from backend.server import GenerateRequest + + +class _FakeModel: + def __init__(self, *, backend: str, model_path: str | None, config) -> None: + self.backend = backend + self.model_path = model_path + self.config = config + self.tiling_config = None + + +def _make_pipeline(default_backend: str = "bonsai-ternary-mlx") -> FluxPipeline: + config = PipelineConfig( + backend=default_backend, # type: ignore[arg-type] + baked_model_path="/tmp/baked", + baked_binary_model_path="/tmp/baked-binary", + te_4bit=False, # _FakeModel lacks model_config; skip the 4-bit-TE load path. + ) + with patch("backend.pipeline._build_model", side_effect=_FakeModel): + return FluxPipeline(config) + + +class EnsureBackendTest(unittest.TestCase): + def test_no_swap_when_backend_matches(self) -> None: + pipeline = _make_pipeline("bonsai-ternary-mlx") + initial_model = pipeline._model + with patch("backend.pipeline._build_model") as build: + pipeline.ensure_backend(backend="bonsai-ternary-mlx", model_path=None) + build.assert_not_called() + self.assertIs(pipeline._model, initial_model) + + def test_swap_to_binary_uses_binary_baked_path(self) -> None: + pipeline = _make_pipeline("bonsai-ternary-mlx") + with patch("backend.pipeline._build_model", side_effect=_FakeModel): + pipeline.ensure_backend(backend="bonsai-binary-mlx", model_path=None) + self.assertEqual(pipeline.backend, "bonsai-binary-mlx") + self.assertEqual(pipeline.model_path, "/tmp/baked-binary") + + def test_swap_back_to_ternary_uses_baked_path(self) -> None: + pipeline = _make_pipeline("bonsai-binary-mlx") + with patch("backend.pipeline._build_model", side_effect=_FakeModel) as build: + pipeline.ensure_backend(backend="bonsai-ternary-mlx", model_path=None) + build.assert_called_once() + self.assertEqual(pipeline.backend, "bonsai-ternary-mlx") + self.assertEqual(pipeline.model_path, "/tmp/baked") + + def test_override_model_path_triggers_swap(self) -> None: + pipeline = _make_pipeline("bonsai-ternary-mlx") + with patch("backend.pipeline._build_model", side_effect=_FakeModel) as build: + pipeline.ensure_backend(backend="bonsai-ternary-mlx", model_path="/tmp/custom") + build.assert_called_once() + self.assertEqual(pipeline.model_path, "/tmp/custom") + + def test_unknown_backend_rejected(self) -> None: + pipeline = _make_pipeline("bonsai-ternary-mlx") + with self.assertRaises(ValueError): + pipeline.ensure_backend(backend="bogus", model_path=None) # type: ignore[arg-type] + + +class DefaultModelPathTest(unittest.TestCase): + def _config(self) -> PipelineConfig: + return PipelineConfig( + baked_model_path="/tmp/baked", + baked_binary_model_path="/tmp/baked-binary", + ) + + def test_ternary_mlx_uses_baked(self) -> None: + self.assertEqual(_default_model_path_for("bonsai-ternary-mlx", self._config()), "/tmp/baked") + + def test_binary_mlx_uses_binary_baked(self) -> None: + self.assertEqual(_default_model_path_for("bonsai-binary-mlx", self._config()), "/tmp/baked-binary") + + def test_remote_backends_return_none(self) -> None: + cfg = self._config() + self.assertIsNone(_default_model_path_for("bonsai-ternary-gemlite", cfg)) + self.assertIsNone(_default_model_path_for("bonsai-binary-gemlite", cfg)) + + +class GenerateRequestDefaultsTest(unittest.TestCase): + def test_no_backend_defaults_to_none(self) -> None: + req = GenerateRequest(prompt="x") + self.assertIsNone(req.backend) + + +class ResolveTilingConfigTest(unittest.TestCase): + # Threshold follows 2 * TilingConfig().vae_decode_tile_size (128 -> 256). + def test_auto_on_threshold_enables_tiling(self) -> None: + threshold = 2 * TilingConfig().vae_decode_tile_size + cfg = _resolve_tiling_config( + request_override=None, server_default="auto", height=threshold, width=threshold + ) + self.assertIsInstance(cfg, TilingConfig) + + def test_auto_below_threshold_disables_tiling(self) -> None: + threshold = 2 * TilingConfig().vae_decode_tile_size + cfg = _resolve_tiling_config( + request_override=None, + server_default="auto", + height=threshold - 1, + width=threshold - 1, + ) + self.assertIsNone(cfg) + + def test_request_override_true_forces_tiling(self) -> None: + cfg = _resolve_tiling_config( + request_override=True, server_default="off", height=64, width=64 + ) + self.assertIsInstance(cfg, TilingConfig) + + def test_request_override_false_disables_tiling(self) -> None: + cfg = _resolve_tiling_config( + request_override=False, server_default="on", height=4096, width=4096 + ) + self.assertIsNone(cfg) + + +class GuidancePassThroughTest(unittest.TestCase): + # Verify the guidance field reaches Flux2Klein.generate_image. + def test_custom_guidance_reaches_model(self) -> None: + pipeline = _make_pipeline("bonsai-ternary-mlx") + + fake_generated = MagicMock() + image = MagicMock() + image.save = lambda buf, format: buf.write(b"\x89PNG\r\n\x1a\n") + fake_generated.image = image + + pipeline._model.generate_image = MagicMock(return_value=fake_generated) # type: ignore[attr-defined] + pipeline.generate_png(prompt="x", guidance=3.5) + + _, kwargs = pipeline._model.generate_image.call_args # type: ignore[attr-defined] + self.assertEqual(kwargs["guidance"], 3.5) + + def test_default_guidance_is_one(self) -> None: + pipeline = _make_pipeline("bonsai-ternary-mlx") + fake_generated = MagicMock() + image = MagicMock() + image.save = lambda buf, format: buf.write(b"\x89PNG\r\n\x1a\n") + fake_generated.image = image + pipeline._model.generate_image = MagicMock(return_value=fake_generated) # type: ignore[attr-defined] + pipeline.generate_png(prompt="x") + + _, kwargs = pipeline._model.generate_image.call_args # type: ignore[attr-defined] + self.assertEqual(kwargs["guidance"], 1.0) + + def test_request_guidance_defaults_to_one(self) -> None: + req = GenerateRequest(prompt="x") + self.assertEqual(req.guidance, 1.0) + + def test_request_guidance_custom(self) -> None: + req = GenerateRequest(prompt="x", guidance=4.2) + self.assertEqual(req.guidance, 4.2) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/image-studio/backend/tests/test_generate_compare.py b/image-studio/backend/tests/test_generate_compare.py new file mode 100644 index 000000000..5aedf053d --- /dev/null +++ b/image-studio/backend/tests/test_generate_compare.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +# Unit tests for POST /generate/compare. FluxPipeline internals are stubbed out — +# no MLX model is constructed. +# +# .venv/bin/python -m unittest backend.tests.test_generate_compare -v + +import base64 +import unittest +from contextlib import contextmanager +from unittest.mock import patch + +from fastapi.testclient import TestClient + +from backend.pipeline import BACKENDS, LOCAL_BACKENDS +from backend.server import CompareRequest, app + + +_PNG_MAGIC = b"\x89PNG\r\n\x1a\n" + + +def _fake_png(backend: str) -> bytes: + return _PNG_MAGIC + f"fake-{backend}".encode() + + +class _FakeModel: + # Stand-in for Flux2Klein inside FluxPipeline — only the surface used by + # generate_png is faked out. generate_image returns a PIL-ish stub whose + # save() writes a unique marker per backend. + def __init__(self, *, backend: str) -> None: + self.backend = backend + self.tiling_config = None + + def generate_image(self, **_kwargs): + backend = self.backend + + class _Image: + def save(self, buf, format): # noqa: A002 — FastAPI signature + buf.write(_fake_png(backend)) + + class _Generated: + image = _Image() + + return _Generated() + + +@contextmanager +def _patched_app(): + # Replace _build_model so FluxPipeline.__init__ (invoked by the lifespan) + # never constructs a real Flux2Klein — we get a lightweight stub instead. + # _FakeModel lacks model_config; opt out of 4-bit-TE so _load works. + with patch.dict("os.environ", {"MFLUX_STUDIO_TE_4BIT": "0"}, clear=False): + with patch("backend.pipeline._build_model", side_effect=lambda *, backend, **_kw: _FakeModel(backend=backend)): + with TestClient(app) as client: + yield client, client.app.state.pipeline + + +class CompareRequestValidationTest(unittest.TestCase): + def test_defaults_to_all_three_backends(self) -> None: + req = CompareRequest(prompt="x") + self.assertEqual(req.backends, list(LOCAL_BACKENDS)) + + def test_empty_backends_rejected(self) -> None: + with self.assertRaises(ValueError): + CompareRequest(prompt="x", backends=[]) + + def test_unknown_backend_rejected(self) -> None: + with self.assertRaises(ValueError): + CompareRequest(prompt="x", backends=["bogus"]) # type: ignore[list-item] + + def test_duplicates_rejected(self) -> None: + with self.assertRaises(ValueError): + CompareRequest(prompt="x", backends=["bonsai-ternary-mlx", "bonsai-ternary-mlx"]) + + def test_subset_accepted(self) -> None: + req = CompareRequest(prompt="x", backends=["bonsai-ternary-mlx", "bonsai-binary-mlx"]) + self.assertEqual(req.backends, ["bonsai-ternary-mlx", "bonsai-binary-mlx"]) + +class CompareEndpointTest(unittest.TestCase): + # End-to-end: POST /generate/compare with a stubbed _build_model, assert + # serialized backend iteration + per-backend result shape. + + def test_default_runs_all_local_backends(self) -> None: + with _patched_app() as (client, pipeline): + response = client.post("/generate/compare", json={"prompt": "a cat"}) + self.assertEqual(response.status_code, 200, response.text) + body = response.json() + self.assertEqual(len(body["results"]), len(LOCAL_BACKENDS)) + self.assertEqual([r["backend"] for r in body["results"]], list(LOCAL_BACKENDS)) + # Final backend is the last one in the request order. + self.assertEqual(pipeline.backend, LOCAL_BACKENDS[-1]) + + def test_result_shape(self) -> None: + with _patched_app() as (client, _): + response = client.post( + "/generate/compare", + json={"prompt": "a cat", "backends": ["bonsai-ternary-mlx", "bonsai-binary-mlx"]}, + ) + self.assertEqual(response.status_code, 200, response.text) + for result in response.json()["results"]: + self.assertIn(result["backend"], BACKENDS) + decoded = base64.b64decode(result["png_b64"]) + self.assertTrue(decoded.startswith(_PNG_MAGIC)) + self.assertTrue(decoded.endswith(f"fake-{result['backend']}".encode())) + self.assertIsInstance(result["wall_seconds"], float) + self.assertIsInstance(result["swap_seconds"], float) + self.assertGreaterEqual(result["wall_seconds"], 0.0) + self.assertGreaterEqual(result["swap_seconds"], 0.0) + + def test_preserves_order(self) -> None: + with _patched_app() as (client, _): + response = client.post( + "/generate/compare", + json={ + "prompt": "x", + "backends": ["bonsai-binary-mlx", "bonsai-ternary-mlx"], + }, + ) + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual( + [r["backend"] for r in response.json()["results"]], + ["bonsai-binary-mlx", "bonsai-ternary-mlx"], + ) + + def test_custom_params_reach_pipeline(self) -> None: + with _patched_app() as (client, pipeline): + captured: dict = {} + original = pipeline._model.generate_image # type: ignore[union-attr] + + def _spy(**kwargs): + captured.update(kwargs) + return original(**kwargs) + + # Replace every lazily-built _FakeModel's generate_image with a spy + # by patching _build_model to wrap the produced model. + with patch("backend.pipeline._build_model", side_effect=lambda *, backend, **_kw: _wrap_with_spy(_FakeModel(backend=backend), captured)): + response = client.post( + "/generate/compare", + json={ + "prompt": "detailed cat", + "seed": 123, + "steps": 8, + "guidance": 3.5, + "height": 1024, + "width": 768, + "backends": ["bonsai-binary-mlx"], + }, + ) + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(captured["seed"], 123) + self.assertEqual(captured["num_inference_steps"], 8) + self.assertEqual(captured["guidance"], 3.5) + self.assertEqual(captured["height"], 1024) + self.assertEqual(captured["width"], 768) + self.assertEqual(captured["prompt"], "detailed cat") + + def test_empty_backends_rejected_by_api(self) -> None: + with _patched_app() as (client, _): + response = client.post( + "/generate/compare", json={"prompt": "x", "backends": []} + ) + self.assertEqual(response.status_code, 422) + + +def _wrap_with_spy(model: _FakeModel, captured: dict) -> _FakeModel: + original = model.generate_image + + def _spy(**kwargs): + captured.update(kwargs) + return original(**kwargs) + + model.generate_image = _spy # type: ignore[method-assign] + return model + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/image-studio/backend/text_encoder_4bit.py b/image-studio/backend/text_encoder_4bit.py new file mode 100644 index 000000000..783a3bf4b --- /dev/null +++ b/image-studio/backend/text_encoder_4bit.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from pathlib import Path + +import mlx.core as mx +import mlx.nn as nn +from huggingface_hub import snapshot_download +from mlx.utils import tree_unflatten + +from mflux.models.flux2.flux2_initializer import Flux2Initializer +from mflux.models.flux2.model.flux2_text_encoder.qwen3_text_encoder import Qwen3TextEncoder + +TE_4BIT_REPO = "mlx-community/Qwen3-4B-4bit" +_QUANT_BITS = 4 +_QUANT_GROUP_SIZE = 64 + + +def load_te_4bit(text_encoder_overrides: dict, model_path: str | None = None) -> Qwen3TextEncoder: + # Prefer the bundled text_encoder-mlx-4bit/ under model_path when present. + # The Bonsai HF repos (prism-ml/bonsai-image-{ternary,binary}-4B-mlx-*bit) + # ship the exact mlx-community Qwen3-4B-4bit weights as a subdir, so using + # them locally avoids a ~3 GB HF download every fresh install AND keeps + # the demo self-contained (no implicit HF dependency for the TE). + local_dir = Path(model_path) / "text_encoder-mlx-4bit" if model_path else None + if local_dir and (local_dir / "model.safetensors").is_file(): + root = local_dir + else: + root = Path( + snapshot_download( + repo_id=TE_4BIT_REPO, + allow_patterns=["*.safetensors", "config.json"], + ) + ) + raw = mx.load(str(root / "model.safetensors")) + stripped = {k[len("model."):]: v for k, v in raw.items() if k.startswith("model.")} + nested = tree_unflatten(list(stripped.items())) + + te = Qwen3TextEncoder(**text_encoder_overrides) + nn.quantize( + te, + class_predicate=lambda _, m: hasattr(m, "to_quantized"), + bits=_QUANT_BITS, + group_size=_QUANT_GROUP_SIZE, + ) + te.update(nested) + return te + + +# Patch Klein's cache-miss reload to honor the 4-bit flag per-instance. Without +# this, eviction + next cache-miss would reinstate bf16 and blow the memory win. +# +# Also short-circuits when pointed at a "slim" packed-mflux checkpoint — i.e. a +# root that has `transformer-packed-mflux/` but no bf16 `transformer/`. mflux's +# stock `_load_weights` would crash trying to read a missing transformer dir +# before the per-instance `_studio_te_4bit` flag is ever set (the flag is set in +# `FluxPipeline._load`, *after* `Flux2Klein()` returns — too late for the init +# path that runs `reload_text_encoder` from inside the constructor). Slim +# checkpoints are exactly what `prism-ml/bonsai-image-*-4B-mlx-*bit` ships, and +# the Bonsai-image-demo download lands here, so this auto-detection lets a fresh +# install generate without any extra knobs. +_ORIGINAL_RELOAD_TEXT_ENCODER = Flux2Initializer.reload_text_encoder + + +def _is_slim_checkpoint(model_path: str | None) -> bool: + if not model_path: + return False + root = Path(model_path) + return (root / "transformer-packed-mflux").exists() and not (root / "transformer").exists() + + +def _reload_text_encoder_with_4bit(model) -> None: + if getattr(model, "_studio_te_4bit", False) or _is_slim_checkpoint(getattr(model, "_model_path", None)): + model.text_encoder = load_te_4bit( + model.model_config.text_encoder_overrides, + model_path=getattr(model, "_model_path", None), + ) + return + _ORIGINAL_RELOAD_TEXT_ENCODER(model) + + +Flux2Initializer.reload_text_encoder = staticmethod(_reload_text_encoder_with_4bit) diff --git a/image-studio/backend_gpu/README.md b/image-studio/backend_gpu/README.md new file mode 100644 index 000000000..1ed7f5134 --- /dev/null +++ b/image-studio/backend_gpu/README.md @@ -0,0 +1,109 @@ +# backend_gpu + +GPU-side companion to `backend/`. Runs on a CUDA host and serves the +`bonsai-ternary-gemlite` backend over the same `/generate` + +`/generate/compare` JSON contract that +`backend.pipeline.RemoteGpuPipeline` POSTs to. + +Pipeline: gemlite transformer + HQQ-int4 text encoder + bf16 VAE on a single +H100. 4-step Klein defaults (`steps=4`, `guidance=1.0`). + +## Layout + +``` +backend_gpu/ + __init__.py + pipeline_gpu.py # GpuPipeline: 5-artifact prewarm + generate_png + diffusion_klein.py # Klein/Qwen3 text→image forward (4-step) + server.py # FastAPI app: /healthz, /generate, /generate/compare + pyproject.toml # deps manifest (gemlite, hqq, transformers, diffusers, torch) + scripts/ + smoke_e2e.py # local CUDA smoke (prewarm + diffusion_forward) + smoke_remote.py # exercise RemoteGpuPipeline against a deployed server + tests/ + test_loaders.py # unit-level coverage of every loader + generate_png contract + test_server.py # FastAPI auth, routing, schema, healthz +``` + +## Artifacts (defaults match `pipeline_gpu.py`) + +| Path | Format | Size | +| --- | --- | --- | +| `` | gemlite-packed ternary transformer | ~1.1 GiB | +| `` | HQQ-packed text encoder | ~2.7 GiB | +| `` | bf16 VAE | ~161 MiB | + +The text encoder artifact bundles its own `tokenizer/` subdir +(Qwen2TokenizerFast). The transformer artifact does NOT carry an HF +`scheduler/` subfolder; `_build_default_scheduler()` in `diffusion_klein.py` +provides the FLUX.2 dynamic-shift defaults (verified against MLX — +`base_shift=0.5, max_shift=1.15, base/max_image_seq_len=256/4096`). + +## Environment + +| Var | Required | Default | Purpose | +| --- | --- | --- | --- | +| `MFLUX_STUDIO_GPU_TOKEN` | yes | — | Bearer token; server refuses to start if unset. | +| `MFLUX_STUDIO_GPU_TERNARY_TRANSFORMER_PATH` | no | (unset) | Packed ternary transformer. | +| `MFLUX_STUDIO_GPU_TRANSFORMER_PATH` | no | (legacy alias for the ternary path) | Retained for backward compatibility. | +| `MFLUX_STUDIO_GPU_TEXT_ENCODER_PATH` | no | (unset) | HQQ-int4 text encoder. | +| `MFLUX_STUDIO_GPU_VAE_PATH` | no | (unset) | bf16 VAE snapshot. | +| `MFLUX_STUDIO_GPU_TOKENIZER_PATH` | no | `/tokenizer/` | Qwen2TokenizerFast directory. | +| `MFLUX_STUDIO_GPU_DEVICE` | no | `cuda:0` | Target device. | + +## Run + +```bash +# Install (assumes torch + gemlite + hqq stack already present in venv) +uv sync # or: pip install -e . + +# Launch +MFLUX_STUDIO_GPU_TOKEN=devtoken \ +uvicorn backend_gpu.server:app --host 0.0.0.0 --port 8801 +``` + +Boot does the full prewarm: 5 artifacts loaded onto `cuda:0`, gemlite +autotune cache restored from `gemlite_autotune.json` in the transformer +artifact dir. First `/generate` call may pay a one-time Triton +compile cost for any image-size / batch shape outside the cached set +(Klein training shapes are covered). + +## Smoke tests + +```bash +# healthz is unauthenticated +curl -s http://localhost:8801/healthz +# {"status":"ok"} + +# /generate requires Bearer; returns image/png +curl -s -o out.png -D - \ + -H "Authorization: Bearer devtoken" \ + -H "Content-Type: application/json" \ + -d '{"prompt":"a cat","steps":4,"width":512,"height":512,"guidance":1.0}' \ + http://localhost:8801/generate + +# unauthenticated → 401 +curl -s -i -H "Content-Type: application/json" \ + -d '{"prompt":"x"}' \ + http://localhost:8801/generate +``` + +Or run the bundled smoke directly on the GPU host (skips the HTTP layer): +```bash +.venv/bin/python -m backend_gpu.scripts.smoke_e2e --prompt "a bonsai" +``` + +## Tests + +```bash +.venv/bin/python -m unittest backend_gpu.tests -v +``` + +Tests stub torch/gemlite/hqq/transformers via `sys.modules` so the suite +runs on macOS without a CUDA stack. + +## Reference perf (H100, single device, fp16 stream) + +- 1024² × 4-step warm: **1.45 s** wall, **6.4 GiB** peak HBM +- 512² × 4-step: smoke target (sub-second per Phase-6 numbers) + diff --git a/image-studio/backend_gpu/__init__.py b/image-studio/backend_gpu/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/image-studio/backend_gpu/diffusion_klein.py b/image-studio/backend_gpu/diffusion_klein.py new file mode 100644 index 000000000..a05d6272f --- /dev/null +++ b/image-studio/backend_gpu/diffusion_klein.py @@ -0,0 +1,262 @@ +"""Klein 4B inference forward (text -> PIL.Image) for the GPU/gemlite backend. + +This is the glue between text encoder, transformer, and VAE. It mirrors the +upstream FLUX.2 inference path but strips out img2img, condition images, real +CFG, T-LoRA, and callback machinery that the GPU backend doesn't need. + +Reused upstream helpers (from diffusers): + Flux2Pipeline._prepare_text_ids # 4-axis RoPE for text tokens + Flux2Pipeline._prepare_latent_ids # 4-axis RoPE for latent grid + Flux2Pipeline._pack_latents # (B,C,H,W) -> (B,H*W,C) + Flux2Pipeline._unpack_latents_with_ids # scatter packed -> (B,C,H,W) + Flux2Pipeline._unpatchify_latents # 2x2 unpack: (B,128,H,W) -> (B,32,2H,2W) + retrieve_timesteps # scheduler.set_timesteps wrapper + +Empirical mu (resolution-dependent shift) is computed locally via +`_mflux_empirical_mu` to match the mflux + iOS Swift port byte-for-byte; +diffusers' built-in linear shift is NOT used. + +Text encode is inlined locally (`_encode_klein_qwen3_prompt`) rather than reusing +`Flux2Pipeline._get_mistral_3_small_prompt_embeds` — that helper is the FLUX.2-dev +Mistral path (layers 10/20/30, Pixtral system message, `add_generation_prompt=False`) +and produces off-distribution embeddings for Klein/Qwen3. + +Klein has `guidance_embeds=False`, so the inference path is single-forward +with a guidance scalar — no two-pass CFG. Default 4 steps. +""" + +from __future__ import annotations + +import logging +from typing import Optional + +import numpy as np +import torch +import torch.nn as nn +from PIL import Image + +from diffusers import FlowMatchEulerDiscreteScheduler, Flux2Pipeline +from diffusers.pipelines.flux2.pipeline_flux2 import retrieve_timesteps + + +log = logging.getLogger(__name__) + + +def _mflux_empirical_mu(image_seq_len: int, num_steps: int) -> float: + """Resolution-dependent shift from mflux. Mirrors + mflux.models.common.schedulers.flow_match_euler_discrete_scheduler.FlowMatchEulerDiscreteScheduler._compute_empirical_mu + (and the iOS port at apple/Bonsai/Pipeline/FlowMatchEulerScheduler.swift::computeEmpiricalMu). + """ + a1, b1 = 8.73809524e-05, 1.89833333 + a2, b2 = 0.00016927, 0.45666666 + if image_seq_len > 4300: + return float(a2 * image_seq_len + b2) + m_200 = a2 * image_seq_len + b2 + m_10 = a1 * image_seq_len + b1 + a = (m_200 - m_10) / 190.0 + b = m_200 - 200.0 * a + return float(a * num_steps + b) + + +# Klein 4B: guidance_embeds=False, so the guidance scalar has no read path +# inside the transformer (Flux2TimestepGuidanceEmbeddings.forward short-circuits +# on `self.guidance_embedder is None`). Defaults: guidance=1.0, steps=4. +DEFAULT_GUIDANCE = 1.0 +DEFAULT_NUM_STEPS = 4 + +# Klein/Qwen3 text-encoder layers stacked into the joint embedding. +KLEIN_OUTPUT_LAYERS = (9, 18, 27) + + +@torch.no_grad() +def _encode_klein_qwen3_prompt( + text_encoder: nn.Module, + tokenizer, + prompt: str, + *, + max_sequence_length: int, +) -> torch.Tensor: + """Klein/Qwen3 prompt encode. + + No system message, plain string content, `add_generation_prompt=True`, + `enable_thinking=False`, hidden states stacked from layers (9, 18, 27). + Returns `(1, max_sequence_length, 3*hidden_dim)` in the encoder's dtype. + """ + device = text_encoder.device + messages = [{"role": "user", "content": prompt}] + text = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True, enable_thinking=False, + ) + inputs = tokenizer( + text, return_tensors="pt", padding="max_length", truncation=True, + max_length=max_sequence_length, + ) + input_ids = inputs["input_ids"].to(device) + attention_mask = inputs["attention_mask"].to(device) + + output = text_encoder( + input_ids=input_ids, attention_mask=attention_mask, + output_hidden_states=True, use_cache=False, + ) + out = torch.stack([output.hidden_states[k] for k in KLEIN_OUTPUT_LAYERS], dim=1) + batch_size, num_channels, seq_len, hidden_dim = out.shape + return out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim) + + +def _build_default_scheduler() -> FlowMatchEulerDiscreteScheduler: + """Default FLUX.2 flow-matching scheduler. + + The Klein model ships a `scheduler/scheduler_config.json` on HF; if the + loader (coder8) wires `_scheduler` into GpuPipeline, prefer that. This + fallback uses diffusers' defaults plus FLUX.2-style dynamic shift so + the mflux empirical mu can flow through `set_timesteps(..., mu=mu)`. + """ + return FlowMatchEulerDiscreteScheduler( + num_train_timesteps=1000, + shift=3.0, + use_dynamic_shifting=True, + base_shift=0.5, + max_shift=1.15, + base_image_seq_len=256, + max_image_seq_len=4096, + ) + + +@torch.no_grad() +def diffusion_forward( + transformer: nn.Module, + text_encoder: nn.Module, + tokenizer, + vae: nn.Module, + prompt: str, + *, + height: int, + width: int, + num_steps: int = DEFAULT_NUM_STEPS, + seed: int = 0, + max_sequence_length: int = 512, + guidance: float = DEFAULT_GUIDANCE, + scheduler: Optional[FlowMatchEulerDiscreteScheduler] = None, +) -> Image.Image: + """Klein 4B text-to-image forward. + + Args: + transformer: gemlite-patched Flux2Transformer2DModel (fp16 internal stream). + text_encoder: Mistral-3 text encoder (bf16). + tokenizer: PixtralProcessor / AutoProcessor for the text encoder. + vae: AutoencoderKLFlux2 (bf16 native). + prompt: text prompt. + height/width: output image size in pixels (must be multiple of 32). + num_steps: flow-matching denoising steps (default 4). + seed: torch CPU generator seed for the initial noise. + max_sequence_length: max text tokens (default 512). + guidance: scalar guidance fed to the transformer (no CFG; single forward). + scheduler: optional pre-loaded scheduler; defaults to FLUX.2 dynamic-shift. + + Returns: PIL.Image.Image, RGB, (height, width). + """ + transformer_device = next(transformer.parameters()).device + vae_device = next(vae.parameters()).device + + if height % 32 != 0 or width % 32 != 0: + raise ValueError(f"height={height} and width={width} must be multiples of 32 (vae_scale_factor*2).") + + if scheduler is None: + scheduler = _build_default_scheduler() + + # 1. Text encode (Klein/Qwen3, bf16 stream). The upstream Mistral helper + # would silently use the wrong layers + Pixtral system message and + # produce off-distribution embeds for Klein. + log.info("encoding prompt (max_seq=%d, klein/qwen3)", max_sequence_length) + prompt_embeds = _encode_klein_qwen3_prompt( + text_encoder=text_encoder, + tokenizer=tokenizer, + prompt=prompt, + max_sequence_length=max_sequence_length, + ) # (1, max_seq, 3*hidden_dim) + text_ids = Flux2Pipeline._prepare_text_ids(prompt_embeds).to(transformer_device) # (1, max_seq, 4) + + # Activation dtype: gemlite int1/int2 kernels require fp16. The bf16 arm + # (`bfl-klein-bf16-gemlite`) wants native bf16 throughout. Loaders mark the + # transformer with `_inference_dtype`; default to fp16 if unset for backwards + # compatibility with callers that build a model outside this package. + activation_dtype = getattr(transformer, "_inference_dtype", torch.float16) + prompt_embeds_t = prompt_embeds.to(device=transformer_device, dtype=activation_dtype) + + # 2. Prepare initial latents in packed (B, image_seq_len, C) form. + # vae_scale_factor=8 (8x VAE compression) and an additional 2x2 patch pack -> 16x total. + vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1) + h_lat = 2 * (int(height) // (vae_scale_factor * 2)) # latent H pre-patchify (= H/8) + w_lat = 2 * (int(width) // (vae_scale_factor * 2)) # latent W pre-patchify + in_channels_latents = transformer.config.in_channels // 4 # = 32 for Klein (in_channels=128) + + # Sample noise on CPU with explicit generator for determinism, then move. + gen = torch.Generator(device="cpu").manual_seed(int(seed)) + noise_shape = (1, in_channels_latents * 4, h_lat // 2, w_lat // 2) # (1, 128, H/16, W/16) + latents_4d = torch.randn(noise_shape, generator=gen, dtype=torch.float32) + latents_4d = latents_4d.to(device=transformer_device, dtype=activation_dtype) + + latent_ids = Flux2Pipeline._prepare_latent_ids(latents_4d).to(transformer_device) # (1, image_seq_len, 4) + latents = Flux2Pipeline._pack_latents(latents_4d) # (1, image_seq_len, 128) fp16 + image_seq_len = latents.shape[1] + log.info("latents: 4D=%s -> packed=%s image_seq_len=%d", + tuple(latents_4d.shape), tuple(latents.shape), image_seq_len) + + # 3. Schedule timesteps with FLUX.2 empirical-mu shift. + mu = _mflux_empirical_mu(image_seq_len=image_seq_len, num_steps=num_steps) + sigmas = np.linspace(1.0, 1.0 / num_steps, num_steps) + if hasattr(scheduler.config, "use_flow_sigmas") and scheduler.config.use_flow_sigmas: + sigmas = None + timesteps, num_steps_eff = retrieve_timesteps( + scheduler, num_steps, transformer_device, sigmas=sigmas, mu=mu, + ) + if hasattr(scheduler, "set_begin_index"): + scheduler.set_begin_index(0) + log.info("scheduling: num_steps=%d mu=%.4f", num_steps_eff, mu) + + # Guidance is a per-batch scalar fed to the transformer (no two-pass CFG). + guidance_t = torch.full([1], guidance, device=transformer_device, dtype=torch.float32) + guidance_t = guidance_t.expand(latents.shape[0]) + + # 4. Denoising loop (single forward per step; gemlite + skip-list both fp16). + for i, t in enumerate(timesteps): + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + noise_pred = transformer( + hidden_states=latents, + timestep=timestep / 1000, + guidance=guidance_t, + encoder_hidden_states=prompt_embeds_t, + txt_ids=text_ids, + img_ids=latent_ids, + return_dict=False, + )[0] + + latents_dtype = latents.dtype + latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0] + if latents.dtype != latents_dtype: + latents = latents.to(latents_dtype) + + # 5. Unpack -> denormalize via VAE batch-norm stats -> unpatchify -> decode. + # Cast to bf16 (vae dtype) on the way out of the transformer stream. + latents = Flux2Pipeline._unpack_latents_with_ids(latents, latent_ids) # (1, 128, H/16, W/16) + latents = latents.to(device=vae_device, dtype=torch.bfloat16) + + bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to( + latents.device, latents.dtype, + ) + latents = latents * bn_std + bn_mean + latents = Flux2Pipeline._unpatchify_latents(latents) # (1, 32, H/8, W/8) + + image = vae.decode(latents, return_dict=False)[0] # (1, 3, H, W) bf16, range [-1, 1] + + # 6. Tensor -> PIL (range conversion, no diffusers VaeImageProcessor dep). + img = image[0].clamp(-1.0, 1.0).float() + img = (img + 1.0) * 127.5 + img = img.clamp(0.0, 255.0).round().to(torch.uint8) + img = img.permute(1, 2, 0).cpu().numpy() # HWC + return Image.fromarray(img, mode="RGB") + + +__all__ = ["diffusion_forward", "DEFAULT_GUIDANCE", "DEFAULT_NUM_STEPS"] diff --git a/image-studio/backend_gpu/pipeline_gpu.py b/image-studio/backend_gpu/pipeline_gpu.py new file mode 100644 index 000000000..4576579b5 --- /dev/null +++ b/image-studio/backend_gpu/pipeline_gpu.py @@ -0,0 +1,547 @@ +from __future__ import annotations + +import io +import json +import logging +import os +import time +from pathlib import Path +from typing import Any, ClassVar, Literal + +GpuBackend = Literal[ + "bonsai-binary-gemlite", + "bonsai-ternary-gemlite", +] + +GPU_BACKENDS: tuple[GpuBackend, ...] = ( + "bonsai-binary-gemlite", + "bonsai-ternary-gemlite", +) + +DEFAULT_GPU_BACKEND: GpuBackend = "bonsai-binary-gemlite" +DEFAULT_SEED = 0 +DEFAULT_STEPS = 4 +DEFAULT_GUIDANCE = 1.0 +DEFAULT_HEIGHT = 512 +DEFAULT_WIDTH = 512 + +# Model artifact paths intentionally have no built-in default — different +# deployments (Evan's laptop, Pasha's Colab demo, primeh200s) lay out the +# model dirs in incompatible places. We require an explicit env var (or +# constructor kwarg) so an unset path fails loudly at startup instead of +# silently fetching weights from HuggingFace or hitting a stale absolute +# path that does not exist on the host. +DEFAULT_BINARY_TRANSFORMER_PATH: str | None = None +DEFAULT_TERNARY_TRANSFORMER_PATH: str | None = None +DEFAULT_TRANSFORMER_PATH: str | None = None # legacy alias +DEFAULT_TEXT_ENCODER_PATH: str | None = None +DEFAULT_VAE_PATH: str | None = None +DEFAULT_TOKENIZER_PATH: str | None = None +DEFAULT_DEVICE = "cuda:0" + + +def _required_path( + explicit: str | Path | None, + env_name: str, + *, + fallbacks: tuple[str | Path | None, ...] = (), + purpose: str, +) -> Path: + """Resolve an artifact path with clear error reporting. + + Tries (in order): the explicit constructor kwarg, the env var, then any + extra fallbacks (e.g. a legacy env var). Returns a Path. Raises + ValueError with a precise hint if none resolve. + """ + candidate = explicit or os.environ.get(env_name) + if candidate is None: + for fb in fallbacks: + if fb is not None: + candidate = fb + break + if candidate is None: + raise ValueError( + f"GpuPipeline {purpose} path is unset. " + f"Set {env_name}= (or pass the matching " + f"kwarg to GpuPipeline(...))." + ) + return Path(candidate) + +def _normalize_gpu_backend(raw: str) -> GpuBackend: + if raw in GPU_BACKENDS: + return raw # type: ignore[return-value] + raise ValueError(f"Unknown GPU backend {raw!r}; expected one of {GPU_BACKENDS}.") + +log = logging.getLogger(__name__) + + +_GEMLITE_LAYER_KEYS = ("W_q", "bias", "scales", "zeros", "metadata", "orig_shape") + + +def _load_gemlite_layers_from_state( + model: Any, + state: dict[str, Any], + *, + bits: int, + group_size: int, + device: str, + DType: Any, + GemLiteLinearTriton: Any, +) -> tuple[int, dict[str, Any]]: + """Bucket gemlite-layer keys in `state`, replace each `nn.Linear` with a + `GemLiteLinearTriton` initialized via gemlite's custom per-layer + `load_state_dict`, then move its tensors to `device`. + + Returns `(n_loaded, remainder_state)` where `remainder_state` excludes the + consumed gemlite keys (caller loads it via `model.load_state_dict`). + + Why per-layer: pack() registers W_q/scales/zeros/metadata/orig_shape as + nn.Parameters, so the saved `state_dict.pt` has those keys, but the + constructor of an empty GemLiteLinearTriton does NOT pre-register the + Parameter slots — `model.load_state_dict(strict=False)` reports them all + as `unexpected`. Gemlite ships a custom `GemLiteLinearTriton.load_state_dict` + that pops these keys and decodes `metadata` into the layer's scalar fields + (W_nbits, group_size, dtypes, …). We dispatch that per layer. + """ + import torch + import torch.nn as nn + + buckets: dict[str, dict[str, torch.Tensor]] = {} + remainder: dict[str, torch.Tensor] = {} + for k, v in state.items(): + fqn, _, leaf = k.rpartition(".") + if leaf in _GEMLITE_LAYER_KEYS and fqn: + buckets.setdefault(fqn, {})[leaf] = v + else: + remainder[k] = v + + n_loaded = 0 + target_device = torch.device(device) + for fqn, layer_state in buckets.items(): + parent_fqn, _, child_name = fqn.rpartition(".") + parent = model.get_submodule(parent_fqn) if parent_fqn else model + child = getattr(parent, child_name) + if not isinstance(child, nn.Linear): + raise RuntimeError( + f"state_dict has gemlite keys at {fqn} but model has {type(child).__name__}" + ) + gl = GemLiteLinearTriton( + W_nbits=bits, + group_size=group_size, + in_features=child.in_features, + out_features=child.out_features, + input_dtype=DType.FP16, + output_dtype=DType.FP16, + ) + gl.load_state_dict(dict(layer_state)) + gl.W_q = gl.W_q.to(target_device) + gl.scales = gl.scales.to(target_device) + gl.zeros = gl.zeros.to(target_device) + if gl.bias is not None: + gl.bias = gl.bias.to(target_device) + gl.device = target_device + setattr(parent, child_name, gl) + n_loaded += 1 + log.info("loaded %d GemLiteLinearTriton layers from state_dict", n_loaded) + return n_loaded, remainder + + +def _null_gemlite_weights(model: Any, GemLiteLinearTriton: Any) -> int: + """Set `.weight = None` on every `GemLiteLinearTriton` in `model`. + + Why: `Flux2AttnProcessor`'s MuonClip-telemetry path eagerly evaluates + `attn.to_q.weight` (and friends). After gemlite replacement those modules + have no real `weight` tensor — but PyTorch still synthesises an empty one + via `nn.Module.__getattr__` lookup against any registered Parameter slot. + Forcing `.weight = None` via `object.__setattr__` (bypassing the Module + parameter machinery) makes the access explicit-None instead of a phantom + tensor; the telemetry path treats None as "skip". + """ + nulled = 0 + for m in model.modules(): + if isinstance(m, GemLiteLinearTriton): + object.__setattr__(m, "weight", None) + nulled += 1 + log.info("nulled .weight on %d GemLiteLinearTriton modules", nulled) + return nulled + + +def _load_gemlite_transformer(path: Path, *, device: str = DEFAULT_DEVICE) -> Any: + """Load the gemlite-packed Klein-4B transformer onto `device`. + + Reads `state_dict.pt`, `config.json`, `quantization_config.json`, + `gemlite_autotune.json` from `path`. Restores the global gemlite autotune + cache via `gemlite.core.load_config(...)`. Calls `set_packing_bitwidth(...)` + BEFORE patching/loading so kernel selection matches the pack run. + + Post-load: + 1. Cast the whole module to fp16 (matches the gemlite forward stream). + 2. Null the `.weight` attribute on every GemLiteLinearTriton (Phase-2 + carryover; see `_null_gemlite_weights` for the reason). + """ + if not path.is_dir(): + raise FileNotFoundError( + f"Gemlite transformer artifact not found at {path} — " + "run scripts/pack_klein_to_gemlite.py to regenerate." + ) + state_path = path / "state_dict.pt" + config_path = path / "config.json" + qcfg_path = path / "quantization_config.json" + autotune_path = path / "gemlite_autotune.json" + for f in (state_path, config_path, qcfg_path, autotune_path): + if not f.is_file(): + raise FileNotFoundError(f"Gemlite transformer missing {f.name} at {f}") + + with config_path.open() as fh: + cfg = json.load(fh) + with qcfg_path.open() as fh: + qcfg = json.load(fh) + bits = int(qcfg.get("bits", 1)) + group_size = int(qcfg.get("group_size", 128)) + packing_bw = int(qcfg.get("packing_bitwidth", 8)) + + import torch + from diffusers import Flux2Transformer2DModel + from gemlite.core import DType, GemLiteLinearTriton, set_packing_bitwidth + + set_packing_bitwidth(packing_bw) + GemLiteLinearTriton.load_config(str(autotune_path)) + + log.info( + "loading gemlite transformer config: bits=%d gs=%d bw=%d", + bits, group_size, packing_bw, + ) + model = Flux2Transformer2DModel.from_config(cfg).to(torch.bfloat16) + state = torch.load(str(state_path), map_location="cpu") + _, remainder = _load_gemlite_layers_from_state( + model, state, + bits=bits, group_size=group_size, device=device, + DType=DType, GemLiteLinearTriton=GemLiteLinearTriton, + ) + missing, unexpected = model.load_state_dict(remainder, strict=False) + if unexpected: + raise RuntimeError(f"unexpected non-gemlite state_dict keys: {unexpected[:8]}") + if missing: + raise RuntimeError(f"missing non-gemlite state_dict keys: {missing[:8]}") + + model = model.to(torch.float16) + _null_gemlite_weights(model, GemLiteLinearTriton) + model = model.to(device).eval() + # Marker consumed by diffusion_klein to pick the activation dtype for the + # transformer-input casts (fp16 here; bf16 for the vanilla bf16 loader). + model._inference_dtype = torch.float16 # type: ignore[attr-defined] + return model + + +def _load_transformer_for_backend( + backend: GpuBackend, + path: Path, + *, + device: str = DEFAULT_DEVICE, +) -> Any: + """Dispatch to the right loader based on the backend's quantization scheme.""" + return _load_gemlite_transformer(path, device=device) + + +def _load_text_encoder(path: Path, *, device: str = DEFAULT_DEVICE) -> Any: + """Load the HQQ-4bit Klein text encoder and gemlite-patch it for inference. + + Output is a `Mistral3ForConditionalGeneration` with HQQLinear modules + converted to gemlite kernels (fp16 stream). + """ + if not path.is_dir(): + raise FileNotFoundError( + f"Text encoder artifact not found at {path} — " + "see scripts/pack_klein_text_encoder_to_gemlite.py." + ) + import torch + from gemlite.core import set_packing_bitwidth + from hqq.models.hf.base import AutoHQQHFModel + from hqq.utils.patching import prepare_for_inference + + set_packing_bitwidth(8) + model = AutoHQQHFModel.from_quantized( + str(path), + compute_dtype=torch.float16, + device=device, + ) + prepare_for_inference(model, backend="gemlite") + return model + + +def _load_vae(path: Path, *, device: str = DEFAULT_DEVICE) -> Any: + """Load the bf16 `AutoencoderKLFlux2` from a local snapshot path.""" + if not path.is_dir(): + raise FileNotFoundError( + f"VAE snapshot not found at {path} — provision via `huggingface-cli download` " + "or rerun the Phase-4 snapshot script." + ) + import torch + from diffusers import AutoencoderKLFlux2 + + vae = AutoencoderKLFlux2.from_pretrained(str(path), torch_dtype=torch.bfloat16) + return vae.to(device).eval() + + +def _load_tokenizer(path_or_repo: str) -> Any: + """Load Klein's Qwen2TokenizerFast (text-encoder side, plain string content). + + `diffusion_klein._encode_klein_qwen3_prompt` only needs `apply_chat_template` + and `__call__`, which `AutoTokenizer` provides. The Klein TE artifact ships + its own `tokenizer/` subdir; default points there. NOT a Pixtral processor — + that's the FLUX.2-dev (Mistral) path and produces wrong embeds for Klein. + """ + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(path_or_repo) + + +def _load_scheduler(transformer_path: Path) -> Any | None: + """Load `FlowMatchEulerDiscreteScheduler` from the transformer snapshot. + + Klein checkpoints ship a `scheduler/` subfolder; if present we load it so + the diffusion forward inherits the trained dynamic-shift settings. Returns + `None` if absent — `diffusion_klein.py` falls back to its FLUX.2 defaults. + """ + sched_path = transformer_path / "scheduler" + if not sched_path.is_dir(): + log.info("no scheduler/ subfolder in %s — diffusion_klein will use defaults", transformer_path) + return None + from diffusers import FlowMatchEulerDiscreteScheduler + + return FlowMatchEulerDiscreteScheduler.from_pretrained(str(transformer_path), subfolder="scheduler") + + +class GpuPipeline: + """Server-side GPU pipeline (gemlite/HQQ on H100). + + `prewarm()` loads 5 artifacts: gemlite transformer, HQQ-gemlite text + encoder, AutoencoderKLFlux2, Qwen2 tokenizer, and an optional scheduler. + `generate_png` calls `backend_gpu.diffusion_klein.diffusion_forward` and + encodes the returned PIL.Image as PNG bytes. + """ + + is_remote: ClassVar[bool] = False + + def __init__( + self, + *, + backend: GpuBackend = DEFAULT_GPU_BACKEND, + transformer_path: str | Path | None = None, + binary_transformer_path: str | Path | None = None, + ternary_transformer_path: str | Path | None = None, + text_encoder_path: str | Path | None = None, + vae_path: str | Path | None = None, + tokenizer_path: str | None = None, + device: str | None = None, + ) -> None: + backend = _normalize_gpu_backend(backend) + self._backend: GpuBackend = backend + self.last_peak_memory_mb: float | None = None + self._ready: bool = False + self._transformer: Any = None + self._text_encoder: Any = None + self._vae: Any = None + self._tokenizer: Any = None + self._scheduler: Any = None + # `transformer_path` (and the matching MFLUX_STUDIO_GPU_TRANSFORMER_PATH + # env) is a legacy alias for the binary path; pre-existing single-backend + # callers keep working without setting the new BINARY-suffixed name. + legacy_default_binary = ( + transformer_path + or os.environ.get("MFLUX_STUDIO_GPU_TRANSFORMER_PATH") + ) + self._transformer_paths: dict[GpuBackend, Path] = { + "bonsai-binary-gemlite": _required_path( + binary_transformer_path, + "MFLUX_STUDIO_GPU_BINARY_TRANSFORMER_PATH", + fallbacks=(legacy_default_binary, DEFAULT_BINARY_TRANSFORMER_PATH), + purpose="binary transformer", + ), + "bonsai-ternary-gemlite": _required_path( + ternary_transformer_path, + "MFLUX_STUDIO_GPU_TERNARY_TRANSFORMER_PATH", + fallbacks=(DEFAULT_TERNARY_TRANSFORMER_PATH,), + purpose="ternary transformer", + ), + } + self.text_encoder_path: Path = _required_path( + text_encoder_path, + "MFLUX_STUDIO_GPU_TEXT_ENCODER_PATH", + fallbacks=(DEFAULT_TEXT_ENCODER_PATH,), + purpose="text encoder", + ) + self.vae_path: Path = _required_path( + vae_path, + "MFLUX_STUDIO_GPU_VAE_PATH", + fallbacks=(DEFAULT_VAE_PATH,), + purpose="VAE", + ) + # tokenizer is loaded by HuggingFace `from_pretrained`, which accepts + # a string (path or repo id). Keep that surface type as `str`. + tok_resolved = _required_path( + tokenizer_path, + "MFLUX_STUDIO_GPU_TOKENIZER_PATH", + fallbacks=(DEFAULT_TOKENIZER_PATH,), + purpose="tokenizer", + ) + self.tokenizer_path: str = str(tok_resolved) + self.device: str = device or os.environ.get("MFLUX_STUDIO_GPU_DEVICE", DEFAULT_DEVICE) + + @property + def backend(self) -> GpuBackend: + return self._backend + + @property + def ready(self) -> bool: + return self._ready + + @property + def transformer_path(self) -> Path: + return self._transformer_paths[self._backend] + + def ensure_backend(self, *, backend: GpuBackend, model_path: str | None = None) -> None: + backend = _normalize_gpu_backend(backend) + if backend == self._backend and model_path is None: + return + if model_path is not None: + self._transformer_paths[backend] = Path(model_path) + # Different backend → drop the resident transformer and reload on next prewarm/use. + if backend != self._backend: + self._backend = backend + self._transformer = None + self._scheduler = None + self._ready = False + if self._text_encoder is not None and self._vae is not None and self._tokenizer is not None: + t0 = time.perf_counter() + self._transformer = _load_transformer_for_backend( + self._backend, self.transformer_path, device=self.device, + ) + log.info( + "swapped transformer to %s in %.2fs", + self._backend, time.perf_counter() - t0, + ) + self._scheduler = _load_scheduler(self.transformer_path) + self._ready = True + + def prewarm(self) -> None: + """Load all 5 artifacts; mark pipeline ready. + + Per-artifact errors surface verbatim so deploy diagnostics are + unambiguous. Each step is timed at INFO. Autotune-warmup forward is + deferred to `diffusion_klein.diffusion_forward`'s first call since the + diffusion-loop shapes live there. + """ + log.info( + "GpuPipeline.prewarm starting backend=%s device=%s " + "transformer=%s text_encoder=%s vae=%s tokenizer=%s", + self._backend, self.device, self.transformer_path, + self.text_encoder_path, self.vae_path, self.tokenizer_path, + ) + t0 = time.perf_counter() + self._transformer = _load_transformer_for_backend( + self._backend, self.transformer_path, device=self.device, + ) + log.info("loaded transformer (%s) in %.2fs", self._backend, time.perf_counter() - t0) + + t0 = time.perf_counter() + self._text_encoder = _load_text_encoder(self.text_encoder_path, device=self.device) + log.info("loaded text encoder in %.2fs", time.perf_counter() - t0) + + t0 = time.perf_counter() + self._vae = _load_vae(self.vae_path, device=self.device) + log.info("loaded VAE in %.2fs", time.perf_counter() - t0) + + t0 = time.perf_counter() + self._tokenizer = _load_tokenizer(self.tokenizer_path) + log.info("loaded tokenizer in %.2fs", time.perf_counter() - t0) + + t0 = time.perf_counter() + self._scheduler = _load_scheduler(self.transformer_path) + log.info("loaded scheduler in %.2fs (present=%s)", + time.perf_counter() - t0, self._scheduler is not None) + + self._ready = True + log.info("GpuPipeline ready: 5 artifacts loaded (scheduler optional)") + + def generate_png( + self, + *, + prompt: str, + seed: int = DEFAULT_SEED, + steps: int = DEFAULT_STEPS, + height: int = DEFAULT_HEIGHT, + width: int = DEFAULT_WIDTH, + guidance: float = DEFAULT_GUIDANCE, + tiled_vae: bool | None = None, # accepted for API parity with MLX; H100 80GiB has no VAE memory pressure + max_sequence_length: int | None = None, + ) -> bytes: + del tiled_vae + if not self._ready: + raise RuntimeError("GpuPipeline.prewarm() must be called before generate_png().") + # Lazy: diffusion_klein imports `diffusers.pipelines.flux2`, which is + # heavy and unavailable in test environments unless mocked. + import torch + from backend_gpu import diffusion_klein + + log.info( + "generate backend=%s size=%dx%d steps=%d seed=%d guidance=%.2f", + self._backend, width, height, steps, seed, guidance, + ) + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + + forward_kwargs: dict[str, Any] = { + "transformer": self._transformer, + "text_encoder": self._text_encoder, + "tokenizer": self._tokenizer, + "vae": self._vae, + "prompt": prompt, + "height": height, + "width": width, + "num_steps": steps, + "seed": seed, + "guidance": guidance, + "scheduler": self._scheduler, + } + if max_sequence_length is not None: + forward_kwargs["max_sequence_length"] = max_sequence_length + + image = diffusion_klein.diffusion_forward(**forward_kwargs) + + self.last_peak_memory_mb = ( + torch.cuda.max_memory_allocated() / (1024**2) + if torch.cuda.is_available() + else 0.0 + ) + + buf = io.BytesIO() + image.save(buf, format="PNG") + return buf.getvalue() + + +__all__ = [ + "GpuBackend", + "GPU_BACKENDS", + "DEFAULT_GPU_BACKEND", + "DEFAULT_SEED", + "DEFAULT_STEPS", + "DEFAULT_GUIDANCE", + "DEFAULT_HEIGHT", + "DEFAULT_WIDTH", + "DEFAULT_BINARY_TRANSFORMER_PATH", + "DEFAULT_TERNARY_TRANSFORMER_PATH", + "DEFAULT_TRANSFORMER_PATH", + "DEFAULT_TEXT_ENCODER_PATH", + "DEFAULT_VAE_PATH", + "DEFAULT_TOKENIZER_PATH", + "DEFAULT_DEVICE", + "GpuPipeline", + "_load_gemlite_transformer", + "_load_text_encoder", + "_load_vae", + "_load_tokenizer", + "_load_scheduler", + "_load_gemlite_layers_from_state", + "_null_gemlite_weights", +] diff --git a/image-studio/backend_gpu/pyproject.toml b/image-studio/backend_gpu/pyproject.toml new file mode 100644 index 000000000..5239c4ec2 --- /dev/null +++ b/image-studio/backend_gpu/pyproject.toml @@ -0,0 +1,25 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "prism-image-studio-backend-gpu" +version = "0.0.1" +description = "GPU (gemlite/HQQ on H100) inference server for prism-image-studio. Phase-5 prep skeleton." +requires-python = ">=3.11" +dependencies = [ + "fastapi>=0.115", + "uvicorn[standard]>=0.30", + "pydantic>=2.7", + "pillow>=10.4", + "diffusers>=0.38.0", +] + +[project.optional-dependencies] +dev = [ + "httpx>=0.27", +] + +[tool.setuptools] +packages = ["backend_gpu"] +package-dir = {"backend_gpu" = "."} diff --git a/image-studio/backend_gpu/scripts/inference_bf16.py b/image-studio/backend_gpu/scripts/inference_bf16.py new file mode 100644 index 000000000..304f05f68 --- /dev/null +++ b/image-studio/backend_gpu/scripts/inference_bf16.py @@ -0,0 +1,145 @@ +"""BF16 inference using local bonsai models, matching the C++ reference.""" + +from __future__ import annotations + +import argparse +import logging +import time +from pathlib import Path + +import torch + + +def _mflux_empirical_mu(image_seq_len: int, num_steps: int) -> float: + a1, b1 = 8.73809524e-05, 1.89833333 + a2, b2 = 0.00016927, 0.45666666 + if image_seq_len > 4300: + return float(a2 * image_seq_len + b2) + m_200 = a2 * image_seq_len + b2 + m_10 = a1 * image_seq_len + b1 + a = (m_200 - m_10) / 190.0 + b = m_200 - 200.0 * a + return float(a * num_steps + b) + + +KLEIN_OUTPUT_LAYERS = (9, 18, 27) + + +@torch.no_grad() +def _encode_prompt(text_encoder, tokenizer, prompt: str, max_sequence_length: int) -> torch.Tensor: + device = text_encoder.device + messages = [{"role": "user", "content": prompt}] + text = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True, enable_thinking=False, + ) + inputs = tokenizer( + text, return_tensors="pt", padding="max_length", truncation=True, + max_length=max_sequence_length, + ) + output = text_encoder( + input_ids=inputs["input_ids"].to(device), + attention_mask=inputs["attention_mask"].to(device), + output_hidden_states=True, use_cache=False, + ) + out = torch.stack([output.hidden_states[k] for k in KLEIN_OUTPUT_LAYERS], dim=1) + _, num_channels, seq_len, hidden_dim = out.shape + return out.permute(0, 2, 1, 3).reshape(1, seq_len, num_channels * hidden_dim) + + +def main() -> int: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") + log = logging.getLogger("inference_bf16") + + parser = argparse.ArgumentParser() + parser.add_argument("--prompt", default="a cat sitting on a window sill") + parser.add_argument("--height", type=int, default=512) + parser.add_argument("--width", type=int, default=512) + parser.add_argument("--steps", type=int, default=4) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--guidance", type=float, default=1.0) + parser.add_argument("--output", type=Path, default=Path("outputs/test_output_python.png")) + parser.add_argument("--model-dir", type=Path, + default=Path("models/bonsai-image-binary-4B-unpacked")) + args = parser.parse_args() + log.info("args: %s", vars(args)) + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA required") + + model_dir = args.model_dir.resolve() + device = "cuda:0" + + # 1. VAE + log.info("loading VAE...") + from diffusers import AutoencoderKLFlux2 + vae = AutoencoderKLFlux2.from_pretrained( + str(model_dir / "vae"), torch_dtype=torch.bfloat16, + ).to(device).eval() + + # 2. Transformer + log.info("loading transformer...") + from diffusers import Flux2Transformer2DModel + transformer = Flux2Transformer2DModel.from_pretrained( + str(model_dir / "transformer"), torch_dtype=torch.bfloat16, + ).to(device).eval() + + # 3. Text encoder + tokenizer + log.info("loading text encoder...") + from transformers import AutoModelForCausalLM, AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained( + str(model_dir / "tokenizer"), trust_remote_code=True, + ) + text_encoder = AutoModelForCausalLM.from_pretrained( + str(model_dir / "text_encoder"), + torch_dtype=torch.bfloat16, + trust_remote_code=True, + low_cpu_mem_usage=True, + output_hidden_states=True, + ).to(device).eval() + + # 4. Scheduler + from diffusers import FlowMatchEulerDiscreteScheduler + scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + str(model_dir), subfolder="scheduler", + ) + + from diffusers import Flux2Pipeline + pipe = Flux2Pipeline( + scheduler=scheduler, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + vae=vae, + ) + pipe.set_progress_bar_config(disable=True) + log.info("all models loaded") + + # Encode prompt (Klein/Qwen3 stacking) + log.info("encoding prompt...") + prompt_embeds = _encode_prompt(text_encoder, tokenizer, args.prompt, max_sequence_length=512) + prompt_embeds = prompt_embeds.to(device=device, dtype=torch.bfloat16) + + # Generate + log.info("generating...") + torch.cuda.reset_peak_memory_stats() + t0 = time.perf_counter() + image = pipe( + prompt_embeds=prompt_embeds, + num_inference_steps=args.steps, + generator=torch.Generator(device="cpu").manual_seed(args.seed), + guidance_scale=args.guidance, + height=args.height, + width=args.width, + ).images[0] + elapsed = time.perf_counter() - t0 + peak_mib = torch.cuda.max_memory_allocated(device) / 1024 / 1024 + + args.output.parent.mkdir(parents=True, exist_ok=True) + image.save(str(args.output)) + log.info("generated %dx%d in %.2fs, peak HBM %.0f MiB -> %s", + image.width, image.height, elapsed, peak_mib, args.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) \ No newline at end of file diff --git a/image-studio/backend_gpu/scripts/smoke_e2e.py b/image-studio/backend_gpu/scripts/smoke_e2e.py new file mode 100644 index 000000000..6efef84bd --- /dev/null +++ b/image-studio/backend_gpu/scripts/smoke_e2e.py @@ -0,0 +1,114 @@ +"""End-to-end smoke for the GPU/gemlite Klein backend (Phase-5c-2). + +Constructs `GpuPipeline`, runs `prewarm()`, then calls `diffusion_forward` to +turn a single prompt into a PIL.Image. Reports per-stage timings, HBM peak, +and writes the result PNG to disk. + +This is the first time the gemlite kernels run with the actual Klein diffusion +loop shapes (transformer expects packed image-token grids, e.g. 4096 tokens at +1024x1024). Cold-start autotune cost is expected on the first forward pass — +the cached `gemlite_autotune.json` only covers shapes that the pack-time +warmup forward exercised. +""" + +from __future__ import annotations + +import argparse +import logging +import sys +import time +from pathlib import Path + + +def main() -> int: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") + p = argparse.ArgumentParser(description=__doc__.split("\n\n", 1)[0]) + p.add_argument("--prompt", default="a serene bonsai tree on a rocky outcrop, dramatic golden-hour lighting, photorealistic") + p.add_argument("--height", type=int, default=512) + p.add_argument("--width", type=int, default=512) + p.add_argument("--steps", type=int, default=4) + p.add_argument("--seed", type=int, default=0) + p.add_argument("--guidance", type=float, default=1.0) + p.add_argument("--max-sequence-length", type=int, default=512) + p.add_argument("--output", type=Path, default=Path("/root/bench_out/phase5c2_smoke.png")) + p.add_argument("--repeat", type=int, default=1, + help="How many forwards to run after prewarm (>=2 separates cold-start from warm).") + args = p.parse_args() + + log = logging.getLogger("smoke_e2e") + log.info("args: %s", vars(args)) + + import torch + + from backend_gpu.diffusion_klein import diffusion_forward + from backend_gpu.pipeline_gpu import GpuPipeline + + if not torch.cuda.is_available(): + raise RuntimeError("smoke is CUDA-only; run on a CUDA host.") + + pipe = GpuPipeline() + log.info("constructed GpuPipeline backend=%s device=%s", pipe.backend, pipe.device) + + t0 = time.perf_counter() + pipe.prewarm() + prewarm_s = time.perf_counter() - t0 + log.info("prewarm complete in %.1fs (ready=%s)", prewarm_s, pipe.ready) + + # Memory after load. + static_alloc_mib = torch.cuda.memory_allocated(pipe.device) / 1024 / 1024 + log.info("post-prewarm HBM allocated: %.1f MiB", static_alloc_mib) + + args.output.parent.mkdir(parents=True, exist_ok=True) + + timings: list[float] = [] + peaks_mib: list[float] = [] + for i in range(args.repeat): + torch.cuda.synchronize(pipe.device) + torch.cuda.reset_peak_memory_stats(pipe.device) + t0 = time.perf_counter() + img = diffusion_forward( + transformer=pipe._transformer, + text_encoder=pipe._text_encoder, + tokenizer=pipe._tokenizer, + vae=pipe._vae, + prompt=args.prompt, + height=args.height, + width=args.width, + num_steps=args.steps, + seed=args.seed, + max_sequence_length=args.max_sequence_length, + guidance=args.guidance, + scheduler=pipe._scheduler, + ) + torch.cuda.synchronize(pipe.device) + forward_s = time.perf_counter() - t0 + peak_mib = torch.cuda.max_memory_allocated(pipe.device) / 1024 / 1024 + timings.append(forward_s) + peaks_mib.append(peak_mib) + log.info("[forward %d/%d] %.2fs peak HBM %.1f MiB output %s", + i + 1, args.repeat, forward_s, peak_mib, img.size) + + if i == 0: + img.save(str(args.output)) + log.info("wrote first-iter image -> %s", args.output) + + print() + print(f"=== Phase 5c-2 E2E smoke ({args.height}x{args.width} | steps={args.steps} | seed={args.seed}) ===") + print(f" prompt: {args.prompt!r}") + print(f" prewarm : {prewarm_s:8.1f} s") + print(f" post-prewarm static HBM : {static_alloc_mib:8.1f} MiB") + print() + for i, (t, peak) in enumerate(zip(timings, peaks_mib)): + print(f" forward[{i}] : {t:8.2f} s peak HBM {peak:8.1f} MiB") + if len(timings) >= 2: + warm = timings[1:] + print(f" warm mean (n={len(warm)}) : {sum(warm) / len(warm):8.2f} s") + print(f" cold-start overhead : {timings[0] - (sum(warm) / len(warm)):+8.2f} s") + print() + print(f" output PNG: {args.output} ({args.output.stat().st_size / 1024:.1f} KiB)") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/image-studio/backend_gpu/scripts/smoke_remote.py b/image-studio/backend_gpu/scripts/smoke_remote.py new file mode 100644 index 000000000..182c4a553 --- /dev/null +++ b/image-studio/backend_gpu/scripts/smoke_remote.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Phase 5b laptop-side smoke: round-trip RemoteGpuPipeline against a live +remote backend_gpu service. + +Reads MFLUX_STUDIO_GPU_HOST + MFLUX_STUDIO_GPU_TOKEN from the environment. +Confirms: + 1. valid PNG bytes come back from /generate + 2. last_peak_memory_mb is parsed from the X-Peak-Memory-MB header + 3. /generate/compare round-trips (via direct httpx call, since + RemoteGpuPipeline.generate_png hits /generate only) + 4. wrong token surfaces a clean RuntimeError + +Run: + MFLUX_STUDIO_GPU_HOST=http://127.0.0.1:8801 \ + MFLUX_STUDIO_GPU_TOKEN= \ + .venv/bin/python backend_gpu/scripts/smoke_remote.py +""" +from __future__ import annotations + +import dataclasses +import os +import sys +from pathlib import Path + +# Allow running from anywhere — pin sys.path to image-studio/. +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT)) + +from backend.pipeline import PipelineConfig, RemoteGpuPipeline # noqa: E402 + +_PNG_MAGIC = b"\x89PNG\r\n\x1a\n" + + +def _config(token_override: str | None = None) -> PipelineConfig: + cfg = PipelineConfig.from_env() + if token_override is not None: + cfg = dataclasses.replace(cfg, gpu_token=token_override) + return cfg + + +def smoke_happy_path() -> None: + pipe = RemoteGpuPipeline(_config()) + try: + png = pipe.generate_png(prompt="a cat", steps=4, width=128, height=128, seed=7) + finally: + pipe.close() + assert png.startswith(_PNG_MAGIC), f"expected PNG magic, got {png[:16].hex()}" + print(f"[OK] /generate happy path: {len(png)} bytes, peak_mb={pipe.last_peak_memory_mb}") + + +def smoke_wrong_token() -> None: + pipe = RemoteGpuPipeline(_config(token_override="bogus-token")) + try: + try: + pipe.generate_png(prompt="x", steps=1, width=64, height=64) + except RuntimeError as exc: + assert "401" in str(exc), f"expected 401 in error, got: {exc}" + print(f"[OK] wrong token surfaces RuntimeError: {exc}") + return + raise AssertionError("expected RuntimeError on wrong token") + finally: + pipe.close() + + +def smoke_compare() -> None: + # /generate/compare is exercised via raw httpx since RemoteGpuPipeline only + # wraps /generate. Confirms the second route plus auth too. + import base64 + + import httpx + + cfg = _config() + with httpx.Client( + base_url=cfg.gpu_host, + headers={"Authorization": f"Bearer {cfg.gpu_token}"}, + timeout=httpx.Timeout(connect=10.0, read=60.0, write=10.0, pool=10.0), + ) as client: + response = client.post( + "/generate/compare", + json={"prompt": "x", "backends": ["bonsai-ternary-gemlite"], "width": 64, "height": 64, "steps": 1}, + ) + response.raise_for_status() + body = response.json() + assert len(body["results"]) == 1 + result = body["results"][0] + decoded = base64.b64decode(result["png_b64"]) + assert decoded.startswith(_PNG_MAGIC) + print( + f"[OK] /generate/compare round-trip: backend={result['backend']} " + f"wall={result['wall_seconds']:.4f}s swap={result['swap_seconds']:.4f}s png={len(decoded)}B" + ) + + +def main() -> int: + if not os.environ.get("MFLUX_STUDIO_GPU_HOST") or not os.environ.get("MFLUX_STUDIO_GPU_TOKEN"): + print("MFLUX_STUDIO_GPU_HOST and MFLUX_STUDIO_GPU_TOKEN must be set.", file=sys.stderr) + return 2 + smoke_happy_path() + smoke_compare() + smoke_wrong_token() + print("\nAll Phase 5b smoke checks PASSED.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/image-studio/backend_gpu/server.py b/image-studio/backend_gpu/server.py new file mode 100644 index 000000000..fcdba3c4c --- /dev/null +++ b/image-studio/backend_gpu/server.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +import base64 +import logging +import os +import time +from contextlib import asynccontextmanager + +from fastapi import Depends, FastAPI, HTTPException, Request, Response +from fastapi.middleware.cors import CORSMiddleware +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from pydantic import BaseModel, Field, model_validator + +from backend_gpu.pipeline_gpu import ( + DEFAULT_GPU_BACKEND, + DEFAULT_GUIDANCE, + DEFAULT_HEIGHT, + DEFAULT_SEED, + DEFAULT_STEPS, + DEFAULT_WIDTH, + GPU_BACKENDS, + GpuBackend, + GpuPipeline, +) + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", +) +log = logging.getLogger(__name__) + + +def _required_token() -> str: + token = os.environ.get("MFLUX_STUDIO_GPU_TOKEN") + if not token: + raise RuntimeError( + "MFLUX_STUDIO_GPU_TOKEN must be set; the GPU server refuses to start unauthenticated." + ) + return token + + +_bearer_scheme = HTTPBearer(auto_error=False) + + +def _verify_bearer( + request: Request, + credentials: HTTPAuthorizationCredentials | None = Depends(_bearer_scheme), +) -> None: + expected: str = request.app.state.token + if credentials is None or credentials.scheme.lower() != "bearer" or credentials.credentials != expected: + raise HTTPException(status_code=401, detail="Invalid or missing bearer token.") + + +@asynccontextmanager +async def lifespan(app: FastAPI): + token = _required_token() + backend_env = os.environ.get("MFLUX_STUDIO_GPU_DEFAULT_BACKEND", DEFAULT_GPU_BACKEND) + from backend_gpu.pipeline_gpu import _normalize_gpu_backend + try: + backend_env = _normalize_gpu_backend(backend_env) + except ValueError as exc: + raise RuntimeError( + f"MFLUX_STUDIO_GPU_DEFAULT_BACKEND={backend_env!r} not in {GPU_BACKENDS}." + ) from exc + pipeline = GpuPipeline(backend=backend_env) + pipeline.prewarm() + app.state.pipeline = pipeline + app.state.token = token + log.info("backend_gpu ready backend=%s", pipeline.backend) + yield + + +app = FastAPI(lifespan=lifespan) +# Why: Bearer auth gates every protected route, so opening CORS is safe and lets +# the laptop frontend hit this directly during development without a proxy. +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_methods=["*"], + allow_headers=["*"], +) + + +class GenerateRequest(BaseModel): + prompt: str = Field(min_length=1) + seed: int = DEFAULT_SEED + steps: int = Field(default=DEFAULT_STEPS, ge=1) + guidance: float = Field(default=DEFAULT_GUIDANCE, ge=0.0) + backend: GpuBackend = DEFAULT_GPU_BACKEND + height: int = Field(default=DEFAULT_HEIGHT, ge=16) + width: int = Field(default=DEFAULT_WIDTH, ge=16) + tiled_vae: bool | None = Field(default=None) + max_sequence_length: int | None = Field(default=None, ge=1) + + +class CompareRequest(BaseModel): + prompt: str = Field(min_length=1) + seed: int = DEFAULT_SEED + steps: int = Field(default=DEFAULT_STEPS, ge=1) + guidance: float = Field(default=DEFAULT_GUIDANCE, ge=0.0) + height: int = Field(default=DEFAULT_HEIGHT, ge=16) + width: int = Field(default=DEFAULT_WIDTH, ge=16) + backends: list[GpuBackend] = Field(default_factory=lambda: list(GPU_BACKENDS)) + tiled_vae: bool | None = Field(default=None) + max_sequence_length: int | None = Field(default=None, ge=1) + + @model_validator(mode="after") + def _validate_backends(self) -> "CompareRequest": + if not self.backends: + raise ValueError("backends must contain at least one entry.") + unknown = [b for b in self.backends if b not in GPU_BACKENDS] + if unknown: + raise ValueError(f"Unknown backend(s): {unknown}; expected subset of {list(GPU_BACKENDS)}.") + if len(set(self.backends)) != len(self.backends): + raise ValueError("backends must not contain duplicates.") + return self + + +@app.get("/healthz") +async def healthz() -> dict[str, str]: + return {"status": "ok"} + + +@app.post( + "/generate", + response_class=Response, + dependencies=[Depends(_verify_bearer)], + responses={ + 200: { + "content": {"image/png": {"schema": {"type": "string", "format": "binary"}}}, + "description": "Generated PNG image.", + } + }, +) +async def generate(request: GenerateRequest) -> Response: + pipeline: GpuPipeline = app.state.pipeline + try: + pipeline.ensure_backend(backend=request.backend) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + gen_start = time.perf_counter() + image_bytes = pipeline.generate_png( + prompt=request.prompt, + seed=request.seed, + steps=request.steps, + height=request.height, + width=request.width, + guidance=request.guidance, + tiled_vae=request.tiled_vae, + max_sequence_length=request.max_sequence_length, + ) + wall_seconds = time.perf_counter() - gen_start + headers = {"X-Wall-Seconds": f"{wall_seconds:.3f}"} + if pipeline.last_peak_memory_mb is not None: + headers["X-Peak-Memory-MB"] = f"{pipeline.last_peak_memory_mb:.1f}" + return Response(content=image_bytes, media_type="image/png", headers=headers) + + +@app.post("/generate/compare", dependencies=[Depends(_verify_bearer)]) +async def generate_compare(request: CompareRequest) -> dict: + pipeline: GpuPipeline = app.state.pipeline + results: list[dict] = [] + for target_backend in request.backends: + swap_start = time.perf_counter() + try: + pipeline.ensure_backend(backend=target_backend) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + swap_seconds = time.perf_counter() - swap_start + + gen_start = time.perf_counter() + image_bytes = pipeline.generate_png( + prompt=request.prompt, + seed=request.seed, + steps=request.steps, + height=request.height, + width=request.width, + guidance=request.guidance, + tiled_vae=request.tiled_vae, + max_sequence_length=request.max_sequence_length, + ) + wall_seconds = time.perf_counter() - gen_start + + results.append( + { + "backend": target_backend, + "png_b64": base64.b64encode(image_bytes).decode("ascii"), + "wall_seconds": wall_seconds, + "swap_seconds": swap_seconds, + } + ) + return {"results": results} + + +__all__ = [ + "app", + "GenerateRequest", + "CompareRequest", + "generate", + "generate_compare", + "healthz", +] diff --git a/image-studio/backend_gpu/tests/__init__.py b/image-studio/backend_gpu/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/image-studio/backend_gpu/tests/test_loaders.py b/image-studio/backend_gpu/tests/test_loaders.py new file mode 100644 index 000000000..1cd2244c5 --- /dev/null +++ b/image-studio/backend_gpu/tests/test_loaders.py @@ -0,0 +1,629 @@ +from __future__ import annotations + +# Loader-internal tests. Mocks `gemlite`, `hqq`, `transformers`, and +# `diffusers` via `sys.modules` so the loaders run on macOS without a +# CUDA stack. +# +# .venv/bin/python -m unittest backend_gpu.tests.test_loaders -v + +import json +import os +import sys +import tempfile +import types +import unittest +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import MagicMock, patch + +# GpuPipeline now requires every artifact path to be explicitly set (no +# baked-in absolute defaults — different hosts have incompatible layouts). +# Inject "fake but set" paths at import time; the loaders themselves are +# patched in each test so these strings are never opened. +os.environ.setdefault("MFLUX_STUDIO_GPU_BINARY_TRANSFORMER_PATH", "/tmp/__binary__") +os.environ.setdefault("MFLUX_STUDIO_GPU_TERNARY_TRANSFORMER_PATH", "/tmp/__ternary__") +os.environ.setdefault("MFLUX_STUDIO_GPU_TEXT_ENCODER_PATH", "/tmp/__te__") +os.environ.setdefault("MFLUX_STUDIO_GPU_VAE_PATH", "/tmp/__vae__") +os.environ.setdefault("MFLUX_STUDIO_GPU_TOKENIZER_PATH", "/tmp/__tok__") + + +@contextmanager +def _inject_module(name: str, module: types.ModuleType): + """Set sys.modules[name]=module then restore on exit (single key only). + + Why not `patch.dict(sys.modules, {...})`: that snapshots and wholesale- + restores the entire dict on exit, which wipes any modules imported during + the test (e.g. torch + all submodules). A second `import torch` then tries + to re-init the C extension and fails ('docstring already set'). + """ + prev = sys.modules.get(name) + sys.modules[name] = module + try: + yield + finally: + if prev is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = prev + +from backend_gpu.pipeline_gpu import ( + GpuPipeline, + _load_gemlite_layers_from_state, + _load_gemlite_transformer, + _load_scheduler, + _load_text_encoder, + _load_tokenizer, + _load_vae, + _null_gemlite_weights, +) + + +class LoadGemliteLayersFromStateTest(unittest.TestCase): + def test_buckets_keys_and_calls_per_layer_load(self) -> None: + # Real torch model with two nn.Linear children at known FQNs. Saved + # state_dict has gemlite-layer keys for one of them and ordinary + # weights for the other; only the gemlite one should be replaced. + import torch + import torch.nn as nn + + class _FakeGemLite(nn.Module): + constructed: list[tuple[int, int]] = [] + loaded: list[dict] = [] + + def __init__(self, *, W_nbits, group_size, in_features, out_features, input_dtype, output_dtype): + super().__init__() + self.W_nbits = W_nbits + self.group_size = group_size + self.in_features = in_features + self.out_features = out_features + _FakeGemLite.constructed.append((in_features, out_features)) + # Tensor attrs filled by load_state_dict; pre-set so .to() works. + self.W_q = torch.zeros(1) + self.scales = torch.zeros(1) + self.zeros = torch.zeros(1) + self.bias = None + + def load_state_dict(self, sd, strict=True): # noqa: ARG002 - mirror gemlite shape + _FakeGemLite.loaded.append(dict(sd)) + # Mimic gemlite: pop tensor keys onto attrs (caller .to()s after). + for k, v in sd.items(): + setattr(self, k, v) + + def forward(self, x): + return x + + class _FakeDType: + FP16 = "fp16" + + model = nn.Module() + model.attn = nn.Module() + model.attn.to_q = nn.Linear(128, 256, bias=False) + model.norm_out = nn.Linear(64, 32, bias=False) # not in state's gemlite bucket + + state = { + # gemlite-layer keys for attn.to_q + "attn.to_q.W_q": torch.ones(2, 4, dtype=torch.uint8), + "attn.to_q.scales": torch.ones(2), + "attn.to_q.zeros": torch.ones(2), + "attn.to_q.metadata": torch.zeros(8, dtype=torch.int32), + "attn.to_q.orig_shape": torch.tensor([256, 128]), + "attn.to_q.bias": torch.zeros(256), + # plain weight for norm_out + "norm_out.weight": torch.randn(32, 64), + # head-level scalar (e.g. time_embed) + "time_embed.weight": torch.randn(8, 8), + } + + n, remainder = _load_gemlite_layers_from_state( + model, state, + bits=1, group_size=128, device="cpu", + DType=_FakeDType, GemLiteLinearTriton=_FakeGemLite, + ) + self.assertEqual(n, 1) + self.assertEqual(_FakeGemLite.constructed, [(128, 256)]) + self.assertEqual(set(remainder.keys()), {"norm_out.weight", "time_embed.weight"}) + # The replaced child is the fake gemlite layer. + self.assertIsInstance(model.attn.to_q, _FakeGemLite) + # Custom load_state_dict was dispatched once with the 6 bucketed keys. + self.assertEqual(len(_FakeGemLite.loaded), 1) + self.assertEqual( + set(_FakeGemLite.loaded[0].keys()), + {"W_q", "bias", "scales", "zeros", "metadata", "orig_shape"}, + ) + + def test_raises_when_state_targets_non_linear(self) -> None: + import torch + import torch.nn as nn + + class _FakeGemLite(nn.Module): + def __init__(self, **kw): + super().__init__() + + def forward(self, x): + return x + + class _FakeDType: + FP16 = "fp16" + + model = nn.Module() + # No nn.Linear at "ghost" — just a Module. + model.ghost = nn.Module() + + state = { + "ghost.W_q": torch.ones(1, dtype=torch.uint8), + "ghost.scales": torch.ones(1), + } + with self.assertRaisesRegex(RuntimeError, "but model has Module"): + _load_gemlite_layers_from_state( + model, state, + bits=1, group_size=128, device="cpu", + DType=_FakeDType, GemLiteLinearTriton=_FakeGemLite, + ) + + +class NullGemliteWeightsTest(unittest.TestCase): + def test_nulls_only_gemlite_modules(self) -> None: + import torch.nn as nn + + class _FakeGemLite(nn.Module): + def __init__(self): + super().__init__() + self.weight = "still-here" + + def forward(self, x): + return x + + model = nn.Module() + model.gl = _FakeGemLite() + model.linear = nn.Linear(4, 4, bias=False) + original_linear_weight = model.linear.weight + + nulled = _null_gemlite_weights(model, _FakeGemLite) + + self.assertEqual(nulled, 1) + self.assertIsNone(model.gl.weight) + self.assertIs(model.linear.weight, original_linear_weight) + + +class TextEncoderLoaderTest(unittest.TestCase): + def test_missing_path_raises_file_not_found(self) -> None: + bogus = Path("/tmp/definitely-does-not-exist-te") + with self.assertRaisesRegex(FileNotFoundError, "Text encoder artifact not found"): + _load_text_encoder(bogus) + + def test_loads_via_hqq_when_path_exists(self) -> None: + with tempfile.TemporaryDirectory() as td: + artifact = Path(td) / "te" + artifact.mkdir() + + fake_torch = MagicMock(name="torch") + fake_torch.float16 = "fp16-sentinel" + fake_gemlite_core = MagicMock(name="gemlite.core") + fake_gemlite = types.ModuleType("gemlite") + fake_gemlite.core = fake_gemlite_core + fake_hqq_models_hf_base = MagicMock(name="hqq.models.hf.base") + fake_hqq_utils_patching = MagicMock(name="hqq.utils.patching") + fake_loaded_model = MagicMock(name="loaded_te_model") + fake_hqq_models_hf_base.AutoHQQHFModel.from_quantized.return_value = fake_loaded_model + + modules = { + "torch": fake_torch, + "gemlite": fake_gemlite, + "gemlite.core": fake_gemlite_core, + "hqq": types.ModuleType("hqq"), + "hqq.models": types.ModuleType("hqq.models"), + "hqq.models.hf": types.ModuleType("hqq.models.hf"), + "hqq.models.hf.base": fake_hqq_models_hf_base, + "hqq.utils": types.ModuleType("hqq.utils"), + "hqq.utils.patching": fake_hqq_utils_patching, + } + with patch.dict(sys.modules, modules): + result = _load_text_encoder(artifact, device="cuda:0") + self.assertIs(result, fake_loaded_model) + fake_gemlite_core.set_packing_bitwidth.assert_called_once_with(8) + fake_hqq_models_hf_base.AutoHQQHFModel.from_quantized.assert_called_once_with( + str(artifact), compute_dtype="fp16-sentinel", device="cuda:0", + ) + fake_hqq_utils_patching.prepare_for_inference.assert_called_once_with( + fake_loaded_model, backend="gemlite", + ) + + +class VaeLoaderTest(unittest.TestCase): + def test_missing_path_raises(self) -> None: + with self.assertRaisesRegex(FileNotFoundError, "VAE snapshot not found"): + _load_vae(Path("/tmp/__no_vae_snapshot__")) + + def test_calls_autoencoder_klflux2_from_pretrained(self) -> None: + with tempfile.TemporaryDirectory() as td: + vae_path = Path(td) / "vae" + vae_path.mkdir() + + fake_torch = MagicMock(name="torch") + fake_torch.bfloat16 = "bf16-sentinel" + fake_diffusers = MagicMock(name="diffusers") + fake_vae = MagicMock(name="vae_model") + fake_diffusers.AutoencoderKLFlux2.from_pretrained.return_value = fake_vae + fake_vae.to.return_value = fake_vae + fake_vae.eval.return_value = fake_vae + + modules = { + "torch": fake_torch, + "diffusers": fake_diffusers, + } + with patch.dict(sys.modules, modules): + result = _load_vae(vae_path, device="cuda:0") + + self.assertIs(result, fake_vae) + fake_diffusers.AutoencoderKLFlux2.from_pretrained.assert_called_once_with( + str(vae_path), torch_dtype="bf16-sentinel", + ) + fake_vae.to.assert_called_once_with("cuda:0") + fake_vae.eval.assert_called_once_with() + + +class TokenizerLoaderTest(unittest.TestCase): + def test_calls_auto_tokenizer_from_pretrained(self) -> None: + # Klein/Qwen3 path: `_encode_klein_qwen3_prompt` only needs + # `apply_chat_template` + `__call__`, which AutoTokenizer provides. + # The TE artifact ships its own `tokenizer/` subdir. + fake_transformers = MagicMock(name="transformers") + fake_tok = MagicMock(name="tokenizer") + fake_transformers.AutoTokenizer.from_pretrained.return_value = fake_tok + + with patch.dict(sys.modules, {"transformers": fake_transformers}): + result = _load_tokenizer("/root/models/klein-4b-text-encoder-hqq-4bit-gemlite/tokenizer/") + + self.assertIs(result, fake_tok) + fake_transformers.AutoTokenizer.from_pretrained.assert_called_once_with( + "/root/models/klein-4b-text-encoder-hqq-4bit-gemlite/tokenizer/", + ) + + +class SchedulerLoaderTest(unittest.TestCase): + def test_returns_none_when_subfolder_missing(self) -> None: + with tempfile.TemporaryDirectory() as td: + tx_path = Path(td) / "tx" + tx_path.mkdir() + self.assertIsNone(_load_scheduler(tx_path)) + + def test_loads_when_subfolder_present(self) -> None: + with tempfile.TemporaryDirectory() as td: + tx_path = Path(td) / "tx" + (tx_path / "scheduler").mkdir(parents=True) + + fake_diffusers = MagicMock(name="diffusers") + fake_sched = MagicMock(name="scheduler") + fake_diffusers.FlowMatchEulerDiscreteScheduler.from_pretrained.return_value = fake_sched + + with patch.dict(sys.modules, {"diffusers": fake_diffusers}): + result = _load_scheduler(tx_path) + + self.assertIs(result, fake_sched) + fake_diffusers.FlowMatchEulerDiscreteScheduler.from_pretrained.assert_called_once_with( + str(tx_path), subfolder="scheduler", + ) + + +class GemliteTransformerLoaderTest(unittest.TestCase): + def _make_artifact(self, td: str, *, qcfg_extra: dict | None = None) -> Path: + path = Path(td) / "transformer" + path.mkdir() + (path / "state_dict.pt").write_bytes(b"fake") + (path / "config.json").write_text(json.dumps({"in_channels": 16})) + qcfg = { + "format": "gemlite-int1-g128", + "bits": 1, + "group_size": 128, + "packing_bitwidth": 8, + } + if qcfg_extra: + qcfg.update(qcfg_extra) + (path / "quantization_config.json").write_text(json.dumps(qcfg)) + (path / "gemlite_autotune.json").write_text("{}") + return path + + def test_missing_path_raises(self) -> None: + with self.assertRaisesRegex(FileNotFoundError, "Gemlite transformer artifact not found"): + _load_gemlite_transformer(Path("/tmp/__nope__")) + + def test_missing_state_dict_raises(self) -> None: + with tempfile.TemporaryDirectory() as td: + path = self._make_artifact(td) + (path / "state_dict.pt").unlink() + with self.assertRaisesRegex(FileNotFoundError, "state_dict.pt"): + _load_gemlite_transformer(path) + + def test_missing_autotune_raises(self) -> None: + with tempfile.TemporaryDirectory() as td: + path = self._make_artifact(td) + (path / "gemlite_autotune.json").unlink() + with self.assertRaisesRegex(FileNotFoundError, "gemlite_autotune.json"): + _load_gemlite_transformer(path) + + def test_full_load_call_sequence(self) -> None: + # Models the chain: + # set_packing_bitwidth(packing_bw) + # GemLiteLinearTriton.load_config(autotune_path) + # Flux2Transformer2DModel.from_config(cfg) -> m_cfg + # m_cfg.to(bf16) -> m_bf16 + # torch.load(state_path) -> state + # _load_gemlite_layers_from_state(m_bf16, state, ...) -> (n, remainder) + # m_bf16.load_state_dict(remainder, strict=False) + # m_bf16.to(fp16) -> m_fp16 + # _null_gemlite_weights(m_fp16, GemLite) + # m_fp16.to(device) -> m_dev + # m_dev.eval() -> m_dev + with tempfile.TemporaryDirectory() as td: + path = self._make_artifact(td) + + fake_torch = MagicMock(name="torch") + fake_torch.bfloat16 = "bf16-sentinel" + fake_torch.float16 = "fp16-sentinel" + fake_state = {"fake.key": "v", "x.W_q": "g"} + fake_torch.load.return_value = fake_state + + fake_gemlite_core = MagicMock(name="gemlite.core") + fake_gemlite = types.ModuleType("gemlite") + fake_gemlite.core = fake_gemlite_core + + fake_diffusers = MagicMock(name="diffusers") + m_cfg = MagicMock(name="m_from_config") + m_bf16 = MagicMock(name="m_bf16") + m_fp16 = MagicMock(name="m_fp16") + m_dev = MagicMock(name="m_dev") + m_cfg.to.return_value = m_bf16 + m_bf16.load_state_dict.return_value = ([], []) + m_bf16.to.return_value = m_fp16 + m_fp16.to.return_value = m_dev + m_dev.eval.return_value = m_dev + fake_diffusers.Flux2Transformer2DModel.from_config.return_value = m_cfg + + modules = { + "torch": fake_torch, + "gemlite": fake_gemlite, + "gemlite.core": fake_gemlite_core, + "diffusers": fake_diffusers, + } + stub_remainder = {"some.weight": "rem"} + with ( + patch.dict(sys.modules, modules), + patch( + "backend_gpu.pipeline_gpu._load_gemlite_layers_from_state", + return_value=(140, stub_remainder), + ) as mock_loader, + patch( + "backend_gpu.pipeline_gpu._null_gemlite_weights", return_value=140, + ) as mock_nuller, + ): + result = _load_gemlite_transformer(path, device="cuda:0") + + self.assertIs(result, m_dev) + # Module-level set_packing_bitwidth precedes classmethod load_config. + fake_gemlite_core.set_packing_bitwidth.assert_called_once_with(8) + fake_gemlite_core.GemLiteLinearTriton.load_config.assert_called_once_with( + str(path / "gemlite_autotune.json"), + ) + + fake_diffusers.Flux2Transformer2DModel.from_config.assert_called_once() + m_cfg.to.assert_called_once_with("bf16-sentinel") + fake_torch.load.assert_called_once_with(str(path / "state_dict.pt"), map_location="cpu") + + mock_loader.assert_called_once() + kw = mock_loader.call_args.kwargs + self.assertEqual(kw["bits"], 1) + self.assertEqual(kw["group_size"], 128) + self.assertEqual(kw["device"], "cuda:0") + # Loader gets the raw state from torch.load. + self.assertIs(mock_loader.call_args.args[1], fake_state) + + # Remainder (non-gemlite keys) goes through model.load_state_dict. + m_bf16.load_state_dict.assert_called_once_with(stub_remainder, strict=False) + # bf16 → fp16 cast happens before device move. + m_bf16.to.assert_called_once_with("fp16-sentinel") + mock_nuller.assert_called_once() + self.assertIs(mock_nuller.call_args.args[0], m_fp16) + m_fp16.to.assert_called_once_with("cuda:0") + m_dev.eval.assert_called_once_with() + + +class PrewarmErrorPathTest(unittest.TestCase): + def test_missing_transformer_artifact_surfaces_error(self) -> None: + pipe = GpuPipeline( + transformer_path="/tmp/__no_transformer__", + text_encoder_path="/tmp/__no_te__", + vae_path="/tmp/__no_vae__", + tokenizer_path="fake", + ) + with self.assertRaisesRegex(FileNotFoundError, "Gemlite transformer artifact not found"): + pipe.prewarm() + self.assertFalse(pipe.ready) + + def test_missing_text_encoder_surfaces_file_not_found(self) -> None: + pipe = GpuPipeline( + transformer_path="/tmp/__no_transformer__", + text_encoder_path="/tmp/__no_te__", + vae_path="/tmp/__no_vae__", + tokenizer_path="fake", + ) + with patch("backend_gpu.pipeline_gpu._load_gemlite_transformer", return_value=MagicMock()): + with self.assertRaisesRegex(FileNotFoundError, "Text encoder artifact not found"): + pipe.prewarm() + self.assertFalse(pipe.ready) + + def test_missing_vae_surfaces_file_not_found(self) -> None: + pipe = GpuPipeline( + transformer_path="/tmp/__no_transformer__", + text_encoder_path="/tmp/__no_te__", + vae_path="/tmp/__no_vae__", + tokenizer_path="fake", + ) + with ( + patch("backend_gpu.pipeline_gpu._load_gemlite_transformer", return_value=MagicMock()), + patch("backend_gpu.pipeline_gpu._load_text_encoder", return_value=MagicMock()), + ): + with self.assertRaisesRegex(FileNotFoundError, "VAE snapshot not found"): + pipe.prewarm() + self.assertFalse(pipe.ready) + + def test_prewarm_succeeds_with_all_loaders_mocked(self) -> None: + pipe = GpuPipeline() + with ( + patch("backend_gpu.pipeline_gpu._load_gemlite_transformer", return_value="tx"), + patch("backend_gpu.pipeline_gpu._load_text_encoder", return_value="te"), + patch("backend_gpu.pipeline_gpu._load_vae", return_value="vae"), + patch("backend_gpu.pipeline_gpu._load_tokenizer", return_value="tok"), + patch("backend_gpu.pipeline_gpu._load_scheduler", return_value="sched"), + ): + pipe.prewarm() + self.assertTrue(pipe.ready) + self.assertEqual(pipe._transformer, "tx") + self.assertEqual(pipe._text_encoder, "te") + self.assertEqual(pipe._vae, "vae") + self.assertEqual(pipe._tokenizer, "tok") + self.assertEqual(pipe._scheduler, "sched") + + def test_prewarm_tolerates_absent_scheduler(self) -> None: + pipe = GpuPipeline() + with ( + patch("backend_gpu.pipeline_gpu._load_gemlite_transformer", return_value="tx"), + patch("backend_gpu.pipeline_gpu._load_text_encoder", return_value="te"), + patch("backend_gpu.pipeline_gpu._load_vae", return_value="vae"), + patch("backend_gpu.pipeline_gpu._load_tokenizer", return_value="tok"), + patch("backend_gpu.pipeline_gpu._load_scheduler", return_value=None), + ): + pipe.prewarm() + self.assertTrue(pipe.ready) + self.assertIsNone(pipe._scheduler) + + +class GeneratePngTest(unittest.TestCase): + """Unit-level coverage for `GpuPipeline.generate_png` after Phase 5c-3 wire-up. + + Server-level coverage lives in test_server.py; these test the kwargs-routing + contract, ready-gate, and PNG roundtrip without going through FastAPI. + """ + + def _make_ready_pipeline(self) -> GpuPipeline: + pipe = GpuPipeline() + pipe._transformer = MagicMock(name="transformer") + pipe._text_encoder = MagicMock(name="text_encoder") + pipe._tokenizer = MagicMock(name="tokenizer") + pipe._vae = MagicMock(name="vae") + pipe._scheduler = MagicMock(name="scheduler") + pipe._ready = True + return pipe + + def _inject_fake_diffusion_klein(self, *, image_size: tuple[int, int] = (32, 32)): + from PIL import Image + + fake_dk = types.ModuleType("backend_gpu.diffusion_klein") + captured: dict = {} + + def fake_forward(**kw): + captured.clear() + captured.update(kw) + return Image.new("RGB", image_size, (200, 100, 50)) + + fake_dk.diffusion_forward = fake_forward + return fake_dk, captured + + def test_raises_when_not_ready(self) -> None: + pipe = GpuPipeline() + with self.assertRaisesRegex(RuntimeError, "prewarm.*before generate_png"): + pipe.generate_png(prompt="x") + + def test_passes_kwargs_to_diffusion_forward_and_returns_png(self) -> None: + pipe = self._make_ready_pipeline() + fake_dk, captured = self._inject_fake_diffusion_klein(image_size=(64, 48)) + + with _inject_module("backend_gpu.diffusion_klein", fake_dk): + result = pipe.generate_png( + prompt="a bonsai", seed=7, steps=12, height=48, width=64, guidance=4.5, + ) + + self.assertIsInstance(result, bytes) + self.assertTrue(result.startswith(b"\x89PNG\r\n\x1a\n")) + # Routing contract: every constructor-injected artifact + the user kwargs + # land in diffusion_forward with the expected names. + self.assertIs(captured["transformer"], pipe._transformer) + self.assertIs(captured["text_encoder"], pipe._text_encoder) + self.assertIs(captured["tokenizer"], pipe._tokenizer) + self.assertIs(captured["vae"], pipe._vae) + self.assertIs(captured["scheduler"], pipe._scheduler) + self.assertEqual(captured["prompt"], "a bonsai") + self.assertEqual(captured["seed"], 7) + self.assertEqual(captured["num_steps"], 12) + self.assertEqual(captured["height"], 48) + self.assertEqual(captured["width"], 64) + self.assertEqual(captured["guidance"], 4.5) + # max_sequence_length omitted when caller passes None — diffusion_forward's + # own default (512) wins, not a hardcoded override. + self.assertNotIn("max_sequence_length", captured) + # No CUDA on macOS ⇒ peak memory recorded as 0.0 (not None). + self.assertEqual(pipe.last_peak_memory_mb, 0.0) + + def test_passes_max_sequence_length_when_set(self) -> None: + pipe = self._make_ready_pipeline() + fake_dk, captured = self._inject_fake_diffusion_klein() + + with _inject_module("backend_gpu.diffusion_klein", fake_dk): + pipe.generate_png(prompt="x", max_sequence_length=256) + self.assertEqual(captured["max_sequence_length"], 256) + + def test_records_cuda_peak_memory_when_available(self) -> None: + # Patch `torch.cuda` attrs in place rather than swapping `sys.modules["torch"]` + # — torch's C-level init state breaks if reimported, so a clean restore is + # not enough; we'd corrupt later test_server tests on the way out. + import torch + + pipe = self._make_ready_pipeline() + fake_dk, _ = self._inject_fake_diffusion_klein() + + with ( + _inject_module("backend_gpu.diffusion_klein", fake_dk), + patch.object(torch.cuda, "is_available", return_value=True), + patch.object(torch.cuda, "max_memory_allocated", return_value=256 * 1024 * 1024), + patch.object(torch.cuda, "reset_peak_memory_stats") as mock_reset, + ): + pipe.generate_png(prompt="x") + mock_reset.assert_called_once_with() + self.assertEqual(pipe.last_peak_memory_mb, 256.0) + + +class ConfigEnvVarTest(unittest.TestCase): + def test_env_vars_picked_up(self) -> None: + # Use the modern, backend-suffixed env name — the legacy unsuffixed + # MFLUX_STUDIO_GPU_TRANSFORMER_PATH is still honored as a fallback + # (covered by test_explicit_kwargs_override_env's legacy probe). + env_overrides = { + "MFLUX_STUDIO_GPU_BINARY_TRANSFORMER_PATH": "/some/tx/path", + "MFLUX_STUDIO_GPU_TEXT_ENCODER_PATH": "/some/te/path", + "MFLUX_STUDIO_GPU_VAE_PATH": "/some/vae/path", + "MFLUX_STUDIO_GPU_TOKENIZER_PATH": "foo/bar-tok", + "MFLUX_STUDIO_GPU_DEVICE": "cuda:1", + } + with patch.dict("os.environ", env_overrides, clear=False): + pipe = GpuPipeline() + self.assertEqual(pipe.transformer_path, Path("/some/tx/path")) + self.assertEqual(pipe.text_encoder_path, Path("/some/te/path")) + self.assertEqual(pipe.vae_path, Path("/some/vae/path")) + self.assertEqual(pipe.tokenizer_path, "foo/bar-tok") + self.assertEqual(pipe.device, "cuda:1") + + def test_explicit_kwargs_override_env(self) -> None: + # Legacy `transformer_path` kwarg must beat the legacy unsuffixed env + # (`MFLUX_STUDIO_GPU_TRANSFORMER_PATH`). Note: the BINARY-suffixed env + # (when set) takes priority over the legacy kwarg by design — see the + # fallback chain in GpuPipeline.__init__. We clear BINARY here so the + # legacy-only path is exercised. + env = os.environ.copy() + env.pop("MFLUX_STUDIO_GPU_BINARY_TRANSFORMER_PATH", None) + env["MFLUX_STUDIO_GPU_TRANSFORMER_PATH"] = "/from/legacy/env" + with patch.dict("os.environ", env, clear=True): + pipe = GpuPipeline(transformer_path="/from/kwarg") + self.assertEqual(pipe.transformer_path, Path("/from/kwarg")) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/image-studio/backend_gpu/tests/test_server.py b/image-studio/backend_gpu/tests/test_server.py new file mode 100644 index 000000000..691cdafa4 --- /dev/null +++ b/image-studio/backend_gpu/tests/test_server.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +# Unit tests for backend_gpu.server. The 4 loaders are patched out so this +# whole suite runs anywhere — no GPU, no gemlite/HQQ/diffusers required. +# +# .venv/bin/python -m unittest backend_gpu.tests.test_server -v + +import base64 +import os +import sys +import types +import unittest +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +from PIL import Image + +from fastapi.testclient import TestClient + + +_TEST_TOKEN = "test-bearer-token" +_PNG_MAGIC = b"\x89PNG\r\n\x1a\n" + + +@contextmanager +def _patched_app(): + # Why: GpuPipeline.prewarm() loads 5 heavy artifacts (gemlite/HQQ/ + # diffusers) and generate_png imports `backend_gpu.diffusion_klein`, + # which itself pulls in `diffusers.pipelines.flux2`. We inject a fake + # diffusion_klein module + patch the loaders so the server tests stay + # GPU-free and don't need the heavy diffusers stack. + os.environ["MFLUX_STUDIO_GPU_TOKEN"] = _TEST_TOKEN + # The artifact paths now require explicit env vars (loaders are patched + # below so the strings never actually get opened). + os.environ.setdefault("MFLUX_STUDIO_GPU_BINARY_TRANSFORMER_PATH", "/root/models/bonsai-binary") + os.environ.setdefault("MFLUX_STUDIO_GPU_TERNARY_TRANSFORMER_PATH", "/root/models/bonsai-ternary") + os.environ.setdefault("MFLUX_STUDIO_GPU_TEXT_ENCODER_PATH", "/root/models/text-encoder") + os.environ.setdefault("MFLUX_STUDIO_GPU_VAE_PATH", "/root/models/vae") + os.environ.setdefault("MFLUX_STUDIO_GPU_TOKENIZER_PATH", "/root/models/text-encoder/tokenizer") + fake_transformer = MagicMock(name="transformer") + fake_te = MagicMock(name="text_encoder") + fake_vae = MagicMock(name="vae") + fake_tokenizer = MagicMock(name="tokenizer") + + fake_dk = types.ModuleType("backend_gpu.diffusion_klein") + fake_dk.diffusion_forward = MagicMock( + side_effect=lambda **kw: Image.new("RGB", (kw["width"], kw["height"]), (12, 34, 56)), + ) + fake_dk.DEFAULT_NUM_STEPS = 4 + fake_dk.DEFAULT_GUIDANCE = 1.0 + + # Why not patch.dict(sys.modules, ...): patch.dict snapshots the dict on + # enter and *wholesale-restores* on exit, which wipes any modules (e.g. + # torch + all submodules) loaded mid-test. A second `import torch` then + # tries to re-init the C extension → "docstring already set". Manual + # set+pop touches only the one key we own. + prev_dk = sys.modules.get("backend_gpu.diffusion_klein") + sys.modules["backend_gpu.diffusion_klein"] = fake_dk + try: + with ( + patch("backend_gpu.pipeline_gpu._load_gemlite_transformer", return_value=fake_transformer), + patch("backend_gpu.pipeline_gpu._load_text_encoder", return_value=fake_te), + patch("backend_gpu.pipeline_gpu._load_vae", return_value=fake_vae), + patch("backend_gpu.pipeline_gpu._load_tokenizer", return_value=fake_tokenizer), + ): + from backend_gpu.server import app + + with TestClient(app) as client: + yield client + finally: + if prev_dk is None: + sys.modules.pop("backend_gpu.diffusion_klein", None) + else: + sys.modules["backend_gpu.diffusion_klein"] = prev_dk + + +class HealthzTest(unittest.TestCase): + def test_healthz_no_auth(self) -> None: + with _patched_app() as client: + response = client.get("/healthz") + self.assertEqual(response.status_code, 200) + self.assertEqual(response.json(), {"status": "ok"}) + + +class AuthTest(unittest.TestCase): + def test_generate_without_token_rejected(self) -> None: + with _patched_app() as client: + response = client.post("/generate", json={"prompt": "x"}) + self.assertEqual(response.status_code, 401) + + def test_generate_with_wrong_token_rejected(self) -> None: + with _patched_app() as client: + response = client.post( + "/generate", + json={"prompt": "x"}, + headers={"Authorization": "Bearer not-the-right-token"}, + ) + self.assertEqual(response.status_code, 401) + + def test_compare_without_token_rejected(self) -> None: + with _patched_app() as client: + response = client.post("/generate/compare", json={"prompt": "x"}) + self.assertEqual(response.status_code, 401) + + +class GenerateTest(unittest.TestCase): + def test_generate_returns_png(self) -> None: + with _patched_app() as client: + response = client.post( + "/generate", + json={"prompt": "a cat", "width": 256, "height": 256}, + headers={"Authorization": f"Bearer {_TEST_TOKEN}"}, + ) + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(response.headers["content-type"], "image/png") + self.assertTrue(response.content.startswith(_PNG_MAGIC)) + self.assertIn("X-Wall-Seconds", response.headers) + self.assertIn("X-Peak-Memory-MB", response.headers) + + def test_generate_rejects_unknown_backend(self) -> None: + with _patched_app() as client: + response = client.post( + "/generate", + json={"prompt": "x", "backend": "bogus-backend"}, + headers={"Authorization": f"Bearer {_TEST_TOKEN}"}, + ) + self.assertEqual(response.status_code, 422) + + def test_generate_rejects_zero_steps(self) -> None: + with _patched_app() as client: + response = client.post( + "/generate", + json={"prompt": "x", "steps": 0}, + headers={"Authorization": f"Bearer {_TEST_TOKEN}"}, + ) + self.assertEqual(response.status_code, 422) + + def test_generate_rejects_empty_prompt(self) -> None: + with _patched_app() as client: + response = client.post( + "/generate", + json={"prompt": ""}, + headers={"Authorization": f"Bearer {_TEST_TOKEN}"}, + ) + self.assertEqual(response.status_code, 422) + + +class CompareTest(unittest.TestCase): + def test_compare_default_backends(self) -> None: + with _patched_app() as client: + response = client.post( + "/generate/compare", + json={"prompt": "a cat"}, + headers={"Authorization": f"Bearer {_TEST_TOKEN}"}, + ) + self.assertEqual(response.status_code, 200, response.text) + body = response.json() + self.assertEqual(len(body["results"]), 2) + self.assertEqual( + [r["backend"] for r in body["results"]], + ["bonsai-binary-gemlite", "bonsai-ternary-gemlite"], + ) + for result in body["results"]: + decoded = base64.b64decode(result["png_b64"]) + self.assertTrue(decoded.startswith(_PNG_MAGIC)) + self.assertGreaterEqual(result["wall_seconds"], 0.0) + self.assertGreaterEqual(result["swap_seconds"], 0.0) + + def test_compare_empty_list_rejected(self) -> None: + with _patched_app() as client: + response = client.post( + "/generate/compare", + json={"prompt": "x", "backends": []}, + headers={"Authorization": f"Bearer {_TEST_TOKEN}"}, + ) + self.assertEqual(response.status_code, 422) + + def test_compare_unknown_backend_rejected(self) -> None: + with _patched_app() as client: + response = client.post( + "/generate/compare", + json={"prompt": "x", "backends": ["bogus"]}, + headers={"Authorization": f"Bearer {_TEST_TOKEN}"}, + ) + self.assertEqual(response.status_code, 422) + + def test_compare_duplicates_rejected(self) -> None: + with _patched_app() as client: + response = client.post( + "/generate/compare", + json={"prompt": "x", "backends": ["bonsai-ternary-gemlite", "bonsai-ternary-gemlite"]}, + headers={"Authorization": f"Bearer {_TEST_TOKEN}"}, + ) + self.assertEqual(response.status_code, 422) + + +class NotReadyTest(unittest.TestCase): + def test_generate_raises_when_pipeline_not_ready(self) -> None: + # Lifespan prewarms by default; flip _ready off after entry to assert + # the GpuPipeline.generate_png guard fires (raises RuntimeError) and + # bubbles up. TestClient defaults to raise_server_exceptions=True so + # the unhandled error surfaces here rather than as a 500. + with _patched_app() as client: + client.app.state.pipeline._ready = False + with self.assertRaises(RuntimeError): + client.post( + "/generate", + json={"prompt": "x"}, + headers={"Authorization": f"Bearer {_TEST_TOKEN}"}, + ) + + +class StartupTest(unittest.TestCase): + def test_lifespan_requires_token(self) -> None: + os.environ.pop("MFLUX_STUDIO_GPU_TOKEN", None) + # Re-import to get a fresh app object whose lifespan has not yet run. + import importlib + + import backend_gpu.server as server_module + + importlib.reload(server_module) + with self.assertRaises(RuntimeError): + with TestClient(server_module.app): + pass + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/image-studio/docs/backend.md b/image-studio/docs/backend.md new file mode 100644 index 000000000..2948d8203 --- /dev/null +++ b/image-studio/docs/backend.md @@ -0,0 +1,101 @@ +# `backend/` — FastAPI Mac backend + +Local FastAPI server that fronts three FLUX.2 Klein generation arms (two resident on Mac MLX, one proxied to a remote GPU). Talks to `frontend/` (Next.js) and the iOS app (Bonsai's remote backends, when an on-device path isn't selected). + +## Layout + +``` +backend/ + server.py # FastAPI app, /generate, /backends, /generate/compare + pipeline.py # Klein/mflux pipeline construction, backend selection, env-var read + pipeline_remote_gpu.py # RemoteGpuPipeline — POSTs /generate to backend_gpu over HTTP + tests/ + test_backends_endpoint.py # /backends contract + GPU probe caching + test_ensure_backend.py # backend swap + locking semantics + test_generate_compare.py # /generate/compare side-by-side + smoke_* # adhoc smoke scripts; not part of pytest collection +``` + +## Run + +Requires Python 3.13 + `uv`. + +```sh +uv venv .venv +uv pip install --python .venv/bin/python -e . +.venv/bin/uvicorn backend.server:app --port 8000 +``` + +## API + +### `POST /generate` + +```json +{ + "prompt": "a small red cube on a table", + "seed": 42, + "steps": 4, + "guidance": 1.0, + "backend": "bonsai-ternary-mlx", + "height": 512, + "width": 512 +} +``` + +Returns `image/png` bytes. + +### `GET /backends` + +Returns `{available: [...], default: "...", gpu: {available: bool, reason: str}}`. `reason` is one of `force_disabled`, `no_gpu_host`, `no_gpu_token`, `healthz_failed:`, `healthz_unreachable`. Cached 30 s. + +### `POST /generate/compare` + +Multi-arm side-by-side. Same request shape with `backends: ["bonsai-ternary-mlx", "bfl-klein-bf16"]`. Returns one entry per arm. + +## The three backends + +| ID | What | Where | +| --- | --- | --- | +| `bonsai-ternary-mlx` | Bonsai (ternary) Klein on MLX | local Mac | +| `bfl-klein-bf16` | upstream mflux on Klein bf16 (`full` VAE) | local Mac | +| `bonsai-ternary-gemlite` | gemlite + HQQ TE + bf16 VAE | remote (`backend_gpu/`) | + +The two Mac arms share a single `Flux2Klein` slot; on switch the server evicts + rebuilds (`asyncio.Lock` serializes). The remote arm holds no in-process state; calls fan out over HTTP. See [`docs/backend_gpu.md`](backend_gpu.md). + +## Checkpoint layout + +The Mac MLX arms each consume a Klein root directory: + +``` +/ + transformer/ # bf16 dense weights + transformer-packed-mflux/ # uint32 packed weights + scales + text_encoder/ # Klein/Qwen3 bf16 TE + tokenizer/ # Qwen3 tokenizer + vae/ # FLUX.2 VAE + scheduler/ # FLUX.2 flow-match scheduler config + model_index.json + LICENSE.md # Apache 2.0 from upstream Klein +``` + +`bfl-klein-bf16` ignores the packed dir and resolves Klein from HF. + +## Env vars + +See [`README.md`](../README.md#env-vars). Most-load-bearing: + +| Var | Required | Purpose | +| --- | --- | --- | +| `MFLUX_STUDIO_BAKED_MODEL_PATH` | for `bonsai-ternary-mlx` | absolute path to the ternary checkpoint root | +| `MFLUX_STUDIO_STOCK_MODEL_PATH` | optional, for `bfl-klein-bf16` | absolute path to a local Klein snapshot; falls back to HF default | +| `MFLUX_STUDIO_DEFAULT_BACKEND` | no | initial backend at boot | +| `MFLUX_STUDIO_GPU_HOST` + `MFLUX_STUDIO_GPU_TOKEN` | for `bonsai-ternary-gemlite` | base URL + bearer token to `backend_gpu/` | +| `MFLUX_STUDIO_FORCE_DISABLE_GPU` | no | hides the GPU arm regardless of probe | + +## Tests + +```sh +.venv/bin/python -m pytest backend/tests/ +``` + +Smoke scripts (`smoke_*.py`) are not collected by pytest; run them directly when comparing baked checkpoints or evicting strategies. diff --git a/image-studio/docs/backend_gpu.md b/image-studio/docs/backend_gpu.md new file mode 100644 index 000000000..bcaf3996b --- /dev/null +++ b/image-studio/docs/backend_gpu.md @@ -0,0 +1,59 @@ +# `backend_gpu/` — GPU arm + +Standalone FastAPI server for the `bonsai-ternary-gemlite` backend. Runs on a CUDA host; `backend.pipeline.RemoteGpuPipeline` POSTs to it over HTTP. + +Pipeline: gemlite transformer + HQQ-int4 text encoder + bf16 VAE on a single H100. 4-step Klein defaults. + +This file is a pointer to the existing detailed docs in the component itself: + +- [`backend_gpu/README.md`](../backend_gpu/README.md) — layout, artifacts, env vars, smoke tests, autotune cache, run instructions. + +## Quick reference + +``` +backend_gpu/ + server.py # FastAPI: /healthz, /generate, /generate/compare + pipeline_gpu.py # GpuPipeline: 5-artifact prewarm + generate_png + diffusion_klein.py # Klein/Qwen3 text→image forward + scripts/smoke_*.py # local CUDA + remote round-trip smokes + tests/ # loader + server unit tests +``` + +## API contract + +Identical JSON shape to `backend/`'s `/generate`, plus a Bearer auth header. `/healthz` is unauthenticated. + +```sh +MFLUX_STUDIO_GPU_TOKEN=devtoken \ +uvicorn backend_gpu.server:app --host 0.0.0.0 --port 8801 + +curl -s http://localhost:8801/healthz +# {"status":"ok"} + +curl -s -o out.png \ + -H 'Authorization: Bearer devtoken' \ + -H 'Content-Type: application/json' \ + -d '{"prompt": "...", "seed": 42, "steps": 4, "guidance": 1.0, "backend": "bonsai-ternary-gemlite", "height": 512, "width": 512}' \ + http://localhost:8801/generate +``` + +## Required env + +| Var | Default | Purpose | +| --- | --- | --- | +| `MFLUX_STUDIO_GPU_TOKEN` | (required) | Bearer; server refuses to start without it | +| `MFLUX_STUDIO_GPU_TERNARY_TRANSFORMER_PATH` | (unset) | ternary transformer pack path | +| `MFLUX_STUDIO_GPU_TRANSFORMER_PATH` | (legacy alias for the ternary path) | retained for backward compatibility | +| `MFLUX_STUDIO_GPU_TEXT_ENCODER_PATH` | `/root/models/klein-4b-text-encoder-hqq-4bit-gemlite/` | HQQ-int4 TE pack | +| `MFLUX_STUDIO_GPU_VAE_PATH` | `/root/models/klein-4b-vae-bf16/` | bf16 VAE snapshot | +| `MFLUX_STUDIO_GPU_DEVICE` | `cuda:0` | target device | + +Full table + artifact regen recipes in [`backend_gpu/README.md`](../backend_gpu/README.md). + +## Tests + +```sh +.venv/bin/python -m pytest backend_gpu/tests/ +``` + +The unit tests stub out CUDA so they pass on any host. Real GPU smoke is `backend_gpu/scripts/smoke_e2e.py` (runs prewarm + a single forward) and `scripts/smoke_remote.py` (exercises `RemoteGpuPipeline` against a deployed server). diff --git a/image-studio/docs/frontend.md b/image-studio/docs/frontend.md new file mode 100644 index 000000000..e36c08e37 --- /dev/null +++ b/image-studio/docs/frontend.md @@ -0,0 +1,61 @@ +# `frontend/` — Next.js studio client + +Next.js 16 (App Router) + Tailwind 4 + Radix UI client for the four-backend Mac/GPU pipeline. Talks to `backend/` over HTTP at `http://localhost:8000` by default. + +## Layout + +``` +frontend/ + app/ # Next.js App Router pages + page.tsx # Studio (single-prompt → single-image) + api/ # Next.js route handlers (proxy to backend) + globals.css + layout.tsx + components/ + studio-client.tsx # Main interactive shell + compare-client.tsx # Three-arm side-by-side (POST /generate/compare) + result-panel.tsx # Right-rail result + metadata chips + history-grid.tsx # Generated-image history (localStorage-backed) + bonsai-background.tsx # Decorative SVG background, theme-aware tint + theme-toggle.tsx # Dark/light toggle (next-themes) + providers.tsx # Theme + history providers + ui/ # Radix-based primitives (Collapsible, etc.) + lib/ + backends.ts # GET /backends + the typed client + use-backends.ts # SWR-style hook around /backends + use-history.ts # localStorage-backed history hook + use-compare-history.ts # Compare-mode history + compare-presets.ts # Preset triples for /generate/compare + resolutions.ts # Resolution tier table (mirrors apple's enum) + utils.ts # className merging, tiny utils + public/ # Static assets (Bonsai SVG, brand) + tokens/ # Generated design tokens (mirror of repo-root tokens/) + next.config.ts # allowedDevOrigins for 127.0.0.1 + tsconfig.json + package.json +``` + +## Run + +```sh +cd frontend +npm install +npm run dev # http://localhost:3000 +``` + +Defaults to `http://localhost:8000` for the backend. Override via the route handlers in `app/api/` if you need a different base. + +## Build / lint + +```sh +npm run build +npm run lint +``` + +## Design tokens + +Tokens live at the repo root in `tokens/design-tokens.json` and are emitted into `frontend/tokens/` by `scripts/gen-design-tokens.py`. Run that script after editing the JSON. + +## Backend probe + +The `/backends` cache TTL is 30 s server-side. The frontend's `useBackends` hook caches per-page-load; refresh the page after toggling `MFLUX_STUDIO_FORCE_DISABLE_GPU` on the backend. diff --git a/image-studio/frontend/.gitignore b/image-studio/frontend/.gitignore new file mode 100644 index 000000000..6d539176e --- /dev/null +++ b/image-studio/frontend/.gitignore @@ -0,0 +1,44 @@ +# See https://help.github.com/articles/ignoring-files/ for more about ignoring files. + +# dependencies +/node_modules +/.pnp +.pnp.* +.yarn/* +!.yarn/patches +!.yarn/plugins +!.yarn/releases +!.yarn/versions + +# testing +/coverage + +# next.js +/.next/ +/out/ + +# production +/build + +# misc +.DS_Store +*.pem + +# debug +npm-debug.log* +yarn-debug.log* +yarn-error.log* +.pnpm-debug.log* + +# env files (can opt-in for committing if needed) +.env* + +# vercel +.vercel + +# playwright / test runs +/test-results + +# typescript +*.tsbuildinfo +next-env.d.ts diff --git a/image-studio/frontend/.gitkeep b/image-studio/frontend/.gitkeep new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/image-studio/frontend/.gitkeep @@ -0,0 +1 @@ + diff --git a/image-studio/frontend/__tests__/prompt-moderator.test.ts b/image-studio/frontend/__tests__/prompt-moderator.test.ts new file mode 100644 index 000000000..4dddc04b9 --- /dev/null +++ b/image-studio/frontend/__tests__/prompt-moderator.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from "vitest"; +import { validatePrompt } from "@/lib/prompt-moderator"; + +describe("validatePrompt", () => { + it("accepts neutral prompts", () => { + const neutral = [ + "A serene Scandinavian woman in her 40s, soft window light", + "A koi fish gliding through lily-pad covered water", + "Cinematic landscape of a forest at dawn", + "Abstract emerald ink swirling through clear water", + "Hand-lettered storefront sign at golden hour", + ]; + for (const prompt of neutral) { + expect(validatePrompt(prompt)).toEqual({ ok: true }); + } + }); + + it("rejects empty / whitespace-only prompts", () => { + expect(validatePrompt("")).toEqual({ ok: false, reason: "empty" }); + expect(validatePrompt(" \n\t ")).toEqual({ ok: false, reason: "empty" }); + }); + + it("rejects 5 representative profane prompts", () => { + // Each entry hits a different category from the bad-words default list. + // Phrasing is plausible image-prompt context to catch bypasses. + const bad = [ + "a cartoon shit emoji", + "a fucking sunset over the lake", + "asshole wearing a tuxedo", + "bitch holding a bouquet", + "dickhead in a bonsai garden", + ]; + for (const prompt of bad) { + expect(validatePrompt(prompt)).toEqual({ ok: false, reason: "moderation" }); + } + }); +}); diff --git a/image-studio/frontend/app/api/backends/route.ts b/image-studio/frontend/app/api/backends/route.ts new file mode 100644 index 000000000..3b3804935 --- /dev/null +++ b/image-studio/frontend/app/api/backends/route.ts @@ -0,0 +1,27 @@ +const BACKEND_URL = process.env.NEXT_PUBLIC_BACKEND_URL ?? "http://127.0.0.1:8000"; + +export async function GET(request: Request) { + const incoming = new URL(request.url); + const upstreamURL = new URL(`${BACKEND_URL}/backends`); + const force = incoming.searchParams.get("force_disable"); + if (force === "1" || force === "true") { + upstreamURL.searchParams.set("force_disable", "1"); + } + + try { + const upstream = await fetch(upstreamURL.toString(), { + method: "GET", + cache: "no-store", + }); + const body = await upstream.text(); + return new Response(body, { + status: upstream.status, + headers: { "Content-Type": upstream.headers.get("content-type") ?? "application/json" }, + }); + } catch { + return Response.json( + { detail: `Could not reach backend at ${BACKEND_URL}.` }, + { status: 502 }, + ); + } +} diff --git a/image-studio/frontend/app/api/generate/compare/route.ts b/image-studio/frontend/app/api/generate/compare/route.ts new file mode 100644 index 000000000..1cadd4c30 --- /dev/null +++ b/image-studio/frontend/app/api/generate/compare/route.ts @@ -0,0 +1,40 @@ +import { MODERATION_REJECT_MESSAGE, validatePrompt } from "@/lib/prompt-moderator"; + +const BACKEND_URL = process.env.NEXT_PUBLIC_BACKEND_URL ?? "http://127.0.0.1:8000"; + +export async function POST(request: Request) { + const payload = await request.text(); + + let parsed: { prompt?: unknown }; + try { + parsed = JSON.parse(payload); + } catch { + return Response.json({ detail: "Invalid JSON body." }, { status: 400 }); + } + if (typeof parsed.prompt === "string") { + const verdict = validatePrompt(parsed.prompt); + if (!verdict.ok && verdict.reason === "moderation") { + return Response.json({ detail: MODERATION_REJECT_MESSAGE }, { status: 400 }); + } + } + + try { + const upstream = await fetch(`${BACKEND_URL}/generate/compare`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: payload, + cache: "no-store", + }); + const contentType = upstream.headers.get("content-type") ?? "application/json"; + const body = await upstream.arrayBuffer(); + return new Response(body, { + status: upstream.status, + headers: { "Content-Type": contentType }, + }); + } catch { + return Response.json( + { detail: `Could not reach backend at ${BACKEND_URL}.` }, + { status: 502 }, + ); + } +} diff --git a/image-studio/frontend/app/api/generate/route.ts b/image-studio/frontend/app/api/generate/route.ts new file mode 100644 index 000000000..87b5eb026 --- /dev/null +++ b/image-studio/frontend/app/api/generate/route.ts @@ -0,0 +1,44 @@ +import { MODERATION_REJECT_MESSAGE, validatePrompt } from "@/lib/prompt-moderator"; + +const BACKEND_URL = process.env.NEXT_PUBLIC_BACKEND_URL ?? "http://127.0.0.1:8000"; + +export async function POST(request: Request) { + const payload = await request.text(); + + // Server-side moderation gate: client may be tampered with, so this is the + // real enforcement. The client also runs validatePrompt for snappier UX. + let parsed: { prompt?: unknown }; + try { + parsed = JSON.parse(payload); + } catch { + return Response.json({ detail: "Invalid JSON body." }, { status: 400 }); + } + if (typeof parsed.prompt === "string") { + const verdict = validatePrompt(parsed.prompt); + if (!verdict.ok && verdict.reason === "moderation") { + return Response.json({ detail: MODERATION_REJECT_MESSAGE }, { status: 400 }); + } + } + + try { + const upstream = await fetch(`${BACKEND_URL}/generate`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: payload, + cache: "no-store", + }); + const contentType = upstream.headers.get("content-type") ?? "application/octet-stream"; + const body = await upstream.arrayBuffer(); + const headers: Record = { "Content-Type": contentType }; + const peak = upstream.headers.get("x-peak-memory-mb"); + const wall = upstream.headers.get("x-wall-seconds"); + if (peak) headers["X-Peak-Memory-MB"] = peak; + if (wall) headers["X-Wall-Seconds"] = wall; + return new Response(body, { status: upstream.status, headers }); + } catch { + return Response.json( + { detail: `Could not reach backend at ${BACKEND_URL}.` }, + { status: 502 }, + ); + } +} diff --git a/image-studio/frontend/app/apple-icon.png b/image-studio/frontend/app/apple-icon.png new file mode 100644 index 000000000..d8d0e21ac Binary files /dev/null and b/image-studio/frontend/app/apple-icon.png differ diff --git a/image-studio/frontend/app/globals.css b/image-studio/frontend/app/globals.css new file mode 100644 index 000000000..ea02fd7c0 --- /dev/null +++ b/image-studio/frontend/app/globals.css @@ -0,0 +1,170 @@ +@import "tailwindcss"; + +@custom-variant dark (&:where(.dark, .dark *)); +@custom-variant light (&:where(.light, .light *)); + +/* ── Light theme (default — Prism ML brand) ─────────────────────────────── */ +:root { + --background: #e9e4df; + --surface: rgba(206, 190, 189, 0.22); + --surface-raised: rgba(243, 239, 234, 0.78); + --surface-strong: rgba(252, 249, 246, 0.88); + --foreground: #1c2030; + --muted: #7a756e; + --muted-strong: #3d3830; + --border: rgba(28, 32, 48, 0.1); + --border-strong: rgba(28, 32, 48, 0.16); + --accent: #1c2030; + --accent-strong: #0e1220; + --accent-soft: rgba(28, 32, 48, 0.07); + --accent-ring: rgba(28, 32, 48, 0.18); + --cta-bg: #1c2030; + --cta-ink: #e9e4df; + --shimmer: rgba(255, 255, 255, 0.2); + --halo-a: rgba(206, 190, 189, 0.55); + --halo-b: rgba(28, 32, 48, 0.04); + --danger: #b91c1c; + --danger-soft: rgba(185, 28, 28, 0.09); + --panel-shadow: 0 20px 56px -38px rgba(28, 32, 48, 0.16), 0 1px 0 rgba(255, 255, 255, 0.82) inset; + --panel-shadow-strong: 0 36px 90px -56px rgba(28, 32, 48, 0.22), 0 1px 0 rgba(255, 255, 255, 0.9) inset; + --ambient-a: rgba(206, 190, 189, 0.45); + --ambient-b: rgba(28, 32, 48, 0.04); + --ambient-c: rgba(28, 32, 48, 0.02); + --grid-line: rgba(28, 32, 48, 0.07); +} + +/* ── Dark theme ─────────────────────────────────────────────────────────── */ +html.dark { + --background: #1a1c24; + --surface: rgba(26, 30, 42, 0.55); + --surface-raised: rgba(32, 36, 50, 0.72); + --surface-strong: rgba(42, 47, 64, 0.84); + --foreground: #e9e4df; + --muted: #9a9590; + --muted-strong: #cfc8be; + --border: rgba(233, 228, 223, 0.1); + --border-strong: rgba(233, 228, 223, 0.16); + --accent: #cebebd; + --accent-strong: #e9e4df; + --accent-soft: rgba(233, 228, 223, 0.1); + --accent-ring: rgba(233, 228, 223, 0.22); + --cta-bg: #e9e4df; + --cta-ink: #1a1c24; + --shimmer: rgba(233, 228, 223, 0.14); + --halo-a: rgba(233, 228, 223, 0.06); + --halo-b: rgba(206, 190, 189, 0.04); + --danger: #ef4444; + --danger-soft: rgba(239, 68, 68, 0.14); + --panel-shadow: 0 30px 72px -54px rgba(0, 0, 0, 0.9), 0 1px 0 rgba(255, 255, 255, 0.05) inset; + --panel-shadow-strong: 0 46px 120px -68px rgba(0, 0, 0, 0.96), 0 1px 0 rgba(255, 255, 255, 0.06) inset; + --ambient-a: rgba(233, 228, 223, 0.08); + --ambient-b: rgba(233, 228, 223, 0.05); + --ambient-c: rgba(233, 228, 223, 0.02); + --grid-line: rgba(233, 228, 223, 0.05); +} + +@theme inline { + --color-background: var(--background); + --color-surface: var(--surface); + --color-surface-raised: var(--surface-raised); + --color-surface-strong: var(--surface-strong); + --color-foreground: var(--foreground); + --color-muted: var(--muted); + --color-muted-strong: var(--muted-strong); + --color-border: var(--border); + --color-border-strong: var(--border-strong); + --color-accent: var(--accent); + --color-accent-strong: var(--accent-strong); + --color-accent-soft: var(--accent-soft); + --color-accent-ring: var(--accent-ring); + --color-cta-bg: var(--cta-bg); + --color-cta-ink: var(--cta-ink); + --color-shimmer: var(--shimmer); + --color-danger: var(--danger); + --color-danger-soft: var(--danger-soft); + --font-sans: var(--font-rethink-sans); + --font-mono: var(--font-plex-mono); +} + +html { + background: var(--background); + color-scheme: light; +} + +html.dark { + color-scheme: dark; +} + +body { + position: relative; + isolation: isolate; + min-height: 100vh; + background: + radial-gradient(circle at 10% 12%, var(--ambient-a), transparent 28%), + radial-gradient(circle at 86% 8%, var(--ambient-b), transparent 24%), + radial-gradient(circle at 52% 100%, var(--ambient-c), transparent 30%), + var(--background); + color: var(--foreground); + font-family: var(--font-sans), ui-sans-serif, system-ui, sans-serif; + overflow-x: hidden; +} + +body::before { + content: ""; + pointer-events: none; + position: fixed; + inset: 0; + z-index: -2; + background: + linear-gradient(90deg, transparent 0, transparent calc(100% - 1px), var(--grid-line) calc(100% - 1px)), + linear-gradient(transparent 0, transparent calc(100% - 1px), var(--grid-line) calc(100% - 1px)); + background-size: 48px 48px; + mask-image: linear-gradient(180deg, rgba(0,0,0,0.5), transparent 80%); + opacity: 0.5; +} + +html.dark body::before { + opacity: 0.45; +} + +body::after { + content: ""; + pointer-events: none; + position: fixed; + inset: 0; + z-index: -1; + background: + radial-gradient(circle at 20% 30%, var(--halo-a), transparent 22%), + radial-gradient(circle at 78% 18%, var(--halo-b), transparent 24%), + radial-gradient(circle at 50% 100%, var(--halo-b), transparent 26%); + filter: blur(16px); + opacity: 0.65; +} + +html.dark body::after { + opacity: 0.7; +} + +::selection { + background: var(--accent-soft); + color: var(--foreground); +} + +@keyframes bonsai-shimmer { + 0% { transform: translateX(-160%); } + 100% { transform: translateX(160%); } +} + +@keyframes bonsai-spin { + to { transform: rotate(360deg); } +} + +@keyframes bonsai-fade { + from { opacity: 0; transform: translateY(6px); } + to { opacity: 1; transform: translateY(0); } +} + +@keyframes bonsai-breathe { + 0%, 100% { opacity: 0.55; transform: scale(1); } + 50% { opacity: 0.95; transform: scale(1.04); } +} diff --git a/image-studio/frontend/app/icon.svg b/image-studio/frontend/app/icon.svg new file mode 100644 index 000000000..d917ed017 --- /dev/null +++ b/image-studio/frontend/app/icon.svg @@ -0,0 +1,157 @@ + + diff --git a/image-studio/frontend/app/layout.tsx b/image-studio/frontend/app/layout.tsx new file mode 100644 index 000000000..3ed72d736 --- /dev/null +++ b/image-studio/frontend/app/layout.tsx @@ -0,0 +1,53 @@ +import type { Metadata } from "next"; +import { IBM_Plex_Mono, Rethink_Sans } from "next/font/google"; +import "./globals.css"; +import { Providers } from "@/components/providers"; + +const rethinkSans = Rethink_Sans({ + variable: "--font-rethink-sans", + subsets: ["latin"], + weight: ["400", "500", "600", "700", "800"], +}); + +const plexMono = IBM_Plex_Mono({ + variable: "--font-plex-mono", + subsets: ["latin"], + weight: ["400", "500"], +}); + +export const metadata: Metadata = { + title: "Bonsai", + description: "Bonsai — an on-device image-generation studio.", + icons: { + icon: "/brand/bonsai-icon-horizontal-dark.svg", + }, + openGraph: { + title: "Bonsai", + description: "Bonsai — an on-device image-generation studio.", + }, + twitter: { + card: "summary", + title: "Bonsai", + description: "Bonsai — an on-device image-generation studio.", + }, +}; + +export default function RootLayout({ + children, +}: Readonly<{ + children: React.ReactNode; +}>) { + return ( + + + + {children} + + + + ); +} diff --git a/image-studio/frontend/app/page.tsx b/image-studio/frontend/app/page.tsx new file mode 100644 index 000000000..ec9fee360 --- /dev/null +++ b/image-studio/frontend/app/page.tsx @@ -0,0 +1,5 @@ +import { StudioClient } from "@/components/studio-client"; + +export default function Home() { + return ; +} diff --git a/image-studio/frontend/components/batch-result-panel.tsx b/image-studio/frontend/components/batch-result-panel.tsx new file mode 100644 index 000000000..4853b3d39 --- /dev/null +++ b/image-studio/frontend/components/batch-result-panel.tsx @@ -0,0 +1,161 @@ +"use client"; + +import { useCallback } from "react"; +import { Download, LoaderCircle, Share2 } from "lucide-react"; +import type { HistoryEntry } from "@/lib/use-history"; +import { resolutionById } from "@/lib/resolutions"; +import { Alert, AlertDescription } from "@/components/ui/alert"; +import { Button } from "@/components/ui/button"; +import { ShareButton } from "@/components/result-panel"; +import { buildMetadata, downloadWithMetadata } from "@/lib/png-metadata"; +import { cn } from "@/lib/utils"; + +interface BatchResultPanelProps { + entries: HistoryEntry[]; + selectedIndex: number | null; + onSelect: (index: number) => void; + progress: { done: number; total: number } | null; + isLoading: boolean; + error: string | null; + // Aspect to use for pending tiles (no entry yet) — matches the resolution + // selected when this batch was kicked off. + pendingAspectRatio: number; +} + +export function BatchResultPanel({ + entries, + selectedIndex, + onSelect, + progress, + isLoading, + error, + pendingAspectRatio, +}: BatchResultPanelProps) { + const selected = selectedIndex !== null ? entries[selectedIndex] ?? null : null; + const downloadName = selected + ? `bonsai-${new Date(selected.timestamp).toISOString().replace(/[:.]/g, "-")}.png` + : "bonsai.png"; + const handleSave = useCallback(async () => { + if (!selected) return; + const res = resolutionById(selected.params.resolutionId); + const meta = buildMetadata(selected, `${res.width}x${res.height}`); + await downloadWithMetadata(selected.imageBlob, downloadName, meta); + }, [selected, downloadName]); + + return ( +
+
+
+ + {error ? ( +
+ + {error} + +
+ ) : ( +
+ {[0, 1, 2, 3].map((i) => { + const entry = entries[i]; + const isSelected = selectedIndex === i; + const showSpinner = !entry && isLoading; + // Each tile sizes to its own image's aspect ratio; pending tiles + // adopt the active generate-time aspect so the grid doesn't + // jump as renders fill in. + const tileAspect = entry + ? (() => { + const r = resolutionById(entry.params.resolutionId); + return r.width / r.height; + })() + : pendingAspectRatio; + return ( + + ); + })} +
+ )} + {progress && progress.done < progress.total ? ( +

+ Rendering {Math.min(progress.done + 1, progress.total)} / {progress.total}… +

+ ) : null} +
+ + {!error ? ( +
+
+

+ {selected?.prompt + || (isLoading + ? "Rendering 4 images — pick one when done to save or share." + : "Pick an image to enable Save and Share.")} +

+
+
+ {selected ? ( + <> + + + + ) : ( + <> + + + + )} +
+
+ ) : null} +
+ ); +} diff --git a/image-studio/frontend/components/bonsai-background.tsx b/image-studio/frontend/components/bonsai-background.tsx new file mode 100644 index 000000000..eb5ca1b65 --- /dev/null +++ b/image-studio/frontend/components/bonsai-background.tsx @@ -0,0 +1,35 @@ +export function BonsaiBackground() { + return ( +
+
+ +
+
+ +
+
+
+ ); +} + +// Silhouette adapted from Logoprojetbonsai.svg (Crazou / Mouagip, Wikimedia +// Commons, CC BY-SA 3.0). The source is a single gradient-filled path, so +// the whole silhouette inherits one theme token instead of our prior two-tone. +function BonsaiSilhouette() { + return ( + + + + + + ); +} diff --git a/image-studio/frontend/components/compare-client.tsx b/image-studio/frontend/components/compare-client.tsx new file mode 100644 index 000000000..0a7662ec0 --- /dev/null +++ b/image-studio/frontend/components/compare-client.tsx @@ -0,0 +1,736 @@ +"use client"; + +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { AlertTriangle, ChevronDown, Dice5, Download, LoaderCircle, Save, Sparkles } from "lucide-react"; +import { + BACKENDS, + MODEL_FAMILIES, + backendFor, + type ModelFamily, +} from "@/lib/backends"; +import { MODERATION_REJECT_MESSAGE, validatePrompt } from "@/lib/prompt-moderator"; +import { useBackends } from "@/lib/use-backends"; +import { COMPARE_PRESETS, DEFAULT_PRESET_ID } from "@/lib/compare-presets"; +import { DEFAULT_RESOLUTION_ID, RESOLUTIONS, resolutionById } from "@/lib/resolutions"; +import { + type CompareEntry, + type CompareSlot, + type ErrorSlot, + type PendingSlot, + type ReadySlot, + useCompareHistory, +} from "@/lib/use-compare-history"; +import { cn } from "@/lib/utils"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Textarea } from "@/components/ui/textarea"; + +const COMPARE_PATH = "/api/generate/compare"; + +interface CompareApiResult { + backend: string; + png_b64: string; + wall_seconds: number; + swap_seconds: number; +} + +interface CompareApiResponse { + results: CompareApiResult[]; +} + +async function parseError(response: Response) { + const text = await response.text().catch(() => ""); + try { + const parsed = JSON.parse(text) as { detail?: string }; + if (typeof parsed.detail === "string" && parsed.detail.length > 0) return parsed.detail; + } catch { + // fall through + } + return text || `Request failed with status ${response.status}.`; +} + +function b64ToBlob(b64: string) { + const binary = atob(b64); + const bytes = new Uint8Array(binary.length); + for (let i = 0; i < binary.length; i++) bytes[i] = binary.charCodeAt(i); + return new Blob([bytes], { type: "image/png" }); +} + +function backendLabel(value: string) { + return BACKENDS.find((b) => b.value === value)?.label ?? value; +} + +function formatSeconds(s: number) { + return `${s.toFixed(2)}s`; +} + +function readToken(name: string, fallback: string): string { + // Canvas can't consume CSS variables directly — read the live value off :root + // so the stitched export picks up whichever theme is active. + if (typeof window === "undefined") return fallback; + const resolved = getComputedStyle(document.documentElement).getPropertyValue(name).trim(); + return resolved || fallback; +} + +async function stitchResults(entry: CompareEntry): Promise { + // Compose ready slots horizontally with a small gutter + caption strip at the + // bottom. Pending/error slots are skipped; caller should gate on ready count. + const readySlots = entry.results.filter((s): s is ReadySlot => s.status === "ready"); + if (readySlots.length === 0) throw new Error("No completed images to stitch yet."); + + const gutter = 16; + const captionHeight = 56; + + const images = await Promise.all( + readySlots.map( + (r) => + new Promise((resolve, reject) => { + const img = new Image(); + img.onload = () => resolve(img); + img.onerror = () => reject(new Error(`Failed to load ${r.backend}`)); + img.src = r.imageUrl; + }), + ), + ); + + const frameWidth = Math.max(...images.map((i) => i.naturalWidth)); + const frameHeight = Math.max(...images.map((i) => i.naturalHeight)); + const n = images.length; + const canvas = document.createElement("canvas"); + canvas.width = frameWidth * n + gutter * (n + 1); + canvas.height = frameHeight + captionHeight + gutter * 2; + const ctx = canvas.getContext("2d"); + if (!ctx) throw new Error("2D canvas context unavailable"); + + const bgColor = readToken("--background", "#101014"); + const mutedStrong = readToken("--muted-strong", "#d4d4d8"); + const muted = readToken("--muted", "#a1a1aa"); + + ctx.fillStyle = bgColor; + ctx.fillRect(0, 0, canvas.width, canvas.height); + + images.forEach((img, i) => { + const x = gutter + i * (frameWidth + gutter); + ctx.drawImage(img, x, gutter, frameWidth, frameHeight); + const result = readySlots[i]; + ctx.fillStyle = mutedStrong; + ctx.font = "14px ui-sans-serif, system-ui, sans-serif"; + ctx.textBaseline = "top"; + ctx.fillText( + backendLabel(result.backend), + x, + gutter + frameHeight + 8, + ); + ctx.fillStyle = muted; + ctx.font = "12px ui-monospace, monospace"; + ctx.fillText( + `wall ${formatSeconds(result.wallSeconds)} · swap ${formatSeconds(result.swapSeconds)}`, + x, + gutter + frameHeight + 30, + ); + }); + + return await new Promise((resolve, reject) => { + canvas.toBlob((blob) => { + if (!blob) reject(new Error("Canvas toBlob returned null")); + else resolve(blob); + }, "image/png"); + }); +} + +export function CompareClient() { + const [presetId, setPresetId] = useState(DEFAULT_PRESET_ID); + const [prompt, setPrompt] = useState( + COMPARE_PRESETS.find((p) => p.id === DEFAULT_PRESET_ID)?.prompt ?? "", + ); + const [seed, setSeed] = useState(42); + const [steps, setSteps] = useState(4); + const [guidance, setGuidance] = useState(1.0); + const [resolutionId, setResolutionId] = useState(DEFAULT_RESOLUTION_ID); + // The relay reports one kind per process; compare picks across families + // within that kind. Switching kinds requires restarting the relay. + const { kind, supportedFamilies, defaultFamily } = useBackends(); + const [selectedFamilies, setSelectedFamilies] = useState([defaultFamily]); + const familyOptions = useMemo( + () => MODEL_FAMILIES.filter((f) => supportedFamilies.includes(f.value)), + [supportedFamilies], + ); + // Sync to the server-resolved family list after /api/backends lands and + // whenever the user's choices drift outside the supported set. + useEffect(() => { + if (!supportedFamilies.length) return; + setSelectedFamilies((prev) => { + const filtered = prev.filter((f) => supportedFamilies.includes(f)); + if (filtered.length === 0) return [defaultFamily]; + return filtered; + }); + }, [supportedFamilies, defaultFamily]); + const effectiveSelectedBackends = useMemo(() => { + // Preserve canonical family order for deterministic slot layout. + return MODEL_FAMILIES + .filter((f) => selectedFamilies.includes(f.value)) + .map((f) => backendFor(kind, f.value)?.value) + .filter((v): v is string => v !== undefined); + }, [selectedFamilies, kind]); + + const [isRunning, setIsRunning] = useState(false); + const [error, setError] = useState(null); + const [activeEntryId, setActiveEntryId] = useState(null); + + const { entries, push, updateEntry, clear } = useCompareHistory(); + const promptRef = useRef(null); + + useEffect(() => { + const el = promptRef.current; + if (!el) return; + el.style.height = "0px"; + el.style.height = `${el.scrollHeight}px`; + }, [prompt]); + + const resolution = useMemo(() => resolutionById(resolutionId), [resolutionId]); + const sizeLabel = `${resolution.width} × ${resolution.height}`; + const activeEntry = useMemo( + () => entries.find((e) => e.id === activeEntryId) ?? entries[0] ?? null, + [entries, activeEntryId], + ); + + const selectPreset = useCallback((id: string) => { + setPresetId(id); + const preset = COMPARE_PRESETS.find((p) => p.id === id); + if (preset) setPrompt(preset.prompt); + }, []); + + const toggleFamily = useCallback((value: ModelFamily) => { + setSelectedFamilies((prev) => { + if (prev.includes(value)) { + if (prev.length === 1) return prev; + return prev.filter((v) => v !== value); + } + // Preserve canonical order from MODEL_FAMILIES. + return MODEL_FAMILIES + .filter((f) => prev.includes(f.value) || f.value === value) + .map((f) => f.value); + }); + }, []); + + const handleRun = useCallback(async () => { + if (isRunning || prompt.trim().length === 0) return; + const verdict = validatePrompt(prompt); + if (!verdict.ok && verdict.reason === "moderation") { + setError(MODERATION_REJECT_MESSAGE); + return; + } + setIsRunning(true); + setError(null); + + // effectiveSelectedBackends is already in canonical family order. + const backends = effectiveSelectedBackends; + + const entry = push({ + prompt, + presetId: presetId && COMPARE_PRESETS.some((p) => p.id === presetId) ? presetId : null, + params: { seed, steps, guidance, resolutionId }, + results: backends.map((backend) => ({ backend, status: "pending" })), + }); + setActiveEntryId(entry.id); + + for (const backend of backends) { + const response = await fetch(COMPARE_PATH, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + prompt, + seed, + steps, + guidance, + height: resolution.height, + width: resolution.width, + backends: [backend], + }), + }).catch(() => null); + + const patchSlot = (slot: CompareSlot): CompareSlot => { + if (slot.backend !== backend) return slot; + if (!response) { + return { backend, status: "error", error: "Could not reach backend." }; + } + return slot; + }; + + if (!response) { + updateEntry(entry.id, (prev) => ({ ...prev, results: prev.results.map(patchSlot) })); + continue; + } + if (!response.ok) { + const message = await parseError(response); + updateEntry(entry.id, (prev) => ({ + ...prev, + results: prev.results.map((s) => + s.backend === backend ? { backend, status: "error", error: message } : s, + ), + })); + continue; + } + + const json = (await response.json()) as CompareApiResponse; + const r = json.results.find((x) => x.backend === backend) ?? json.results[0]; + if (!r) { + updateEntry(entry.id, (prev) => ({ + ...prev, + results: prev.results.map((s) => + s.backend === backend ? { backend, status: "error", error: "Empty response." } : s, + ), + })); + continue; + } + const imageUrl = URL.createObjectURL(b64ToBlob(r.png_b64)); + updateEntry(entry.id, (prev) => ({ + ...prev, + results: prev.results.map((s) => + s.backend === backend + ? { + backend, + status: "ready", + imageUrl, + wallSeconds: r.wall_seconds, + swapSeconds: r.swap_seconds, + } + : s, + ), + })); + } + + setIsRunning(false); + }, [ + guidance, + isRunning, + presetId, + prompt, + push, + resolution.height, + resolution.width, + resolutionId, + seed, + effectiveSelectedBackends, + steps, + updateEntry, + ]); + + const readyCount = activeEntry?.results.filter((s) => s.status === "ready").length ?? 0; + + const handleSaveAll = useCallback(async () => { + if (!activeEntry || readyCount === 0) return; + try { + const blob = await stitchResults(activeEntry); + const url = URL.createObjectURL(blob); + const a = document.createElement("a"); + a.href = url; + a.download = `bonsai-compare-${new Date(activeEntry.timestamp) + .toISOString() + .replace(/[:.]/g, "-")}.png`; + document.body.appendChild(a); + a.click(); + a.remove(); + // Give the download a tick before revoking. + setTimeout(() => URL.revokeObjectURL(url), 1000); + } catch (err) { + setError(err instanceof Error ? err.message : "Failed to stitch results."); + } + }, [activeEntry, readyCount]); + + return ( +
+
+
+
+ + Compare prompt across backends +
+ + internal · power user + +
+ +
+

Preset

+
+ {COMPARE_PRESETS.map((preset) => { + const selected = preset.id === presetId; + return ( + + ); + })} +
+
+ +