All files / api/src/ai-providers/stability-ai image-generate.ts

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
import {
  generateErrorResponse,
  generateInvalidProviderResponseError,
} from '@api/utils/ai-provider';
import type { AIProviderFunctionConfig } from '@shared/types/ai-providers/config';
import type { ErrorResponseBody } from '@shared/types/api/response/body';
import type {
  GenerateImageRequestBody,
  GenerateImageResponseBody,
} from '@shared/types/api/routes/images-api';
import { AIProvider } from '@shared/types/constants';
 
export const StabilityAIImageGenerateV1Config: AIProviderFunctionConfig = {
  prompt: {
    param: 'text_prompts',
    required: true,
    transform: (saRequestBody: GenerateImageRequestBody) => {
      return [
        {
          text: saRequestBody.prompt,
          weight: 1,
        },
      ];
    },
  },
  n: {
    param: 'samples',
    min: 1,
    max: 10,
  },
  size: [
    {
      param: 'height',
      transform: (saRequestBody: GenerateImageRequestBody): number =>
        parseInt(saRequestBody.size?.toLowerCase().split('x')[1] || '0', 10),
      min: 320,
    },
    {
      param: 'width',
      transform: (saRequestBody: GenerateImageRequestBody): number =>
        parseInt(saRequestBody.size?.toLowerCase().split('x')[0] || '0', 10),
      min: 320,
    },
  ],
  style: {
    param: 'style_preset',
  },
  cfg_scale: {
    param: 'cfg_scale',
  },
  clip_guidance_preset: {
    param: 'clip_guidance_preset',
  },
  sampler: {
    param: 'sampler',
  },
  seed: {
    param: 'seed',
  },
  steps: {
    param: 'steps',
  },
  extras: {
    param: 'extras',
  },
};
 
interface StabilityAIImageGenerateV1Response extends GenerateImageResponseBody {
  artifacts: ImageArtifact[];
}
 
interface StabilityAIImageGenerateV1ErrorResponse {
  id: string;
  name: string;
  message: string;
}
 
interface ImageArtifact {
  base64: string; // Image encoded in base64
  finishReason: 'CONTENT_FILTERED' | 'ERROR' | 'SUCCESS'; // Enum for finish reason
  seed: number; // The seed associated with this image
}
 
export const StabilityAIImageGenerateV1ResponseTransform: (
  response:
    | StabilityAIImageGenerateV1Response
    | StabilityAIImageGenerateV1ErrorResponse,
  responseStatus: number,
) => GenerateImageResponseBody | ErrorResponseBody = (
  response,
  responseStatus,
) => {
  if (responseStatus !== 200 && 'message' in response) {
    return generateErrorResponse(
      {
        message: response.message,
        type: response.name,
        param: '',
        code: '',
      },
      AIProvider.STABILITY_AI,
    );
  }
 
  if ('artifacts' in response) {
    return {
      created: Math.floor(Date.now() / 1000), // Corrected method call
      data: response.artifacts.map((art) => ({ b64_json: art.base64 })), // Corrected object creation within map
      provider: AIProvider.STABILITY_AI,
    };
  }
 
  return generateInvalidProviderResponseError(
    response as unknown as Record<string, unknown>,
    AIProvider.STABILITY_AI,
  );
};