Stability AI (Stable Diffusion)

STABILITY_AI_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
STABILITY_AI_BASE_URL=https://api.stability.ai
STABILITY_AI_DEFAULT_ENGINE=stable-image-core   # See model options below

import axios from 'axios';
import FormData from 'form-data';
 
class StabilityAITool implements ImageGenerationTool {
  readonly name = 'stability-ai';
 
  constructor(
    private apiKey:  string,
    private baseUrl: string = 'https://api.stability.ai',
  ) {}
 
  private headers(contentType?: string) {
    return {
      Authorization:  `Bearer ${this.apiKey}`,
      Accept:         'application/json',
      ...(contentType ? { 'Content-Type': contentType } : {}),
    };
  }
 
  // ── Text-to-image (Stable Image Core / SD3) ──────────────────────────────────
 
  async generateImage(options: {
    prompt:          string;
    negativePrompt?: string;
    aspectRatio?:    '16:9' | '1:1' | '21:9' | '2:3' | '3:2' | '4:5' | '5:4' | '9:16' | '9:21';
    model?:          'core' | 'sd3-large' | 'sd3-medium' | 'sd3-large-turbo';
    style?:          string;  // Stable Image Core style presets — see below
    seed?:           number;  // 0 = random
    outputFormat?:   'jpeg' | 'png' | 'webp';
  }): Promise<GeneratedImage[]> {
    const model      = options.model ?? 'core';
    const endpoint   = model === 'core'
      ? `/v2beta/stable-image/generate/core`
      : `/v2beta/stable-image/generate/sd3`;
 
    const form = new FormData();
    form.append('prompt',         options.prompt);
    form.append('aspect_ratio',   options.aspectRatio  ?? '16:9');
    form.append('output_format',  options.outputFormat ?? 'jpeg');
    if (options.negativePrompt) form.append('negative_prompt', options.negativePrompt);
    if (options.seed !== undefined) form.append('seed', String(options.seed));
    if (model === 'core' && options.style) form.append('style_preset', options.style);
    if (model !== 'core') form.append('model', model);
 
    const response = await axios.post(`${this.baseUrl}${endpoint}`, form, {
      headers: {
        ...this.headers(),
        ...form.getHeaders(),
        Accept: 'image/*',   // Stability returns raw image bytes
      },
      responseType: 'arraybuffer',
    });
 
    const imageBase64 = Buffer.from(response.data).toString('base64');
    const mimeType    = response.headers['content-type'] ?? 'image/jpeg';
 
    return [{
      provider:      'stability-ai',
      model:         endpoint.includes('sd3') ? `sd3-${model}` : 'stable-image-core',
      imageBase64,
      mimeType,
      revisedPrompt: null,
    }];
  }
 
  // ── Image-to-image ───────────────────────────────────────────────────────────
 
  async imageToImage(options: {
    prompt:          string;
    image:           Buffer;
    strength?:       number;   // 0–1: how much to deviate from input image (default 0.7)
    negativePrompt?: string;
    outputFormat?:   'jpeg' | 'png' | 'webp';
  }): Promise<GeneratedImage[]> {
    const form = new FormData();
    form.append('prompt',        options.prompt);
    form.append('image',         options.image, { filename: 'input.jpg', contentType: 'image/jpeg' });
    form.append('strength',      String(options.strength ?? 0.7));
    form.append('output_format', options.outputFormat ?? 'jpeg');
    if (options.negativePrompt) form.append('negative_prompt', options.negativePrompt);
 
    const response = await axios.post(
      `${this.baseUrl}/v2beta/stable-image/generate/sd3`,
      form,
      {
        headers: { ...this.headers(), ...form.getHeaders(), Accept: 'image/*' },
        responseType: 'arraybuffer',
      },
    );
 
    return [{
      provider:      'stability-ai',
      model:         'sd3-image-to-image',
      imageBase64:   Buffer.from(response.data).toString('base64'),
      mimeType:      response.headers['content-type'] ?? 'image/jpeg',
      revisedPrompt: null,
    }];
  }
 
  // ── Background removal ───────────────────────────────────────────────────────
 
  async removeBackground(imageBuffer: Buffer): Promise<Buffer> {
    const form = new FormData();
    form.append('image', imageBuffer, { filename: 'input.png', contentType: 'image/png' });
    form.append('output_format', 'png');
 
    const response = await axios.post(
      `${this.baseUrl}/v2beta/stable-image/edit/remove-background`,
      form,
      {
        headers: { ...this.headers(), ...form.getHeaders(), Accept: 'image/*' },
        responseType: 'arraybuffer',
      },
    );
    return Buffer.from(response.data);
  }
 
  // ── Upscale ──────────────────────────────────────────────────────────────────
 
  async upscale(options: {
    image:  Buffer;
    prompt: string;   // Description of the image for guided upscaling
  }): Promise<Buffer> {
    const form = new FormData();
    form.append('image',  options.image, { filename: 'input.jpg', contentType: 'image/jpeg' });
    form.append('prompt', options.prompt);
    form.append('output_format', 'jpeg');
 
    const response = await axios.post(
      `${this.baseUrl}/v2beta/stable-image/upscale/conservative`,
      form,
      {
        headers: { ...this.headers(), ...form.getHeaders(), Accept: 'image/*' },
        responseType: 'arraybuffer',
      },
    );
    return Buffer.from(response.data);
  }
 
  async verify(): Promise<void> {
    // Check account balance / API key validity
    const response = await axios.get(`${this.baseUrl}/v1/user/account`, {
      headers: this.headers('application/json'),
    });
    if (!response.data.email) throw new Error('Stability AI verification failed');
  }
}

// packages/tools/src/image-generation.ts
function resolveImageProvider(options: {
  useCase:    'blog_hero' | 'social_post' | 'illustration' | 'product';
  tenantTier: 'free' | 'pro' | 'agency' | 'enterprise';
  preferredProvider?: string;
}): ImageGenerationTool {
  const preferred = options.preferredProvider ?? config.DEFAULT_IMAGE_PROVIDER;
 
  switch (preferred) {
    case 'stability-ai': return new StabilityAITool(config.STABILITY_AI_API_KEY);
    case 'dall-e':       return new DallETool(config.OPENAI_API_KEY);
    case 'gemini':       return new GeminiImageTool(config.GOOGLE_AI_API_KEY);
    default:
      // Cost-based default: Stability AI for bulk; DALL-E for premium
      if (options.tenantTier === 'free') return new StabilityAITool(config.STABILITY_AI_API_KEY);
      return new DallETool(config.OPENAI_API_KEY);
  }
}