aigen/utils/imageV2.ts

export const IMAGE_V2_MODES = ['edit', 'compose', 'refine', 'generate'] as const
export const IMAGE_V2_TASKS = ['scene', 'identity', 'outfit', 'face_lock', 'refine', 't2i'] as const

export type ImageV2Mode = (typeof IMAGE_V2_MODES)[number]
export type ImageV2Task = (typeof IMAGE_V2_TASKS)[number]

export const IMAGE_V2_STEPS_DEFAULT = 24
export const IMAGE_V2_CFG_DEFAULT = 4
export const IMAGE_V2_TURBO_STEPS = 8
export const IMAGE_V2_TURBO_CFG = 1
export const IMAGE_V2_SNOFS_MODEL = 0.65
export const IMAGE_V2_SNOFS_CLIP = 0.35
export const IMAGE_V2_CONSISTENCY_MODEL = 0.7
export const IMAGE_V2_CONSISTENCY_CLIP = 0.7
export const IMAGE_V2_STRENGTH_MIN = 0
export const IMAGE_V2_STRENGTH_MAX = 2
export const IMAGE_V2_STRENGTH_STEP = 0.05

export const IMAGE_V2_SNOFS_LORA = 'klein_snofs_v1_4.safetensors'
export const IMAGE_V2_CONSISTENCY_LORA = 'Flux2-Klein-9B-consistency-V2.safetensors'
export const IMAGE_V2_DENOISE_DEFAULT = 0.35
export const IMAGE_V2_DENOISE_MIN = 0.15
export const IMAGE_V2_DENOISE_MAX = 0.75
export const IMAGE_V2_DENOISE_STEP = 0.01
export const IMAGE_V2_REFINE_FACE = {
  strength: 0.28,
  snofsModel: 0.55,
  snofsClip: 0.25,
  consistencyModel: 0.75,
  consistencyClip: 0.8
}
export const IMAGE_V2_REFINE_HAND = {
  strength: 0.4,
  snofsModel: 0.65,
  snofsClip: 0.3,
  consistencyModel: 0.7,
  consistencyClip: 0.7
}
export const IMAGE_V2_REFINE_HEAVY = { strength: 0.55 }
export const IMAGE_V2_REFINE_FACE_PROMPT = 'same face as the canvas, same glasses, same cheeks and jaw. Change only the face in the masked area.'
export const IMAGE_V2_SIZE_SIDES = [768, 1024, 1280] as const
export const IMAGE_V2_GENERATE_WIDTH = 1024
export const IMAGE_V2_GENERATE_HEIGHT = 1024
export const IMAGE_V2_GENERATE_ASPECTS = {
  '1:1': [1024, 1024],
  '3:4': [768, 1024],
  '4:3': [1024, 768],
  '16:9': [1280, 768],
  '9:16': [768, 1280]
} as const
export type ImageV2GenerateAspect = keyof typeof IMAGE_V2_GENERATE_ASPECTS
export const IMAGE_V2_GENERATE_SFW = { snofsModel: 0, snofsClip: 0 }
export const IMAGE_V2_GENERATE_SNOFS = { snofsModel: 0.65, snofsClip: 0.35 }

export const IMAGE_V2_ROLE_HEADERS: Record<Exclude<ImageV2Task, 'scene' | 'refine' | 't2i'>, string> = {
  outfit: 'Person, face, body, pose, and background from image 1. Clothing only from image 2. Fit the outfit from image 2 to the body in image 1. Do not copy image 2’s face, body shape, or pose.',
  face_lock: 'Body, pose, and scene from image 1. Exact face from image 2.',
  identity: 'Same person as image 1. Use image 2 only to reinforce the face. Follow the user’s pose/scene prompt.'
}

export function parseImageV2Mode(raw: unknown): ImageV2Mode | null {
  const value = String(raw || '').trim().toLowerCase()
  return IMAGE_V2_MODES.includes(value as ImageV2Mode) ? value as ImageV2Mode : null
}

export function parseImageV2Task(raw: unknown, fallback: ImageV2Task = 'scene'): ImageV2Task {
  const value = String(raw || '').trim().toLowerCase()
  return IMAGE_V2_TASKS.includes(value as ImageV2Task) ? value as ImageV2Task : fallback
}

export function clampImageV2Strength(raw: unknown, fallback: number) {
  const value = Number(raw)
  if (!Number.isFinite(value)) return fallback
  const snapped = Math.round(value / IMAGE_V2_STRENGTH_STEP) * IMAGE_V2_STRENGTH_STEP
  return Math.min(IMAGE_V2_STRENGTH_MAX, Math.max(IMAGE_V2_STRENGTH_MIN, Math.round(snapped * 100) / 100))
}

export function clampImageV2Denoise(raw: unknown, fallback = IMAGE_V2_DENOISE_DEFAULT) {
  const value = Number(raw)
  if (!Number.isFinite(value)) return fallback
  const snapped = Math.round(value / IMAGE_V2_DENOISE_STEP) * IMAGE_V2_DENOISE_STEP
  return Math.min(IMAGE_V2_DENOISE_MAX, Math.max(IMAGE_V2_DENOISE_MIN, Math.round(snapped * 100) / 100))
}

export function snapImageV2Side(raw: unknown, fallback = IMAGE_V2_GENERATE_WIDTH) {
  const value = Number(raw)
  if (!Number.isFinite(value)) return fallback
  return IMAGE_V2_SIZE_SIDES.reduce((best, side) => (
    Math.abs(side - value) < Math.abs(best - value) ? side : best
  ), IMAGE_V2_SIZE_SIDES[0])
}

export function parseImageV2Aspect(raw: unknown): ImageV2GenerateAspect | null {
  const value = String(raw || '').trim()
  return value in IMAGE_V2_GENERATE_ASPECTS ? value as ImageV2GenerateAspect : null
}

export function clampImageV2Size(width: unknown, height: unknown, aspect?: unknown) {
  const preset = parseImageV2Aspect(aspect)
  if (preset) {
    const [w, h] = IMAGE_V2_GENERATE_ASPECTS[preset]
    return { width: w, height: h, aspect: preset }
  }
  return {
    width: snapImageV2Side(width, IMAGE_V2_GENERATE_WIDTH),
    height: snapImageV2Side(height, IMAGE_V2_GENERATE_HEIGHT),
    aspect: null as ImageV2GenerateAspect | null
  }
}

export function composeImageV2Prompt(mode: ImageV2Mode, task: ImageV2Task, prompt: string) {
  const body = String(prompt || '').trim()
  if (mode === 'generate' || mode === 'refine' || mode !== 'compose' || task === 'scene' || task === 't2i') return body
  const header = IMAGE_V2_ROLE_HEADERS[task as Exclude<ImageV2Task, 'scene' | 'refine' | 't2i'>]
  return header ? `${header}\n\n${body}` : body
}

export type ImageV2PresetSettings = {
  mode?: ImageV2Mode
  task?: ImageV2Task
  negative?: string
  snofsModel?: number
  snofsClip?: number
  consistencyModel?: number
  consistencyClip?: number
  steps?: number
  cfg?: number
  megapixels?: number
  turbo?: boolean
  strength?: number
  width?: number
  height?: number
}