Merge pull request #61 from vicentefelipechile/master

Added support to Stability Model

.env.example

@@ -6,12 +6,13 @@ API_KEY_OPENAI=ADD_YOUR_KEY
 API_KEY_DREAMSTUDIO=ADD_YOUR_KEY
 API_KEY_GEMINI=ADD_YOUR_KEY
 API_KEY_HF=HUGGING_FACE_TOKEN
+API_KEY_STABILITY=ADD_YOUR_KEY
 
 # Optional
 # API_KEY_OPENAI_DALLE=ADD_YOUR_KEY
 
 # MongoDB
-MONGO_ENABLED=True
+MONGO_ENABLED=False
 MONGO_URL=YOUR_MONGO_DB_URL
 
 
@@ -34,4 +35,14 @@ GEMINI_ICON_PREFIX=🔮
 # # Hugging Face Flux
 HF_PREFIX=!flux
 HF_ENABLED=False
-HF_ICON_PREFIX=🤗
+HF_ICON_PREFIX=🤗
+
+# # Stability
+STABILITY_PREFIX=!stability
+STABILITY_ENABLED=False
+STABILITY_ICON_PREFIX=🧠
+STABILITY_MODEL=core
+
+# https://platform.stability.ai/docs/api-reference#tag/Generate
+# Available Models:
+# core | ultra | sd3

README.md

@@ -10,7 +10,7 @@
 
 </div>
 
-The WhatsApp AI Bot is a chatbot that uses AI models APIs to generate responses to user input. The bot supports several AI models, including **`Gemini`**, **`ChatGPT`** and **`DALL-E`**..
+The WhatsApp AI Bot is a chatbot that uses AI models APIs to generate responses to user input. The bot supports several AI models, including **`Gemini`**, **`ChatGPT`**, **`DALL-E`**, **`Flux`**, and **`Stability AI`**.
 
 > [!NOTE]
 > Custom models aren't supported for now, but we are working on it.
@@ -35,6 +35,7 @@ The WhatsApp AI Bot is a chatbot that uses AI models APIs to generate responses
 | Gemini Vision | [Google](https://ai.google.dev/gemini-api/docs/vision?lang=node#upload-image) | Image to Text | none |
 | Dalle 2 & 3 | [OpenAI](https://platform.openai.com/docs/api-reference/images/create) | Text to Image | !dalle |
 | Flux   | [Hugging Face](https://huggingface.co/black-forest-labs/FLUX.1-dev) | Text to Image | !flux |
+| Stability AI | [Stability AI](https://platform.stability.ai/docs/getting-started/stable-image) | Text to Image | !stability |
 
 # Demo
 
@@ -85,7 +86,8 @@ The WhatsApp AI Bot is a chatbot that uses AI models APIs to generate responses
 
 - [Gemini API Key](https://aistudio.google.com/app/apikey)
 - [OpenAI API Key](https://platform.openai.com/api-keys)
-- [Hugging Face API Key](https://huggingface.co/)
+- [Hugging Face API Key](https://huggingface.co/settings/tokens)
+- [Stability AI API Key](https://platform.stability.ai/account/keys)
 
 ### 3. Add API Keys
 
@@ -103,6 +105,7 @@ Copy the file `.env.example` and rename it to `.env`, then set any settings you
 - `!chatgpt` use chat-gpt.
 - `!dalle` use Dalle.
 - `!flux` use flux.
+- `!stability` use stability.
 
 **Note! open `src/whatsapp-ai.config.ts` to edit config.**
 
@@ -119,7 +122,7 @@ Copy the file `.env.example` and rename it to `.env`, then set any settings you
 
 # Disclaimer
 
-This bot utilizes Puppeteer to operate an actual instance of Whatsapp Web to prevent blocking. However, it is essential to note that these operations come at a cost charged by OpenAI and Stability AI for every request made. Please be aware that WhatsApp does not support bots or unofficial clients on its platform, so using this method is not entirely secure and could lead to getting blocked.
+This bot utilizes [baileys](https://github.com/WhiskeySockets/Baileys) to operate an actual instance of Whatsapp Web to prevent blocking. However, it is essential to note that these operations come at a cost charged by OpenAI and Stability AI for every request made. Please be aware that WhatsApp does not support bots or unofficial clients on its platform, so using this method is not entirely secure and could lead to getting blocked.
 
 ## Contributors
 

src/baileys/env.ts

@@ -13,6 +13,7 @@ interface EnvInterface {
   API_KEY_DREAMSTUDIO?: string;
   API_KEY_GEMINI?: string;
   API_KEY_HF?: string;
+  API_KEY_STABILITY?: string;
 
   // MongoDB
   MONGO_ENABLED: boolean;
@@ -38,6 +39,12 @@ interface EnvInterface {
   HF_PREFIX?: string;
   HF_ENABLED: boolean;
   HF_ICON_PREFIX?: string;
+
+  // // Stability
+  STABILITY_PREFIX?: string;
+  STABILITY_ENABLED: boolean;
+  STABILITY_ICON_PREFIX?: string;
+  STABILITY_MODEL: string;
 }
 
 export const ENV: EnvInterface = {
@@ -49,6 +56,7 @@ export const ENV: EnvInterface = {
   API_KEY_DREAMSTUDIO: process.env.API_KEY_DREAMSTUDIO,
   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,
 
@@ -64,5 +72,9 @@ export const ENV: EnvInterface = {
   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
+  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,
+  STABILITY_MODEL: process.env.STABILITY_MODEL || 'core'
 };

src/baileys/handlers/message.ts

@@ -4,6 +4,7 @@ import { AIModels } from './../../types/AiModels';
 import { Util } from './../../util/Util';
 
 /* Models */
+import { StabilityModel } from '../../models/StabilityModel';
 import { ChatGPTModel } from './../../models/OpenAIModel';
 import { GeminiModel } from './../../models/GeminiModel';
 import { FluxModel } from './../../models/FluxModel';
@@ -13,7 +14,8 @@ import { ENV } from '../env';
 const modelTable: Record<AIModels, any> = {
   ChatGPT: ENV.OPENAI_ENABLED ? new ChatGPTModel() : null,
   Gemini: ENV.GEMINI_ENABLED ? new GeminiModel() : null,
-  FLUX: ENV.HF_ENABLED ? new FluxModel() : null
+  FLUX: ENV.HF_ENABLED ? new FluxModel() : null,
+  Stability: ENV.STABILITY_ENABLED ? new StabilityModel() : null
 };
 
 // handles message

src/models/StabilityModel.ts

@@ -0,0 +1,56 @@
+/* Third-party modules */
+import FormData from 'form-data';
+import axios from 'axios';
+
+/* Local modules */
+import { AIArguments, AIHandle, AIModel } from './BaseAiModel';
+import { ENV } from '../baileys/env';
+
+/* Stability Mode */
+class StabilityModel extends AIModel<AIArguments, AIHandle> {
+  public endPointGenerate: string = 'https://api.stability.ai/v2beta/stable-image/generate/';
+  public headers;
+
+  constructor() {
+    super(ENV.API_KEY_STABILITY, 'Stability');
+
+    this.endPointGenerate = this.endPointGenerate + ENV.STABILITY_MODEL;
+    this.headers = {
+      Authorization: `Bearer ${this.getApiKey()}`,
+      Accept: 'image/*'
+    };
+  }
+
+  public async generateImage(prompt: string): Promise<Buffer> {
+    const payload = {
+      prompt: prompt,
+      output_format: 'jpeg'
+    };
+
+    const formData = new FormData();
+
+    const response = await axios.postForm(this.endPointGenerate, axios.toFormData(payload, formData), {
+      validateStatus: undefined,
+      responseType: 'arraybuffer',
+      headers: this.headers
+    });
+
+    if ( response.status === 200 ) {
+      return Buffer.from(response.data);
+    } else {
+      throw new Error(response.data.errors[0]);
+    }
+  }
+
+  async sendMessage({ prompt }: AIArguments, handle: AIHandle): Promise<any> {
+    try {
+      const imageData = await this.generateImage(prompt);
+
+      await handle({ image: imageData });
+    } catch (err) {
+      await handle('', err as string);
+    }
+  }
+}
+
+export { StabilityModel };

src/types/AiModels.d.ts

@@ -1,2 +1,2 @@
-export type AIModels = 'ChatGPT' | 'Gemini' | 'FLUX';
+export type AIModels = 'ChatGPT' | 'Gemini' | 'FLUX' | 'Stability';
 export type AIModelsName = Exclude<AIModels, 'Custom'>;

src/whatsapp-ai.config.ts

@@ -22,6 +22,10 @@ const config: Config = {
     FLUX: {
       prefix: ENV.HF_PREFIX,
       enable: ENV.HF_ENABLED
+    },
+    Stability: {
+      prefix: ENV.STABILITY_PREFIX,
+      enable: ENV.STABILITY_ENABLED,
     }
     /*
     Custom: [