aigen/utils/loras.ts

export const XAIGEN_LORA_MESSAGE = 'This LoRA is only available on xAIGen'

export type LoraKind = 'image' | 'video'

export type LoraStackItem = {
  name: string
  strengthModel: number
  strengthClip: number
}

export type LoraStack = LoraStackItem[]

export const LORA_STRENGTH_MIN = -2
export const LORA_STRENGTH_MAX = 2
export const LORA_STRENGTH_DEFAULT = 1
export const LORA_STRENGTH_STEP = 0.05
export const LORA_STACK_MAX = 8

/** Graph system LoRA for MiniMax turbo. Not a user pick. */
export const MINIMAX_TURBO_LORA = 'minimax_h3_fl2v_turbo_8step_v1.0_comfyui_bf16.safetensors'

export function loraBasename(name: string) {
  const trimmed = String(name || '').trim().replace(/\\/g, '/')
  const parts = trimmed.split('/').filter(Boolean)
  return parts[parts.length - 1] || trimmed
}

export function loraIdentityKey(name: string) {
  return loraBasename(name).replace(/^xaigen-?/i, '').toLowerCase()
}

export function isXaigenOnlyLora(name: string) {
  const base = loraBasename(name)
  return /^xaigen/i.test(base) || /klein_snofs|snofs_krea/i.test(base)
}

/** Pick the real Comfy filename: exact path, then same basename ignoring `xaigen-`, then suffix/contains. */
export function resolveComfyLoraName(wanted: string, available: string[]) {
  const target = String(wanted || '').trim().replace(/\\/g, '/')
  if (!target) return ''
  if (!available.length) return target
  const targetLower = target.toLowerCase()
  const exact = available.find(item => item.replace(/\\/g, '/') === target)
  if (exact) return exact
  const exactCi = available.find(item => item.replace(/\\/g, '/').toLowerCase() === targetLower)
  if (exactCi) return exactCi

  const identity = loraIdentityKey(target)
  const identityHit = available.find(item => loraIdentityKey(item) === identity)
  if (identityHit) return identityHit

  const stem = identity.replace(/\.(safetensors|sft|ckpt|pt|bin)$/i, '')
  if (!stem) return target
  const containsHits = available.filter((item) => {
    const base = loraBasename(item).toLowerCase()
    return base.includes(stem) || loraIdentityKey(item).includes(stem)
  })
  if (!containsHits.length) return target
  const suffix = containsHits.find((item) => {
    const base = loraBasename(item).toLowerCase()
    return base.endsWith(identity) || base.endsWith(stem) || loraIdentityKey(item).endsWith(identity)
  })
  return suffix || containsHits[0]
}

export function loraDisplayName(name: string) {
  return loraBasename(name).replace(/\.(safetensors|sft|ckpt|pt|bin)$/i, '')
}

export function normalizeLoraName(raw: unknown) {
  const value = String(raw || '').trim().replace(/\\/g, '/')
  if (!value) return ''
  if (value.length > 240) return ''
  if (value.includes('..') || value.startsWith('/') || /[\0\n\r]/.test(value)) return ''
  if (!/\.(safetensors|sft|ckpt|pt|bin)$/i.test(value)) return ''
  return value
}

export function clampLoraStrength(raw: unknown) {
  const value = Number(raw)
  if (!Number.isFinite(value)) return LORA_STRENGTH_DEFAULT
  const snapped = Math.round(value / LORA_STRENGTH_STEP) * LORA_STRENGTH_STEP
  return Math.min(LORA_STRENGTH_MAX, Math.max(LORA_STRENGTH_MIN, Math.round(snapped * 100) / 100))
}

export function formatLoraStrength(value: number) {
  const n = clampLoraStrength(value)
  return Number.isInteger(n) ? String(n) : n.toFixed(2).replace(/0$/, '')
}

function loraKey(name: string) {
  return name.replace(/\\/g, '/').toLowerCase()
}

export function makeLoraStackItem(
  name: string,
  strengthModel: unknown = LORA_STRENGTH_DEFAULT,
  strengthClip: unknown = LORA_STRENGTH_DEFAULT
): LoraStackItem | null {
  const normalized = normalizeLoraName(name)
  if (!normalized) return null
  return {
    name: normalized,
    strengthModel: clampLoraStrength(strengthModel),
    strengthClip: clampLoraStrength(strengthClip)
  }
}

function parseStackEntry(raw: unknown): LoraStackItem | null {
  if (typeof raw === 'string') return makeLoraStackItem(raw)
  if (!raw || typeof raw !== 'object') return null
  const rec = raw as Record<string, unknown>
  const name = rec.name ?? rec.loraName ?? rec.lora_name
  return makeLoraStackItem(
    name,
    rec.strengthModel ?? rec.strength_model ?? LORA_STRENGTH_DEFAULT,
    rec.strengthClip ?? rec.strength_clip ?? LORA_STRENGTH_DEFAULT
  )
}

export function normalizeLoraStack(raw: unknown): LoraStackItem[] {
  if (raw == null || raw === '') return []
  let parsed: unknown = raw
  if (typeof raw === 'string') {
    const trimmed = raw.trim()
    if (!trimmed) return []
    if (trimmed.startsWith('[') || trimmed.startsWith('{')) {
      try {
        parsed = JSON.parse(trimmed)
      } catch {
        const item = makeLoraStackItem(trimmed)
        return item ? [item] : []
      }
    } else {
      const item = makeLoraStackItem(trimmed)
      return item ? [item] : []
    }
  }
  const entries = Array.isArray(parsed) ? parsed : [parsed]
  const seen = new Set<string>()
  const out: LoraStackItem[] = []
  for (const entry of entries) {
    if (out.length >= LORA_STACK_MAX) break
    const item = parseStackEntry(entry)
    if (!item) continue
    const key = loraIdentityKey(item.name)
    if (seen.has(key)) continue
    seen.add(key)
    out.push(item)
  }
  return out
}

export function readLoraStack(value?: { loraStack?: unknown; loraName?: unknown } | string | unknown[] | null) {
  if (value == null || value === '') return [] as LoraStackItem[]
  if (typeof value === 'string' || Array.isArray(value)) return normalizeLoraStack(value)
  const stack = normalizeLoraStack(value.loraStack)
  if (stack.length) return stack
  return normalizeLoraStack(value.loraName)
}

export function resolveLoraName(fallback?: string | null, override?: string | null) {
  const over = normalizeLoraName(override)
  if (over) return over
  return normalizeLoraName(fallback)
}

export function resolveLoraStack(fallback?: unknown, override?: unknown) {
  const over = normalizeLoraStack(override)
  if (over.length) return over
  return normalizeLoraStack(fallback)
}

export function primaryLoraName(stack?: LoraStackItem[] | null) {
  return stack?.[0]?.name || ''
}

export function loraStackLabel(stack?: LoraStackItem[] | null) {
  return (stack || []).map(item => loraDisplayName(item.name)).join(' · ')
}

export function persistLoraFields(stack?: unknown) {
  const items = normalizeLoraStack(stack)
  if (!items.length) return { loraName: undefined as string | undefined, loraStack: undefined as LoraStackItem[] | undefined }
  return { loraName: items[0].name, loraStack: items }
}

export function loraStacksEqual(a?: unknown, b?: unknown) {
  const left = normalizeLoraStack(a)
  const right = normalizeLoraStack(b)
  if (left.length !== right.length) return false
  return left.every((item, index) => (
    loraKey(item.name) === loraKey(right[index].name)
    && item.strengthModel === right[index].strengthModel
    && item.strengthClip === right[index].strengthClip
  ))
}

/** User LoRA picker only. `xaigen-*` names are visible only on xAIGen. */
export function filterLorasForStudio(names: string[], xaigenStudio: boolean) {
  const seen = new Set<string>()
  const out: string[] = []
  for (const raw of names) {
    const name = String(raw || '').trim()
    if (!name || seen.has(name)) continue
    if (!xaigenStudio && isXaigenOnlyLora(name)) continue
    seen.add(name)
    out.push(name)
  }
  return out
}

export function filterLoraStackForStudio(stack: unknown, xaigenStudio: boolean) {
  return normalizeLoraStack(stack).filter(item => xaigenStudio || !isXaigenOnlyLora(item.name))
}

/** Klein / Flux.2 Klein adapters. Must not load on a Krea graph. */
export function isKleinImageLora(name: string) {
  const id = loraIdentityKey(name)
  return /flux-?2-klein|klein_snofs|flux2-klein-9b-consistency|(^|[^a-z0-9])klein([^a-z0-9]|$)/i.test(id)
}

export function filterLorasForImageEngine(names: string[], engine: 'flux' | 'krea') {
  if (engine !== 'krea') return names
  return names.filter(name => !isKleinImageLora(name))
}

export function filterLoraStackForImageEngine(stack: unknown, engine: 'flux' | 'krea') {
  const items = normalizeLoraStack(stack)
  if (engine !== 'krea') return items
  return items.filter(item => !isKleinImageLora(item.name))
}

export function kleinLoraOnKreaMessage(name: string) {
  return `${loraDisplayName(name)} is a Flux LoRA and cannot be used with Krea.`
}