Genarrative/scripts/check-bgfilter-live.mjs

import { deflateSync } from 'node:zlib';
import { existsSync, readFileSync } from 'node:fs';
import { dirname, resolve } from 'node:path';
import { fileURLToPath } from 'node:url';

const __dirname = dirname(fileURLToPath(import.meta.url));
const root = resolve(__dirname, '..');

const DEFAULT_BASE_URL = 'http://58.87.105.82/bgfilter';
const DEFAULT_SCREEN_COLOR = '#CFEFFF';
const DEFAULT_SEG_MODEL = 'birefnet';
const DEFAULT_TIMEOUT_MS = 60_000;

function loadEnvFile(path) {
  if (!existsSync(path)) {
    return {};
  }
  const env = {};
  const content = readFileSync(path, 'utf8');
  for (const line of content.split(/\r?\n/)) {
    const trimmed = line.trim();
    if (!trimmed || trimmed.startsWith('#')) {
      continue;
    }
    const eqIndex = trimmed.indexOf('=');
    if (eqIndex < 0) {
      continue;
    }
    const key = trimmed.slice(0, eqIndex).trim();
    let value = trimmed.slice(eqIndex + 1).trim();
    if (
      (value.startsWith('"') && value.endsWith('"')) ||
      (value.startsWith("'") && value.endsWith("'"))
    ) {
      value = value.slice(1, -1);
    }
    env[key] = value;
  }
  return env;
}

const fileEnv = {
  ...loadEnvFile(resolve(root, '.env')),
  ...loadEnvFile(resolve(root, '.env.local')),
  ...loadEnvFile(resolve(root, '.env.secrets.local')),
};

function readConfig(name) {
  return process.env[name] || fileEnv[name] || '';
}

function readArg(name) {
  const prefix = `--${name}=`;
  const value = process.argv.find((arg) => arg.startsWith(prefix));
  return value ? value.slice(prefix.length) : '';
}

function readPositiveIntArg(name, fallback) {
  const rawValue = readArg(name);
  if (!rawValue) {
    return fallback;
  }
  const value = Number.parseInt(rawValue, 10);
  if (!Number.isFinite(value) || value <= 0) {
    throw new Error(`--${name} 必须是正整数，当前为 ${rawValue}`);
  }
  return value;
}

function normalizeHexColor(value) {
  const color = value.trim();
  if (!/^#[0-9a-fA-F]{6}$/.test(color)) {
    throw new Error(`screen_color 必须是 #RRGGBB，当前为 ${value}`);
  }
  return color.toUpperCase();
}

function parseHexColor(value) {
  const color = normalizeHexColor(value);
  return [
    Number.parseInt(color.slice(1, 3), 16),
    Number.parseInt(color.slice(3, 5), 16),
    Number.parseInt(color.slice(5, 7), 16),
  ];
}

function crc32(buffer) {
  let crc = 0xffffffff;
  for (const byte of buffer) {
    crc ^= byte;
    for (let bit = 0; bit < 8; bit += 1) {
      crc = crc & 1 ? (crc >>> 1) ^ 0xedb88320 : crc >>> 1;
    }
  }
  return (crc ^ 0xffffffff) >>> 0;
}

function pngChunk(type, data) {
  const typeBuffer = Buffer.from(type, 'ascii');
  const length = Buffer.alloc(4);
  length.writeUInt32BE(data.length, 0);
  const checksum = Buffer.alloc(4);
  checksum.writeUInt32BE(crc32(Buffer.concat([typeBuffer, data])), 0);
  return Buffer.concat([length, typeBuffer, data, checksum]);
}

function makeProbePng(screenColor, width, height) {
  const [red, green, blue] = parseHexColor(screenColor);
  const rowSize = 1 + width * 4;
  const raw = Buffer.alloc(rowSize * height);
  const subjectLeft = Math.floor(width * 0.31);
  const subjectRight = Math.ceil(width * 0.69);
  const subjectTop = Math.floor(height * 0.25);
  const subjectBottom = Math.ceil(height * 0.75);

  for (let y = 0; y < height; y += 1) {
    const row = y * rowSize;
    raw[row] = 0;
    for (let x = 0; x < width; x += 1) {
      const offset = row + 1 + x * 4;
      const inSubject =
        x >= subjectLeft && x < subjectRight && y >= subjectTop && y < subjectBottom;
      raw[offset] = inSubject ? 35 : red;
      raw[offset + 1] = inSubject ? 35 : green;
      raw[offset + 2] = inSubject ? 35 : blue;
      raw[offset + 3] = 255;
    }
  }

  const ihdr = Buffer.alloc(13);
  ihdr.writeUInt32BE(width, 0);
  ihdr.writeUInt32BE(height, 4);
  ihdr[8] = 8;
  ihdr[9] = 6;
  ihdr[10] = 0;
  ihdr[11] = 0;
  ihdr[12] = 0;

  return Buffer.concat([
    Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]),
    pngChunk('IHDR', ihdr),
    pngChunk('IDAT', deflateSync(raw)),
    pngChunk('IEND', Buffer.alloc(0)),
  ]);
}

function buildEndpoint(baseUrl) {
  const base = baseUrl.trim().replace(/\/+$/, '');
  if (!base) {
    throw new Error('BgFilter base url 为空');
  }
  return `${base}/remove-background`;
}

const baseUrl =
  readArg('base-url') ||
  readConfig('GENARRATIVE_EDITOR_BGFILTER_BASE_URL') ||
  DEFAULT_BASE_URL;
const screenColor = normalizeHexColor(
  readArg('screen-color') ||
    readConfig('GENARRATIVE_EDITOR_BGFILTER_SCREEN_COLOR') ||
    DEFAULT_SCREEN_COLOR,
);
const segModel =
  readArg('seg-model') ||
  readConfig('GENARRATIVE_EDITOR_BGFILTER_SEG_MODEL') ||
  DEFAULT_SEG_MODEL;
const token =
  readConfig('GENARRATIVE_EDITOR_BGFILTER_TOKEN') ||
  readConfig('GENARRATIVE_EDITOR_BACKGROUND_REMOVAL_TOKEN');
const timeoutMs = Number.parseInt(
  readArg('timeout-ms') ||
    readConfig('GENARRATIVE_EDITOR_BGFILTER_REQUEST_TIMEOUT_MS') ||
    `${DEFAULT_TIMEOUT_MS}`,
  10,
);
const width = readPositiveIntArg('width', 96);
const height = readPositiveIntArg('height', width);

const endpoint = buildEndpoint(baseUrl);
const inputImage = makeProbePng(screenColor, width, height);
const form = new FormData();
form.append(
  'file',
  new Blob([inputImage], { type: 'image/png' }),
  'bgfilter-live-probe.png',
);
form.append('screen_color', screenColor);
form.append('seg_model', segModel);

const controller = new AbortController();
const timer = setTimeout(() => controller.abort(), timeoutMs);
const startedAt = Date.now();

console.log('BgFilter 探活');
console.log(`目标: ${endpoint}`);
console.log(`screen_color: ${screenColor}`);
console.log(`seg_model: ${segModel}`);
console.log(`image: ${width}x${height}, ${inputImage.length} bytes`);
console.log(`timeout_ms: ${timeoutMs}`);
console.log(`token: ${token ? '已配置' : '未配置'}\n`);

try {
  const headers = {};
  if (token) {
    headers['X-Genarrative-Image-Token'] = token;
  }
  const response = await fetch(endpoint, {
    method: 'POST',
    headers,
    body: form,
    signal: controller.signal,
  });
  clearTimeout(timer);

  const elapsedMs = Date.now() - startedAt;
  const contentType = response.headers.get('content-type') || '';
  const upstreamElapsedMs = response.headers.get('x-bgfilter-elapsed-ms') || '';
  const upstreamSegModel = response.headers.get('x-bgfilter-seg-model') || '';
  const upstreamScreenColor = response.headers.get('x-bgfilter-screen-color') || '';
  const body = Buffer.from(await response.arrayBuffer());

  if (!response.ok) {
    console.error(`失败: HTTP ${response.status} ${response.statusText}`);
    console.error(`耗时: ${elapsedMs}ms`);
    console.error(body.toString('utf8').slice(0, 500));
    process.exit(1);
  }

  if (!contentType.toLowerCase().startsWith('image/')) {
    console.error(`失败: 返回 content-type 不是图片: ${contentType || '(empty)'}`);
    console.error(body.toString('utf8').slice(0, 500));
    process.exit(1);
  }

  if (body.length === 0) {
    console.error('失败: 返回图片为空');
    process.exit(1);
  }

  console.log(`成功: HTTP ${response.status} ${response.statusText}`);
  console.log(`耗时: ${elapsedMs}ms`);
  console.log(`content-type: ${contentType}`);
  console.log(`bytes: ${body.length}`);
  if (upstreamElapsedMs) {
    console.log(`x-bgfilter-elapsed-ms: ${upstreamElapsedMs}`);
  }
  if (upstreamSegModel) {
    console.log(`x-bgfilter-seg-model: ${upstreamSegModel}`);
  }
  if (upstreamScreenColor) {
    console.log(`x-bgfilter-screen-color: ${upstreamScreenColor}`);
  }
} catch (error) {
  clearTimeout(timer);
  const elapsedMs = Date.now() - startedAt;
  const message =
    error?.name === 'AbortError'
      ? `请求超时 (${timeoutMs}ms)`
      : error?.message || String(error);
  console.error(`失败: ${message}`);
  console.error(`耗时: ${elapsedMs}ms`);
  process.exit(1);
}