Initial inference engine build out
This commit is contained in:
Jaret Burkett committed
1 parent
f56b5a1d40
commit
449488f6b7
33 files changed
+4253
-4
No files matched your search
+2
-1
@@ -192,4 +192,5 @@ aitk_db.db-shm
|
||||
.claude
|
||||
original_repo
|
||||
.next
|
||||
testing/.model_test_outputs
|
||||
testing/.model_test_outputs
|
||||
testing/.engine_test_outputs
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 <job_folder>/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
|
||||
@@ -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 <config>` 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
|
||||
`<training_folder>/<job_name>/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 `<training_folder>/<job.name>` 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: <training_folder>/<name>/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/<arch>`), 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)?
|
||||
@@ -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
|
||||
`<job_folder>/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: <token>` or
|
||||
`Authorization: Bearer <token>`).
|
||||
|
||||
| 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).
|
||||
@@ -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,
|
||||
]
|
||||
@@ -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"
|
||||
@@ -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],
|
||||
],
|
||||
},
|
||||
}
|
||||
@@ -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"],
|
||||
}
|
||||
@@ -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("<I")
|
||||
_LEN = struct.Struct("<Q")
|
||||
|
||||
|
||||
def encode_frame(header: dict, payload: bytes = b"") -> 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
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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())]
|
||||
@@ -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
|
||||
|
||||
@@ -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]),
|
||||
}
|
||||
@@ -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)
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
@@ -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));
|
||||
|
||||
@@ -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=<id>` 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);
|
||||
}
|
||||
@@ -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<string, string> => {
|
||||
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 (
|
||||
<div className="bg-gray-900 rounded-xl border border-gray-800">
|
||||
<button type="button" onClick={onToggle} className="w-full flex items-center justify-between px-4 py-2.5 text-sm text-gray-200 font-semibold">
|
||||
<span>{title}</span>
|
||||
{open ? <ChevronDown className="w-4 h-4 text-gray-400" /> : <ChevronLeft className="w-4 h-4 text-gray-400" />}
|
||||
</button>
|
||||
<div className={open ? 'px-4 pb-4' : 'hidden'}>{children}</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function GeneratePageInner() {
|
||||
const searchParams = useSearchParams();
|
||||
const preferredJob = searchParams.get('job');
|
||||
const { gpuList } = useGPUInfo();
|
||||
|
||||
// ---- engine lifecycle ----
|
||||
const [engineStatus, setEngineStatus] = useState<EngineStatus | null>(null);
|
||||
const [health, setHealth] = useState<any>(null);
|
||||
const persisted = useMemo(loadPersisted, []);
|
||||
const [gpuId, setGpuId] = useState<string>(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<string | null>(null);
|
||||
const [stoppingJobId, setStoppingJobId] = useState<string | null>(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<string>(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<string>('');
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [results, setResults] = useState<ResultItem[]>(persisted.results || []);
|
||||
const [selected, setSelected] = useState<ResultItem | null>(persisted.results?.[0] || null);
|
||||
// the stage shows the live latent canvas while generating; a finished (or
|
||||
// clicked) result replaces it
|
||||
const [showPreview, setShowPreview] = useState<boolean>(!persisted.results?.length);
|
||||
const [sidebarOpen, setSidebarOpen] = useState<boolean>(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<HTMLCanvasElement | null>(null);
|
||||
const framesRef = useRef<ImageData[]>([]);
|
||||
const abortRef = useRef<AbortController | null>(null);
|
||||
const requestIdRef = useRef<string | null>(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<HTMLCanvasElement | null>(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 (
|
||||
<>
|
||||
<TopBar>
|
||||
<div className="flex items-center gap-2">
|
||||
<Sparkles className="w-5 h-5 text-blue-400" />
|
||||
<h1 className="text-base sm:text-lg">Generate</h1>
|
||||
</div>
|
||||
<div className="flex-1" />
|
||||
<div className="flex items-center gap-2 text-xs sm:text-sm">
|
||||
{isStopping ? (
|
||||
<Button disabled className="px-2 py-1 rounded-md bg-red-900 text-white/70 flex items-center gap-1 cursor-wait">
|
||||
<Loader2 className="w-4 h-4 animate-spin" /> Stopping engine…
|
||||
</Button>
|
||||
) : ready ? (
|
||||
<>
|
||||
<span className="text-green-400">Engine running</span>
|
||||
{health?.active?.arch && <span className="text-gray-400 hidden sm:inline">· {health.active.arch}</span>}
|
||||
{vram && <span className="text-gray-400 hidden md:inline">· VRAM {vram}</span>}
|
||||
<Button onClick={stopEngine} disabled={engineBusy} className="ml-2 px-2 py-1 rounded-md bg-red-700 hover:bg-red-600 disabled:opacity-60 text-white flex items-center gap-1">
|
||||
{engineBusy ? <Loader2 className="w-4 h-4 animate-spin" /> : <Square className="w-4 h-4" />} Stop
|
||||
</Button>
|
||||
</>
|
||||
) : null}
|
||||
</div>
|
||||
</TopBar>
|
||||
<MainContent belowTopBar className="overflow-hidden">
|
||||
{!ready && (
|
||||
<div className="absolute inset-0 z-20 flex items-center justify-center p-4">
|
||||
<div className="bg-gray-900/95 backdrop-blur border border-gray-700 rounded-2xl shadow-2xl px-8 py-8 max-w-md w-full text-center">
|
||||
<div className="mx-auto w-16 h-16 rounded-full bg-blue-500/10 flex items-center justify-center mb-4">
|
||||
<Sparkles className="w-8 h-8 text-blue-400" />
|
||||
</div>
|
||||
<h2 className="text-lg text-gray-100 font-semibold mb-2">Inference Engine</h2>
|
||||
<p className="text-sm text-gray-400 mb-6">
|
||||
{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.'}
|
||||
</p>
|
||||
{isStarting ? (
|
||||
<Button disabled className="w-full px-3 py-2 rounded-md bg-blue-900 text-white/70 flex items-center justify-center gap-2 cursor-wait">
|
||||
<Loader2 className="w-4 h-4 animate-spin" />
|
||||
{engineStatus?.engine?.status === 'running' ? 'Engine loading…' : engineStatus?.engine ? `Engine ${engineStatus.engine.status}…` : 'Starting engine…'}
|
||||
</Button>
|
||||
) : (
|
||||
<div className="flex flex-col gap-3">
|
||||
{!isMac() && (
|
||||
<select value={gpuId} onChange={e => setGpuId(e.target.value)} className="w-full bg-gray-800 text-gray-200 rounded-md px-3 py-2 border border-gray-700">
|
||||
{gpuList.map(g => (
|
||||
<option key={g.index} value={`${g.index}`}>
|
||||
GPU #{g.index} {g.name}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
)}
|
||||
<Button onClick={startEngine} className="w-full px-3 py-2 rounded-md bg-blue-700 hover:bg-blue-600 text-white flex items-center justify-center gap-2">
|
||||
<Play className="w-4 h-4" /> Start engine
|
||||
</Button>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
<div className={`absolute inset-0 pb-9 flex ${ready ? '' : 'opacity-30 pointer-events-none select-none'}`}>
|
||||
{/* ---- center: stage + history strip ---- */}
|
||||
<div className="flex-1 min-w-0 flex flex-col">
|
||||
<div className="flex-1 min-h-0 relative bg-black/40 m-2 mb-0 rounded-xl border border-gray-800 overflow-hidden">
|
||||
{/* status strip over the stage */}
|
||||
<div className="absolute top-0 left-0 right-0 z-10 px-3 py-1.5 flex items-center gap-2 text-xs text-gray-300 bg-gradient-to-b from-gray-950/80 to-transparent">
|
||||
{running && <Loader2 className="w-3.5 h-3.5 animate-spin text-blue-400" />}
|
||||
<span className="truncate">
|
||||
{running
|
||||
? statusLine
|
||||
: selected
|
||||
? `${selected.arch} · seed ${selected.seed}${selected.seconds ? ` · ${selected.seconds.toFixed(1)}s${selected.steps ? ` / ${selected.steps} steps` : ''}` : ''}`
|
||||
: ''}
|
||||
</span>
|
||||
{progress ? (
|
||||
<span className="text-gray-400 shrink-0">
|
||||
step {progress.step}
|
||||
{progress.total ? `/${progress.total}` : ''} · {progress.elapsed.toFixed(1)}s
|
||||
</span>
|
||||
) : null}
|
||||
<div className="flex-1" />
|
||||
{running && (
|
||||
<Button onClick={cancel} className="px-2 py-0.5 rounded bg-gray-800/80 hover:bg-gray-700 text-gray-200 text-xs">
|
||||
Cancel
|
||||
</Button>
|
||||
)}
|
||||
</div>
|
||||
{progress?.total ? (
|
||||
<div className="absolute top-0 left-0 right-0 h-0.5 bg-gray-800 z-10">
|
||||
<div className="h-0.5 bg-blue-500 transition-all" style={{ width: `${Math.min(100, (progress.step / progress.total) * 100)}%` }} />
|
||||
</div>
|
||||
) : null}
|
||||
{error && <div className="absolute bottom-3 left-3 right-3 z-10 text-rose-300 text-sm bg-rose-950/70 rounded px-3 py-2 whitespace-pre-wrap">{error}</div>}
|
||||
|
||||
{/* live latent preview: shown while generating (or when nothing is selected) */}
|
||||
<div className={`absolute inset-3 flex items-center justify-center ${showPreview ? '' : 'hidden'}`}>
|
||||
<canvas ref={canvasRef} className="max-w-full max-h-full rounded" style={{ imageRendering: 'auto' }} />
|
||||
{!running && framesRef.current.length === 0 && (
|
||||
<div className="absolute text-gray-600 text-sm">Your generation will appear here</div>
|
||||
)}
|
||||
</div>
|
||||
{frameCount > 1 && showPreview && (
|
||||
<input type="range" min={0} max={frameCount - 1} value={frameIdx} onChange={e => 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 && (
|
||||
<div className="absolute inset-3 flex items-center justify-center">
|
||||
{selected.kind === 'video' ? (
|
||||
<video key={selected.path} src={resultUrl(selected)} controls autoPlay loop className="max-w-full max-h-full rounded" />
|
||||
) : selected.kind === 'audio' ? (
|
||||
<div className="w-full max-w-xl bg-gray-900 rounded-xl p-6 text-center">
|
||||
<div className="text-gray-300 text-sm mb-4 truncate" title={selected.prompt}>
|
||||
{selected.prompt}
|
||||
</div>
|
||||
<audio key={selected.path} src={resultUrl(selected)} controls autoPlay className="w-full" />
|
||||
</div>
|
||||
) : (
|
||||
<a href={resultUrl(selected)} target="_blank" rel="noreferrer" className="w-full h-full flex items-center justify-center">
|
||||
{/* eslint-disable-next-line @next/next/no-img-element */}
|
||||
<img key={selected.path} src={resultUrl(selected)} alt={selected.prompt} className="max-w-full max-h-full rounded object-contain" />
|
||||
</a>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* history strip */}
|
||||
<div className="h-20 shrink-0 m-2 flex items-center gap-2">
|
||||
<div className="flex-1 min-w-0 h-full flex items-center gap-2 overflow-x-auto overflow-y-hidden">
|
||||
{results.length === 0 ? (
|
||||
<div className="text-[11px] text-gray-600 px-1">History</div>
|
||||
) : (
|
||||
results.map(r => {
|
||||
const isSel = selected?.path === r.path && !showPreview;
|
||||
return (
|
||||
<div key={`${r.request_id}-${r.path}`} className="relative shrink-0 group">
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
setSelected(r);
|
||||
setShowPreview(false);
|
||||
}}
|
||||
title={`${r.arch} · seed ${r.seed} · ${r.prompt}`}
|
||||
className={`h-16 w-16 rounded-md overflow-hidden border-2 ${isSel ? 'border-blue-500' : 'border-transparent hover:border-gray-600'} bg-gray-800 block`}
|
||||
>
|
||||
{r.kind === 'audio' ? (
|
||||
<div className="w-full h-full flex items-center justify-center text-gray-400 text-[10px]">audio</div>
|
||||
) : r.kind === 'video' ? (
|
||||
<video src={resultUrl(r)} muted className="w-full h-full object-cover" />
|
||||
) : (
|
||||
// eslint-disable-next-line @next/next/no-img-element
|
||||
<img src={resultUrl(r)} alt="" className="w-full h-full object-cover" />
|
||||
)}
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
onClick={e => {
|
||||
e.stopPropagation();
|
||||
deleteResults([r]);
|
||||
}}
|
||||
title="Delete"
|
||||
className="absolute -top-1 -right-1 w-5 h-5 rounded-full bg-gray-900 border border-gray-600 text-gray-300 hover:bg-red-700 hover:text-white items-center justify-center hidden group-hover:flex"
|
||||
>
|
||||
<X className="w-3 h-3" />
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
})
|
||||
)}
|
||||
</div>
|
||||
{results.length > 0 && (
|
||||
<button
|
||||
type="button"
|
||||
onClick={deleteAll}
|
||||
title="Delete all results"
|
||||
className="shrink-0 h-16 w-10 rounded-md border border-gray-800 text-gray-500 hover:text-red-400 hover:border-red-900 flex items-center justify-center"
|
||||
>
|
||||
<Trash2 className="w-4 h-4" />
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* ---- right: collapsible model + prompt sidebar ---- */}
|
||||
<div className={`relative shrink-0 transition-[width] duration-200 ${sidebarOpen ? 'w-[380px]' : 'w-0'}`}>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setSidebarOpen(o => !o)}
|
||||
className="absolute top-1/2 -translate-y-1/2 -left-4 z-20 w-4 h-16 rounded-l-md bg-gray-800 border border-r-0 border-gray-700 text-gray-300 hover:bg-gray-700 flex items-center justify-center"
|
||||
title={sidebarOpen ? 'Collapse settings' : 'Expand settings'}
|
||||
>
|
||||
{sidebarOpen ? <ChevronRight className="w-3.5 h-3.5" /> : <ChevronLeft className="w-3.5 h-3.5" />}
|
||||
</button>
|
||||
<div className={`h-full flex flex-col border-l border-gray-800 bg-gray-900/60 ${sidebarOpen ? '' : 'hidden'}`}>
|
||||
<div className="flex-1 min-h-0 overflow-y-auto p-3 space-y-3">
|
||||
|
||||
<Card title="Model" open={cardOpen('model')} onToggle={() => toggleCard('model')}>
|
||||
<SelectInput label="Architecture" value={arch} onChange={v => applyArch(v as string)} options={archOptions} disabled={!ready} />
|
||||
<CreatableSelectInput
|
||||
label="Name or Path"
|
||||
value={model.name_or_path || ''}
|
||||
onChange={(v: string | null) => 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 || {}) && (
|
||||
<TextInput label="Extras Name or Path" value={model.extras_name_or_path || ''} onChange={v => setModel(m => ({ ...m, extras_name_or_path: v }))} />
|
||||
)}
|
||||
<div className="grid grid-cols-2 gap-3">
|
||||
<SelectInput
|
||||
label="Quantize Transformer"
|
||||
value={model.quantize ? model.qtype || 'convrot8' : ''}
|
||||
onChange={v => setModel(m => ({ ...m, quantize: !!v, qtype: v || undefined }))}
|
||||
options={[{ value: '', label: 'None (bf16)' }, ...qtypeOptions]}
|
||||
/>
|
||||
<SelectInput
|
||||
label="Quantize Text Encoder"
|
||||
value={model.quantize_te ? model.qtype_te || 'convrot8' : ''}
|
||||
onChange={v => setModel(m => ({ ...m, quantize_te: !!v, qtype_te: v || undefined }))}
|
||||
options={[{ value: '', label: 'None (bf16)' }, ...qtypeOptions]}
|
||||
/>
|
||||
</div>
|
||||
<div className="grid grid-cols-2 gap-3 mt-3">
|
||||
<Checkbox label="Low VRAM" checked={!!model.low_vram} onChange={v => setModel(m => ({ ...m, low_vram: v }))} />
|
||||
<Checkbox label="Layer offloading" checked={!!model.layer_offloading} onChange={v => setModel(m => ({ ...m, layer_offloading: v }))} />
|
||||
</div>
|
||||
</Card>
|
||||
|
||||
<Card title="Prompt" open={cardOpen('prompt')} onToggle={() => toggleCard('prompt')}>
|
||||
<TextAreaInput label="Prompt" value={sample.prompt || ''} onChange={v => setSample(s => ({ ...s, prompt: v }))} placeholder={modality === 'audio' ? 'song description, lyrics, bpm…' : 'describe what to generate'} />
|
||||
{(sample.guidance_scale ?? 4) > 1 && (
|
||||
<TextAreaInput label="Negative prompt" value={sample.negative_prompt || ''} onChange={v => setSample(s => ({ ...s, negative_prompt: v }))} />
|
||||
)}
|
||||
<div className="grid grid-cols-2 sm:grid-cols-3 gap-3">
|
||||
<NumberInput label="Width" value={sample.width ?? 1024} onChange={v => setSample(s => ({ ...s, width: v }))} min={64} max={4096} />
|
||||
<NumberInput label="Height" value={sample.height ?? 1024} onChange={v => setSample(s => ({ ...s, height: v }))} min={64} max={4096} />
|
||||
<NumberInput label="Steps" value={sample.num_inference_steps ?? 25} onChange={v => setSample(s => ({ ...s, num_inference_steps: v }))} min={1} max={200} />
|
||||
<NumberInput label="Guidance" value={sample.guidance_scale ?? 4} onChange={v => setSample(s => ({ ...s, guidance_scale: v }))} min={1} max={30} />
|
||||
<NumberInput label="Seed (-1 random)" value={sample.seed ?? -1} onChange={v => setSample(s => ({ ...s, seed: v }))} min={-1} max={4294967295} />
|
||||
{modality === 'video' && (
|
||||
<>
|
||||
<NumberInput label="Frames" value={sample.num_frames ?? 33} onChange={v => setSample(s => ({ ...s, num_frames: v }))} min={1} max={1000} />
|
||||
<NumberInput label="FPS" value={sample.fps ?? 16} onChange={v => setSample(s => ({ ...s, fps: v }))} min={1} max={60} />
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
</Card>
|
||||
|
||||
</div>
|
||||
{/* always-visible action bar */}
|
||||
<div className="shrink-0 p-3 border-t border-gray-800 bg-gray-900 flex gap-2">
|
||||
<Button
|
||||
onClick={generate}
|
||||
disabled={!ready || running || !model.name_or_path}
|
||||
className="flex-1 px-3 py-2 rounded-md bg-blue-700 hover:bg-blue-600 disabled:opacity-40 text-white flex items-center justify-center gap-2"
|
||||
title="Ctrl/Cmd + Enter"
|
||||
>
|
||||
{running ? <Loader2 className="w-4 h-4 animate-spin" /> : <Sparkles className="w-4 h-4" />} Generate
|
||||
</Button>
|
||||
{running && (
|
||||
<Button onClick={cancel} className="px-3 py-2 rounded-md bg-gray-700 hover:bg-gray-600 text-white">
|
||||
Cancel
|
||||
</Button>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</MainContent>
|
||||
<GenerateFooter jobId={engineJobId} status={footerStatus} busy={running || isStarting} progress={running ? progress : null} />
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
export default function GeneratePage() {
|
||||
return (
|
||||
<Suspense fallback={null}>
|
||||
<GeneratePageInner />
|
||||
</Suspense>
|
||||
);
|
||||
}
|
||||
@@ -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 (
|
||||
<>
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
};
|
||||
@@ -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;
|
||||
|
||||
@@ -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 = (
|
||||
<>
|
||||
<small className="opacity-50">INFERENCE ENGINE: </small> {row.name}
|
||||
</>
|
||||
);
|
||||
}
|
||||
if (row.job_type === 'caption') {
|
||||
let splits = row.job_ref.split(/[/\\]/);
|
||||
const datasetPath = `${splits[splits.length - 1]}`;
|
||||
|
||||
@@ -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.icon className="w-5 h-5 mr-3" />
|
||||
{item.name}
|
||||
{'badge' in item && item.badge && (
|
||||
<span className="ml-2 px-1.5 py-0.5 text-[10px] leading-none uppercase tracking-wide rounded bg-blue-600/20 text-blue-400 border border-blue-500/40">
|
||||
{item.badge}
|
||||
</span>
|
||||
)}
|
||||
</Link>
|
||||
</li>
|
||||
))}
|
||||
|
||||
@@ -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<HTMLDivElement | null>(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 (
|
||||
<div className="absolute left-0 right-0 bottom-0 z-30 border-t border-gray-800 bg-gray-900 shadow-[0_-4px_16px_rgba(0,0,0,0.4)]">
|
||||
{/* progress hairline */}
|
||||
{pct !== null && (
|
||||
<div className="h-0.5 bg-gray-800">
|
||||
<div className="h-0.5 bg-blue-500 transition-all" style={{ width: `${pct}%` }} />
|
||||
</div>
|
||||
)}
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setOpen(o => !o)}
|
||||
className="w-full flex items-center gap-2 px-3 h-9 text-xs text-gray-300 hover:bg-gray-800/60"
|
||||
title={open ? 'Hide log' : 'Show log'}
|
||||
>
|
||||
{busy ? <Loader2 className="w-3.5 h-3.5 animate-spin text-blue-400 shrink-0" /> : <span className="w-2 h-2 rounded-full bg-green-500 shrink-0" />}
|
||||
<span className="truncate flex-1 text-left font-mono">{status}</span>
|
||||
{progress?.total ? (
|
||||
<span className="text-gray-500 shrink-0">
|
||||
{progress.step}/{progress.total}
|
||||
</span>
|
||||
) : null}
|
||||
{open ? <ChevronDown className="w-4 h-4 shrink-0" /> : <ChevronUp className="w-4 h-4 shrink-0" />}
|
||||
</button>
|
||||
{open && (
|
||||
<div className="h-[40vh] bg-gray-950 border-t border-gray-800 relative">
|
||||
<div ref={logRef} onScroll={handleScroll} className="absolute inset-0 overflow-y-auto p-3 text-[11px] leading-4 text-gray-300 font-mono">
|
||||
{!jobId && <div className="text-gray-500">No engine job.</div>}
|
||||
{jobId && logStatus === 'loading' && !log && <div className="text-gray-500">Loading log…</div>}
|
||||
{jobId && logStatus === 'error' && <div className="text-rose-400">Error loading log</div>}
|
||||
{lines.map((line, i) => (
|
||||
<pre key={i} className="whitespace-pre-wrap break-all">
|
||||
{line}
|
||||
</pre>
|
||||
))}
|
||||
</div>
|
||||
{!follow && (
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
setFollow(true);
|
||||
if (logRef.current) logRef.current.scrollTop = logRef.current.scrollHeight;
|
||||
}}
|
||||
className="absolute bottom-3 right-4 px-2 py-1 rounded bg-gray-800 border border-gray-700 text-[11px] text-gray-300 hover:bg-gray-700"
|
||||
>
|
||||
Follow
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -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',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
@@ -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 <job folder>/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<EngineHandle[]> {
|
||||
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<EngineHandle | null> {
|
||||
const handles = await listEngines();
|
||||
if (jobId) return handles.find(h => h.jobId === jobId) || null;
|
||||
return handles.find(h => h.endpoint) || handles[0] || null;
|
||||
}
|
||||
Reference in new issue
Block a user