Happy Image Gen

{
  "name": "happy-image-gen",
  "version": "0.1.0",
  "description": "Universal AI image generation skill — OpenAI, Google, Replicate, Stability, FAL, Ark, Bailian, SiliconFlow",
  "type": "module",
  "private": true,
  "engines": {
    "bun": ">=1.1.0"
  },
  "scripts": {
    "gen": "bun scripts/main.ts"
  }
}

import path from "node:path";
import { readFile } from "node:fs/promises";
import { fileURLToPath } from "node:url";
import type { CliArgs, Provider, ProviderModule, Quality } from "./types";
import { ProviderError, ConfigError } from "./types";
import { writeBytesToFile, validateImageMagic } from "./utils/output";

const KNOWN_PROVIDERS: Provider[] = [
  "openai",
  "google",
  "replicate",
  "stability",
  "fal",
  "ark",
  "bailian",
  "siliconflow",
];

const IMPLEMENTED: Set<Provider> = new Set([
  "openai",
  "google",
  "replicate",
  "stability",
  "fal",
  "ark",
  "bailian",
  "siliconflow",
]);

function printUsage(): void {
  console.log(`happy-image-gen — universal AI image generation CLI

Usage:
  bun scripts/main.ts --prompt "A cat on grass" --image out.png
  bun scripts/main.ts --promptfiles system.md user.md --image out.png
  bun scripts/main.ts --prompt "..." --ref src.png --image out.png

Options:
  -p, --prompt <text>        Prompt text
  --promptfiles <files...>   Read prompt from files (concatenated with blank lines)
  --image <path>             Output image path (REQUIRED)
  --ref <path>               Reference image (repeat to pass multiple)
  --provider <id>            openai | google | replicate | stability | fal | ark | bailian | siliconflow
  -m, --model <id>           Provider-specific model id
  --ar <ratio>               Aspect ratio (e.g., 16:9, 1:1, 4:3)
  --size <WxH>               Explicit size (e.g., 1024x1024), overrides --ar
  --quality <preset>         draft | hd (default) | ultra
  --n <count>                Number of variations (default 1, not all providers support)
  --json                     Emit JSON result to stdout instead of plain text
  --setup                    Print first-time setup guide and exit
  -h, --help                 Show this message

Environment variables (single-provider MVP uses env for API keys):
  OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_IMAGE_MODEL
  (more providers coming — see references/providers.md)

See references/config/first-time-setup.md and assets/EXTEND.template.md for the EXTEND.md config mechanism.
`);
}

async function printSetupGuide(): Promise<void> {
  const here = path.dirname(fileURLToPath(import.meta.url));
  const guidePath = path.resolve(here, "../references/config/first-time-setup.md");
  try {
    const guide = await readFile(guidePath, "utf8");
    console.log(guide);
  } catch {
    console.log(
      "No first-time-setup.md found. Start with:\n  export OPENAI_API_KEY=sk-...\n  bun scripts/main.ts --prompt 'A cat' --image cat.png"
    );
  }
}

function parseArgs(argv: string[]): CliArgs {
  const args: CliArgs = {
    prompt: null,
    promptFiles: [],
    imagePath: null,
    provider: null,
    model: null,
    aspectRatio: null,
    size: null,
    quality: null,
    referenceImages: [],
    n: 1,
    json: false,
    setup: false,
    help: false,
  };

  const tokens = [...argv];
  while (tokens.length > 0) {
    const token = tokens.shift()!;
    switch (token) {
      case "-p":
      case "--prompt":
        args.prompt = tokens.shift() ?? null;
        break;
      case "--promptfiles":
        while (tokens.length > 0 && !tokens[0]!.startsWith("--") && tokens[0] !== "-p") {
          args.promptFiles.push(tokens.shift()!);
        }
        break;
      case "--image":
        args.imagePath = tokens.shift() ?? null;
        break;
      case "--ref":
        args.referenceImages.push(tokens.shift() ?? "");
        break;
      case "--provider":
        args.provider = tokens.shift() as Provider | null;
        break;
      case "-m":
      case "--model":
        args.model = tokens.shift() ?? null;
        break;
      case "--ar":
      case "--aspect":
        args.aspectRatio = tokens.shift() ?? null;
        break;
      case "--size":
        args.size = tokens.shift() ?? null;
        break;
      case "--quality":
        args.quality = tokens.shift() as Quality | null;
        break;
      case "--n":
        args.n = Math.max(1, Number(tokens.shift() || 1));
        break;
      case "--json":
        args.json = true;
        break;
      case "--setup":
        args.setup = true;
        break;
      case "-h":
      case "--help":
        args.help = true;
        break;
      default:
        if (!token.startsWith("-") && args.prompt === null && args.promptFiles.length === 0) {
          args.prompt = token;
        } else {
          console.error(`Unknown argument: ${token}`);
          args.help = true;
        }
    }
  }

  return args;
}

async function resolvePrompt(args: CliArgs): Promise<string> {
  if (args.prompt && args.promptFiles.length === 0) return args.prompt;
  if (args.promptFiles.length > 0) {
    const parts: string[] = [];
    if (args.prompt) parts.push(args.prompt);
    for (const file of args.promptFiles) {
      parts.push(await readFile(file, "utf8"));
    }
    return parts.join("\n\n").trim();
  }
  throw new ConfigError("No prompt provided. Use --prompt or --promptfiles.");
}

function detectProvider(explicit: Provider | null): Provider {
  if (explicit) {
    if (!KNOWN_PROVIDERS.includes(explicit)) {
      throw new ConfigError(
        `Unknown provider: ${explicit}. Known: ${KNOWN_PROVIDERS.join(", ")}`
      );
    }
    return explicit;
  }
  if (process.env.OPENAI_API_KEY) return "openai";
  if (process.env.GOOGLE_API_KEY || process.env.GEMINI_API_KEY) return "google";
  if (process.env.REPLICATE_API_TOKEN) return "replicate";
  if (process.env.STABILITY_API_KEY) return "stability";
  if (process.env.FAL_KEY) return "fal";
  if (process.env.ARK_API_KEY) return "ark";
  if (process.env.DASHSCOPE_API_KEY) return "bailian";
  if (process.env.SILICONFLOW_API_KEY) return "siliconflow";
  throw new ConfigError(
    "No provider selected and no provider-specific API key found in env. Pass --provider or set OPENAI_API_KEY / GOOGLE_API_KEY / FAL_KEY / ARK_API_KEY / DASHSCOPE_API_KEY / SILICONFLOW_API_KEY / STABILITY_API_KEY / REPLICATE_API_TOKEN."
  );
}

async function loadProvider(provider: Provider): Promise<ProviderModule> {
  if (!IMPLEMENTED.has(provider)) {
    throw new ConfigError(
      `Provider '${provider}' is declared but not yet implemented in this release. Available: ${[
        ...IMPLEMENTED,
      ].join(", ")}.`
    );
  }
  const mod = await import(`./providers/${provider}.ts`);
  const module = mod as ProviderModule;
  if (typeof module.defaultModel !== "function" || typeof module.generate !== "function") {
    throw new ConfigError(`Provider module '${provider}' is missing defaultModel/generate exports.`);
  }
  return module;
}

async function run(argv: string[]): Promise<number> {
  const args = parseArgs(argv);

  if (args.help) {
    printUsage();
    return 0;
  }
  if (args.setup) {
    await printSetupGuide();
    return 0;
  }

  if (!args.imagePath) {
    console.error("--image is required (output path).");
    printUsage();
    return 2;
  }

  const prompt = await resolvePrompt(args);
  const provider = detectProvider(args.provider);
  const mod = await loadProvider(provider);
  const model = args.model ?? mod.defaultModel();

  if (!args.quality) args.quality = "hd";

  const bytes = await mod.generate(prompt, model, args);
  await writeBytesToFile(bytes, args.imagePath);

  const kind = validateImageMagic(bytes);

  if (args.json) {
    console.log(
      JSON.stringify(
        {
          success: true,
          provider,
          model,
          image: path.resolve(args.imagePath),
          size_bytes: bytes.byteLength,
          format: kind,
        },
        null,
        2
      )
    );
  } else {
    console.log(
      `✓ Generated ${kind.toUpperCase()} (${bytes.byteLength} bytes) via ${provider}/${model}`
    );
    console.log(`  ${path.resolve(args.imagePath)}`);
  }

  return 0;
}

run(process.argv.slice(2))
  .then((code) => process.exit(code))
  .catch((err) => {
    if (err instanceof ProviderError) {
      console.error(`[${err.provider}] ${err.message}`);
    } else if (err instanceof ConfigError) {
      console.error(`config: ${err.message}`);
    } else if (err instanceof Error) {
      console.error(err.stack || err.message);
    } else {
      console.error(String(err));
    }
    process.exit(1);
  });

import type { CliArgs } from "../types";
import { ProviderError } from "../types";
import { orientationOf } from "../utils/aspect_ratio";

const API_BASE =
  process.env.ARK_BASE_URL ||
  "https://ark.cn-beijing.volces.com/api/v3";

// Seedream only accepts a fixed set of sizes.
const SEEDREAM_SIZES = {
  square: "2048x2048",
  landscape16x9: "2048x1152",
  portrait16x9: "1152x2048",
  landscape4x3: "1728x1296",
  portrait4x3: "1296x1728",
};

export function defaultModel(): string {
  return process.env.ARK_IMAGE_MODEL || "doubao-seedream-4-0-250528";
}

function requireApiKey(): string {
  const key = process.env.ARK_API_KEY;
  if (!key) {
    throw new ProviderError(
      "ARK_API_KEY not set. Get one at https://console.volcengine.com/ark",
      "ark"
    );
  }
  return key;
}

function pickSize(args: CliArgs): string {
  if (args.size) return args.size;
  const ar = args.aspectRatio;
  if (!ar || ar === "1:1") return SEEDREAM_SIZES.square;
  if (ar === "4:3") return SEEDREAM_SIZES.landscape4x3;
  if (ar === "3:4") return SEEDREAM_SIZES.portrait4x3;
  const orientation = orientationOf(ar);
  if (orientation === "landscape") return SEEDREAM_SIZES.landscape16x9;
  if (orientation === "portrait") return SEEDREAM_SIZES.portrait16x9;
  return SEEDREAM_SIZES.square;
}

type ArkResponse = {
  data?: Array<{ url?: string; b64_json?: string }>;
  error?: { message?: string; code?: string };
};

export async function generate(prompt: string, model: string, args: CliArgs): Promise<Uint8Array> {
  const apiKey = requireApiKey();

  if (args.referenceImages.length > 0) {
    // Seedream supports i2i via a separate endpoint / parameter. Not covered in MVP.
    throw new ProviderError(
      "Ark Seedream reference-image workflow (SeeEdit) is not wired in this release. Use --provider openai or google with --ref, or drop --ref.",
      "ark"
    );
  }

  const body: Record<string, unknown> = {
    model,
    prompt,
    size: pickSize(args),
    response_format: "url",
    watermark: false,
  };

  const res = await fetch(`${API_BASE}/images/generations`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${apiKey}`,
    },
    body: JSON.stringify(body),
  });

  if (!res.ok) {
    const text = await res.text();
    throw new ProviderError(`Ark ${res.status}: ${text}`, "ark");
  }

  const payload = (await res.json()) as ArkResponse;
  const item = payload.data?.[0];
  if (item?.b64_json) {
    return Uint8Array.from(Buffer.from(item.b64_json, "base64"));
  }
  if (item?.url) {
    const imgRes = await fetch(item.url);
    if (!imgRes.ok) {
      throw new ProviderError(`Failed to download Ark output (${imgRes.status})`, "ark");
    }
    return new Uint8Array(await imgRes.arrayBuffer());
  }
  throw new ProviderError("Ark response contained no image data", "ark");
}

import type { CliArgs } from "../types";
import { ProviderError } from "../types";
import { orientationOf } from "../utils/aspect_ratio";

const DEFAULT_BASE =
  process.env.DASHSCOPE_BASE_URL || "https://dashscope.aliyuncs.com";
const NATIVE_BASE = `${DEFAULT_BASE}/api/v1`;

// qwen-image supports these fixed sizes; other models accept "Wx H" with an asterisk.
const QWEN_IMAGE_SIZES = {
  square: "1024*1024",
  landscape16x9: "1664*928",
  portrait16x9: "928*1664",
  landscape4x3: "1472*1140",
  portrait4x3: "1140*1472",
};

export function defaultModel(): string {
  return process.env.DASHSCOPE_IMAGE_MODEL || "wanx2.1-t2i-turbo";
}

function requireApiKey(): string {
  const key = process.env.DASHSCOPE_API_KEY;
  if (!key) {
    throw new ProviderError(
      "DASHSCOPE_API_KEY not set. Get one at https://dashscope.console.aliyun.com/apiKey",
      "bailian"
    );
  }
  return key;
}

function pickSize(args: CliArgs): string {
  if (args.size) return args.size.replace(/x/i, "*");
  const ar = args.aspectRatio;
  if (!ar || ar === "1:1") return QWEN_IMAGE_SIZES.square;
  if (ar === "4:3") return QWEN_IMAGE_SIZES.landscape4x3;
  if (ar === "3:4") return QWEN_IMAGE_SIZES.portrait4x3;
  const orientation = orientationOf(ar);
  if (orientation === "landscape") return QWEN_IMAGE_SIZES.landscape16x9;
  if (orientation === "portrait") return QWEN_IMAGE_SIZES.portrait16x9;
  return QWEN_IMAGE_SIZES.square;
}

type SubmitResponse = {
  output?: { task_id?: string; task_status?: string };
  code?: string;
  message?: string;
};
type QueryResponse = {
  output?: {
    task_id?: string;
    task_status?: string;
    results?: Array<{ url?: string }>;
    code?: string;
    message?: string;
  };
  code?: string;
  message?: string;
};

async function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

export async function generate(prompt: string, model: string, args: CliArgs): Promise<Uint8Array> {
  const apiKey = requireApiKey();
  if (args.referenceImages.length > 0) {
    throw new ProviderError(
      "Bailian reference-image workflow not wired in MVP. Drop --ref to use text-to-image.",
      "bailian"
    );
  }

  const submitBody = {
    model,
    input: { prompt },
    parameters: {
      size: pickSize(args),
      n: Math.max(1, args.n),
    },
  };

  const submitRes = await fetch(`${NATIVE_BASE}/services/aigc/text2image/image-synthesis`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${apiKey}`,
      "X-DashScope-Async": "enable",
    },
    body: JSON.stringify(submitBody),
  });

  if (!submitRes.ok) {
    const text = await submitRes.text();
    throw new ProviderError(`Bailian submit ${submitRes.status}: ${text}`, "bailian");
  }

  const submit = (await submitRes.json()) as SubmitResponse;
  const taskId = submit.output?.task_id;
  if (!taskId) {
    throw new ProviderError(
      `Bailian submit returned no task_id: ${submit.message ?? JSON.stringify(submit)}`,
      "bailian"
    );
  }

  const deadline = Date.now() + 180_000;
  while (Date.now() < deadline) {
    await sleep(2000);
    const queryRes = await fetch(`${NATIVE_BASE}/tasks/${taskId}`, {
      headers: { Authorization: `Bearer ${apiKey}` },
    });
    if (!queryRes.ok) {
      const text = await queryRes.text();
      throw new ProviderError(`Bailian query ${queryRes.status}: ${text}`, "bailian");
    }
    const payload = (await queryRes.json()) as QueryResponse;
    const status = payload.output?.task_status;
    if (status === "SUCCEEDED") {
      const url = payload.output?.results?.[0]?.url;
      if (!url) throw new ProviderError("Bailian succeeded but no url in results", "bailian");
      const imgRes = await fetch(url);
      if (!imgRes.ok) {
        throw new ProviderError(`Failed to download Bailian output (${imgRes.status})`, "bailian");
      }
      return new Uint8Array(await imgRes.arrayBuffer());
    }
    if (status === "FAILED" || status === "CANCELED" || status === "UNKNOWN") {
      throw new ProviderError(
        `Bailian task ${status}: ${payload.output?.message ?? "(no detail)"}`,
        "bailian"
      );
    }
  }
  throw new ProviderError("Bailian task timed out after 180s", "bailian");
}

import path from "node:path";
import { readFile } from "node:fs/promises";
import type { CliArgs } from "../types";
import { ProviderError } from "../types";
import { orientationOf } from "../utils/aspect_ratio";
import { getMimeType } from "../utils/output";

const SYNC_BASE = "https://fal.run";

export function defaultModel(): string {
  return process.env.FAL_IMAGE_MODEL || "fal-ai/nano-banana";
}

function requireApiKey(): string {
  const key = process.env.FAL_KEY;
  if (!key) {
    throw new ProviderError(
      "FAL_KEY not set. Get one at https://fal.ai/dashboard/keys",
      "fal"
    );
  }
  return key;
}

type FalImageItem = { url?: string; content_type?: string };
type FalResponse = { images?: FalImageItem[]; image?: FalImageItem; error?: string };

function pickImageSize(args: CliArgs): string {
  if (args.size) return args.size;
  const orientation = orientationOf(args.aspectRatio);
  if (args.aspectRatio === "4:3") return "landscape_4_3";
  if (args.aspectRatio === "3:4") return "portrait_4_3";
  if (orientation === "landscape") return "landscape_16_9";
  if (orientation === "portrait") return "portrait_9_16";
  return "square_hd";
}

async function referenceAsDataUri(refPath: string): Promise<string> {
  const bytes = await readFile(refPath);
  const mime = getMimeType(path.basename(refPath));
  return `data:${mime};base64,${Buffer.from(bytes).toString("base64")}`;
}

export async function generate(prompt: string, model: string, args: CliArgs): Promise<Uint8Array> {
  const apiKey = requireApiKey();
  const input: Record<string, unknown> = {
    prompt,
    image_size: pickImageSize(args),
    num_images: Math.max(1, args.n),
  };

  if (args.referenceImages.length > 0) {
    input.image_url = await referenceAsDataUri(args.referenceImages[0]!);
  }

  const url = `${SYNC_BASE}/${model.replace(/^\/+/, "")}`;
  const res = await fetch(url, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Key ${apiKey}`,
    },
    body: JSON.stringify(input),
  });

  if (!res.ok) {
    const text = await res.text();
    throw new ProviderError(`FAL ${res.status}: ${text}`, "fal");
  }

  const payload = (await res.json()) as FalResponse;
  const item = payload.images?.[0] ?? payload.image;
  const outputUrl = item?.url;
  if (!outputUrl) {
    throw new ProviderError("FAL response contained no image URL", "fal");
  }

  const imgRes = await fetch(outputUrl);
  if (!imgRes.ok) {
    throw new ProviderError(`Failed to download FAL output (${imgRes.status})`, "fal");
  }
  return new Uint8Array(await imgRes.arrayBuffer());
}

import path from "node:path";
import { readFile } from "node:fs/promises";
import type { CliArgs } from "../types";
import { ProviderError } from "../types";
import { getMimeType } from "../utils/output";

const API_BASE = "https://generativelanguage.googleapis.com/v1beta";
const GEMINI_IMAGE_RATIOS = new Set(["1:1", "16:9", "9:16", "3:4", "4:3"]);

export function defaultModel(): string {
  return process.env.GOOGLE_IMAGE_MODEL || "gemini-3-pro-image-preview";
}

function requireApiKey(): string {
  const key = process.env.GOOGLE_API_KEY || process.env.GEMINI_API_KEY;
  if (!key) {
    throw new ProviderError(
      "GOOGLE_API_KEY (or GEMINI_API_KEY) not set. Get one at https://aistudio.google.com/app/apikey",
      "google"
    );
  }
  return key;
}

function pickAspect(ar: string | null): string {
  if (!ar) return "1:1";
  if (GEMINI_IMAGE_RATIOS.has(ar)) return ar;
  return "1:1";
}

export async function generate(prompt: string, model: string, args: CliArgs): Promise<Uint8Array> {
  const apiKey = requireApiKey();
  if (model.startsWith("imagen")) {
    return generateWithImagen(apiKey, prompt, model, args);
  }
  return generateWithGemini(apiKey, prompt, model, args);
}

type GeminiInlineData = { mime_type?: string; mimeType?: string; data: string };
type GeminiPart = {
  text?: string;
  inline_data?: GeminiInlineData;
  inlineData?: GeminiInlineData;
};
type GeminiResponse = {
  candidates?: Array<{ content?: { parts?: GeminiPart[] } }>;
  promptFeedback?: { blockReason?: string };
  error?: { message?: string; code?: number };
};

async function generateWithGemini(
  apiKey: string,
  prompt: string,
  model: string,
  args: CliArgs
): Promise<Uint8Array> {
  const parts: GeminiPart[] = [{ text: prompt }];

  for (const refPath of args.referenceImages) {
    const bytes = await readFile(refPath);
    parts.push({
      inline_data: {
        mime_type: getMimeType(path.basename(refPath)),
        data: Buffer.from(bytes).toString("base64"),
      },
    });
  }

  const body = {
    contents: [{ role: "user", parts }],
    generationConfig: {
      responseModalities: ["IMAGE"],
      imageConfig: { aspectRatio: pickAspect(args.aspectRatio) },
    },
  };

  const url = `${API_BASE}/models/${model}:generateContent?key=${encodeURIComponent(apiKey)}`;
  const res = await fetch(url, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(body),
  });

  if (!res.ok) {
    const text = await res.text();
    throw new ProviderError(`Gemini generateContent ${res.status}: ${text}`, "google");
  }

  const payload = (await res.json()) as GeminiResponse;
  if (payload.promptFeedback?.blockReason) {
    throw new ProviderError(
      `Gemini blocked the prompt (${payload.promptFeedback.blockReason})`,
      "google"
    );
  }

  const candidateParts = payload.candidates?.[0]?.content?.parts ?? [];
  for (const part of candidateParts) {
    const inline = part.inline_data ?? part.inlineData;
    if (inline?.data) {
      return Uint8Array.from(Buffer.from(inline.data, "base64"));
    }
  }

  throw new ProviderError("Gemini response contained no image part", "google");
}

type ImagenPrediction = { bytesBase64Encoded?: string; mimeType?: string };
type ImagenResponse = {
  predictions?: ImagenPrediction[];
  error?: { message?: string };
};

async function generateWithImagen(
  apiKey: string,
  prompt: string,
  model: string,
  args: CliArgs
): Promise<Uint8Array> {
  if (args.referenceImages.length > 0) {
    throw new ProviderError(
      "Imagen endpoint does not accept reference images. Use gemini-3-pro-image-preview for multimodal.",
      "google"
    );
  }

  const body = {
    instances: [{ prompt }],
    parameters: {
      sampleCount: 1,
      aspectRatio: pickAspect(args.aspectRatio),
    },
  };

  const url = `${API_BASE}/models/${model}:predict?key=${encodeURIComponent(apiKey)}`;
  const res = await fetch(url, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(body),
  });

  if (!res.ok) {
    const text = await res.text();
    throw new ProviderError(`Imagen predict ${res.status}: ${text}`, "google");
  }

  const payload = (await res.json()) as ImagenResponse;
  const b64 = payload.predictions?.[0]?.bytesBase64Encoded;
  if (!b64) {
    throw new ProviderError("Imagen response contained no prediction bytes", "google");
  }
  return Uint8Array.from(Buffer.from(b64, "base64"));
}

import path from "node:path";
import { readFile } from "node:fs/promises";
import type { CliArgs } from "../types";
import { ProviderError } from "../types";
import { orientationOf } from "../utils/aspect_ratio";
import { getMimeType } from "../utils/output";

type OpenAIImageResponse = { data: Array<{ url?: string; b64_json?: string }> };

export function defaultModel(): string {
  return process.env.OPENAI_IMAGE_MODEL || "gpt-image-1";
}

function pickSize(model: string, args: CliArgs): string {
  if (args.size) return args.size;

  const orientation = orientationOf(args.aspectRatio);
  if (model.includes("dall-e-3")) {
    if (orientation === "landscape") return "1792x1024";
    if (orientation === "portrait") return "1024x1792";
    return "1024x1024";
  }
  if (model.includes("dall-e-2")) return "1024x1024";
  if (orientation === "landscape") return "1536x1024";
  if (orientation === "portrait") return "1024x1536";
  return "1024x1024";
}

function pickQuality(model: string, quality: CliArgs["quality"]): string | null {
  if (model.includes("dall-e-3")) {
    return quality === "ultra" ? "hd" : "standard";
  }
  if (model.includes("gpt-image")) {
    if (quality === "draft") return "low";
    if (quality === "ultra") return "high";
    return "medium";
  }
  return null;
}

export async function generate(prompt: string, model: string, args: CliArgs): Promise<Uint8Array> {
  const baseUrl = process.env.OPENAI_BASE_URL || "https://api.openai.com/v1";
  const apiKey = process.env.OPENAI_API_KEY;

  if (!apiKey) {
    throw new ProviderError(
      "OPENAI_API_KEY is not set. Export it in your shell or add to EXTEND.md providers.openai.api_key_env.",
      "openai"
    );
  }

  const size = pickSize(model, args);
  const quality = pickQuality(model, args.quality);

  if (args.referenceImages.length > 0) {
    if (model.includes("dall-e-3") || model.includes("dall-e-2")) {
      throw new ProviderError(
        "Reference images require a gpt-image-* model. Pass --model gpt-image-1 or set OPENAI_IMAGE_MODEL.",
        "openai"
      );
    }
    return generateWithEdits(baseUrl, apiKey, prompt, model, size, quality, args.referenceImages);
  }

  return generateFromText(baseUrl, apiKey, prompt, model, size, quality);
}

async function generateFromText(
  baseUrl: string,
  apiKey: string,
  prompt: string,
  model: string,
  size: string,
  quality: string | null
): Promise<Uint8Array> {
  const body: Record<string, unknown> = { model, prompt, size, n: 1 };
  if (quality) body.quality = quality;

  const res = await fetch(`${baseUrl}/images/generations`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${apiKey}`,
    },
    body: JSON.stringify(body),
  });

  if (!res.ok) {
    const text = await res.text();
    throw new ProviderError(`OpenAI images API ${res.status}: ${text}`, "openai");
  }

  const payload = (await res.json()) as OpenAIImageResponse;
  return extractImage(payload);
}

async function generateWithEdits(
  baseUrl: string,
  apiKey: string,
  prompt: string,
  model: string,
  size: string,
  quality: string | null,
  references: string[]
): Promise<Uint8Array> {
  const form = new FormData();
  form.append("model", model);
  form.append("prompt", prompt);
  form.append("size", size);
  if (quality) form.append("quality", quality);

  for (const refPath of references) {
    const bytes = await readFile(refPath);
    const name = path.basename(refPath);
    form.append("image[]", new Blob([bytes], { type: getMimeType(name) }), name);
  }

  const res = await fetch(`${baseUrl}/images/edits`, {
    method: "POST",
    headers: { Authorization: `Bearer ${apiKey}` },
    body: form,
  });

  if (!res.ok) {
    const text = await res.text();
    throw new ProviderError(`OpenAI edits API ${res.status}: ${text}`, "openai");
  }

  const payload = (await res.json()) as OpenAIImageResponse;
  return extractImage(payload);
}

async function extractImage(result: OpenAIImageResponse): Promise<Uint8Array> {
  const item = result.data?.[0];
  if (item?.b64_json) {
    return Uint8Array.from(Buffer.from(item.b64_json, "base64"));
  }
  if (item?.url) {
    const res = await fetch(item.url);
    if (!res.ok) {
      throw new ProviderError(`Failed to download image from OpenAI URL (${res.status})`, "openai");
    }
    return new Uint8Array(await res.arrayBuffer());
  }
  throw new ProviderError("OpenAI response contained no image data", "openai");
}

import path from "node:path";
import { readFile } from "node:fs/promises";
import type { CliArgs } from "../types";
import { ProviderError } from "../types";
import { getMimeType } from "../utils/output";

const API_BASE = "https://api.replicate.com/v1";

export function defaultModel(): string {
  return process.env.REPLICATE_IMAGE_MODEL || "black-forest-labs/flux-1.1-pro";
}

function requireApiKey(): string {
  const key = process.env.REPLICATE_API_TOKEN;
  if (!key) {
    throw new ProviderError(
      "REPLICATE_API_TOKEN not set. Get one at https://replicate.com/account/api-tokens",
      "replicate"
    );
  }
  return key;
}

async function referenceAsDataUri(refPath: string): Promise<string> {
  const bytes = await readFile(refPath);
  const mime = getMimeType(path.basename(refPath));
  return `data:${mime};base64,${Buffer.from(bytes).toString("base64")}`;
}

type Prediction = {
  id?: string;
  status?: "starting" | "processing" | "succeeded" | "failed" | "canceled";
  output?: string | string[] | null;
  error?: string | null;
  urls?: { get?: string };
};

function buildInput(prompt: string, args: CliArgs, refDataUris: string[]): Record<string, unknown> {
  const input: Record<string, unknown> = { prompt };
  if (args.aspectRatio) input.aspect_ratio = args.aspectRatio;
  if (args.size) {
    const match = args.size.match(/^(\d+)x(\d+)$/i);
    if (match) {
      input.width = Number(match[1]);
      input.height = Number(match[2]);
    }
  }
  input.num_outputs = Math.max(1, args.n);
  if (refDataUris.length > 0) {
    // Most Replicate image models accept a single `image` input; some accept `image_reference` / `image_prompt`.
    // We pass both common names to maximize compatibility.
    input.image = refDataUris[0];
    input.image_prompt = refDataUris[0];
  }
  return input;
}

function parsePredictionOutput(pred: Prediction): string {
  const output = pred.output;
  if (Array.isArray(output) && output.length > 0 && typeof output[0] === "string") {
    return output[0];
  }
  if (typeof output === "string" && output.length > 0) {
    return output;
  }
  throw new ProviderError("Replicate prediction returned no output URL", "replicate");
}

async function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

async function pollPrediction(apiKey: string, pred: Prediction, deadlineMs: number): Promise<Prediction> {
  const getUrl = pred.urls?.get;
  if (!getUrl) return pred;
  let current = pred;
  const start = Date.now();
  while (current.status !== "succeeded" && current.status !== "failed" && current.status !== "canceled") {
    if (Date.now() - start > deadlineMs) break;
    await sleep(1500);
    const r = await fetch(getUrl, { headers: { Authorization: `Bearer ${apiKey}` } });
    if (!r.ok) {
      const text = await r.text();
      throw new ProviderError(`Replicate poll ${r.status}: ${text}`, "replicate");
    }
    current = (await r.json()) as Prediction;
  }
  return current;
}

export async function generate(prompt: string, model: string, args: CliArgs): Promise<Uint8Array> {
  const apiKey = requireApiKey();
  const refs: string[] = [];
  for (const refPath of args.referenceImages) {
    refs.push(await referenceAsDataUri(refPath));
  }

  const [owner, name] = model.split("/");
  if (!owner || !name) {
    throw new ProviderError(
      `Replicate model id must be 'owner/name' (got '${model}'). Example: black-forest-labs/flux-1.1-pro`,
      "replicate"
    );
  }

  const body = { input: buildInput(prompt, args, refs) };
  const createUrl = `${API_BASE}/models/${owner}/${name}/predictions`;
  const res = await fetch(createUrl, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${apiKey}`,
      Prefer: "wait=60",
    },
    body: JSON.stringify(body),
  });

  if (!res.ok) {
    const text = await res.text();
    throw new ProviderError(`Replicate create ${res.status}: ${text}`, "replicate");
  }

  let pred = (await res.json()) as Prediction;
  if (pred.status !== "succeeded" && pred.status !== "failed" && pred.status !== "canceled") {
    pred = await pollPrediction(apiKey, pred, 180_000);
  }
  if (pred.status === "failed" || pred.status === "canceled") {
    throw new ProviderError(`Replicate prediction ${pred.status}: ${pred.error ?? "(no detail)"}`, "replicate");
  }

  const outputUrl = parsePredictionOutput(pred);
  const imgRes = await fetch(outputUrl);
  if (!imgRes.ok) {
    throw new ProviderError(`Failed to download Replicate output (${imgRes.status})`, "replicate");
  }
  return new Uint8Array(await imgRes.arrayBuffer());
}

import type { CliArgs } from "../types";
import { ProviderError } from "../types";
import { orientationOf } from "../utils/aspect_ratio";

const API_BASE =
  process.env.SILICONFLOW_BASE_URL || "https://api.siliconflow.cn/v1";

// Kolors-family defaults. Flux-family on SiliconFlow accepts arbitrary W x H.
const KOLORS_SIZES = {
  square: "1024x1024",
  landscape4x3: "1024x768",
  portrait4x3: "768x1024",
  landscape16x9: "1280x720",
  portrait16x9: "720x1280",
};

export function defaultModel(): string {
  return process.env.SILICONFLOW_IMAGE_MODEL || "Kwai-Kolors/Kolors";
}

function requireApiKey(): string {
  const key = process.env.SILICONFLOW_API_KEY;
  if (!key) {
    throw new ProviderError(
      "SILICONFLOW_API_KEY not set. Get one at https://cloud.siliconflow.cn/account/ak",
      "siliconflow"
    );
  }
  return key;
}

function pickImageSize(args: CliArgs, model: string): string {
  if (args.size) return args.size;
  const ar = args.aspectRatio;
  const isFlux = /flux/i.test(model);
  if (!ar || ar === "1:1") return isFlux ? "1024x1024" : KOLORS_SIZES.square;
  if (ar === "4:3") return KOLORS_SIZES.landscape4x3;
  if (ar === "3:4") return KOLORS_SIZES.portrait4x3;
  const orientation = orientationOf(ar);
  if (orientation === "landscape") return KOLORS_SIZES.landscape16x9;
  if (orientation === "portrait") return KOLORS_SIZES.portrait16x9;
  return KOLORS_SIZES.square;
}

type SiliconFlowResponse = {
  images?: Array<{ url?: string }>;
  data?: Array<{ url?: string; b64_json?: string }>;
  error?: { message?: string };
};

export async function generate(prompt: string, model: string, args: CliArgs): Promise<Uint8Array> {
  const apiKey = requireApiKey();
  if (args.referenceImages.length > 0) {
    throw new ProviderError(
      "SiliconFlow reference-image workflow not wired in MVP. Drop --ref to use text-to-image.",
      "siliconflow"
    );
  }

  const body: Record<string, unknown> = {
    model,
    prompt,
    image_size: pickImageSize(args, model),
    batch_size: Math.max(1, args.n),
    num_inference_steps: 20,
    guidance_scale: 7.5,
  };

  const res = await fetch(`${API_BASE}/images/generations`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${apiKey}`,
    },
    body: JSON.stringify(body),
  });

  if (!res.ok) {
    const text = await res.text();
    throw new ProviderError(`SiliconFlow ${res.status}: ${text}`, "siliconflow");
  }

  const payload = (await res.json()) as SiliconFlowResponse;
  const item = payload.images?.[0] ?? payload.data?.[0];
  if (item?.b64_json) {
    return Uint8Array.from(Buffer.from(item.b64_json, "base64"));
  }
  if (item?.url) {
    const imgRes = await fetch(item.url);
    if (!imgRes.ok) {
      throw new ProviderError(`Failed to download SiliconFlow output (${imgRes.status})`, "siliconflow");
    }
    return new Uint8Array(await imgRes.arrayBuffer());
  }
  throw new ProviderError("SiliconFlow response contained no image", "siliconflow");
}

import type { CliArgs } from "../types";
import { ProviderError } from "../types";

const API_BASE = "https://api.stability.ai/v2beta/stable-image";

const STABILITY_RATIOS = new Set(["21:9", "16:9", "3:2", "5:4", "1:1", "4:5", "2:3", "9:16", "9:21"]);

export function defaultModel(): string {
  return process.env.STABILITY_IMAGE_MODEL || "ultra";
}

function requireApiKey(): string {
  const key = process.env.STABILITY_API_KEY;
  if (!key) {
    throw new ProviderError(
      "STABILITY_API_KEY not set. Get one at https://platform.stability.ai/account/keys",
      "stability"
    );
  }
  return key;
}

function pickEndpoint(model: string): string {
  const normalized = model.trim().toLowerCase();
  const known: Record<string, string> = {
    ultra: `${API_BASE}/generate/ultra`,
    core: `${API_BASE}/generate/core`,
    "sd3": `${API_BASE}/generate/sd3`,
    "stable-image-ultra": `${API_BASE}/generate/ultra`,
    "stable-image-core": `${API_BASE}/generate/core`,
    "stable-diffusion-3": `${API_BASE}/generate/sd3`,
  };
  return known[normalized] ?? `${API_BASE}/generate/${normalized}`;
}

function pickAspect(ar: string | null): string {
  if (!ar) return "1:1";
  if (STABILITY_RATIOS.has(ar)) return ar;
  return "1:1";
}

export async function generate(prompt: string, model: string, args: CliArgs): Promise<Uint8Array> {
  const apiKey = requireApiKey();
  if (args.referenceImages.length > 0) {
    throw new ProviderError(
      "Stability provider in this release only supports text-to-image. Reference image workflows require /control/sketch or /control/structure endpoints.",
      "stability"
    );
  }

  const form = new FormData();
  form.append("prompt", prompt);
  form.append("aspect_ratio", pickAspect(args.aspectRatio));
  form.append("output_format", "png");

  const res = await fetch(pickEndpoint(model), {
    method: "POST",
    headers: {
      Authorization: `Bearer ${apiKey}`,
      Accept: "image/*",
    },
    body: form,
  });

  if (!res.ok) {
    const text = await res.text();
    throw new ProviderError(`Stability ${res.status}: ${text}`, "stability");
  }

  const buf = await res.arrayBuffer();
  return new Uint8Array(buf);
}

export type Provider =
  | "openai"
  | "google"
  | "replicate"
  | "stability"
  | "fal"
  | "ark"
  | "bailian"
  | "siliconflow";

export type Quality = "draft" | "hd" | "ultra";

export type CliArgs = {
  prompt: string | null;
  promptFiles: string[];
  imagePath: string | null;
  provider: Provider | null;
  model: string | null;
  aspectRatio: string | null;
  size: string | null;
  quality: Quality | null;
  referenceImages: string[];
  n: number;
  json: boolean;
  setup: boolean;
  help: boolean;
};

export type ExtendConfig = {
  version?: number;
  default_provider?: Provider;
  default_quality?: Quality;
  default_aspect_ratio?: string;
  default_save_dir?: string;
  default_model?: Partial<Record<Provider, string>>;
  providers?: Partial<
    Record<
      Provider,
      {
        api_key_env?: string;
        api_key_source?: string;
        base_url?: string;
      }
    >
  >;
};

export type ProviderModule = {
  defaultModel: () => string;
  generate: (prompt: string, model: string, args: CliArgs) => Promise<Uint8Array>;
};

export class ConfigError extends Error {
  code = "CONFIG_ERROR";
}

export class ProviderError extends Error {
  constructor(message: string, public provider: Provider) {
    super(message);
  }
  code = "PROVIDER_ERROR";
}

export type Orientation = "square" | "landscape" | "portrait";

export function parseAspectRatio(ar: string): { width: number; height: number } | null {
  const match = ar.match(/^(\d+(?:\.\d+)?):(\d+(?:\.\d+)?)$/);
  if (!match) return null;
  const w = parseFloat(match[1]!);
  const h = parseFloat(match[2]!);
  if (!(w > 0) || !(h > 0)) return null;
  return { width: w, height: h };
}

export function orientationOf(ar: string | null): Orientation {
  if (!ar) return "square";
  const parsed = parseAspectRatio(ar);
  if (!parsed) return "square";
  const ratio = parsed.width / parsed.height;
  if (ratio > 1.15) return "landscape";
  if (ratio < 0.87) return "portrait";
  return "square";
}

import path from "node:path";
import { mkdir, writeFile } from "node:fs/promises";

export async function writeBytesToFile(bytes: Uint8Array, outputPath: string): Promise<void> {
  const dir = path.dirname(outputPath);
  await mkdir(dir, { recursive: true });
  await writeFile(outputPath, bytes);
}

export function getMimeType(filename: string): string {
  const ext = path.extname(filename).toLowerCase();
  if (ext === ".jpg" || ext === ".jpeg") return "image/jpeg";
  if (ext === ".webp") return "image/webp";
  if (ext === ".gif") return "image/gif";
  return "image/png";
}

export function validateImageMagic(bytes: Uint8Array): "png" | "jpeg" | "webp" | "gif" | "unknown" {
  if (bytes.length < 12) return "unknown";
  if (bytes[0] === 0x89 && bytes[1] === 0x50 && bytes[2] === 0x4e && bytes[3] === 0x47) return "png";
  if (bytes[0] === 0xff && bytes[1] === 0xd8) return "jpeg";
  if (
    bytes[0] === 0x52 && bytes[1] === 0x49 && bytes[2] === 0x46 && bytes[3] === 0x46 &&
    bytes[8] === 0x57 && bytes[9] === 0x45 && bytes[10] === 0x42 && bytes[11] === 0x50
  ) {
    return "webp";
  }
  if (bytes[0] === 0x47 && bytes[1] === 0x49 && bytes[2] === 0x46) return "gif";
  return "unknown";
}