aigen/server/utils/workflow.ts

// Nitro bundles this JSON into the production server output.
import workflowTemplate from '../assets/workflow_minimax_video.json'

export interface GenerateParams {
  prompt: string
  imageName: string
  width: number
  height: number
  steps: number
  seed: number
  turbo: boolean
  length: number
  cfg: number
  fps: number
  samplerName: string
  scheduler: string
  filenamePrefix?: string
  sound?: boolean
}

type WorkflowNode = { class_type: string; inputs: Record<string, unknown>; _meta?: { title?: string } }
type WorkflowGraph = Record<string, WorkflowNode>

const PROMPT_CLASSES = new Set([
  'MiniMaxH3ImageToVideo',
  'CLIPTextEncode',
  'CLIPTextEncodeQwen3VL',
  'TextEncodeQwen3VL'
])
const SEED_CLASSES = new Set(['RandomNoise', 'KSampler', 'SamplerMiniMax', 'KSamplerAdvanced'])
const IMAGE_CLASSES = new Set(['LoadImage'])
const LORA_CLASSES = new Set(['LoraLoaderModelOnly', 'LoraLoader'])

function entriesOf(graph: WorkflowGraph, classType: string) {
  return Object.entries(graph).filter(([, node]) => node.class_type === classType)
}

function linkId(value: unknown): string | null {
  return Array.isArray(value) && typeof value[0] === 'string' ? value[0] : null
}

function snap32(value: number) {
  return Math.max(32, Math.round(value / 32) * 32)
}

function titleOf(node: WorkflowNode) {
  return node._meta?.title || ''
}

export function buildWorkflow(params: GenerateParams) {
  const graph = structuredClone(workflowTemplate) as WorkflowGraph
  const width = snap32(params.width)
  const height = snap32(params.height)

  for (const [, node] of Object.entries(graph)) {
    if (IMAGE_CLASSES.has(node.class_type) && 'image' in node.inputs) {
      node.inputs.image = params.imageName
    }

    if (PROMPT_CLASSES.has(node.class_type) || node.class_type.includes('Qwen3')) {
      if ('prompt' in node.inputs) node.inputs.prompt = params.prompt
      if ('text' in node.inputs) node.inputs.text = params.prompt
    }

    if (node.class_type === 'MiniMaxH3ImageToVideo') {
      node.inputs.prompt = params.prompt
      node.inputs.width = width
      node.inputs.height = height
      node.inputs.length = params.length
    }

    if (node.class_type === 'CreateVideo' && 'fps' in node.inputs) {
      node.inputs.fps = params.fps
    }

    if (node.class_type === 'KSamplerSelect' && 'sampler_name' in node.inputs) {
      node.inputs.sampler_name = params.samplerName
    }

    if (node.class_type === 'BasicScheduler' && 'scheduler' in node.inputs) {
      node.inputs.scheduler = params.scheduler
    }

    if ('cfg' in node.inputs && typeof node.inputs.cfg === 'number') {
      node.inputs.cfg = params.cfg
    }
    if ('guidance' in node.inputs && typeof node.inputs.guidance === 'number') {
      node.inputs.guidance = params.cfg
    }

    if (SEED_CLASSES.has(node.class_type)) {
      if ('noise_seed' in node.inputs) node.inputs.noise_seed = params.seed
      if ('seed' in node.inputs) node.inputs.seed = params.seed
    }

    if (LORA_CLASSES.has(node.class_type)) {
      node.inputs.strength_model = params.turbo ? 1 : 0
    }

    if (node.class_type === 'SaveVideo' && 'filename_prefix' in node.inputs) {
      node.inputs.filename_prefix = params.filenamePrefix || node.inputs.filename_prefix || 'video/MiniMax_H3'
    }

    if (node.class_type === 'PrimitiveBoolean') {
      node.inputs.value = params.turbo
    }

    if (node.class_type === 'ComfySwitchNode' && typeof node.inputs.switch === 'boolean') {
      node.inputs.switch = params.turbo
    }

    if (node.class_type === 'BasicScheduler' && typeof node.inputs.steps === 'number') {
      node.inputs.steps = params.steps
    }
  }

  if (graph['128']) {
    graph['128'] = {
      class_type: 'ImageScale',
      inputs: {
        upscale_method: 'nearest-exact',
        width,
        height,
        crop: 'center',
        image: ['114', 0]
      },
      _meta: { title: 'Scaling image' }
    }
  }

  graph['105:cfg'] = {
    class_type: 'FluxGuidance',
    inputs: {
      guidance: params.cfg,
      conditioning: ['105:104', 0]
    },
    _meta: { title: 'Guidance scale' }
  }
  if (graph['105:16']?.inputs) {
    graph['105:16'].inputs.conditioning = ['105:cfg', 0]
  }

  for (const [, node] of entriesOf(graph, 'ComfySwitchNode')) {
    if (!/step/i.test(titleOf(node))) continue
    const qualityId = linkId(node.inputs.on_false)
    const turboId = linkId(node.inputs.on_true)
    if (qualityId && graph[qualityId]?.class_type === 'PrimitiveInt') {
      graph[qualityId].inputs.value = params.turbo ? 20 : params.steps
    }
    if (turboId && graph[turboId]?.class_type === 'PrimitiveInt') {
      graph[turboId].inputs.value = params.turbo ? params.steps : 8
    }
  }

  if (params.sound === false) {
    delete graph['105:23']
    delete graph['105:24']
    if (graph['105:91']?.inputs) delete graph['105:91'].inputs.audio
  }

  return graph
}

const LABEL_OVERRIDES: Record<string, string> = {
  '114': 'Loading image',
  '128': 'Scaling image',
  '127': 'Reading image size',
  '119': 'Scaling image',
  '120': 'Reading image size',
  '92': 'Saving MP4',
  '105:6': 'Loading MiniMax H3 model',
  '105:11': 'Loading video VAE',
  '105:13': 'Loading Qwen3-VL encoder',
  '105:15': 'Preparing seed',
  '105:9': 'Building schedule',
  '105:10': 'Decoding video frames',
  '105:14': 'Sampling',
  '105:16': 'Preparing guider',
  '105:cfg': 'Applying guidance',
  '105:17': 'Selecting sampler',
  '105:23': 'Decoding audio',
  '105:24': 'Loading audio VAE',
  '105:91': 'Encoding video',
  '105:104': 'Encoding prompt & latent',
  '105:107': 'Computing duration',
  '105:111': 'Setting duration',
  '105:121': 'Applying turbo LoRA',
  '105:122': 'Selecting model path',
  '105:123': 'Selecting steps',
  '105:124': 'Quality steps',
  '105:125': 'Turbo steps',
  '105:126': 'Selecting turbo LoRA'
}

export const NODE_LABELS: Record<string, string> = {
  ...Object.fromEntries(
    Object.entries(workflowTemplate as WorkflowGraph).map(([id, node]) => [id, node._meta?.title || `Node ${id}`])
  ),
  ...LABEL_OVERRIDES
}

export function isEncodingNode(node: string) {
  return node === '92' || node === '105:91' || /encoding|saving mp4/i.test(NODE_LABELS[node] || '')
}