feat: sd3 api

Signed-off-by: yihong0618 <zouzou0208@gmail.com>

.gitignore

@@ -160,6 +160,7 @@ cython_debug/
 #.idea/
 *.jpg
 *.png
+*.jpeg
 OUT_FOLDER/
 cache/
 nohup.out

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.
 

handlers/chatgpt.py

@@ -28,7 +28,7 @@ chatgpt_pro_player_dict = {}
 
 
 def chatgpt_handler(message: Message, bot: TeleBot) -> None:
-    """chatgpt : /chatgpt <question>"""
+    """gpt : /gpt <question>"""
     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 <question>"""
+    """gpt_pro : /gpt_pro <question>"""
     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, commands=["gpt"], pass_bot=True)
-    bot.register_message_handler(chatgpt_handler, regexp="^chatgpt:", 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
+        chatgpt_pro_handler, commands=["gpt_pro"], pass_bot=True
-    )
-    bot.register_message_handler(
-        chatgpt_pro_handler, regexp="^chatgpt_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,
     )

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)

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 <address>"""
+        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)

requirements.txt

@@ -7,4 +7,5 @@ anthropic
 telegramify-markdown
 openai
 requests
-urlextract
+urlextract
+requests