bar/scripts/generate-images.mjs

#!/usr/bin/env node
/**
 * generate-images.mjs
 * Generates AI cocktail images via Replicate flux-2-pro and saves them as
 * content/recipes/{slug}/cocktail.webp
 *
 * Usage:
 *   REPLICATE_API_TOKEN=r8_... node scripts/generate-images.mjs
 *   REPLICATE_API_TOKEN=r8_... node scripts/generate-images.mjs --limit 10
 *   REPLICATE_API_TOKEN=r8_... node scripts/generate-images.mjs --slug margarita
 *
 * Options:
 *   --limit N      Only process the first N cocktails
 *   --slug NAME    Process only the cocktail with this slug
 *   --concurrency N  Parallel requests (default: 3)
 *   --dry-run      Log what would be generated without calling API
 */

import { readFileSync, writeFileSync, existsSync, mkdirSync } from "fs";
import { join, dirname } from "path";
import { fileURLToPath } from "url";
import { parse } from "csv-parse/sync";
import Replicate from "replicate";

const __dirname = dirname(fileURLToPath(import.meta.url));
const ROOT = join(__dirname, "..");
const CSV_PATH = join(ROOT, "prototype/uploads/final_cocktails.csv");
const CONTENT_DIR = join(ROOT, "content/recipes");

const PROMPT = `A high-end, professional commercial photograph of the cocktail from the reference image. The drink retains its original vibrant color hue, served in the exact same glass type with identical decorations details. The liquid has a crisp, realistic texture with subtle condensation droplets on the outside of the glassware. The cocktail is set against a modern, moody, and minimalist bar background with soft, cinematic bokeh. Elegant, dramatic studio lighting casts a gentle glow on the glass, creating sharp reflections and a premium, editorial aesthetic fit for a luxury cocktail book. Photorealistic, 8k resolution, crisp focus.`;

// ── CLI args ──────────────────────────────────────────────────────────────────

const args = process.argv.slice(2);
const limitIdx = args.indexOf("--limit");
const LIMIT = limitIdx !== -1 ? parseInt(args[limitIdx + 1], 10) : Infinity;
const slugIdx = args.indexOf("--slug");
const ONLY_SLUG = slugIdx !== -1 ? args[slugIdx + 1] : null;
const concurrencyIdx = args.indexOf("--concurrency");
const CONCURRENCY = concurrencyIdx !== -1 ? parseInt(args[concurrencyIdx + 1], 10) : 3;
const DRY_RUN = args.includes("--dry-run");

// ── Helpers ──────────────────────────────────────────────────────────────────

function slugify(name) {
  return name
    .toLowerCase()
    .replace(/&/g, "and")
    .replace(/[^a-z0-9]+/g, "-")
    .replace(/^-+|-+$/g, "");
}

async function downloadToBuffer(url) {
  const res = await fetch(url);
  if (!res.ok) throw new Error(`HTTP ${res.status} fetching ${url}`);
  const ab = await res.arrayBuffer();
  return Buffer.from(ab);
}

// ── Replicate call ────────────────────────────────────────────────────────────

async function generateImage(replicate, cocktail, index, total) {
  const { slug, name, thumbnail } = cocktail;
  const outPath = join(CONTENT_DIR, slug, "cocktail.webp");

  if (existsSync(outPath)) {
    console.log(`[${index}/${total}] SKIP  ${name} (image exists)`);
    return;
  }

  if (!existsSync(join(CONTENT_DIR, slug))) {
    console.warn(`[${index}/${total}] WARN  ${name}: content dir missing, skipping`);
    return;
  }

  if (DRY_RUN) {
    console.log(`[${index}/${total}] DRY   ${name} → ${outPath}`);
    return;
  }

  console.log(`[${index}/${total}] GEN   ${name}…`);

  try {
    const input = {
      prompt: PROMPT,
      output_format: "webp",
      aspect_ratio: "1:1",
      safety_tolerance: 6,
    };

    // Add reference image if thumbnail URL exists
    if (thumbnail) {
      input.image = thumbnail;
      input.image_prompt_strength = 0.15;
    }

    const output = await replicate.run("black-forest-labs/flux-2-pro", { input });

    // Output is either a URL string or a ReadableStream
    let imageData;
    if (typeof output === "string") {
      imageData = await downloadToBuffer(output);
    } else if (Array.isArray(output) && output.length > 0) {
      const first = output[0];
      if (typeof first === "string") {
        imageData = await downloadToBuffer(first);
      } else {
        // Replicate stream
        const chunks = [];
        for await (const chunk of first) {
          chunks.push(chunk);
        }
        imageData = Buffer.concat(chunks);
      }
    } else if (output && typeof output[Symbol.asyncIterator] === "function") {
      const chunks = [];
      for await (const chunk of output) {
        chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
      }
      imageData = Buffer.concat(chunks);
    } else {
      throw new Error(`Unexpected output type: ${typeof output}`);
    }

    writeFileSync(outPath, imageData);
    console.log(`[${index}/${total}] DONE  ${name} → ${imageData.length} bytes`);
  } catch (err) {
    console.error(`[${index}/${total}] ERR   ${name}: ${err.message}`);
  }
}

// ── Concurrency limiter ───────────────────────────────────────────────────────

async function runWithConcurrency(tasks, concurrency) {
  const results = [];
  let idx = 0;

  async function worker() {
    while (idx < tasks.length) {
      const current = idx++;
      await tasks[current]();
    }
  }

  const workers = Array.from({ length: Math.min(concurrency, tasks.length) }, worker);
  await Promise.all(workers);
  return results;
}

// ── Main ──────────────────────────────────────────────────────────────────────

const apiToken = process.env.REPLICATE_API_TOKEN;
if (!apiToken && !DRY_RUN) {
  console.error("Error: REPLICATE_API_TOKEN environment variable is not set.");
  process.exit(1);
}

const replicate = new Replicate({ auth: apiToken });

const csv = readFileSync(CSV_PATH, "utf-8");
const records = parse(csv, {
  columns: true,
  skip_empty_lines: true,
  relax_column_count: true,
  trim: true,
});

// Build cocktail list
const slugMap = new Map();
const cocktails = [];

for (const row of records) {
  const name = (row.name || "").trim();
  if (!name) continue;

  let slug = slugify(name);
  if (slugMap.has(slug)) {
    slug = `${slug}-${row.id || slugMap.size}`;
  }
  slugMap.set(slug, true);

  if (ONLY_SLUG && slug !== ONLY_SLUG) continue;

  cocktails.push({
    slug,
    name,
    thumbnail: (row.drinkThumbnail || "").trim(),
  });
}

const limited = isFinite(LIMIT) ? cocktails.slice(0, LIMIT) : cocktails;
console.log(`Processing ${limited.length} cocktails (concurrency: ${CONCURRENCY})${DRY_RUN ? " [DRY RUN]" : ""}`);

const tasks = limited.map((cocktail, i) => () =>
  generateImage(replicate, cocktail, i + 1, limited.length)
);

await runWithConcurrency(tasks, CONCURRENCY);
console.log("\nAll done.");