From 6e1b39c7a84cfdb48e7e91e09767ab2f4aa51a15 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 10 Sep 2026 15:44:00 -0600 Subject: [PATCH] Improve ui and functionality --- .../minimax_h3/src/pipeline.py | 5 + .../inference_engine/InferenceEngine.py | 21 +- .../inference_engine/engine.py | 132 ++++- .../inference_engine/server.py | 24 +- ui/src/app/api/jobs/[jobID]/kill/route.ts | 45 ++ ui/src/app/generate/page.tsx | 456 ++++++++++++++---- ui/src/app/globals.css | 22 + ui/src/server/inferenceEngine.ts | 28 +- 8 files changed, 610 insertions(+), 123 deletions(-) create mode 100644 ui/src/app/api/jobs/[jobID]/kill/route.ts diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py b/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py index 2e081f9..3e48c56 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py +++ b/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py @@ -249,6 +249,11 @@ class MiniMaxH3Pipeline: v_audio = audio_pred[:, layout.num_condition_audio_rows :].float() denoised_v = video_rows + sv * v_video + # inference engine preview: the video x0 estimate as (1, C, T, H, W); + # the audio stream is not previewed + emit = getattr(model, "_emit_sample_step", None) + if emit is not None and getattr(model, "sample_step_hook", None) is not None: + emit(unpatchify_video_tokens(denoised_v, t_lat, h_lat, w_lat), i, num_steps) ratio_v = sv_next / sv video_rows = ratio_v * video_rows + (1.0 - ratio_v) * denoised_v diff --git a/extensions_built_in/inference_engine/InferenceEngine.py b/extensions_built_in/inference_engine/InferenceEngine.py index b63a715..8665dc1 100644 --- a/extensions_built_in/inference_engine/InferenceEngine.py +++ b/extensions_built_in/inference_engine/InferenceEngine.py @@ -76,13 +76,28 @@ class InferenceEngine(BaseExtensionProcess): 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)) + def _on_signal(signum, frame): + # first signal: stop after the current request; if a request (or a + # model load) is in flight, interrupt it right here on the main + # thread so a long load cannot hold the shutdown hostage + self.engine.stop("stopped") + if self.engine.current is not None or self.engine.loading: + raise KeyboardInterrupt() + for sig in (signal.SIGINT, signal.SIGTERM): try: - signal.signal(sig, lambda *_: self.engine.stop("stopped")) + signal.signal(sig, _on_signal) except Exception: pass + def _on_db_stop(reason): + self.engine.stop(reason) + # deliver a real SIGINT to the main thread (never os.kill: on + # Windows that is TerminateProcess) + signal.raise_signal(signal.SIGINT) + + self.ui.start_stop_watcher(_on_db_stop) + if self.engine_config.default_model: with torch.no_grad(): self.engine.preload(dict(self.engine_config.default_model)) @@ -133,7 +148,7 @@ class InferenceEngine(BaseExtensionProcess): self.engine.unload() except Exception: pass - self.ui.update_status("queued" if reason == "queued" else "stopped", "Engine stopped") + self.ui.update_status("stopped", "Engine stopped") def on_error(self, e: Exception): super().on_error(e) diff --git a/extensions_built_in/inference_engine/engine.py b/extensions_built_in/inference_engine/engine.py index da3df21..d90d744 100644 --- a/extensions_built_in/inference_engine/engine.py +++ b/extensions_built_in/inference_engine/engine.py @@ -41,7 +41,13 @@ class GenerationJob: 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() + # every frame so far (latents: only the newest kept) so a client can + # attach late / reattach after a reload and see the same stream + self.history: List[bytes] = [] + self._latent_index: Optional[int] = None + self._subscribers: List["queue.Queue[Optional[bytes]]"] = [] + self._sub_lock = threading.Lock() + self.finished = False self.cancel = threading.Event() self.status = "queued" self.created_at = time.time() @@ -54,12 +60,41 @@ class GenerationJob: 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)) + data = encode_frame(header, payload) + with self._sub_lock: + if type_ == "latent": + if self._latent_index is not None: + self.history[self._latent_index] = data + else: + self._latent_index = len(self.history) + self.history.append(data) + else: + self.history.append(data) + for q in self._subscribers: + q.put(data) + + def subscribe(self, replay: bool = True) -> "queue.Queue[Optional[bytes]]": + """A queue that receives this job's frames (history first when + replay) and END once the job is finished.""" + q: "queue.Queue[Optional[bytes]]" = queue.Queue() + with self._sub_lock: + if replay: + for data in self.history: + q.put(data) + if self.finished: + q.put(END) + else: + self._subscribers.append(q) + return q def finish(self): self.finished_at = time.time() self.emit("end", status=self.status) - self.frames.put(END) + with self._sub_lock: + self.finished = True + for q in self._subscribers: + q.put(END) + self._subscribers = [] def info(self) -> dict: return { @@ -81,6 +116,15 @@ 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) +# settings that can change on a loaded holder without reloading any weights: +# the memory manager is detached and re-attached at the new fractions +OFFLOAD_KEYS = ("layer_offloading", "layer_offloading_transformer_percent", "layer_offloading_text_encoder_percent") + + +def _without(model: dict, keys) -> dict: + return {k: v for k, v in model.items() if k not in keys} + + class Engine: def __init__( self, @@ -102,6 +146,7 @@ class Engine: self.current: Optional[GenerationJob] = None self.stop_event = threading.Event() self.stop_reason = "stopped" + self.loading = False self.holder = None self.holder_key: Optional[str] = None self.holder_model: Optional[dict] = None @@ -118,9 +163,18 @@ class Engine: job = GenerationJob(model, sample, stream) with self._lock: self.jobs[job.request_id] = job + # bounded memory: drop the oldest finished jobs' frame history + finished = [j for j in self.jobs.values() if j.finished] + if len(finished) > 100: + finished.sort(key=lambda j: j.created_at) + for old in finished[: len(finished) - 100]: + self.jobs.pop(old.request_id, None) self.queue.put(job) return job + def get_job(self, request_id: str) -> Optional[GenerationJob]: + return self.jobs.get(request_id) + def cancel(self, request_id: str) -> bool: job = self.jobs.get(request_id) if job is None: @@ -197,6 +251,7 @@ class Engine: run_forever.""" job = GenerationJob(dict(model), {}, {"latents": "none"}) job.status = "running" + self.current = job try: self._ensure_holder(job) job.status = "done" @@ -205,6 +260,7 @@ class Engine: job.error = f"{type(e).__name__}: {e}" traceback.print_exc() finally: + self.current = None job.finish() return job @@ -268,11 +324,21 @@ class Engine: key = _model_key(model) if self.holder is not None and key == self.holder_key: return self.holder + if ( + self.holder is not None + and _model_key(_without(model, OFFLOAD_KEYS)) == _model_key(_without(self.holder_model, OFFLOAD_KEYS)) + ): + # same weights, different offloading: re-stage in place + self._apply_offload(job, model) + self.holder_model = model + self.holder_key = key + return self.holder if self.holder is not None: self._drop_holder() self.pool.begin_request() ComponentPool.current = self.pool + self.loading = True try: try: holder, is_legacy = self._build_holder(job, model) @@ -287,6 +353,7 @@ class Engine: holder, is_legacy = self._build_holder(job, model) finally: ComponentPool.current = None + self.loading = False freed = self.pool.release_unused() if freed: job.emit("status", message=f"Released {freed / 1e9:.1f} GB of components not used by {model['arch']}") @@ -297,6 +364,58 @@ class Engine: self.holder_is_legacy = is_legacy return holder + def _apply_offload(self, job: GenerationJob, model: dict): + """Detach/re-attach layer offloading on the loaded holder's components + at the requested fractions. No weights are reloaded.""" + from toolkit.memory_management import MemoryManager + + holder = self.holder + mc = holder.model_config + enabled = bool(model.get("layer_offloading", False)) + t_pct = float(model.get("layer_offloading_transformer_percent", 1.0)) + te_pct = float(model.get("layer_offloading_text_encoder_percent", 1.0)) + mc.layer_offloading = enabled + mc.layer_offloading_transformer_percent = t_pct + mc.layer_offloading_text_encoder_percent = te_pct + device = torch.device(self.device) + + targets = [] + model_module = getattr(holder, "model", None) + subs = [getattr(model_module, n, None) for n in ("transformer_1", "transformer_2")] + if model_module is not None and all(m is not None for m in subs): + targets += [(m, t_pct) for m in subs] + elif model_module is not None: + targets.append((model_module, t_pct)) + tes = getattr(holder, "text_encoder", None) + for te in tes if isinstance(tes, list) else [tes]: + targets.append((te, te_pct)) + + desc = f"offload {'on' if enabled else 'off'} (transformer {t_pct:.0%}, text encoder {te_pct:.0%})" + job.emit("status", message=f"Applying {desc}") + self.status_cb(f"Applying {desc}") + changed = 0 + for module, pct in targets: + if not isinstance(module, torch.nn.Module) or type(module).__name__.startswith("Fake"): + continue + if next(module.parameters(), None) is None: + continue + if hasattr(module, "_memory_manager"): + MemoryManager.detach(module) + if enabled and pct > 0: + get_ignore = getattr(module, "get_offload_ignore_modules", None) + ignore = get_ignore() if callable(get_ignore) else None + MemoryManager.attach(module, device, offload_percent=pct, ignore_modules=list(ignore or [])) + else: + module.to(device) + # keep the pool's policy record in sync so a later identical + # request is a no-op rather than a re-attach + policy = getattr(module, "_aitk_policy", None) + if isinstance(policy, dict): + policy["offload"] = pct if enabled else 0.0 + changed += 1 + flush() + job.emit("status", message=f"Applied {desc} to {changed} components") + def _build_holder(self, job: GenerationJob, model: dict): model_config = self._build_model_config(model) ModelClass = get_model_class(model_config) @@ -473,12 +592,15 @@ class Engine: job.emit("result", **result) job.status = "done" self.stats["generations"] += 1 - except GenerationCancelled: + except (GenerationCancelled, KeyboardInterrupt) as e: job.status = "cancelled" job.error = "cancelled" self.stats["cancelled"] += 1 - job.emit("error", message="cancelled", cancelled=True) + job.emit("error", message="cancelled" if isinstance(e, GenerationCancelled) else "engine stopping", cancelled=True) flush() + if isinstance(e, KeyboardInterrupt): + job.finish() + raise except Exception as e: job.status = "error" job.error = f"{type(e).__name__}: {e}" diff --git a/extensions_built_in/inference_engine/server.py b/extensions_built_in/inference_engine/server.py index b1c4106..3ec81b6 100644 --- a/extensions_built_in/inference_engine/server.py +++ b/extensions_built_in/inference_engine/server.py @@ -98,23 +98,41 @@ def create_app(engine: Engine, token: Optional[str] = None) -> FastAPI: if body.get("wait"): # plain JSON mode for scripts: block until the job is done, drop frames job = engine.submit(model, sample, {"latents": "none"}) + q = job.subscribe(replay=False) while True: - frame = await asyncio.to_thread(job.frames.get) + frame = await asyncio.to_thread(q.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")) + # a dropped connection (page reload) must not cancel: the client + # reattaches via /stream/{id}; cancelling is explicit (/cancel) + return _stream_job(job, cancel_on_disconnect=False) + + @app.get("/stream/{request_id}") + async def stream(request: Request, request_id: str): + """(Re)attach to a request's frame stream: replays everything so far + (status/progress/result frames and the newest latent) then follows + live until the end frame. Disconnecting does not cancel the job.""" + check_auth(request) + job = engine.get_job(request_id) + if job is None: + raise HTTPException(status_code=404, detail="unknown request") + return _stream_job(job, cancel_on_disconnect=False) + + def _stream_job(job, cancel_on_disconnect: bool): + q = job.subscribe(replay=True) async def frames(): try: while True: - frame = await asyncio.to_thread(job.frames.get) + frame = await asyncio.to_thread(q.get) if frame is END: break yield frame finally: - if job.status in ("queued", "running"): + if cancel_on_disconnect and job.status in ("queued", "running"): engine.cancel(job.request_id) return StreamingResponse( diff --git a/ui/src/app/api/jobs/[jobID]/kill/route.ts b/ui/src/app/api/jobs/[jobID]/kill/route.ts new file mode 100644 index 0000000..4939778 --- /dev/null +++ b/ui/src/app/api/jobs/[jobID]/kill/route.ts @@ -0,0 +1,45 @@ +import { NextRequest, NextResponse } from 'next/server'; +import prisma from '@/server/prisma'; +import { listEngines } from '@/server/inferenceEngine'; +import { exec } from 'child_process'; +import { promisify } from 'util'; + +const execAsync = promisify(exec); +const isWindows = process.platform === 'win32'; + +// Force stop: SIGKILL the job process (tree-kill on Windows) and mark the row +// stopped. For jobs that ignored the graceful stop (a hung model load). +export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) { + const { jobID } = await params; + const job = await prisma.job.findUnique({ where: { id: jobID } }); + // the pid from the row, else (inference engines) the one the live process + // published in its engine.json + let pid: number | null = job?.pid ?? null; + if (pid == null) { + const engine = (await listEngines(true)).find(e => e.jobId === jobID); + pid = engine?.endpoint?.pid ?? null; + } + if (!job && pid == null) { + return NextResponse.json({ error: 'Job not found' }, { status: 404 }); + } + let killed = false; + if (pid != null) { + try { + if (isWindows) { + await execAsync(`taskkill /PID ${pid} /T /F`, { windowsHide: true }); + } else { + process.kill(pid, 'SIGKILL'); + } + killed = true; + } catch (e: any) { + console.warn(`kill ${jobID} pid ${pid}: ${e?.message || e}`); + } + } + if (job) { + await prisma.job.update({ + where: { id: jobID }, + data: { stop: true, status: 'stopped', info: 'Job force stopped', pid: null }, + }); + } + return NextResponse.json({ ok: true, killed }); +} diff --git a/ui/src/app/generate/page.tsx b/ui/src/app/generate/page.tsx index a2c1e49..a5a3159 100644 --- a/ui/src/app/generate/page.tsx +++ b/ui/src/app/generate/page.tsx @@ -3,10 +3,10 @@ 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 { ChevronDown, ChevronLeft, ChevronRight, Loader2, OctagonX, 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 { Checkbox, CreatableSelectInput, NumberInput, SelectInput, SliderInput, TextAreaInput, TextInput } from '@/components/formInputs'; import { apiClient } from '@/utils/api'; import { startJob, stopJob } from '@/utils/jobs'; import { startQueue } from '@/utils/queue'; @@ -49,6 +49,7 @@ const qtypeOptions = [ // everything the user set on the page survives a reload const STORAGE_KEY = 'aitk_generate_page'; +const ACTIVE_REQUEST_KEY = 'aitk_generate_active_request'; interface PersistedState { arch?: string; model?: { [key: string]: any }; @@ -73,14 +74,29 @@ const authHeaders = (): Record => { return token ? { Authorization: `Bearer ${token}` } : {}; }; -function Card({ title, open, onToggle, children }: { title: string; open: boolean; onToggle: () => void; children: React.ReactNode }) { +function Card({ + title, + subtitle, + open, + onToggle, + children, +}: { + title: string; + subtitle?: string; + open: boolean; + onToggle: () => void; + children: React.ReactNode; +}) { return (
- -
{children}
+
{children}
); } @@ -100,6 +116,12 @@ function GeneratePageInner() { // the engine is actually up (or the job row reports a failure) const [startingJobId, setStartingJobId] = useState(null); const [stoppingJobId, setStoppingJobId] = useState(null); + const [stoppingSince, setStoppingSince] = useState(null); + const [now, setNow] = useState(Date.now()); + useEffect(() => { + const t = setInterval(() => setNow(Date.now()), 1000); + return () => clearInterval(t); + }, []); const engineJobId = preferredJob || engineStatus?.engine?.jobId || null; const proxy = useCallback((path: string) => `/api/inference/${path}${engineJobId ? `${path.includes('?') ? '&' : '?'}job=${engineJobId}` : ''}`, [engineJobId]); @@ -125,6 +147,7 @@ function GeneratePageInner() { } if (stoppingJobId && !data.engines.some(e => e.jobId === stoppingJobId)) { setStoppingJobId(null); + setStoppingSince(null); } if (data.running) { return apiClient @@ -167,12 +190,28 @@ function GeneratePageInner() { } }; + // Stop = graceful stop, then a hard kill if the process is still around a + // few seconds later (a model load in progress can hold the shutdown up) const stopEngine = async () => { - if (!engineStatus?.engine) return; + const eng = engineStatus?.engine; + if (!eng) return; + const id = eng.jobId; setEngineBusy(true); + setStoppingJobId(id); + setStoppingSince(Date.now()); try { - await stopJob(engineStatus.engine.jobId); - setStoppingJobId(engineStatus.engine.jobId); + await stopJob(id).catch(e => console.warn('graceful stop failed, will kill', e?.message || e)); + const deadline = Date.now() + 4000; + let alive = true; + while (Date.now() < deadline) { + await new Promise(r => setTimeout(r, 750)); + const st = await apiClient.get(`/api/inference/status?job=${id}`).then(r => r.data).catch(() => null); + alive = !!st?.engines?.some((e: any) => e.jobId === id && (e.ready || e.status === 'stopping')); + if (!alive) break; + } + if (alive) { + await apiClient.get(`/api/jobs/${id}/kill`); + } } catch (e: any) { alert(`Failed to stop the engine: ${e?.response?.data?.error || e?.message || e}`); } finally { @@ -180,8 +219,43 @@ function GeneratePageInner() { } }; - const isStarting = engineBusy || !!startingJobId || (!!engineStatus?.engine && !engineStatus.running && !stoppingJobId); - const isStopping = !!stoppingJobId; + // a row that is stopped but whose process is still alive shows as 'stopping' from the status route + const isStopping = !!stoppingJobId || engineStatus?.engine?.status === 'stopping'; + const isStarting = !isStopping && (engineBusy || !!startingJobId || (!!engineStatus?.engine && !engineStatus.running)); + const stoppingForMs = stoppingSince ? now - stoppingSince : isStopping ? 999_999 : 0; + // cancel a start that is not getting anywhere: a queued job (no process) + // is just marked stopped; a launched one gets the normal stop + const [startingSince, setStartingSince] = useState(null); + useEffect(() => { + if (isStarting && startingSince === null) setStartingSince(Date.now()); + if (!isStarting && startingSince !== null) setStartingSince(null); + }, [isStarting, startingSince]); + const cancelStart = async () => { + const eng = engineStatus?.engine; + const id = startingJobId || eng?.jobId; + if (!id) return; + try { + if (!eng || eng.status === 'queued') { + await apiClient.get(`/api/jobs/${id}/mark_stopped`); + } else { + await stopJob(id); + } + } catch (e: any) { + alert(`Cancel failed: ${e?.response?.data?.error || e?.message || e}`); + } + setStartingJobId(null); + setStartingSince(null); + }; + const forceStop = async () => { + const id = stoppingJobId || engineStatus?.engine?.jobId; + if (!id) return; + if (!confirm('Force stop kills the engine process immediately. Continue?')) return; + try { + await apiClient.get(`/api/jobs/${id}/kill`); + } catch (e: any) { + alert(`Force stop failed: ${e?.response?.data?.error || e?.message || e}`); + } + }; // ---- models: the same arch list the training UI uses (jobs/new/options.tsx) ---- const archs: GenerateDefaults[] = useMemo(() => modelArchs.map(getGenerateDefaults), []); @@ -302,15 +376,121 @@ function GeneratePageInner() { ctx.imageSmoothingEnabled = true; ctx.imageSmoothingQuality = 'high'; ctx.drawImage(scratch, 0, 0, w, h); + // the live tile in the history strip mirrors the preview + const thumb = thumbRef.current; + if (thumb) { + const size = 64; + const scaleT = Math.min(size / w, size / h); + const tw = Math.max(1, Math.round(w * scaleT)); + const th = Math.max(1, Math.round(h * scaleT)); + if (thumb.width !== tw || thumb.height !== th) { + thumb.width = tw; + thumb.height = th; + } + const tctx = thumb.getContext('2d'); + if (tctx) { + tctx.imageSmoothingEnabled = true; + tctx.drawImage(scratch, 0, 0, tw, th); + } + } }; + const thumbRef = useRef(null); useEffect(() => drawFrame(frameIdx), [frameIdx]); - const generate = async () => { - if (!ready || running) return; + // temporal latents auto-play: T latent frames span num_frames / fps + // seconds, so each latent frame shows for that fraction of the clip + const clipRef = useRef<{ num_frames: number; fps: number } | null>(null); + useEffect(() => { + if (frameCount <= 1 || !showPreview) return; + const clip = clipRef.current; + const seconds = clip && clip.num_frames > 1 && clip.fps > 0 ? clip.num_frames / clip.fps : frameCount / 4; + const msPerFrame = Math.max(30, (seconds * 1000) / frameCount); + const t = setInterval(() => setFrameIdx(i => (i + 1) % frameCount), msPerFrame); + return () => clearInterval(t); + }, [frameCount, showPreview]); + + // one consumer for both a fresh /generate response and a /stream/{id} + // reattach: everything the stage needs comes from the frames themselves + const consumeStream = async (res: Response, abort: AbortController) => { + let preview: PreviewInfo | null = null; + let startInfo: { arch: string; prompt: string } = { arch, prompt: sample.prompt || '' }; + // only the first latent of a run pulls the stage to the preview; after + // that the user may browse history and come back via the live tile + let shownPreview = false; + await readEngineFrames( + res, + ({ header, payload }) => { + requestIdRef.current = header.request_id; + switch (header.type) { + case 'start': + preview = header.preview; + startInfo = { arch: header.model?.arch || arch, prompt: header.sample?.prompt ?? '' }; + if (header.sample?.width && header.sample?.height) { + previewSizeRef.current = { width: header.sample.width, height: header.sample.height }; + } + clipRef.current = { num_frames: header.sample?.num_frames ?? 1, fps: header.sample?.fps ?? 16 }; + setStatusLine(`Generating with ${startInfo.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); + if (!shownPreview) { + shownPreview = true; + 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: startInfo.prompt, + arch: startInfo.arch, + }; + setResults(r => (r.some(x => x.path === item.path) ? 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, + ); + }; + + const beginRun = () => { setRunning(true); setError(null); setProgress(null); - setStatusLine('Submitting'); framesRef.current = []; setFrameCount(1); setFrameIdx(0); @@ -318,14 +498,31 @@ function GeneratePageInner() { // 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; + return abort; + }; + + const endRun = () => { + setRunning(false); + abortRef.current = null; + requestIdRef.current = null; + try { + localStorage.removeItem(ACTIVE_REQUEST_KEY); + } catch { + // storage unavailable + } + }; + + const generate = async () => { + if (!ready || running) return; + const abort = beginRun(); + setStatusLine('Submitting'); + previewSizeRef.current = sample.width && sample.height ? { width: sample.width, height: sample.height } : 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 }, + stream: { latents: 'raw', every_n_steps: 1, max_frames: 0 }, }; try { const res = await fetch(proxy('generate'), { @@ -335,74 +532,73 @@ function GeneratePageInner() { 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, - ); + // remembered so a reload can reattach to this run + const id = res.headers.get('x-request-id'); + if (id) { + try { + localStorage.setItem(ACTIVE_REQUEST_KEY, id); + } catch { + // storage unavailable + } + } + await consumeStream(res, abort); } catch (e: any) { if (e?.name !== 'AbortError') setError(e?.message || String(e)); } finally { - setRunning(false); - abortRef.current = null; + endRun(); + } + }; + + // attach to a request already running on the engine (after a reload, or a + // run started from another tab): replays its frames, then follows live + const attachedRef = useRef>(new Set()); + const attachTo = async (id: string) => { + if (running || attachedRef.current.has(id)) return; + attachedRef.current.add(id); + const abort = beginRun(); + requestIdRef.current = id; + setStatusLine('Reattaching…'); + try { + const res = await fetch(proxy(`stream/${id}`), { headers: authHeaders(), signal: abort.signal }); + if (!res.ok) throw new Error(res.status === 404 ? 'previous generation is gone' : `HTTP ${res.status}`); + await consumeStream(res, abort); + } catch (e: any) { + if (e?.name !== 'AbortError') setError(e?.message || String(e)); + } finally { + endRun(); + } + }; + const attachRef = useRef(attachTo); + attachRef.current = attachTo; + + // 1) the id we stored when we started a run (covers a job still queued) + useEffect(() => { + if (!ready || running) return; + let id: string | null = null; + try { + id = localStorage.getItem(ACTIVE_REQUEST_KEY); + } catch { + id = null; + } + if (id) attachRef.current(id); + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [ready]); + + // 2) whatever the engine says it is generating right now + const busyRequestId: string | null = health?.current?.request_id || null; + useEffect(() => { + if (ready && !running && busyRequestId) attachRef.current(busyRequestId); + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [ready, running, busyRequestId]); + + // engine busy with a run we could not attach to (e.g. an engine older than + // /stream): still let the user cancel it + const cancelBusy = async () => { + if (!busyRequestId) return; + try { + await apiClient.post(proxy(`cancel/${busyRequestId}`)); + } catch (e: any) { + alert(`Cancel failed: ${e?.response?.data?.error || e?.message || e}`); } }; @@ -459,9 +655,16 @@ function GeneratePageInner() {
{isStopping ? ( - + <> + + {stoppingForMs > 8000 && ( + + )} + ) : ready ? ( <> Engine running @@ -488,10 +691,20 @@ function GeneratePageInner() { : 'Start the inference engine on a GPU to generate images, video, and audio with any model.'}

{isStarting ? ( - +
+ + {engineStatus?.engine?.status === 'queued' && ( +
Waiting for the GPU queue to pick the job up. If the queue is stopped, start it from the Queue page or cancel here.
+ )} + {startingSince !== null && now - startingSince > 5000 && ( + + )} +
) : (
{!isMac() && ( @@ -537,6 +750,17 @@ function GeneratePageInner() { Cancel )} + {!running && busyRequestId && ( + <> + + Engine busy: {health.current.arch} + {health.current.total_steps ? ` ${health.current.step}/${health.current.total_steps}` : ''} + + + + )}
{progress?.total ? (
@@ -552,9 +776,6 @@ function GeneratePageInner() {
Your generation will appear here
)}
- {frameCount > 1 && showPreview && ( - setFrameIdx(parseInt(e.target.value))} className="absolute bottom-3 left-1/2 -translate-x-1/2 w-1/2 z-10" /> - )} {/* selected result replaces the preview */} {!showPreview && selected && ( @@ -581,7 +802,20 @@ function GeneratePageInner() { {/* history strip */}
- {results.length === 0 ? ( + {running && ( + + )} + {results.length === 0 && !running ? (
History
) : ( results.map(r => { @@ -636,7 +870,7 @@ function GeneratePageInner() {
{/* ---- right: collapsible model + prompt sidebar ---- */} -
+
-
+
- toggleCard('model')}> + toggleCard('model')}> applyArch(v as string)} options={archOptions} disabled={!ready} /> setModel(m => ({ ...m, extras_name_or_path: v }))} /> )} -
+
-
+
setModel(m => ({ ...m, low_vram: v }))} /> setModel(m => ({ ...m, layer_offloading: v }))} />
+ {model.layer_offloading && ( +
+ setModel(m => ({ ...m, layer_offloading_transformer_percent: v * 0.01 }))} + min={0} + max={100} + step={1} + /> + setModel(m => ({ ...m, layer_offloading_text_encoder_percent: v * 0.01 }))} + min={0} + max={100} + step={1} + /> +
+ )} toggleCard('prompt')}> @@ -685,7 +939,7 @@ function GeneratePageInner() { {(sample.guidance_scale ?? 4) > 1 && ( setSample(s => ({ ...s, negative_prompt: v }))} /> )} -
+
setSample(s => ({ ...s, width: v }))} min={64} max={4096} /> setSample(s => ({ ...s, height: v }))} min={64} max={4096} /> setSample(s => ({ ...s, num_inference_steps: v }))} min={1} max={200} /> @@ -702,11 +956,11 @@ function GeneratePageInner() {
{/* always-visible action bar */} -
+