@ai-sdk/replicate 0.0.0-70e0935a-20260114150030 → 0.0.0-98261322-20260122142521

package/src/replicate-image-model.ts

@@ -0,0 +1,268 @@
+import type { ImageModelV3, SharedV3Warning } from '@ai-sdk/provider';
+import type { Resolvable } from '@ai-sdk/provider-utils';
+import {
+  combineHeaders,
+  convertImageModelFileToDataUri,
+  createBinaryResponseHandler,
+  createJsonResponseHandler,
+  FetchFunction,
+  getFromApi,
+  InferSchema,
+  lazySchema,
+  parseProviderOptions,
+  postJsonToApi,
+  resolve,
+  zodSchema,
+} from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+import { replicateFailedResponseHandler } from './replicate-error';
+import { ReplicateImageModelId } from './replicate-image-settings';
+
+interface ReplicateImageModelConfig {
+  provider: string;
+  baseURL: string;
+  headers?: Resolvable<Record<string, string | undefined>>;
+  fetch?: FetchFunction;
+  _internal?: {
+    currentDate?: () => Date;
+  };
+}
+
+// Flux-2 models support up to 8 input images with input_image, input_image_2, etc.
+const FLUX_2_MODEL_PATTERN = /^black-forest-labs\/flux-2-/;
+const MAX_FLUX_2_INPUT_IMAGES = 8;
+
+export class ReplicateImageModel implements ImageModelV3 {
+  readonly specificationVersion = 'v3';
+
+  get maxImagesPerCall(): number {
+    // Flux-2 models support up to 8 input images
+    return this.isFlux2Model ? MAX_FLUX_2_INPUT_IMAGES : 1;
+  }
+
+  get provider(): string {
+    return this.config.provider;
+  }
+
+  private get isFlux2Model(): boolean {
+    return FLUX_2_MODEL_PATTERN.test(this.modelId);
+  }
+
+  constructor(
+    readonly modelId: ReplicateImageModelId,
+    private readonly config: ReplicateImageModelConfig,
+  ) {}
+
+  async doGenerate({
+    prompt,
+    n,
+    aspectRatio,
+    size,
+    seed,
+    providerOptions,
+    headers,
+    abortSignal,
+    files,
+    mask,
+  }: Parameters<ImageModelV3['doGenerate']>[0]): Promise<
+    Awaited<ReturnType<ImageModelV3['doGenerate']>>
+  > {
+    const warnings: Array<SharedV3Warning> = [];
+
+    const [modelId, version] = this.modelId.split(':');
+
+    const currentDate = this.config._internal?.currentDate?.() ?? new Date();
+
+    // Parse provider options
+    const replicateOptions = await parseProviderOptions({
+      provider: 'replicate',
+      providerOptions,
+      schema: replicateImageProviderOptionsSchema,
+    });
+
+    // Handle image input from files
+    let imageInputs: Record<string, string> = {};
+    if (files != null && files.length > 0) {
+      if (this.isFlux2Model) {
+        // Flux-2 models use input_image, input_image_2, input_image_3, etc.
+        for (
+          let i = 0;
+          i < Math.min(files.length, MAX_FLUX_2_INPUT_IMAGES);
+          i++
+        ) {
+          const key = i === 0 ? 'input_image' : `input_image_${i + 1}`;
+          imageInputs[key] = convertImageModelFileToDataUri(files[i]);
+        }
+        if (files.length > MAX_FLUX_2_INPUT_IMAGES) {
+          warnings.push({
+            type: 'other',
+            message: `Flux-2 models support up to ${MAX_FLUX_2_INPUT_IMAGES} input images. Additional images are ignored.`,
+          });
+        }
+      } else {
+        // Other models use single 'image' parameter
+        imageInputs = { image: convertImageModelFileToDataUri(files[0]) };
+        if (files.length > 1) {
+          warnings.push({
+            type: 'other',
+            message:
+              'This Replicate model only supports a single input image. Additional images are ignored.',
+          });
+        }
+      }
+    }
+
+    // Handle mask input (not supported by Flux-2 models)
+    let maskInput: string | undefined;
+    if (mask != null) {
+      if (this.isFlux2Model) {
+        warnings.push({
+          type: 'other',
+          message:
+            'Flux-2 models do not support mask input. The mask will be ignored.',
+        });
+      } else {
+        maskInput = convertImageModelFileToDataUri(mask);
+      }
+    }
+
+    // Extract maxWaitTimeInSeconds from provider options and prepare the rest for the request body
+    const { maxWaitTimeInSeconds, ...inputOptions } = replicateOptions ?? {};
+
+    // Build the prefer header based on maxWaitTimeInSeconds:
+    // - undefined/null: use default sync wait (prefer: wait)
+    // - positive number: use custom wait duration (prefer: wait=N)
+    const preferHeader: Record<string, string> =
+      maxWaitTimeInSeconds != null
+        ? { prefer: `wait=${maxWaitTimeInSeconds}` }
+        : { prefer: 'wait' };
+
+    const {
+      value: { output },
+      responseHeaders,
+    } = await postJsonToApi({
+      url:
+        // different endpoints for versioned vs unversioned models:
+        version != null
+          ? `${this.config.baseURL}/predictions`
+          : `${this.config.baseURL}/models/${modelId}/predictions`,
+
+      headers: combineHeaders(
+        await resolve(this.config.headers),
+        headers,
+        preferHeader,
+      ),
+
+      body: {
+        input: {
+          prompt,
+          aspect_ratio: aspectRatio,
+          size,
+          seed,
+          num_outputs: n,
+          ...imageInputs,
+          ...(maskInput != null ? { mask: maskInput } : {}),
+          ...inputOptions,
+        },
+        // for versioned models, include the version in the body:
+        ...(version != null ? { version } : {}),
+      },
+
+      successfulResponseHandler: createJsonResponseHandler(
+        replicateImageResponseSchema,
+      ),
+      failedResponseHandler: replicateFailedResponseHandler,
+      abortSignal,
+      fetch: this.config.fetch,
+    });
+
+    // download the images:
+    const outputArray = Array.isArray(output) ? output : [output];
+    const images = await Promise.all(
+      outputArray.map(async url => {
+        const { value: image } = await getFromApi({
+          url,
+          successfulResponseHandler: createBinaryResponseHandler(),
+          failedResponseHandler: replicateFailedResponseHandler,
+          abortSignal,
+          fetch: this.config.fetch,
+        });
+        return image;
+      }),
+    );
+
+    return {
+      images,
+      warnings,
+      response: {
+        timestamp: currentDate,
+        modelId: this.modelId,
+        headers: responseHeaders,
+      },
+    };
+  }
+}
+
+const replicateImageResponseSchema = z.object({
+  output: z.union([z.array(z.string()), z.string()]),
+});
+
+/**
+ * Provider options schema for Replicate image generation.
+ *
+ * Note: Different Replicate models support different parameters.
+ * This schema includes common parameters, but you can pass any
+ * model-specific parameters through the passthrough.
+ */
+export const replicateImageProviderOptionsSchema = lazySchema(() =>
+  zodSchema(
+    z
+      .object({
+        /**
+         * Maximum time in seconds to wait for the prediction to complete in sync mode.
+         * By default, Replicate uses sync mode with a 60-second timeout.
+         *
+         * - When not specified: Uses default 60-second sync wait (`prefer: wait`)
+         * - When set to a positive number: Uses that duration (`prefer: wait=N`)
+         */
+        maxWaitTimeInSeconds: z.number().positive().nullish(),
+
+        /**
+         * Guidance scale for classifier-free guidance.
+         * Higher values make the output more closely match the prompt.
+         */
+        guidance_scale: z.number().nullish(),
+
+        /**
+         * Number of denoising steps. More steps = higher quality but slower.
+         */
+        num_inference_steps: z.number().nullish(),
+
+        /**
+         * Negative prompt to guide what to avoid in the generation.
+         */
+        negative_prompt: z.string().nullish(),
+
+        /**
+         * Output image format.
+         */
+        output_format: z.enum(['png', 'jpg', 'webp']).nullish(),
+
+        /**
+         * Output image quality (1-100). Only applies to jpg and webp.
+         */
+        output_quality: z.number().min(1).max(100).nullish(),
+
+        /**
+         * Strength of the transformation for img2img (0-1).
+         * Lower values keep more of the original image.
+         */
+        strength: z.number().min(0).max(1).nullish(),
+      })
+      .passthrough(),
+  ),
+);
+
+export type ReplicateImageProviderOptions = InferSchema<
+  typeof replicateImageProviderOptionsSchema
+>;

package/src/replicate-image-settings.ts

@@ -0,0 +1,36 @@
+export type ReplicateImageModelId =
+  // Text-to-image models
+  | 'black-forest-labs/flux-1.1-pro'
+  | 'black-forest-labs/flux-1.1-pro-ultra'
+  | 'black-forest-labs/flux-dev'
+  | 'black-forest-labs/flux-pro'
+  | 'black-forest-labs/flux-schnell'
+  | 'bytedance/sdxl-lightning-4step'
+  | 'fofr/aura-flow'
+  | 'fofr/latent-consistency-model'
+  | 'fofr/realvisxl-v3-multi-controlnet-lora'
+  | 'fofr/sdxl-emoji'
+  | 'fofr/sdxl-multi-controlnet-lora'
+  | 'ideogram-ai/ideogram-v2'
+  | 'ideogram-ai/ideogram-v2-turbo'
+  | 'lucataco/dreamshaper-xl-turbo'
+  | 'lucataco/open-dalle-v1.1'
+  | 'lucataco/realvisxl-v2.0'
+  | 'lucataco/realvisxl2-lcm'
+  | 'luma/photon'
+  | 'luma/photon-flash'
+  | 'nvidia/sana'
+  | 'playgroundai/playground-v2.5-1024px-aesthetic'
+  | 'recraft-ai/recraft-v3'
+  | 'recraft-ai/recraft-v3-svg'
+  | 'stability-ai/stable-diffusion-3.5-large'
+  | 'stability-ai/stable-diffusion-3.5-large-turbo'
+  | 'stability-ai/stable-diffusion-3.5-medium'
+  | 'tstramer/material-diffusion'
+  // Inpainting and image editing models
+  | 'black-forest-labs/flux-fill-pro'
+  | 'black-forest-labs/flux-fill-dev'
+  // Flux-2 models (support up to 8 reference images)
+  | 'black-forest-labs/flux-2-pro'
+  | 'black-forest-labs/flux-2-dev'
+  | (string & {});

package/src/replicate-provider.test.ts

@@ -0,0 +1,24 @@
+import { describe, it, expect } from 'vitest';
+import { createReplicate } from './replicate-provider';
+import { ReplicateImageModel } from './replicate-image-model';
+
+describe('createReplicate', () => {
+  it('creates a provider with required settings', () => {
+    const provider = createReplicate({ apiToken: 'test-token' });
+    expect(provider.image).toBeDefined();
+  });
+
+  it('creates a provider with custom settings', () => {
+    const provider = createReplicate({
+      apiToken: 'test-token',
+      baseURL: 'https://custom.replicate.com',
+    });
+    expect(provider.image).toBeDefined();
+  });
+
+  it('creates an image model instance', () => {
+    const provider = createReplicate({ apiToken: 'test-token' });
+    const model = provider.image('black-forest-labs/flux-schnell');
+    expect(model).toBeInstanceOf(ReplicateImageModel);
+  });
+});

package/src/replicate-provider.ts

@@ -0,0 +1,99 @@
+import { NoSuchModelError, ProviderV3 } from '@ai-sdk/provider';
+import type { FetchFunction } from '@ai-sdk/provider-utils';
+import { loadApiKey, withUserAgentSuffix } from '@ai-sdk/provider-utils';
+import { ReplicateImageModel } from './replicate-image-model';
+import { ReplicateImageModelId } from './replicate-image-settings';
+import { VERSION } from './version';
+
+export interface ReplicateProviderSettings {
+  /**
+API token that is being send using the `Authorization` header.
+It defaults to the `REPLICATE_API_TOKEN` environment variable.
+   */
+  apiToken?: string;
+
+  /**
+Use a different URL prefix for API calls, e.g. to use proxy servers.
+The default prefix is `https://api.replicate.com/v1`.
+   */
+  baseURL?: string;
+
+  /**
+Custom headers to include in the requests.
+     */
+  headers?: Record<string, string>;
+
+  /**
+Custom fetch implementation. You can use it as a middleware to intercept requests,
+or to provide a custom fetch implementation for e.g. testing.
+    */
+  fetch?: FetchFunction;
+}
+
+export interface ReplicateProvider extends ProviderV3 {
+  /**
+   * Creates a Replicate image generation model.
+   */
+  image(modelId: ReplicateImageModelId): ReplicateImageModel;
+
+  /**
+   * Creates a Replicate image generation model.
+   */
+  imageModel(modelId: ReplicateImageModelId): ReplicateImageModel;
+
+  /**
+   * @deprecated Use `embeddingModel` instead.
+   */
+  textEmbeddingModel(modelId: string): never;
+}
+
+/**
+ * Create a Replicate provider instance.
+ */
+export function createReplicate(
+  options: ReplicateProviderSettings = {},
+): ReplicateProvider {
+  const createImageModel = (modelId: ReplicateImageModelId) =>
+    new ReplicateImageModel(modelId, {
+      provider: 'replicate',
+      baseURL: options.baseURL ?? 'https://api.replicate.com/v1',
+      headers: withUserAgentSuffix(
+        {
+          Authorization: `Bearer ${loadApiKey({
+            apiKey: options.apiToken,
+            environmentVariableName: 'REPLICATE_API_TOKEN',
+            description: 'Replicate',
+          })}`,
+          ...options.headers,
+        },
+        `ai-sdk/replicate/${VERSION}`,
+      ),
+      fetch: options.fetch,
+    });
+
+  const embeddingModel = (modelId: string) => {
+    throw new NoSuchModelError({
+      modelId,
+      modelType: 'embeddingModel',
+    });
+  };
+
+  return {
+    specificationVersion: 'v3' as const,
+    image: createImageModel,
+    imageModel: createImageModel,
+    languageModel: (modelId: string) => {
+      throw new NoSuchModelError({
+        modelId,
+        modelType: 'languageModel',
+      });
+    },
+    embeddingModel,
+    textEmbeddingModel: embeddingModel,
+  };
+}
+
+/**
+ * Default Replicate provider instance.
+ */
+export const replicate = createReplicate();

package/src/version.ts

@@ -0,0 +1,6 @@
+// Version string of this package injected at build time.
+declare const __PACKAGE_VERSION__: string | undefined;
+export const VERSION: string =
+  typeof __PACKAGE_VERSION__ !== 'undefined'
+    ? __PACKAGE_VERSION__
+    : '0.0.0-test';