aigen/tests/studio2-extend-fit.test.mjs

import test from 'node:test'
import assert from 'node:assert/strict'
import { createRequire } from 'node:module'
import { readFileSync, mkdtempSync, writeFileSync, existsSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import ts from 'typescript'

const require = createRequire(import.meta.url)
function load(path, resolve = require) {
  const source = readFileSync(new URL(path, import.meta.url), 'utf8')
  const code = ts.transpileModule(source, { compilerOptions: { target: ts.ScriptTarget.ES2022, module: ts.ModuleKind.CommonJS } }).outputText
  const exports = {}
  new Function('require', 'exports', code)(resolve, exports)
  return exports
}
const { runFfmpeg } = load('../server/utils/ffmpeg.ts', id => id.startsWith('node:') ? require(id) : {})
const { imageDimensions } = load('../server/utils/resolution.ts')
const { fitStill, needsFit } = load('../server/utils/studio2/media.ts', id =>
  id.startsWith('node:') ? require(id) : id === '../ffmpeg' ? { runFfmpeg } : id === '../resolution' ? { imageDimensions } : {}
)
const { isValidAnchorFile } = load('../server/utils/extensionFrame.ts', id =>
  id.endsWith('.json') ? { default: {} } : id.startsWith('node:') ? require(id) : {}
)

// A fake PNG carrying a real signature + IHDR width/height for imageDimensions() to read.
// This is not a decodable image; tests never run ffmpeg against it (fit is either skipped
// or mocked), they only need imageDimensions() to report the right size.
function fakePng(width, height) {
  const buf = Buffer.alloc(64) // isValidAnchorFile requires >= 64 bytes
  buf.set([0x89, 0x50, 0x4e, 0x47], 0) // PNG signature start
  buf.writeUInt32BE(width, 16)
  buf.writeUInt32BE(height, 20)
  return buf
}
const PNG_32x32 = fakePng(32, 32)

const text = readFileSync(new URL('../server/utils/studio2/runner.ts', import.meta.url), 'utf8')
const tree = ts.createSourceFile('runner.ts', text, ts.ScriptTarget.Latest, true)
function extractUpload(scope) {
  const nodes = tree.statements.filter(n =>
    (ts.isFunctionDeclaration(n) && n.name?.text === 'upload') ||
    (ts.isVariableStatement(n) && n.getText(tree).startsWith('const NEVER_FIT'))
  )
  const js = nodes.map(n => ts.transpileModule(n.getText(tree), { compilerOptions: { target: ts.ScriptTarget.ES2022 } }).outputText).join('\n')
  return new Function(...Object.keys(scope), `${js};return upload`)(...Object.values(scope))
}

function fixture() {
  const uploaded = []
  const record = { id: 'r', index: 0, files: [], request: { settings: { width: 32, height: 32 } } }
  const scope = {
    fitStill: async (data) => { throw new Error('fitStill must not be called for an already-canvas-sized handoff frame') },
    needsFit,
    saveRecord() {},
    useRuntimeConfig: () => ({ comfyFilenamePrefix: 'preview/MiniMax_H3' }),
    comfyFetch: async (_url, opts) => {
      uploaded.push(opts.body)
      return { ok: true, json: async () => ({ name: 'start.png', subfolder: 'preview/MiniMax_H3/studio2/r/0' }) }
    },
    scopedFile: () => true
  }
  return { scope, record, uploaded }
}

test('upload() never fits the start/hero labels when the still already matches the job canvas', async () => {
  const f = fixture()
  const upload = extractUpload(f.scope)
  await upload(f.record, 'start', PNG_32x32)
  assert.equal(f.record.handoffFitted, false)
  assert.equal(f.record.files.length, 1)
})

test('upload() never fits the hero label even when it mismatches the canvas aspect', async () => {
  const f = fixture()
  f.record.request.settings = { width: 1344, height: 768 }
  const upload = extractUpload(f.scope)
  // PNG_32x32 is square, canvas is 16:9 -- would trigger a fit for any other label.
  await upload(f.record, 'hero', PNG_32x32)
  assert.equal(f.record.handoffFitted, undefined) // only tracked for 'start'
})

test('upload() still fits a source still whose aspect truly differs from the job canvas', async () => {
  const f = fixture()
  f.record.request.settings = { width: 1344, height: 768 }
  let fitCalled = false
  f.scope.fitStill = async (data) => { fitCalled = true; return data }
  const upload = extractUpload(f.scope)
  await upload(f.record, 'compose', PNG_32x32)
  assert.equal(fitCalled, true)
})

test('needsFit is false for an exact canvas match and true for a real aspect mismatch', () => {
  assert.equal(needsFit(PNG_32x32, 32, 32), false)
  assert.equal(needsFit(PNG_32x32, 1344, 768), true)
})

test('two extends at identical width/height: the anchor PNG reaches Comfy byte-identical, no ffmpeg fit', async () => {
  const dir = mkdtempSync(join(tmpdir(), 'studio2-anchor-'))
  const anchorPath = join(dir, 'anchor-last.png')
  writeFileSync(anchorPath, PNG_32x32)
  assert.ok(isValidAnchorFile(anchorPath))

  for (const hop of [0, 1]) {
    const f = fixture()
    const upload = extractUpload(f.scope)
    const anchorBytes = readFileSync(anchorPath)
    const key = await upload(f.record, 'start', anchorBytes)
    assert.equal(f.record.handoffFitted, false, `hop ${hop} must not resample the anchor frame`)
    const sentBody = f.uploaded[0]
    const sentBlob = sentBody.get('image')
    const sentBytes = Buffer.from(await sentBlob.arrayBuffer())
    assert.deepEqual(sentBytes, anchorBytes, `hop ${hop} must upload the anchor PNG byte-identical`)
    assert.ok(key.length > 0)
  }
})