aigen/server/utils/loras.ts

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
  }
}