Add LoRA stacks with per-item strengths and named video/image generation presets so setups survive reload without mixing Klein onto MiniMax. Keep LTX Pink Cherry and xAIGen LoRAs on xAIGen, and run Recommend on Beast.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-28 15:24:44 -05:00
parent 0af3cddb14
commit f8db0d626a
36 changed files with 1767 additions and 246 deletions
@@ -0,0 +1,305 @@
+import { comfyConfigured, comfyFetch, getComfyHost } from '~/server/utils/comfy'
+import { getBeastImageHost, imageComfyFetch, sameImageHost } from '~/server/utils/imageComfy'
+import {
+  XAIGEN_LORA_MESSAGE,
+  filterLorasForStudio,
+  isXaigenOnlyLora,
+  normalizeLoraName,
+  normalizeLoraStack,
+  persistLoraFields,
+  type LoraKind,
+  type LoraStackItem
+} from '~/utils/loras'
+import { isXaigenStudio } from '~/utils/videoModels'
+
+type WorkflowNode = { class_type: string; inputs: Record<string, unknown>; _meta?: { title?: string } }
+type WorkflowGraph = Record<string, WorkflowNode>
+
+const LORA_LOADER = 'LoraLoader'
+const LORA_MODEL_ONLY = 'LoraLoaderModelOnly'
+const OBJECT_INFO_NODES = [
+  LORA_LOADER,
+  LORA_MODEL_ONLY,
+  'Power Lora Loader (rgthree)',
+  'Lora Loader Stack (rgthree)',
+  'LoraLoaderModelOnly [10]',
+  'WanVideoLoraSelect',
+  'LoraLoaderStacked'
+]
+const MODEL_FEED_CLASSES = new Set(['BasicGuider', 'CFGGuider', 'KSampler', 'KSamplerAdvanced'])
+const CACHE_MS = 20_000
+
+type LoraCache = {
+  at: number
+  image: string[]
+  video: string[]
+}
+
+let cache: LoraCache | null = null
+let inflight: Promise<LoraCache> | null = null
+
+function uniqueNames(values: unknown[]) {
+  const seen = new Set<string>()
+  const out: string[] = []
+  for (const value of values) {
+    const name = String(value || '').trim()
+    if (!name || seen.has(name)) continue
+    seen.add(name)
+    out.push(name)
+  }
+  return out
+}
+
+function comboFromSpec(spec: unknown): string[] {
+  if (!Array.isArray(spec)) return []
+  const first = spec[0]
+  if (Array.isArray(first)) return first.filter((item): item is string => typeof item === 'string' && item.trim().length > 0)
+  if (typeof first === 'string' && first.trim()) return spec.filter((item): item is string => typeof item === 'string' && item.trim().length > 0)
+  return []
+}
+
+function namesFromObjectInfoNode(info: unknown) {
+  if (!info || typeof info !== 'object') return [] as string[]
+  const input = (info as { input?: { required?: Record<string, unknown>; optional?: Record<string, unknown> } }).input
+  const buckets = [input?.required, input?.optional]
+  const names: string[] = []
+  for (const bucket of buckets) {
+    if (!bucket) continue
+    for (const [key, spec] of Object.entries(bucket)) {
+      if (!/lora/i.test(key)) continue
+      names.push(...comboFromSpec(spec))
+    }
+  }
+  return uniqueNames(names)
+}
+
+async function fetchJson(path: string, via: 'video' | 'image') {
+  const res = via === 'image'
+    ? await imageComfyFetch(path, { signal: AbortSignal.timeout(8000) })
+    : await comfyFetch(path, { signal: AbortSignal.timeout(8000) })
+  if (!res.ok) return null
+  return res.json().catch(() => null)
+}
+
+async function fetchModelsLoras(via: 'video' | 'image') {
+  const payload = await fetchJson('/models/loras', via)
+  if (Array.isArray(payload)) return uniqueNames(payload)
+  if (payload && typeof payload === 'object' && Array.isArray((payload as { loras?: unknown[] }).loras)) {
+    return uniqueNames((payload as { loras: unknown[] }).loras)
+  }
+  return [] as string[]
+}
+
+async function fetchObjectInfoLoras(via: 'video' | 'image') {
+  const image: string[] = []
+  const video: string[] = []
+  let any = false
+  for (const node of OBJECT_INFO_NODES) {
+    const info = await fetchJson(`/object_info/${encodeURIComponent(node)}`, via)
+    if (!info || typeof info !== 'object') continue
+    const record = info as Record<string, unknown>
+    const body = record[node] || (Object.keys(record).length === 1 ? record[Object.keys(record)[0] as string] : record)
+    const names = namesFromObjectInfoNode(body)
+    if (!names.length) continue
+    any = true
+    if (node === LORA_LOADER) image.push(...names)
+    else video.push(...names)
+    if (node !== LORA_LOADER && node !== LORA_MODEL_ONLY) {
+      image.push(...names)
+      video.push(...names)
+    }
+  }
+  if (any) {
+    return { image: uniqueNames(image), video: uniqueNames(video) }
+  }
+
+  const all = await fetchJson('/object_info', via)
+  if (!all || typeof all !== 'object') return { image: [] as string[], video: [] as string[] }
+  const imageAll: string[] = []
+  const videoAll: string[] = []
+  for (const [classType, info] of Object.entries(all as Record<string, unknown>)) {
+    if (!/lora/i.test(classType)) continue
+    const names = namesFromObjectInfoNode(info)
+    if (classType === LORA_LOADER) imageAll.push(...names)
+    else videoAll.push(...names)
+  }
+  return { image: uniqueNames(imageAll), video: uniqueNames(videoAll) }
+}
+
+async function discoverFromHost(via: 'video' | 'image') {
+  const models = await fetchModelsLoras(via).catch(() => [] as string[])
+  const fromInfo = await fetchObjectInfoLoras(via).catch(() => ({ image: [] as string[], video: [] as string[] }))
+  const image = uniqueNames([...fromInfo.image, ...models])
+  const video = uniqueNames([...fromInfo.video, ...models])
+  if (!image.length && video.length) return { image: video, video }
+  if (!video.length && image.length) return { image, video: image }
+  return { image, video }
+}
+
+async function loadLoraCache(): Promise<LoraCache> {
+  const now = Date.now()
+  if (cache && now - cache.at < CACHE_MS) return cache
+  if (inflight) return inflight
+  inflight = (async () => {
+    const videoHost = comfyConfigured() ? getComfyHost() : ''
+    const imageHost = getBeastImageHost()
+    const same = Boolean(videoHost && imageHost && sameImageHost(videoHost, imageHost))
+    const video = videoHost
+      ? await discoverFromHost('video').catch(() => ({ image: [] as string[], video: [] as string[] }))
+      : { image: [] as string[], video: [] as string[] }
+    const image = imageHost && !same
+      ? await discoverFromHost('image').catch(() => ({ image: [] as string[], video: [] as string[] }))
+      : video
+    const xaigen = isXaigenStudio()
+    const next: LoraCache = {
+      at: Date.now(),
+      image: filterLorasForStudio(uniqueNames([...image.image, ...video.image]), xaigen),
+      video: filterLorasForStudio(uniqueNames([...video.video, ...image.video]), xaigen)
+    }
+    cache = next
+    return next
+  })().finally(() => {
+    inflight = null
+  })
+  return inflight
+}
+
+export async function listStudioLoras() {
+  try {
+    return await loadLoraCache()
+  } catch {
+    return { at: 0, image: [] as string[], video: [] as string[] }
+  }
+}
+
+export function assertLoraAllowed(raw: unknown, kind: LoraKind) {
+  const name = normalizeLoraName(raw)
+  if (!name) return ''
+  if (isXaigenOnlyLora(name) && !isXaigenStudio()) {
+    throw createError({ statusCode: 400, statusMessage: XAIGEN_LORA_MESSAGE })
+  }
+  const known = cache?.[kind]
+  if (known?.length && !known.includes(name) && !known.some(item => item.replace(/\\/g, '/') === name)) {
+    throw createError({ statusCode: 400, statusMessage: `Unknown ${kind} LoRA` })
+  }
+  return name
+}
+
+export function parsePostedLora(raw: unknown, kind: LoraKind) {
+  return parsePostedLoraStack(raw, kind)[0]?.name || ''
+}
+
+export function parsePostedLoraStack(raw: unknown, kind: LoraKind): LoraStackItem[] {
+  const out: LoraStackItem[] = []
+  for (const item of normalizeLoraStack(raw)) {
+    const name = assertLoraAllowed(item.name, kind)
+    if (!name) continue
+    out.push({ ...item, name })
+  }
+  return out
+}
+
+export function parseShotLoras(raw: unknown, shotCount: number, kind: LoraKind = 'video') {
+  return parseShotLoraStacks(raw, shotCount, kind).map(stack => stack[0]?.name || '')
+}
+
+export function parseShotLoraStacks(raw: unknown, shotCount: number, kind: LoraKind = 'video'): LoraStackItem[][] {
+  const empty = Array.from({ length: shotCount }, () => [] as LoraStackItem[])
+  if (!raw) return empty
+  let parsed: unknown = raw
+  if (typeof raw === 'string') {
+    try {
+      parsed = JSON.parse(raw)
+    } catch {
+      return empty
+    }
+  }
+  if (!Array.isArray(parsed)) return empty
+  return Array.from({ length: shotCount }, (_, index) => parsePostedLoraStack(parsed[index], kind))
+}
+
+export { persistLoraFields }
+
+function linkSource(value: unknown): string | null {
+  return Array.isArray(value) && typeof value[0] === 'string' ? value[0] : null
+}
+
+function alreadyHasLora(graph: WorkflowGraph, name: string) {
+  const wanted = name.replace(/\\/g, '/').toLowerCase()
+  return Object.values(graph).some((node) => {
+    if (node.class_type !== LORA_LOADER && node.class_type !== LORA_MODEL_ONLY) return false
+    return String(node.inputs.lora_name || '').replace(/\\/g, '/').toLowerCase() === wanted
+  })
+}
+
+function injectAfter(
+  graph: WorkflowGraph,
+  sourceId: string,
+  nodeId: string,
+  node: WorkflowNode
+) {
+  if (!graph[sourceId] || graph[nodeId]) return
+  graph[nodeId] = node
+  for (const [id, other] of Object.entries(graph)) {
+    if (id === nodeId) continue
+    for (const [key, value] of Object.entries(other.inputs)) {
+      const src = linkSource(value)
+      if (src === sourceId && Array.isArray(value)) {
+        other.inputs[key] = [nodeId, value[1]]
+      }
+    }
+  }
+}
+
+function findModelFeed(graph: WorkflowGraph) {
+  for (const node of Object.values(graph)) {
+    if (!MODEL_FEED_CLASSES.has(node.class_type)) continue
+    const source = linkSource(node.inputs.model)
+    if (source && graph[source]) return source
+  }
+  return ''
+}
+
+export function applyUserLoraToGraph(graph: WorkflowGraph, stack?: unknown) {
+  const items = normalizeLoraStack(stack).filter(item => !alreadyHasLora(graph, item.name))
+  if (!items.length) return
+
+  const clipLoaders = Object.entries(graph).filter(([, node]) => node.class_type === LORA_LOADER)
+  if (clipLoaders.length) {
+    for (const [id] of clipLoaders) {
+      let sourceId = id
+      for (const [index, item] of items.entries()) {
+        const nodeId = `user:lora:${id}:${index}`
+        injectAfter(graph, sourceId, nodeId, {
+          class_type: LORA_LOADER,
+          inputs: {
+            lora_name: item.name,
+            strength_model: item.strengthModel,
+            strength_clip: item.strengthClip,
+            model: [sourceId, 0],
+            clip: [sourceId, 1]
+          },
+          _meta: { title: items.length === 1 ? 'User LoRA' : `User LoRA ${index + 1}` }
+        })
+        sourceId = nodeId
+      }
+    }
+    return
+  }
+
+  let source = findModelFeed(graph)
+  if (!source) return
+  for (const [index, item] of items.entries()) {
+    const nodeId = index === 0 ? 'user:lora' : `user:lora:${index}`
+    injectAfter(graph, source, nodeId, {
+      class_type: LORA_MODEL_ONLY,
+      inputs: {
+        lora_name: item.name,
+        strength_model: item.strengthModel,
+        model: [source, 0]
+      },
+      _meta: { title: items.length === 1 ? 'User LoRA' : `User LoRA ${index + 1}` }
+    })
+    source = nodeId
+  }
+}