aigen/server/utils/generationPresets.ts

import { existsSync, mkdirSync, readFileSync, renameSync, writeFileSync } from 'node:fs'
import { join } from 'node:path'
import { ownerRoot } from '~/server/utils/library'
import { filterLoraStackForStudio, normalizeLoraStack } from '~/utils/loras'
import {
  clampImageCfg,
  clampImageScaleMegapixels,
  clampImageSteps,
  clampVideoCfg,
  normalizePresetName,
  parseGenerationPresetKind,
  PRESET_PROMPT_MAX,
  type GenerationPreset,
  type GenerationPresetKind,
  type ImagePresetSettings,
  type VideoPresetSettings
} from '~/utils/generationPresets'
import {
  IMAGE_V2_CFG_DEFAULT,
  IMAGE_V2_CONSISTENCY_CLIP,
  IMAGE_V2_CONSISTENCY_MODEL,
  IMAGE_V2_SNOFS_CLIP,
  IMAGE_V2_SNOFS_MODEL,
  IMAGE_V2_STEPS_DEFAULT,
  IMAGE_V2_KREA_CFG,
  IMAGE_V2_KREA_COMPOSE_DENOISE,
  IMAGE_V2_KREA_EDIT_DENOISE,
  IMAGE_V2_KREA_STEPS,
  IMAGE_V2_DENOISE_DEFAULT,
  clampImageV2Denoise,
  clampImageV2Size,
  clampImageV2Strength,
  parseImageV2Engine,
  parseImageV2Mode,
  parseImageV2Task,
  type ImageV2PresetSettings
} from '~/utils/imageV2'
import { coerceVideoWorkflow, isXaigenStudio, parseVideoWorkflow } from '~/utils/videoModels'

const writeChains = new Map<string, Promise<unknown>>()

function presetsPath(owner: string) {
  return join(ownerRoot(owner), 'generation-presets.json')
}

function ensureOwner(owner: string) {
  mkdirSync(ownerRoot(owner), { recursive: true })
}

function readPresets(owner: string): GenerationPreset[] {
  ensureOwner(owner)
  const path = presetsPath(owner)
  if (!existsSync(path)) return []
  try {
    const parsed = JSON.parse(readFileSync(path, 'utf8'))
    return Array.isArray(parsed) ? parsed.map(normalizeStoredPreset).filter((item): item is GenerationPreset => Boolean(item)) : []
  } catch {
    return []
  }
}

function writePresets(owner: string, presets: GenerationPreset[]) {
  ensureOwner(owner)
  const path = presetsPath(owner)
  const tmp = `${path}.tmp`
  writeFileSync(tmp, JSON.stringify(presets, null, 2))
  renameSync(tmp, path)
}

function mutate<T>(owner: string, fn: (presets: GenerationPreset[]) => T): Promise<T> {
  const prev = writeChains.get(owner) || Promise.resolve()
  const run = prev.then(() => {
    const presets = readPresets(owner)
    const result = fn(presets)
    writePresets(owner, presets)
    return result
  })
  writeChains.set(owner, run.then(() => undefined, () => undefined))
  return run
}

function clampNum(raw: unknown, fallback: number, min: number, max: number, step = 0.1) {
  const value = Number(raw)
  if (!Number.isFinite(value)) return fallback
  const snapped = Math.round(value / step) * step
  return Math.min(max, Math.max(min, Math.round(snapped * 100) / 100))
}

function sanitizeVideoSettings(raw: unknown): VideoPresetSettings {
  const rec = raw && typeof raw === 'object' ? raw as Record<string, unknown> : {}
  const workflow = coerceVideoWorkflow(parseVideoWorkflow(rec.workflow))
  const turbo = rec.turbo !== false
  const sampler = String(rec.samplerName || '')
  const scheduler = String(rec.scheduler || '')
  const aspect = String(rec.aspect || 'auto')
  return {
    workflow,
    turbo,
    cfg: clampVideoCfg(rec.cfg, turbo ? 1.5 : 4),
    fps: rec.fps === 12 || rec.fps === 30 || rec.fps === 24 ? Number(rec.fps) : 24,
    samplerName: sampler === 'euler' || sampler === 'dpmpp_2m' || sampler === 'res_multistep' ? sampler : 'res_multistep',
    scheduler: scheduler === 'ddim_uniform' || scheduler === 'sgm_uniform' || scheduler === 'simple' ? scheduler : 'simple',
    duration: clampNum(rec.duration, 5, 0.5, 120, 0.1),
    aspect: aspect === '16:9' || aspect === '9:16' || aspect === '1:1' || aspect === 'auto' ? aspect : 'auto',
    sound: rec.sound !== false
  }
}

function sanitizePrompt(raw: unknown) {
  return typeof raw === 'string' ? raw.slice(0, PRESET_PROMPT_MAX) : undefined
}

function sanitizeImageSettings(raw: unknown): ImagePresetSettings {
  const rec = raw && typeof raw === 'object' ? raw as Record<string, unknown> : {}
  return {
    prompt: sanitizePrompt(rec.prompt),
    steps: clampImageSteps(rec.steps),
    cfg: clampImageCfg(rec.cfg),
    negative: String(rec.negative || '').slice(0, 2000),
    scaleToTotalPixels: rec.scaleToTotalPixels === true,
    scaleMegapixels: clampImageScaleMegapixels(rec.scaleMegapixels),
    aspect: rec.aspect === '16:9' || rec.aspect === '9:16' || rec.aspect === '1:1' || rec.aspect === 'auto' ? rec.aspect : 'auto'
  }
}

function sanitizeImageV2Settings(raw: unknown): ImageV2PresetSettings {
  const rec = raw && typeof raw === 'object' ? raw as Record<string, unknown> : {}
  const mode = parseImageV2Mode(rec.mode) || 'edit'
  const turbo = rec.turbo === true
  const engine = parseImageV2Engine(rec.engine, 'flux')
  return {
    mode,
    engine,
    task: mode === 'compose' ? parseImageV2Task(rec.task, 'scene') : mode === 'refine' ? 'refine' : mode === 'generate' ? 't2i' : 'scene',
    prompt: sanitizePrompt(rec.prompt),
    negative: String(rec.negative || '').slice(0, 2000),
    snofsModel: isXaigenStudio() ? clampImageV2Strength(rec.snofsModel ?? rec.snofs_model, IMAGE_V2_SNOFS_MODEL) : 0,
    snofsClip: isXaigenStudio() ? clampImageV2Strength(rec.snofsClip ?? rec.snofs_clip, IMAGE_V2_SNOFS_CLIP) : 0,
    consistencyModel: clampImageV2Strength(rec.consistencyModel ?? rec.consistency_model, mode === 'generate' ? 0 : IMAGE_V2_CONSISTENCY_MODEL),
    consistencyClip: clampImageV2Strength(rec.consistencyClip ?? rec.consistency_clip, mode === 'generate' ? 0 : IMAGE_V2_CONSISTENCY_CLIP),
    steps: engine === 'krea' ? clampImageSteps(rec.steps, IMAGE_V2_KREA_STEPS) : turbo ? 8 : clampImageSteps(rec.steps, IMAGE_V2_STEPS_DEFAULT),
    cfg: engine === 'krea' ? clampImageCfg(rec.cfg, IMAGE_V2_KREA_CFG) : turbo ? 1 : clampImageCfg(rec.cfg, IMAGE_V2_CFG_DEFAULT),
    megapixels: mode === 'generate' ? undefined : clampImageScaleMegapixels(rec.megapixels ?? rec.scaleMegapixels, 1),
    turbo,
    strength: mode === 'generate' || (mode !== 'refine' && engine !== 'krea')
      ? undefined
      : clampImageV2Denoise(
        rec.strength,
        mode === 'compose' ? IMAGE_V2_KREA_COMPOSE_DENOISE : mode === 'edit' ? IMAGE_V2_KREA_EDIT_DENOISE : IMAGE_V2_DENOISE_DEFAULT
      ),
    width: mode === 'generate' ? clampImageV2Size(rec.width, rec.height, rec.aspect).width : undefined,
    height: mode === 'generate' ? clampImageV2Size(rec.width, rec.height, rec.aspect).height : undefined,
    aspect: rec.aspect === '16:9' || rec.aspect === '9:16' || rec.aspect === '1:1' || rec.aspect === 'auto' ? String(rec.aspect) : (mode === 'generate' ? '1:1' : 'auto')
  }
}

function sanitizeSettings(kind: GenerationPresetKind, raw: unknown) {
  if (kind === 'image') return sanitizeImageSettings(raw)
  if (kind === 'imagev2') return sanitizeImageV2Settings(raw)
  return sanitizeVideoSettings(raw)
}

function normalizeStoredPreset(raw: unknown): GenerationPreset | null {
  if (!raw || typeof raw !== 'object') return null
  const rec = raw as Record<string, unknown>
  const name = normalizePresetName(rec.name)
  if (!name) return null
  const kind = parseGenerationPresetKind(rec.kind)
  const id = String(rec.id || '').trim() || crypto.randomUUID()
  const createdAt = Number(rec.createdAt) || Date.now()
  return {
    id,
    name,
    kind,
    createdAt,
    updatedAt: Number(rec.updatedAt) || createdAt,
    loraStack: normalizeLoraStack(rec.loraStack),
    settings: sanitizeSettings(kind, rec.settings)
  }
}

function publicPreset(preset: GenerationPreset): GenerationPreset {
  const xaigen = isXaigenStudio()
  const settings = preset.kind === 'video'
    ? { ...preset.settings, workflow: coerceVideoWorkflow((preset.settings as VideoPresetSettings).workflow) }
    : preset.kind === 'imagev2' && !xaigen
      ? { ...preset.settings, snofsModel: 0, snofsClip: 0 }
      : preset.settings
  return {
    ...preset,
    loraStack: filterLoraStackForStudio(preset.loraStack, xaigen),
    settings
  }
}

export function listGenerationPresets(owner: string, kind?: GenerationPresetKind) {
  return readPresets(owner)
    .filter(item => !kind || item.kind === kind)
    .map(publicPreset)
    .sort((a, b) => a.name.localeCompare(b.name) || b.updatedAt - a.updatedAt)
}

export async function saveGenerationPreset(owner: string, input: {
  name: string
  kind: GenerationPresetKind
  loraStack?: unknown
  settings?: unknown
  overwrite?: boolean
}) {
  const name = normalizePresetName(input.name)
  if (!name) {
    throw createError({ statusCode: 400, statusMessage: 'A preset name is required' })
  }
  const kind = parseGenerationPresetKind(input.kind)
  const xaigen = isXaigenStudio()
  const loraStack = filterLoraStackForStudio(input.loraStack, xaigen)
  const settings = sanitizeSettings(kind, input.settings)
  const now = Date.now()
  return mutate(owner, (presets) => {
    const existing = presets.find(item => item.kind === kind && item.name.toLowerCase() === name.toLowerCase())
    if (existing && !input.overwrite) {
      throw createError({ statusCode: 409, statusMessage: `Overwrite preset “${existing.name}”?` })
    }
    if (existing) {
      existing.updatedAt = now
      existing.loraStack = loraStack
      existing.settings = settings
      return publicPreset(existing)
    }
    const created: GenerationPreset = {
      id: crypto.randomUUID(),
      name,
      kind,
      createdAt: now,
      updatedAt: now,
      loraStack,
      settings
    }
    presets.unshift(created)
    return publicPreset(created)
  })
}

export async function deleteGenerationPreset(owner: string, id: string) {
  return mutate(owner, (presets) => {
    const index = presets.findIndex(item => item.id === id)
    if (index < 0) throw createError({ statusCode: 404, statusMessage: 'Preset not found' })
    const [removed] = presets.splice(index, 1)
    return removed
  })
}