diff --git a/extensions_built_in/inference_engine/engine.py b/extensions_built_in/inference_engine/engine.py index d90d744..113aa42 100644 --- a/extensions_built_in/inference_engine/engine.py +++ b/extensions_built_in/inference_engine/engine.py @@ -22,6 +22,7 @@ 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 @@ -119,6 +120,10 @@ def _model_key(model: dict) -> str: # 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: @@ -151,6 +156,7 @@ class Engine: 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) @@ -212,7 +218,12 @@ class Engine: "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")} + "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, @@ -224,6 +235,9 @@ class Engine: 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 @@ -321,8 +335,12 @@ class Engine: 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 @@ -332,10 +350,20 @@ class Engine: 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 @@ -362,8 +390,64 @@ class Engine: 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.""" @@ -510,7 +594,7 @@ class Engine: job.total_steps = gen.num_inference_steps job.emit( "start", - model=self.holder_model, + 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, diff --git a/toolkit/inference_lora.py b/toolkit/inference_lora.py new file mode 100644 index 0000000..9abc324 --- /dev/null +++ b/toolkit/inference_lora.py @@ -0,0 +1,342 @@ +"""Inference-time LoRA loading for the engine: any common LoRA file format -> +per-module low-rank deltas, applied either as forward hooks (dynamic strength, +removable) or merged into the weights (quantized weights re-quantized with +stochastic rounding so small deltas survive). + +Independent of the training network classes on purpose. + +Supported key formats (auto-detected per key): + peft / diffusers transformer.blocks.0.attn.to_q.lora_A.weight / lora_B + comfy diffusion_model.<...>.lora_A|lora_down (+ .alpha) + kohya / lycoris lora_unet_.lora_down.weight / lora_up + alpha, + lora_te_ / lora_te1_ / lora_te2_ for text encoders + full diffs .diff (weight delta), .diff_b (bias delta) +The holder's convert_lora_weights_before_load hook runs first, so archs with +their own conversions (anima, ltx2, hidream_o1) keep working. +""" + +import os +from typing import Dict, List, Optional, Tuple + +import torch +from safetensors.torch import load_file + +_A_SUFFIXES = (".lora_A.weight", ".lora_down.weight", ".lora_A.default.weight") +_B_SUFFIXES = (".lora_B.weight", ".lora_up.weight", ".lora_B.default.weight") +_TE_PREFIXES = { + "lora_te1_": 0, "lora_te_": 0, "lora_te2_": 1, "lora_te3_": 2, + "text_encoder.": 0, "text_encoder_2.": 1, "text_encoder_3.": 2, + "te.": 0, "te1.": 0, "te2.": 1, +} +_DIT_PREFIXES = ( + "model.diffusion_model.", "diffusion_model.", "transformer.", "unet.", + "lora_unet_", "lora_transformer_", "base_model.model.", +) + + +class LoRAEntry: + def __init__(self, module: torch.nn.Module, name: str): + self.module = module + self.name = name + self.A: Optional[torch.Tensor] = None # (r, in) + self.B: Optional[torch.Tensor] = None # (out, r) + self.alpha: Optional[float] = None + self.diff: Optional[torch.Tensor] = None + self.diff_b: Optional[torch.Tensor] = None + + @property + def scale(self) -> float: + if self.A is None or self.alpha is None: + return 1.0 + r = self.A.shape[0] + return float(self.alpha) / r if r else 1.0 + + def delta(self, strength: float, device, dtype=torch.float32) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """(weight delta, bias delta) at this strength, full precision.""" + dw = None + if self.A is not None and self.B is not None: + A = self.A.to(device=device, dtype=dtype) + B = self.B.to(device=device, dtype=dtype) + if A.dim() == 4: # conv lora: (r, in, kh, kw) @ (out, r, 1, 1) + dw = torch.einsum("or,rikl->oikl", B.flatten(1), A) * (self.scale * strength) + else: + dw = (B @ A) * (self.scale * strength) + if self.diff is not None: + d = self.diff.to(device=device, dtype=dtype) * strength + dw = d if dw is None else dw + d + db = self.diff_b.to(device=device, dtype=dtype) * strength if self.diff_b is not None else None + return dw, db + + +class InferenceLoRA: + def __init__(self, path: str, strength: float = 1.0, name: Optional[str] = None): + self.path = path + self.strength = float(strength) + self.name = name or os.path.splitext(os.path.basename(path))[0] + self.entries: List[LoRAEntry] = [] + self.unmatched: List[str] = [] + self._hooks: List = [] + + # ---- loading / key resolution ---- + def load(self, holder) -> "InferenceLoRA": + if not os.path.isfile(self.path): + raise FileNotFoundError(f"LoRA not found: {self.path}") + sd = load_file(self.path) + convert = getattr(holder, "convert_lora_weights_before_load", None) + if callable(convert): + try: + sd = convert(sd) + except Exception: + pass + roots = self._roots(holder) + grouped: Dict[Tuple[int, str], dict] = {} + for key, tensor in sd.items(): + base, part = self._split(key) + if part is None: + continue + grouped.setdefault(base, {})[part] = tensor + self.entries = [] + self.unmatched = [] + for base, parts in grouped.items(): + module, name = self._resolve(base, roots) + if module is None: + self.unmatched.append(base) + continue + e = LoRAEntry(module, name) + e.A = parts.get("A") + e.B = parts.get("B") + e.diff = parts.get("diff") + e.diff_b = parts.get("diff_b") + alpha = parts.get("alpha") + if alpha is not None: + e.alpha = float(alpha.flatten()[0]) + if e.A is not None and e.B is not None and e.A.dim() == 2: + w_shape = getattr(getattr(module, "weight", None), "shape", None) + if w_shape is not None and (e.B.shape[0] != w_shape[0] or e.A.shape[1] != w_shape[1]): + self.unmatched.append(f"{base} (shape {tuple(e.B.shape[0:1]) + tuple(e.A.shape[1:])} vs {tuple(w_shape)})") + continue + if e.A is None and e.diff is None and e.diff_b is None: + continue + self.entries.append(e) + return self + + @staticmethod + def _split(key: str): + for suf in _A_SUFFIXES: + if key.endswith(suf): + return key[: -len(suf)], "A" + for suf in _B_SUFFIXES: + if key.endswith(suf): + return key[: -len(suf)], "B" + if key.endswith(".alpha"): + return key[: -len(".alpha")], "alpha" + if key.endswith(".diff"): + return key[: -len(".diff")], "diff" + if key.endswith(".diff_b"): + return key[: -len(".diff_b")], "diff_b" + return key, None + + @staticmethod + def _roots(holder) -> List[Tuple[str, torch.nn.Module, dict, dict]]: + """(kind, root module, dotted-name map, flat-name map) for the DiT and TEs.""" + roots = [] + dit = getattr(holder, "model", None) + if isinstance(dit, torch.nn.Module): + names = dict(dit.named_modules()) + roots.append(("dit", dit, names, {n.replace(".", "_"): n for n in names})) + tes = getattr(holder, "text_encoder", None) + tes = tes if isinstance(tes, list) else [tes] + for i, te in enumerate(tes): + if isinstance(te, torch.nn.Module) and not type(te).__name__.startswith("Fake"): + names = dict(te.named_modules()) + roots.append((f"te{i}", te, names, {n.replace(".", "_"): n for n in names})) + return roots + + @classmethod + def _resolve(cls, base: str, roots): + te_index = None + rest = base + for pref, idx in _TE_PREFIXES.items(): + if base.startswith(pref): + te_index = idx + rest = base[len(pref):] + break + if te_index is None: + for pref in _DIT_PREFIXES: + if base.startswith(pref): + rest = base[len(pref):] + break + candidates = [r for r in roots if (r[0] == f"te{te_index}" if te_index is not None else r[0] == "dit")] + if not candidates: + candidates = roots + for _, root, names, flat in candidates: + if rest in names: + return names[rest], rest + flat_key = rest.replace(".", "_") + if flat_key in flat: + n = flat[flat_key] + return names[n], n + return None, None + + # ---- hook mode ---- + def attach(self): + """Forward hooks that add strength * delta(x) to each module's output.""" + self.detach() + for e in self.entries: + if e.A is None or e.A.dim() != 2: + if e.diff is None and e.diff_b is None: + continue + lora = self + + def make(entry: LoRAEntry): + cache = {} + + def hook(module, inputs, output): + if lora.strength == 0: + return output + x = inputs[0] if inputs else None + if x is None: + return output + key = (x.device, x.dtype) + if key not in cache: + cache.clear() + cache[key] = ( + entry.A.to(device=x.device, dtype=x.dtype) if entry.A is not None and entry.A.dim() == 2 else None, + entry.B.to(device=x.device, dtype=x.dtype) if entry.B is not None and entry.B.dim() == 2 else None, + entry.diff.to(device=x.device, dtype=x.dtype) if entry.diff is not None and entry.diff.dim() == 2 else None, + entry.diff_b.to(device=x.device, dtype=x.dtype) if entry.diff_b is not None else None, + ) + A, B, diff, diff_b = cache[key] + add = None + if A is not None and B is not None: + add = ((x @ A.t()) @ B.t()) * (entry.scale * lora.strength) + if diff is not None: + d = (x @ diff.t()) * lora.strength + add = d if add is None else add + d + if diff_b is not None: + add = (diff_b * lora.strength) if add is None else add + diff_b * lora.strength + if add is None: + return output + if isinstance(output, tuple): + return (output[0] + add.to(output[0].dtype),) + tuple(output[1:]) + return output + add.to(output.dtype) + + return hook + + self._hooks.append(e.module.register_forward_hook(make(e))) + return self + + def detach(self): + for h in self._hooks: + h.remove() + self._hooks = [] + + # ---- merge mode ---- + @torch.no_grad() + def merge(self, stochastic: bool = True) -> int: + """Add strength * delta into each module's weight. Quantized (Ostris) + weights: dequantize -> add -> requantize on the stored grid with + stochastic rounding (int8 backends), nearest otherwise. Plain weights + in bf16/fp16 on cuda: stochastic rounding via copy_stochastic.""" + from toolkit.optimizers.optimizer_utils import copy_stochastic + + merged = 0 + for e in self.entries: + m = e.module + w = getattr(m, "weight", None) + if w is None: + continue + device = w.device + dw, db = e.delta(self.strength, device) + if dw is not None: + if getattr(m, "is_ostris_quantized", False): + fp = m.dequantize_weight().to(torch.float32) + dw.reshape(w.shape) + q = m.ostris_quantizer + rq_codes = getattr(q, "requantize_codes_", None) + if rq_codes is not None and stochastic: + try: + rq_codes(m, fp, stochastic=True) + except TypeError: + rq_codes(m, fp) + elif rq_codes is not None: + rq_codes(m, fp) + else: + m.requantize_(fp) + elif isinstance(w, torch.nn.Parameter) or isinstance(w, torch.Tensor): + fp = w.data.to(torch.float32) + dw.reshape(w.shape) + if stochastic and w.is_cuda and w.dtype in (torch.bfloat16, torch.float16): + copy_stochastic(w.data, fp) + else: + w.data.copy_(fp.to(w.dtype)) + else: + continue + if db is not None and getattr(m, "bias", None) is not None: + m.bias.data.add_(db.to(m.bias.dtype)) + merged += 1 + m_marker = getattr(self, "_marker", None) + return merged + + def summary(self) -> dict: + return { + "name": self.name, + "path": self.path, + "strength": self.strength, + "modules": len(self.entries), + "unmatched": len(self.unmatched), + } + + +class LoRAStack: + """The LoRAs applied to one holder, in one mode.""" + + def __init__(self, holder, mode: str = "hook"): + self.holder = holder + self.mode = mode + self.loras: List[InferenceLoRA] = [] + + def key(self) -> str: + return "|".join(f"{l.path}@{l.strength:.4f}" for l in self.loras) + f"#{self.mode}" + + def load(self, specs: List[dict], status_fn=None): + for spec in specs: + lora = InferenceLoRA(spec["path"], spec.get("strength", 1.0), spec.get("name")).load(self.holder) + if status_fn: + status_fn( + f"LoRA {lora.name}: {len(lora.entries)} modules" + + (f", {len(lora.unmatched)} unmatched" if lora.unmatched else "") + ) + if not lora.entries: + raise ValueError(f"LoRA {lora.name} matched no modules of this model (unmatched: {lora.unmatched[:3]})") + self.loras.append(lora) + return self + + def apply(self, status_fn=None): + touched = set() + for lora in self.loras: + if self.mode == "merge": + n = lora.merge(stochastic=True) + if status_fn: + status_fn(f"Merged LoRA {lora.name} into {n} modules (stochastic rounding)") + else: + lora.attach() + for e in lora.entries: + touched.add(id(e.module)) + return touched + + def remove(self): + """Hook mode: detach every hook. Merge mode: cannot un-merge; the + engine reloads the affected components instead.""" + for lora in self.loras: + lora.detach() + self.loras = [] + + def set_strengths(self, specs: List[dict]) -> bool: + """Hook mode only: update strengths in place for the same paths.""" + if self.mode != "hook" or len(specs) != len(self.loras): + return False + for lora, spec in zip(self.loras, specs): + if spec["path"] != lora.path: + return False + for lora, spec in zip(self.loras, specs): + lora.strength = float(spec.get("strength", 1.0)) + return True diff --git a/toolkit/models/v2/pool.py b/toolkit/models/v2/pool.py index f31379f..831f3b4 100644 --- a/toolkit/models/v2/pool.py +++ b/toolkit/models/v2/pool.py @@ -160,6 +160,16 @@ class ComponentPool: self._free(entry) return freed + def evict_where(self, predicate) -> int: + """Free every entry whose module satisfies predicate(module).""" + with self._lock: + victims = [e for e in self.entries.values() if predicate(e.module)] + freed = 0 + for entry in victims: + freed += entry.bytes + self._free(entry) + return freed + def clear(self): with self._lock: for entry in list(self.entries.values()): diff --git a/toolkit/util/convrot_quant.py b/toolkit/util/convrot_quant.py index 6fc5b76..231d24e 100644 --- a/toolkit/util/convrot_quant.py +++ b/toolkit/util/convrot_quant.py @@ -1419,13 +1419,20 @@ class ConvRotInt8Quantizer(OstrisQuantizer): module.cr8_scales = scales.view(torch.uint8) @torch.no_grad() - def requantize_codes_(self, module, fp_weight: torch.Tensor) -> None: + def requantize_codes_(self, module, fp_weight: torch.Tensor, stochastic: bool = False) -> None: """Re-quantize only the codes on the module's STORED scales — the grid a - QAT run trains against (see ConvRotQuantizer.requantize_codes_).""" + QAT run trains against (see ConvRotQuantizer.requantize_codes_). + stochastic: floor(z + u) instead of round(z) — unbiased, so a small + delta (a merged LoRA) survives instead of rounding away.""" w = fp_weight.to(device=module.cr8_qdata.device, dtype=torch.float32) w_rot = rotate(w, self._rot(module)) s = self._scales(module).unsqueeze(1) - module.cr8_qdata = torch.round(w_rot / s).clamp_(-127, 127).to(torch.int8) + z = w_rot / s + if stochastic: + z = torch.floor(z + torch.rand_like(z)) + else: + z = torch.round(z) + module.cr8_qdata = z.clamp_(-127, 127).to(torch.int8) def _gemv_args(self, module): """(qdata, gratio fp32 or None, bits) for the fused decode gemv, or None diff --git a/ui/src/app/api/loras/route.ts b/ui/src/app/api/loras/route.ts new file mode 100644 index 0000000..32747ef --- /dev/null +++ b/ui/src/app/api/loras/route.ts @@ -0,0 +1,78 @@ +import { NextResponse } from 'next/server'; +import fs from 'fs'; +import path from 'path'; +import prisma from '@/server/prisma'; +import { getModelsPath, getTrainingFolder } from '@/server/settings'; + +export const dynamic = 'force-dynamic'; + +interface LoraFile { + name: string; + path: string; + size: number; + mtime: number; +} + +const listSafetensors = async (dir: string): Promise => { + try { + const entries = await fs.promises.readdir(dir, { withFileTypes: true }); + const files = entries.filter(e => e.isFile() && e.name.endsWith('.safetensors')); + return ( + await Promise.all( + files.map(async e => { + const p = path.join(dir, e.name); + const st = await fs.promises.stat(p); + return { name: e.name, path: p, size: st.size, mtime: st.mtimeMs }; + }), + ) + ).sort((a, b) => b.mtime - a.mtime); + } catch { + return []; + } +}; + +const walk = async (root: string, rel = '', depth = 0, out: (LoraFile & { relpath: string })[] = []) => { + if (depth > 5) return out; + let entries: fs.Dirent[]; + try { + entries = await fs.promises.readdir(path.join(root, rel), { withFileTypes: true }); + } catch { + return out; + } + for (const e of entries) { + if (e.name.startsWith('.')) continue; + const r = rel ? path.join(rel, e.name) : e.name; + if (e.isDirectory()) { + await walk(root, r, depth + 1, out); + } else if (e.isFile() && e.name.endsWith('.safetensors')) { + const p = path.join(root, r); + const st = await fs.promises.stat(p); + out.push({ name: e.name, relpath: r, path: p, size: st.size, mtime: st.mtimeMs }); + } + } + return out; +}; + +// LoRA files available to the Generate page: every training job's saved +// checkpoints, and everything under /loras. +export async function GET() { + const trainingFolder = await getTrainingFolder(); + const jobs = await prisma.job.findMany({ where: { job_type: 'train' }, orderBy: { updated_at: 'desc' } }); + const jobEntries = ( + await Promise.all( + jobs.map(async job => ({ + id: job.id, + name: job.name, + status: job.status, + updated_at: job.updated_at, + files: await listSafetensors(path.join(trainingFolder, job.name)), + })), + ) + ).filter(j => j.files.length > 0); + + const modelsPath = await getModelsPath(); + const lorasRoot = path.join(modelsPath, 'loras'); + const models = (await walk(lorasRoot)).sort((a, b) => a.relpath.localeCompare(b.relpath)); + + return NextResponse.json({ jobs: jobEntries, models, lorasRoot }); +} diff --git a/ui/src/app/generate/page.tsx b/ui/src/app/generate/page.tsx index a5a3159..6c08c63 100644 --- a/ui/src/app/generate/page.tsx +++ b/ui/src/app/generate/page.tsx @@ -3,7 +3,7 @@ import React, { Suspense, useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { useSearchParams } from 'next/navigation'; import { Button } from '@headlessui/react'; -import { ChevronDown, ChevronLeft, ChevronRight, Loader2, OctagonX, Play, Square, Sparkles, Trash2, X } from 'lucide-react'; +import { ChevronDown, ChevronLeft, ChevronRight, Loader2, OctagonX, Play, Plus, Square, Sparkles, Trash2, X } from 'lucide-react'; import { openConfirm } from '@/components/ConfirmModal'; import { TopBar, MainContent } from '@/components/layout'; import { Checkbox, CreatableSelectInput, NumberInput, SelectInput, SliderInput, TextAreaInput, TextInput } from '@/components/formInputs'; @@ -18,6 +18,7 @@ import { latentToImage, payloadToFloat32, readEngineFrames, PreviewInfo } from ' import { isMac } from '@/helpers/basic'; import { modelArchs, getGenerateDefaults, GenerateDefaults } from '@/app/jobs/new/options'; import GenerateFooter from '@/components/generate/GenerateFooter'; +import LoraBrowserModal, { LoraPick } from '@/components/generate/LoraBrowserModal'; interface EngineStatus { @@ -298,6 +299,13 @@ function GeneratePageInner() { const [showPreview, setShowPreview] = useState(!persisted.results?.length); const [sidebarOpen, setSidebarOpen] = useState(persisted.sidebarOpen ?? true); const [cards, setCards] = useState<{ [key: string]: boolean }>(persisted.cards || {}); + const [loraModalOpen, setLoraModalOpen] = useState(false); + const loras: { path: string; name: string; strength: number }[] = model.loras || []; + const setLoras = (next: { path: string; name: string; strength: number }[]) => setModel(m => ({ ...m, loras: next })); + const addLora = (pick: LoraPick) => { + if (loras.some(l => l.path === pick.path)) return; + setLoras([...loras, { path: pick.path, name: pick.name, strength: 1.0 }]); + }; const cardOpen = (key: string) => cards[key] ?? true; const toggleCard = (key: string) => setCards(c => ({ ...c, [key]: !(c[key] ?? true) })); const resultUrl = (r: ResultItem) => `/api/files/${encodeFilePathForUrl(r.path)}`; @@ -934,6 +942,70 @@ function GeneratePageInner() { )} + toggleCard('loras')}> +
+ {loras.map((l, i) => ( +
+
+ + {l.name} + + { + const v = parseFloat(e.target.value); + const next = [...loras]; + next[i] = { ...l, strength: isNaN(v) ? 0 : v }; + setLoras(next); + }} + className="w-16 text-xs px-1.5 py-0.5 bg-gray-950 border border-gray-700 rounded text-gray-100 text-right" + /> + +
+
+ ))} +
+ +
+ setModel(m => ({ ...m, lora_mode: v ? 'merge' : 'hook' }))} + doc={{ + title: 'Merge LoRAs into weights', + description: ( + <> +

+ Off (default): each LoRA is applied dynamically as an extra term on the layers it targets. Strength + changes take effect on the next generation and adding or removing a LoRA never reloads the model. +

+

+ On: the LoRA deltas are added into the model weights. Quantized weights are dequantized, merged, and + re-quantized with stochastic rounding so small deltas are not rounded away. Merged weights cannot be un-merged, so + changing the LoRA set or strengths reloads the affected components. Slightly faster per step than the dynamic path. +

+ + ), + }} + /> +
+ toggleCard('prompt')}> setSample(s => ({ ...s, prompt: v }))} placeholder={modality === 'audio' ? 'song description, lyrics, bpm…' : 'describe what to generate'} /> {(sample.guidance_scale ?? 4) > 1 && ( @@ -976,6 +1048,7 @@ function GeneratePageInner() { + setLoraModalOpen(false)} onPick={addLora} /> ); } diff --git a/ui/src/app/globals.css b/ui/src/app/globals.css index 44b07cd..54d4582 100644 --- a/ui/src/app/globals.css +++ b/ui/src/app/globals.css @@ -147,7 +147,34 @@ body { .compact-form textarea { min-height: 3.5rem; } -.compact-form [class*='control'] { - min-height: 30px !important; +.compact-form .aitk-react-select-container .aitk-react-select__control { + height: 26px !important; + min-height: 26px !important; font-size: 12px !important; + line-height: 1.2 !important; +} +.compact-form .aitk-react-select-container .aitk-react-select__value-container { + padding: 0 6px !important; +} +.compact-form .aitk-react-select-container .aitk-react-select__single-value, +.compact-form .aitk-react-select-container .aitk-react-select__placeholder, +.compact-form .aitk-react-select-container .aitk-react-select__input-container, +.compact-form .aitk-react-select-container .aitk-react-select__input, +.compact-form .aitk-react-select-container .aitk-react-select__option { + font-size: 12px !important; + line-height: 1.2 !important; +} +.compact-form .aitk-react-select-container .aitk-react-select__input-container { + margin: 0 !important; + padding: 0 !important; +} +.compact-form .aitk-react-select-container .aitk-react-select__option { + padding: 4px 8px !important; +} +.compact-form .aitk-react-select-container .aitk-react-select__indicator { + padding: 2px !important; +} +.compact-form .aitk-react-select-container .aitk-react-select__indicator svg { + width: 14px; + height: 14px; } diff --git a/ui/src/components/generate/LoraBrowserModal.tsx b/ui/src/components/generate/LoraBrowserModal.tsx new file mode 100644 index 0000000..118bde8 --- /dev/null +++ b/ui/src/components/generate/LoraBrowserModal.tsx @@ -0,0 +1,175 @@ +'use client'; + +import { useEffect, useMemo, useState } from 'react'; +import { ChevronDown, ChevronRight, FolderOpen, Loader2, Search } from 'lucide-react'; +import { Modal } from '@/components/Modal'; +import { apiClient } from '@/utils/api'; + +export interface LoraPick { + path: string; + name: string; +} + +interface LoraFile { + name: string; + path: string; + size: number; + mtime: number; + relpath?: string; +} + +interface JobEntry { + id: string; + name: string; + status: string; + updated_at: string; + files: LoraFile[]; +} + +interface Props { + isOpen: boolean; + onClose: () => void; + onPick: (lora: LoraPick) => void; +} + +const fmtSize = (n: number) => (n > 1e9 ? `${(n / 1e9).toFixed(2)} GB` : `${(n / 1e6).toFixed(0)} MB`); +const fmtDate = (ms: number) => new Date(ms).toLocaleString(); + +export default function LoraBrowserModal({ isOpen, onClose, onPick }: Props) { + const [tab, setTab] = useState<'jobs' | 'models'>('jobs'); + const [loading, setLoading] = useState(false); + const [jobs, setJobs] = useState([]); + const [models, setModels] = useState([]); + const [lorasRoot, setLorasRoot] = useState(''); + const [filter, setFilter] = useState(''); + const [open, setOpen] = useState>(new Set()); + + useEffect(() => { + if (!isOpen) return; + setLoading(true); + apiClient + .get('/api/loras') + .then(r => { + setJobs(r.data.jobs || []); + setModels(r.data.models || []); + setLorasRoot(r.data.lorasRoot || ''); + }) + .catch(() => {}) + .finally(() => setLoading(false)); + }, [isOpen]); + + const q = filter.trim().toLowerCase(); + const filteredJobs = useMemo( + () => (q ? jobs.filter(j => j.name.toLowerCase().includes(q) || j.files.some(f => f.name.toLowerCase().includes(q))) : jobs), + [jobs, q], + ); + const filteredModels = useMemo(() => (q ? models.filter(m => (m.relpath || m.name).toLowerCase().includes(q)) : models), [models, q]); + + const pick = (f: LoraFile, jobName?: string) => { + onPick({ path: f.path, name: jobName ? `${jobName} / ${f.name.replace(/\.safetensors$/, '')}` : (f.relpath || f.name).replace(/\.safetensors$/, '') }); + onClose(); + }; + + const tabClass = (t: 'jobs' | 'models') => + `px-3 py-1.5 text-sm border-b-2 ${tab === t ? 'border-blue-500 text-blue-400' : 'border-transparent text-gray-400 hover:text-gray-200'}`; + + return ( + +
+
+ + +
+
+ + setFilter(e.target.value)} + placeholder="filter" + className="pl-7 pr-2 py-1 text-xs bg-gray-950 border border-gray-700 rounded text-gray-100 w-48" + /> +
+
+
+ {loading && ( +
+ Loading… +
+ )} + {!loading && tab === 'jobs' && ( +
+ {filteredJobs.length === 0 &&
No training jobs with saved LoRAs.
} + {filteredJobs.map(j => { + const isOpen = open.has(j.id) || !!q; + return ( +
+ + {isOpen && ( +
+ {j.files + .filter(f => !q || f.name.toLowerCase().includes(q) || j.name.toLowerCase().includes(q)) + .map(f => ( + + ))} +
+ )} +
+ ); + })} +
+ )} + {!loading && tab === 'models' && ( +
+
+ {lorasRoot || 'MODELS_PATH/loras'} +
+ {filteredModels.length === 0 &&
No .safetensors files found under the loras folder.
} +
+ {filteredModels.map(f => ( + + ))} +
+
+ )} +
+
+ + ); +} diff --git a/ui/src/server/settings.ts b/ui/src/server/settings.ts index 469d0a8..aa3a439 100644 --- a/ui/src/server/settings.ts +++ b/ui/src/server/settings.ts @@ -90,3 +90,11 @@ export const getDataRoot = async () => { myCache.set(key, dataRoot); return dataRoot; }; + +// MODELS_PATH: env wins, then the setting, then /models +export const getModelsPath = async () => { + if (process.env.MODELS_PATH && process.env.MODELS_PATH.trim() !== '') return process.env.MODELS_PATH; + const row = await prisma.settings.findFirst({ where: { key: 'MODELS_PATH' } }); + if (row?.value && row.value !== '') return row.value; + return path.join(process.cwd(), '..', 'models'); +};