src/lib/getStableDiffusionParams.ts

import { StableDiffusionParams } from "@/types"

import { CharacterExpression, CharacterGender, CharacterPose, GenerationMode, HairColor, HairStyle } from './config'
import { generateSeed } from "./generateSeed"

export function getStableDiffusionParams({
  generationMode,
  negativePrompt,
  characterAge,
  characterExpression,
  characterGender,
  characterHairColor,
  characterHairStyle,
  characterName,
  characterPose
}: {
  generationMode: GenerationMode
  negativePrompt: string
  characterAge: number
  characterExpression: CharacterExpression
  characterGender: CharacterGender
  characterHairColor: HairColor
  characterHairStyle: HairStyle
  characterName: string
  characterPose: CharacterPose
}): StableDiffusionParams {

  const baseParams: StableDiffusionParams = {
    prompt: "",
    negativePrompt: "",
    guidanceScale: 9,
    seed: generateSeed(),
    width: 1024,
    height: 1024,
    numInferenceSteps: 20,
    hfApiKey: "",
  }

  const params: StableDiffusionParams = {
    ...baseParams,
    ...generationMode === "characters" ? {
      prompt: [
        `${characterAge} years old $${characterGender}`,
        `${characterHairColor} ${characterHairStyle}`,
        `looking ${characterExpression}`,
        `${characterPose}`,
        `named ${characterName}`,
        'beautiful',
        'award winning',
        'sharp',
        'crisp',
        'centered',
        'aligned'
      ].filter(x => x).join(", "),
      negativePrompt: [
        negativePrompt,
        'drawing',
        'painting',
        'unrealistic',
        'sitting',
        'chair',
        '3D render',
        'unaligned',
        'cropped',
        'bad hands',
        'wrong hands',
        'deformed',
        'glitch',
        'blurry',
        'overexposed'
      ].join(", "),
    } : {
      prompt: [
        'picture of a single',
        prompt,
        '3D render',
        'logo',
        'ios icon',
        'illustration',
        'vector graphics',
        'svg',
        'beautiful',
        'award winning',
        'sharp',
        'crisp',
        'centered',
        'aligned',
      ].join(", "),
      negativePrompt: [
        negativePrompt,
        'photo',
        'gradient',
        'many',
        'realistic',
        'shadow',
        'multiple',
        'various',
        'unaligned',
        'cropped',
        'bad hands',
        'wrong hands',
        'deformed',
        'glitch',
        'blurry',
        'overexposed'
      ].join(", "),
    }
  }
  return params
}