"""Model loading, Settings and the custom denoising loop.

Every rule the workshop talks about is a line in this file. The pipeline object
is only a container for the components and a few static helpers. We never call
pipe(...) during the workshop.
"""

from __future__ import annotations

import math
import os
import threading
import time
from collections import OrderedDict
from contextlib import contextmanager
from dataclasses import dataclass, field

import numpy as np
import torch
from PIL import Image

from family import ACTIVE as FAMILY

# Defaults come from the model family (family.py), environment variables override them.
MODEL_ID = os.environ.get("MODEL_ID", FAMILY.model_id)
DEFAULT_SIZE = int(os.environ.get("DEFAULT_SIZE", FAMILY.default_size))
DEFAULT_STEPS = int(os.environ.get("DEFAULT_STEPS", FAMILY.default_steps))
GRID_N = int(os.environ.get("GRID_N", FAMILY.grid_n))     # FLUX: 9 images take 30 s on the L40S, 4 take 13 s
MAX_BATCH = int(os.environ.get("MAX_BATCH", FAMILY.max_batch))
MAX_SEQ = int(os.environ.get("MAX_SEQ", FAMILY.max_seq))
GRID_SIZES = tuple(k * k for k in range(2, 9))             # 2 x 2 to 8 x 8, one choice for the whole app

STANDARD_LAYERS = FAMILY.standard_layers
EXPERIMENTS = ("seed_grid", "guidance_sweep", "interpolation")
# Pixels per VAE decode call. Decoding in float32 is memory hungry at 1024 px.
DECODE_PIXELS = int(os.environ.get("DECODE_PIXELS", str(4 * 512 * 512)))
EMB_CACHE_SIZE = 32

BEND_OPS = ("none", "amplify", "invert", "ablate", "rotate", "flip", "shift", "noise", "threshold")
SPATIAL_OPS = ("rotate", "flip", "shift")   # image stream only
# Offsets from the seed, so each source of randomness has its own stream.
SEED_SHUFFLE, SEED_EMB_NOISE, SEED_ABLATE, SEED_BEND_NOISE, SEED_WEIGHTS = 3, 5, 7, 8, 9


@dataclass(frozen=True)
class Settings:
    prompt_a: str
    prompt_b: str
    experiment: str                 # "seed_grid" | "guidance_sweep" | "interpolation"
    seed: int
    steps: int
    t: float = 0.0                  # blend A -> B, used by seed_grid and guidance_sweep
    guidance: float = 4.0           # used by seed_grid and interpolation
    g_from: float = 1.5             # guidance_sweep
    g_to: float = 7.0
    t_from: float = 0.0             # interpolation
    t_to: float = 1.0
    width: int = DEFAULT_SIZE
    height: int = DEFAULT_SIZE
    emb_scale: float = 1.0          # R4
    emb_noise: float = 0.0          # R4
    emb_shuffle: bool = False       # R4
    te_layers: tuple = STANDARD_LAYERS  # R5
    sigma_mult: float = 1.0         # R6
    sigma_offset: float = 0.0       # R6
    stop_at: float = 1.0            # R7, fraction of steps actually run
    show_x0: bool = False           # R7, decode predicted clean image instead of x_t
    bend_block: str | None = None   # R8, "double.0".."double.4" | "single.0".."single.19"
    bend_stream: str = "image"      # "image" | "text"
    bend_op: str = "none"
    bend_strength: float = 0.0
    bend_from: float = 0.0          # active window as fraction of the schedule
    bend_to: float = 1.0
    weight_block: str | None = None # R9
    weight_noise: float = 0.0       # R9, relative to each weight tensor's std
    grid: int | None = None         # images in the grid, None = GRID_N
    tools: tuple = ()               # Tab 4: tools_sd15.ToolUse entries, SD 1.5 only


@dataclass
class Result:
    items: list                     # [(PIL.Image, caption)] for gr.Gallery
    grid: Image.Image               # composite for the board
    warnings: list = field(default_factory=list)
    seconds: float = 0.0
    timings: dict = field(default_factory=dict)


class Clock:
    """Progress through the schedule, read by the bending hooks."""

    progress = 0.0


def fmt_num(x: float, decimals: int) -> str:
    """German number formatting with a decimal comma."""
    return f"{x:.{decimals}f}".replace(".", ",")


def clamp_size(x: int) -> int:
    return max(16, int(round(x / 16)) * 16)


def sample_specs(s: Settings, n: int | None = None) -> list[tuple[int, float, float, str]]:
    """One (seed, t, guidance, caption) tuple per image of the grid."""
    n = n or s.grid or GRID_N
    if s.experiment == "seed_grid":
        return [(s.seed + i, s.t, s.guidance, f"Seed {s.seed + i}") for i in range(n)]
    if s.experiment == "guidance_sweep":
        gs = np.linspace(s.g_from, s.g_to, n)
        return [(s.seed, s.t, float(g), f"Guidance {fmt_num(g, 1)}") for g in gs]
    if s.experiment == "interpolation":
        ts = np.linspace(s.t_from, s.t_to, n)
        return [(s.seed, float(t), s.guidance, f"t = {fmt_num(t, 3)}") for t in ts]
    raise ValueError(f"unknown experiment {s.experiment!r}")


def make_grid(images: list[Image.Image], cols: int | None = None, cell: int = 256,
              gap: int = 4, bg=(255, 255, 255)) -> Image.Image:
    """Composite of all images, the long side of each cell is `cell` px (smaller for wide grids)."""
    n = len(images)
    cols = cols or math.ceil(math.sqrt(n))
    rows = math.ceil(n / cols)
    cell = min(cell, 1040 // cols)                 # a 6 x 6 grid stays about as wide as a 4 x 4 one
    w0, h0 = images[0].size
    scale = cell / max(w0, h0)
    cw, ch = max(1, round(w0 * scale)), max(1, round(h0 * scale))
    grid = Image.new("RGB", (cols * cw + (cols - 1) * gap, rows * ch + (rows - 1) * gap), bg)
    for i, im in enumerate(images):
        r, c = divmod(i, cols)
        grid.paste(im.convert("RGB").resize((cw, ch), Image.LANCZOS), (c * (cw + gap), r * (ch + gap)))
    return grid


def block_ids(transformer) -> list[str]:
    """"double.0".."double.4" (text and image still separate), then "single.0".."single.19"."""
    return ([f"double.{k}" for k in range(len(transformer.transformer_blocks))]
            + [f"single.{k}" for k in range(len(transformer.single_transformer_blocks))])


def get_block(transformer, block_id: str):
    kind, k = block_id.split(".")
    blocks = {"double": transformer.transformer_blocks, "single": transformer.single_transformer_blocks}[kind]
    return blocks[int(k)]


def bend_tensor(h: torch.Tensor, op: str, s: float, lat_hw: tuple[int, int] | None = None,
                mask: torch.Tensor | None = None, gen: torch.Generator | None = None) -> torch.Tensor:
    """R8: one bending operation on tokens h (B, L, D). Spatial ops need lat_hw = (lat_h, lat_w)."""
    if op == "amplify":
        return h * (1 + 4 * s)
    if op == "invert":
        return h * (1 - 2 * s)
    if op == "ablate":
        return h * mask
    if op == "noise":
        z = torch.randn(h.shape, generator=gen, device=h.device, dtype=torch.float32)
        return h + (2 * s * h.float().std() * z).to(h.dtype)
    if op == "threshold":
        return torch.where(h.abs() > 2 * s * h.float().std(), h, torch.zeros_like(h))
    B, L, D = h.shape
    lat_h, lat_w = lat_hw
    x = h.reshape(B, lat_h, lat_w, D).permute(0, 3, 1, 2)          # spatial view (B, D, h, w)
    if op == "rotate":
        # 90 degrees keeps the shape only for square latents, otherwise turn by 180.
        x = torch.lerp(x, torch.rot90(x, 1 if lat_h == lat_w else 2, dims=(-2, -1)), s)
    elif op == "flip":
        x = torch.lerp(x, torch.flip(x, dims=(-1,)), s)
    elif op == "shift":
        x = torch.roll(x, shifts=round(s * lat_w / 2), dims=-1)
    else:
        raise ValueError(f"unknown bending op {op!r}")
    return x.permute(0, 2, 3, 1).reshape(B, L, D)


class Engine:
    """FLUX.2 [klein] engine. engine_sd15.SD15Engine reuses render, conditioning and the interventions."""

    latent_factor = 16              # pixels per packed latent token: VAE scale 8 x patch 2
    standard_layers = (9, 18, 27)   # Qwen3 hidden states the pipeline reads

    def __init__(self, pipe, device: str | torch.device = "cuda"):
        from diffusers import FlowMatchEulerDiscreteScheduler
        from diffusers.pipelines.flux2.pipeline_flux2_klein import (
            Flux2KleinPipeline, compute_empirical_mu, retrieve_timesteps)

        self.pipe = pipe
        self.device = torch.device(device)
        self.P = Flux2KleinPipeline
        self._Scheduler = FlowMatchEulerDiscreteScheduler
        self._mu = compute_empirical_mu
        self._retrieve_timesteps = retrieve_timesteps
        self.clock = Clock()
        self._emb_cache: OrderedDict = OrderedDict()
        self._lock = threading.Lock()
        self.post_step: list = []       # functions latent -> latent after every sampler step (Tab 4)
        self.encode_tag = ""            # part of the encoding cache key while a tool changes the text encoder
        self.run_notes: list = []       # notes for the status line from the current render call
        for m in (pipe.transformer, pipe.text_encoder, pipe.vae):
            m.requires_grad_(False)
            m.eval()

    @classmethod
    def from_pretrained(cls, model_id: str = MODEL_ID, device: str = "cuda") -> "Engine":
        from diffusers import Flux2KleinPipeline

        pipe = Flux2KleinPipeline.from_pretrained(model_id, dtype=torch.bfloat16).to(device)
        # Extreme latents overflow less in float32. The VAE config sets force_upcast anyway.
        pipe.vae.to(torch.float32)
        # Tiles only above the VAE's sample_size (1024 px), so smaller images are unchanged.
        # Without it a float32 decode at 1536 px runs out of memory on a 24 GB card.
        pipe.vae.enable_tiling()
        return cls(pipe, device)

    def block(self, block_id: str):
        return get_block(self.pipe.transformer, block_id)

    def text_bend_note(self, s: Settings) -> str | None:
        """A status line when a text-stream bend cannot act, else None."""
        if s.bend_stream == "text" and s.bend_block == block_ids(self.pipe.transformer)[-1]:
            # The transformer drops the text tokens after its last block.
            return ("Nach dem letzten Block liest das Netz den Text-Strom nicht mehr. "
                    "Dieser Eingriff ändert deshalb nichts. Wählt einen früheren Block.")
        return None

    def warmup(self, size: int = DEFAULT_SIZE, n: int = 2, steps: int = 2) -> float:
        s = Settings("warm-up", "warm-up", "seed_grid", seed=0, steps=steps, width=size, height=size,
                     te_layers=self.standard_layers)
        return self.render(s, sample_specs(s, n)).seconds

    # ---------------------------------------------------------------- text

    def encode(self, prompt: str, layers: tuple = STANDARD_LAYERS, max_len: int = MAX_SEQ) -> torch.Tensor:
        """Stacked Qwen3 hidden states (1, max_len, 3 * 2560), cached."""
        key = (prompt, tuple(layers), max_len)
        with self._lock:
            if key in self._emb_cache:
                self._emb_cache.move_to_end(key)
                return self._emb_cache[key]
        with torch.inference_mode():
            emb = self.P._get_qwen3_prompt_embeds(
                self.pipe.text_encoder, self.pipe.tokenizer, prompt, dtype=torch.bfloat16,
                device=self.device, max_sequence_length=max_len, hidden_states_layers=tuple(layers))
        with self._lock:
            self._emb_cache[key] = emb
            while len(self._emb_cache) > EMB_CACHE_SIZE:
                self._emb_cache.popitem(last=False)
        return emb

    def clear_cache(self) -> None:
        with self._lock:
            self._emb_cache.clear()

    def conditioning(self, s: Settings, ts: list[float], max_len: int) -> tuple[torch.Tensor, torch.Tensor]:
        """cond (n, L, D) and uncond (1, L, D), both bfloat16."""
        n_layers = self.pipe.text_encoder.config.num_hidden_layers   # hidden_states has n_layers + 1 entries
        if any(not 0 <= k <= n_layers for k in s.te_layers):
            raise ValueError(f"text encoder layers must lie in 0..{n_layers}, got {s.te_layers}")
        emb_a, emb_b = self.prompt_embeddings(s, max_len)
        uncond = self.encode("", s.te_layers, max_len)
        # R2: a linear blend in the stacked hidden-state space, padding positions included.
        t = torch.tensor(ts, device=self.device, dtype=torch.float32).view(-1, 1, 1)
        cond = emb_a + t * (emb_b - emb_a)
        # R4 acts on cond only, in this order: shuffle, scale, noise. uncond stays untouched.
        if s.emb_shuffle:
            cond = self.shuffle_cond(s, cond, max_len)
        if s.emb_scale != 1.0:
            cond = cond * s.emb_scale
        if s.emb_noise != 0:
            # The same noise for every sample, so a seed grid stays comparable.
            z = torch.randn(cond.shape[1:], generator=torch.Generator("cpu").manual_seed(s.seed + SEED_EMB_NOISE))
            cond = cond + s.emb_noise * z.to(self.device) * cond.std(dim=-1, keepdim=True)
        return cond.to(torch.bfloat16), uncond

    def prompt_embeddings(self, s: Settings, max_len: int) -> tuple[torch.Tensor, torch.Tensor]:
        return (self.encode(s.prompt_a, s.te_layers, max_len).float(),
                self.encode(s.prompt_b, s.te_layers, max_len).float())

    def shuffle_cond(self, s: Settings, cond: torch.Tensor, max_len: int) -> torch.Tensor:
        """R4: permute the hidden states at the prompt's token positions (template and padding stay)."""
        pos = sorted(set(self.content_positions(s.prompt_a, max_len))
                     | set(self.content_positions(s.prompt_b, max_len)))
        if len(pos) > 1:
            perm = torch.randperm(len(pos), generator=torch.Generator("cpu").manual_seed(s.seed + SEED_SHUFFLE))
            idx = torch.tensor(pos, device=self.device)
            cond[:, idx] = cond[:, idx[perm.to(self.device)]]
        return cond

    def content_positions(self, prompt: str, max_len: int = MAX_SEQ) -> list[int]:
        """Token positions of the user's prompt inside the chat template (no template tokens, no padding)."""
        tok = self.pipe.tokenizer

        def template(content: str) -> str:
            return tok.apply_chat_template([{"role": "user", "content": content}], tokenize=False,
                                           add_generation_prompt=True, enable_thinking=False)

        text = template(prompt)
        start = len(template("\x00").split("\x00")[0])
        end = start + len(prompt)
        if text[start:end] != prompt:
            raise ValueError("chat template does not contain the prompt verbatim")
        enc = tok(text, return_offsets_mapping=True, padding="max_length", truncation=True, max_length=max_len)
        return [i for i, (a, b) in enumerate(enc["offset_mapping"]) if b > a and a >= start and b <= end]

    # ---------------------------------------------------------------- loop

    @contextmanager
    def interventions(self, s: Settings, lat_h: int, lat_w: int, txt_len: int):
        """Hooks (R8) and weight noise (R9) for exactly one render call, removed on any exit.

        Safe only because the Space runs one GPU job at a time.
        """
        with self.weight_noise(s), self.bending(s, lat_h, lat_w, txt_len):
            yield

    @contextmanager
    def bending(self, s: Settings, lat_h: int, lat_w: int, txt_len: int):
        """R8: a forward hook on one block changes its image or text tokens inside a phase window."""
        if not s.bend_block or s.bend_op == "none" or s.bend_strength == 0:
            yield
            return
        op, st, image = s.bend_op, float(s.bend_strength), s.bend_stream == "image"
        if op not in BEND_OPS:
            raise ValueError(f"unknown bending op {op!r}")
        if not image and op in SPATIAL_OPS:
            raise ValueError(f"{op} works on the image stream only")
        block = self.block(s.bend_block)
        double = s.bend_block.startswith("double")
        mask = gen = None
        if op == "ablate":
            dim = self.pipe.transformer.config.num_attention_heads * self.pipe.transformer.config.attention_head_dim
            off = torch.randperm(dim, generator=torch.Generator("cpu").manual_seed(s.seed + SEED_ABLATE))
            mask = torch.ones(dim)
            mask[off[:round(min(max(st, 0.0), 1.0) * dim)]] = 0.0   # typed values may leave 0..1
            mask = mask.to(self.device, torch.bfloat16)
        elif op == "noise":
            gen = torch.Generator(self.device).manual_seed(s.seed + SEED_BEND_NOISE)
        clock = self.clock

        def bend(h, spatial):
            return bend_tensor(h, op, st, (lat_h, lat_w) if spatial else None, mask, gen)

        def hook(module, args, output):
            if not (s.bend_from <= clock.progress <= s.bend_to):
                return None                                    # outside the window: unchanged
            if double:                                         # (text, image)
                txt, img = output
                return (txt, bend(img, True)) if image else (bend(txt, False), img)
            txt, img = output[:, :txt_len], output[:, txt_len:]   # single stream: [text, image]
            return torch.cat([txt, bend(img, True)] if image else [bend(txt, False), img], dim=1)

        handle = block.register_forward_hook(hook)
        try:
            yield
        finally:
            handle.remove()

    @contextmanager
    def weight_noise(self, s: Settings):
        """R9: add noise to one block's weight matrices, relative to each tensor's std. Restored on exit."""
        if not s.weight_block or s.weight_noise == 0:
            yield
            return
        block = self.block(s.weight_block)
        gen = torch.Generator(self.device).manual_seed(s.seed + SEED_WEIGHTS)
        saved = []
        try:
            with torch.no_grad():
                for p in block.parameters():
                    if p.dim() < 2:
                        continue
                    saved.append((p, p.detach().clone()))
                    z = torch.randn(p.shape, generator=gen, device=p.device, dtype=torch.float32)
                    w = p.float()
                    p.copy_((w + s.weight_noise * w.std() * z).to(p.dtype))
            yield
        finally:
            with torch.no_grad():
                for p, orig in saved:
                    p.copy_(orig)

    def denoise(self, s: Settings, cond: torch.Tensor, uncond: torch.Tensor, seeds: list[int],
                guidances: list[float], on_step=None) -> torch.Tensor:
        """Run the sampler for one chunk. Returns packed latents (n, h*w, 128) in float32."""
        P, dev, N = self.P, self.device, s.steps
        n = len(seeds)
        lat_h, lat_w = s.height // 16, s.width // 16          # vae scale 8 x patch 2
        noise = torch.stack([torch.randn((128, lat_h, lat_w), generator=torch.Generator("cpu").manual_seed(sd))
                             for sd in seeds])                 # float32, device independent
        img_ids = P._prepare_latent_ids(noise).to(dev)
        lat = P._pack_latents(noise).to(dev)                   # (n, h*w, 128), stays float32

        sched = self._Scheduler.from_config(self.pipe.scheduler.config)   # fresh: the scheduler is stateful
        sigmas = np.linspace(1.0, 1.0 / N, N)
        mu = self._mu(image_seq_len=lat.shape[1], num_steps=N)
        timesteps, _ = self._retrieve_timesteps(sched, N, dev, sigmas=sigmas, mu=mu)
        sched.set_begin_index(0)
        sig = sched.sigmas.float().cpu().tolist()              # shifted sigmas, N + 1 values ending in 0
        n_run = max(1, round(N * s.stop_at))

        enc = torch.cat([cond, uncond.expand(n, -1, -1)]).to(torch.bfloat16)
        txt_ids = P._prepare_text_ids(enc).to(dev)
        ids2 = torch.cat([img_ids, img_ids])
        g = torch.tensor(guidances, device=dev, dtype=torch.float32).view(-1, 1, 1)
        x0_pred = None
        for i, t in enumerate(timesteps[:n_run]):
            self.clock.progress = i / N
            sigma = sig[i]
            t_model = sigma * s.sigma_mult + s.sigma_offset    # R6: the model is told a different time
            x = torch.cat([lat, lat]).to(torch.bfloat16)
            v = self.pipe.transformer(
                hidden_states=x, timestep=torch.full((2 * n,), t_model, device=dev), guidance=None,
                encoder_hidden_states=enc, txt_ids=txt_ids, img_ids=ids2, return_dict=False)[0].float()
            v_c, v_u = v.chunk(2)
            v = v_u + g * (v_c - v_u)                          # R3: per-sample guidance
            x0_pred = lat - sigma * v
            lat = sched.step(v, t, lat, return_dict=False)[0]  # float32 stays float32 because v is float32
            if on_step is not None:
                on_step()
        return x0_pred if (s.show_x0 and n_run < N) else lat

    def decode(self, final: torch.Tensor, lat_h: int, lat_w: int) -> tuple[list[Image.Image], list[bool]]:
        vae = self.pipe.vae
        n = final.shape[0]
        finite = torch.isfinite(final).flatten(1).all(1).tolist()
        final = torch.nan_to_num(final.float(), 0.0, 0.0, 0.0)
        lat = final.view(n, lat_h, lat_w, 128).permute(0, 3, 1, 2)   # row-major, as _prepare_latent_ids
        mean = vae.bn.running_mean.view(1, -1, 1, 1).to(lat.device, torch.float32)
        std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(lat.device, torch.float32)
        lat = self.P._unpatchify_latents(lat * std + mean)
        per_call = max(1, DECODE_PIXELS // (lat_h * lat_w * 256))
        images = []
        for k in range(0, n, per_call):
            img = vae.decode(lat[k:k + per_call].to(vae.dtype), return_dict=False)[0]
            img = torch.nan_to_num(img.float(), 0.0, 0.0, 0.0)
            images += self.pipe.image_processor.postprocess(img, output_type="pil")
        return images, finite

    # ---------------------------------------------------------------- entry points

    def render(self, s: Settings, specs: list[tuple[int, float, float, str]], max_seq: int = MAX_SEQ,
               per_chunk: int | None = None, progress=None) -> Result:
        """Render explicit per-sample specs. `progress(done, total)` is called once per step.

        Images per forward pass default to MAX_BATCH at DEFAULT_SIZE, scaled by pixel count.
        """
        s = _checked(s)
        t0 = time.perf_counter()
        lat_h, lat_w = s.height // self.latent_factor, s.width // self.latent_factor
        n = len(specs)
        per_chunk = per_chunk or max(1, math.floor(MAX_BATCH * DEFAULT_SIZE ** 2 / (s.width * s.height)))
        chunks = [specs[k:k + per_chunk] for k in range(0, n, per_chunk)]
        n_run = max(1, round(s.steps * s.stop_at))
        total, done = len(chunks) * n_run, 0

        def on_step():
            nonlocal done
            done += 1
            if progress is not None:
                progress(done, total)

        timings = {}
        self.run_notes = []
        try:
            with self.interventions(s, lat_h, lat_w, max_seq), torch.inference_mode():
                _sync()
                t1 = time.perf_counter()
                finals = []
                cond_all, uncond = self.conditioning(s, [sp[1] for sp in specs], max_seq)
                _sync()
                timings["encode"] = time.perf_counter() - t1
                t1 = time.perf_counter()
                for k, chunk in enumerate(chunks):
                    lo = k * per_chunk
                    finals.append(self.denoise(s, cond_all[lo:lo + len(chunk)], uncond,
                                               [sp[0] for sp in chunk], [sp[2] for sp in chunk], on_step))
                _sync()
                timings["denoise"] = time.perf_counter() - t1
                t1 = time.perf_counter()
                images, finite = self.decode(torch.cat(finals), lat_h, lat_w)
                _sync()
                timings["decode"] = time.perf_counter() - t1
        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            import gradio as gr
            raise gr.Error("Zu wenig GPU-Speicher. Wählt eine kleinere Bildgröße.")

        items, warnings = [], []
        for im, ok, sp in zip(images, finite, specs):
            items.append((im, sp[3] if ok else f"{sp[3]} (numerischer Überlauf)"))
        note = self.text_bend_note(s) if (s.bend_block and s.bend_op != "none" and s.bend_strength != 0) else None
        if note:
            warnings.append(note)
        warnings += self.run_notes
        n_bad = finite.count(False)
        if n_bad:
            warnings.append(f"Bei {n_bad} von {n} Bildern ist die Rechnung numerisch übergelaufen. "
                            "Diese Bilder zeigen nur einen Ersatz für die ungültigen Werte.")
        cols = math.ceil(math.sqrt(n))
        return Result(items=items, grid=make_grid(images, cols), warnings=warnings,
                      seconds=time.perf_counter() - t0, timings=timings)

    def generate(self, s: Settings, progress=None) -> Result:
        return self.render(s, sample_specs(s), progress=progress)


def load_engine(device: str = "cuda") -> Engine:
    """The engine of the active model family."""
    if FAMILY.key == "sd15":
        from engine_sd15 import SD15Engine

        return SD15Engine.from_pretrained(MODEL_ID, device)
    return Engine.from_pretrained(MODEL_ID, device)


def _checked(s: Settings) -> Settings:
    from dataclasses import replace

    if s.experiment not in EXPERIMENTS:
        raise ValueError(f"unknown experiment {s.experiment!r}")
    if s.bend_op not in BEND_OPS:
        raise ValueError(f"unknown bending op {s.bend_op!r}")
    if s.bend_stream not in ("image", "text"):
        raise ValueError(f"unknown stream {s.bend_stream!r}")
    if s.tools:
        from tools_sd15 import TOOLS

        unknown = [u.key for u in s.tools if u.key not in TOOLS]
        if unknown:
            raise ValueError(f"unknown tools {unknown}")
    return replace(s, prompt_a=s.prompt_a or "", prompt_b=s.prompt_b or "",
                   width=clamp_size(s.width), height=clamp_size(s.height),
                   steps=max(1, int(s.steps)), seed=int(s.seed),
                   te_layers=tuple(int(k) for k in s.te_layers),
                   stop_at=min(1.0, max(0.0, float(s.stop_at))),
                   grid=None if s.grid is None else max(1, min(max(GRID_SIZES), int(s.grid))))


def _sync() -> None:
    if torch.cuda.is_available():
        torch.cuda.synchronize()