Improve ui and functionality

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
 

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)

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}"

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(

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 });
+}

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<string, string> => {
   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 (
     <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>
+      <button type="button" onClick={onToggle} className="w-full flex items-center justify-between gap-2 px-3 py-1.5 text-xs text-gray-200 font-semibold">
+        <span className="flex items-baseline gap-2 min-w-0">
+          <span>{title}</span>
+          {subtitle && <span className="text-[11px] font-normal text-gray-500 truncate">{subtitle}</span>}
+        </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 className={open ? 'px-3 pb-3' : 'hidden'}>{children}</div>
     </div>
   );
 }
@@ -100,6 +116,12 @@ function GeneratePageInner() {
   // 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 [stoppingSince, setStoppingSince] = useState<number | null>(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<number | null>(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<HTMLCanvasElement | null>(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<Set<string>>(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() {
         <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>
+            <>
+              <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>
+              {stoppingForMs > 8000 && (
+                <Button onClick={forceStop} className="px-2 py-1 rounded-md bg-red-700 hover:bg-red-600 text-white flex items-center gap-1" title="Kill the engine process">
+                  <OctagonX className="w-4 h-4" /> Force stop
+                </Button>
+              )}
+            </>
           ) : ready ? (
             <>
               <span className="text-green-400">Engine running</span>
@@ -488,10 +691,20 @@ function GeneratePageInner() {
                   : '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-2">
+                  <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>
+                  {engineStatus?.engine?.status === 'queued' && (
+                    <div className="text-[11px] text-gray-500">Waiting for the GPU queue to pick the job up. If the queue is stopped, start it from the Queue page or cancel here.</div>
+                  )}
+                  {startingSince !== null && now - startingSince > 5000 && (
+                    <Button onClick={cancelStart} className="w-full px-3 py-1.5 rounded-md bg-gray-800 hover:bg-gray-700 text-gray-200 text-sm">
+                      Cancel
+                    </Button>
+                  )}
+                </div>
               ) : (
                 <div className="flex flex-col gap-3">
                   {!isMac() && (
@@ -537,6 +750,17 @@ function GeneratePageInner() {
                     Cancel
                   </Button>
                 )}
+                {!running && busyRequestId && (
+                  <>
+                    <span className="text-amber-300 truncate">
+                      Engine busy: {health.current.arch}
+                      {health.current.total_steps ? ` ${health.current.step}/${health.current.total_steps}` : ''}
+                    </span>
+                    <Button onClick={cancelBusy} 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">
@@ -552,9 +776,6 @@ function GeneratePageInner() {
                   <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 && (
@@ -581,7 +802,20 @@ function GeneratePageInner() {
             {/* 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 ? (
+                {running && (
+                  <button
+                    type="button"
+                    onClick={() => setShowPreview(true)}
+                    title="Back to the live preview"
+                    className={`h-16 w-16 shrink-0 rounded-md overflow-hidden border-2 ${showPreview ? 'border-blue-500' : 'border-transparent hover:border-gray-600'} bg-gray-800 relative flex items-center justify-center`}
+                  >
+                    <canvas ref={thumbRef} className="max-w-full max-h-full" />
+                    <span className="absolute bottom-0.5 right-0.5 rounded bg-gray-950/80 p-0.5">
+                      <Loader2 className="w-3 h-3 animate-spin text-blue-400" />
+                    </span>
+                  </button>
+                )}
+                {results.length === 0 && !running ? (
                   <div className="text-[11px] text-gray-600 px-1">History</div>
                 ) : (
                   results.map(r => {
@@ -636,7 +870,7 @@ function GeneratePageInner() {
           </div>
 
           {/* ---- right: collapsible model + prompt sidebar ---- */}
-          <div className={`relative shrink-0 transition-[width] duration-200 ${sidebarOpen ? 'w-[380px]' : 'w-0'}`}>
+          <div className={`relative shrink-0 transition-[width] duration-200 ${sidebarOpen ? 'w-[340px]' : 'w-0'}`}>
             <button
               type="button"
               onClick={() => setSidebarOpen(o => !o)}
@@ -646,9 +880,9 @@ function GeneratePageInner() {
               {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">
+              <div className="compact-form flex-1 min-h-0 overflow-y-auto p-2 space-y-2">
 
-                <Card title="Model" open={cardOpen('model')} onToggle={() => toggleCard('model')}>
+                <Card title="Model" subtitle={entry?.label || arch} 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"
@@ -660,7 +894,7 @@ function GeneratePageInner() {
                   {'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">
+                  <div className="grid grid-cols-2 gap-2">
                     <SelectInput
                       label="Quantize Transformer"
                       value={model.quantize ? model.qtype || 'convrot8' : ''}
@@ -674,10 +908,30 @@ function GeneratePageInner() {
                       options={[{ value: '', label: 'None (bf16)' }, ...qtypeOptions]}
                     />
                   </div>
-                  <div className="grid grid-cols-2 gap-3 mt-3">
+                  <div className="grid grid-cols-2 gap-2 mt-2">
                     <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>
+                  {model.layer_offloading && (
+                    <div className="mt-2 space-y-1">
+                      <SliderInput
+                        label="Transformer Offload %"
+                        value={Math.round((model.layer_offloading_transformer_percent ?? 1) * 100)}
+                        onChange={v => setModel(m => ({ ...m, layer_offloading_transformer_percent: v * 0.01 }))}
+                        min={0}
+                        max={100}
+                        step={1}
+                      />
+                      <SliderInput
+                        label="Text Encoder Offload %"
+                        value={Math.round((model.layer_offloading_text_encoder_percent ?? 1) * 100)}
+                        onChange={v => setModel(m => ({ ...m, layer_offloading_text_encoder_percent: v * 0.01 }))}
+                        min={0}
+                        max={100}
+                        step={1}
+                      />
+                    </div>
+                  )}
                 </Card>
     
                 <Card title="Prompt" open={cardOpen('prompt')} onToggle={() => toggleCard('prompt')}>
@@ -685,7 +939,7 @@ function GeneratePageInner() {
                   {(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">
+                  <div className="grid grid-cols-3 gap-2">
                     <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} />
@@ -702,11 +956,11 @@ function GeneratePageInner() {
     
               </div>
               {/* always-visible action bar */}
-              <div className="shrink-0 p-3 border-t border-gray-800 bg-gray-900 flex gap-2">
+              <div className="shrink-0 p-2 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"
+                  className="flex-1 px-3 py-1.5 rounded-md bg-blue-700 hover:bg-blue-600 disabled:opacity-40 text-white text-sm 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

ui/src/app/globals.css

@@ -129,3 +129,25 @@ body {
     @apply text-sm text-gray-200;
   }
 }
+
+/* Generate page sidebar: tighter version of the shared form inputs */
+.compact-form label {
+  margin-top: 0.25rem !important;
+  margin-bottom: 0.125rem !important;
+  font-size: 11px !important;
+}
+.compact-form input[type='text'],
+.compact-form input[type='number'],
+.compact-form input[type='password'],
+.compact-form textarea {
+  padding-top: 0.25rem !important;
+  padding-bottom: 0.25rem !important;
+  font-size: 12px !important;
+}
+.compact-form textarea {
+  min-height: 3.5rem;
+}
+.compact-form [class*='control'] {
+  min-height: 30px !important;
+  font-size: 12px !important;
+}

ui/src/server/inferenceEngine.ts

@@ -52,18 +52,24 @@ 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' },
-  });
+  // every inference job: a row already marked stopped can still be a live
+  // process finishing (or ignoring) its shutdown, which engine.json's pid tells us
+  const jobs = await prisma.job.findMany({ where: { job_type: 'inference' }, 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,
-  }));
+  const handles: EngineHandle[] = [];
+  for (const job of jobs) {
+    const active = ['running', 'queued', 'stopping'].includes(job.status);
+    const endpoint = readEndpoint(path.join(trainingFolder, job.name));
+    if (!active && !endpoint) continue;
+    handles.push({
+      jobId: job.id,
+      jobName: job.name,
+      gpuIds: job.gpu_ids,
+      // process still alive after the row was stopped: report it as stopping
+      status: !active && endpoint ? 'stopping' : job.status,
+      endpoint,
+    });
+  }
   cache = { ts: Date.now(), handles };
   return handles;
 }