From 449488f6b75a500fbb96298db3674d7cb35844a3 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 10 Sep 2026 14:22:49 -0600 Subject: [PATCH] Initial inference engine build out --- .gitignore | 3 +- .../ace_step/ace_step_15_model.py | 1 + .../audio_models/ace_step/src/pipeline.py | 4 + .../diffusion_models/krea2/src/pipeline.py | 10 +- .../inference_engine/InferenceEngine.py | 147 ++++ .../inference_engine/PLANNING.md | 366 +++++++++ .../inference_engine/README.md | 110 +++ .../inference_engine/__init__.py | 18 + .../inference_engine/engine.py | 546 +++++++++++++ .../inference_engine/latent_format_tables.py | 543 +++++++++++++ .../inference_engine/latent_preview.py | 86 ++ .../inference_engine/protocol.py | 96 +++ .../inference_engine/server.py | 152 ++++ requirements_base.txt | 3 + testing/test_inference_engine.py | 384 +++++++++ toolkit/models/base_model.py | 19 + toolkit/models/registry.py | 85 ++ toolkit/models/v2/_mixin.py | 48 ++ toolkit/models/v2/pool.py | 194 +++++ toolkit/sample_step_hook.py | 134 ++++ toolkit/stable_diffusion_model.py | 8 + toolkit/ui_job_status.py | 130 ++++ ui/cron/actions/startJob.ts | 8 + ui/src/app/api/inference/[...path]/route.ts | 125 +++ ui/src/app/generate/page.tsx | 735 ++++++++++++++++++ ui/src/app/jobs/[jobID]/page.tsx | 3 + ui/src/app/jobs/new/options.tsx | 93 +++ ui/src/components/ActiveJobWidget.tsx | 4 + ui/src/components/JobsTable.tsx | 8 + ui/src/components/Sidebar.tsx | 8 +- ui/src/components/generate/GenerateFooter.tsx | 92 +++ ui/src/helpers/inferenceJobConfig.ts | 18 + ui/src/server/inferenceEngine.ts | 76 ++ 33 files changed, 4253 insertions(+), 4 deletions(-) create mode 100644 extensions_built_in/inference_engine/InferenceEngine.py create mode 100644 extensions_built_in/inference_engine/PLANNING.md create mode 100644 extensions_built_in/inference_engine/README.md create mode 100644 extensions_built_in/inference_engine/__init__.py create mode 100644 extensions_built_in/inference_engine/engine.py create mode 100644 extensions_built_in/inference_engine/latent_format_tables.py create mode 100644 extensions_built_in/inference_engine/latent_preview.py create mode 100644 extensions_built_in/inference_engine/protocol.py create mode 100644 extensions_built_in/inference_engine/server.py create mode 100644 testing/test_inference_engine.py create mode 100644 toolkit/models/registry.py create mode 100644 toolkit/models/v2/pool.py create mode 100644 toolkit/sample_step_hook.py create mode 100644 toolkit/ui_job_status.py create mode 100644 ui/src/app/api/inference/[...path]/route.ts create mode 100644 ui/src/app/generate/page.tsx create mode 100644 ui/src/components/generate/GenerateFooter.tsx create mode 100644 ui/src/helpers/inferenceJobConfig.ts create mode 100644 ui/src/server/inferenceEngine.ts diff --git a/.gitignore b/.gitignore index d54f259..6d552f4 100644 --- a/.gitignore +++ b/.gitignore @@ -192,4 +192,5 @@ aitk_db.db-shm .claude original_repo .next -testing/.model_test_outputs \ No newline at end of file +testing/.model_test_outputs +testing/.engine_test_outputs diff --git a/extensions_built_in/audio_models/ace_step/ace_step_15_model.py b/extensions_built_in/audio_models/ace_step/ace_step_15_model.py index e91c444..ad93d47 100644 --- a/extensions_built_in/audio_models/ace_step/ace_step_15_model.py +++ b/extensions_built_in/audio_models/ace_step/ace_step_15_model.py @@ -264,6 +264,7 @@ class AceStep15Model(BaseAudioModel): time_sig=time_sig, language=language, guidance_scale=gen_config.guidance_scale, + step_callback=lambda i, n, latents: self._emit_sample_step(latents, i, n), ) return output diff --git a/extensions_built_in/audio_models/ace_step/src/pipeline.py b/extensions_built_in/audio_models/ace_step/src/pipeline.py index 57edff6..72f94d8 100644 --- a/extensions_built_in/audio_models/ace_step/src/pipeline.py +++ b/extensions_built_in/audio_models/ace_step/src/pipeline.py @@ -87,6 +87,7 @@ class AceStep15Pipeline: time_sig="N/A", language="en", guidance_scale=1.0, + step_callback=None, # step_callback(i, num_steps, latents) after each update ): t_sched = compute_timesteps(num_inference_steps, 3.0) latent_len = int(duration * self.LATENT_RATE) @@ -153,6 +154,9 @@ class AceStep15Pipeline: else: vt = vt_cond + if step_callback is not None: + # x0 estimate at this step (flow matching: x_t - t * v) + step_callback(i, len(t_sched_t), xt - vt * tv) if i == len(t_sched_t) - 1: xt = xt - vt * tv else: diff --git a/extensions_built_in/diffusion_models/krea2/src/pipeline.py b/extensions_built_in/diffusion_models/krea2/src/pipeline.py index 2e901e6..f1d5b57 100644 --- a/extensions_built_in/diffusion_models/krea2/src/pipeline.py +++ b/extensions_built_in/diffusion_models/krea2/src/pipeline.py @@ -357,7 +357,8 @@ class Krea2Pipeline: ref_cache = {"kv": None, "mask": None} # Euler integration of the flow ODE (with optional CFG). - for tcurr, tprev in zip(ts[:-1], ts[1:]): + num_steps = len(ts) - 1 + for i, (tcurr, tprev) in enumerate(zip(ts[:-1], ts[1:])): t = torch.full((latents.shape[0],), tcurr, dtype=dtype, device=device) v_cond = predict_velocity( transformer, @@ -383,7 +384,12 @@ class Krea2Pipeline: v = v_cond + guidance_scale * (v_cond - v_uncond) else: v = v_cond - latents = latents + (tprev - tcurr) * v.to(torch.float32) + v = v.to(torch.float32) + # inference engine preview: x0 estimate for this step + emit = getattr(model, "_emit_sample_step", None) + if emit is not None and getattr(model, "sample_step_hook", None) is not None: + emit(latents - float(tcurr) * v, i, num_steps) + latents = latents + (tprev - tcurr) * v images = model.decode_latents(latents, device=device, dtype=dtype) images = images.float().clamp(-1.0, 1.0) diff --git a/extensions_built_in/inference_engine/InferenceEngine.py b/extensions_built_in/inference_engine/InferenceEngine.py new file mode 100644 index 0000000..b63a715 --- /dev/null +++ b/extensions_built_in/inference_engine/InferenceEngine.py @@ -0,0 +1,147 @@ +"""The inference engine as an AI Toolkit job: a resident generation server +that stays up until the job is stopped. + + job: extension + config: + name: inference_engine + process: + - type: InferenceEngine + sqlite_db_path: ./aitk_db.db + device: cuda + engine: + job_folder: output/inference_engine # engine.json + outputs/ + assets/ + default_model: {arch: zimage} # optional warm load + +The server binds 127.0.0.1 on an ephemeral port and publishes it (with a +per-launch token) to /engine.json; the UI proxies to it. +""" + +import json +import os +import secrets +import signal +import time +from collections import OrderedDict + +import torch + +from jobs.process import BaseExtensionProcess +from toolkit.paths import TOOLKIT_ROOT +from toolkit.ui_job_status import UIJobStatus + + +class InferenceEngineConfig: + def __init__(self, **kwargs): + self.job_folder = kwargs.get("job_folder", None) + self.output_folder = kwargs.get("output_folder", None) + self.default_model = kwargs.get("default_model", None) + self.dtype = kwargs.get("dtype", "bf16") + self.host = kwargs.get("host", "127.0.0.1") + # 0 = ephemeral (recommended); a fixed port is for external clients + self.port = int(kwargs.get("port", 0)) + self.token = kwargs.get("token", None) + + +class InferenceEngine(BaseExtensionProcess): + def __init__(self, process_id: int, job, config: OrderedDict, **kwargs): + super().__init__(process_id, job, config, **kwargs) + self.ui = UIJobStatus(self.config.get("sqlite_db_path", "./aitk_db.db")) + self.engine_config = InferenceEngineConfig(**self.get_conf("engine", {})) + self.device = self.get_conf("device", "cuda") + job_folder = self.engine_config.job_folder or os.path.join(TOOLKIT_ROOT, "output", self.name) + self.job_folder = os.path.abspath(job_folder) + self.output_folder = os.path.abspath( + self.engine_config.output_folder or os.path.join(self.job_folder, "outputs") + ) + self.endpoint_path = os.path.join(self.job_folder, "engine.json") + self.engine = None + self.server = None + self.server_thread = None + self.endpoint = None + + def run(self): + super().run() + from .engine import Engine + from .server import create_app + + self.ui.update_status("running", "Starting server") + os.makedirs(self.job_folder, exist_ok=True) + self.engine = Engine( + device=self.device, + output_folder=self.output_folder, + assets_folder=os.path.join(self.job_folder, "assets"), + status_cb=self.ui.update_info, + dtype=self.engine_config.dtype, + ) + token = self.engine_config.token or secrets.token_urlsafe(24) + self.start_server(create_app(self.engine, token), token) + + self.ui.start_stop_watcher(lambda reason: self.engine.stop(reason)) + for sig in (signal.SIGINT, signal.SIGTERM): + try: + signal.signal(sig, lambda *_: self.engine.stop("stopped")) + except Exception: + pass + + if self.engine_config.default_model: + with torch.no_grad(): + self.engine.preload(dict(self.engine_config.default_model)) + + self.engine._idle_status() + try: + with torch.no_grad(): + self.engine.run_forever() + finally: + self.shutdown() + + def start_server(self, app, token: str): + from .server import serve_in_thread + + self.server, self.server_thread, host, port = serve_in_thread( + app, self.engine_config.host, self.engine_config.port + ) + self.endpoint = { + "host": host, + "port": port, + "token": token, + "pid": os.getpid(), + "job_id": self.ui.job_id, + "device": self.device, + "job_folder": self.job_folder, + "output_folder": self.output_folder, + "started_at": time.time(), + } + tmp = self.endpoint_path + ".tmp" + with open(tmp, "w") as f: + json.dump(self.endpoint, f, indent=2) + os.replace(tmp, self.endpoint_path) + print(f"[AITK] Inference engine listening on http://{host}:{port} (endpoint file: {self.endpoint_path})") + + def shutdown(self): + reason = self.engine.stop_reason if self.engine else "stopped" + print(f"[AITK] Inference engine shutting down ({reason})") + if self.server is not None: + self.server.should_exit = True + if self.server_thread is not None: + self.server_thread.join(timeout=10) + try: + os.remove(self.endpoint_path) + except FileNotFoundError: + pass + if self.engine is not None: + try: + self.engine.unload() + except Exception: + pass + self.ui.update_status("queued" if reason == "queued" else "stopped", "Engine stopped") + + def on_error(self, e: Exception): + super().on_error(e) + if isinstance(e, KeyboardInterrupt): + self.ui.update_status("stopped", "Engine stopped") + else: + self.ui.update_status("error", str(e)) + try: + os.remove(self.endpoint_path) + except Exception: + pass diff --git a/extensions_built_in/inference_engine/PLANNING.md b/extensions_built_in/inference_engine/PLANNING.md new file mode 100644 index 0000000..0622b2b --- /dev/null +++ b/extensions_built_in/inference_engine/PLANNING.md @@ -0,0 +1,366 @@ +# Inference Engine — Planning + +## Goal + +A resident generation server that runs as a normal AI Toolkit job (like a +captioner job: `job: extension`, `process[0].type: InferenceEngine`), stays up +until stopped from the UI, and serves generation requests for every arch the +toolkit knows (image, video, audio). Each request names its model config; the +engine loads/reloads only the components that differ from what is resident, +drops what the request does not need, and streams progress + per-step latents +back over the request's connection, ending with the finished file. + +This is Phase 3 of `toolkit/models/v2/PLANNING.md` ("live server with model +hot-swap"); everything below builds on the unified v2 load API +(`ModelClass.load(...)`, `aitk_post_load`, `component_load_kwargs`) that is +already in place for every non-legacy arch. + +## Survey (2026-09-10) + +What exists and is reused as-is: + +- **Job plumbing**: `jobs/ExtensionJob.py` -> `get_all_extensions_process_dict` + -> our process class. The UI launches jobs via `ui/cron/actions/startJob.ts` + (`python run.py ` with `AITK_JOB_ID`, `CUDA_VISIBLE_DEVICES`, + detached, stdout to the job log). Stop = `Job.stop` flag in sqlite, which + `BaseCaptioner` polls in a daemon thread and turns into a SIGINT. The DB + status/step/info writers (`_update_status`, `_update_key`, retry-on-lock, + stop watcher) live in `extensions_built_in/captioner/BaseCaptioner.py`. +- **Generation path**: `BaseModel.generate_images([GenerateImageConfig])` does + prompt encoding, control images, LoRA/network handling, then + `generate_single_image` (per arch), then `gen_config.save_image_atomic` + (png/jpg/webp image, animated webp video, wav/mp3 audio; `toolkit/audio/ + make_video.py` has the mp4 writer). `testing/test_model_loading.py` proves + this flow for 34 archs and holds realistic per-arch sample settings. +- **Per-step hook point**: every arch samples through a diffusers-style + scheduler object owned by the generation pipeline (`pipeline.scheduler.step` + — diffusers pipelines, the vendored wan22/ltx/minimax/ace pipelines, the + legacy monolith). Wrapping that one method yields (step, timestep, latents) + for all of them without touching per-arch code. +- **Component loading**: `OstrisModelMixin.load` / `load_model` is the single + funnel every v2 component goes through (DiTs, TEs, VAEs, vision encoders, + connectors, vocoders). Legacy monolith archs adopt their pipeline + components into v2 wrappers at one post-load funnel (`adopt_component`). +- **UI front**: `ui/cron/fileServer.ts` owns the public port and already + spawns Next.js on an ephemeral loopback port and proxies to it (incl. + websocket upgrades). Next.js route handlers stream `Response` bodies (the + `/api/monitor` SSE route is the pattern). Bearer auth is a middleware on + `/api/*`. `/api/files/*` serves job output folders. The job detail page + already has a per-job `plugin.html` iframe tab. +- **Latent previews**: ComfyUI's `comfy/latent_formats.py` (read-only + reference) has per-family linear `latent_rgb_factors` tables for SD1/SDXL/ + SD3/Flux/Flux2/Wan21/Wan22/LTXV/LTXAV/MiniMaxH3/Cosmos/HunyuanVideo. +- **HTTP server libs**: `fastapi` 0.140 and `uvicorn` 0.51 are importable in + this env but only as transitive deps of gradio; they are NOT pinned in + `requirements.txt`. + +## Decisions (proposed — confirm before build) + +1. **Transport = one chunked HTTP response per generation request** (not + WebSocket). A `POST /generate` stays open for the whole generation and + streams binary frames; the Next.js proxy route passes the body through as + a `ReadableStream`, so it works with or without the fileServer front, in + docker, on Windows/macOS, and from `curl`. WebSocket would require the + fileServer to be mandatory (Next route handlers cannot upgrade). Cancel is + a separate `POST /cancel/{request_id}` (also fired when the client aborts + the stream). +2. **Frame format**: length-prefixed binary frames + `u32 header_len | json header | u64 payload_len | payload`. Header + `{type, request_id, ...}`; types: `start` (resolved config, output kind, + latent layout + rgb factors), `status` (load/quantize/encode stages, + from `print_and_status_update`), `progress` (step/total/eta), `latent` + (raw fp16 tensor bytes + shape), `result` (file path/url + media meta), + `error`, `end`. Text frames have empty payload. +3. **Latents stream raw, projection happens client-side.** The `start` frame + carries `latent_layout` (`BCHW` / `BCFHW` / `audio`), the arch's + `latent_rgb_factors` (+bias) and any scaling, so the browser projects + latents to a canvas with a channels x 3 matmul per pixel and needs no + per-arch code. Server options `stream.latents = "raw" | "none"`, + `stream.every_n_steps`, and for video `stream.max_frames` / + `stream.downsample` to cap frame size (a 30-step video run at full latent + res is hundreds of MB; images are ~0.5 MB/step). Archs without known + factors still stream latents; the client falls back to a per-channel + grayscale/PCA preview. A server-side `preview="jpeg"` mode (full or + TAE decode) is a later add, not v1. +4. **Endpoint discovery = ephemeral loopback port + a file in the job + folder.** Engine binds `127.0.0.1:0`, writes + `//engine.json` (`host, port, token, pid, + started_at`) and logs it. No fixed port, nothing network-exposed, no + collisions, works identically in docker (same container as the UI) and + on Windows (uvicorn has no named-pipe support, so a loopback TCP port is + the portable stand-in for "no port"). No DB schema change (a new Job + column would crash the jobs list for anyone who has not run `update_db`). + The UI already resolves `/` for `plugin.html`. + Requests carry the per-launch `token` (proxy adds it) so other local + processes cannot drive the engine. +5. **One engine process per GPU set, one generation at a time, FIFO queue** + (`GET /queue`, `POST /cancel`). Multiple engine jobs on different GPUs are + independent; the UI picks by job. +6. **Hot-swap via a component pool hooked into the mixin funnel** (no + rewrites of the 35 holder `load_model`s). See "Component pool". +7. **LoRA in v1 = part of the transformer's pool identity** (`lora_path`, + `assistant_lora_path`, network multiplier apply through the existing + holder paths; a different LoRA means the transformer entry reloads; TEs/ + VAEs stay shared). Hot-swappable unmerged LoRA on a resident transformer + is v2. +8. **Server = FastAPI + uvicorn**, run on a thread inside the job process + (generation stays on the main thread). Requires adding `fastapi` and + `uvicorn` to `requirements.txt` (Jaret's call; fallback is a stdlib + `ThreadingHTTPServer` with chunked responses — strictly worse, but zero + deps). +9. **Model registry served by the engine.** `GET /models` returns every + registered arch with capabilities + defaults (modality, control-image + requirement, size divisibility, native res, default steps/CFG/frames/fps, + latent layout, rgb factors). The table lives in python + (`toolkit/models/registry.py`, seeded from `MODEL_TESTS` in + `testing/test_model_loading.py` + per-holder class attrs) so the UI reads + one source of truth at runtime instead of a second copy in `options.tsx`. + +## Architecture + +``` +extensions_built_in/inference_engine/ + __init__.py # InferenceEngineExtension (uid "InferenceEngine") + InferenceEngine.py # process: UI status/stop plumbing, server lifecycle + server.py # FastAPI app: /health /models /generate /cancel /queue /unload + engine.py # Engine: request queue, holder cache, pool orchestration + protocol.py # frame encoder (python) — mirrored by ui/src/lib/engineStream.ts + latent_preview.py # latent layout + rgb factor tables per arch family +toolkit/models/v2/pool.py # ComponentPool (identity, residency, eviction) +toolkit/models/registry.py # arch capabilities/defaults (shared with tests + UI) +toolkit/ui_job_status.py # DB status/stop helpers lifted from BaseCaptioner + # (engine uses it; captioner can adopt later, optional) +testing/test_inference_engine.py +ui/src/app/api/inference/[...path]/route.ts # proxy to the engine (streams) +ui/src/app/generate/page.tsx # Generate page +ui/src/components/generate/* # model picker, prompt form, preview canvas, gallery +ui/src/lib/engineStream.ts # frame parser + latent->canvas projection +``` + +### Process (`InferenceEngine.py`) + +- Config: + ```yaml + job: extension + config: + name: inference_engine + process: + - type: InferenceEngine + sqlite_db_path: ./aitk_db.db + device: cuda + engine: + output_folder: //outputs # results, served by /api/files + keep_resident_gb: 0 # 0 = keep only what the last request needs + idle_unload_seconds: 0 # 0 = never auto-unload + default_model: {arch: zimage, name_or_path: Tongyi-MAI/Z-Image-Turbo} # optional warm load + ``` +- `run()`: status "running / Starting server" -> start uvicorn on a thread + (port 0) -> write `engine.json` -> optional warm load -> main thread loops + on the request queue, executing generations; `info` reflects + "Idle — zimage (Tongyi-MAI/Z-Image-Turbo)" / "Generating 12/30" / "Loading + qwen_image: Quantizing transformer". Stop flag -> SIGINT (captioner + pattern) -> cancel in-flight request, shut the server, delete + `engine.json`, free the pool, status "stopped". +- Also runnable headless: `python run.py config/inference_engine.yaml` + prints the endpoint; no UI needed (useful for tests). + +### Engine (`engine.py`) + +- `GenerateRequest`: `{model: ModelConfig kwargs (arch, name_or_path, + extras_name_or_path, quantize, qtype, quantize_te, qtype_te, low_vram, + layer_offloading..., lora_path, dtype), sample: GenerateImageConfig kwargs + (prompt, negative_prompt, width, height, num_inference_steps, + guidance_scale, seed, num_frames, fps, ctrl_img[_1..3] as server path or + uploaded asset id, output_ext), stream: {latents, every_n_steps, + max_frames}}`. Uploads: `POST /assets` stores control images under the job + folder and returns an id/path. +- Holder handling: `get_model_class(model_config)`, build the holder with + the pool active, `load_model()`, then `pool.release_unused()`. The holder + is kept as "active" and reused verbatim for the next request when its + ModelConfig is identical (fast path: no rebuild). A new ModelConfig + rebuilds the holder; the pool makes the rebuild cost = only the changed + components. +- Sampling: `sd.generate_images([gen])` with the step hook installed; + `add_status_update_hook` feeds `status` frames; the `_after_sample_img` + hook plus the saved file feed `result`. Output goes to + `engine.output_folder` with a request-id filename; the frame carries the + path and the `/api/files` url. Cancel = flag checked in the step hook, + raising a `GenerationCancelled` that `generate_images` cleanup tolerates + (it already restores device/rng state in its epilogue; verify the + exception path leaves the holder usable). + +### Step hook (`BaseModel`) + +- `BaseModel.sample_step_hook: Optional[Callable]` and, inside + `generate_images`, a context manager that wraps `pipeline.scheduler.step` + to call `hook(step_index, num_steps, timestep, latents)` with the + scheduler's returned sample. Zero cost when no hook is set. Custom loops + that bypass the scheduler (found by the per-arch streaming test) get an + explicit `self._emit_sample_step(...)` call. Legacy monolith: same wrap on + its pipeline scheduler. +- Latent extraction from the DTO/carrier: video latents `BCFHW`, ltx/minimax + AV extras (audio rows) ride the DTO — the `latent` frame sends the main + tensor; audio extras optional. + +### Component pool (`toolkit/models/v2/pool.py`) + +- Identity key = `(component class, resolved source [repo/subfolder or file + path after comfy resolution], dtype, qtype (incl. ARA suffix), offload + fraction, final device, lora identity for transformers)`. Computed inside + `OstrisModelMixin.load` from exactly the args it already receives, so a + holder's `component_load_kwargs(role)` output determines sharing (a TE + loaded quantized convrot8 on cuda:0 is one entry; the same TE bf16 on cpu + is another). +- Wiring: a process-global `ComponentPool.current` context; `load` checks + it first (hit = return resident module, skip source+post_load), registers + on miss. Legacy `adopt_component` funnel registers too (consume side stays + full reload for legacy archs). Tokenizers/processors pooled by repo. +- Residency ops: `mark_used(key)` per request; `release_unused(keep)` = + MemoryManager.detach/free, drop refs, `gc.collect()`, + `torch.cuda.empty_cache()`; `bytes(key)` for accounting; + `keep_resident_gb` LRU keep-warm; eviction before a load when + `torch.cuda.mem_get_info` says the incoming component will not fit. +- Safety: entries must not be mutated by holders in ways that change their + identity. Audit as part of step 3: holders that merge LoRAs into weights, + `.to(device)` moves under low_vram (device is in the key — a moved module + must be moved back or re-keyed on release), `MemoryManager.attach` (offload + is in the key), FakeTextEncoder swaps (`toolkit/unloader.py`), per-arch + patches applied in `load_model` after `load` (e.g. vision-tower drops: + fine, they are class-level config and thus part of the class identity). +- Counters exposed on `/health`: hits/misses/evictions, resident entries with + bytes and last-used — the tests assert on these. + +### Server (`server.py`) + +- `GET /health` `{ok, busy, queue, active: {arch, name_or_path, config}, + pool: [...], vram: {used, total}}`; `GET /models`; `POST /generate` + (stream); `POST /cancel/{id}`; `GET /queue`; `POST /unload` + (`{all: true}` or component keys); `POST /assets` (multipart upload); + `GET /outputs` (list results). Token check on everything except `/health`. + +### UI + +- Proxy: `ui/src/app/api/inference/[...path]/route.ts` — finds the running + `job_type = "inference"` job (query param `job` to pick one when several + GPUs run engines), reads its `engine.json` (cached a few seconds), + forwards method/headers/body with the engine token, returns + `new Response(upstream.body)` so chunked frames flow through untouched. + `export const dynamic = 'force-dynamic'`, no body size limit issues + (control images go through `/assets`, multipart). +- `/generate` page (sidebar "Generate"): engine strip (pick/start/stop the + engine job on a GPU — creates a job via `/api/jobs` with + `job_type: 'inference'` and the config above, then start + startQueue, + exactly like `CaptionDatasetModal.saveJob`; shows health), model panel + (arch select from `/models`, name_or_path creatable select with the arch's + defaults, quantize/qtype/low_vram/LoRA), prompt panel (prompt/negative, + size with divisibility snapping, steps/CFG/seed/frames/fps, control image + upload), preview panel (canvas fed by the stream, step/eta, cancel), + results gallery (files from the engine output folder via `/api/files`, + click to open; video/audio players). +- Jobs list / action bar: `job_type === 'inference'` rows show "Inference + Engine" with a link to `/generate`; stop only (no edit/restart-from-step + UI). Dashboard active-job widget same treatment as caption. +- `ui/src/lib/engineStream.ts`: async frame reader over `fetch` + + `ReadableStream`, event emitter, `projectLatentToImageData(latent, layout, + rgbFactors)`; video: shows a chosen frame (scrubber) — audio: heatmap. + +### Tests (`testing/test_inference_engine.py`) + +Mirrors `test_model_loading.py` conventions (`--arch`, `--all` = subprocess +per arch, `--allow-download`, HF offline default, SKIP classification, +outputs under `testing/.engine_test_outputs/`), reusing its +`MODEL_TESTS` entries for per-arch model/sample settings: + +1. `--arch X`: start the engine in-process on an ephemeral port, POST + `/generate`, consume the stream, assert `start` -> >=1 `progress` -> + >=1 `latent` (shape/dtype/layout sanity, finite values) -> `result` + (file exists, > 1 KB, right extension for the modality) -> `end`; + record timings + VRAM like the loading test. +2. `--swap A,B[,C]`: sequential requests; assert pool hits for shared + components (e.g. zimage -> zimage_l2p shares Qwen3 TE + KLVAE; wan21 -> + wan22_5b shares UMT5; qwen_image -> qwen_image_edit shares TE + VAE), + assert VRAM after B does not include A's DiT, assert a second A request + is a holder fast path (no loads). +3. `--cancel`: cancel mid-run, assert the stream ends with `error/cancelled` + and the next request on the same holder succeeds. +4. `--all`: (1) for every registry arch, then a summary + `report.md`. + +## Steps + +Each step lands runnable + tested before the next. + +- [x] **1. Skeleton job + server** (2026-09-10) — `InferenceEngineExtension` + (uid `InferenceEngine`), process with the sqlite status/stop plumbing + lifted into `toolkit/ui_job_status.py`, FastAPI + uvicorn on + `127.0.0.1:0` (`server.serve_in_thread`), `engine.json` in the job + folder, `/health` `/models` `/generate` `/cancel` `/queue` `/unload` + `/assets` `/outputs`, headless yaml. `fastapi`/`uvicorn` pinned in + requirements_base.txt. +- [x] **2. Streaming** — binary frame protocol (`protocol.py`, mirrored by + `ui/src/lib/engineStream.ts`), `BaseModel.sample_step_hook` via a + scheduler.step wrap in `generate_images` (+ `_emit_sample_step` for + hand-rolled loops; ace_step wired), latent layouts BCHW/BCFHW/BLC with + flux-style packed-token unpacking and video frame subsampling, per-arch + preview tables (`latent_preview.py`, numbers from ComfyUI's + latent_formats), cancel (step hook raises; holder verified reusable: + the rerun after a cancel reproduced the same file byte-for-byte). + Verified: zimage — 8 progress + 8 latent frames (1x16x64x64, Flux + preview) + result; `--cancel` PASS. +- [x] **3. Component pool + hot-swap** — `toolkit/models/v2/pool.py`; + `OstrisModelMixin.load_model` consults/registers the active pool by + source key (class, name_or_path, dtype, config, subfolder, comfy flag, + kwargs); `aitk_post_load` is idempotent for an identical policy and + detaches offloading before re-applying a different one. Engine: drop + holder → begin_request → build → release_unused; CUDA OOM during a + build frees untouched entries and retries once. Verified: + zimage→zimage_l2p→zimage PASS (TE shared on the way back; l2p's + transformer/VAE bypass the mixin so they reload — see step 4); + wan21→wan22_5b: UMT5 TE shared (pool hit), wan21 DiT+VAE released. +- [ ] **4. Coverage** — engine test across the registry. PASS so far + (2026-09-10, 5090 with ~16 GB free next to a training job, convrot8): + zimage (image, Flux preview), wan21 + wan22_5b (video BCFHW, animated + webp), ace_step_15 (audio BLC, mp3), sd1 (legacy monolith: the hook + is installed in its own generate_images via toolkit/sample_step_hook), + zimage `--cancel`. Known gaps: holders whose component loads + bypass `load_model` (zimage_l2p's from_pretrained + pixel conversion, + legacy monolith adopt path) do not pool; omnigen2 builds its scheduler + inside get_generation_pipeline (hook still applies — wrapping happens + after); custom loops without scheduler.step need `_emit_sample_step` + as they are found. Full `--all` sweep not run (shared GPU); run it + when a free card is available. +- [x] **5. UI** — `/api/inference/[...path]` streaming proxy + (`ui/src/server/inferenceEngine.ts` resolves the running inference + job's engine.json, checks the pid, injects the token), `/generate` + page (engine start/stop per GPU, model + prompt forms fed by + `/models`, live latent canvas with frame scrubber, results gallery via + `/api/files`), sidebar entry, inference rows in the queue table / + dashboard widget / job page, `startJob.ts` injects + `engine.job_folder`. Type-checked and built in an isolated copy of the + UI (the live `next start` on this box was not rebuilt — run + `npm run build` to pick it up, the worker change needs it too). + Proxy verified end to end against that isolated build with a copy of + the db: status → health → models → streamed generate (first frame + arrives immediately, progress frames live, result at the end) → stop + via the Job.stop flag (engine exits 0, row 'stopped', engine.json + removed, status reports no engine). +- [ ] **6. Hardening** — done: FIFO queue + `/queue`, OOM-retry eviction, + stale engine.json ignored when its pid is dead, `wait=1` JSON mode. + Open: idle unload, `keep_resident_gb` LRU keep-warm, pre-load VRAM-fit + eviction (today untouched entries are only freed after the build or + on OOM, so a tight card pays a reload instead of sharing when the + DiT loads first), Windows relay launch check, docker compose smoke, + pooling for the custom-path holders. + +## Open questions + +1. Chunked HTTP stream vs WebSocket (decision 1). +2. Raw latents + client-side projection vs server-rendered previews + (decision 3). Both can coexist later; which is v1? +3. `engine.json` in the job folder vs a new Job DB column (decision 4). +4. OK to add `fastapi` + `uvicorn` to `requirements.txt`? +5. LoRA as transformer identity in v1 (decision 7) — acceptable? +6. Should the per-arch registry (`toolkit/models/registry.py`) also become + the source for `ui/src/app/jobs/new/options.tsx` in the same work, or + stay engine-only for now? +7. Does the engine also need a non-streamed `POST /generate?wait=1` JSON + mode for external scripts (cheap to add alongside)? diff --git a/extensions_built_in/inference_engine/README.md b/extensions_built_in/inference_engine/README.md new file mode 100644 index 0000000..d2716ba --- /dev/null +++ b/extensions_built_in/inference_engine/README.md @@ -0,0 +1,110 @@ +# Inference Engine + +A resident generation server that runs as an AI Toolkit job and serves every +registered model arch (image, video, audio). It hot-swaps models per request: +components the new model shares with the previous one (a text encoder, a +VAE) stay resident, the rest is freed and loaded. Progress and the raw +per-step latents stream back over the request's HTTP connection. + +Design notes and status: `PLANNING.md` in this folder. + +## Running + +From the UI: **Generate** in the sidebar, pick a GPU, **Start engine**. The +engine shows up in the Queue as an `inference` job and runs until stopped. + +Headless: + +```yaml +# inference_engine.yaml +job: extension +config: + name: inference_engine + process: + - type: InferenceEngine + device: cuda + engine: + job_folder: output/inference_engine # engine.json + outputs/ + assets/ + # port: 8866 # fixed port for external clients (default: ephemeral) + # token: secret # request token (default: random per launch) + # default_model: {arch: zimage} # warm load +``` + +``` +python run.py inference_engine.yaml +``` + +The server binds `127.0.0.1` on an ephemeral port and writes +`/engine.json` (`host`, `port`, `token`, `pid`). The UI proxies +`/api/inference/*` to it; nothing else is exposed. + +## HTTP API + +All routes except `/health` need the token (`X-Engine-Token: ` or +`Authorization: Bearer `). + +| route | | +|---|---| +| `GET /health` | busy/queue, active model, VRAM, pool residents + hit/miss stats | +| `GET /models` | every arch with modality, defaults, control-image requirement (`toolkit/models/registry.py`) | +| `POST /generate` | body `{model, sample, stream}` → binary frame stream (below). `{"wait": true}` returns JSON when done instead | +| `POST /cancel/{request_id}` | cancel a queued/running request | +| `GET /queue` | queued + recent requests | +| `POST /unload` | drop the holder and every pooled component | +| `POST /assets?name=x.png` | raw-body upload; returns `{path}` usable as `sample.ctrl_img` | +| `GET /outputs/{relpath}` | fetch a produced file | + +`model` = `ModelConfig` kwargs (`arch` required; `name_or_path`, `quantize`, +`qtype`, `quantize_te`, `qtype_te`, `low_vram`, `layer_offloading`, +`lora_path`, `extras_name_or_path`, ...). Registry defaults fill in what is +omitted. `sample` = `GenerateImageConfig` kwargs (`prompt`, `negative_prompt`, +`width`, `height`, `num_inference_steps`, `guidance_scale`, `seed`, +`num_frames`, `fps`, `ctrl_img`, `output_ext`). `stream` = +`{latents: "raw"|"none", every_n_steps, max_frames}`. + +### Frame stream + +`u32 header_len | json header | u64 payload_len | payload` (little endian). +Header always has `type` and `request_id`: + +- `start` — resolved model/sample, `modality`, `preview` (latent→RGB factors) +- `status` — holder stage lines (loading, quantizing, ...) +- `progress` — `{step, total, elapsed}` +- `latent` — raw fp16 tensor payload; `shape`, `dtype`, `layout` (`BCHW`, + `BCFHW`, `BLC`). Video latents are subsampled to `max_frames` frames. +- `result` — `{path, kind, ext, seed, width, height, ...}` per produced file +- `error` — `{message, cancelled}` +- `end` — `{status}` + +Python reader: `protocol.FrameReader`; browser reader: +`ui/src/lib/engineStream.ts` (also does the latent→RGB projection from the +`preview` block, so the client needs no per-arch code). + +### Per-step latents + +`BaseModel.generate_images` wraps `scheduler.step` on the sampling scheduler +while `sample_step_hook` is set, so every diffusers-style pipeline reports +its latents without per-arch code. Holders with a hand-rolled loop call +`self._emit_sample_step(latents, i, n)` (ace_step does). + +## Component pool + +`toolkit/models/v2/pool.py`. While a pool is active, `OstrisModelMixin. +load_model` returns the resident module for an identical source request and +`aitk_post_load` is a no-op for a module already carrying that policy. After +each holder build the engine frees every pooled component the new model did +not touch. On CUDA OOM during a load it frees the untouched components and +retries once. Components loaded outside the mixin's `load_model` (custom +holder paths, legacy monolith archs) are not pooled and reload each time. + +## Tests + +``` +python testing/test_inference_engine.py --arch zimage --quantize +python testing/test_inference_engine.py --arch zimage --cancel +python testing/test_inference_engine.py --swap zimage,zimage_l2p,zimage --quantize +python testing/test_inference_engine.py --all +``` + +`--quantize` forces convrot8 on transformer + TE for small GPUs; +`--allow-download` permits hub downloads (default offline → SKIP). diff --git a/extensions_built_in/inference_engine/__init__.py b/extensions_built_in/inference_engine/__init__.py new file mode 100644 index 0000000..8860e3a --- /dev/null +++ b/extensions_built_in/inference_engine/__init__.py @@ -0,0 +1,18 @@ +from toolkit.extension import Extension + + +class InferenceEngineExtension(Extension): + uid = "InferenceEngine" + name = "Inference Engine" + + @classmethod + def get_process(cls): + # import your process class here so it is only loaded when needed and return it + from .InferenceEngine import InferenceEngine + + return InferenceEngine + + +AI_TOOLKIT_EXTENSIONS = [ + InferenceEngineExtension, +] diff --git a/extensions_built_in/inference_engine/engine.py b/extensions_built_in/inference_engine/engine.py new file mode 100644 index 0000000..da3df21 --- /dev/null +++ b/extensions_built_in/inference_engine/engine.py @@ -0,0 +1,546 @@ +"""Engine: request queue + holder lifecycle + generation with streamed frames. + +Runs on the process main thread (`run_forever`); the HTTP server hands it +GenerationJobs from its own thread and reads the frames each job emits. +""" + +import gc +import glob +import json +import os +import queue +import shutil +import threading +import time +import traceback +import uuid +from typing import Callable, Dict, List, Optional + +import torch + +from toolkit.basic import flush +from toolkit.config_modules import GenerateImageConfig, ModelConfig +from toolkit.models.registry import OUTPUT_EXT, get_arch_entry +from toolkit.models.v2.pool import ComponentPool +from toolkit.util.get_model import LEGACY_ARCHS, get_model_class + +from .latent_preview import preview_info +from .protocol import encode_frame, tensor_payload + +END = None # frames-queue terminator (after the "end" frame) + + +class GenerationCancelled(Exception): + pass + + +class GenerationJob: + def __init__(self, model: dict, sample: dict, stream: Optional[dict] = None): + self.request_id = uuid.uuid4().hex[:12] + self.model = model + self.sample = sample + self.stream = {"latents": "raw", "every_n_steps": 1, "max_frames": 4} + self.stream.update(stream or {}) + self.frames: "queue.Queue[Optional[bytes]]" = queue.Queue() + self.cancel = threading.Event() + self.status = "queued" + self.created_at = time.time() + self.started_at: Optional[float] = None + self.finished_at: Optional[float] = None + self.results: List[dict] = [] + self.error: Optional[str] = None + self.step = 0 + self.total_steps: Optional[int] = None + + def emit(self, type_: str, payload: bytes = b"", **fields): + header = {"type": type_, "request_id": self.request_id, **fields} + self.frames.put(encode_frame(header, payload)) + + def finish(self): + self.finished_at = time.time() + self.emit("end", status=self.status) + self.frames.put(END) + + def info(self) -> dict: + return { + "request_id": self.request_id, + "status": self.status, + "arch": self.model.get("arch"), + "prompt": self.sample.get("prompt", ""), + "step": self.step, + "total_steps": self.total_steps, + "created_at": self.created_at, + "started_at": self.started_at, + "finished_at": self.finished_at, + "results": self.results, + "error": self.error, + } + + +def _model_key(model: dict) -> str: + return json.dumps({k: v for k, v in sorted(model.items()) if v is not None}, sort_keys=True) + + +class Engine: + def __init__( + self, + device: str, + output_folder: str, + assets_folder: Optional[str] = None, + status_cb: Optional[Callable[[str], None]] = None, + dtype: str = "bf16", + ): + self.device = device + self.dtype = dtype + self.output_folder = output_folder + self.assets_folder = assets_folder or os.path.join(os.path.dirname(output_folder), "assets") + os.makedirs(self.output_folder, exist_ok=True) + os.makedirs(self.assets_folder, exist_ok=True) + self.status_cb = status_cb or (lambda s: None) + self.queue: "queue.Queue[GenerationJob]" = queue.Queue() + self.jobs: Dict[str, GenerationJob] = {} + self.current: Optional[GenerationJob] = None + self.stop_event = threading.Event() + self.stop_reason = "stopped" + self.holder = None + self.holder_key: Optional[str] = None + self.holder_model: Optional[dict] = None + self.holder_is_legacy = False + self.stats = {"holder_loads": 0, "generations": 0, "errors": 0, "cancelled": 0, "oom_retries": 0} + self._lock = threading.Lock() + # resident components shared across holders (same TE/VAE under two archs) + self.pool = ComponentPool() + + # ------------------------------------------------------------------ api + def submit(self, model: dict, sample: dict, stream: Optional[dict] = None) -> GenerationJob: + if "arch" not in model: + raise ValueError("model.arch is required") + job = GenerationJob(model, sample, stream) + with self._lock: + self.jobs[job.request_id] = job + self.queue.put(job) + return job + + def cancel(self, request_id: str) -> bool: + job = self.jobs.get(request_id) + if job is None: + return False + job.cancel.set() + if job.status == "queued": + job.status = "cancelled" + job.emit("error", message="cancelled", cancelled=True) + job.finish() + return True + + def stop(self, reason: str = "stopped"): + self.stop_reason = reason + self.stop_event.set() + if self.current is not None: + self.current.cancel.set() + + def queued(self) -> List[dict]: + return [j.info() for j in self.jobs.values() if j.status in ("queued", "running")] + + def recent(self, limit: int = 50) -> List[dict]: + jobs = sorted(self.jobs.values(), key=lambda j: j.created_at, reverse=True) + return [j.info() for j in jobs[:limit]] + + def health(self) -> dict: + vram = None + if torch.cuda.is_available(): + try: + free, total = torch.cuda.mem_get_info(torch.device(self.device)) + vram = {"used": total - free, "total": total} + except Exception: + vram = None + return { + "ok": True, + "busy": self.current is not None, + "queue": len([j for j in self.jobs.values() if j.status == "queued"]), + "current": self.current.info() if self.current else None, + "active": {"model": self.holder_model, "arch": (self.holder_model or {}).get("arch")} + if self.holder is not None + else None, + "device": self.device, + "vram": vram, + "stats": dict(self.stats), + "pool": self.pool.stats(), + } + + def _drop_holder(self): + """Drop the holder object; its components stay resident in the pool + until release_unused()/clear() decides.""" + holder = self.holder + self.holder = None + self.holder_key = None + self.holder_model = None + if holder is not None: + for attr in ("pipeline", "model", "vae", "text_encoder", "image_encoder", "adapter", "network"): + try: + setattr(holder, attr, None) + except Exception: + pass + del holder + gc.collect() + + def unload(self): + if self.holder is None and not self.pool.entries: + return + self.status_cb("Unloading model") + self._drop_holder() + self.pool.clear() + flush() + + def preload(self, model: dict): + """Load a model's holder ahead of any request (warm start). Runs on + the caller's thread; meant for the engine main thread before + run_forever.""" + job = GenerationJob(dict(model), {}, {"latents": "none"}) + job.status = "running" + try: + self._ensure_holder(job) + job.status = "done" + except Exception as e: + job.status = "error" + job.error = f"{type(e).__name__}: {e}" + traceback.print_exc() + finally: + job.finish() + return job + + # ------------------------------------------------------------------ loop + def run_forever(self): + while not self.stop_event.is_set(): + try: + job = self.queue.get(timeout=0.5) + except queue.Empty: + continue + except KeyboardInterrupt: + self.stop("stopped") + break + if job.status != "queued": + continue + self.current = job + try: + self._run(job) + except KeyboardInterrupt: + job.status = "cancelled" + job.emit("error", message="engine stopping", cancelled=True) + job.finish() + self.stop("stopped") + break + finally: + self.current = None + self._idle_status() + # drain anything left queued + while True: + try: + job = self.queue.get_nowait() + except queue.Empty: + break + self.cancel(job.request_id) + + def _idle_status(self): + if self.holder_model is not None: + self.status_cb(f"Idle - {self.holder_model.get('arch')} ({self.holder_model.get('name_or_path')})") + else: + self.status_cb("Idle - no model loaded") + + # ------------------------------------------------------------------ holder + def _build_model_config(self, model: dict) -> ModelConfig: + kwargs = dict(model) + arch = kwargs.pop("arch") + entry = get_arch_entry(arch) + for k, v in entry["model"].items(): + kwargs.setdefault(k, v) + kwargs.setdefault("dtype", self.dtype) + if kwargs.get("quantize"): + kwargs.setdefault("qtype", "convrot8") + if kwargs.get("quantize_te"): + kwargs.setdefault("qtype_te", "convrot8") + return ModelConfig(arch=arch, **kwargs) + + def _ensure_holder(self, job: GenerationJob): + model = dict(job.model) + entry = get_arch_entry(model["arch"]) + for k, v in entry["model"].items(): + model.setdefault(k, v) + key = _model_key(model) + if self.holder is not None and key == self.holder_key: + return self.holder + if self.holder is not None: + self._drop_holder() + + self.pool.begin_request() + ComponentPool.current = self.pool + try: + try: + holder, is_legacy = self._build_holder(job, model) + except torch.OutOfMemoryError: + # components the previous model needed but this one does not + # were still resident: free them and try once more + self.stats["oom_retries"] += 1 + job.emit("status", message="Out of memory: freeing unused components and retrying") + freed = self.pool.release_unused() + flush() + job.emit("status", message=f"Freed {freed / 1e9:.1f} GB, retrying load") + holder, is_legacy = self._build_holder(job, model) + finally: + ComponentPool.current = None + freed = self.pool.release_unused() + if freed: + job.emit("status", message=f"Released {freed / 1e9:.1f} GB of components not used by {model['arch']}") + flush() + self.holder = holder + self.holder_key = key + self.holder_model = model + self.holder_is_legacy = is_legacy + return holder + + def _build_holder(self, job: GenerationJob, model: dict): + model_config = self._build_model_config(model) + ModelClass = get_model_class(model_config) + is_legacy = not hasattr(ModelClass, "get_train_scheduler") + if is_legacy: + from toolkit.sampler import get_sampler + + legacy_arch = "sd" + if model_config.is_pixart: + legacy_arch = "pixart" + elif model_config.is_flux: + legacy_arch = "flux" + elif model_config.is_lumina2: + legacy_arch = "lumina2" + sampler = get_sampler( + "ddpm", + {"prediction_type": "v_prediction" if model_config.is_v_pred else "epsilon"}, + arch=legacy_arch, + ) + else: + sampler = ModelClass.get_train_scheduler() + + job.emit("status", message=f"Loading {model['arch']}") + self.status_cb(f"Loading {model['arch']}: {model.get('name_or_path')}") + holder = ModelClass( + device=self.device, + model_config=model_config, + dtype=self.dtype, + noise_scheduler=sampler, + ) + + def _status_hook(msg): + job.emit("status", message=str(msg)) + self.status_cb(f"Loading {model['arch']}: {msg}") + + if hasattr(holder, "add_status_update_hook"): + holder.add_status_update_hook(_status_hook) + t0 = time.perf_counter() + holder.load_model() + self.stats["holder_loads"] += 1 + p = self.pool.stats() + job.emit( + "status", + message=f"Loaded {model['arch']} in {time.perf_counter() - t0:.1f}s (pool hits {p['hits']}, misses {p['misses']})", + ) + return holder, is_legacy + + # ------------------------------------------------------------------ run + def _resolve_asset(self, value): + if value is None: + return None + if os.path.isabs(value) and os.path.exists(value): + return value + cand = os.path.join(self.assets_folder, os.path.basename(value)) + if os.path.exists(cand): + return cand + if os.path.exists(value): + return os.path.abspath(value) + raise FileNotFoundError(f"control image not found: {value}") + + def _run(self, job: GenerationJob): + job.status = "running" + job.started_at = time.time() + arch = job.model.get("arch") + entry = get_arch_entry(arch) + out_dir = os.path.join(self.output_folder, job.request_id) + os.makedirs(out_dir, exist_ok=True) + try: + holder = self._ensure_holder(job) + if job.cancel.is_set(): + raise GenerationCancelled() + + sample = dict(entry["sample"]) + sample.update({k: v for k, v in job.sample.items() if v is not None}) + for key in ("ctrl_img", "ctrl_img_1", "ctrl_img_2", "ctrl_img_3"): + if sample.get(key): + sample[key] = self._resolve_asset(sample[key]) + if entry["needs_control_image"] and not sample.get("ctrl_img"): + raise ValueError(f"{arch} requires a control image (ctrl_img)") + output_ext = sample.pop("output_ext", None) or OUTPUT_EXT[entry["modality"]] + if entry["modality"] == "video" and sample.get("num_frames", 1) > 1 and output_ext == "mp4": + # BaseModel's video saver writes animated webp; archs with their + # own mp4 writer (ltx2) override output_ext themselves + output_ext = "webp" + prompt = sample.pop("prompt", "") + gen = GenerateImageConfig( + prompt=prompt, + output_folder=out_dir, + output_ext=output_ext, + **sample, + ) + job.total_steps = gen.num_inference_steps + job.emit( + "start", + model=self.holder_model, + sample={ + "prompt": gen.prompt, + "negative_prompt": gen.negative_prompt, + "width": gen.width, + "height": gen.height, + "num_inference_steps": gen.num_inference_steps, + "guidance_scale": gen.guidance_scale, + "seed": gen.seed, + "num_frames": gen.num_frames, + "fps": gen.fps, + }, + modality=entry["modality"], + preview=preview_info(arch), + ) + + t_start = time.perf_counter() + every_n = max(1, int(job.stream.get("every_n_steps", 1))) + want_latents = job.stream.get("latents", "raw") == "raw" + max_frames = int(job.stream.get("max_frames", 4)) + + def step_hook(idx, num_steps, latents): + if job.cancel.is_set(): + raise GenerationCancelled() + job.step = idx + 1 + if num_steps: + job.total_steps = num_steps + elapsed = time.perf_counter() - t_start + job.emit("progress", step=job.step, total=job.total_steps, elapsed=round(elapsed, 3)) + self.status_cb(f"Generating {job.step}/{job.total_steps or '?'}") + if want_latents and (idx % every_n == 0 or (job.total_steps and job.step == job.total_steps)): + self._emit_latent(job, latents, gen, max_frames) + + def status_hook(msg): + job.emit("status", message=str(msg)) + + holder.sample_step_hook = step_hook + if hasattr(holder, "add_status_update_hook"): + holder._status_update_hooks.append(status_hook) + gen_kwargs = {"sampler": "ddpm"} if self.holder_is_legacy else {} + self.status_cb(f"Generating with {arch}") + before = set(_list_outputs(out_dir)) + t_gen = time.perf_counter() + try: + holder.generate_images([gen], **gen_kwargs) + finally: + gen_seconds = time.perf_counter() - t_gen + holder.sample_step_hook = None + if hasattr(holder, "_status_update_hooks"): + try: + holder._status_update_hooks.remove(status_hook) + except ValueError: + pass + produced = [p for p in _list_outputs(out_dir) if p not in before] + if not produced: + raise RuntimeError("generation produced no output file") + steps = job.total_steps or gen.num_inference_steps + per_step = gen_seconds / steps if steps else 0 + size = f"{gen.width}x{gen.height}" + (f"x{gen.num_frames}f" if gen.num_frames > 1 else "") + timing = f"Generated {arch} {size} in {gen_seconds:.2f}s ({steps} steps, {per_step:.2f}s/step, seed {gen.seed})" + print(f"[AITK] {timing}", flush=True) + job.emit("status", message=timing) + for path in sorted(produced): + ext = os.path.splitext(path)[1].lower().lstrip(".") + result = { + "path": path, + "relpath": os.path.relpath(path, self.output_folder), + "kind": _kind_for_ext(ext), + "ext": ext, + "bytes": os.path.getsize(path), + "seed": gen.seed, + "width": gen.width, + "height": gen.height, + "num_frames": gen.num_frames, + "fps": gen.fps, + "seconds": round(gen_seconds, 3), + "steps": steps, + } + job.results.append(result) + job.emit("result", **result) + job.status = "done" + self.stats["generations"] += 1 + except GenerationCancelled: + job.status = "cancelled" + job.error = "cancelled" + self.stats["cancelled"] += 1 + job.emit("error", message="cancelled", cancelled=True) + flush() + except Exception as e: + job.status = "error" + job.error = f"{type(e).__name__}: {e}" + self.stats["errors"] += 1 + traceback.print_exc() + job.emit("error", message=job.error, cancelled=False, traceback=traceback.format_exc()) + flush() + finally: + if job.status != "done": + shutil.rmtree(out_dir, ignore_errors=True) + job.finish() + + def _emit_latent(self, job: GenerationJob, latents, gen: GenerateImageConfig, max_frames: int): + try: + t = latents + if hasattr(t, "tensor"): + t = t.tensor + if not isinstance(t, torch.Tensor): + return + t = t.detach() + layout = None + if t.dim() == 3 and t.shape[-1] % 4 == 0: + # packed 2x2 patch tokens (flux family): (B, h/2*w/2, C*4) + h = gen.height // 16 + w = gen.width // 16 + if t.shape[1] == h * w: + c = t.shape[-1] // 4 + t = t.view(t.shape[0], h, w, c, 2, 2).permute(0, 3, 1, 4, 2, 5).reshape(t.shape[0], c, h * 2, w * 2) + if t.dim() == 4: + layout = "BCHW" + elif t.dim() == 5: + layout = "BCFHW" + if max_frames and t.shape[2] > max_frames: + idx = torch.linspace(0, t.shape[2] - 1, max_frames).round().long().to(t.device) + t = t.index_select(2, idx) + elif t.dim() == 3: + layout = "BLC" # sequence latents (audio: batch, length, channels) + else: + layout = "raw" + fields, payload = tensor_payload(t[:1] if t.dim() >= 3 else t) + extra = {} + if not getattr(job, "_preview_sent", False) and layout in ("BCHW", "BCFHW"): + # now that the real latent shape is known, resolve the preview + # table (arch table, else a channel-count guess) once per job + job._preview_sent = True + extra["preview"] = preview_info(job.model.get("arch"), int(t.shape[1]), 2 if layout == "BCHW" else 3) + job.emit("latent", payload, step=job.step, total=job.total_steps, layout=layout, **fields, **extra) + except Exception as e: + job.emit("status", message=f"latent preview skipped: {e}") + + +def _list_outputs(folder: str) -> List[str]: + return [ + p + for p in glob.glob(os.path.join(folder, "*")) + if os.path.isfile(p) and not os.path.basename(p).startswith(".") and not p.endswith(".txt") + ] + + +def _kind_for_ext(ext: str) -> str: + if ext in ("mp4", "webm", "mov", "webp"): + return "video" if ext != "webp" else "image" + if ext in ("wav", "mp3", "flac", "ogg"): + return "audio" + return "image" diff --git a/extensions_built_in/inference_engine/latent_format_tables.py b/extensions_built_in/inference_engine/latent_format_tables.py new file mode 100644 index 0000000..b02bce4 --- /dev/null +++ b/extensions_built_in/inference_engine/latent_format_tables.py @@ -0,0 +1,543 @@ +"""Latent -> RGB preview tables. + +Linear per-pixel projections (channels x 3, plus bias) so a client can turn a +raw latent into a rough RGB preview without a VAE. Numeric tables come from +ComfyUI comfy/latent_formats.py (same VAE families). Layouts: dims 2 = BCHW +image latents, dims 3 = BCFHW video latents. "reshape" names a packing the +client must undo first (flux2: 32 channels x 2x2 pixel shuffle). +""" + +LATENT_FORMATS = { + "SD15": { + "channels": 4, "dims": 2, "spatial": 8, "temporal": 1, + "reshape": None, + "bias": None, + "factors": [ + [0.3512, 0.2297, 0.3227], + [0.325, 0.4974, 0.235], + [-0.2829, 0.1762, 0.2721], + [-0.212, -0.2616, -0.7177], + ], + }, + "SDXL": { + "channels": 4, "dims": 2, "spatial": 8, "temporal": 1, + "reshape": None, + "bias": [0.1084, -0.0175, -0.0011], + "factors": [ + [0.3651, 0.4232, 0.4341], + [-0.2533, -0.0042, 0.1068], + [0.1076, 0.1111, -0.0362], + [-0.3165, -0.2492, -0.2188], + ], + }, + "SD3": { + "channels": 16, "dims": 2, "spatial": 8, "temporal": 1, + "reshape": None, + "bias": [0.2394, 0.2135, 0.1925], + "factors": [ + [-0.0922, -0.0175, 0.0749], + [0.0311, 0.0633, 0.0954], + [0.1994, 0.0927, 0.0458], + [0.0856, 0.0339, 0.0902], + [0.0587, 0.0272, -0.0496], + [-0.0006, 0.1104, 0.0309], + [0.0978, 0.0306, 0.0427], + [-0.0042, 0.1038, 0.1358], + [-0.0194, 0.002, 0.0669], + [-0.0488, 0.013, -0.0268], + [0.0922, 0.0988, 0.0951], + [-0.0278, 0.0524, -0.0542], + [0.0332, 0.0456, 0.0895], + [-0.0069, -0.003, -0.081], + [-0.0596, -0.0465, -0.0293], + [-0.1448, -0.1463, -0.1189], + ], + }, + "Flux": { + "channels": 16, "dims": 2, "spatial": 8, "temporal": 1, + "reshape": None, + "bias": [-0.0329, -0.0718, -0.0851], + "factors": [ + [-0.0346, 0.0244, 0.0681], + [0.0034, 0.021, 0.0687], + [0.0275, -0.0668, -0.0433], + [-0.0174, 0.016, 0.0617], + [0.0859, 0.0721, 0.0329], + [0.0004, 0.0383, 0.0115], + [0.0405, 0.0861, 0.0915], + [-0.0236, -0.0185, -0.0259], + [-0.0245, 0.025, 0.118], + [0.1008, 0.0755, -0.0421], + [-0.0515, 0.0201, 0.0011], + [0.0428, -0.0012, -0.0036], + [0.0817, 0.0765, 0.0749], + [-0.1264, -0.0522, -0.1103], + [-0.028, -0.0881, -0.0499], + [-0.1262, -0.0982, -0.0778], + ], + }, + "Flux2": { + "channels": 128, "dims": 2, "spatial": 16, "temporal": 1, + "reshape": "flux2_2x2", + "bias": [-0.0329, -0.0718, -0.0851], + "factors": [ + [0.0058, 0.0113, 0.0073], + [0.0495, 0.0443, 0.0836], + [-0.0099, 0.0096, 0.0644], + [0.2144, 0.3009, 0.3652], + [0.0166, -0.0039, -0.0054], + [0.0157, 0.0103, -0.016], + [-0.0398, 0.0902, -0.0235], + [-0.0052, 0.0095, 0.0109], + [-0.3527, -0.2712, -0.1666], + [-0.0301, -0.0356, -0.018], + [-0.0107, 0.0078, 0.0013], + [0.0746, 0.009, -0.0941], + [0.0156, 0.0169, 0.007], + [-0.0034, -0.004, -0.0114], + [0.0032, 0.0181, 0.008], + [-0.0939, -0.0008, 0.0186], + [0.0018, 0.0043, 0.0104], + [0.0284, 0.0056, -0.0127], + [-0.0024, -0.0022, -0.003], + [0.1207, -0.0026, 0.0065], + [0.0128, 0.0101, 0.0142], + [0.0137, -0.0072, -0.0007], + [0.0095, 0.0092, -0.0059], + [0.0, -0.0077, -0.0049], + [-0.0465, -0.0204, -0.0312], + [0.0095, 0.0012, -0.0066], + [0.029, -0.0034, 0.0025], + [0.022, 0.0169, -0.0048], + [-0.0332, -0.0457, -0.0468], + [-0.0085, 0.0389, 0.0609], + [-0.0076, 0.0003, -0.0043], + [-0.0111, -0.046, -0.0614], + ], + }, + "Wan21": { + "channels": 16, "dims": 3, "spatial": 8, "temporal": 4, + "reshape": None, + "bias": [-0.1835, -0.0868, -0.336], + "factors": [ + [-0.1299, -0.1692, 0.2932], + [0.0671, 0.0406, 0.0442], + [0.3568, 0.2548, 0.1747], + [0.0372, 0.2344, 0.142], + [0.0313, 0.0189, -0.0328], + [0.0296, -0.0956, -0.0665], + [-0.3477, -0.4059, -0.2925], + [0.0166, 0.1902, 0.1975], + [-0.0412, 0.0267, -0.1364], + [-0.1293, 0.074, 0.1636], + [0.068, 0.3019, 0.1128], + [0.0032, 0.0581, 0.0639], + [-0.1251, 0.0927, 0.1699], + [0.006, -0.0633, 0.0005], + [0.3477, 0.2275, 0.295], + [0.1984, 0.0913, 0.1861], + ], + }, + "Wan22": { + "channels": 48, "dims": 3, "spatial": 16, "temporal": 4, + "reshape": None, + "bias": [0.0317, -0.0878, -0.1388], + "factors": [ + [0.0119, 0.0103, 0.0046], + [-0.1062, -0.0504, 0.0165], + [0.014, 0.0409, 0.0491], + [-0.0813, -0.0677, 0.0607], + [0.0656, 0.0851, 0.0808], + [0.0264, 0.0463, 0.0912], + [0.0295, 0.0326, 0.059], + [-0.0244, -0.027, 0.0025], + [0.0443, -0.0102, 0.0288], + [-0.0465, -0.009, -0.0205], + [0.0359, 0.0236, 0.0082], + [-0.0776, 0.0854, 0.1048], + [0.0564, 0.0264, 0.0561], + [0.0006, 0.0594, 0.0418], + [-0.0319, -0.0542, -0.0637], + [-0.0268, 0.0024, 0.026], + [0.0539, 0.0265, 0.0358], + [-0.0359, -0.0312, -0.0287], + [-0.0285, -0.1032, -0.1237], + [0.1041, 0.0537, 0.0622], + [-0.0086, -0.0374, -0.0051], + [0.039, 0.067, 0.2863], + [0.0069, 0.0144, 0.0082], + [0.0006, -0.0167, 0.0079], + [0.0313, -0.0574, -0.0232], + [-0.1454, -0.0902, -0.0481], + [0.0714, 0.0827, 0.0447], + [-0.0304, -0.0574, -0.0196], + [0.0401, 0.0384, 0.0204], + [-0.0758, -0.0297, -0.0014], + [0.0568, 0.1307, 0.1372], + [-0.0055, -0.031, -0.038], + [0.0239, -0.0305, 0.0325], + [-0.0663, -0.0673, -0.014], + [-0.0416, -0.0047, -0.0023], + [0.0166, 0.0112, -0.0093], + [-0.0211, 0.0011, 0.0331], + [0.1833, 0.1466, 0.225], + [-0.0368, 0.037, 0.0295], + [-0.3441, -0.3543, -0.2008], + [-0.0479, -0.0489, -0.042], + [-0.066, -0.0153, 0.08], + [-0.0101, 0.0068, 0.0156], + [-0.069, -0.0452, -0.0927], + [-0.0145, 0.0041, 0.0015], + [0.0421, 0.0451, 0.0373], + [0.0504, -0.0483, -0.0356], + [-0.0837, 0.0168, 0.0055], + ], + }, + "LTXV": { + "channels": 128, "dims": 3, "spatial": 32, "temporal": 8, + "reshape": None, + "bias": [-0.0571, -0.1657, -0.2512], + "factors": [ + [0.011202, -0.00063815, -0.010021], + [0.086031, 0.065813, 0.00095409], + [-0.012576, -0.0075734, -0.0040528], + [0.0094063, -0.0021688, 0.0026093], + [0.0037636, 0.012765, 0.0091548], + [0.021024, -0.0052973, 0.0034373], + [-0.0088896, -0.019703, -0.018761], + [-0.01316, -0.010523, 0.0019709], + [-0.0015152, -0.0069891, -0.007581], + [-0.0017247, 0.0004656, -0.0033839], + [0.013617, 0.0047077, -0.0020045], + [0.010256, 0.0077318, 0.013948], + [-0.016108, -0.0062151, 0.0011561], + [0.0073407, 0.015628, 0.00044865], + [0.00095357, -0.0029518, -0.01476], + [0.019143, 0.010868, 0.012264], + [0.0044575, 3.6682e-05, -0.0068508], + [-0.00045681, 0.003257, 0.0077929], + [0.033902, 0.033405, 0.037454], + [-0.023001, -0.0024877, -0.0031033], + [0.050265, 0.038841, 0.033539], + [-0.0041018, -0.0011095, 0.0015859], + [-0.12689, -0.13107, -0.21005], + [0.026276, 0.014189, -0.0035963], + [-0.0048679, 0.0088486, 0.0078029], + [-0.001661, -0.0048597, -0.005206], + [-0.002101, 0.002361, 0.0093796], + [-0.022482, -0.021305, -0.015087], + [-0.015753, -0.010646, -0.0065083], + [-0.0046975, 0.0050288, -0.006739], + [0.011951, 0.020712, 0.016191], + [-0.0063704, -0.0084827, -0.0095483], + [0.007261, -0.0099326, -0.022978], + [-0.00091904, 0.0062882, 0.009572], + [-0.037178, -0.037123, -0.056713], + [-0.13373, -0.1072, -0.053801], + [-0.0053702, 0.0081256, 0.0088397], + [-0.15247, -0.21437, -0.21843], + [0.031441, 0.0070335, -0.0097541], + [0.0021528, -0.0089817, -0.021023], + [0.0038461, -0.0058957, -0.015014], + [-0.004347, -0.01294, -0.015972], + [-0.0054781, -0.010842, -0.0030204], + [-0.0065347, 0.0030806, -0.010163], + [-0.0050414, -0.0071503, -0.00089686], + [-0.0085851, -0.0024351, 0.0010674], + [-0.0090016, -0.0096493, 0.0015692], + [0.0050914, 0.012099, 0.019968], + [0.013758, 0.011669, 0.0081958], + [-0.010518, -0.011575, -0.0041307], + [-0.02841, -0.031266, -0.022149], + [0.0029336, 0.036511, 0.018717], + [-0.016703, -0.016696, -0.0044529], + [0.048818, 0.040063, 0.008741], + [-0.015066, -0.00057328, 0.0029785], + [-0.017613, -0.0081034, 0.013086], + [-0.0092633, 0.010803, -0.0063489], + [0.0030851, 0.0004775, 0.012347], + [-0.022785, -0.023043, -0.026005], + [-0.024787, -0.015389, -0.022104], + [-0.023572, 0.0010544, 0.012361], + [-0.0078915, -0.0012271, -0.0060968], + [-0.011478, -0.0012543, 0.0062679], + [-0.054229, 0.026644, 0.0063394], + [0.0044216, -0.0073338, -0.010464], + [-0.0045013, 0.0016082, 0.01442], + [0.013673, 0.0088877, 0.0041253], + [-0.010145, 0.0090072, 0.015695], + [-0.0056234, 0.0011847, 0.0081261], + [-0.0037171, -0.0053538, 0.001259], + [0.029476, 0.021424, 0.030424], + [-0.034925, -0.02434, -0.025316], + [-0.034127, -0.022406, -0.010589], + [-0.017342, -0.013249, -0.010719], + [-0.0021478, -0.0086051, -0.0029878], + [0.0012089, -0.0042391, -0.0068569], + [0.00090411, -0.0066886, -6.7547e-05], + [0.016048, -0.010057, -0.028929], + [0.001229, 0.010163, 0.018861], + [0.017264, 0.00027257, 0.013785], + [-0.013482, -0.0036427, 0.00067481], + [0.0046782, -0.0052423, 0.0024467], + [-0.0059113, -0.0062244, -0.0018162], + [0.015496, 0.014582, 0.0019514], + [0.0074958, 0.0015886, -0.0082305], + [0.019086, 0.001636, -0.0039674], + [-0.0057021, -0.0027307, -0.0041066], + [0.001745, 0.014602, 0.025794], + [-0.00082788, 0.0022902, 0.0045161], + [0.011632, 0.0089193, -0.0072813], + [0.0075721, 0.0026784, 0.011393], + [0.0051939, 0.0036903, 0.014049], + [-0.018383, -0.022529, -0.024477], + [0.00058842, -0.0057874, -0.01477], + [-0.016125, -0.0086101, -0.014533], + [0.02054, 0.020729, 0.0064338], + [0.0033587, -0.011226, -0.016444], + [-0.0014742, -0.010489, 0.0017097], + [0.02813, 0.023546, 0.032791], + [-0.018532, -0.012842, -0.0087756], + [-0.0080533, -0.010771, -0.017536], + [-0.0039009, 0.01615, 0.033359], + [-0.0074554, -0.014154, -0.006191], + [0.0034734, -0.01137, -0.010581], + [0.011476, 0.0039281, 0.0028231], + [0.0071639, -0.0014741, -0.0038066], + [0.002225, -0.0087552, -0.0095719], + [0.024146, 0.021696, 0.028056], + [-0.0054365, -0.024291, -0.017802], + [0.0074263, 0.01051, 0.012705], + [0.0062669, 0.0062658, 0.019211], + [0.016378, 0.0094933, 0.0066971], + [0.017173, 0.023601, 0.023296], + [-0.014568, -0.0098279, -0.011556], + [0.014431, 0.01443, 0.0066362], + [-0.006823, 0.018863, 0.014555], + [0.0061156, 0.00347, -0.0026662], + [-0.0026983, -0.0059402, -0.0092276], + [0.010235, 0.0074173, -0.0076243], + [-0.013255, 0.019322, -0.00092153], + [0.0024222, -0.0048039, -0.015759], + [0.026244, 0.025951, 0.020249], + [0.015711, 0.018498, 0.0027407], + [-0.0021714, 0.0047214, -0.022443], + [-0.0074747, 0.0074166, 0.01443], + [-0.0083906, -0.0079776, 0.0097927], + [0.038321, 0.0096622, -0.019268], + [-0.014605, -0.0067032, 0.0039675], + ], + }, + "LTXAV": { + "channels": 128, "dims": 3, "spatial": 32, "temporal": 8, + "reshape": None, + "bias": [-0.347892, -0.363814, -0.370287], + "factors": [ + [0.001135, -0.010555, -0.004925], + [-0.008019, -0.006231, -0.005564], + [0.012637, 0.005605, 0.012713], + [0.023454, 0.020771, 0.017844], + [-0.01194, -0.000932, 0.009292], + [0.018602, 0.011018, 0.013969], + [-0.036369, -0.046631, -0.057898], + [-0.031919, 0.000131, 0.015214], + [0.014519, 0.021041, 0.015325], + [0.018889, 0.016149, -0.002836], + [-0.003784, -0.006057, -0.008195], + [0.013262, 0.030259, 0.029775], + [0.050465, 0.050366, 0.025255], + [0.018628, 0.007691, 0.002893], + [-0.015698, -0.008451, -0.000676], + [-0.0136, -0.012587, -0.004437], + [0.012482, 0.021469, 0.027913], + [-0.018241, -0.013488, -0.010975], + [0.013828, 0.012568, 0.021984], + [0.017911, 0.006552, 0.005567], + [0.026769, 0.006803, -0.00936], + [-0.006794, -0.008447, -0.013921], + [0.029708, 0.018671, 0.022811], + [-0.014732, -0.019169, 0.000903], + [0.019607, 0.032595, 0.053409], + [-0.003721, 0.003976, 0.010364], + [-0.020193, -0.026076, -0.036068], + [-0.002328, 0.006527, 0.013052], + [0.017171, 0.009224, 0.006548], + [0.001104, -0.000591, 0.000147], + [-0.000217, 0.011834, 0.017945], + [-0.015329, -0.012463, -0.006178], + [-0.009478, -0.00868, -0.004107], + [-0.005565, -0.006006, -0.001493], + [0.009451, 0.008794, 0.013207], + [-0.009989, -0.008027, -0.009568], + [-0.001505, -0.008805, -0.006828], + [0.001105, 0.008999, 0.009079], + [0.025935, 0.016426, 0.008036], + [0.006313, 0.000694, -0.006039], + [-0.001893, -0.006951, -0.00956], + [-0.007082, -0.002566, -0.007152], + [-0.005231, 0.004829, 0.00822], + [-0.004333, 0.001251, -0.004852], + [-0.017024, -0.01273, -0.007457], + [0.024988, 0.032963, 0.036556], + [0.013697, 0.012278, 0.009979], + [-0.013751, -0.008369, -0.015446], + [-0.009348, -0.001047, 0.007622], + [-0.003135, -0.00335, -0.003766], + [0.007436, 0.004957, 0.01048], + [0.018315, 0.022066, 0.021104], + [-0.005621, -0.00677, -0.008219], + [-0.007427, 0.001911, -0.001231], + [-0.007413, 0.000486, -0.006039], + [-0.014698, -0.00716, 0.006509], + [0.013775, 0.014185, 0.008203], + [0.060246, 0.069787, 0.072833], + [0.009861, 0.00487, 0.001194], + [-0.00366, 0.003251, 0.008015], + [0.003696, -0.00368, -0.008851], + [0.014924, 0.006196, 0.005282], + [-0.00674, -0.004319, -0.006729], + [0.020635, 0.015163, 0.012385], + [-0.032623, -0.006105, 0.010436], + [-0.058988, -0.030162, -0.037961], + [-0.035614, -0.021929, -0.011062], + [-0.023412, -0.011305, -0.005054], + [-0.002716, -0.005184, -0.004084], + [0.014591, 0.015294, 0.014045], + [0.00831, 0.002466, -0.003225], + [0.005176, 0.001119, 0.000695], + [-0.021569, -0.030886, -0.044732], + [0.007517, 0.003891, 0.000551], + [-0.006793, 0.004059, 0.010184], + [-0.086481, -0.082033, -0.083414], + [0.004192, 0.000762, -0.008658], + [0.01097, 0.009002, 0.007384], + [0.004042, -0.006732, -0.011031], + [0.012164, 0.006401, 0.007483], + [0.029252, 0.01399, 0.011128], + [0.048452, 0.034648, 0.016269], + [0.024104, 0.012647, 0.011754], + [-0.013216, -0.020192, -0.019752], + [-0.010799, -0.008535, -0.005467], + [0.005823, 0.001403, 0.00189], + [0.052393, 0.044771, 0.032777], + [0.007576, -0.00808, -0.012453], + [0.00983, 0.004244, 0.001213], + [-0.025867, -0.013169, -0.010636], + [0.008494, 0.003135, 0.00079], + [0.003969, -0.002625, -0.010204], + [0.006509, 0.008272, 0.020819], + [-0.004943, -0.013424, -0.015351], + [0.005541, 0.009136, -0.003666], + [-0.0143, -0.015864, -0.016853], + [0.00265, 0.028393, 0.014125], + [-0.027661, -0.045422, -0.064995], + [0.00922, 0.015522, 0.010574], + [-0.002236, 0.002915, 0.004557], + [-0.020269, -0.008212, -0.000532], + [0.019294, 0.003655, -0.002809], + [0.007116, -0.002784, 1.7e-05], + [0.057277, 0.07327, 0.074401], + [-0.002616, -0.001696, -0.000498], + [0.007248, 0.009793, 0.022829], + [-0.00259, -0.005601, -0.000436], + [-0.007681, 0.003893, -0.004119], + [-0.057392, -0.045545, -0.02529], + [0.045188, 0.047985, 0.054059], + [0.000937, -0.008861, -0.038406], + [-0.010192, -0.008036, -0.005385], + [-0.030222, -0.027498, -0.030765], + [-0.008359, 0.013247, 0.010918], + [0.004102, 0.002093, 0.006934], + [0.039461, 0.027339, 0.008284], + [-0.075747, -0.07634, -0.071625], + [0.002692, 0.005096, -0.002247], + [-0.002453, -0.002785, -0.010483], + [0.012265, 0.005481, 0.001729], + [0.017755, 0.008655, 0.003532], + [0.05556, 0.049128, 0.044137], + [-0.025861, -0.023798, -0.018815], + [-0.014876, -0.01077, -0.010713], + [-0.017315, -0.012599, -0.008661], + [-0.008461, -0.00621, -0.007744], + [-0.040175, -0.042255, -0.048119], + [-0.019355, -0.021055, -0.021919], + ], + }, + "MiniMaxH3Video": { + "channels": 24, "dims": 3, "spatial": 16, "temporal": 4, + "reshape": None, + "bias": [0.057426, -0.022078, -0.071449], + "factors": [ + [-0.018555, 0.024344, -0.017536], + [0.150164, 0.137244, 0.129221], + [0.027367, -0.050369, -0.208606], + [-0.000793, -0.164622, -0.323161], + [-0.048556, 0.01397, -0.074286], + [0.01174, 0.014172, -0.006906], + [0.061517, 0.061212, 0.110025], + [0.035321, 0.086879, 0.110059], + [-0.017426, 0.002997, 0.035356], + [0.531539, 0.548819, 0.624404], + [-0.024968, -0.040234, -0.034302], + [-0.032549, -0.029096, -0.017221], + [0.022609, 0.020286, 0.050661], + [-0.084001, -0.038131, -0.020805], + [-0.01883, 0.010412, 0.06112], + [0.020777, 0.011196, -0.030994], + [-0.00839, -0.012201, -0.025687], + [-0.013281, -0.002924, 0.006331], + [0.00026, 0.001833, -0.011038], + [0.105471, 0.100482, 0.132106], + [0.016529, 0.015213, 0.009999], + [-0.014015, -0.017438, -0.019134], + [-0.033787, -0.009984, -0.019725], + [0.004224, 0.017284, 0.027196], + ], + }, + "HunyuanVideo": { + "channels": 16, "dims": 3, "spatial": 8, "temporal": 4, + "reshape": None, + "bias": [0.0259, -0.0192, -0.0761], + "factors": [ + [-0.0395, -0.0331, 0.0445], + [0.0696, 0.0795, 0.0518], + [0.0135, -0.0945, -0.0282], + [0.0108, -0.025, -0.0765], + [-0.0209, 0.0032, 0.0224], + [-0.0804, -0.0254, -0.0639], + [-0.0991, 0.0271, -0.0669], + [-0.0646, -0.0422, -0.04], + [-0.0696, -0.0595, -0.0894], + [-0.0799, -0.0208, -0.0375], + [0.1166, 0.1627, 0.0962], + [0.1165, 0.0432, 0.0407], + [-0.2315, -0.192, -0.1355], + [-0.027, 0.0401, -0.0821], + [-0.0616, -0.0997, -0.0727], + [0.0249, -0.0469, -0.1703], + ], + }, + "Cosmos1CV8x8x8": { + "channels": 16, "dims": 3, "spatial": 8, "temporal": 8, + "reshape": None, + "bias": [-0.1223, -0.1889, -0.1976], + "factors": [ + [0.1817, 0.2284, 0.2423], + [-0.0586, -0.0862, -0.3108], + [-0.4703, -0.4255, -0.3995], + [0.0803, 0.1963, 0.1001], + [-0.082, -0.105, 0.04], + [0.2511, 0.3098, 0.2787], + [-0.183, -0.2117, -0.004], + [-0.0621, -0.2187, -0.0939], + [0.3619, 0.1082, 0.1455], + [0.3164, 0.3922, 0.2575], + [0.1152, 0.0231, -0.0462], + [-0.1434, -0.3609, -0.3665], + [0.0635, 0.1471, 0.168], + [-0.3635, -0.1963, -0.3248], + [-0.1865, 0.0365, 0.2346], + [0.0447, 0.0994, 0.0881], + ], + }, +} diff --git a/extensions_built_in/inference_engine/latent_preview.py b/extensions_built_in/inference_engine/latent_preview.py new file mode 100644 index 0000000..a29f1ed --- /dev/null +++ b/extensions_built_in/inference_engine/latent_preview.py @@ -0,0 +1,86 @@ +from typing import Optional + +from .latent_format_tables import LATENT_FORMATS + +# arch -> preview table. Archs sharing a VAE family share a table; anything +# not listed falls back to a channel-count guess in `guess_format`. +ARCH_LATENT_FORMAT = { + "sd1": "SD15", + "sd2": "SD15", + "sdxl": "SDXL", + "ssd": "SDXL", + "vega": "SDXL", + "sd3": "SD3", + "flux": "Flux", + "flux_kontext": "Flux", + "chroma": "Flux", + "chroma_radiance": None, # pixel space + "zeta_chroma": "Flux", + "zimage": "Flux", + "zimage_l2p": "Flux", + "krea2": "Wan21", # Qwen-Image VAE + "hidream": "Flux", + "hidream_e1": "Flux", + "hidream_o1": "Flux", + "omnigen2": "Flux", + "f-lite": "Flux", + "nucleus_image": "Wan21", # Qwen-Image VAE + "qwen_image": "Wan21", + "qwen_image_edit": "Wan21", + "qwen_image_edit_plus": "Wan21", + "boogu_image": "Wan21", + "boogu_image_edit": "Wan21", + "flux2": "Flux2", + "flux2_klein_4b": "Flux2", + "flux2_klein_9b": "Flux2", + "ideogram4": "Flux2", + "mageflow": "Wan21", + "mageflow_edit": "Wan21", + "wan21": "Wan21", + "wan21_i2v": "Wan21", + "wan22_5b": "Wan22", + "wan22_14b": "Wan21", + "wan22_14b_i2v": "Wan21", + "ltx2": "LTXAV", + "ltx2.3": "LTXAV", + "ltx2.5": "LTXAV", + "minimax_h3": "MiniMaxH3Video", + "minimax_h3_ref2va": "MiniMaxH3Video", + "minimax_h3_vsa": "MiniMaxH3Video", + "anima": "Wan21", + "cogview4": "Flux", + "prx_pixel": None, # pixel space +} + + +def guess_format(channels: int, dims: int) -> Optional[str]: + if dims == 2: + return {4: "SD15", 16: "Flux", 128: "Flux2"}.get(channels) + if dims == 3: + return {16: "Wan21", 48: "Wan22", 128: "LTXAV", 24: "MiniMaxH3Video"}.get(channels) + return None + + +def preview_info(arch: str, channels: Optional[int] = None, dims: Optional[int] = None) -> Optional[dict]: + """The `start` frame's preview block: layout + factors, or None when the + latent cannot be previewed linearly (pixel-space models, audio).""" + # UI arch names may carry a variant suffix ("zimage:turbo", "wan21:1b") + arch = (arch or "").split(":")[0] + name = ARCH_LATENT_FORMAT.get(arch, "__unset__") + if name is None: + return None + if name == "__unset__" or (channels is not None and LATENT_FORMATS[name]["channels"] != channels): + name = guess_format(channels, dims) if channels is not None else None + if name is None: + return None + fmt = LATENT_FORMATS[name] + return { + "format": name, + "channels": fmt["channels"], + "dims": fmt["dims"], + "spatial": fmt["spatial"], + "temporal": fmt["temporal"], + "reshape": fmt["reshape"], + "factors": fmt["factors"], + "bias": fmt["bias"], + } diff --git a/extensions_built_in/inference_engine/protocol.py b/extensions_built_in/inference_engine/protocol.py new file mode 100644 index 0000000..9142a28 --- /dev/null +++ b/extensions_built_in/inference_engine/protocol.py @@ -0,0 +1,96 @@ +"""Binary stream framing for engine responses. + +One generation request is answered with a chunked HTTP body made of frames: + + u32 header_len (little endian) | json header (utf-8) | + u64 payload_len (little endian) | payload bytes + +Every frame header has "type" and "request_id". Text-only frames carry an +empty payload. The javascript reader (ui/src/lib/engineStream.ts) mirrors +this exactly. Frame types: + + start resolved model/sample config, output kind, latent preview info + status a holder stage line (loading / quantizing / encoding ...) + progress {step, total, elapsed} + latent raw tensor payload; header carries shape, dtype, layout + result {path, url, kind, ext, ...} for each produced file + error {message, cancelled} + end terminal frame +""" + +import json +import struct +from typing import Iterator, Optional, Tuple + +import numpy as np + +_HDR = struct.Struct(" bytes: + hdr = json.dumps(header, separators=(",", ":")).encode("utf-8") + return _HDR.pack(len(hdr)) + hdr + _LEN.pack(len(payload)) + payload + + +def tensor_payload(tensor, dtype="float16") -> Tuple[dict, bytes]: + """Serialize a tensor to raw bytes + the header fields that describe it.""" + import torch + + t = tensor.detach() + if t.is_cuda or t.device.type != "cpu": + t = t.to("cpu") + if dtype == "float16": + t = t.to(torch.float16) + elif dtype == "float32": + t = t.to(torch.float32) + arr = t.contiguous().numpy() + return {"shape": list(arr.shape), "dtype": str(arr.dtype)}, arr.tobytes() + + +def payload_to_array(header: dict, payload: bytes) -> np.ndarray: + return np.frombuffer(payload, dtype=np.dtype(header["dtype"])).reshape(header["shape"]) + + +class FrameReader: + """Decode frames from an iterator of byte chunks (tests / python clients).""" + + def __init__(self, chunks: Iterator[bytes]): + self._chunks = chunks + self._buf = bytearray() + self._eof = False + + def _fill(self, n: int) -> bool: + while len(self._buf) < n and not self._eof: + try: + chunk = next(self._chunks) + except StopIteration: + self._eof = True + break + if chunk: + self._buf.extend(chunk) + return len(self._buf) >= n + + def _take(self, n: int) -> bytes: + out = bytes(self._buf[:n]) + del self._buf[:n] + return out + + def read(self) -> Optional[Tuple[dict, bytes]]: + if not self._fill(_HDR.size): + return None + (hlen,) = _HDR.unpack(self._take(_HDR.size)) + if not self._fill(hlen + _LEN.size): + raise EOFError("truncated frame header") + header = json.loads(self._take(hlen).decode("utf-8")) + (plen,) = _LEN.unpack(self._take(_LEN.size)) + if not self._fill(plen): + raise EOFError("truncated frame payload") + return header, self._take(plen) + + def __iter__(self): + while True: + frame = self.read() + if frame is None: + return + yield frame diff --git a/extensions_built_in/inference_engine/server.py b/extensions_built_in/inference_engine/server.py new file mode 100644 index 0000000..b1c4106 --- /dev/null +++ b/extensions_built_in/inference_engine/server.py @@ -0,0 +1,152 @@ +"""HTTP surface of the engine (FastAPI). Runs on its own thread; every call +hands work to the Engine and returns, except /generate which streams the +job's frames for as long as the generation runs.""" + +import asyncio +import os +import re +import socket +import threading +import time +import uuid +from typing import Optional, Tuple + +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import FileResponse, JSONResponse, StreamingResponse + +from toolkit.models.registry import describe_archs + +from .engine import END, Engine + +_SAFE_NAME = re.compile(r"[^A-Za-z0-9._-]+") + + +def create_app(engine: Engine, token: Optional[str] = None) -> FastAPI: + app = FastAPI(title="AI Toolkit Inference Engine", docs_url=None, redoc_url=None) + + def check_auth(request: Request): + if not token: + return + supplied = request.headers.get("x-engine-token") + if supplied is None: + auth = request.headers.get("authorization", "") + if auth.lower().startswith("bearer "): + supplied = auth[7:] + if supplied != token: + raise HTTPException(status_code=401, detail="bad engine token") + + @app.get("/health") + async def health(): + return engine.health() + + @app.get("/models") + async def models(request: Request): + check_auth(request) + return {"archs": await asyncio.to_thread(describe_archs)} + + @app.get("/queue") + async def queue_(request: Request): + check_auth(request) + return {"queued": engine.queued(), "recent": engine.recent()} + + @app.post("/cancel/{request_id}") + async def cancel(request: Request, request_id: str): + check_auth(request) + return {"ok": engine.cancel(request_id)} + + @app.post("/unload") + async def unload(request: Request): + check_auth(request) + if engine.current is not None: + raise HTTPException(status_code=409, detail="busy") + await asyncio.to_thread(engine.unload) + return {"ok": True} + + @app.post("/assets") + async def upload_asset(request: Request, name: str = "asset.png"): + """Raw-body upload (no multipart dependency). Returns the stored path, + usable as ctrl_img in a generate request.""" + check_auth(request) + body = await request.body() + if not body: + raise HTTPException(status_code=400, detail="empty body") + base = _SAFE_NAME.sub("_", os.path.basename(name)) or "asset" + stem, ext = os.path.splitext(base) + asset_id = f"{stem}_{uuid.uuid4().hex[:8]}{ext or '.bin'}" + path = os.path.join(engine.assets_folder, asset_id) + with open(path, "wb") as f: + f.write(body) + return {"id": asset_id, "path": path, "bytes": len(body)} + + @app.get("/outputs/{relpath:path}") + async def get_output(request: Request, relpath: str): + check_auth(request) + root = os.path.realpath(engine.output_folder) + path = os.path.realpath(os.path.join(root, relpath)) + if not path.startswith(root + os.sep) or not os.path.isfile(path): + raise HTTPException(status_code=404, detail="not found") + return FileResponse(path) + + @app.post("/generate") + async def generate(request: Request): + check_auth(request) + body = await request.json() + model = body.get("model") + sample = body.get("sample") or {} + if not isinstance(model, dict) or not model.get("arch"): + raise HTTPException(status_code=400, detail="model.arch is required") + if body.get("wait"): + # plain JSON mode for scripts: block until the job is done, drop frames + job = engine.submit(model, sample, {"latents": "none"}) + while True: + frame = await asyncio.to_thread(job.frames.get) + if frame is END: + break + return JSONResponse(job.info(), status_code=200 if job.status == "done" else 500) + + job = engine.submit(model, sample, body.get("stream")) + + async def frames(): + try: + while True: + frame = await asyncio.to_thread(job.frames.get) + if frame is END: + break + yield frame + finally: + if job.status in ("queued", "running"): + engine.cancel(job.request_id) + + return StreamingResponse( + frames(), + media_type="application/octet-stream", + headers={"X-Request-Id": job.request_id, "Cache-Control": "no-store"}, + ) + + return app + + +def serve_in_thread(app, host: str = "127.0.0.1", port: int = 0, timeout: float = 30.0) -> Tuple[object, threading.Thread, str, int]: + """Bind the socket ourselves (port 0 = ephemeral) and run uvicorn on a + daemon thread. Returns (server, thread, host, port); the caller stops it + with server.should_exit = True. uvicorn only installs signal handlers on + the main thread, so the process keeps its own.""" + import uvicorn + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind((host, port)) + sock.listen(64) + host, port = sock.getsockname()[:2] + config = uvicorn.Config(app, log_level="warning", access_log=False, timeout_keep_alive=600) + server = uvicorn.Server(config) + thread = threading.Thread(target=server.run, kwargs={"sockets": [sock]}, daemon=True, name="engine-http") + thread.start() + deadline = time.time() + timeout + while not server.started and time.time() < deadline: + if not thread.is_alive(): + raise RuntimeError("engine http server died on startup") + time.sleep(0.05) + if not server.started: + raise RuntimeError("engine http server failed to start") + return server, thread, host, port diff --git a/requirements_base.txt b/requirements_base.txt index f9bb7fb..280a2f0 100644 --- a/requirements_base.txt +++ b/requirements_base.txt @@ -31,6 +31,9 @@ sentencepiece huggingface_hub==1.23.0 peft==0.18.1 gradio +# inference engine http server (extensions_built_in/inference_engine) +fastapi>=0.115 +uvicorn>=0.30 python-slugify opencv-python pytorch-wavelets==1.3.0 diff --git a/testing/test_inference_engine.py b/testing/test_inference_engine.py new file mode 100644 index 0000000..8e6a02e --- /dev/null +++ b/testing/test_inference_engine.py @@ -0,0 +1,384 @@ +"""Inference engine end-to-end test (extensions_built_in/inference_engine). + +Starts the engine in-process on an ephemeral loopback port, sends real +generate requests over HTTP, decodes the streamed frames and asserts on +them. Model/sample settings come from testing/test_model_loading.py's +MODEL_TESTS (tiny sizes: this proves the path, not quality). + +Usage: + python testing/test_inference_engine.py --arch zimage # one arch + python testing/test_inference_engine.py --swap zimage,wan21 # sequential + # model switches + python testing/test_inference_engine.py --arch zimage --cancel # cancel mid-run + python testing/test_inference_engine.py --all # every arch, + # subprocess each + --allow-download permit hub downloads (default HF_HUB_OFFLINE=1 -> SKIP) + --device cuda:0 +""" + +import argparse +import http.client +import json +import os +import shutil +import subprocess +import sys +import time + +TOOLKIT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, TOOLKIT_ROOT) + +from dotenv import load_dotenv + +load_dotenv(os.path.join(TOOLKIT_ROOT, ".env")) + +OUTPUT_ROOT = os.path.join(TOOLKIT_ROOT, "testing", ".engine_test_outputs") + +from testing.test_model_loading import MODEL_TESTS, SKIP_MARKERS, classify_error # noqa: E402 + + +class Client: + def __init__(self, host, port, token): + self.host, self.port, self.token = host, port, token + + def _conn(self): + return http.client.HTTPConnection(self.host, self.port, timeout=3600) + + def _headers(self): + return {"Content-Type": "application/json", "X-Engine-Token": self.token} + + def get(self, path): + c = self._conn() + c.request("GET", path, headers=self._headers()) + r = c.getresponse() + return r.status, json.loads(r.read()) + + def post_json(self, path, body=None): + c = self._conn() + c.request("POST", path, body=json.dumps(body or {}), headers=self._headers()) + r = c.getresponse() + return r.status, json.loads(r.read()) + + def upload(self, path, name): + c = self._conn() + with open(path, "rb") as f: + data = f.read() + c.request("POST", f"/assets?name={name}", body=data, headers={"X-Engine-Token": self.token}) + r = c.getresponse() + return r.status, json.loads(r.read()) + + def generate_stream(self, body): + """Yields (header, payload) frames as they arrive.""" + from extensions_built_in.inference_engine.protocol import FrameReader + + c = self._conn() + c.request("POST", "/generate", body=json.dumps(body), headers=self._headers()) + r = c.getresponse() + if r.status != 200: + raise RuntimeError(f"/generate -> {r.status}: {r.read()[:400]}") + + def chunks(): + while True: + data = r.read1(65536) # read1 returns the available chunk instead of blocking for 64KB + if not data: + return + yield data + + reader = FrameReader(chunks()) + for frame in reader: + yield frame + c.close() + + +def start_engine(device: str, out_root: str): + from extensions_built_in.inference_engine.engine import Engine + from extensions_built_in.inference_engine.server import create_app, serve_in_thread + + import threading + + engine = Engine(device=device, output_folder=os.path.join(out_root, "outputs"), assets_folder=os.path.join(out_root, "assets")) + token = "test-token" + server, thread, host, port = serve_in_thread(create_app(engine, token)) + loop = threading.Thread(target=engine.run_forever, daemon=True, name="engine-loop") + loop.start() + return engine, server, Client(host, port, token) + + +def stop_engine(engine, server): + engine.stop() + server.should_exit = True + engine.unload() + + +QUANTIZE_ALL = False # --quantize: force convrot8 on transformer + TE (small GPUs) + + +def request_body(arch: str, entry: dict, client: Client, out_dir: str, tiny=True, stream=None): + model = dict(entry["model"], arch=arch) + if QUANTIZE_ALL: + model["quantize"] = True + model["quantize_te"] = True + if model.get("quantize"): + model.setdefault("qtype", "convrot8") + if model.get("quantize_te"): + model.setdefault("qtype_te", "convrot8") + sample = dict(entry["sample"]) + sample["prompt"] = "a photo of a cat sitting on a wooden table" + if tiny and not entry.get("size_locked"): + # 384 divides every bucket size (16/32/64); video runs also keep the + # frame count small so the swap tests fit next to a resident TE + cap = 384 if tiny == "tiny" else 512 + sample["width"] = min(sample["width"], cap) + sample["height"] = min(sample["height"], cap) + if tiny == "tiny" and "num_frames" in sample: + sample["num_frames"] = min(sample["num_frames"], 9) + if entry.get("needs_control_image"): + from PIL import Image + + ctrl_path = os.path.join(out_dir, "ctrl.png") + Image.new("RGB", (sample["width"], sample["height"]), (128, 128, 128)).save(ctrl_path) + status, asset = client.upload(ctrl_path, "ctrl.png") + assert status == 200, asset + sample["ctrl_img"] = asset["path"] + return {"model": model, "sample": sample, "stream": stream or {"latents": "raw", "every_n_steps": 1, "max_frames": 2}} + + +def run_request(client: Client, body: dict, cancel_after: int = 0) -> dict: + """Consume one streamed generation; returns a summary of the frames.""" + import numpy as np + + from extensions_built_in.inference_engine.protocol import payload_to_array + + seen = {"start": 0, "status": 0, "progress": 0, "latent": 0, "result": 0, "error": 0, "end": 0} + summary = {"frames": seen, "results": [], "latent_shapes": [], "error": None, "cancelled": False, "request_id": None, "preview": None} + for header, payload in client.generate_stream(body): + t = header["type"] + seen[t] = seen.get(t, 0) + 1 + summary["request_id"] = header.get("request_id") + if t == "start": + summary["preview"] = header.get("preview") + summary["modality"] = header.get("modality") + elif t == "status": + print(f" [status] {header.get('message')}") + elif t == "progress": + print(f" [progress] {header['step']}/{header['total']} {header['elapsed']}s", flush=True) + if cancel_after and header["step"] >= cancel_after: + st, res = client.post_json(f"/cancel/{summary['request_id']}") + assert st == 200 and res["ok"], res + cancel_after = 0 + elif t == "latent": + arr = payload_to_array(header, payload) + assert np.isfinite(arr.astype(np.float32)).all(), "non-finite latent" + summary["latent_shapes"].append((header["layout"], list(arr.shape))) + elif t == "result": + summary["results"].append(header) + print(f" [result] {header['kind']} {header['path']} ({header['bytes']} bytes)") + elif t == "error": + summary["error"] = header.get("message") + summary["cancelled"] = bool(header.get("cancelled")) + if header.get("traceback"): + print(header["traceback"]) + elif t == "end": + summary["status"] = header.get("status") + return summary + + +def assert_success(summary: dict, modality: str): + assert summary["frames"]["start"] == 1, summary + assert summary["frames"]["end"] == 1, summary + assert summary["error"] is None, summary["error"] + assert summary["frames"]["progress"] >= 1, "no progress frames" + assert summary["frames"]["latent"] >= 1, "no latent frames" + assert summary["results"], "no result frames" + for r in summary["results"]: + assert os.path.isfile(r["path"]) and r["bytes"] > 1024, r + kinds = {r["kind"] for r in summary["results"]} + assert modality in kinds or (modality == "video" and "image" in kinds), (modality, kinds) + + +def run_one(arch: str, device: str, cancel: bool = False) -> dict: + from toolkit.models.registry import get_arch_entry + + entry = MODEL_TESTS[arch] + out_root = os.path.join(OUTPUT_ROOT, arch.replace("/", "_")) + shutil.rmtree(out_root, ignore_errors=True) + os.makedirs(out_root, exist_ok=True) + engine, server, client = start_engine(device, out_root) + try: + st, health = client.get("/health") + assert st == 200 and health["ok"], health + st, models = client.get("/models") + assert st == 200 and any(m["arch"] == arch for m in models["archs"]), f"{arch} missing from /models" + modality = get_arch_entry(arch)["modality"] + body = request_body(arch, entry, client, out_root) + + t0 = time.perf_counter() + if cancel: + s = run_request(client, body, cancel_after=2) + assert s["cancelled"], f"expected a cancelled stream, got {s}" + assert s["frames"]["end"] == 1 + print(" cancel: ok, re-running the same request") + s = run_request(client, body) + gen_seconds = time.perf_counter() - t0 + assert_success(s, modality) + st, health = client.get("/health") + assert health["active"]["arch"] == arch, health + return { + "arch": arch, + "status": "PASS", + "seconds": round(gen_seconds, 2), + "frames": s["frames"], + "latent_shapes": s["latent_shapes"][:2], + "results": [r["path"] for r in s["results"]], + "preview_format": (s["preview"] or {}).get("format"), + "stats": engine.stats, + } + finally: + stop_engine(engine, server) + + +def run_swap(archs, device: str) -> dict: + out_root = os.path.join(OUTPUT_ROOT, "swap_" + "_".join(a.replace("/", "_") for a in archs)) + shutil.rmtree(out_root, ignore_errors=True) + os.makedirs(out_root, exist_ok=True) + engine, server, client = start_engine(device, out_root) + from toolkit.models.registry import get_arch_entry + + timeline = [] + try: + for i, arch in enumerate(archs): + body = request_body(arch, MODEL_TESTS[arch], client, out_root, tiny="tiny") + t0 = time.perf_counter() + s = run_request(client, body) + assert_success(s, get_arch_entry(arch)["modality"]) + pool = engine.pool.stats() + resident = [(e["cls"], round(e["bytes"] / 1e9, 2), e["device"]) for e in pool["resident"]] + entry = { + "arch": arch, + "seconds": round(time.perf_counter() - t0, 2), + "holder_loads": engine.stats["holder_loads"], + "oom_retries": engine.stats["oom_retries"], + "pool_hits": pool["hits"], + "pool_misses": pool["misses"], + "pool_evictions": pool["evictions"], + "resident_gb": round(pool["resident_bytes"] / 1e9, 2), + "resident": resident, + } + timeline.append(entry) + print(f" {arch}: {json.dumps(entry)}") + # after a request only this holder's components may stay resident + assert pool["untouched"] == 0, f"stale components left resident: {pool}" + if i > 0 and archs[i - 1] == arch: + assert entry["holder_loads"] == timeline[i - 1]["holder_loads"], "same model should reuse the holder" + if len(set(archs)) > 1: + assert engine.pool.hits > 0, "expected shared components between the swapped archs" + return {"arch": ",".join(archs), "status": "PASS", "timeline": timeline} + finally: + stop_engine(engine, server) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--arch", type=str, default=None) + parser.add_argument("--swap", type=str, default=None, help="comma-separated archs to run back to back") + parser.add_argument("--cancel", action="store_true") + parser.add_argument("--all", action="store_true") + parser.add_argument("--list", action="store_true") + parser.add_argument("--device", type=str, default="cuda:0") + parser.add_argument("--allow-download", action="store_true") + parser.add_argument("--json-result", type=str, default=None) + parser.add_argument("--quantize", action="store_true", help="force convrot8 quantization of transformer + TE") + args = parser.parse_args() + global QUANTIZE_ALL + QUANTIZE_ALL = args.quantize + + if args.list: + for arch in MODEL_TESTS: + print(arch) + return + + os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") + if not args.allow_download: + os.environ.setdefault("HF_HUB_OFFLINE", "1") + + def _record(result): + if args.json_result: + with open(args.json_result, "w") as f: + json.dump(result, f) + print(f"[{result['status']}] {result['arch']}" + (f" - {result.get('error', '')}" if result["status"] != "PASS" else f" {json.dumps({k: v for k, v in result.items() if k not in ('arch', 'status')})}")) + if result["status"] == "FAIL": + sys.exit(1) + + if args.swap: + archs = [a.strip() for a in args.swap.split(",") if a.strip()] + try: + result = run_swap(archs, args.device) + except BaseException as err: + import traceback + + traceback.print_exc() + result = {"arch": args.swap, "status": classify_error(err), "error": f"{type(err).__name__}: {err}"} + _record(result) + return + + if args.arch: + if args.arch not in MODEL_TESTS: + raise SystemExit(f"arch {args.arch!r} is not registered; --list shows options") + try: + result = run_one(args.arch, args.device, cancel=args.cancel) + except BaseException as err: + status = classify_error(err) + if status == "FAIL": + import traceback + + traceback.print_exc() + result = {"arch": args.arch, "status": status, "error": f"{type(err).__name__}: {err}"} + _record(result) + return + + if not args.all: + parser.print_help() + return + + results = [] + total = len(MODEL_TESTS) + t0 = time.perf_counter() + for i, arch in enumerate(MODEL_TESTS, start=1): + print(f"\n===== {arch} ({i}/{total}) =====", flush=True) + result_path = os.path.join(OUTPUT_ROOT, f".{arch.replace('/', '_')}.result.json") + os.makedirs(OUTPUT_ROOT, exist_ok=True) + cmd = [sys.executable, os.path.abspath(__file__), "--arch", arch, "--device", args.device, "--json-result", result_path] + if args.allow_download: + cmd.append("--allow-download") + if args.quantize: + cmd.append("--quantize") + proc = subprocess.run(cmd, cwd=TOOLKIT_ROOT) + if os.path.exists(result_path): + with open(result_path) as f: + results.append(json.load(f)) + os.remove(result_path) + else: + results.append({"arch": arch, "status": "FAIL", "error": f"subprocess died (exit {proc.returncode})"}) + tally = {"PASS": 0, "FAIL": 0, "SKIP": 0} + for r in results: + tally[r["status"]] = tally.get(r["status"], 0) + 1 + elapsed = time.perf_counter() - t0 + print(f">>> progress: {i}/{total} - {tally['PASS']} pass, {tally['FAIL']} fail, {tally['SKIP']} skip - elapsed {elapsed / 60:.0f}m", flush=True) + + print("\n===== summary =====") + counts = {"PASS": 0, "FAIL": 0, "SKIP": 0} + for r in results: + counts[r["status"]] = counts.get(r["status"], 0) + 1 + line = f"[{r['status']}] {r['arch']}" + if r["status"] != "PASS": + line += f" - {r.get('error', '')[:160]}" + print(line) + print(f"\n{counts['PASS']} passed, {counts['FAIL']} failed, {counts['SKIP']} skipped") + with open(os.path.join(OUTPUT_ROOT, "report.json"), "w") as f: + json.dump(results, f, indent=2) + if counts["FAIL"]: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 98f7cb2..4f314f5 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -167,6 +167,9 @@ class BaseModel: self.invert_assistant_lora = False self._after_sample_img_hooks = [] self._status_update_hooks = [] + # inference engine: called as hook(step_index, num_steps, latents) + # after every scheduler step while generating samples + self.sample_step_hook = None self.is_transformer = False self.sample_prompts_cache = None @@ -401,6 +404,18 @@ class BaseModel: def add_status_update_hook(self, func): self._status_update_hooks.append(func) + def _emit_sample_step(self, latents, step_index=None, num_steps=None): + """For holders whose sampling loop bypasses scheduler.step: report one + denoised latent to sample_step_hook (no-op when unset).""" + from toolkit.sample_step_hook import emit_sample_step + + emit_sample_step(self, latents, step_index, num_steps) + + def _install_sample_step_hooks(self, pipeline): + from toolkit.sample_step_hook import install_sample_step_hooks + + return install_sample_step_hooks(self, pipeline) + @torch.no_grad() def generate_images( self, @@ -457,6 +472,8 @@ class BaseModel: except: pass + unwrap_step_hooks = self._install_sample_step_hooks(pipeline) + start_multiplier = 1.0 if network is not None: start_multiplier = network.multiplier @@ -517,6 +534,7 @@ class BaseModel: if network is not None: network.multiplier = gen_config.network_multiplier + self._sample_step_index = 0 torch.manual_seed(gen_config.seed) torch.cuda.manual_seed(gen_config.seed) @@ -725,6 +743,7 @@ class BaseModel: if self.adapter is not None and isinstance(self.adapter, ReferenceAdapter): self.adapter.clear_memory() + unwrap_step_hooks() # clear pipeline and cache to reduce vram usage del pipeline torch.cuda.empty_cache() diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py new file mode 100644 index 0000000..64b8fa8 --- /dev/null +++ b/toolkit/models/registry.py @@ -0,0 +1,85 @@ +"""Per-arch capabilities + realistic defaults, one table shared by the +inference engine (`GET /models`), its tests, and eventually the UI. + +Every arch registered through get_all_models()/LEGACY_ARCHS is served; an +entry here adds defaults and capability flags on top. Archs without an entry +get GENERIC. Keep `model` limited to ModelConfig kwargs and `sample` to +GenerateImageConfig kwargs. +""" + +from typing import Dict, List, Optional + +IMG = {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0} +VID = {"width": 832, "height": 480, "num_inference_steps": 20, "guidance_scale": 5.0, "num_frames": 33, "fps": 16} + +GENERIC = {"modality": "image", "model": {}, "sample": dict(IMG)} + +# modality: image | video | audio +# needs_control_image: the arch requires ctrl_img (edit / i2v models) +# size_locked: only the default resolution is valid +ARCH_REGISTRY: Dict[str, dict] = { + "zimage": {"modality": "image", "model": {"name_or_path": "Tongyi-MAI/Z-Image-Turbo"}, "sample": {**IMG, "num_inference_steps": 8, "guidance_scale": 1.0}}, + "zimage_l2p": {"modality": "image", "model": {"name_or_path": "zhen-nan/L2P/model-1k-merge.safetensors", "extras_name_or_path": "Tongyi-MAI/Z-Image-Turbo", "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 8, "guidance_scale": 1.0}}, + "qwen_image": {"modality": "image", "model": {"name_or_path": "Qwen/Qwen-Image", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 20}}, + "qwen_image_edit": {"modality": "image", "model": {"name_or_path": "Qwen/Qwen-Image-Edit", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 20}, "needs_control_image": True}, + "qwen_image_edit_plus": {"modality": "image", "model": {"name_or_path": "Qwen/Qwen-Image-Edit-2509", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 20}, "needs_control_image": True}, + "krea2": {"modality": "image", "model": {"name_or_path": "krea/Krea-2-Turbo", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 8, "guidance_scale": 1.0}}, + "boogu_image": {"modality": "image", "model": {"name_or_path": "Boogu/Boogu-Image-0.1-Base", "quantize": True, "quantize_te": True}, "sample": dict(IMG)}, + "boogu_image_edit": {"modality": "image", "model": {"name_or_path": "Boogu/Boogu-Image-0.1-Edit", "quantize": True, "quantize_te": True}, "sample": dict(IMG), "needs_control_image": True}, + "ernie_image": {"modality": "image", "model": {"name_or_path": "baidu/ERNIE-Image", "quantize": True, "quantize_te": True}, "sample": dict(IMG)}, + "mageflow": {"modality": "image", "model": {"name_or_path": "microsoft/Mage-Flow-Base", "quantize": True, "quantize_te": True}, "sample": dict(IMG)}, + "mageflow_edit": {"modality": "image", "model": {"name_or_path": "microsoft/Mage-Flow-Edit", "quantize": True, "quantize_te": True}, "sample": dict(IMG), "needs_control_image": True}, + "ideogram4": {"modality": "image", "model": {"name_or_path": "ideogram-ai/ideogram-4-fp8", "quantize": True, "quantize_te": True}, "sample": dict(IMG)}, + "hidream": {"modality": "image", "model": {"name_or_path": "HiDream-ai/HiDream-I1-Full", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 28, "guidance_scale": 5.0}}, + "hidream_e1": {"modality": "image", "model": {"name_or_path": "HiDream-ai/HiDream-E1-1", "quantize": True, "quantize_te": True}, "sample": {"width": 768, "height": 768, "num_inference_steps": 28, "guidance_scale": 5.0}, "needs_control_image": True, "size_locked": True}, + "hidream_o1": {"modality": "image", "model": {"name_or_path": "HiDream-ai/HiDream-O1-Image", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 28, "guidance_scale": 5.0}}, + "anima": {"modality": "image", "model": {"name_or_path": "circlestone-labs/Anima-Base-v1.0-Diffusers"}, "sample": {**IMG, "guidance_scale": 4.5}}, + "nucleus_image": {"modality": "image", "model": {"name_or_path": "NucleusAI/Nucleus-Image", "quantize": True, "quantize_te": True}, "sample": dict(IMG)}, + "omnigen2": {"modality": "image", "model": {"name_or_path": "OmniGen2/OmniGen2", "quantize": True, "quantize_te": True}, "sample": dict(IMG)}, + "chroma": {"modality": "image", "model": {"name_or_path": "lodestones/Chroma1-HD", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 26}}, + "chroma_radiance": {"modality": "image", "model": {"name_or_path": "lodestones/Chroma1-Radiance", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 26}}, + "zeta_chroma": {"modality": "image", "model": {"name_or_path": "lodestones/Zeta-Chroma/zeta-chroma-base-x0-pixel-dino-distance.safetensors", "extras_name_or_path": "Tongyi-MAI/Z-Image-Turbo", "quantize": True, "quantize_te": True}, "sample": dict(IMG)}, + "flux_kontext": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.1-Kontext-dev", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 2.5}, "needs_control_image": True}, + "flux2": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.2-dev", "quantize": True, "quantize_te": True}, "sample": dict(IMG)}, + "flux2_klein_4b": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-4B", "quantize_te": True}, "sample": dict(IMG)}, + "flux2_klein_9b": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-9B", "quantize": True, "quantize_te": True}, "sample": dict(IMG)}, + "prx_pixel": {"modality": "image", "model": {"name_or_path": "Photoroom/prxpixel-t2i", "quantize_te": True}, "sample": dict(IMG)}, + "f-lite": {"modality": "image", "model": {"name_or_path": "Freepik/F-Lite", "quantize": True, "quantize_te": True}, "sample": dict(IMG)}, + "sd1": {"modality": "image", "model": {"name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 7.5}}, + "sdxl": {"modality": "image", "model": {"name_or_path": "stabilityai/stable-diffusion-xl-base-1.0"}, "sample": {**IMG, "guidance_scale": 6.0}}, + "wan21": {"modality": "video", "model": {"name_or_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"}, "sample": dict(VID)}, + "wan21_i2v": {"modality": "video", "model": {"name_or_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", "quantize": True, "quantize_te": True}, "sample": dict(VID), "needs_control_image": True}, + "wan22_5b": {"modality": "video", "model": {"name_or_path": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "quantize": True, "quantize_te": True}, "sample": {**VID, "fps": 24}}, + "wan22_14b": {"modality": "video", "model": {"name_or_path": "ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True}, "sample": {**VID, "guidance_scale": 3.5}}, + "wan22_14b_i2v": {"modality": "video", "model": {"name_or_path": "ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True}, "sample": {**VID, "guidance_scale": 3.5}, "needs_control_image": True}, + "ltx2.3": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}}, + "ltx2.5": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}}, + "ace_step_15": {"modality": "audio", "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0}}, +} + +OUTPUT_EXT = {"image": "png", "video": "mp4", "audio": "mp3"} + + +def get_arch_entry(arch: str) -> dict: + # variant suffixes ("zimage:turbo") share the base arch's entry + entry = ARCH_REGISTRY.get(arch) or ARCH_REGISTRY.get(arch.split(":")[0], GENERIC) + return { + "arch": arch, + "modality": entry.get("modality", "image"), + "model": dict(entry.get("model", {})), + "sample": dict(entry.get("sample", IMG)), + "needs_control_image": bool(entry.get("needs_control_image", False)), + "size_locked": bool(entry.get("size_locked", False)), + } + + +def list_archs() -> List[str]: + from toolkit.util.get_model import LEGACY_ARCHS, get_all_models + + archs = {m.arch for m in get_all_models() if getattr(m, "arch", None)} + archs |= LEGACY_ARCHS + return sorted(archs) + + +def describe_archs(archs: Optional[List[str]] = None) -> List[dict]: + return [get_arch_entry(a) for a in (archs or list_archs())] diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index 6be23f0..dc16118 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -255,6 +255,31 @@ class OstrisModelMixin: base_model.print_and_status_update if base_model is not None else print_acc ) + # live server: a pooled module already carrying this exact policy is + # done; a different policy re-applies on top (requantize / dequantize + # paths below handle the weights, offloading is detached first) + from toolkit.models.v2.pool import ComponentPool, _stable + + policy = _stable( + dict( + qtype=qtype, + offload=offload, + dtype=dtype, + device=device, + quantize_device=quantize_device, + exclude=exclude_quant_modules, + ) + ) + if ComponentPool.current is not None: + if getattr(self, "_aitk_policy", None) == policy: + ComponentPool.current.touch_module(self) + return self + if hasattr(self, "_memory_manager"): + from toolkit.memory_management import MemoryManager + + MemoryManager.detach(self) + self._aitk_policy = policy + ara_path = None if qtype is not None and "|" in qtype: qtype, ara_path = qtype.split("|", 1) @@ -428,6 +453,27 @@ class OstrisModelMixin: elif subfolder == "": subfolder = None + # live server: an identical source request returns the resident module + from toolkit.models.v2.pool import ComponentPool + + pool = ComponentPool.current + pool_key = None + if pool is not None: + pool_key = ComponentPool.make_key( + cls, + name_or_path, + dtype=dtype, + config_path=config_path, + config=config, + subfolder=subfolder, + use_comfy_weights=use_comfy_weights, + qtype=qtype if quantize_on_load else None, + kwargs=kwargs, + ) + resident = pool.get(pool_key) + if resident is not None: + return resident + if ( use_comfy_weights and not name_or_path.endswith(".safetensors") @@ -473,6 +519,8 @@ class OstrisModelMixin: if device is not None: model.to(device) + if pool is not None and pool_key is not None: + pool.put(pool_key, model) return model @staticmethod diff --git a/toolkit/models/v2/pool.py b/toolkit/models/v2/pool.py new file mode 100644 index 0000000..f31379f --- /dev/null +++ b/toolkit/models/v2/pool.py @@ -0,0 +1,194 @@ +"""Process-wide component pool for the live inference server. + +While a pool is active (`ComponentPool.current`), `OstrisModelMixin.load_model` +returns an already-resident module for an identical source request instead +of loading it again, and `aitk_post_load` becomes a no-op when the module +already carries the requested policy (qtype / offload / device). Holders +keep calling the normal load API; component reuse across model switches +(same TE or VAE under two archs) falls out of that. + +Lifecycle per request: `begin_request()` -> holder loads (hits mark entries +as touched) -> `release_unused()` frees every entry the request did not +touch. `clear()` drops everything. +""" + +import gc +import json +import threading +import time +from typing import Dict, Optional + +import torch + + +def _stable(value): + """JSON-safe, deterministic rendering of a key part.""" + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, (torch.dtype, torch.device)): + return str(value) + if isinstance(value, dict): + return {str(k): _stable(v) for k, v in sorted(value.items(), key=lambda kv: str(kv[0]))} + if isinstance(value, (list, tuple)): + return [_stable(v) for v in value] + to_dict = getattr(value, "to_dict", None) + if callable(to_dict): + try: + return _stable(to_dict()) + except Exception: + pass + if hasattr(value, "__dataclass_fields__"): + return _stable({k: getattr(value, k) for k in value.__dataclass_fields__}) + return repr(value) + + +def module_bytes(module: torch.nn.Module) -> int: + total = 0 + seen = set() + for t in list(module.parameters()) + list(module.buffers()): + try: + ptr = t.untyped_storage().data_ptr() + except Exception: + ptr = id(t) + if ptr in seen or t.device.type == "meta": + continue + seen.add(ptr) + total += t.numel() * t.element_size() + return total + + +class PoolEntry: + def __init__(self, key: str, module: torch.nn.Module, generation: int): + self.key = key + self.module = module + self.cls_name = type(module).__name__ + self.created_at = time.time() + self.last_used = self.created_at + self.uses = 0 + self.touched_gen = generation + self._bytes: Optional[int] = None + + @property + def bytes(self) -> int: + if self._bytes is None: + try: + self._bytes = module_bytes(self.module) + except Exception: + self._bytes = 0 + return self._bytes + + def device(self) -> str: + try: + p = next(self.module.parameters()) + return str(p.device) + except StopIteration: + return "none" + + def info(self) -> dict: + try: + parsed = json.loads(self.key) + except Exception: + parsed = {"key": self.key} + return { + "cls": self.cls_name, + "src": parsed.get("src"), + "subfolder": parsed.get("subfolder"), + "dtype": parsed.get("dtype"), + "policy": getattr(self.module, "_aitk_policy", None), + "bytes": self.bytes, + "device": self.device(), + "uses": self.uses, + "last_used": self.last_used, + } + + +class ComponentPool: + current: Optional["ComponentPool"] = None + + def __init__(self): + self.entries: Dict[str, PoolEntry] = {} + self.generation = 0 + self.hits = 0 + self.misses = 0 + self.evictions = 0 + self._lock = threading.RLock() + + # ---- keys ---- + @staticmethod + def make_key(cls, name_or_path: str, **parts) -> str: + key = {"cls": f"{cls.__module__}.{cls.__qualname__}", "src": name_or_path} + key.update({k: _stable(v) for k, v in parts.items()}) + return json.dumps(key, sort_keys=True, default=repr) + + # ---- request lifecycle ---- + def begin_request(self): + with self._lock: + self.generation += 1 + + def get(self, key: str) -> Optional[torch.nn.Module]: + with self._lock: + entry = self.entries.get(key) + if entry is None: + self.misses += 1 + return None + self.hits += 1 + entry.uses += 1 + entry.last_used = time.time() + entry.touched_gen = self.generation + return entry.module + + def put(self, key: str, module: torch.nn.Module): + with self._lock: + entry = PoolEntry(key, module, self.generation) + entry.uses = 1 + self.entries[key] = entry + + def touch_module(self, module: torch.nn.Module): + with self._lock: + for entry in self.entries.values(): + if entry.module is module: + entry.touched_gen = self.generation + entry.last_used = time.time() + + def release_unused(self) -> int: + """Free every entry not touched during the current request.""" + with self._lock: + stale = [e for e in self.entries.values() if e.touched_gen != self.generation] + freed = 0 + for entry in stale: + freed += entry.bytes + self._free(entry) + return freed + + def clear(self): + with self._lock: + for entry in list(self.entries.values()): + self._free(entry) + + def _free(self, entry: PoolEntry): + from toolkit.memory_management import MemoryManager + + self.entries.pop(entry.key, None) + self.evictions += 1 + try: + MemoryManager.free(entry.module) + except Exception: + try: + torch.nn.Module.to(entry.module, "meta") + except Exception: + pass + entry.module = None + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def stats(self) -> dict: + with self._lock: + return { + "hits": self.hits, + "misses": self.misses, + "evictions": self.evictions, + "resident": [e.info() for e in self.entries.values()], + "resident_bytes": sum(e.bytes for e in self.entries.values()), + "untouched": len([e for e in self.entries.values() if e.touched_gen != self.generation]), + } diff --git a/toolkit/sample_step_hook.py b/toolkit/sample_step_hook.py new file mode 100644 index 0000000..d618d24 --- /dev/null +++ b/toolkit/sample_step_hook.py @@ -0,0 +1,134 @@ +"""Per-step latent reporting for sample generation (inference engine). + +A holder exposes `sample_step_hook(step_index, num_steps, latents)`; while it +is set, the sampling scheduler's `step` is wrapped on the instance so every +step reaches the hook without per-pipeline code. The latent handed to the +hook is the step's predicted CLEAN sample (x0) — what a preview should show +— falling back to the scheduler's next sample when x0 cannot be derived. Shared by +BaseModel and the legacy StableDiffusion monolith. +""" + +from typing import Callable, List + + +def install_sample_step_hooks(holder, pipeline) -> Callable[[], None]: + """Wrap scheduler.step on the sampling scheduler(s). Returns an unwrap + callable. No-op without a hook. Instance-level wrap: the class method is + untouched and a stale wrapper is replaced rather than stacked.""" + if getattr(holder, "sample_step_hook", None) is None: + return lambda: None + schedulers: List = [] + for cand in ( + getattr(pipeline, "scheduler", None), + getattr(getattr(holder, "pipeline", None), "scheduler", None), + getattr(holder, "noise_scheduler", None), + ): + if cand is not None and hasattr(cand, "step") and all(cand is not s for s in schedulers): + schedulers.append(cand) + unwraps = [] + for scheduler in schedulers: + current = scheduler.step + orig = getattr(current, "_aitk_orig_step", None) or current + + def make(orig, scheduler): + def step(*args, **kwargs): + # capture the pre-step state: x0 is derived from (model_output, + # timestep, sample) with the sigma/alpha of THIS step, which + # the scheduler advances past inside orig() + model_output = args[0] if len(args) > 0 else kwargs.get("model_output") + timestep = args[1] if len(args) > 1 else kwargs.get("timestep") + sample_in = args[2] if len(args) > 2 else kwargs.get("sample") + coeffs = _x0_coefficients(scheduler, timestep) + out = orig(*args, **kwargs) + if isinstance(out, tuple): + prev_sample = out[0] + pred_x0 = None + else: + prev_sample = getattr(out, "prev_sample", None) + pred_x0 = getattr(out, "pred_original_sample", None) + hook = getattr(holder, "sample_step_hook", None) + if hook is not None and prev_sample is not None: + if pred_x0 is None: + pred_x0 = _predict_x0(coeffs, model_output, sample_in) + timesteps = getattr(scheduler, "timesteps", None) + num_steps = len(timesteps) if timesteps is not None else None + idx = getattr(holder, "_sample_step_index", 0) + holder._sample_step_index = idx + 1 + hook(idx, num_steps, pred_x0 if pred_x0 is not None else prev_sample) + return out + + step._aitk_orig_step = orig + return step + + scheduler.step = make(orig, scheduler) + + def unwrap(scheduler=scheduler): + if "step" in scheduler.__dict__: + del scheduler.__dict__["step"] + + unwraps.append(unwrap) + return lambda: [u() for u in unwraps] + + +def _x0_coefficients(scheduler, timestep): + """Per-step numbers needed to turn the model output into a clean-sample + estimate, read BEFORE scheduler.step advances its index. + Returns ("flow", sigma) | ("ddpm", prediction_type, alpha_bar) | None.""" + try: + import torch + + sigmas = getattr(scheduler, "sigmas", None) + if sigmas is not None and hasattr(scheduler, "index_for_timestep"): + idx = getattr(scheduler, "step_index", None) + if idx is None: + idx = scheduler.index_for_timestep(timestep) + return ("flow", float(sigmas[idx])) + alphas_cumprod = getattr(scheduler, "alphas_cumprod", None) + if alphas_cumprod is not None: + t = timestep + if isinstance(t, torch.Tensor): + t = int(t.flatten()[0].item()) + prediction_type = getattr(getattr(scheduler, "config", None), "prediction_type", "epsilon") + return ("ddpm", prediction_type, float(alphas_cumprod[int(t)])) + except Exception: + return None + return None + + +def _predict_x0(coeffs, model_output, sample): + """Clean-sample estimate for the preview: flow matching x0 = x_t - sigma*v; + DDPM-style eps / v / sample predictions via alpha_bar. None when unknown.""" + if coeffs is None or model_output is None or sample is None: + return None + try: + x = sample.tensor if hasattr(sample, "tensor") else sample + v = model_output.tensor if hasattr(model_output, "tensor") else model_output + if v.shape != x.shape: + return None + v = v.to(x.dtype) + if coeffs[0] == "flow": + sigma = coeffs[1] + return x - sigma * v + _, prediction_type, alpha_bar = coeffs + a = alpha_bar**0.5 + b = (1.0 - alpha_bar) ** 0.5 + if prediction_type == "epsilon": + return (x - b * v) / a + if prediction_type == "v_prediction": + return a * x - b * v + if prediction_type == "sample": + return v + except Exception: + return None + return None + + +def emit_sample_step(holder, latents, step_index=None, num_steps=None): + """For sampling loops that bypass scheduler.step: report one denoised + latent to the holder's hook (no-op when unset).""" + hook = getattr(holder, "sample_step_hook", None) + if hook is None: + return + idx = getattr(holder, "_sample_step_index", 0) if step_index is None else step_index + holder._sample_step_index = idx + 1 + hook(idx, num_steps, latents) diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index 52cb81b..3768578 100644 --- a/toolkit/stable_diffusion_model.py +++ b/toolkit/stable_diffusion_model.py @@ -191,6 +191,8 @@ class StableDiffusion: self.te_torch_dtype = get_torch_dtype(model_config.te_dtype) self.model_config = model_config + # inference engine: hook(step_index, num_steps, latents) per scheduler step + self.sample_step_hook = None self.prediction_type = "v_prediction" if self.model_config.is_v_pred else "epsilon" self.arch = model_config.arch @@ -1417,6 +1419,10 @@ class StableDiffusion: # disable progress bar pipeline.set_progress_bar_config(disable=True) + from toolkit.sample_step_hook import install_sample_step_hooks + + unwrap_step_hooks = install_sample_step_hooks(self, pipeline) + refiner_pipeline = None if self.refiner_unet: # build refiner pipeline @@ -1495,6 +1501,7 @@ class StableDiffusion: if network is not None: network.multiplier = gen_config.network_multiplier + self._sample_step_index = 0 torch.manual_seed(gen_config.seed) torch.cuda.manual_seed(gen_config.seed) @@ -1792,6 +1799,7 @@ class StableDiffusion: if self.adapter is not None and isinstance(self.adapter, ReferenceAdapter): self.adapter.clear_memory() + unwrap_step_hooks() # clear pipeline and cache to reduce vram usage del pipeline if refiner_pipeline is not None: diff --git a/toolkit/ui_job_status.py b/toolkit/ui_job_status.py new file mode 100644 index 0000000..5fefee3 --- /dev/null +++ b/toolkit/ui_job_status.py @@ -0,0 +1,130 @@ +"""UI job bookkeeping for long-running extension processes. + +When a process is launched by the UI it gets AITK_JOB_ID and a sqlite db +path; status / step / info are written to the Job row and the UI's stop +button sets Job.stop. This mirrors what BaseCaptioner does inline, as a +standalone helper so other UI-driven processes (the inference engine) can +share it. Outside the UI (no job id or no db file) every call is a no-op. +""" + +import os +import sqlite3 +import threading +import time +from typing import Callable, Literal, Optional + +AITK_Status = Literal["running", "stopped", "error", "completed", "queued"] + + +class UIJobStatus: + def __init__(self, sqlite_db_path: str = "./aitk_db.db"): + self.sqlite_db_path = sqlite_db_path + job_id = os.environ.get("AITK_JOB_ID", None) + self.job_id = job_id.strip() if job_id else None + self.is_ui_job = self.job_id is not None and os.path.exists(sqlite_db_path) + if self.is_ui_job: + print(f"Using SQLite database at {sqlite_db_path}") + print(f'Job ID: "{self.job_id}"') + self._lock = threading.Lock() + self._watcher = None + # collapse bursts of identical writes (progress ticks) + self._last_info = None + + # ---- db plumbing ---- + def _connect(self): + conn = sqlite3.connect(self.sqlite_db_path, timeout=30.0) + conn.isolation_level = None + return conn + + def _retry(self, fn, max_retries=3, base_delay=2.0): + last = None + for attempt in range(max_retries + 1): + try: + return fn() + except sqlite3.OperationalError as e: + if "database is locked" not in str(e): + raise + last = e + if attempt < max_retries: + time.sleep(base_delay * (2**attempt)) + raise last + + def _write(self, query: str, params: tuple): + if not self.is_ui_job: + return + + def _do(): + with self._lock, self._connect() as conn: + cur = conn.cursor() + cur.execute("BEGIN IMMEDIATE") + try: + cur.execute(query, params) + finally: + cur.execute("COMMIT") + + try: + self._retry(_do) + except Exception as e: + print(f"[AITK] Warning: db write failed: {e}") + + def _read_flag(self, column: str) -> bool: + if not self.is_ui_job: + return False + + def _do(): + with self._lock, self._connect() as conn: + cur = conn.cursor() + cur.execute(f"SELECT {column} FROM Job WHERE id = ?", (self.job_id,)) + row = cur.fetchone() + return False if row is None else row[0] == 1 + + return self._retry(_do) + + # ---- public api ---- + def update_status(self, status: AITK_Status, info: Optional[str] = None): + if info is None: + self._write("UPDATE Job SET status = ? WHERE id = ?", (status, self.job_id)) + else: + self._last_info = info + self._write( + "UPDATE Job SET status = ?, info = ? WHERE id = ?", + (status, info, self.job_id), + ) + + def update_info(self, info: str): + if info == self._last_info: + return + self._last_info = info + self._write("UPDATE Job SET info = ? WHERE id = ?", (info, self.job_id)) + + def update_key(self, key: str, value): + # key is a column name from our own code, never user input + self._write(f"UPDATE Job SET {key} = ? WHERE id = ?", (str(value), self.job_id)) + + def should_stop(self) -> bool: + return self._read_flag("stop") + + def should_return_to_queue(self) -> bool: + return self._read_flag("return_to_queue") + + def start_stop_watcher(self, on_stop: Callable[[str], None], interval_sec: float = 2.0): + """Poll the stop / return_to_queue flags on a daemon thread and call + on_stop("stopped" | "queued") once when one is set.""" + if not self.is_ui_job or self._watcher is not None: + return + + def _loop(): + while True: + try: + if self.should_stop(): + on_stop("stopped") + return + if self.should_return_to_queue(): + on_stop("queued") + return + except Exception: + pass + time.sleep(interval_sec) + + self._watcher = threading.Thread(target=_loop, daemon=True) + self._watcher.start() diff --git a/ui/cron/actions/startJob.ts b/ui/cron/actions/startJob.ts index f0f3534..4ab33a7 100644 --- a/ui/cron/actions/startJob.ts +++ b/ui/cron/actions/startJob.ts @@ -217,6 +217,14 @@ const startAndWatchJob = (job: Job) => { // update the config dataset path const jobConfig = JSON.parse(job.job_config); jobConfig.config.process[0].sqlite_db_path = path.join(TOOLKIT_ROOT, 'aitk_db.db'); + if (job.job_type === 'inference') { + // the engine publishes engine.json (endpoint + token) and writes its + // outputs under the job folder, which the UI proxy/file routes know + jobConfig.config.process[0].engine = { + ...(jobConfig.config.process[0].engine || {}), + job_folder: trainingFolder, + }; + } // write the config file fs.writeFileSync(configPath, JSON.stringify(jobConfig, null, 2)); diff --git a/ui/src/app/api/inference/[...path]/route.ts b/ui/src/app/api/inference/[...path]/route.ts new file mode 100644 index 0000000..4664fe4 --- /dev/null +++ b/ui/src/app/api/inference/[...path]/route.ts @@ -0,0 +1,125 @@ +import { NextRequest, NextResponse } from 'next/server'; +import fs from 'fs'; +import path from 'path'; +import prisma from '@/server/prisma'; +import { getTrainingFolder } from '@/server/settings'; +import { listEngines, resolveEngine } from '@/server/inferenceEngine'; + +// Reverse proxy to the resident inference engine (a python job that binds an +// ephemeral loopback port and publishes it in its job folder). Everything +// under /api/inference/* is forwarded verbatim; a streaming upstream body is +// passed through untouched, so /generate frames reach the browser as they +// are produced. `?job=` picks an engine when several run. +// +// /api/inference/status is answered here: it never touches the engine. +export const dynamic = 'force-dynamic'; +export const runtime = 'nodejs'; + +const HOP_BY_HOP = new Set(['connection', 'keep-alive', 'transfer-encoding', 'host', 'content-length', 'authorization']); + +async function proxy(request: NextRequest, segments: string[]) { + const upstreamPath = segments.join('/'); + const jobId = request.nextUrl.searchParams.get('job'); + + if (upstreamPath === 'status') { + const engines = await listEngines(true); + const active = jobId ? engines.find(e => e.jobId === jobId) || null : engines.find(e => e.endpoint) || engines[0] || null; + return NextResponse.json({ + running: !!active?.endpoint, + engine: active + ? { jobId: active.jobId, jobName: active.jobName, gpuIds: active.gpuIds, status: active.status, ready: !!active.endpoint } + : null, + engines: engines.map(e => ({ jobId: e.jobId, jobName: e.jobName, gpuIds: e.gpuIds, status: e.status, ready: !!e.endpoint })), + }); + } + + if (upstreamPath === 'outputs/delete' && request.method === 'POST') { + // delete generated files from disk. Handled here (not by the engine) so + // history can be cleaned up while no engine is running. Only files under + // an inference job's outputs folder are touched. + const body = await request.json().catch(() => ({})); + const paths: string[] = Array.isArray(body?.paths) ? body.paths : []; + const trainingFolder = path.resolve(await getTrainingFolder()); + const jobs = await prisma.job.findMany({ where: { job_type: 'inference' }, select: { name: true } }); + const roots = jobs.map(j => path.join(trainingFolder, j.name, 'outputs')); + const deleted: string[] = []; + const refused: string[] = []; + for (const p of paths) { + const resolved = path.resolve(String(p)); + const root = roots.find(r => resolved.startsWith(r + path.sep)); + if (!root) { + refused.push(p); + continue; + } + try { + await fs.promises.rm(resolved, { force: true }); + deleted.push(p); + // request folders hold one generation each; drop the folder once empty + const dir = path.dirname(resolved); + if (dir !== root && (await fs.promises.readdir(dir)).length === 0) { + await fs.promises.rmdir(dir); + } + } catch (e: any) { + refused.push(`${p}: ${e?.message || e}`); + } + } + return NextResponse.json({ deleted, refused }); + } + + const engine = await resolveEngine(jobId); + if (!engine || !engine.endpoint) { + return NextResponse.json({ error: 'no inference engine is running', engine: engine ? { status: engine.status } : null }, { status: 503 }); + } + + const search = new URLSearchParams(request.nextUrl.searchParams); + search.delete('job'); + const qs = search.toString(); + const url = `http://${engine.endpoint.host}:${engine.endpoint.port}/${upstreamPath}${qs ? `?${qs}` : ''}`; + + const headers = new Headers(); + request.headers.forEach((value, key) => { + if (!HOP_BY_HOP.has(key.toLowerCase())) headers.set(key, value); + }); + headers.set('x-engine-token', engine.endpoint.token); + + const method = request.method.toUpperCase(); + const init: RequestInit & { duplex?: 'half' } = { method, headers, redirect: 'manual' }; + if (method !== 'GET' && method !== 'HEAD') { + // buffer the request body (JSON or a control-image upload); the response + // side is what needs to stream + init.body = Buffer.from(await request.arrayBuffer()); + } + + let upstream: Response; + try { + upstream = await fetch(url, init); + } catch (e: any) { + return NextResponse.json({ error: `engine unreachable: ${e?.message || e}` }, { status: 502 }); + } + + const responseHeaders = new Headers(); + upstream.headers.forEach((value, key) => { + if (!HOP_BY_HOP.has(key.toLowerCase())) responseHeaders.set(key, value); + }); + responseHeaders.set('Cache-Control', 'no-store'); + // keep intermediaries from buffering the frame stream + responseHeaders.set('X-Accel-Buffering', 'no'); + return new Response(upstream.body, { status: upstream.status, headers: responseHeaders }); +} + +type Ctx = { params: Promise<{ path: string[] }> | { path: string[] } }; + +export async function GET(request: NextRequest, ctx: Ctx) { + const { path } = await ctx.params; + return proxy(request, path); +} + +export async function POST(request: NextRequest, ctx: Ctx) { + const { path } = await ctx.params; + return proxy(request, path); +} + +export async function DELETE(request: NextRequest, ctx: Ctx) { + const { path } = await ctx.params; + return proxy(request, path); +} diff --git a/ui/src/app/generate/page.tsx b/ui/src/app/generate/page.tsx new file mode 100644 index 0000000..a2c1e49 --- /dev/null +++ b/ui/src/app/generate/page.tsx @@ -0,0 +1,735 @@ +'use client'; + +import React, { Suspense, useCallback, useEffect, useMemo, useRef, useState } from 'react'; +import { useSearchParams } from 'next/navigation'; +import { Button } from '@headlessui/react'; +import { ChevronDown, ChevronLeft, ChevronRight, Loader2, Play, Square, Sparkles, Trash2, X } from 'lucide-react'; +import { openConfirm } from '@/components/ConfirmModal'; +import { TopBar, MainContent } from '@/components/layout'; +import { Checkbox, CreatableSelectInput, NumberInput, SelectInput, TextAreaInput, TextInput } from '@/components/formInputs'; +import { apiClient } from '@/utils/api'; +import { startJob, stopJob } from '@/utils/jobs'; +import { startQueue } from '@/utils/queue'; +import useGPUInfo from '@/hooks/useGPUInfo'; +import usePollLoop from '@/hooks/usePollLoop'; +import { defaultInferenceJobConfig } from '@/helpers/inferenceJobConfig'; +import { encodeFilePathForUrl } from '@/utils/basic'; +import { latentToImage, payloadToFloat32, readEngineFrames, PreviewInfo } from '@/lib/engineStream'; +import { isMac } from '@/helpers/basic'; +import { modelArchs, getGenerateDefaults, GenerateDefaults } from '@/app/jobs/new/options'; +import GenerateFooter from '@/components/generate/GenerateFooter'; + + +interface EngineStatus { + running: boolean; + engine: { jobId: string; jobName: string; gpuIds: string; status: string; ready: boolean } | null; + engines: { jobId: string; jobName: string; gpuIds: string; status: string; ready: boolean }[]; +} + +interface ResultItem { + request_id: string; + path: string; + seconds?: number; + steps?: number; + kind: string; + ext: string; + seed: number; + width: number; + height: number; + prompt: string; + arch: string; +} + +const qtypeOptions = [ + { value: 'convrot8', label: 'ConvRot 8-bit' }, + { value: 'float8', label: 'Float8' }, + { value: 'qfloat8', label: 'Quanto float8' }, + { value: 'nvfp4', label: 'NVFP4' }, +]; + +// everything the user set on the page survives a reload +const STORAGE_KEY = 'aitk_generate_page'; +interface PersistedState { + arch?: string; + model?: { [key: string]: any }; + sample?: { [key: string]: any }; + gpuId?: string; + results?: ResultItem[]; + sidebarOpen?: boolean; + cards?: { [key: string]: boolean }; +} +const loadPersisted = (): PersistedState => { + if (typeof window === 'undefined') return {}; + try { + const raw = localStorage.getItem(STORAGE_KEY); + return raw ? (JSON.parse(raw) as PersistedState) : {}; + } catch { + return {}; + } +}; + +const authHeaders = (): Record => { + const token = typeof window !== 'undefined' ? localStorage.getItem('AI_TOOLKIT_AUTH') : null; + return token ? { Authorization: `Bearer ${token}` } : {}; +}; + +function Card({ title, open, onToggle, children }: { title: string; open: boolean; onToggle: () => void; children: React.ReactNode }) { + return ( +
+ +
{children}
+
+ ); +} + +function GeneratePageInner() { + const searchParams = useSearchParams(); + const preferredJob = searchParams.get('job'); + const { gpuList } = useGPUInfo(); + + // ---- engine lifecycle ---- + const [engineStatus, setEngineStatus] = useState(null); + const [health, setHealth] = useState(null); + const persisted = useMemo(loadPersisted, []); + const [gpuId, setGpuId] = useState(persisted.gpuId || '0'); + const [engineBusy, setEngineBusy] = useState(false); + // job ids we are waiting on: the button stays disabled from the click until + // the engine is actually up (or the job row reports a failure) + const [startingJobId, setStartingJobId] = useState(null); + const [stoppingJobId, setStoppingJobId] = useState(null); + const engineJobId = preferredJob || engineStatus?.engine?.jobId || null; + const proxy = useCallback((path: string) => `/api/inference/${path}${engineJobId ? `${path.includes('?') ? '&' : '?'}job=${engineJobId}` : ''}`, [engineJobId]); + + usePollLoop( + () => + apiClient + .get(`/api/inference/status${preferredJob ? `?job=${preferredJob}` : ''}`) + .then(res => res.data) + .then(async (data: EngineStatus) => { + setEngineStatus(data); + if (startingJobId) { + const mine = data.engines.find(e => e.jobId === startingJobId); + if (mine?.ready) { + setStartingJobId(null); + } else if (!mine) { + // no longer queued/running: the launch failed (or was stopped) + const job = await apiClient.get('/api/jobs', { params: { id: startingJobId } }).then(r => r.data).catch(() => null); + if (!job || ['error', 'stopped', 'completed'].includes(job.status)) { + setStartingJobId(null); + alert(`The engine did not start${job?.info ? `: ${job.info}` : ''}. Check the job log in the Queue.`); + } + } + } + if (stoppingJobId && !data.engines.some(e => e.jobId === stoppingJobId)) { + setStoppingJobId(null); + } + if (data.running) { + return apiClient + .get(proxy('health')) + .then(r => setHealth(r.data)) + .catch(() => setHealth(null)); + } + setHealth(null); + }) + .catch(() => {}), + 2000, + [preferredJob, proxy, startingJobId, stoppingJobId], + ); + + useEffect(() => { + if (gpuList.length && !gpuList.some(g => `${g.index}` === gpuId)) setGpuId(`${gpuList[0].index}`); + }, [gpuList, gpuId]); + + const startEngine = async () => { + setEngineBusy(true); + try { + const gpu = isMac() ? 'mps' : gpuId; + const name = `inference_engine_gpu${gpu}`; + const existing = await apiClient.get('/api/jobs', { params: { job_type: 'inference' } }).then(r => r.data.jobs || []); + const match = existing.find((j: any) => j.name === name); + const res = await apiClient.post('/api/jobs', { + id: match ? match.id : null, + name, + gpu_ids: gpu, + job_config: { ...defaultInferenceJobConfig, config: { ...defaultInferenceJobConfig.config, name } }, + job_type: 'inference', + }); + await startJob(res.data.id); + await startQueue(gpu); + setStartingJobId(res.data.id); + } catch (e: any) { + alert(`Failed to start the engine: ${e?.response?.data?.error || e?.message || e}`); + } finally { + setEngineBusy(false); + } + }; + + const stopEngine = async () => { + if (!engineStatus?.engine) return; + setEngineBusy(true); + try { + await stopJob(engineStatus.engine.jobId); + setStoppingJobId(engineStatus.engine.jobId); + } catch (e: any) { + alert(`Failed to stop the engine: ${e?.response?.data?.error || e?.message || e}`); + } finally { + setEngineBusy(false); + } + }; + + const isStarting = engineBusy || !!startingJobId || (!!engineStatus?.engine && !engineStatus.running && !stoppingJobId); + const isStopping = !!stoppingJobId; + + // ---- models: the same arch list the training UI uses (jobs/new/options.tsx) ---- + const archs: GenerateDefaults[] = useMemo(() => modelArchs.map(getGenerateDefaults), []); + const ready = !!engineStatus?.running; + + const [arch, setArch] = useState(persisted.arch || 'zimage'); + const [model, setModel] = useState<{ [key: string]: any }>(persisted.model || {}); + const [sample, setSample] = useState<{ [key: string]: any }>(persisted.sample || { prompt: '', negative_prompt: '', seed: -1 }); + const entry = useMemo(() => archs.find(a => a.arch === arch) || null, [archs, arch]); + + const applyArch = (name: string) => { + setArch(name); + const e = archs.find(a => a.arch === name); + if (!e) return; + setModel({ ...e.model }); + setSample(s => ({ ...s, ...e.sample, prompt: s.prompt, negative_prompt: s.negative_prompt, seed: s.seed ?? -1, ctrl_img: undefined })); + }; + useEffect(() => { + if (archs.length && !Object.keys(model).length) applyArch(archs.some(a => a.arch === arch) ? arch : archs[0].arch); + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [archs]); + + const archOptions = useMemo(() => { + const groups: { [k: string]: { value: string; label: string }[] } = {}; + for (const a of archs) { + (groups[a.group] ||= []).push({ value: a.arch, label: a.label }); + } + return Object.entries(groups).map(([label, options]) => ({ label, options })); + }, [archs]); + + // ---- generation ---- + const [running, setRunning] = useState(false); + const [progress, setProgress] = useState<{ step: number; total: number | null; elapsed: number } | null>(null); + const [statusLine, setStatusLine] = useState(''); + const [error, setError] = useState(null); + const [results, setResults] = useState(persisted.results || []); + const [selected, setSelected] = useState(persisted.results?.[0] || null); + // the stage shows the live latent canvas while generating; a finished (or + // clicked) result replaces it + const [showPreview, setShowPreview] = useState(!persisted.results?.length); + const [sidebarOpen, setSidebarOpen] = useState(persisted.sidebarOpen ?? true); + const [cards, setCards] = useState<{ [key: string]: boolean }>(persisted.cards || {}); + const cardOpen = (key: string) => cards[key] ?? true; + const toggleCard = (key: string) => setCards(c => ({ ...c, [key]: !(c[key] ?? true) })); + const resultUrl = (r: ResultItem) => `/api/files/${encodeFilePathForUrl(r.path)}`; + + const deleteResults = async (items: ResultItem[]) => { + const paths = items.map(i => i.path); + try { + await apiClient.post('/api/inference/outputs/delete', { paths }); + } catch (e: any) { + alert(`Failed to delete: ${e?.response?.data?.error || e?.message || e}`); + return; + } + setResults(prev => { + const next = prev.filter(r => !paths.includes(r.path)); + if (selected && paths.includes(selected.path)) { + const replacement = next[0] || null; + setSelected(replacement); + if (!replacement) setShowPreview(true); + } + return next; + }); + }; + + const deleteAll = () => { + if (!results.length) return; + openConfirm({ + title: 'Delete all results', + message: `Delete all ${results.length} generated files from disk? This cannot be undone.`, + type: 'warning', + confirmText: 'Delete all', + onConfirm: () => deleteResults(results), + }); + }; + + // persist on every change (small payload; results capped) + useEffect(() => { + try { + const state: PersistedState = { arch, model, sample, gpuId, results: results.slice(0, 60), sidebarOpen, cards }; + localStorage.setItem(STORAGE_KEY, JSON.stringify(state)); + } catch { + // storage full or unavailable: nothing to do + } + }, [arch, model, sample, gpuId, results, sidebarOpen, cards]); + const [frameIdx, setFrameIdx] = useState(0); + const [frameCount, setFrameCount] = useState(1); + const canvasRef = useRef(null); + const framesRef = useRef([]); + const abortRef = useRef(null); + const requestIdRef = useRef(null); + + // the canvas is sized to the requested image dimensions (from the engine's + // start frame) so the preview occupies exactly the space the image will + const previewSizeRef = useRef<{ width: number; height: number } | null>(null); + const scratchRef = useRef(null); + const drawFrame = (idx: number) => { + const canvas = canvasRef.current; + const frames = framesRef.current; + if (!canvas || !frames.length) return; + const img = frames[Math.min(idx, frames.length - 1)]; + const target = previewSizeRef.current; + const w = target?.width || img.width; + const h = target?.height || img.height; + if (canvas.width !== w || canvas.height !== h) { + canvas.width = w; + canvas.height = h; + } + if (!scratchRef.current) scratchRef.current = document.createElement('canvas'); + const scratch = scratchRef.current; + if (scratch.width !== img.width || scratch.height !== img.height) { + scratch.width = img.width; + scratch.height = img.height; + } + scratch.getContext('2d')?.putImageData(img, 0, 0); + const ctx = canvas.getContext('2d'); + if (!ctx) return; + ctx.imageSmoothingEnabled = true; + ctx.imageSmoothingQuality = 'high'; + ctx.drawImage(scratch, 0, 0, w, h); + }; + useEffect(() => drawFrame(frameIdx), [frameIdx]); + + const generate = async () => { + if (!ready || running) return; + setRunning(true); + setError(null); + setProgress(null); + setStatusLine('Submitting'); + framesRef.current = []; + setFrameCount(1); + setFrameIdx(0); + // keep whatever is on the stage until the first latent of this run + // arrives; wipe the stale bitmap so an empty canvas never shows old data + const cv = canvasRef.current; + if (cv) cv.getContext('2d')?.clearRect(0, 0, cv.width, cv.height); + previewSizeRef.current = sample.width && sample.height ? { width: sample.width, height: sample.height } : null; + const abort = new AbortController(); + abortRef.current = abort; + let preview: PreviewInfo | null = null; + const body: { model: any; sample: any; stream: any } = { + model: { ...model, arch }, + sample: { ...sample, seed: sample.seed === '' ? -1 : sample.seed }, + stream: { latents: 'raw', every_n_steps: 1, max_frames: 8 }, + }; + try { + const res = await fetch(proxy('generate'), { + method: 'POST', + headers: { ...authHeaders(), 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + signal: abort.signal, + }); + if (!res.ok) throw new Error((await res.text()) || `HTTP ${res.status}`); + await readEngineFrames( + res, + ({ header, payload }) => { + requestIdRef.current = header.request_id; + switch (header.type) { + case 'start': + preview = header.preview; + if (header.sample?.width && header.sample?.height) { + previewSizeRef.current = { width: header.sample.width, height: header.sample.height }; + } + setStatusLine(`Generating with ${header.model?.arch}`); + break; + case 'status': + setStatusLine(header.message); + break; + case 'progress': + setProgress({ step: header.step, total: header.total, elapsed: header.elapsed }); + break; + case 'latent': { + try { + if (header.preview) preview = header.preview; + const data = payloadToFloat32(header, payload); + const img = latentToImage(header, data, preview); + if (img) { + framesRef.current = img.frameData.map(px => new ImageData(px, img.width, img.height)); + setFrameCount(img.frames); + drawFrame(frameIdx); + setShowPreview(true); + } + } catch (e) { + console.warn('latent preview failed', e); + } + break; + } + case 'result': { + const item: ResultItem = { + request_id: header.request_id, + path: header.path, + kind: header.kind, + ext: header.ext, + seed: header.seed, + width: header.width, + height: header.height, + seconds: header.seconds, + steps: header.steps, + prompt: body.sample.prompt, + arch, + }; + setResults(r => [item, ...r]); + setSelected(item); + setShowPreview(false); + break; + } + case 'error': + setError(header.cancelled ? 'Cancelled' : header.message); + break; + case 'end': + setStatusLine(header.status === 'done' ? 'Done' : header.status); + break; + } + }, + abort.signal, + ); + } catch (e: any) { + if (e?.name !== 'AbortError') setError(e?.message || String(e)); + } finally { + setRunning(false); + abortRef.current = null; + } + }; + + // ctrl/cmd+enter anywhere on the page generates (the prompt textarea included) + const generateRef = useRef(generate); + generateRef.current = generate; + useEffect(() => { + const onKey = (e: KeyboardEvent) => { + if (e.key === 'Enter' && (e.ctrlKey || e.metaKey) && !e.repeat) { + e.preventDefault(); + generateRef.current(); + } + }; + window.addEventListener('keydown', onKey); + return () => window.removeEventListener('keydown', onKey); + }, []); + + const cancel = async () => { + const id = requestIdRef.current; + if (id) { + try { + await apiClient.post(proxy(`cancel/${id}`)); + } catch { + // engine gone; abort below still tears the stream down + } + } + abortRef.current?.abort(); + }; + + const modality = entry?.modality || 'image'; + const footerStatus = running + ? statusLine || 'Generating' + : error + ? error + : ready + ? health?.current + ? `Generating ${health.current.arch}` + : health?.active?.arch + ? `Idle - ${health.active.arch} (${health.active.model?.name_or_path || ''})` + : 'Idle - no model loaded' + : engineStatus?.engine + ? `Engine ${engineStatus.engine.status}` + : 'Engine off'; + const vram = health?.vram ? `${(health.vram.used / 1e9).toFixed(1)} / ${(health.vram.total / 1e9).toFixed(0)} GB` : null; + const pool = health?.pool; + + return ( + <> + +
+ +

Generate

+
+
+
+ {isStopping ? ( + + ) : ready ? ( + <> + Engine running + {health?.active?.arch && · {health.active.arch}} + {vram && · VRAM {vram}} + + + ) : null} +
+ + + {!ready && ( +
+
+
+ +
+

Inference Engine

+

+ {isStarting + ? 'The engine is starting. Model options appear once its server is up.' + : 'Start the inference engine on a GPU to generate images, video, and audio with any model.'} +

+ {isStarting ? ( + + ) : ( +
+ {!isMac() && ( + + )} + +
+ )} +
+
+ )} +
+ {/* ---- center: stage + history strip ---- */} +
+
+ {/* status strip over the stage */} +
+ {running && } + + {running + ? statusLine + : selected + ? `${selected.arch} · seed ${selected.seed}${selected.seconds ? ` · ${selected.seconds.toFixed(1)}s${selected.steps ? ` / ${selected.steps} steps` : ''}` : ''}` + : ''} + + {progress ? ( + + step {progress.step} + {progress.total ? `/${progress.total}` : ''} · {progress.elapsed.toFixed(1)}s + + ) : null} +
+ {running && ( + + )} +
+ {progress?.total ? ( +
+
+
+ ) : null} + {error &&
{error}
} + + {/* live latent preview: shown while generating (or when nothing is selected) */} +
+ + {!running && framesRef.current.length === 0 && ( +
Your generation will appear here
+ )} +
+ {frameCount > 1 && showPreview && ( + setFrameIdx(parseInt(e.target.value))} className="absolute bottom-3 left-1/2 -translate-x-1/2 w-1/2 z-10" /> + )} + + {/* selected result replaces the preview */} + {!showPreview && selected && ( +
+ {selected.kind === 'video' ? ( +
+ )} +
+ + {/* history strip */} +
+
+ {results.length === 0 ? ( +
History
+ ) : ( + results.map(r => { + const isSel = selected?.path === r.path && !showPreview; + return ( +
+ + +
+ ); + }) + )} +
+ {results.length > 0 && ( + + )} +
+
+ + {/* ---- right: collapsible model + prompt sidebar ---- */} +
+ +
+
+ + toggleCard('model')}> + applyArch(v as string)} options={archOptions} disabled={!ready} /> + setModel(m => ({ ...m, name_or_path: v || '' }))} + options={entry?.model?.name_or_path ? [{ value: entry.model.name_or_path, label: entry.model.name_or_path }] : []} + placeholder="hub repo, local folder, or .safetensors" + /> + {'extras_name_or_path' in (entry?.model || {}) && ( + setModel(m => ({ ...m, extras_name_or_path: v }))} /> + )} +
+ setModel(m => ({ ...m, quantize: !!v, qtype: v || undefined }))} + options={[{ value: '', label: 'None (bf16)' }, ...qtypeOptions]} + /> + setModel(m => ({ ...m, quantize_te: !!v, qtype_te: v || undefined }))} + options={[{ value: '', label: 'None (bf16)' }, ...qtypeOptions]} + /> +
+
+ setModel(m => ({ ...m, low_vram: v }))} /> + setModel(m => ({ ...m, layer_offloading: v }))} /> +
+ + + toggleCard('prompt')}> + setSample(s => ({ ...s, prompt: v }))} placeholder={modality === 'audio' ? 'song description, lyrics, bpm…' : 'describe what to generate'} /> + {(sample.guidance_scale ?? 4) > 1 && ( + setSample(s => ({ ...s, negative_prompt: v }))} /> + )} +
+ setSample(s => ({ ...s, width: v }))} min={64} max={4096} /> + setSample(s => ({ ...s, height: v }))} min={64} max={4096} /> + setSample(s => ({ ...s, num_inference_steps: v }))} min={1} max={200} /> + setSample(s => ({ ...s, guidance_scale: v }))} min={1} max={30} /> + setSample(s => ({ ...s, seed: v }))} min={-1} max={4294967295} /> + {modality === 'video' && ( + <> + setSample(s => ({ ...s, num_frames: v }))} min={1} max={1000} /> + setSample(s => ({ ...s, fps: v }))} min={1} max={60} /> + + )} +
+ + +
+ {/* always-visible action bar */} +
+ + {running && ( + + )} +
+
+
+
+ + + + ); +} + +export default function GeneratePage() { + return ( + + + + ); +} diff --git a/ui/src/app/jobs/[jobID]/page.tsx b/ui/src/app/jobs/[jobID]/page.tsx index 80a01f3..7199c38 100644 --- a/ui/src/app/jobs/[jobID]/page.tsx +++ b/ui/src/app/jobs/[jobID]/page.tsx @@ -97,6 +97,9 @@ export default function JobPage({ params }: { params: { jobID: string } }) { if (jobType === 'caption') { title = `Captioning: ${job?.job_ref || 'Loading...'}`; } + if (jobType === 'inference') { + title = `Inference Engine: ${job?.name || 'Loading...'}`; + } return ( <> diff --git a/ui/src/app/jobs/new/options.tsx b/ui/src/app/jobs/new/options.tsx index af95196..b4e2e78 100644 --- a/ui/src/app/jobs/new/options.tsx +++ b/ui/src/app/jobs/new/options.tsx @@ -58,10 +58,38 @@ export interface SampleTags { [key: string]: SampleTag; } +export type GenerateModality = 'image' | 'video' | 'audio'; + +// Per-arch overrides for the Generate page. Everything it needs is derived +// from the training entry (name_or_path / quantize defaults, video/audio +// group, ctrl_img section); set these only where the derivation is wrong. +export interface GenerateOptions { + modality?: GenerateModality; + model?: { [key: string]: any }; // extra ModelConfig kwargs + sample?: { [key: string]: any }; // GenerateImageConfig kwargs + needsControlImage?: boolean; + sizeLocked?: boolean; +} + +export interface GenerateDefaults { + arch: string; + label: string; + group: ModelGroup; + modality: GenerateModality; + model: { [key: string]: any }; + sample: { [key: string]: any }; + needsControlImage: boolean; + sizeLocked: boolean; +} + export interface ModelArch { name: string; label: string; + /** label shown by the Generate page instead of `label` (training-specific + * wording like "w/ Training Adapter" does not apply to inference) */ + generateNameOverride?: string; group: ModelGroup; + generate?: GenerateOptions; controls?: Control[]; isVideoModel?: boolean; hasMultiLinePrompts?: boolean; @@ -658,6 +686,7 @@ export const modelArchs: ModelArch[] = [ { name: 'zimage:turbo', label: 'Z-Image Turbo (w/ Training Adapter)', + generateNameOverride: 'Z-Image Turbo', group: 'image', defaults: { // default updates when [selected, unselected] in the UI @@ -1519,6 +1548,7 @@ export const modelArchs: ModelArch[] = [ { name: 'krea2:turbo', label: 'Krea 2 Turbo (w/ Training Adapter)', + generateNameOverride: 'Krea 2 Turbo', group: 'image', gateUrl: 'https://huggingface.co/krea/Krea-2-Turbo', defaults: { @@ -1575,6 +1605,7 @@ export const modelArchs: ModelArch[] = [ { name: 'krea2:o_edit_turbo', label: 'Krea 2 Turbo (w/ Training Adapter) [Edit Training]', + generateNameOverride: 'Krea 2 Turbo (Edit)', gateUrl: 'https://huggingface.co/krea/Krea-2-Turbo', group: 'experimental', defaults: { @@ -1770,3 +1801,65 @@ export const jobTypeOptions: JobTypeOption[] = [ }, }, ]; + + +const MODEL_PREFIX = 'config.process[0].model.'; +const SAMPLE_PREFIX = 'config.process[0].sample'; +// training-only model settings: the training adapter (e.g. Z-Image Turbo's +// de-distill LoRA) and the unconditional LoRA must not load for inference +const TRAINING_ONLY_MODEL_KEYS = new Set(['assistant_lora_path', 'unconditional_lora_path', 'inference_lora_path']); + +/** What the Generate page sends the inference engine for an arch: ModelConfig + * kwargs + GenerateImageConfig kwargs, derived from the training defaults. */ +export const getGenerateDefaults = (arch: ModelArch): GenerateDefaults => { + const defaults = arch.defaults || {}; + const model: { [key: string]: any } = {}; + const sample: { [key: string]: any } = { width: 1024, height: 1024, num_inference_steps: 25, guidance_scale: 4 }; + for (const [key, pair] of Object.entries(defaults)) { + const value = Array.isArray(pair) ? pair[0] : pair; + if (key.startsWith(MODEL_PREFIX)) { + const field = key.slice(MODEL_PREFIX.length); + if (value === '' || value === undefined || value === null) continue; + if (field.includes('.')) continue; // nested model_kwargs etc. + if (TRAINING_ONLY_MODEL_KEYS.has(field)) continue; + model[field] = value; + } else if (key === SAMPLE_PREFIX && value && typeof value === 'object') { + // whole SampleConfig object (audio models) + const sc = value as any; + if (sc.width) sample.width = sc.width; + if (sc.height) sample.height = sc.height; + if (sc.sample_steps) sample.num_inference_steps = sc.sample_steps; + if (sc.guidance_scale !== undefined) sample.guidance_scale = sc.guidance_scale; + if (sc.num_frames) sample.num_frames = sc.num_frames; + if (sc.fps) sample.fps = sc.fps; + if (sc.neg) sample.negative_prompt = sc.neg; + } else if (key.startsWith(SAMPLE_PREFIX + '.')) { + const field = key.slice(SAMPLE_PREFIX.length + 1); + if (value === undefined || value === null || value === '') continue; + if (field === 'sample_steps') sample.num_inference_steps = value; + else if (field === 'neg') sample.negative_prompt = value; + else if (['width', 'height', 'guidance_scale', 'num_frames', 'fps'].includes(field)) sample[field] = value; + } + } + // the engine defaults to convrot8; the training default qtype is the + // quanto one, which is the slower inference choice + if (model.quantize && (!model.qtype || model.qtype === 'qfloat8')) model.qtype = 'convrot8'; + if (model.quantize_te && (!model.qtype_te || model.qtype_te === 'qfloat8')) model.qtype_te = 'convrot8'; + const sections = arch.additionalSections || []; + const gen = arch.generate || {}; + const modality: GenerateModality = gen.modality || (arch.group === 'audio' ? 'audio' : arch.isVideoModel ? 'video' : 'image'); + if (modality !== 'video') { + delete sample.num_frames; + delete sample.fps; + } + return { + arch: arch.name, + label: arch.generateNameOverride || arch.label, + group: arch.group, + modality, + model: { ...model, ...(gen.model || {}) }, + sample: { ...sample, ...(gen.sample || {}) }, + needsControlImage: gen.needsControlImage ?? (sections.includes('sample.ctrl_img') || sections.includes('sample.multi_ctrl_imgs')), + sizeLocked: gen.sizeLocked ?? false, + }; +}; diff --git a/ui/src/components/ActiveJobWidget.tsx b/ui/src/components/ActiveJobWidget.tsx index 4513e5b..43a8050 100644 --- a/ui/src/components/ActiveJobWidget.tsx +++ b/ui/src/components/ActiveJobWidget.tsx @@ -26,6 +26,10 @@ export default function ActiveJobWidget() { let label = job.name; let href = `/jobs/${job.id}`; + if (job.job_type === 'inference') { + label = 'Inference Engine'; + href = `/generate?job=${job.id}`; + } if (job.job_type === 'caption') { const splits = (job.job_ref ?? '').split(/[/\\]/); label = splits[splits.length - 1] || job.name; diff --git a/ui/src/components/JobsTable.tsx b/ui/src/components/JobsTable.tsx index cb6f761..2c35401 100644 --- a/ui/src/components/JobsTable.tsx +++ b/ui/src/components/JobsTable.tsx @@ -128,6 +128,14 @@ export default function JobsTable({ onlyActive = false, job_type = null }: JobsT let title = row.name; let href = `/jobs/${row.id}`; // if (row.job_type === 'train') title = `Train: ${title}`; + if (row.job_type === 'inference') { + href = `/generate?job=${row.id}`; + title = ( + <> + INFERENCE ENGINE: {row.name} + + ); + } if (row.job_type === 'caption') { let splits = row.job_ref.split(/[/\\]/); const datasetPath = `${splits[splits.length - 1]}`; diff --git a/ui/src/components/Sidebar.tsx b/ui/src/components/Sidebar.tsx index 6882dfc..5277f23 100644 --- a/ui/src/components/Sidebar.tsx +++ b/ui/src/components/Sidebar.tsx @@ -3,7 +3,7 @@ import { useEffect } from 'react'; import Link from 'next/link'; import { usePathname } from 'next/navigation'; -import { Home, Settings, BrainCircuit, Images, Plus, X } from 'lucide-react'; +import { Home, Settings, BrainCircuit, Images, Plus, X, Sparkles } from 'lucide-react'; import { FaXTwitter, FaDiscord, FaYoutube } from 'react-icons/fa6'; import { createGlobalState } from 'react-global-hooks'; import ThemeToggle from './ThemeToggle'; @@ -36,6 +36,7 @@ const Sidebar = () => { const navigation = [ { name: 'Dashboard', href: '/dashboard', icon: Home }, + { name: 'Generate', href: '/generate', icon: Sparkles, badge: 'alpha' }, { name: 'New Job', href: '/jobs/new', icon: Plus }, { name: 'Queue', href: '/jobs', icon: BrainCircuit }, { name: 'Datasets', href: '/datasets', icon: Images }, @@ -73,6 +74,11 @@ const Sidebar = () => { > {item.name} + {'badge' in item && item.badge && ( + + {item.badge} + + )} ))} diff --git a/ui/src/components/generate/GenerateFooter.tsx b/ui/src/components/generate/GenerateFooter.tsx new file mode 100644 index 0000000..94335bf --- /dev/null +++ b/ui/src/components/generate/GenerateFooter.tsx @@ -0,0 +1,92 @@ +'use client'; + +import { useEffect, useMemo, useRef, useState } from 'react'; +import { ChevronUp, ChevronDown, Loader2 } from 'lucide-react'; +import useJobLog from '@/hooks/useJobLog'; + +interface Props { + /** engine job id: the log panel follows its log.txt */ + jobId: string | null; + /** one-line status shown in the collapsed bar */ + status: string; + busy?: boolean; + progress?: { step: number; total: number | null } | null; +} + +export default function GenerateFooter({ jobId, status, busy, progress }: Props) { + const [open, setOpen] = useState(false); + const { log, status: logStatus } = useJobLog(jobId ?? '', jobId && open ? 1500 : null); + const logRef = useRef(null); + const [follow, setFollow] = useState(true); + + const lines = useMemo(() => { + const splits = log.split('\n'); + return splits.length > 1500 ? splits.slice(splits.length - 1500) : splits; + }, [log]); + + const handleScroll = () => { + const el = logRef.current; + if (!el) return; + setFollow(el.scrollHeight - el.scrollTop - el.clientHeight < 10); + }; + + useEffect(() => { + if (open && follow && logRef.current) { + logRef.current.scrollTop = logRef.current.scrollHeight; + } + }, [log, open, follow]); + + const pct = progress?.total ? Math.min(100, (progress.step / progress.total) * 100) : null; + + return ( +
+ {/* progress hairline */} + {pct !== null && ( +
+
+
+ )} + + {open && ( +
+
+ {!jobId &&
No engine job.
} + {jobId && logStatus === 'loading' && !log &&
Loading log…
} + {jobId && logStatus === 'error' &&
Error loading log
} + {lines.map((line, i) => ( + 
+                {line}
+
+ ))} +
+ {!follow && ( + + )} +
+ )} +
+ ); +} diff --git a/ui/src/helpers/inferenceJobConfig.ts b/ui/src/helpers/inferenceJobConfig.ts new file mode 100644 index 0000000..8647933 --- /dev/null +++ b/ui/src/helpers/inferenceJobConfig.ts @@ -0,0 +1,18 @@ +// Job config for the resident inference engine (extensions_built_in/inference_engine). +// The worker fills engine.job_folder in when it launches the job. +export const defaultInferenceJobConfig = { + job: 'extension', + config: { + name: 'inference_engine', + process: [ + { + type: 'InferenceEngine', + sqlite_db_path: './aitk_db.db', + device: 'cuda', + engine: { + dtype: 'bf16', + }, + }, + ], + }, +}; diff --git a/ui/src/server/inferenceEngine.ts b/ui/src/server/inferenceEngine.ts new file mode 100644 index 0000000..d32430c --- /dev/null +++ b/ui/src/server/inferenceEngine.ts @@ -0,0 +1,76 @@ +import fs from 'fs'; +import path from 'path'; +import prisma from '@/server/prisma'; +import { getTrainingFolder } from '@/server/settings'; + +// What the engine job writes to /engine.json when its server is up. +export interface EngineEndpoint { + host: string; + port: number; + token: string; + pid: number; + job_id?: string | null; + device?: string; + job_folder: string; + output_folder: string; + started_at: number; +} + +export interface EngineHandle { + jobId: string; + jobName: string; + gpuIds: string; + status: string; + endpoint: EngineEndpoint | null; +} + +const isProcessAlive = (pid: number): boolean => { + try { + process.kill(pid, 0); + return true; + } catch (e: any) { + return e?.code === 'EPERM'; + } +}; + +const readEndpoint = (jobFolder: string): EngineEndpoint | null => { + try { + const raw = fs.readFileSync(path.join(jobFolder, 'engine.json'), 'utf8'); + const ep = JSON.parse(raw) as EngineEndpoint; + if (!ep.host || !ep.port) return null; + // a stale file from a crashed engine must not be proxied to + if (ep.pid && !isProcessAlive(ep.pid)) return null; + return ep; + } catch { + return null; + } +}; + +let cache: { ts: number; handles: EngineHandle[] } | null = null; +const CACHE_MS = 2000; + +/** Every inference-engine job that is running (or starting), with its endpoint when published. */ +export async function listEngines(forceFresh = false): Promise { + if (!forceFresh && cache && Date.now() - cache.ts < CACHE_MS) return cache.handles; + const jobs = await prisma.job.findMany({ + where: { job_type: 'inference', status: { in: ['running', 'queued', 'stopping'] } }, + orderBy: { updated_at: 'desc' }, + }); + const trainingFolder = await getTrainingFolder(); + const handles: EngineHandle[] = jobs.map(job => ({ + jobId: job.id, + jobName: job.name, + gpuIds: job.gpu_ids, + status: job.status, + endpoint: job.status === 'running' ? readEndpoint(path.join(trainingFolder, job.name)) : null, + })); + cache = { ts: Date.now(), handles }; + return handles; +} + +/** The engine to talk to: the requested job, else the first running one with an endpoint. */ +export async function resolveEngine(jobId?: string | null): Promise { + const handles = await listEngines(); + if (jobId) return handles.find(h => h.jobId === jobId) || null; + return handles.find(h => h.endpoint) || handles[0] || null; +}