Initial inference engine build out

toolkit/models/base_model.py

@@ -167,6 +167,9 @@ class BaseModel:
         self.invert_assistant_lora = False
         self._after_sample_img_hooks = []
         self._status_update_hooks = []
+        # inference engine: called as hook(step_index, num_steps, latents)
+        # after every scheduler step while generating samples
+        self.sample_step_hook = None
         self.is_transformer = False
 
         self.sample_prompts_cache = None
@@ -401,6 +404,18 @@ class BaseModel:
     def add_status_update_hook(self, func):
         self._status_update_hooks.append(func)
 
+    def _emit_sample_step(self, latents, step_index=None, num_steps=None):
+        """For holders whose sampling loop bypasses scheduler.step: report one
+        denoised latent to sample_step_hook (no-op when unset)."""
+        from toolkit.sample_step_hook import emit_sample_step
+
+        emit_sample_step(self, latents, step_index, num_steps)
+
+    def _install_sample_step_hooks(self, pipeline):
+        from toolkit.sample_step_hook import install_sample_step_hooks
+
+        return install_sample_step_hooks(self, pipeline)
+
     @torch.no_grad()
     def generate_images(
             self,
@@ -457,6 +472,8 @@ class BaseModel:
             except:
                 pass
 
+        unwrap_step_hooks = self._install_sample_step_hooks(pipeline)
+
         start_multiplier = 1.0
         if network is not None:
             start_multiplier = network.multiplier
@@ -517,6 +534,7 @@ class BaseModel:
 
                     if network is not None:
                         network.multiplier = gen_config.network_multiplier
+                    self._sample_step_index = 0
                     torch.manual_seed(gen_config.seed)
                     torch.cuda.manual_seed(gen_config.seed)
 
@@ -725,6 +743,7 @@ class BaseModel:
                 if self.adapter is not None and isinstance(self.adapter, ReferenceAdapter):
                     self.adapter.clear_memory()
 
+        unwrap_step_hooks()
         # clear pipeline and cache to reduce vram usage
         del pipeline
         torch.cuda.empty_cache()

toolkit/models/registry.py

@@ -0,0 +1,85 @@
+"""Per-arch capabilities + realistic defaults, one table shared by the
+inference engine (`GET /models`), its tests, and eventually the UI.
+
+Every arch registered through get_all_models()/LEGACY_ARCHS is served; an
+entry here adds defaults and capability flags on top. Archs without an entry
+get GENERIC. Keep `model` limited to ModelConfig kwargs and `sample` to
+GenerateImageConfig kwargs.
+"""
+
+from typing import Dict, List, Optional
+
+IMG = {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0}
+VID = {"width": 832, "height": 480, "num_inference_steps": 20, "guidance_scale": 5.0, "num_frames": 33, "fps": 16}
+
+GENERIC = {"modality": "image", "model": {}, "sample": dict(IMG)}
+
+# modality: image | video | audio
+# needs_control_image: the arch requires ctrl_img (edit / i2v models)
+# size_locked: only the default resolution is valid
+ARCH_REGISTRY: Dict[str, dict] = {
+    "zimage": {"modality": "image", "model": {"name_or_path": "Tongyi-MAI/Z-Image-Turbo"}, "sample": {**IMG, "num_inference_steps": 8, "guidance_scale": 1.0}},
+    "zimage_l2p": {"modality": "image", "model": {"name_or_path": "zhen-nan/L2P/model-1k-merge.safetensors", "extras_name_or_path": "Tongyi-MAI/Z-Image-Turbo", "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 8, "guidance_scale": 1.0}},
+    "qwen_image": {"modality": "image", "model": {"name_or_path": "Qwen/Qwen-Image", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 20}},
+    "qwen_image_edit": {"modality": "image", "model": {"name_or_path": "Qwen/Qwen-Image-Edit", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 20}, "needs_control_image": True},
+    "qwen_image_edit_plus": {"modality": "image", "model": {"name_or_path": "Qwen/Qwen-Image-Edit-2509", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 20}, "needs_control_image": True},
+    "krea2": {"modality": "image", "model": {"name_or_path": "krea/Krea-2-Turbo", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 8, "guidance_scale": 1.0}},
+    "boogu_image": {"modality": "image", "model": {"name_or_path": "Boogu/Boogu-Image-0.1-Base", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
+    "boogu_image_edit": {"modality": "image", "model": {"name_or_path": "Boogu/Boogu-Image-0.1-Edit", "quantize": True, "quantize_te": True}, "sample": dict(IMG), "needs_control_image": True},
+    "ernie_image": {"modality": "image", "model": {"name_or_path": "baidu/ERNIE-Image", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
+    "mageflow": {"modality": "image", "model": {"name_or_path": "microsoft/Mage-Flow-Base", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
+    "mageflow_edit": {"modality": "image", "model": {"name_or_path": "microsoft/Mage-Flow-Edit", "quantize": True, "quantize_te": True}, "sample": dict(IMG), "needs_control_image": True},
+    "ideogram4": {"modality": "image", "model": {"name_or_path": "ideogram-ai/ideogram-4-fp8", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
+    "hidream": {"modality": "image", "model": {"name_or_path": "HiDream-ai/HiDream-I1-Full", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 28, "guidance_scale": 5.0}},
+    "hidream_e1": {"modality": "image", "model": {"name_or_path": "HiDream-ai/HiDream-E1-1", "quantize": True, "quantize_te": True}, "sample": {"width": 768, "height": 768, "num_inference_steps": 28, "guidance_scale": 5.0}, "needs_control_image": True, "size_locked": True},
+    "hidream_o1": {"modality": "image", "model": {"name_or_path": "HiDream-ai/HiDream-O1-Image", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 28, "guidance_scale": 5.0}},
+    "anima": {"modality": "image", "model": {"name_or_path": "circlestone-labs/Anima-Base-v1.0-Diffusers"}, "sample": {**IMG, "guidance_scale": 4.5}},
+    "nucleus_image": {"modality": "image", "model": {"name_or_path": "NucleusAI/Nucleus-Image", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
+    "omnigen2": {"modality": "image", "model": {"name_or_path": "OmniGen2/OmniGen2", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
+    "chroma": {"modality": "image", "model": {"name_or_path": "lodestones/Chroma1-HD", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 26}},
+    "chroma_radiance": {"modality": "image", "model": {"name_or_path": "lodestones/Chroma1-Radiance", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 26}},
+    "zeta_chroma": {"modality": "image", "model": {"name_or_path": "lodestones/Zeta-Chroma/zeta-chroma-base-x0-pixel-dino-distance.safetensors", "extras_name_or_path": "Tongyi-MAI/Z-Image-Turbo", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
+    "flux_kontext": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.1-Kontext-dev", "quantize": True, "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 2.5}, "needs_control_image": True},
+    "flux2": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.2-dev", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
+    "flux2_klein_4b": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-4B", "quantize_te": True}, "sample": dict(IMG)},
+    "flux2_klein_9b": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-9B", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
+    "prx_pixel": {"modality": "image", "model": {"name_or_path": "Photoroom/prxpixel-t2i", "quantize_te": True}, "sample": dict(IMG)},
+    "f-lite": {"modality": "image", "model": {"name_or_path": "Freepik/F-Lite", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
+    "sd1": {"modality": "image", "model": {"name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 7.5}},
+    "sdxl": {"modality": "image", "model": {"name_or_path": "stabilityai/stable-diffusion-xl-base-1.0"}, "sample": {**IMG, "guidance_scale": 6.0}},
+    "wan21": {"modality": "video", "model": {"name_or_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"}, "sample": dict(VID)},
+    "wan21_i2v": {"modality": "video", "model": {"name_or_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", "quantize": True, "quantize_te": True}, "sample": dict(VID), "needs_control_image": True},
+    "wan22_5b": {"modality": "video", "model": {"name_or_path": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "quantize": True, "quantize_te": True}, "sample": {**VID, "fps": 24}},
+    "wan22_14b": {"modality": "video", "model": {"name_or_path": "ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True}, "sample": {**VID, "guidance_scale": 3.5}},
+    "wan22_14b_i2v": {"modality": "video", "model": {"name_or_path": "ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True}, "sample": {**VID, "guidance_scale": 3.5}, "needs_control_image": True},
+    "ltx2.3": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}},
+    "ltx2.5": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}},
+    "ace_step_15": {"modality": "audio", "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0}},
+}
+
+OUTPUT_EXT = {"image": "png", "video": "mp4", "audio": "mp3"}
+
+
+def get_arch_entry(arch: str) -> dict:
+    # variant suffixes ("zimage:turbo") share the base arch's entry
+    entry = ARCH_REGISTRY.get(arch) or ARCH_REGISTRY.get(arch.split(":")[0], GENERIC)
+    return {
+        "arch": arch,
+        "modality": entry.get("modality", "image"),
+        "model": dict(entry.get("model", {})),
+        "sample": dict(entry.get("sample", IMG)),
+        "needs_control_image": bool(entry.get("needs_control_image", False)),
+        "size_locked": bool(entry.get("size_locked", False)),
+    }
+
+
+def list_archs() -> List[str]:
+    from toolkit.util.get_model import LEGACY_ARCHS, get_all_models
+
+    archs = {m.arch for m in get_all_models() if getattr(m, "arch", None)}
+    archs |= LEGACY_ARCHS
+    return sorted(archs)
+
+
+def describe_archs(archs: Optional[List[str]] = None) -> List[dict]:
+    return [get_arch_entry(a) for a in (archs or list_archs())]

toolkit/models/v2/_mixin.py

@@ -255,6 +255,31 @@ class OstrisModelMixin:
             base_model.print_and_status_update if base_model is not None else print_acc
         )
 
+        # live server: a pooled module already carrying this exact policy is
+        # done; a different policy re-applies on top (requantize / dequantize
+        # paths below handle the weights, offloading is detached first)
+        from toolkit.models.v2.pool import ComponentPool, _stable
+
+        policy = _stable(
+            dict(
+                qtype=qtype,
+                offload=offload,
+                dtype=dtype,
+                device=device,
+                quantize_device=quantize_device,
+                exclude=exclude_quant_modules,
+            )
+        )
+        if ComponentPool.current is not None:
+            if getattr(self, "_aitk_policy", None) == policy:
+                ComponentPool.current.touch_module(self)
+                return self
+            if hasattr(self, "_memory_manager"):
+                from toolkit.memory_management import MemoryManager
+
+                MemoryManager.detach(self)
+        self._aitk_policy = policy
+
         ara_path = None
         if qtype is not None and "|" in qtype:
             qtype, ara_path = qtype.split("|", 1)
@@ -428,6 +453,27 @@ class OstrisModelMixin:
         elif subfolder == "":
             subfolder = None
 
+        # live server: an identical source request returns the resident module
+        from toolkit.models.v2.pool import ComponentPool
+
+        pool = ComponentPool.current
+        pool_key = None
+        if pool is not None:
+            pool_key = ComponentPool.make_key(
+                cls,
+                name_or_path,
+                dtype=dtype,
+                config_path=config_path,
+                config=config,
+                subfolder=subfolder,
+                use_comfy_weights=use_comfy_weights,
+                qtype=qtype if quantize_on_load else None,
+                kwargs=kwargs,
+            )
+            resident = pool.get(pool_key)
+            if resident is not None:
+                return resident
+
         if (
             use_comfy_weights
             and not name_or_path.endswith(".safetensors")
@@ -473,6 +519,8 @@ class OstrisModelMixin:
 
         if device is not None:
             model.to(device)
+        if pool is not None and pool_key is not None:
+            pool.put(pool_key, model)
         return model
 
     @staticmethod

toolkit/models/v2/pool.py

@@ -0,0 +1,194 @@
+"""Process-wide component pool for the live inference server.
+
+While a pool is active (`ComponentPool.current`), `OstrisModelMixin.load_model`
+returns an already-resident module for an identical source request instead
+of loading it again, and `aitk_post_load` becomes a no-op when the module
+already carries the requested policy (qtype / offload / device). Holders
+keep calling the normal load API; component reuse across model switches
+(same TE or VAE under two archs) falls out of that.
+
+Lifecycle per request: `begin_request()` -> holder loads (hits mark entries
+as touched) -> `release_unused()` frees every entry the request did not
+touch. `clear()` drops everything.
+"""
+
+import gc
+import json
+import threading
+import time
+from typing import Dict, Optional
+
+import torch
+
+
+def _stable(value):
+    """JSON-safe, deterministic rendering of a key part."""
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    if isinstance(value, (torch.dtype, torch.device)):
+        return str(value)
+    if isinstance(value, dict):
+        return {str(k): _stable(v) for k, v in sorted(value.items(), key=lambda kv: str(kv[0]))}
+    if isinstance(value, (list, tuple)):
+        return [_stable(v) for v in value]
+    to_dict = getattr(value, "to_dict", None)
+    if callable(to_dict):
+        try:
+            return _stable(to_dict())
+        except Exception:
+            pass
+    if hasattr(value, "__dataclass_fields__"):
+        return _stable({k: getattr(value, k) for k in value.__dataclass_fields__})
+    return repr(value)
+
+
+def module_bytes(module: torch.nn.Module) -> int:
+    total = 0
+    seen = set()
+    for t in list(module.parameters()) + list(module.buffers()):
+        try:
+            ptr = t.untyped_storage().data_ptr()
+        except Exception:
+            ptr = id(t)
+        if ptr in seen or t.device.type == "meta":
+            continue
+        seen.add(ptr)
+        total += t.numel() * t.element_size()
+    return total
+
+
+class PoolEntry:
+    def __init__(self, key: str, module: torch.nn.Module, generation: int):
+        self.key = key
+        self.module = module
+        self.cls_name = type(module).__name__
+        self.created_at = time.time()
+        self.last_used = self.created_at
+        self.uses = 0
+        self.touched_gen = generation
+        self._bytes: Optional[int] = None
+
+    @property
+    def bytes(self) -> int:
+        if self._bytes is None:
+            try:
+                self._bytes = module_bytes(self.module)
+            except Exception:
+                self._bytes = 0
+        return self._bytes
+
+    def device(self) -> str:
+        try:
+            p = next(self.module.parameters())
+            return str(p.device)
+        except StopIteration:
+            return "none"
+
+    def info(self) -> dict:
+        try:
+            parsed = json.loads(self.key)
+        except Exception:
+            parsed = {"key": self.key}
+        return {
+            "cls": self.cls_name,
+            "src": parsed.get("src"),
+            "subfolder": parsed.get("subfolder"),
+            "dtype": parsed.get("dtype"),
+            "policy": getattr(self.module, "_aitk_policy", None),
+            "bytes": self.bytes,
+            "device": self.device(),
+            "uses": self.uses,
+            "last_used": self.last_used,
+        }
+
+
+class ComponentPool:
+    current: Optional["ComponentPool"] = None
+
+    def __init__(self):
+        self.entries: Dict[str, PoolEntry] = {}
+        self.generation = 0
+        self.hits = 0
+        self.misses = 0
+        self.evictions = 0
+        self._lock = threading.RLock()
+
+    # ---- keys ----
+    @staticmethod
+    def make_key(cls, name_or_path: str, **parts) -> str:
+        key = {"cls": f"{cls.__module__}.{cls.__qualname__}", "src": name_or_path}
+        key.update({k: _stable(v) for k, v in parts.items()})
+        return json.dumps(key, sort_keys=True, default=repr)
+
+    # ---- request lifecycle ----
+    def begin_request(self):
+        with self._lock:
+            self.generation += 1
+
+    def get(self, key: str) -> Optional[torch.nn.Module]:
+        with self._lock:
+            entry = self.entries.get(key)
+            if entry is None:
+                self.misses += 1
+                return None
+            self.hits += 1
+            entry.uses += 1
+            entry.last_used = time.time()
+            entry.touched_gen = self.generation
+            return entry.module
+
+    def put(self, key: str, module: torch.nn.Module):
+        with self._lock:
+            entry = PoolEntry(key, module, self.generation)
+            entry.uses = 1
+            self.entries[key] = entry
+
+    def touch_module(self, module: torch.nn.Module):
+        with self._lock:
+            for entry in self.entries.values():
+                if entry.module is module:
+                    entry.touched_gen = self.generation
+                    entry.last_used = time.time()
+
+    def release_unused(self) -> int:
+        """Free every entry not touched during the current request."""
+        with self._lock:
+            stale = [e for e in self.entries.values() if e.touched_gen != self.generation]
+            freed = 0
+            for entry in stale:
+                freed += entry.bytes
+                self._free(entry)
+            return freed
+
+    def clear(self):
+        with self._lock:
+            for entry in list(self.entries.values()):
+                self._free(entry)
+
+    def _free(self, entry: PoolEntry):
+        from toolkit.memory_management import MemoryManager
+
+        self.entries.pop(entry.key, None)
+        self.evictions += 1
+        try:
+            MemoryManager.free(entry.module)
+        except Exception:
+            try:
+                torch.nn.Module.to(entry.module, "meta")
+            except Exception:
+                pass
+        entry.module = None
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    def stats(self) -> dict:
+        with self._lock:
+            return {
+                "hits": self.hits,
+                "misses": self.misses,
+                "evictions": self.evictions,
+                "resident": [e.info() for e in self.entries.values()],
+                "resident_bytes": sum(e.bytes for e in self.entries.values()),
+                "untouched": len([e for e in self.entries.values() if e.touched_gen != self.generation]),
+            }

toolkit/sample_step_hook.py

@@ -0,0 +1,134 @@
+"""Per-step latent reporting for sample generation (inference engine).
+
+A holder exposes `sample_step_hook(step_index, num_steps, latents)`; while it
+is set, the sampling scheduler's `step` is wrapped on the instance so every
+step reaches the hook without per-pipeline code. The latent handed to the
+hook is the step's predicted CLEAN sample (x0) — what a preview should show
+— falling back to the scheduler's next sample when x0 cannot be derived. Shared by
+BaseModel and the legacy StableDiffusion monolith.
+"""
+
+from typing import Callable, List
+
+
+def install_sample_step_hooks(holder, pipeline) -> Callable[[], None]:
+    """Wrap scheduler.step on the sampling scheduler(s). Returns an unwrap
+    callable. No-op without a hook. Instance-level wrap: the class method is
+    untouched and a stale wrapper is replaced rather than stacked."""
+    if getattr(holder, "sample_step_hook", None) is None:
+        return lambda: None
+    schedulers: List = []
+    for cand in (
+        getattr(pipeline, "scheduler", None),
+        getattr(getattr(holder, "pipeline", None), "scheduler", None),
+        getattr(holder, "noise_scheduler", None),
+    ):
+        if cand is not None and hasattr(cand, "step") and all(cand is not s for s in schedulers):
+            schedulers.append(cand)
+    unwraps = []
+    for scheduler in schedulers:
+        current = scheduler.step
+        orig = getattr(current, "_aitk_orig_step", None) or current
+
+        def make(orig, scheduler):
+            def step(*args, **kwargs):
+                # capture the pre-step state: x0 is derived from (model_output,
+                # timestep, sample) with the sigma/alpha of THIS step, which
+                # the scheduler advances past inside orig()
+                model_output = args[0] if len(args) > 0 else kwargs.get("model_output")
+                timestep = args[1] if len(args) > 1 else kwargs.get("timestep")
+                sample_in = args[2] if len(args) > 2 else kwargs.get("sample")
+                coeffs = _x0_coefficients(scheduler, timestep)
+                out = orig(*args, **kwargs)
+                if isinstance(out, tuple):
+                    prev_sample = out[0]
+                    pred_x0 = None
+                else:
+                    prev_sample = getattr(out, "prev_sample", None)
+                    pred_x0 = getattr(out, "pred_original_sample", None)
+                hook = getattr(holder, "sample_step_hook", None)
+                if hook is not None and prev_sample is not None:
+                    if pred_x0 is None:
+                        pred_x0 = _predict_x0(coeffs, model_output, sample_in)
+                    timesteps = getattr(scheduler, "timesteps", None)
+                    num_steps = len(timesteps) if timesteps is not None else None
+                    idx = getattr(holder, "_sample_step_index", 0)
+                    holder._sample_step_index = idx + 1
+                    hook(idx, num_steps, pred_x0 if pred_x0 is not None else prev_sample)
+                return out
+
+            step._aitk_orig_step = orig
+            return step
+
+        scheduler.step = make(orig, scheduler)
+
+        def unwrap(scheduler=scheduler):
+            if "step" in scheduler.__dict__:
+                del scheduler.__dict__["step"]
+
+        unwraps.append(unwrap)
+    return lambda: [u() for u in unwraps]
+
+
+def _x0_coefficients(scheduler, timestep):
+    """Per-step numbers needed to turn the model output into a clean-sample
+    estimate, read BEFORE scheduler.step advances its index.
+    Returns ("flow", sigma) | ("ddpm", prediction_type, alpha_bar) | None."""
+    try:
+        import torch
+
+        sigmas = getattr(scheduler, "sigmas", None)
+        if sigmas is not None and hasattr(scheduler, "index_for_timestep"):
+            idx = getattr(scheduler, "step_index", None)
+            if idx is None:
+                idx = scheduler.index_for_timestep(timestep)
+            return ("flow", float(sigmas[idx]))
+        alphas_cumprod = getattr(scheduler, "alphas_cumprod", None)
+        if alphas_cumprod is not None:
+            t = timestep
+            if isinstance(t, torch.Tensor):
+                t = int(t.flatten()[0].item())
+            prediction_type = getattr(getattr(scheduler, "config", None), "prediction_type", "epsilon")
+            return ("ddpm", prediction_type, float(alphas_cumprod[int(t)]))
+    except Exception:
+        return None
+    return None
+
+
+def _predict_x0(coeffs, model_output, sample):
+    """Clean-sample estimate for the preview: flow matching x0 = x_t - sigma*v;
+    DDPM-style eps / v / sample predictions via alpha_bar. None when unknown."""
+    if coeffs is None or model_output is None or sample is None:
+        return None
+    try:
+        x = sample.tensor if hasattr(sample, "tensor") else sample
+        v = model_output.tensor if hasattr(model_output, "tensor") else model_output
+        if v.shape != x.shape:
+            return None
+        v = v.to(x.dtype)
+        if coeffs[0] == "flow":
+            sigma = coeffs[1]
+            return x - sigma * v
+        _, prediction_type, alpha_bar = coeffs
+        a = alpha_bar**0.5
+        b = (1.0 - alpha_bar) ** 0.5
+        if prediction_type == "epsilon":
+            return (x - b * v) / a
+        if prediction_type == "v_prediction":
+            return a * x - b * v
+        if prediction_type == "sample":
+            return v
+    except Exception:
+        return None
+    return None
+
+
+def emit_sample_step(holder, latents, step_index=None, num_steps=None):
+    """For sampling loops that bypass scheduler.step: report one denoised
+    latent to the holder's hook (no-op when unset)."""
+    hook = getattr(holder, "sample_step_hook", None)
+    if hook is None:
+        return
+    idx = getattr(holder, "_sample_step_index", 0) if step_index is None else step_index
+    holder._sample_step_index = idx + 1
+    hook(idx, num_steps, latents)

toolkit/stable_diffusion_model.py

@@ -191,6 +191,8 @@ class StableDiffusion:
         self.te_torch_dtype = get_torch_dtype(model_config.te_dtype)
 
         self.model_config = model_config
+        # inference engine: hook(step_index, num_steps, latents) per scheduler step
+        self.sample_step_hook = None
         self.prediction_type = "v_prediction" if self.model_config.is_v_pred else "epsilon"
         self.arch = model_config.arch
 
@@ -1417,6 +1419,10 @@ class StableDiffusion:
             # disable progress bar
             pipeline.set_progress_bar_config(disable=True)
 
+        from toolkit.sample_step_hook import install_sample_step_hooks
+
+        unwrap_step_hooks = install_sample_step_hooks(self, pipeline)
+
         refiner_pipeline = None
         if self.refiner_unet:
             # build refiner pipeline
@@ -1495,6 +1501,7 @@ class StableDiffusion:
 
                     if network is not None:
                         network.multiplier = gen_config.network_multiplier
+                    self._sample_step_index = 0
                     torch.manual_seed(gen_config.seed)
                     torch.cuda.manual_seed(gen_config.seed)
                     
@@ -1792,6 +1799,7 @@ class StableDiffusion:
                 if self.adapter is not None and isinstance(self.adapter, ReferenceAdapter):
                     self.adapter.clear_memory()
 
+        unwrap_step_hooks()
         # clear pipeline and cache to reduce vram usage
         del pipeline
         if refiner_pipeline is not None:

toolkit/ui_job_status.py

@@ -0,0 +1,130 @@
+"""UI job bookkeeping for long-running extension processes.
+
+When a process is launched by the UI it gets AITK_JOB_ID and a sqlite db
+path; status / step / info are written to the Job row and the UI's stop
+button sets Job.stop. This mirrors what BaseCaptioner does inline, as a
+standalone helper so other UI-driven processes (the inference engine) can
+share it. Outside the UI (no job id or no db file) every call is a no-op.
+"""
+
+import os
+import sqlite3
+import threading
+import time
+from typing import Callable, Literal, Optional
+
+AITK_Status = Literal["running", "stopped", "error", "completed", "queued"]
+
+
+class UIJobStatus:
+    def __init__(self, sqlite_db_path: str = "./aitk_db.db"):
+        self.sqlite_db_path = sqlite_db_path
+        job_id = os.environ.get("AITK_JOB_ID", None)
+        self.job_id = job_id.strip() if job_id else None
+        self.is_ui_job = self.job_id is not None and os.path.exists(sqlite_db_path)
+        if self.is_ui_job:
+            print(f"Using SQLite database at {sqlite_db_path}")
+            print(f'Job ID: "{self.job_id}"')
+        self._lock = threading.Lock()
+        self._watcher = None
+        # collapse bursts of identical writes (progress ticks)
+        self._last_info = None
+
+    # ---- db plumbing ----
+    def _connect(self):
+        conn = sqlite3.connect(self.sqlite_db_path, timeout=30.0)
+        conn.isolation_level = None
+        return conn
+
+    def _retry(self, fn, max_retries=3, base_delay=2.0):
+        last = None
+        for attempt in range(max_retries + 1):
+            try:
+                return fn()
+            except sqlite3.OperationalError as e:
+                if "database is locked" not in str(e):
+                    raise
+                last = e
+                if attempt < max_retries:
+                    time.sleep(base_delay * (2**attempt))
+        raise last
+
+    def _write(self, query: str, params: tuple):
+        if not self.is_ui_job:
+            return
+
+        def _do():
+            with self._lock, self._connect() as conn:
+                cur = conn.cursor()
+                cur.execute("BEGIN IMMEDIATE")
+                try:
+                    cur.execute(query, params)
+                finally:
+                    cur.execute("COMMIT")
+
+        try:
+            self._retry(_do)
+        except Exception as e:
+            print(f"[AITK] Warning: db write failed: {e}")
+
+    def _read_flag(self, column: str) -> bool:
+        if not self.is_ui_job:
+            return False
+
+        def _do():
+            with self._lock, self._connect() as conn:
+                cur = conn.cursor()
+                cur.execute(f"SELECT {column} FROM Job WHERE id = ?", (self.job_id,))
+                row = cur.fetchone()
+                return False if row is None else row[0] == 1
+
+        return self._retry(_do)
+
+    # ---- public api ----
+    def update_status(self, status: AITK_Status, info: Optional[str] = None):
+        if info is None:
+            self._write("UPDATE Job SET status = ? WHERE id = ?", (status, self.job_id))
+        else:
+            self._last_info = info
+            self._write(
+                "UPDATE Job SET status = ?, info = ? WHERE id = ?",
+                (status, info, self.job_id),
+            )
+
+    def update_info(self, info: str):
+        if info == self._last_info:
+            return
+        self._last_info = info
+        self._write("UPDATE Job SET info = ? WHERE id = ?", (info, self.job_id))
+
+    def update_key(self, key: str, value):
+        # key is a column name from our own code, never user input
+        self._write(f"UPDATE Job SET {key} = ? WHERE id = ?", (str(value), self.job_id))
+
+    def should_stop(self) -> bool:
+        return self._read_flag("stop")
+
+    def should_return_to_queue(self) -> bool:
+        return self._read_flag("return_to_queue")
+
+    def start_stop_watcher(self, on_stop: Callable[[str], None], interval_sec: float = 2.0):
+        """Poll the stop / return_to_queue flags on a daemon thread and call
+        on_stop("stopped" | "queued") once when one is set."""
+        if not self.is_ui_job or self._watcher is not None:
+            return
+
+        def _loop():
+            while True:
+                try:
+                    if self.should_stop():
+                        on_stop("stopped")
+                        return
+                    if self.should_return_to_queue():
+                        on_stop("queued")
+                        return
+                except Exception:
+                    pass
+                time.sleep(interval_sec)
+
+        self._watcher = threading.Thread(target=_loop, daemon=True)
+        self._watcher.start()