Feat: Stable diffusion v3

src/globals.ts

@@ -13,6 +13,7 @@ export const HEADER_KEYS: Record<string, string> = {
   CUSTOM_HOST: `x-${POWERED_BY}-custom-host`,
   REQUEST_TIMEOUT: `x-${POWERED_BY}-request-timeout`,
   STRICT_OPEN_AI_COMPLIANCE: `x-${POWERED_BY}-strict-open-ai-compliance`,
+  CONTENT_TYPE: `Content-Type`,
 };
 
 export const RESPONSE_HEADER_KEYS: Record<string, string> = {

src/handlers/handlerUtils.ts

@@ -11,9 +11,9 @@ import {
   OPEN_AI,
   AZURE_AI_INFERENCE,
   ANTHROPIC,
-  MULTIPART_FORM_DATA_ENDPOINTS,
   CONTENT_TYPES,
   HUGGING_FACE,
+  STABILITY_AI,
 } from '../globals';
 import Providers from '../providers';
 import { ProviderAPIConfig, endpointStrings } from '../providers/types';
@@ -524,6 +524,7 @@ export async function tryPost(
     fn,
     transformedRequestBody,
     transformedRequestUrl: url,
+    gatewayRequestBody: params,
   });
 
   // Construct the base object for the POST request
@@ -535,9 +536,10 @@ export async function tryPost(
     requestHeaders
   );
 
-  fetchOptions.body = MULTIPART_FORM_DATA_ENDPOINTS.includes(fn)
-    ? (transformedRequestBody as FormData)
-    : JSON.stringify(transformedRequestBody);
+  fetchOptions.body =
+    headers[HEADER_KEYS.CONTENT_TYPE] === CONTENT_TYPES.MULTIPART_FORM_DATA
+      ? (transformedRequestBody as FormData)
+      : JSON.stringify(transformedRequestBody);
 
   providerOption.retry = {
     attempts: providerOption.retry?.attempts ?? 0,
@@ -1012,6 +1014,14 @@ export function constructConfigFromRequestHeaders(
     azureModelName: requestHeaders[`x-${POWERED_BY}-azure-model-name`],
   };
 
+  const stabilityAiConfig = {
+    stabilityClientId: requestHeaders[`x-${POWERED_BY}-stability-client-id`],
+    stabilityClientUserId:
+      requestHeaders[`x-${POWERED_BY}-stability-client-user-id`],
+    stabilityClientVersion:
+      requestHeaders[`x-${POWERED_BY}-stability-client-version`],
+  };
+
   const azureAiInferenceConfig = {
     azureDeploymentName:
       requestHeaders[`x-${POWERED_BY}-azure-deployment-name`],
@@ -1128,6 +1138,12 @@ export function constructConfigFromRequestHeaders(
           ...anthropicConfig,
         };
       }
+      if (parsedConfigJson.provider === STABILITY_AI) {
+        parsedConfigJson = {
+          ...parsedConfigJson,
+          ...stabilityAiConfig,
+        };
+      }
     }
     return convertKeysToCamelCase(parsedConfigJson, [
       'override_params',
@@ -1158,6 +1174,8 @@ export function constructConfigFromRequestHeaders(
       huggingfaceConfig),
     mistralFimCompletion:
       requestHeaders[`x-${POWERED_BY}-mistral-fim-completion`],
+    ...(requestHeaders[`x-${POWERED_BY}-provider`] === STABILITY_AI &&
+      stabilityAiConfig),
   };
 }
 

src/providers/bedrock/constants.ts

@@ -34,3 +34,8 @@ export const MISTRAL_CONTROL_TOKENS = {
   MIDDLE: '[MIDDLE]',
   SUFFIX: '[SUFFIX]',
 };
+
+export const BEDROCK_STABILITY_V1_MODELS = [
+  'stable-diffusion-xl-v0',
+  'stable-diffusion-xl-v1',
+];

src/providers/bedrock/imageGenerate.ts

@@ -1,10 +1,11 @@
 import { BEDROCK } from '../../globals';
+import { StabilityAIImageGenerateV2Config } from '../stability-ai/imageGenerateV2';
 import { ErrorResponse, ImageGenerateResponse, ProviderConfig } from '../types';
 import { generateInvalidProviderResponseError } from '../utils';
 import { BedrockErrorResponseTransform } from './chatComplete';
 import { BedrockErrorResponse } from './embed';
 
-export const BedrockStabilityAIImageGenerateConfig: ProviderConfig = {
+export const BedrockStabilityAIImageGenerateV1Config: ProviderConfig = {
   prompt: {
     param: 'text_prompts',
     required: true,
@@ -47,29 +48,60 @@ interface ImageArtifact {
   seed: number;
 }
 
-interface BedrockStabilityAIImageGenerateResponse {
+interface BedrockStabilityAIImageGenerateV1Response {
   result: string;
   artifacts: ImageArtifact[];
 }
 
-export const BedrockStabilityAIImageGenerateResponseTransform: (
-  response: BedrockStabilityAIImageGenerateResponse | BedrockErrorResponse,
+export const BedrockStabilityAIImageGenerateV1ResponseTransform: (
+  response: BedrockStabilityAIImageGenerateV1Response | BedrockErrorResponse,
   responseStatus: number
 ) => ImageGenerateResponse | ErrorResponse = (response, responseStatus) => {
   if (responseStatus !== 200) {
-    const errorResposne = BedrockErrorResponseTransform(
+    const errorResponse = BedrockErrorResponseTransform(
       response as BedrockErrorResponse
     );
-    if (errorResposne) return errorResposne;
+    if (errorResponse) return errorResponse;
   }
 
   if ('artifacts' in response) {
     return {
-      created: `${new Date().getTime()}`,
+      created: Math.floor(Date.now() / 1000),
       data: response.artifacts.map((art) => ({ b64_json: art.base64 })),
       provider: BEDROCK,
     };
   }
 
   return generateInvalidProviderResponseError(response, BEDROCK);
 };
+
+interface BedrockStabilityAIImageGenerateV2Response {
+  seeds: number[];
+  finish_reasons: string[];
+  images: string[];
+}
+
+export const BedrockStabilityAIImageGenerateV2Config =
+  StabilityAIImageGenerateV2Config;
+
+export const BedrockStabilityAIImageGenerateV2ResponseTransform: (
+  response: BedrockStabilityAIImageGenerateV2Response | BedrockErrorResponse,
+  responseStatus: number
+) => ImageGenerateResponse | ErrorResponse = (response, responseStatus) => {
+  if (responseStatus !== 200) {
+    const errorResponse = BedrockErrorResponseTransform(
+      response as BedrockErrorResponse
+    );
+    if (errorResponse) return errorResponse;
+  }
+
+  if ('images' in response) {
+    return {
+      created: Math.floor(Date.now() / 1000),
+      data: response.images.map((image) => ({ b64_json: image })),
+      provider: BEDROCK,
+    };
+  }
+
+  return generateInvalidProviderResponseError(response, BEDROCK);
+};

src/providers/bedrock/index.ts

@@ -42,15 +42,18 @@ import {
   BedrockTitanCompleteResponseTransform,
   BedrockTitanCompleteStreamChunkTransform,
 } from './complete';
+import { BEDROCK_STABILITY_V1_MODELS } from './constants';
 import {
   BedrockCohereEmbedConfig,
   BedrockCohereEmbedResponseTransform,
   BedrockTitanEmbedConfig,
   BedrockTitanEmbedResponseTransform,
 } from './embed';
 import {
-  BedrockStabilityAIImageGenerateConfig,
-  BedrockStabilityAIImageGenerateResponseTransform,
+  BedrockStabilityAIImageGenerateV1Config,
+  BedrockStabilityAIImageGenerateV1ResponseTransform,
+  BedrockStabilityAIImageGenerateV2Config,
+  BedrockStabilityAIImageGenerateV2ResponseTransform,
 } from './imageGenerate';
 
 const BedrockConfig: ProviderConfigs = {
@@ -63,8 +66,9 @@ const BedrockConfig: ProviderConfigs = {
     // To remove the region in case its a cross-region inference profile ID
     // https://docs.aws.amazon.com/bedrock/latest/userguide/cross-region-inference-support.html
     const providerModel = params.model.replace(/^(us\.|eu\.)/, '');
-    const provider = providerModel?.split('.')[0];
-    const model = providerModel?.split('.')[1];
+    const providerModelArray = providerModel.split('.');
+    const provider = providerModelArray[0];
+    const model = providerModelArray.slice(1).join('.');
     switch (provider) {
       case ANTHROPIC:
         return {
@@ -148,11 +152,20 @@ const BedrockConfig: ProviderConfigs = {
           },
         };
       case 'stability':
+        if (model && BEDROCK_STABILITY_V1_MODELS.includes(model)) {
+          return {
+            imageGenerate: BedrockStabilityAIImageGenerateV1Config,
+            api: BedrockAPIConfig,
+            responseTransforms: {
+              imageGenerate: BedrockStabilityAIImageGenerateV1ResponseTransform,
+            },
+          };
+        }
         return {
-          imageGenerate: BedrockStabilityAIImageGenerateConfig,
+          imageGenerate: BedrockStabilityAIImageGenerateV2Config,
           api: BedrockAPIConfig,
           responseTransforms: {
-            imageGenerate: BedrockStabilityAIImageGenerateResponseTransform,
+            imageGenerate: BedrockStabilityAIImageGenerateV2ResponseTransform,
           },
         };
       default:

src/providers/fireworks-ai/imageGenerate.ts

@@ -100,7 +100,7 @@ export const FireworksAIImageGenerateResponseTransform: (
   }
   if (response instanceof Array) {
     return {
-      created: `${new Date().getTime()}`, // Corrected method call
+      created: Math.floor(Date.now() / 1000), // Corrected method call
       data: response?.map((r) => ({
         b64_json: r.base64,
         seed: r.seed,

src/providers/segmind/imageGenerate.ts

@@ -142,7 +142,7 @@ export const SegmindImageGenerateResponseTransform: (
     let dataObj: object[] = imageArr.map((img) => ({ b64_json: img }));
 
     return {
-      created: `${new Date().getTime()}`, // Corrected method call
+      created: Math.floor(Date.now() / 1000), // Corrected method call
       data: dataObj,
       provider: SEGMIND,
     } as ImageGenerateResponse;

src/providers/stability-ai/api.ts

@@ -1,9 +1,17 @@
+import { CONTENT_TYPES } from '../../globals';
 import { ProviderAPIConfig } from '../types';
+import { isStabilityV1Model } from './utils';
 
 const StabilityAIAPIConfig: ProviderAPIConfig = {
-  getBaseURL: () => 'https://api.stability.ai/v1',
-  headers: ({ providerOptions }) => {
-    return { Authorization: `Bearer ${providerOptions.apiKey}` };
+  getBaseURL: () => 'https://api.stability.ai',
+  headers: ({ providerOptions, gatewayRequestBody }) => {
+    const headers: Record<string, string> = {
+      Authorization: `Bearer ${providerOptions.apiKey}`,
+    };
+    if (isStabilityV1Model(gatewayRequestBody?.model)) return headers;
+    headers['Content-Type'] = CONTENT_TYPES.MULTIPART_FORM_DATA;
+    headers['Accept'] = CONTENT_TYPES.APPLICATION_JSON;
+    return headers;
   },
   getEndpoint: ({ fn, gatewayRequestBody, providerOptions }) => {
     let mappedFn = fn;
@@ -18,12 +26,18 @@ const StabilityAIAPIConfig: ProviderAPIConfig = {
 
     switch (mappedFn) {
       case 'imageGenerate': {
-        return `/generation/${gatewayRequestBody.model}/text-to-image`;
+        if (isStabilityV1Model(gatewayRequestBody.model))
+          return `/v1/generation/${gatewayRequestBody.model}/text-to-image`;
+        return `/v2beta/stable-image/generate/${gatewayRequestBody.model}`;
       }
       default:
         return '';
     }
   },
+  transformToFormData: ({ gatewayRequestBody }) => {
+    if (isStabilityV1Model(gatewayRequestBody.model)) return false;
+    return true;
+  },
 };
 
 export default StabilityAIAPIConfig;

src/providers/stability-ai/constants.ts

@@ -0,0 +1,4 @@
+export const STABILITY_V1_MODELS = [
+  'stable-diffusion-xl-1024-v1-0',
+  'stable-diffusion-v1-6',
+];

src/providers/stability-ai/imageGenerate.ts

@@ -5,7 +5,7 @@ import {
   generateInvalidProviderResponseError,
 } from '../utils';
 
-export const StabilityAIImageGenerateConfig: ProviderConfig = {
+export const StabilityAIImageGenerateV1Config: ProviderConfig = {
   prompt: {
     param: 'text_prompts',
     required: true,
@@ -60,15 +60,11 @@ export const StabilityAIImageGenerateConfig: ProviderConfig = {
   },
 };
 
-interface StabilityAIImageGenerateResponse extends ImageGenerateResponse {
+interface StabilityAIImageGenerateV1Response extends ImageGenerateResponse {
   artifacts: ImageArtifact[];
 }
 
-interface StabilityAIImageGenerateResponse extends ImageGenerateResponse {
-  artifacts: ImageArtifact[];
-}
-
-interface StabilityAIImageGenerateErrorResponse {
+interface StabilityAIImageGenerateV1ErrorResponse {
   id: string;
   name: string;
   message: string;
@@ -80,10 +76,10 @@ interface ImageArtifact {
   seed: number; // The seed associated with this image
 }
 
-export const StabilityAIImageGenerateResponseTransform: (
+export const StabilityAIImageGenerateV1ResponseTransform: (
   response:
-    | StabilityAIImageGenerateResponse
-    | StabilityAIImageGenerateErrorResponse,
+    | StabilityAIImageGenerateV1Response
+    | StabilityAIImageGenerateV1ErrorResponse,
   responseStatus: number
 ) => ImageGenerateResponse | ErrorResponse = (response, responseStatus) => {
   if (responseStatus !== 200 && 'message' in response) {
@@ -100,7 +96,7 @@ export const StabilityAIImageGenerateResponseTransform: (
 
   if ('artifacts' in response) {
     return {
-      created: `${new Date().getTime()}`, // Corrected method call
+      created: Math.floor(Date.now() / 1000), // Corrected method call
       data: response.artifacts.map((art) => ({ b64_json: art.base64 })), // Corrected object creation within map
       provider: STABILITY_AI,
     };