feat (provider): Add maxImagesPerCall setting to all image models/providers. (#4364)

Author: shaper

.changeset/nervous-gifts-bake.md

@@ -0,0 +1,8 @@
+---
+'@ai-sdk/google-vertex': patch
+'@ai-sdk/fireworks': patch
+'@ai-sdk/replicate': patch
+'@ai-sdk/openai': patch
+---
+
+feat (provider): Add maxImagesPerCall setting to all image providers.

content/docs/03-ai-sdk-core/35-image-generation.mdx

@@ -83,6 +83,22 @@ const { images } = await generateImage({
   parallel) to generate the requested number of images.
 </Note>
 
+Each image model has an internal limit on how many images it can generate in a single API call. The AI SDK manages this automatically by batching requests appropriately when you request multiple images using the `n` parameter. By default, the SDK uses provider-documented limits (for example, DALL-E 3 can only generate 1 image per call, while DALL-E 2 supports up to 10).
+
+If needed, you can override this behavior using the `maxImagesPerCall` setting when configuring your model. This is particularly useful when working with new or custom models where the default batch size might not be optimal:
+
+```tsx
+const model = openai.image('dall-e-2', {
+  maxImagesPerCall: 5, // Override the default batch size
+});
+
+const { images } = await generateImage({
+  model,
+  prompt: 'Santa Claus driving a Cadillac',
+  n: 10, // Will make 2 calls of 5 images each
+});
+```
+
 ### Providing a Seed
 
 You can provide a seed to the `generateImage` function to control the output of the image generation process.

packages/fireworks/src/fireworks-image-model.test.ts

@@ -2,27 +2,35 @@ import { FetchFunction } from '@ai-sdk/provider-utils';
 import { createTestServer } from '@ai-sdk/provider-utils/test';
 import { describe, expect, it } from 'vitest';
 import { FireworksImageModel } from './fireworks-image-model';
+import { FireworksImageSettings } from './fireworks-image-settings';
 
 const prompt = 'A cute baby sea otter';
 
 function createBasicModel({
   headers,
   fetch,
+  settings,
 }: {
   headers?: () => Record<string, string>;
   fetch?: FetchFunction;
+  settings?: FireworksImageSettings;
 } = {}) {
-  return new FireworksImageModel('accounts/fireworks/models/flux-1-dev-fp8', {
-    provider: 'fireworks',
-    baseURL: 'https://api.example.com',
-    headers: headers ?? (() => ({ 'api-key': 'test-key' })),
-    fetch,
-  });
+  return new FireworksImageModel(
+    'accounts/fireworks/models/flux-1-dev-fp8',
+    settings ?? {},
+    {
+      provider: 'fireworks',
+      baseURL: 'https://api.example.com',
+      headers: headers ?? (() => ({ 'api-key': 'test-key' })),
+      fetch,
+    },
+  );
 }
 
 function createSizeModel() {
   return new FireworksImageModel(
     'accounts/fireworks/models/playground-v2-5-1024px-aesthetic',
+    {},
     {
       provider: 'fireworks',
       baseURL: 'https://api.size-example.com',
@@ -298,5 +306,21 @@ describe('FireworksImageModel', () => {
       expect(model.specificationVersion).toBe('v1');
       expect(model.maxImagesPerCall).toBe(1);
     });
+
+    it('should use maxImagesPerCall from settings', () => {
+      const model = createBasicModel({
+        settings: {
+          maxImagesPerCall: 4,
+        },
+      });
+
+      expect(model.maxImagesPerCall).toBe(4);
+    });
+
+    it('should default maxImagesPerCall to 1 when not specified', () => {
+      const model = createBasicModel();
+
+      expect(model.maxImagesPerCall).toBe(1);
+    });
   });
 });

packages/fireworks/src/fireworks-image-model.ts

@@ -10,17 +10,10 @@ import {
   postJsonToApi,
   ResponseHandler,
 } from '@ai-sdk/provider-utils';
-
-// https://fireworks.ai/models?type=image
-export type FireworksImageModelId =
-  | 'accounts/fireworks/models/flux-1-dev-fp8'
-  | 'accounts/fireworks/models/flux-1-schnell-fp8'
-  | 'accounts/fireworks/models/playground-v2-5-1024px-aesthetic'
-  | 'accounts/fireworks/models/japanese-stable-diffusion-xl'
-  | 'accounts/fireworks/models/playground-v2-1024px-aesthetic'
-  | 'accounts/fireworks/models/SSD-1B'
-  | 'accounts/fireworks/models/stable-diffusion-xl-1024-v1-0'
-  | (string & {});
+import {
+  FireworksImageModelId,
+  FireworksImageSettings,
+} from './fireworks-image-settings';
 
 interface FireworksImageModelBackendConfig {
   urlFormat: 'workflows' | 'image_generation';
@@ -141,10 +134,13 @@ export class FireworksImageModel implements ImageModelV1 {
     return this.config.provider;
   }
 
-  readonly maxImagesPerCall = 1;
+  get maxImagesPerCall(): number {
+    return this.settings.maxImagesPerCall ?? 1;
+  }
 
   constructor(
     readonly modelId: FireworksImageModelId,
+    readonly settings: FireworksImageSettings,
     private config: FireworksImageModelConfig,
   ) {}
 

packages/fireworks/src/fireworks-image-settings.ts

@@ -0,0 +1,17 @@
+// https://fireworks.ai/models?type=image
+export type FireworksImageModelId =
+  | 'accounts/fireworks/models/flux-1-dev-fp8'
+  | 'accounts/fireworks/models/flux-1-schnell-fp8'
+  | 'accounts/fireworks/models/playground-v2-5-1024px-aesthetic'
+  | 'accounts/fireworks/models/japanese-stable-diffusion-xl'
+  | 'accounts/fireworks/models/playground-v2-1024px-aesthetic'
+  | 'accounts/fireworks/models/SSD-1B'
+  | 'accounts/fireworks/models/stable-diffusion-xl-1024-v1-0'
+  | (string & {});
+
+export interface FireworksImageSettings {
+  /**
+Override the maximum number of images per call (default 1)
+   */
+  maxImagesPerCall?: number;
+}

packages/fireworks/src/fireworks-provider.test.ts

@@ -1,13 +1,13 @@
 import { describe, it, expect, vi, beforeEach, Mock } from 'vitest';
 import { createFireworks } from './fireworks-provider';
-import { createOpenAICompatible } from '@ai-sdk/openai-compatible';
 import { LanguageModelV1, EmbeddingModelV1 } from '@ai-sdk/provider';
 import { loadApiKey } from '@ai-sdk/provider-utils';
 import {
   OpenAICompatibleChatLanguageModel,
   OpenAICompatibleCompletionLanguageModel,
   OpenAICompatibleEmbeddingModel,
 } from '@ai-sdk/openai-compatible';
+import { FireworksImageModel } from './fireworks-image-model';
 
 // Add type assertion for the mocked class
 const OpenAICompatibleChatLanguageModelMock =
@@ -24,6 +24,10 @@ vi.mock('@ai-sdk/provider-utils', () => ({
   withoutTrailingSlash: vi.fn(url => url),
 }));
 
+vi.mock('./fireworks-image-model', () => ({
+  FireworksImageModel: vi.fn(),
+}));
+
 describe('FireworksProvider', () => {
   let mockLanguageModel: LanguageModelV1;
   let mockEmbeddingModel: EmbeddingModelV1<string>;
@@ -125,4 +129,54 @@ describe('FireworksProvider', () => {
       expect(model).toBeInstanceOf(OpenAICompatibleEmbeddingModel);
     });
   });
+
+  describe('image', () => {
+    it('should construct an image model with correct configuration', () => {
+      const provider = createFireworks();
+      const modelId = 'accounts/fireworks/models/flux-1-dev-fp8';
+      const settings = { maxImagesPerCall: 2 };
+
+      const model = provider.image(modelId, settings);
+
+      expect(model).toBeInstanceOf(FireworksImageModel);
+      expect(FireworksImageModel).toHaveBeenCalledWith(
+        modelId,
+        settings,
+        expect.objectContaining({
+          provider: 'fireworks.image',
+          baseURL: 'https://api.fireworks.ai/inference/v1',
+        }),
+      );
+    });
+
+    it('should use default settings when none provided', () => {
+      const provider = createFireworks();
+      const modelId = 'accounts/fireworks/models/flux-1-dev-fp8';
+
+      const model = provider.image(modelId);
+
+      expect(model).toBeInstanceOf(FireworksImageModel);
+      expect(FireworksImageModel).toHaveBeenCalledWith(
+        modelId,
+        {},
+        expect.any(Object),
+      );
+    });
+
+    it('should respect custom baseURL', () => {
+      const customBaseURL = 'https://custom.api.fireworks.ai';
+      const provider = createFireworks({ baseURL: customBaseURL });
+      const modelId = 'accounts/fireworks/models/flux-1-dev-fp8';
+
+      const model = provider.image(modelId);
+
+      expect(FireworksImageModel).toHaveBeenCalledWith(
+        modelId,
+        expect.any(Object),
+        expect.objectContaining({
+          baseURL: customBaseURL,
+        }),
+      );
+    });
+  });
 });

packages/fireworks/src/fireworks-provider.ts

@@ -27,10 +27,11 @@ import {
   FireworksEmbeddingModelId,
   FireworksEmbeddingSettings,
 } from './fireworks-embedding-settings';
+import { FireworksImageModel } from './fireworks-image-model';
 import {
-  FireworksImageModel,
   FireworksImageModelId,
-} from './fireworks-image-model';
+  FireworksImageSettings,
+} from './fireworks-image-settings';
 
 export type FireworksErrorData = z.infer<typeof fireworksErrorSchema>;
 
@@ -100,7 +101,10 @@ Creates a text embedding model for text generation.
   /**
 Creates a model for image generation.
 */
-  image(modelId: FireworksImageModelId): ImageModelV1;
+  image(
+    modelId: FireworksImageModelId,
+    settings?: FireworksImageSettings,
+  ): ImageModelV1;
 }
 
 const defaultBaseURL = 'https://api.fireworks.ai/inference/v1';
@@ -161,8 +165,11 @@ export function createFireworks(
       errorStructure: fireworksErrorStructure,
     });
 
-  const createImageModel = (modelId: FireworksImageModelId) =>
-    new FireworksImageModel(modelId, {
+  const createImageModel = (
+    modelId: FireworksImageModelId,
+    settings: FireworksImageSettings = {},
+  ) =>
+    new FireworksImageModel(modelId, settings, {
       ...getCommonModelConfig('image'),
       baseURL: baseURL ?? defaultBaseURL,
     });

packages/google-vertex/src/google-vertex-image-model.test.ts

@@ -4,11 +4,15 @@ import { GoogleVertexImageModel } from './google-vertex-image-model';
 
 const prompt = 'A cute baby sea otter';
 
-const model = new GoogleVertexImageModel('imagen-3.0-generate-001', {
-  provider: 'google-vertex',
-  baseURL: 'https://api.example.com',
-  headers: { 'api-key': 'test-key' },
-});
+const model = new GoogleVertexImageModel(
+  'imagen-3.0-generate-001',
+  {},
+  {
+    provider: 'google-vertex',
+    baseURL: 'https://api.example.com',
+    headers: { 'api-key': 'test-key' },
+  },
+);
 
 describe('GoogleVertexImageModel', () => {
   describe('doGenerate', () => {
@@ -53,6 +57,7 @@ describe('GoogleVertexImageModel', () => {
 
       const modelWithHeaders = new GoogleVertexImageModel(
         'imagen-3.0-generate-001',
+        {},
         {
           provider: 'google-vertex',
           baseURL: 'https://api.example.com',
@@ -83,6 +88,34 @@ describe('GoogleVertexImageModel', () => {
       });
     });
 
+    it('should respect maxImagesPerCall setting', () => {
+      const customModel = new GoogleVertexImageModel(
+        'imagen-3.0-generate-001',
+        { maxImagesPerCall: 2 },
+        {
+          provider: 'google-vertex',
+          baseURL: 'https://api.example.com',
+          headers: { 'api-key': 'test-key' },
+        },
+      );
+
+      expect(customModel.maxImagesPerCall).toBe(2);
+    });
+
+    it('should use default maxImagesPerCall when not specified', () => {
+      const defaultModel = new GoogleVertexImageModel(
+        'imagen-3.0-generate-001',
+        {},
+        {
+          provider: 'google-vertex',
+          baseURL: 'https://api.example.com',
+          headers: { 'api-key': 'test-key' },
+        },
+      );
+
+      expect(defaultModel.maxImagesPerCall).toBe(4);
+    });
+
     it('should extract the generated images', async () => {
       prepareJsonResponse();
 

packages/google-vertex/src/google-vertex-image-model.ts

@@ -8,11 +8,10 @@ import {
 } from '@ai-sdk/provider-utils';
 import { z } from 'zod';
 import { googleVertexFailedResponseHandler } from './google-vertex-error';
-
-export type GoogleVertexImageModelId =
-  | 'imagen-3.0-generate-001'
-  | 'imagen-3.0-fast-generate-001'
-  | (string & {});
+import {
+  GoogleVertexImageModelId,
+  GoogleVertexImageSettings,
+} from './google-vertex-image-settings';
 
 interface GoogleVertexImageModelConfig {
   provider: string;
@@ -25,15 +24,18 @@ interface GoogleVertexImageModelConfig {
 export class GoogleVertexImageModel implements ImageModelV1 {
   readonly specificationVersion = 'v1';
 
-  // https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/imagen-api#parameter_list
-  readonly maxImagesPerCall = 4;
-
   get provider(): string {
     return this.config.provider;
   }
 
+  get maxImagesPerCall(): number {
+    // https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/imagen-api#parameter_list
+    return this.settings.maxImagesPerCall ?? 4;
+  }
+
   constructor(
     readonly modelId: GoogleVertexImageModelId,
+    readonly settings: GoogleVertexImageSettings,
     private config: GoogleVertexImageModelConfig,
   ) {}
 

packages/google-vertex/src/google-vertex-image-settings.ts

@@ -0,0 +1,11 @@
+export type GoogleVertexImageModelId =
+  | 'imagen-3.0-generate-001'
+  | 'imagen-3.0-fast-generate-001'
+  | (string & {});
+
+export interface GoogleVertexImageSettings {
+  /**
+Override the maximum number of images per call (default 4)
+   */
+  maxImagesPerCall?: number;
+}