753 lines
31 KiB
Python
753 lines
31 KiB
Python
"""Engine: request queue + holder lifecycle + generation with streamed frames.
|
|
|
|
Runs on the process main thread (`run_forever`); the HTTP server hands it
|
|
GenerationJobs from its own thread and reads the frames each job emits.
|
|
"""
|
|
|
|
import gc
|
|
import glob
|
|
import json
|
|
import os
|
|
import queue
|
|
import shutil
|
|
import threading
|
|
import time
|
|
import traceback
|
|
import uuid
|
|
from typing import Callable, Dict, List, Optional
|
|
|
|
import torch
|
|
|
|
from toolkit.basic import flush
|
|
from toolkit.config_modules import GenerateImageConfig, ModelConfig
|
|
from toolkit.models.registry import OUTPUT_EXT, get_arch_entry
|
|
from toolkit.models.v2.pool import ComponentPool
|
|
from toolkit.inference_lora import InferenceLoRA, LoRAStack
|
|
from toolkit.util.get_model import LEGACY_ARCHS, get_model_class
|
|
|
|
from .latent_preview import preview_info
|
|
from .protocol import encode_frame, tensor_payload
|
|
|
|
END = None # frames-queue terminator (after the "end" frame)
|
|
|
|
|
|
class GenerationCancelled(Exception):
|
|
pass
|
|
|
|
|
|
class GenerationJob:
|
|
def __init__(self, model: dict, sample: dict, stream: Optional[dict] = None):
|
|
self.request_id = uuid.uuid4().hex[:12]
|
|
self.model = model
|
|
self.sample = sample
|
|
self.stream = {"latents": "raw", "every_n_steps": 1, "max_frames": 4}
|
|
self.stream.update(stream or {})
|
|
# 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()
|
|
self.started_at: Optional[float] = None
|
|
self.finished_at: Optional[float] = None
|
|
self.results: List[dict] = []
|
|
self.error: Optional[str] = None
|
|
self.step = 0
|
|
self.total_steps: Optional[int] = None
|
|
|
|
def emit(self, type_: str, payload: bytes = b"", **fields):
|
|
header = {"type": type_, "request_id": self.request_id, **fields}
|
|
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)
|
|
with self._sub_lock:
|
|
self.finished = True
|
|
for q in self._subscribers:
|
|
q.put(END)
|
|
self._subscribers = []
|
|
|
|
def info(self) -> dict:
|
|
return {
|
|
"request_id": self.request_id,
|
|
"status": self.status,
|
|
"arch": self.model.get("arch"),
|
|
"prompt": self.sample.get("prompt", ""),
|
|
"step": self.step,
|
|
"total_steps": self.total_steps,
|
|
"created_at": self.created_at,
|
|
"started_at": self.started_at,
|
|
"finished_at": self.finished_at,
|
|
"results": self.results,
|
|
"error": self.error,
|
|
}
|
|
|
|
|
|
def _model_key(model: dict) -> str:
|
|
return json.dumps({k: v for k, v in sorted(model.items()) if v is not None}, sort_keys=True)
|
|
|
|
|
|
# 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")
|
|
# LoRAs ride the request's model block but are applied on top of the loaded
|
|
# holder (hook mode: dynamic; merge mode: into the weights, which reloads
|
|
# the touched components when the merged set changes)
|
|
LORA_KEYS = ("loras", "lora_mode")
|
|
|
|
|
|
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,
|
|
device: str,
|
|
output_folder: str,
|
|
assets_folder: Optional[str] = None,
|
|
status_cb: Optional[Callable[[str], None]] = None,
|
|
dtype: str = "bf16",
|
|
):
|
|
self.device = device
|
|
self.dtype = dtype
|
|
self.output_folder = output_folder
|
|
self.assets_folder = assets_folder or os.path.join(os.path.dirname(output_folder), "assets")
|
|
os.makedirs(self.output_folder, exist_ok=True)
|
|
os.makedirs(self.assets_folder, exist_ok=True)
|
|
self.status_cb = status_cb or (lambda s: None)
|
|
self.queue: "queue.Queue[GenerationJob]" = queue.Queue()
|
|
self.jobs: Dict[str, GenerationJob] = {}
|
|
self.current: Optional[GenerationJob] = None
|
|
self.stop_event = threading.Event()
|
|
self.stop_reason = "stopped"
|
|
self.loading = False
|
|
self.holder = None
|
|
self.holder_key: Optional[str] = None
|
|
self.holder_model: Optional[dict] = None
|
|
self.holder_is_legacy = False
|
|
self.lora_stack = None
|
|
self.stats = {"holder_loads": 0, "generations": 0, "errors": 0, "cancelled": 0, "oom_retries": 0}
|
|
self._lock = threading.Lock()
|
|
# resident components shared across holders (same TE/VAE under two archs)
|
|
self.pool = ComponentPool()
|
|
|
|
# ------------------------------------------------------------------ api
|
|
def submit(self, model: dict, sample: dict, stream: Optional[dict] = None) -> GenerationJob:
|
|
if "arch" not in model:
|
|
raise ValueError("model.arch is required")
|
|
job = GenerationJob(model, sample, stream)
|
|
with self._lock:
|
|
self.jobs[job.request_id] = job
|
|
# 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:
|
|
return False
|
|
job.cancel.set()
|
|
if job.status == "queued":
|
|
job.status = "cancelled"
|
|
job.emit("error", message="cancelled", cancelled=True)
|
|
job.finish()
|
|
return True
|
|
|
|
def stop(self, reason: str = "stopped"):
|
|
self.stop_reason = reason
|
|
self.stop_event.set()
|
|
if self.current is not None:
|
|
self.current.cancel.set()
|
|
|
|
def queued(self) -> List[dict]:
|
|
return [j.info() for j in self.jobs.values() if j.status in ("queued", "running")]
|
|
|
|
def recent(self, limit: int = 50) -> List[dict]:
|
|
jobs = sorted(self.jobs.values(), key=lambda j: j.created_at, reverse=True)
|
|
return [j.info() for j in jobs[:limit]]
|
|
|
|
def health(self) -> dict:
|
|
vram = None
|
|
if torch.cuda.is_available():
|
|
try:
|
|
free, total = torch.cuda.mem_get_info(torch.device(self.device))
|
|
vram = {"used": total - free, "total": total}
|
|
except Exception:
|
|
vram = None
|
|
return {
|
|
"ok": True,
|
|
"busy": self.current is not None,
|
|
"queue": len([j for j in self.jobs.values() if j.status == "queued"]),
|
|
"current": self.current.info() if self.current else None,
|
|
"active": {
|
|
"model": self.holder_model,
|
|
"arch": (self.holder_model or {}).get("arch"),
|
|
"loras": [l.summary() for l in self.lora_stack.loras] if self.lora_stack else [],
|
|
"lora_mode": self.lora_stack.mode if self.lora_stack else None,
|
|
}
|
|
if self.holder is not None
|
|
else None,
|
|
"device": self.device,
|
|
"vram": vram,
|
|
"stats": dict(self.stats),
|
|
"pool": self.pool.stats(),
|
|
}
|
|
|
|
def _drop_holder(self):
|
|
"""Drop the holder object; its components stay resident in the pool
|
|
until release_unused()/clear() decides."""
|
|
if self.lora_stack is not None and self.lora_stack.mode == "hook":
|
|
self.lora_stack.remove()
|
|
self.lora_stack = None
|
|
holder = self.holder
|
|
self.holder = None
|
|
self.holder_key = None
|
|
self.holder_model = None
|
|
if holder is not None:
|
|
for attr in ("pipeline", "model", "vae", "text_encoder", "image_encoder", "adapter", "network"):
|
|
try:
|
|
setattr(holder, attr, None)
|
|
except Exception:
|
|
pass
|
|
del holder
|
|
gc.collect()
|
|
|
|
def unload(self):
|
|
if self.holder is None and not self.pool.entries:
|
|
return
|
|
self.status_cb("Unloading model")
|
|
self._drop_holder()
|
|
self.pool.clear()
|
|
flush()
|
|
|
|
def preload(self, model: dict):
|
|
"""Load a model's holder ahead of any request (warm start). Runs on
|
|
the caller's thread; meant for the engine main thread before
|
|
run_forever."""
|
|
job = GenerationJob(dict(model), {}, {"latents": "none"})
|
|
job.status = "running"
|
|
self.current = job
|
|
try:
|
|
self._ensure_holder(job)
|
|
job.status = "done"
|
|
except Exception as e:
|
|
job.status = "error"
|
|
job.error = f"{type(e).__name__}: {e}"
|
|
traceback.print_exc()
|
|
finally:
|
|
self.current = None
|
|
job.finish()
|
|
return job
|
|
|
|
# ------------------------------------------------------------------ loop
|
|
def run_forever(self):
|
|
while not self.stop_event.is_set():
|
|
try:
|
|
job = self.queue.get(timeout=0.5)
|
|
except queue.Empty:
|
|
continue
|
|
except KeyboardInterrupt:
|
|
self.stop("stopped")
|
|
break
|
|
if job.status != "queued":
|
|
continue
|
|
self.current = job
|
|
try:
|
|
self._run(job)
|
|
except KeyboardInterrupt:
|
|
job.status = "cancelled"
|
|
job.emit("error", message="engine stopping", cancelled=True)
|
|
job.finish()
|
|
self.stop("stopped")
|
|
break
|
|
finally:
|
|
self.current = None
|
|
self._idle_status()
|
|
# drain anything left queued
|
|
while True:
|
|
try:
|
|
job = self.queue.get_nowait()
|
|
except queue.Empty:
|
|
break
|
|
self.cancel(job.request_id)
|
|
|
|
def _idle_status(self):
|
|
if self.holder_model is not None:
|
|
self.status_cb(f"Idle - {self.holder_model.get('arch')} ({self.holder_model.get('name_or_path')})")
|
|
else:
|
|
self.status_cb("Idle - no model loaded")
|
|
|
|
# ------------------------------------------------------------------ holder
|
|
def _build_model_config(self, model: dict) -> ModelConfig:
|
|
kwargs = dict(model)
|
|
arch = kwargs.pop("arch")
|
|
entry = get_arch_entry(arch)
|
|
for k, v in entry["model"].items():
|
|
kwargs.setdefault(k, v)
|
|
kwargs.setdefault("dtype", self.dtype)
|
|
if kwargs.get("quantize"):
|
|
kwargs.setdefault("qtype", "convrot8")
|
|
if kwargs.get("quantize_te"):
|
|
kwargs.setdefault("qtype_te", "convrot8")
|
|
return ModelConfig(arch=arch, **kwargs)
|
|
|
|
def _ensure_holder(self, job: GenerationJob):
|
|
model = dict(job.model)
|
|
entry = get_arch_entry(model["arch"])
|
|
for k, v in entry["model"].items():
|
|
model.setdefault(k, v)
|
|
lora_specs = [l for l in (model.get("loras") or []) if l.get("path")]
|
|
lora_mode = model.get("lora_mode") or "hook"
|
|
model = _without(model, LORA_KEYS)
|
|
key = _model_key(model)
|
|
if self.holder is not None and key == self.holder_key:
|
|
self._apply_loras(job, lora_specs, lora_mode)
|
|
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
|
|
self._apply_loras(job, lora_specs, lora_mode)
|
|
return self.holder
|
|
if self.holder is not None:
|
|
self._drop_holder()
|
|
|
|
# merged LoRAs live inside pooled weights: a component merged with a
|
|
# different set than this request wants must reload from disk
|
|
wanted = LoRAStack(None, lora_mode)
|
|
wanted.loras = [InferenceLoRA(l["path"], l.get("strength", 1.0)) for l in lora_specs] if lora_mode == "merge" else []
|
|
wanted_key = wanted.key() if lora_mode == "merge" else None
|
|
freed = self.pool.evict_where(lambda m: getattr(m, "_aitk_merged_loras", None) not in (None, wanted_key))
|
|
if freed:
|
|
job.emit("status", message=f"Reloading {freed / 1e9:.1f} GB of components that carried other merged LoRAs")
|
|
|
|
self.pool.begin_request()
|
|
ComponentPool.current = self.pool
|
|
self.loading = True
|
|
try:
|
|
try:
|
|
holder, is_legacy = self._build_holder(job, model)
|
|
except torch.OutOfMemoryError:
|
|
# components the previous model needed but this one does not
|
|
# were still resident: free them and try once more
|
|
self.stats["oom_retries"] += 1
|
|
job.emit("status", message="Out of memory: freeing unused components and retrying")
|
|
freed = self.pool.release_unused()
|
|
flush()
|
|
job.emit("status", message=f"Freed {freed / 1e9:.1f} GB, retrying load")
|
|
holder, is_legacy = self._build_holder(job, model)
|
|
finally:
|
|
ComponentPool.current = None
|
|
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']}")
|
|
flush()
|
|
self.holder = holder
|
|
self.holder_key = key
|
|
self.holder_model = model
|
|
self.holder_is_legacy = is_legacy
|
|
self.lora_stack = None
|
|
self._apply_loras(job, lora_specs, lora_mode)
|
|
return holder
|
|
|
|
def _apply_loras(self, job: GenerationJob, specs: List[dict], mode: str):
|
|
"""Bring the holder's LoRA state to `specs` in `mode`. Hook mode:
|
|
strengths update in place, a different set re-attaches. Merge mode:
|
|
a different set than what the weights carry evicts and reloads them."""
|
|
wanted = LoRAStack(self.holder, mode)
|
|
wanted.loras = [InferenceLoRA(l["path"], l.get("strength", 1.0)) for l in specs]
|
|
wanted_key = wanted.key()
|
|
current = self.lora_stack
|
|
if current is not None and current.key() == wanted_key:
|
|
return
|
|
if current is not None and current.mode == "hook" and mode == "hook" and current.set_strengths(specs):
|
|
self.lora_stack = current
|
|
return
|
|
|
|
def status(msg):
|
|
job.emit("status", message=msg)
|
|
self.status_cb(msg)
|
|
|
|
if current is not None:
|
|
if current.mode == "merge":
|
|
# weights carry the old set: reload the merged components
|
|
status("Reloading weights to drop merged LoRAs")
|
|
merged_key = current.key()
|
|
self._drop_holder()
|
|
self.pool.evict_where(lambda m: getattr(m, "_aitk_merged_loras", None) == merged_key)
|
|
flush()
|
|
self.lora_stack = None
|
|
self._ensure_holder_for_loras(job, specs, mode)
|
|
return
|
|
current.remove()
|
|
self.lora_stack = None
|
|
if not specs:
|
|
return
|
|
stack = LoRAStack(self.holder, mode).load(specs, status_fn=status)
|
|
touched = stack.apply(status_fn=status)
|
|
if mode == "merge":
|
|
for root in [getattr(self.holder, "model", None)] + (
|
|
self.holder.text_encoder if isinstance(getattr(self.holder, "text_encoder", None), list) else [getattr(self.holder, "text_encoder", None)]
|
|
):
|
|
if isinstance(root, torch.nn.Module) and any(id(m) in touched for m in root.modules()):
|
|
root._aitk_merged_loras = stack.key()
|
|
self.lora_stack = stack
|
|
|
|
def _ensure_holder_for_loras(self, job: GenerationJob, specs: List[dict], mode: str):
|
|
model = dict(self.holder_model or job.model)
|
|
model["loras"] = specs
|
|
model["lora_mode"] = mode
|
|
job_model = job.model
|
|
job.model = model
|
|
try:
|
|
self._ensure_holder(job)
|
|
finally:
|
|
job.model = job_model
|
|
|
|
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)
|
|
is_legacy = not hasattr(ModelClass, "get_train_scheduler")
|
|
if is_legacy:
|
|
from toolkit.sampler import get_sampler
|
|
|
|
legacy_arch = "sd"
|
|
if model_config.is_pixart:
|
|
legacy_arch = "pixart"
|
|
elif model_config.is_flux:
|
|
legacy_arch = "flux"
|
|
elif model_config.is_lumina2:
|
|
legacy_arch = "lumina2"
|
|
sampler = get_sampler(
|
|
"ddpm",
|
|
{"prediction_type": "v_prediction" if model_config.is_v_pred else "epsilon"},
|
|
arch=legacy_arch,
|
|
)
|
|
else:
|
|
sampler = ModelClass.get_train_scheduler()
|
|
|
|
job.emit("status", message=f"Loading {model['arch']}")
|
|
self.status_cb(f"Loading {model['arch']}: {model.get('name_or_path')}")
|
|
holder = ModelClass(
|
|
device=self.device,
|
|
model_config=model_config,
|
|
dtype=self.dtype,
|
|
noise_scheduler=sampler,
|
|
)
|
|
|
|
def _status_hook(msg):
|
|
job.emit("status", message=str(msg))
|
|
self.status_cb(f"Loading {model['arch']}: {msg}")
|
|
|
|
if hasattr(holder, "add_status_update_hook"):
|
|
holder.add_status_update_hook(_status_hook)
|
|
t0 = time.perf_counter()
|
|
holder.load_model()
|
|
self.stats["holder_loads"] += 1
|
|
p = self.pool.stats()
|
|
job.emit(
|
|
"status",
|
|
message=f"Loaded {model['arch']} in {time.perf_counter() - t0:.1f}s (pool hits {p['hits']}, misses {p['misses']})",
|
|
)
|
|
return holder, is_legacy
|
|
|
|
# ------------------------------------------------------------------ run
|
|
def _resolve_asset(self, value):
|
|
if value is None:
|
|
return None
|
|
if os.path.isabs(value) and os.path.exists(value):
|
|
return value
|
|
cand = os.path.join(self.assets_folder, os.path.basename(value))
|
|
if os.path.exists(cand):
|
|
return cand
|
|
if os.path.exists(value):
|
|
return os.path.abspath(value)
|
|
raise FileNotFoundError(f"control image not found: {value}")
|
|
|
|
def _run(self, job: GenerationJob):
|
|
job.status = "running"
|
|
job.started_at = time.time()
|
|
arch = job.model.get("arch")
|
|
entry = get_arch_entry(arch)
|
|
out_dir = os.path.join(self.output_folder, job.request_id)
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
try:
|
|
holder = self._ensure_holder(job)
|
|
if job.cancel.is_set():
|
|
raise GenerationCancelled()
|
|
|
|
sample = dict(entry["sample"])
|
|
sample.update({k: v for k, v in job.sample.items() if v is not None})
|
|
for key in ("ctrl_img", "ctrl_img_1", "ctrl_img_2", "ctrl_img_3"):
|
|
if sample.get(key):
|
|
sample[key] = self._resolve_asset(sample[key])
|
|
if entry["needs_control_image"] and not sample.get("ctrl_img"):
|
|
raise ValueError(f"{arch} requires a control image (ctrl_img)")
|
|
output_ext = sample.pop("output_ext", None) or OUTPUT_EXT[entry["modality"]]
|
|
if entry["modality"] == "video" and sample.get("num_frames", 1) > 1 and output_ext == "mp4":
|
|
# BaseModel's video saver writes animated webp; archs with their
|
|
# own mp4 writer (ltx2) override output_ext themselves
|
|
output_ext = "webp"
|
|
prompt = sample.pop("prompt", "")
|
|
gen = GenerateImageConfig(
|
|
prompt=prompt,
|
|
output_folder=out_dir,
|
|
output_ext=output_ext,
|
|
**sample,
|
|
)
|
|
job.total_steps = gen.num_inference_steps
|
|
job.emit(
|
|
"start",
|
|
model={**(self.holder_model or {}), "loras": job.model.get("loras") or [], "lora_mode": job.model.get("lora_mode") or "hook"},
|
|
sample={
|
|
"prompt": gen.prompt,
|
|
"negative_prompt": gen.negative_prompt,
|
|
"width": gen.width,
|
|
"height": gen.height,
|
|
"num_inference_steps": gen.num_inference_steps,
|
|
"guidance_scale": gen.guidance_scale,
|
|
"seed": gen.seed,
|
|
"num_frames": gen.num_frames,
|
|
"fps": gen.fps,
|
|
},
|
|
modality=entry["modality"],
|
|
preview=preview_info(arch),
|
|
)
|
|
|
|
t_start = time.perf_counter()
|
|
every_n = max(1, int(job.stream.get("every_n_steps", 1)))
|
|
want_latents = job.stream.get("latents", "raw") == "raw"
|
|
max_frames = int(job.stream.get("max_frames", 4))
|
|
|
|
def step_hook(idx, num_steps, latents):
|
|
if job.cancel.is_set():
|
|
raise GenerationCancelled()
|
|
job.step = idx + 1
|
|
if num_steps:
|
|
job.total_steps = num_steps
|
|
elapsed = time.perf_counter() - t_start
|
|
job.emit("progress", step=job.step, total=job.total_steps, elapsed=round(elapsed, 3))
|
|
self.status_cb(f"Generating {job.step}/{job.total_steps or '?'}")
|
|
if want_latents and (idx % every_n == 0 or (job.total_steps and job.step == job.total_steps)):
|
|
self._emit_latent(job, latents, gen, max_frames)
|
|
|
|
def status_hook(msg):
|
|
job.emit("status", message=str(msg))
|
|
|
|
holder.sample_step_hook = step_hook
|
|
if hasattr(holder, "add_status_update_hook"):
|
|
holder._status_update_hooks.append(status_hook)
|
|
gen_kwargs = {"sampler": "ddpm"} if self.holder_is_legacy else {}
|
|
self.status_cb(f"Generating with {arch}")
|
|
before = set(_list_outputs(out_dir))
|
|
t_gen = time.perf_counter()
|
|
try:
|
|
holder.generate_images([gen], **gen_kwargs)
|
|
finally:
|
|
gen_seconds = time.perf_counter() - t_gen
|
|
holder.sample_step_hook = None
|
|
if hasattr(holder, "_status_update_hooks"):
|
|
try:
|
|
holder._status_update_hooks.remove(status_hook)
|
|
except ValueError:
|
|
pass
|
|
produced = [p for p in _list_outputs(out_dir) if p not in before]
|
|
if not produced:
|
|
raise RuntimeError("generation produced no output file")
|
|
steps = job.total_steps or gen.num_inference_steps
|
|
per_step = gen_seconds / steps if steps else 0
|
|
size = f"{gen.width}x{gen.height}" + (f"x{gen.num_frames}f" if gen.num_frames > 1 else "")
|
|
timing = f"Generated {arch} {size} in {gen_seconds:.2f}s ({steps} steps, {per_step:.2f}s/step, seed {gen.seed})"
|
|
print(f"[AITK] {timing}", flush=True)
|
|
job.emit("status", message=timing)
|
|
for path in sorted(produced):
|
|
ext = os.path.splitext(path)[1].lower().lstrip(".")
|
|
result = {
|
|
"path": path,
|
|
"relpath": os.path.relpath(path, self.output_folder),
|
|
"kind": _kind_for_ext(ext),
|
|
"ext": ext,
|
|
"bytes": os.path.getsize(path),
|
|
"seed": gen.seed,
|
|
"width": gen.width,
|
|
"height": gen.height,
|
|
"num_frames": gen.num_frames,
|
|
"fps": gen.fps,
|
|
"seconds": round(gen_seconds, 3),
|
|
"steps": steps,
|
|
}
|
|
job.results.append(result)
|
|
job.emit("result", **result)
|
|
job.status = "done"
|
|
self.stats["generations"] += 1
|
|
except (GenerationCancelled, KeyboardInterrupt) as e:
|
|
job.status = "cancelled"
|
|
job.error = "cancelled"
|
|
self.stats["cancelled"] += 1
|
|
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}"
|
|
self.stats["errors"] += 1
|
|
traceback.print_exc()
|
|
job.emit("error", message=job.error, cancelled=False, traceback=traceback.format_exc())
|
|
flush()
|
|
finally:
|
|
if job.status != "done":
|
|
shutil.rmtree(out_dir, ignore_errors=True)
|
|
job.finish()
|
|
|
|
def _emit_latent(self, job: GenerationJob, latents, gen: GenerateImageConfig, max_frames: int):
|
|
try:
|
|
t = latents
|
|
if hasattr(t, "tensor"):
|
|
t = t.tensor
|
|
if not isinstance(t, torch.Tensor):
|
|
return
|
|
t = t.detach()
|
|
layout = None
|
|
if t.dim() == 3 and t.shape[-1] % 4 == 0:
|
|
# packed 2x2 patch tokens (flux family): (B, h/2*w/2, C*4)
|
|
h = gen.height // 16
|
|
w = gen.width // 16
|
|
if t.shape[1] == h * w:
|
|
c = t.shape[-1] // 4
|
|
t = t.view(t.shape[0], h, w, c, 2, 2).permute(0, 3, 1, 4, 2, 5).reshape(t.shape[0], c, h * 2, w * 2)
|
|
if t.dim() == 4:
|
|
layout = "BCHW"
|
|
elif t.dim() == 5:
|
|
layout = "BCFHW"
|
|
if max_frames and t.shape[2] > max_frames:
|
|
idx = torch.linspace(0, t.shape[2] - 1, max_frames).round().long().to(t.device)
|
|
t = t.index_select(2, idx)
|
|
elif t.dim() == 3:
|
|
layout = "BLC" # sequence latents (audio: batch, length, channels)
|
|
else:
|
|
layout = "raw"
|
|
fields, payload = tensor_payload(t[:1] if t.dim() >= 3 else t)
|
|
extra = {}
|
|
if not getattr(job, "_preview_sent", False) and layout in ("BCHW", "BCFHW"):
|
|
# now that the real latent shape is known, resolve the preview
|
|
# table (arch table, else a channel-count guess) once per job
|
|
job._preview_sent = True
|
|
extra["preview"] = preview_info(job.model.get("arch"), int(t.shape[1]), 2 if layout == "BCHW" else 3)
|
|
job.emit("latent", payload, step=job.step, total=job.total_steps, layout=layout, **fields, **extra)
|
|
except Exception as e:
|
|
job.emit("status", message=f"latent preview skipped: {e}")
|
|
|
|
|
|
def _list_outputs(folder: str) -> List[str]:
|
|
return [
|
|
p
|
|
for p in glob.glob(os.path.join(folder, "*"))
|
|
if os.path.isfile(p) and not os.path.basename(p).startswith(".") and not p.endswith(".txt")
|
|
]
|
|
|
|
|
|
def _kind_for_ext(ext: str) -> str:
|
|
if ext in ("mp4", "webm", "mov", "webp"):
|
|
return "video" if ext != "webp" else "image"
|
|
if ext in ("wav", "mp3", "flac", "ogg"):
|
|
return "audio"
|
|
return "image"
|