From 385df6c8211d0b972e7e22b706353f94c90ad8dc Mon Sep 17 00:00:00 2001 From: yihong0618 Date: Thu, 18 Apr 2024 11:01:37 +0800 Subject: [PATCH] feat: sd3 api Signed-off-by: yihong0618 --- .gitignore | 1 + README.md | 2 +- handlers/chatgpt.py | 16 +++++------ handlers/gemini.py | 1 - handlers/sd.py | 70 +++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 3 +- 6 files changed, 81 insertions(+), 12 deletions(-) create mode 100644 handlers/sd.py diff --git a/.gitignore b/.gitignore index d44afcb..4a571f8 100644 --- a/.gitignore +++ b/.gitignore @@ -160,6 +160,7 @@ cython_debug/ #.idea/ *.jpg *.png +*.jpeg OUT_FOLDER/ cache/ nohup.out diff --git a/README.md b/README.md index 84c8a50..abb8332 100644 --- a/README.md +++ b/README.md @@ -45,7 +45,7 @@ Note, if you are using third party service, you need to `export ANTHROPIC_BASE_U 1. visit https://platform.openai.com/account/api-keys get the key 2. export OPENAI_API_KEY=${the_key} -3. use `chatgpt: ${message}` to ask +3. use `gpt: ${message}` to ask Note, if you are using third party service, you need to `export OPENAI_API_BASE=${the_url}` to change the url. diff --git a/handlers/chatgpt.py b/handlers/chatgpt.py index 5ed9e93..24578a7 100644 --- a/handlers/chatgpt.py +++ b/handlers/chatgpt.py @@ -28,7 +28,7 @@ chatgpt_pro_player_dict = {} def chatgpt_handler(message: Message, bot: TeleBot) -> None: - """chatgpt : /chatgpt """ + """gpt : /gpt """ m = message.text.strip() player_message = [] @@ -90,7 +90,7 @@ def chatgpt_handler(message: Message, bot: TeleBot) -> None: def chatgpt_pro_handler(message: Message, bot: TeleBot) -> None: - """chatgpt_pro : /chatgpt_pro """ + """gpt_pro : /gpt_pro """ m = message.text.strip() player_message = [] @@ -206,17 +206,15 @@ def chatgpt_photo_handler(message: Message, bot: TeleBot) -> None: def register(bot: TeleBot) -> None: - bot.register_message_handler(chatgpt_handler, commands=["chatgpt"], pass_bot=True) - bot.register_message_handler(chatgpt_handler, regexp="^chatgpt:", pass_bot=True) + bot.register_message_handler(chatgpt_handler, commands=["gpt"], pass_bot=True) + bot.register_message_handler(chatgpt_handler, regexp="^gpt:", pass_bot=True) bot.register_message_handler( - chatgpt_pro_handler, commands=["chatgpt_pro"], pass_bot=True - ) - bot.register_message_handler( - chatgpt_pro_handler, regexp="^chatgpt_pro:", pass_bot=True + chatgpt_pro_handler, commands=["gpt_pro"], pass_bot=True ) + bot.register_message_handler(chatgpt_pro_handler, regexp="^gpt_pro:", pass_bot=True) bot.register_message_handler( chatgpt_photo_handler, content_types=["photo"], - func=lambda m: m.caption and m.caption.startswith(("chatgpt:", "/chatgpt")), + func=lambda m: m.caption and m.caption.startswith(("gpt:", "/gpt")), pass_bot=True, ) diff --git a/handlers/gemini.py b/handlers/gemini.py index fc22644..5a6490c 100644 --- a/handlers/gemini.py +++ b/handlers/gemini.py @@ -225,7 +225,6 @@ def gemini_audio_handler(message: Message, bot: TeleBot) -> None: start = time.time() for e in r: s += e.text - print(s) if time.time() - start > 1.7: start = time.time() bot_reply_markdown(reply_id, who, s, bot, split_text=False) diff --git a/handlers/sd.py b/handlers/sd.py new file mode 100644 index 0000000..cbd2a5b --- /dev/null +++ b/handlers/sd.py @@ -0,0 +1,70 @@ +from telebot import TeleBot +from telebot.types import Message +import requests +from os import environ + +from . import * + + +SD_API_KEY = environ.get("SD3_KEY") + +if SD_API_KEY is not None: + + def get_user_balance(): + api_host = "https://api.stability.ai" + url = f"{api_host}/v1/user/balance" + + response = requests.get(url, headers={"Authorization": f"Bearer {SD_API_KEY}"}) + + if response.status_code != 200: + print("Non-200 response: " + str(response.text)) + + # Do something with the payload... + payload = response.json() + return payload["credits"] + + def generate_sd3_image(prompt): + response = requests.post( + f"https://api.stability.ai/v2beta/stable-image/generate/sd3", + headers={"authorization": f"Bearer {SD_API_KEY}", "accept": "image/*"}, + files={"none": ""}, + data={ + "prompt": prompt, + "model": "sd3-turbo", + "output_format": "jpeg", + }, + ) + + if response.status_code == 200: + with open("sd3.jpeg", "wb") as file: + file.write(response.content) + return True + else: + print(str(response.json())) + return False + + def sd_handler(message: Message, bot: TeleBot): + """pretty sd3: /sd3
""" + credits = get_user_balance() + bot.reply_to( + message, + f"Generating pretty sd3-turbo image may take some time please left credits {credits} every try will cost 4 criedits wait:", + ) + m = message.text.strip() + prompt = m.strip() + try: + r = generate_sd3_image(prompt) + if r: + with open(f"sd3.jpeg", "rb") as photo: + bot.send_photo( + message.chat.id, photo, reply_to_message_id=message.message_id + ) + else: + bot.reply_to(message, "prompt error") + except Exception as e: + print(e) + bot.reply_to(message, "sd3 error") + + def register(bot: TeleBot) -> None: + bot.register_message_handler(sd_handler, commands=["sd3"], pass_bot=True) + bot.register_message_handler(sd_handler, regexp="^sd3:", pass_bot=True) diff --git a/requirements.txt b/requirements.txt index d8e27ee..e9e4923 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,5 @@ anthropic telegramify-markdown openai requests -urlextract \ No newline at end of file +urlextract +requests \ No newline at end of file