Merge pull request #64 from vicentefelipechile/master

Fixed Dalle sending Image messages
Zain-ul-din · Oct 24, 2024 · ca06c7e · ca06c7e
2 parents 0fb8b2a + 6b64895
commit ca06c7e
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 11 deletions.
diff --git a/.env.example b/.env.example
@@ -26,6 +26,8 @@ DALLE_PREFIX=!dalle
 DALLE_ENABLED=False
 DALLE_ICON_PREFIX=🎨
 DALLE_USE_3=False
+# See all supported sizes here https://platform.openai.com/docs/api-reference/images/create#images-create-size
+DALLE_SIZE=512x512
 
 # # Google Gemini
 GEMINI_PREFIX=!gemini

diff --git a/src/baileys/env.ts b/src/baileys/env.ts
@@ -30,6 +30,7 @@ interface EnvInterface {
   DALLE_ENABLED: boolean;
   DALLE_ICON_PREFIX?: string;
   DALLE_USE_3: boolean;
+  DALLE_SIZE: string;
 
   // // Google Gemini
   GEMINI_PREFIX?: string;
@@ -58,22 +59,29 @@ export const ENV: EnvInterface = {
   API_KEY_GEMINI: process.env.API_KEY_GEMINI,
   API_KEY_HF: process.env.API_KEY_HF,
   API_KEY_STABILITY: process.env.API_KEY_STABILITY,
+
   MONGO_ENABLED: process.env.MONGO_ENABLED === 'True',
   MONGO_URL: process.env.MONGO_URL,
+
   OPENAI_MODEL: process.env.OPENAI_MODEL || 'gpt-3.5-turbo',
   OPENAI_PREFIX: process.env.OPENAI_PREFIX,
   OPENAI_ENABLED: process.env.OPENAI_ENABLED === 'True',
   OPENAI_ICON_PREFIX: process.env.OPENAI_ICON_PREFIX,
+
   DALLE_PREFIX: process.env.DALLE_PREFIX,
   DALLE_ENABLED: process.env.DALLE_ENABLED === 'True',
   DALLE_ICON_PREFIX: process.env.DALLE_ICON_PREFIX,
   DALLE_USE_3: process.env.DALLE_USE_3 === 'True',
+  DALLE_SIZE: process.env.DALLE_SIZE || '512x512',
+
   GEMINI_PREFIX: process.env.GEMINI_PREFIX,
   GEMINI_ENABLED: process.env.GEMINI_ENABLED === 'True',
   GEMINI_ICON_PREFIX: process.env.GEMINI_ICON_PREFIX,
+
   HF_PREFIX: process.env.HF_PREFIX,
   HF_ENABLED: process.env.HF_ENABLED === 'True',
   HF_ICON_PREFIX: process.env.HF_ICON_PREFIX,
+
   STABILITY_PREFIX: process.env.STABILITY_PREFIX,
   STABILITY_ENABLED: process.env.STABILITY_ENABLED === 'True',
   STABILITY_ICON_PREFIX: process.env.STABILITY_ICON_PREFIX,

diff --git a/src/models/OpenAIModel.ts b/src/models/OpenAIModel.ts
@@ -1,5 +1,6 @@
 /* Third-party modules */
 import { ChatCompletionMessage } from 'openai/resources/chat/completions';
+import { Image, ImagesResponse } from 'openai/resources';
 import OpenAI from 'openai';
 
 /* Local modules */
@@ -9,16 +10,19 @@ import config from '../whatsapp-ai.config';
 
 /* Util */
 interface BotImageResponse {
-  url: string;
+  image: Buffer;
   caption: string;
 }
 
+type DalleSizeImage = '256x256' | '512x512' | '1024x1024' | '1792x1024' | '1024x1792' | null;
+
 /* ChatGPT Model */
 class ChatGPTModel extends AIModel<AIArguments, AIHandle> {
   /* Variables */
   private Dalle3: boolean;
   private Dalle: OpenAI;
   private OpenAI: OpenAI;
+  public DalleSize: DalleSizeImage;
 
   public constructor() {
     super(ENV.API_KEY_OPENAI, 'ChatGPT');
@@ -32,6 +36,7 @@ class ChatGPTModel extends AIModel<AIArguments, AIHandle> {
     });
 
     this.Dalle3 = ENV.DALLE_USE_3;
+    this.DalleSize = ENV.DALLE_SIZE as DalleSizeImage;
   }
 
   /* Methods */
@@ -48,22 +53,21 @@ class ChatGPTModel extends AIModel<AIArguments, AIHandle> {
   }
 
   public async generateImage(prompt: string): Promise<BotImageResponse> {
-    const res: OpenAI.Images.ImagesResponse = await this.Dalle.images.generate({
+    const res: ImagesResponse = await this.Dalle.images.generate({
       model: this.Dalle3 ? 'dall-e-3' : 'dall-e-2',
       n: 1,
-      size: '512x512',
-      prompt: prompt
+      size: this.DalleSize,
+      prompt: prompt,
+      response_format: 'b64_json'
     });
 
-    const resInfo: OpenAI.Images.Image = res.data[0];
+    const image: Image = res.data[0];
+    const base64Img = Buffer.from(image.b64_json as string, 'base64');
 
-    return { url: resInfo.url as string, caption: resInfo.revised_prompt as string };
+    return { image: base64Img, caption: image.revised_prompt || prompt };
   }
 
-  public async sendMessage(
-    { sender, prompt, prefix }: AIArguments,
-    handle: AIHandle
-  ): Promise<any> {
+  public async sendMessage({ sender, prompt, prefix }: AIArguments, handle: AIHandle) {
     try {
       // Use Dalle
       if (ENV.DALLE_ENABLED && prefix === ENV.DALLE_PREFIX) {

diff --git a/src/types/Config.d.ts b/src/types/Config.d.ts
@@ -1,4 +1,4 @@
-import { AIModels, AIModelsName } from './AiModels';
+import { AIModels } from './AiModels';
 
 export interface IModelConfig {
   prefix: string | undefined;