Blame for server.mjs

/**
 * Flux Studio — Zero GPU
 * ------------------------------------------------------------------
 * A CPU-only web app (Hugging Face Docker Space) that generates
 * images by proxying to the Black Forest Labs (BFL) FLUX API.
 *
 * The BFL_API_KEY never reaches the browser: every generation call
 * goes through this server, which attaches the key server-side.
 *
 * Configuration (env vars):
 *   BFL_API_KEY            Black Forest Labs API key (secret).
 *                          Without it the UI still loads and reports
 *                          its config state, but generation is refused.
 *   BFL_MODEL_ID           Preferred model  (default: flux-2-pro-preview)
 *   BFL_FALLBACK_MODEL_ID  Fallback model   (default: flux-2-klein-9b)
 *   BFL_API_BASE_URL       API base URL     (default: https://api.bfl.ai)
 *   PORT                   Listen port      (default: 7860, HF Spaces)
 */

import express from 'express';
import path from 'node:path';
import { fileURLToPath, pathToFileURL } from 'node:url';
import crypto from 'node:crypto';

const __dirname = path.dirname(fileURLToPath(import.meta.url));

/* ------------------------------------------------------------------ */
/* Configuration                                                       */
/* ------------------------------------------------------------------ */

const PORT = Number(process.env.PORT || 7860, 10);
const HOST = process.env.HOST || '0.0.0.0';

const BFL_API_KEY = (process.env.BFL_API_KEY || '').trim();
const BFL_MODEL_ID = (process.env.BFL_MODEL_ID || 'flux-2-pro-preview').trim();
const BFL_FALLBACK_MODEL_ID = (process.env.BFL_FALLBACK_MODEL_ID || 'flux-2-klein-9b').trim();
const BFL_API_BASE_URL = (process.env.BFL_API_BASE_URL || 'https://api.bfl.ai').trim().replace(/\/+$/, '');

const POLL_INTERVAL_MS = 1000;
const POLL_TIMEOUT_MS = 300_000; // give a single generation up to 5 minutes
const JOB_MAX = 100;
const JOB_TTL_MS = 30 * 60 * 1000;

/** A key that is obviously a placeholder must not be treated as configured. */
function isPlaceholderKey(key) {
  return !key || /placeholder|replace.?me|dummy|example|changeme|xxx|todo/i.test(key);
}

const apiKeyPresent = Boolean(BFL_API_KEY);
const apiKeyConfigured = apiKeyPresent && !isPlaceholderKey(BFL_API_KEY);

const config = {
  appName: 'Flux Studio — Zero GPU',
  port: PORT,
  apiKeyPresent,
  apiKeyConfigured,
  placeholderKey: apiKeyPresent && !apiKeyConfigured,
  modelId: BFL_MODEL_ID,
  fallbackModelId: BFL_FALLBACK_MODEL_ID,
  apiBaseHost: (() => {
    try {
      return new URL(BFL_API_BASE_URL).host;
    } catch {
      return BFL_API_BASE_URL;
    }
  })(),
  ready: apiKeyConfigured,
};

/* ------------------------------------------------------------------ */
/* BFL API client helpers                                              */
/* ------------------------------------------------------------------ */

function bflHeaders() {
  return {
    accept: 'application/json',
    'content-type': 'application/json',
    'x-key': BFL_API_KEY,
    authorization: `Bearer ${BFL_API_KEY}`,
  };
}

/**
 * Submit a generation task to a BFL model endpoint.
 * Returns { id, polling_url }.
 */
async function bflSubmit(model, payload) {
  const url = `${BFL_API_BASE_URL}/v1/${encodeURIComponent(model)}`;
  let res;
  try {
    res = await fetch(url, {
      method: 'POST',
      headers: bflHeaders(),
      body: JSON.stringify(payload),
      signal: AbortSignal.timeout(60_000),
    });
  } catch (err) {
    const e = new Error(`Could not reach the BFL API at ${config.apiBaseHost}: ${err.message}`);
    e.code = 'UPSTREAM_UNREACHABLE';
    throw e;
  }

  const text = await res.text();
  let data = null;
  try {
    data = text ? JSON.parse(text) : null;
  } catch {
    /* non-JSON body */
  }

  if (!res.ok) {
    const detail =
      (data && (data.detail || data.error || (Array.isArray(data.errors) && data.errors.join('; ')) || JSON.stringify(data))) ||
      text.slice(0, 300) ||
      res.statusText;
    const e = new Error(`BFL API rejected the request (HTTP ${res.status}) for model "${model}": ${detail}`);
    e.code = res.status === 401 || res.status === 403 ? 'AUTH' : `HTTP_${res.status}`;
    throw e;
  }

  if (!data || typeof data.polling_url !== 'string') {
    const e = new Error(`Unexpected response from BFL API for model "${model}" (missing polling_url): ${JSON.stringify(data).slice(0, 300)}`);
    e.code = 'BAD_RESPONSE';
    throw e;
  }
  return { id: data.id, pollingUrl: data.polling_url };
}

/** Poll a BFL polling_url until Ready / terminal error / timeout. */
async function bflPoll(pollingUrl) {
  const deadline = Date.now() + POLL_TIMEOUT_MS;
  for (;;) {
    await new Promise((r) => setTimeout(r, POLL_INTERVAL_MS));
    let res;
    try {
      res = await fetch(pollingUrl, {
        headers: { accept: 'application/json', 'x-key': BFL_API_KEY, authorization: `Bearer ${BFL_API_KEY}` },
        signal: AbortSignal.timeout(30_000),
      });
    } catch (err) {
      const e = new Error(`Lost contact with the BFL API while polling: ${err.message}`);
      e.code = 'UPSTREAM_UNREACHABLE';
      throw e;
    }

    let data = null;
    try {
      data = await res.json();
    } catch {
      /* transient bad body — keep polling until deadline */
    }

    if (data) {
      const status = data.status;
      if (status === 'Ready') {
        const sample = data?.result?.sample || data?.sample;
        if (typeof sample !== 'string') {
          const e = new Error(`Generation reported Ready but no image URL was returned: ${JSON.stringify(data).slice(0, 300)}`);
          e.code = 'BAD_RESPONSE';
          throw e;
        }
        return { imageUrl: sample, result: data.result || data };
      }
      if (status === 'Error' || status === 'Failed' || status === 'Cancelled' || status === 'Content Moderated') {
        const e = new Error(
          `Generation failed with status "${status}": ${JSON.stringify(data?.result || data).slice(0, 300)}`
        );
        e.code = status === 'Content Moderated' ? 'MODERATED' : 'GENERATION_FAILED';
        throw e;
      }
      if (typeof status === 'string') {
        // e.g. "Queue", "Pending", "Generating", "Requesting" — keep going
        if (typeof bflPoll.onProgress === 'function') bflPoll.onProgress(status);
      }
    }

    if (Date.now() > deadline) {
      const e = new Error(`Generation timed out after ${POLL_TIMEOUT_MS / 1000}s of polling.`);
      e.code = 'TIMEOUT';
      throw e;
    }
  }
}

/* ------------------------------------------------------------------ */
/* In-memory job store (single process — enough for a demo Space)      */
/* ------------------------------------------------------------------ */

const jobs = new Map();

function pruneJobs() {
  const now = Date.now();
  // Drop jobs that finished (or were created) too long ago.
  for (const [id, job] of jobs) {
    if (now - (job.finishedAt || job.createdAt) > JOB_TTL_MS) jobs.delete(id);
  }
  // If still over capacity, evict the oldest finished jobs first.
  if (jobs.size > JOB_MAX) {
    const finished = [...jobs.values()]
      .filter((j) => j.finishedAt)
      .sort((a, b) => a.finishedAt - b.finishedAt);
    for (const j of finished) {
      if (jobs.size <= JOB_MAX) break;
      jobs.delete(j.id);
    }
  }
}

function publicJob(job) {
  return {
    id: job.id,
    status: job.status,
    detail: job.detail || null,
    model: job.model || null,
    fallbackUsed: Boolean(job.fallbackUsed),
    prompt: job.prompt,
    params: job.params,
    imageUrl: job.imageUrl || null,
    error: job.error || null,
    code: job.code || null,
    createdAt: job.createdAt,
    finishedAt: job.finishedAt || null,
  };
}

/* ------------------------------------------------------------------ */
/* Request validation                                                  */
/* ------------------------------------------------------------------ */

const ASPECT_PRESETS = {
  '1:1': [1024, 1024],
  '16:9': [1344, 768],
  '9:16': [768, 1344],
  '4:3': [1152, 896],
  '3:4': [896, 1152],
  '3:2': [1216, 832],
  '2:3': [832, 1216],
};

function coerceInt(v, fallback) {
  if (v === undefined || v === null || v === '') return fallback;
  const n = Number(v);
  return Number.isFinite(n) && Math.trunc(n) === n ? Math.trunc(n) : NaN;
}

function validateBody(body) {
  const errors = [];
  const prompt = typeof body?.prompt === 'string' ? body.prompt.trim() : '';
  if (!prompt) errors.push('prompt is required');
  if (prompt.length > 2000) errors.push('prompt must be at most 2000 characters');

  let width = coerceInt(body?.width, 1024);
  let height = coerceInt(body?.height, 1024);
  const aspect = typeof body?.aspectRatio === 'string' ? body.aspectRatio : null;
  if (aspect) {
    if (!ASPECT_PRESETS[aspect]) {
      errors.push(`aspectRatio must be one of ${Object.keys(ASPECT_PRESETS).join(', ')}`);
    } else {
      [width, height] = ASPECT_PRESETS[aspect];
    }
  }
  if (!Number.isInteger(width) || width < 256 || width > 2048) errors.push('width must be an integer between 256 and 2048');
  if (!Number.isInteger(height) || height < 256 || height > 2048) errors.push('height must be an integer between 256 and 2048');

  const seed = body?.seed === undefined || body?.seed === null || body?.seed === '' ? undefined : coerceInt(body.seed, NaN);
  if (seed !== undefined && (!Number.isInteger(seed) || seed < 0 || seed > 4294967295)) {
    errors.push('seed must be an integer between 0 and 4294967295');
  }

  const outputFormat = body?.outputFormat || 'jpeg';
  if (!['jpeg', 'png'].includes(outputFormat)) errors.push('outputFormat must be "jpeg" or "png"');

  if (errors.length) {
    const e = new Error(`Invalid request: ${errors.join('; ')}`);
    e.code = 'INVALID_REQUEST';
    throw e;
  }
  return { prompt, width, height, seed, outputFormat };
}

/* ------------------------------------------------------------------ */
/* Generation pipeline: primary model, then fallback                   */
/* ------------------------------------------------------------------ */

function buildPayload(p) {
  const payload = { prompt: p.prompt, width: p.width, height: p.height, output_format: p.outputFormat };
  if (p.seed !== undefined) payload.seed = p.seed;
  return payload;
}

async function runGeneration(job) {
  const models = [BFL_MODEL_ID, BFL_FALLBACK_MODEL_ID].filter((m, i, a) => m && a.indexOf(m) === i);
  const failures = [];

  for (let i = 0; i < models.length; i++) {
    const model = models[i];
    try {
      job.detail = `Submitting to ${model}…`;
      const { pollingUrl } = await bflSubmit(model, buildPayload(job.params));
      job.model = model;
      job.fallbackUsed = i > 0;
      job.status = 'processing';

      bflPoll.onProgress = (status) => {
        job.detail = `${model}: ${status}…`;
      };
      const { imageUrl } = await bflPoll(pollingUrl);

      job.status = 'succeeded';
      job.imageUrl = imageUrl;
      job.detail = `Done via ${model}`;
      job.finishedAt = Date.now();
      return;
    } catch (err) {
      failures.push({ model, message: err.message, code: err.code || 'UNKNOWN' });
      job.detail = `${model} failed — ${failures.length < models.length ? 'trying fallback model…' : 'no more models to try'}`;
      // An auth failure will fail on every model — stop early.
      if (err.code === 'AUTH') break;
    } finally {
      bflPoll.onProgress = null;
    }
  }

  job.status = 'failed';
  job.finishedAt = Date.now();
  const first = failures[0];
  job.error = failures.map((f) => `[${f.model}] ${f.message}`).join(' | ');
  job.code = first ? first.code : 'UNKNOWN';
  if (job.code === 'AUTH') {
    job.error =
      'The BFL API rejected the API key (HTTP 401/403). If this Space uses a placeholder, set a real BFL_API_KEY under Settings → Secrets and restart. Underlying detail: ' +
      job.error;
  }
}

/* ------------------------------------------------------------------ */
/* Express app                                                         */
/* ------------------------------------------------------------------ */

const app = express();
app.disable('x-powered-by');
app.use(express.json({ limit: '64kb' }));

/* --- API: health & config ----------------------------------------- */

app.get('/api/health', (_req, res) => {
  res.json({ ok: true, uptime: process.uptime(), ts: Date.now() });
});

app.get('/api/config', (_req, res) => {
  res.json({
    appName: config.appName,
    modelId: config.modelId,
    fallbackModelId: config.fallbackModelId,
    apiBaseHost: config.apiBaseHost,
    apiKeyPresent: config.apiKeyPresent,
    apiKeyConfigured: config.apiKeyConfigured,
    placeholderKey: config.placeholderKey,
    ready: config.ready,
  });
});

/* --- API: generate (async job) ------------------------------------ */

app.post('/api/generate', (req, res) => {
  if (!config.apiKeyConfigured) {
    const msg = config.apiKeyPresent
      ? 'Generation is disabled: BFL_API_KEY is set but is a placeholder. Replace it with a real Black Forest Labs key under Space Settings → Secrets, then restart the Space.'
      : 'Generation is disabled: BFL_API_KEY is not configured. Add it under Space Settings → Secrets, then restart the Space.';
    return res.status(503).json({ error: msg, code: 'API_KEY_NOT_CONFIGURED' });
  }

  let params;
  try {
    params = validateBody(req.body);
  } catch (err) {
    return res.status(400).json({ error: err.message, code: err.code });
  }

  pruneJobs();
  const id = crypto.randomUUID();
  const job = {
    id,
    status: 'submitting',
    detail: 'Queued…',
    prompt: params.prompt,
    params,
    createdAt: Date.now(),
  };
  jobs.set(id, job);

  runGeneration(job).catch((err) => {
    job.status = 'failed';
    job.error = err.message;
    job.code = 'INTERNAL';
    job.finishedAt = Date.now();
  });

  res.status(202).json({ jobId: id, job: publicJob(job) });
});

app.get('/api/jobs/:id', (req, res) => {
  const job = jobs.get(req.params.id);
  if (!job) return res.status(404).json({ error: 'Unknown job id.', code: 'NOT_FOUND' });
  res.json(publicJob(job));
});

/* --- API: image proxy (restricted to BFL delivery hosts) ---------- */

const DELIVERY_HOST_SUFFIXES = ['.bfl.ai', '.blob.core.windows.net', '.cdn.bfl.ml'];

function isAllowedDeliveryHost(host) {
  const h = (host || '').toLowerCase();
  return h === 'bfl.ai' || h === 'api.bfl.ai' || DELIVERY_HOST_SUFFIXES.some((s) => h.endsWith(s));
}

app.get('/api/image', async (req, res) => {
  const raw = String(req.query.u || '');
  let target;
  try {
    target = new URL(raw);
  } catch {
    return res.status(400).json({ error: 'Invalid url parameter.', code: 'INVALID_REQUEST' });
  }
  if (!['http:', 'https:'].includes(target.protocol) || !isAllowedDeliveryHost(target.host)) {
    return res.status(400).json({ error: 'Only BFL delivery URLs may be proxied.', code: 'FORBIDDEN_URL' });
  }
  try {
    const upstream = await fetch(target, { signal: AbortSignal.timeout(60_000) });
    if (!upstream.ok || !upstream.body) {
      return res.status(502).json({ error: `Upstream returned HTTP ${upstream.status}.`, code: 'UPSTREAM' });
    }
    res.status(upstream.status);
    const ct = upstream.headers.get('content-type');
    if (ct) res.setHeader('content-type', ct);
    res.setHeader('cache-control', 'public, max-age=3600');
    upstream.body.pipe(res);
  } catch (err) {
    res.status(502).json({ error: `Proxy fetch failed: ${err.message}`, code: 'UPSTREAM' });
  }
});

/* --- Static client ------------------------------------------------- */

app.use(express.static(path.join(__dirname, 'public'), { extensions: ['html'] }));
app.get(/^\/(?!api\/).*/, (_req, res) => {
  res.sendFile(path.join(__dirname, 'public', 'index.html'));
});
app.use('/api', (_req, res) => res.status(404).json({ error: 'Unknown API route.', code: 'NOT_FOUND' }));

/* eslint-disable no-unused-vars */
app.use((err, _req, res, _next) => {
  const status = err.type === 'entity.parse.failed' || err.type === 'entity.too.large' ? 400 : 500;
  res.status(status).json({ error: err.message || 'Internal server error', code: 'INTERNAL' });
});

/* ------------------------------------------------------------------ */
/* Boot (only when run as the main module)                             */
/* ------------------------------------------------------------------ */

export { app, config, isPlaceholderKey, validateBody, buildPayload, ASPECT_PRESETS, isAllowedDeliveryHost };

const isMain =
  process.argv[1] && import.meta.url === pathToFileURL(path.resolve(process.argv[1])).href;

if (isMain) {
  const server = app.listen(PORT, HOST, () => {
    console.log(`[flux-studio] listening on http://${HOST}:${PORT}`);
    console.log(
      `[flux-studio] model=${config.modelId} fallback=${config.fallbackModelId} base=${config.apiBaseHost} keyConfigured=${config.apiKeyConfigured}${config.placeholderKey ? ' (placeholder detected)' : ''}`
    );
  });
  const shutdown = () => server.close(() => process.exit(0));
  process.on('SIGTERM', shutdown);
  process.on('SIGINT', shutdown);
}