onw-Image 0.3: Panther Lake ~77 -> ~48 s per 1024 image (fused RMSNorm kernels, FlashAttention-2, async zero-copy blocks; same model files)
6b8ea09 verified | """onw-Image runtime v2: conv-form DiT graphs (zit.cdit) on an OpenVINO device; host does the small parts. | |
| eng = Engine(conv_dir, device="NPU", size=(1024, 1024)) | |
| lat = eng.denoise("prompt", seed=42) | |
| NPU graphs (DiT blocks, text encoder segments, VAE) are compiled once into exported blobs in <conv>/npu_blobs3 and | |
| imported from there (with NPU_TURBO). With onw's NPU kit (zit.npukit) the blobs are weightless: the weights are rebuilt | |
| on the NPU from the model's .bin files at import, so the blobs hold the graphs only (0.1 kept the weights in every | |
| blob: ~36 GB next to a 10 GB model on Panther Lake). | |
| """ | |
| import json, os, time | |
| import numpy as np | |
| import openvino as ov | |
| from . import host, dit, te, cdit, npukit | |
| # (0.1: npu_blobs2 = full blobs, and an OpenVINO CACHE_DIR for the text encoder and VAE; ONW_IMAGE_BLOB_DIR: tests) | |
| BLOB_DIR = os.environ.get("ONW_IMAGE_BLOB_DIR", "npu_blobs3") | |
| CHAT = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" | |
| D = cdit.D | |
| CAP_BUCKETS = (64, 256, 512) # caption token buckets (each needs its own compiled DiT set; Japanese ~1 token/char) | |
| def cap_bucket(n): | |
| for b in CAP_BUCKETS: | |
| if n <= b: | |
| return b | |
| raise ValueError("prompt too long (max 512 tokens)") | |
| def f16(a): | |
| return np.ascontiguousarray(a, np.float16) | |
| class HostData(dict): | |
| """Host-side arrays: <conv>/host/<key>.npy memory-mapped on first use (the 778 MB token table is only touched | |
| row by row), else everything from host.npz.""" | |
| def __init__(self, conv): | |
| super().__init__() | |
| self.dir = os.path.join(conv, "host") | |
| if not os.path.isdir(self.dir): | |
| self.update(dict(np.load(os.path.join(conv, "host.npz")))) | |
| self.dir = None | |
| def __missing__(self, k): | |
| if self.dir is None: | |
| raise KeyError(k) | |
| v = np.load(os.path.join(self.dir, k.replace(":", "@") + ".npy"), mmap_mode="r") | |
| self[k] = v | |
| return v | |
| def __contains__(self, k): | |
| return dict.__contains__(self, k) or (self.dir is not None and | |
| os.path.exists(os.path.join(self.dir, k.replace(":", "@") + ".npy"))) | |
| def _arr(x): | |
| """numpy view of a block output (an ov.Tensor with zero-copy chaining, else already numpy).""" | |
| return x.data if isinstance(x, ov.Tensor) else x | |
| class Engine: | |
| def __init__(self, conv, device="NPU", size=(1024, 1024), cache_dir=None, log=print, te_device=None, props=None, | |
| lazy=None): | |
| self.conv, self.device, self.log = conv, device, log | |
| self.lazy = (device == "CPU") if lazy is None else lazy # CPU checks: compile per use (f32 copies are big) | |
| self.Hh = HostData(conv) | |
| self.V = dict(np.load(os.path.join(conv, "dit2_vec.npz"))) | |
| self.core = npukit.core() | |
| if cache_dir and device != "NPU": # (NPU graphs: exported blobs, see _npu_request) | |
| self.core.set_property({"CACHE_DIR": cache_dir}) | |
| self.props = dict(props or {}) | |
| self.export_only = False # precompile children: compile + export, no import | |
| # block outputs stay in the plugin's tensors and feed the next block directly ($ONW_IMAGE_ZEROCOPY=0: copies) | |
| self.zero_copy = device == "NPU" and os.environ.get("ONW_IMAGE_ZEROCOPY", "1") != "0" | |
| self._wmodels = {} # graph key -> its ov.Model (a weightless import reads it) | |
| # exported DiT blobs are imported with NPU_TURBO (higher NPU clock: ~15% faster convolutions on Panther Lake); | |
| # it is a runtime setting at import - passed to compile_model it would change the cache key and recompile | |
| turbo = os.environ.get("ONW_IMAGE_TURBO", "1") != "0" | |
| self.import_props = {**self.props, "NPU_TURBO": "YES"} if turbo else dict(self.props) | |
| self.te_device = te_device or device | |
| self.H, self.W = int(size[0]), int(size[1]) | |
| self.ht, self.wt = self.H // 16, self.W // 16 | |
| self.n_img = self.ht * self.wt | |
| self._te, self._blocks, self._bucket, self._models = {}, {}, None, [] | |
| # zit.cdit_u builds (dit2/uniform.json): every block of a kind is the same graph -> one weightless blob per kind | |
| self.uniform = os.path.exists(os.path.join(conv, "dit2", "uniform.json")) | |
| # 0.3: uniform blocks are rebuilt in memory in a faster graph form (zit.cdit_u2: same weight constants, no new | |
| # model files): G4 (default; FlashAttention-2 with the kit's compiler) or G3; $ONW_IMAGE_GRAPH=0.2 keeps the | |
| # graphs as stored | |
| gf = os.environ.get("ONW_IMAGE_GRAPH", "g4") | |
| self.graph_form = gf if self.uniform and gf in ("g3", "g4") else "" | |
| from tokenizers import Tokenizer | |
| self.tok = Tokenizer.from_file(os.path.join(conv, "tokenizer.json")) | |
| # ---------- compile ---------- | |
| def _shapes(self, streams, sfxs=("",)): | |
| sh = {} | |
| for s, (h, w) in streams.items(): | |
| n = h * w | |
| sh[f"x_{s}"] = [1, D, h, w] | |
| sh[f"mask_{s}"] = [1, 1, 1, n] | |
| for x in sfxs: | |
| sh.update({f"{t}_{s}{x}": [1, 1, n, cdit.HD] for t in ("cq", "sq", "ck", "sk")}) | |
| for x in sfxs: | |
| sh.update({n + x: [1, D, 1, 1] for n in ("m_a", "g_a", "m_f", "g_f")}) | |
| return sh | |
| def groups(self): | |
| """[[layer names]] when the main layers are chained in groups (dit2/groups.json), else one per layer.""" | |
| p = os.path.join(self.conv, "dit2", "groups.json") | |
| return json.load(open(p)) if os.path.exists(p) else [[f"layers.{i}"] for i in range(30)] | |
| def _graph_name(self, gi, grp): | |
| return f"group.{gi}" if len(grp) > 1 or grp[0].startswith("group") else grp[0] | |
| def blob_path(self, b, shapes): | |
| """Exported NPU blob for graph b at these shapes (<conv>/npu_blobs3/<b>_<HxW>[_<HxW>].blob); uniform builds | |
| with weightless blobs: one per kind (<kind>.all_...: the 30 main layers, the noise / context refiners).""" | |
| tag = "_".join(f"{v[2]}x{v[3]}" for k, v in sorted(shapes.items()) if k.startswith("x_")) | |
| if self.uniform and npukit.weightless() and "." in b: | |
| b = b.split(".")[0] + ".all" | |
| form = f"-{self.graph_form}" if self.graph_form else "" | |
| return os.path.join(self.conv, BLOB_DIR, f"{b}{form}_{tag}.blob") | |
| def _npu_request(self, path, model, device, cfg=None): | |
| """Infer request for an NPU graph from its exported blob `path` - compiled and exported first when missing | |
| (weightless with the kit). Imports carry NPU_TURBO (a runtime setting: given to the compile it would be part | |
| of the blob). export_only: compile + export, return None.""" | |
| cfg = dict(self.props if cfg is None else cfg) | |
| if npukit.blob_file(path) is None: | |
| flash = self.graph_form == "g4" and "_s" not in os.path.basename(path) and not os.path.basename(path).startswith("vae") | |
| if flash: # (the kit's compiler: FlashAttention-2 for the DiT's SDPA) | |
| os.environ["ONW_FLASH_SDPA"] = "1" | |
| try: | |
| npukit.compile_blob(self.core, model, device, cfg, path, log=self.log) | |
| finally: | |
| if flash: | |
| os.environ.pop("ONW_FLASH_SDPA", None) | |
| if self.export_only: | |
| return None | |
| self._wmodels[path] = model | |
| imp = {**cfg, **{k: v for k, v in self.import_props.items() if k == "NPU_TURBO"}} | |
| return npukit.import_blob(self.core, model, device, imp, path).create_infer_request() | |
| def _compile_block(self, b, streams, n_sub=1): | |
| sh = self._shapes(streams, [f"@{j}" for j in range(n_sub)] if n_sub > 1 else ("",)) | |
| xml = os.path.join(self.conv, "dit2", b + ".xml") | |
| if self.graph_form: | |
| from . import cdit_u2 | |
| m = cdit_u2.build(xml, cdit_u2.G4 if self.graph_form == "g4" else cdit_u2.G3, self.core) | |
| else: | |
| m = self.core.read_model(xml) | |
| m.reshape(sh) | |
| if self.device == "NPU": | |
| return self._npu_request(self.blob_path(b, sh), m, self.device) | |
| self._models.append(m) # CPU graphs may reference the model constants | |
| return self.core.compile_model(m, self.device, self.props).create_infer_request() | |
| def load_blocks(self, bucket): | |
| if self._bucket == bucket: | |
| return | |
| self._blocks, self._models = {}, [] | |
| t = time.time() | |
| img, cap = (self.ht, self.wt), (1, bucket) | |
| self._groups = self.groups() | |
| self._nsub = {} | |
| self._streams = {**{f"noise_refiner.{i}": {"img": img} for i in range(2)}, | |
| **{f"context_refiner.{i}": {"cap": cap} for i in range(2)}} | |
| for gi, grp in enumerate(self._groups): | |
| name = self._graph_name(gi, grp) | |
| self._streams[name] = {"img": img, "cap": cap} | |
| self._nsub[name] = len(grp) if name.startswith("group") else 1 | |
| if self.lazy: | |
| self._blocks = {b: None for b in self._streams} | |
| self._bucket = bucket | |
| return | |
| for i in range(2): | |
| self._blocks[f"noise_refiner.{i}"] = self._compile_block(f"noise_refiner.{i}", {"img": img}) | |
| self._blocks[f"context_refiner.{i}"] = self._compile_block(f"context_refiner.{i}", {"cap": cap}) | |
| for gi, grp in enumerate(self._groups): | |
| name = self._graph_name(gi, grp) | |
| self._blocks[name] = self._compile_block(name, {"img": img, "cap": cap}, self._nsub[name]) | |
| self.log(f" {name} ready ({time.time()-t:.0f}s)") | |
| self._bucket = bucket | |
| # ---------- text ---------- | |
| def _te_paths(self): | |
| d = os.path.join(self.conv, "te_seg") | |
| if os.path.isdir(d): | |
| n = len([f for f in os.listdir(d) if f.endswith(".xml")]) | |
| return [os.path.join(d, f"te_{i}.xml") for i in range(n)] | |
| return [os.path.join(self.conv, "te", "te.xml")] | |
| def encode_compile(self, S): | |
| """Compile (or import from the cache) the text encoder segments for prompt bucket S.""" | |
| if S in self._te: | |
| return | |
| reqs = [] | |
| for path in self._te_paths(): | |
| m = self.core.read_model(path) | |
| m.reshape({"emb": [1, S, te.HID], "cos": [1, S, 1, 64], "sin": [1, S, 1, 64], "mask": [1, 1, S, S]}) | |
| if self.te_device == "NPU": | |
| name = os.path.splitext(os.path.basename(path))[0] | |
| reqs.append(self._npu_request(os.path.join(self.conv, BLOB_DIR, f"{name}_s{S}.blob"), m, "NPU")) | |
| else: | |
| reqs.append(self.core.compile_model(m, self.te_device, self.props).create_infer_request()) | |
| del m | |
| self._te = {S: reqs} | |
| def embed(self, ids): | |
| """Token embeddings [n, 2560] f32: from the INT8 table (te.embed_tokens_q + per-row scale te.embed_tokens_s, | |
| 0.39 GB) when the model has it, else the f16 one (0.78 GB).""" | |
| if "te.embed_tokens_q" in self.Hh: | |
| return self.Hh["te.embed_tokens_q"][ids].astype(np.float32) * \ | |
| self.Hh["te.embed_tokens_s"][ids].astype(np.float32)[:, None] | |
| return self.Hh["te.embed_tokens"][ids].astype(np.float32) | |
| def encode(self, prompt): | |
| """Prompt -> caption features [n, 2560] (text encoder segments compiled per prompt bucket, then kept until | |
| free_text_encoder()).""" | |
| ids = self.tok.encode(CHAT.format(prompt)).ids[:512] | |
| S = next(b for b in (64, 128, 256, 512) if len(ids) <= b) | |
| self.encode_compile(S) | |
| x = np.zeros((1, S, te.HID), np.float32) | |
| x[0, :len(ids)] = self.embed(ids) | |
| cos, sin, mask = te.host_inputs(ids, S) | |
| for rq in self._te[S]: | |
| rq.infer({"emb": x, "cos": cos, "sin": sin, "mask": mask}) | |
| x = rq.get_output_tensor(0).data.astype(np.float32).copy() | |
| return x[0, :len(ids)].copy() | |
| def free_text_encoder(self): | |
| self._te = {} | |
| import gc | |
| gc.collect() | |
| # ---------- DiT ---------- | |
| def _run(self, b, feed): | |
| return self._finish(b, self._start(b, feed)) | |
| def _start(self, b, feed): | |
| """Start block b on its own infer request (async on the NPU: the host prepares the next block meanwhile).""" | |
| rq = self._blocks[b] | |
| if rq is None: | |
| self._models = [] | |
| rq = self._compile_block(b, self._streams[b], self._nsub.get(b, 1)) | |
| if self.zero_copy: | |
| # NPU: a block's output tensor (the plugin's own device-visible buffer) goes straight into the next block - | |
| # no copy out to numpy and back in (~60 MB per block at 1024); numpy inputs are shared, not copied | |
| rq.start_async(feed, share_inputs=True) | |
| else: | |
| rq.infer(feed) | |
| return rq | |
| def _finish(self, b, rq): | |
| if self.zero_copy: | |
| rq.wait() | |
| out = {o.get_any_name(): rq.get_tensor(o) for o in rq.model_outputs} | |
| else: | |
| out = {o.get_any_name(): rq.get_tensor(o).data.copy() for o in rq.model_outputs} | |
| if self.lazy: | |
| self._models = [] | |
| return out | |
| def _vecs(self, b, mods=None): | |
| V = self.V | |
| a_in, a_out, f_in, f_out = (V[f"{b}:{k}"] for k in ("a_in", "a_out", "f_in", "f_out")) | |
| if mods is not None: | |
| s1, g1, s2, g2 = (m.ravel() for m in mods) | |
| a_in, a_out, f_in, f_out = a_in * s1, a_out * g1, f_in * s2, f_out * g2 | |
| out = {n: f16(v.reshape(1, D, 1, 1)) for n, v in zip(("m_a", "g_a", "m_f", "g_f"), (a_in, a_out, f_in, f_out))} | |
| for lin in ("qkv", "out", "w1", "w3", "w2"): # zit.cdit_u shared weights: the caption stream's guards | |
| v = V.get(f"{b}:p_{lin}") | |
| if v is not None: | |
| out[f"p_{lin}"] = f16(v.reshape(1, -1, 1, 1)) | |
| return out | |
| def _sublayers(self, b): | |
| if b.startswith("group"): | |
| return self._groups[int(b.split(".")[1])] | |
| return [b] | |
| def _tables(self, b, s, ang): | |
| cq, sq, ck, sk = cdit.rope_tables(ang, self.V[f"{b}:q_g"], self.V[f"{b}:k_g"]) | |
| return {f"cq_{s}": cq, f"sq_{s}": sq, f"ck_{s}": ck, f"sk_{s}": sk} | |
| def prepare(self, cap_feats): | |
| Hh = self.Hh | |
| n = len(cap_feats) | |
| cp = host.cap_padded_len(n) | |
| bucket = cap_bucket(cp) | |
| self.load_blocks(bucket) | |
| c = dit.rms_np(cap_feats, Hh["cap_embedder.0.weight"], 1e-5) | |
| c = c @ Hh["cap_embedder.1.weight"].T + Hh["cap_embedder.1.bias"] | |
| cap = np.empty((bucket, D), np.float32) | |
| cap[:n] = c | |
| cap[n:] = Hh["cap_pad_token"] # pad and bucket rows (bucket rows: masked keys, never zero) | |
| img_ids, cap_ids = host.position_ids(n, self.ht, self.wt) | |
| ids_b = np.zeros((bucket, 3), np.int64) | |
| ids_b[:cp] = cap_ids | |
| a_img, a_cap = host.rope_angles(img_ids), host.rope_angles(ids_b) | |
| mask_cap = np.zeros((1, 1, 1, bucket), np.float16) | |
| mask_cap[..., cp:] = cdit.NEG | |
| mask_img = np.zeros((1, 1, 1, self.n_img), np.float16) | |
| self.cp = cp | |
| self.const = {} | |
| self._fbc, self.fbcache_skipped, self.fbcache_log = None, 0, [] # (First Block Cache: per image) | |
| for b in self._blocks: | |
| f = {} | |
| subs = self._sublayers(b) | |
| for j, lb in enumerate(subs): | |
| sfx = f"@{j}" if len(subs) > 1 or b.startswith("group") else "" | |
| if not b.startswith("context"): | |
| f.update({k + sfx: v for k, v in self._tables(lb, "img", a_img).items()}) | |
| if not b.startswith("noise"): | |
| f.update({k + sfx: v for k, v in self._tables(lb, "cap", a_cap).items()}) | |
| if not b.startswith("context"): | |
| f["mask_img"] = mask_img | |
| if not b.startswith("noise"): | |
| f["mask_cap"] = mask_cap | |
| self.const[b] = f | |
| k = cdit.RESID_K | |
| x = f16((cap / k["context_refiner"]).T.reshape(1, D, 1, bucket)) | |
| for i in range(2): | |
| b = f"context_refiner.{i}" | |
| x = self._run(b, dict(self.const[b], x_cap=x, **self._vecs(b)))["out_cap"] | |
| self.cap = f16(_arr(x).astype(np.float32) * (k["context_refiner"] / k["layers"])) | |
| def dit(self, lat, t, trace=None): | |
| Hh, k = self.Hh, cdit.RESID_K | |
| temb = dit.t_embed(Hh, host.timestep_freq(t)) | |
| xi = host.patchify(lat) @ Hh["all_x_embedder.2-1.weight"].T + Hh["all_x_embedder.2-1.bias"] | |
| xi = f16((xi / k["noise_refiner"]).T.reshape(1, D, self.ht, self.wt)) | |
| for i in range(2): | |
| b = f"noise_refiner.{i}" | |
| xi = self._run(b, dict(self.const[b], x_img=xi, **self._vecs(b, dit.modulation(Hh, b + ".", temb))))["out_img"] | |
| if trace is not None: | |
| trace[b] = _arr(xi)[0].reshape(D, -1).T.astype(np.float32) * k["noise_refiner"] | |
| xi = f16(_arr(xi).astype(np.float32) * (k["noise_refiner"] / k["layers"])) | |
| xc = self.cap | |
| # First Block Cache ($ONW_IMAGE_FBCACHE = threshold, e.g. 0.1): after the first main block, when its residual | |
| # changed little since the previous step, the other blocks are skipped and their previous residual reused | |
| thr = float(os.environ.get("ONW_IMAGE_FBCACHE", "0") or 0) | |
| fb = getattr(self, "_fbc", None) if thr > 0 else None | |
| pending = None # (block, request) running on the NPU while the host prepares the next block | |
| x_first_in = xi | |
| skipped = False | |
| for gi, grp in enumerate(self._groups + [None]): | |
| vecs = {} | |
| if grp is not None: | |
| b = self._graph_name(gi, grp) | |
| for j, lb in enumerate(grp): | |
| sfx = f"@{j}" if b.startswith("group") else "" | |
| vecs.update({n + sfx: v for n, v in self._vecs(lb, dit.modulation(Hh, lb + ".", temb)).items()}) | |
| if pending is not None: | |
| r = self._finish(*pending) | |
| xi, xc = r["out_img"], r["out_cap"] | |
| pending = None | |
| if trace is not None: | |
| trace[self._groups[gi - 1][-1]] = np.concatenate([_arr(xi)[0].reshape(D, -1).T, _arr(xc)[0].reshape(D, -1).T[:self.cp]]).astype(np.float32) * k["layers"] | |
| if thr > 0 and gi == 1: | |
| r1 = _arr(xi).astype(np.float32) - _arr(x_first_in).astype(np.float32) | |
| x_first = _arr(xi).astype(np.float32) | |
| if fb is not None and fb["r1"].shape == r1.shape: | |
| change = float(np.abs(r1 - fb["r1"]).mean() / max(np.abs(fb["r1"]).mean(), 1e-12)) | |
| self.fbcache_log = getattr(self, "fbcache_log", []) + [change] | |
| if change < thr: | |
| xi = f16(x_first + fb["rest"]) | |
| self.fbcache_skipped = getattr(self, "fbcache_skipped", 0) + 1 | |
| skipped = True # (the reference residuals stay: errors do not drift) | |
| break | |
| if grp is None: | |
| break | |
| pending = (b, self._start(b, dict(self.const[b], x_img=xi, x_cap=xc, **vecs))) | |
| if thr > 0 and not skipped: | |
| self._fbc = {"r1": r1, "rest": _arr(xi).astype(np.float32) - x_first} | |
| out = dit.final_layer(Hh, _arr(xi)[0].reshape(D, -1).T.astype(np.float32) * k["layers"], temb) | |
| return host.unpatchify(out, self.H // 8, self.W // 8) | |
| # ---------- VAE ---------- | |
| # On the NPU the VAE decodes in tiles of VAE_TILE x VAE_TILE latents (512x512 pixels) overlapping by >= VAE_OVERLAP, | |
| # blended linearly: ONE compiled graph for every image size, and a 1024x1024 graph alone was 189 MB of NPU task | |
| # descriptors and a 14-minute compile (512: 52 MB, 99 s). ONW_IMAGE_VAE_TILE=0: the whole image at once. | |
| VAE_TILE, VAE_OVERLAP = 64, 16 | |
| def _tile_starts(n, t, ov): | |
| """Evenly spread tile starts covering [0, n) with tiles of t and overlaps >= ov.""" | |
| if n <= t: | |
| return [0] | |
| k = -(-(n - t) // (t - ov)) + 1 | |
| return [round(i * (n - t) / (k - 1)) for i in range(k)] | |
| def _vae_graph(self, dev, h, w): | |
| key = (dev, h, w) | |
| if getattr(self, "_vae_key", None) != key: | |
| m = self.core.read_model(os.path.join(self.conv, "vae", "vae_decoder.xml")) | |
| m.reshape({0: [1, 16, h, w]}) | |
| self._vae_m = m | |
| if dev == "NPU": | |
| bp = os.path.join(self.conv, BLOB_DIR, f"vae_{h}x{w}.blob") | |
| self._vae = self._npu_request(bp, m, "NPU", {"INFERENCE_PRECISION_HINT": "f16"}) | |
| else: | |
| self._vae = self.core.compile_model(m, dev).create_infer_request() | |
| self._vae_key = key if self._vae is not None else None | |
| return self._vae | |
| def decode(self, lat, device=None): | |
| """latent [16,H/8,W/8] -> PIL image (Flux VAE decoder IR in <conv>/vae).""" | |
| from PIL import Image | |
| dev = device or getattr(self, "vae_device", None) or self.device | |
| _, lh, lw = lat.shape | |
| t = self.VAE_TILE | |
| tiled = os.environ.get("ONW_IMAGE_VAE_TILE", "1") # (force: tiles on any device - checks) | |
| if (dev != "NPU" and tiled != "force") or tiled == "0" or (lh <= t and lw <= t): | |
| rq = self._vae_graph(dev, lh, lw) | |
| if rq is None: # (export_only) | |
| return None | |
| rq.infer({0: lat[None].astype(np.float32)}) | |
| img = rq.get_output_tensor(0).data[0] | |
| else: | |
| rq = self._vae_graph(dev, t, t) | |
| if rq is None: | |
| return None | |
| ys, xs = self._tile_starts(lh, t, self.VAE_OVERLAP), self._tile_starts(lw, t, self.VAE_OVERLAP) | |
| acc = np.zeros((3, lh * 8, lw * 8), np.float32) | |
| wsum = np.zeros((1, lh * 8, lw * 8), np.float32) | |
| def ramp(starts, i, n): | |
| """Weights along one axis for tile i: linear ramps over the overlaps with its neighbours.""" | |
| w = np.ones(t * 8, np.float32) | |
| s = starts[i] | |
| if i > 0: | |
| ov = (starts[i - 1] + t - s) * 8 | |
| w[:ov] = np.minimum(w[:ov], (np.arange(ov) + 0.5) / ov) | |
| if i < len(starts) - 1: | |
| ov = (s + t - starts[i + 1]) * 8 | |
| w[t * 8 - ov:] = np.minimum(w[t * 8 - ov:], (ov - np.arange(ov) - 0.5) / ov) | |
| return w | |
| for iy, y0 in enumerate(ys): | |
| wy = ramp(ys, iy, lh) | |
| for ix, x0 in enumerate(xs): | |
| wx = ramp(xs, ix, lw) | |
| rq.infer({0: np.ascontiguousarray(lat[None, :, y0:y0 + t, x0:x0 + t], np.float32)}) | |
| tile = rq.get_output_tensor(0).data[0].astype(np.float32) | |
| w2 = wy[:, None] * wx[None, :] | |
| acc[:, y0 * 8:(y0 + t) * 8, x0 * 8:(x0 + t) * 8] += tile * w2 | |
| wsum[:, y0 * 8:(y0 + t) * 8, x0 * 8:(x0 + t) * 8] += w2 | |
| img = acc / np.maximum(wsum, 1e-6) | |
| img = ((np.clip(img, -1, 1) + 1) * 127.5).round().astype(np.uint8).transpose(1, 2, 0) | |
| return Image.fromarray(img) | |
| def denoise(self, prompt, seed=42, steps=9, latents=None, callback=None): | |
| t0 = time.time() | |
| cap = self.encode(prompt) | |
| self.free_text_encoder() # 16 GB machines: the text encoder and the DiT do not fit together | |
| t1 = time.time() | |
| self.prepare(cap) | |
| if latents is None: | |
| latents = np.random.default_rng(seed).standard_normal((16, self.H // 8, self.W // 8)).astype(np.float32) | |
| lat = latents.astype(np.float32) | |
| sig = host.sigmas(steps) | |
| for i in range(steps): | |
| lat = host.euler(lat, self.dit(lat, host.model_t(sig[i])), sig[i], sig[i + 1]) | |
| if callback: | |
| callback(i, steps, lat) | |
| self.timing = {"text": t1 - t0, "denoise": time.time() - t1} | |
| return lat | |