Add lora on generate

This commit is contained in:
Jaret Burkett committed 2026-09-10 18:23:23 -06:00
1 parent 6e1b39c7a8
commit 732c761231
9 files changed
+812 -8

No files matched your search

+86 -2
View File
@@ -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,
+342
View File
@@ -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_<flat_name>.lora_down.weight / lora_up + alpha,
lora_te_ / lora_te1_ / lora_te2_ for text encoders
full diffs <module>.diff (weight delta), <module>.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
+10
View File
@@ -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()):
+10 -3
View File
@@ -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
+78
View File
@@ -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<LoraFile[]> => {
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 <MODELS_PATH>/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 });
}
+74 -1
View File
@@ -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<boolean>(!persisted.results?.length);
const [sidebarOpen, setSidebarOpen] = useState<boolean>(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() {
)}
</Card>
<Card title="LoRAs" subtitle={loras.length ? `${loras.length} loaded` : 'none'} open={cardOpen('loras')} onToggle={() => toggleCard('loras')}>
<div className="space-y-2">
{loras.map((l, i) => (
<div key={l.path} className="bg-gray-950/60 border border-gray-800 rounded-md px-2 py-1.5">
<div className="flex items-center gap-2">
<span className="text-xs text-gray-200 truncate flex-1" title={l.path}>
{l.name}
</span>
<input
type="number"
step={0.05}
min={-2}
max={3}
value={l.strength}
onChange={e => {
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"
/>
<button
type="button"
onClick={() => setLoras(loras.filter(x => x.path !== l.path))}
className="text-gray-500 hover:text-red-400"
title="Remove"
>
<X className="w-3.5 h-3.5" />
</button>
</div>
</div>
))}
<div className="flex items-center gap-2">
<Button
onClick={() => setLoraModalOpen(true)}
className="flex-1 px-2 py-1 rounded-md bg-gray-800 hover:bg-gray-700 text-gray-200 text-xs flex items-center justify-center gap-1"
>
<Plus className="w-3.5 h-3.5" /> Add LoRA
</Button>
</div>
<Checkbox
label="Merge into weights"
checked={model.lora_mode === 'merge'}
onChange={v => setModel(m => ({ ...m, lora_mode: v ? 'merge' : 'hook' }))}
doc={{
title: 'Merge LoRAs into weights',
description: (
<>
<p>
<strong>Off</strong> (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.
</p>
<p className="mt-2">
<strong>On</strong>: 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.
</p>
</>
),
}}
/>
</div>
</Card>
<Card title="Prompt" open={cardOpen('prompt')} onToggle={() => toggleCard('prompt')}>
<TextAreaInput label="Prompt" value={sample.prompt || ''} onChange={v => 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() {
</div>
</MainContent>
<GenerateFooter jobId={engineJobId} status={footerStatus} busy={running || isStarting} progress={running ? progress : null} />
<LoraBrowserModal isOpen={loraModalOpen} onClose={() => setLoraModalOpen(false)} onPick={addLora} />
</>
);
}
+29 -2
View File
@@ -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;
}
@@ -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<JobEntry[]>([]);
const [models, setModels] = useState<LoraFile[]>([]);
const [lorasRoot, setLorasRoot] = useState('');
const [filter, setFilter] = useState('');
const [open, setOpen] = useState<Set<string>>(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 (
<Modal isOpen={isOpen} onClose={onClose} title="Add LoRA" size="lg">
<div className="text-gray-200">
<div className="flex items-center gap-2 border-b border-gray-800 mb-3">
<button className={tabClass('jobs')} onClick={() => setTab('jobs')}>
Training jobs
</button>
<button className={tabClass('models')} onClick={() => setTab('models')}>
Models folder
</button>
<div className="flex-1" />
<div className="relative mb-1">
<Search className="w-3.5 h-3.5 absolute left-2 top-2 text-gray-500" />
<input
value={filter}
onChange={e => 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"
/>
</div>
</div>
<div className="max-h-[60vh] overflow-y-auto pr-1">
{loading && (
<div className="text-gray-500 text-sm flex items-center gap-2">
<Loader2 className="w-4 h-4 animate-spin" /> Loading…
</div>
)}
{!loading && tab === 'jobs' && (
<div className="space-y-1">
{filteredJobs.length === 0 && <div className="text-xs text-gray-500">No training jobs with saved LoRAs.</div>}
{filteredJobs.map(j => {
const isOpen = open.has(j.id) || !!q;
return (
<div key={j.id} className="rounded-md bg-gray-900 border border-gray-800">
<button
className="w-full flex items-center gap-2 px-3 py-2 text-sm text-left hover:bg-gray-800/60"
onClick={() =>
setOpen(prev => {
const n = new Set(prev);
if (n.has(j.id)) n.delete(j.id);
else n.add(j.id);
return n;
})
}
>
{isOpen ? <ChevronDown className="w-4 h-4 text-gray-500" /> : <ChevronRight className="w-4 h-4 text-gray-500" />}
<span className="truncate flex-1">{j.name}</span>
<span className="text-[11px] text-gray-500">
{j.files.length} file{j.files.length === 1 ? '' : 's'}
</span>
</button>
{isOpen && (
<div className="px-2 pb-2 space-y-0.5">
{j.files
.filter(f => !q || f.name.toLowerCase().includes(q) || j.name.toLowerCase().includes(q))
.map(f => (
<button
key={f.path}
onClick={() => pick(f, j.name)}
className="w-full flex items-center gap-2 px-2 py-1 rounded text-xs text-left hover:bg-blue-900/40"
title={f.path}
>
<span className="truncate flex-1 font-mono">{f.name}</span>
<span className="text-gray-500 shrink-0">{fmtSize(f.size)}</span>
<span className="text-gray-600 shrink-0 hidden sm:inline">{fmtDate(f.mtime)}</span>
</button>
))}
</div>
)}
</div>
);
})}
</div>
)}
{!loading && tab === 'models' && (
<div>
<div className="text-[11px] text-gray-500 mb-2 flex items-center gap-1">
<FolderOpen className="w-3.5 h-3.5" /> {lorasRoot || 'MODELS_PATH/loras'}
</div>
{filteredModels.length === 0 && <div className="text-xs text-gray-500">No .safetensors files found under the loras folder.</div>}
<div className="space-y-0.5">
{filteredModels.map(f => (
<button
key={f.path}
onClick={() => pick(f)}
className="w-full flex items-center gap-2 px-2 py-1 rounded text-xs text-left hover:bg-blue-900/40"
title={f.path}
>
<span className="truncate flex-1 font-mono">{f.relpath || f.name}</span>
<span className="text-gray-500 shrink-0">{fmtSize(f.size)}</span>
</button>
))}
</div>
</div>
)}
</div>
</div>
</Modal>
);
}
+8
View File
@@ -90,3 +90,11 @@ export const getDataRoot = async () => {
myCache.set(key, dataRoot);
return dataRoot;
};
// MODELS_PATH: env wins, then the setting, then <toolkit>/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');
};