From 7dce12f6ab4cfa3aa3ead5aa1ea98ff0ad8fcdfe Mon Sep 17 00:00:00 2001 From: vicentefelipechile Date: Sat, 19 Oct 2024 12:29:23 -0300 Subject: [PATCH] Added support to Stability Model --- .env.example | 15 +++++++-- README.md | 9 ++++-- src/baileys/env.ts | 14 ++++++++- src/baileys/handlers/message.ts | 4 ++- src/models/StabilityModel.ts | 56 +++++++++++++++++++++++++++++++++ src/types/AiModels.d.ts | 2 +- src/whatsapp-ai.config.ts | 4 +++ 7 files changed, 96 insertions(+), 8 deletions(-) create mode 100644 src/models/StabilityModel.ts diff --git a/.env.example b/.env.example index 8539ee3..cbfc5e7 100644 --- a/.env.example +++ b/.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=🤗 \ No newline at end of file +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 \ No newline at end of file diff --git a/README.md b/README.md index 3177a9e..a7aa14c 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ -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 diff --git a/src/baileys/env.ts b/src/baileys/env.ts index 9543232..b2b6552 100644 --- a/src/baileys/env.ts +++ b/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' }; diff --git a/src/baileys/handlers/message.ts b/src/baileys/handlers/message.ts index 785407f..850201b 100644 --- a/src/baileys/handlers/message.ts +++ b/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 = { 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 diff --git a/src/models/StabilityModel.ts b/src/models/StabilityModel.ts new file mode 100644 index 0000000..2511e25 --- /dev/null +++ b/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 { + 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 { + 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 { + try { + const imageData = await this.generateImage(prompt); + + await handle({ image: imageData }); + } catch (err) { + await handle('', err as string); + } + } +} + +export { StabilityModel }; diff --git a/src/types/AiModels.d.ts b/src/types/AiModels.d.ts index d96c646..938cc0e 100644 --- a/src/types/AiModels.d.ts +++ b/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; diff --git a/src/whatsapp-ai.config.ts b/src/whatsapp-ai.config.ts index e65be69..e8f74ee 100644 --- a/src/whatsapp-ai.config.ts +++ b/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: [