ADD choose model

.gitignore

@@ -1,2 +1,3 @@
 *.png
 .env
+*.session

main.py

@@ -1,30 +1,41 @@
-from pyrogram import Client, filters
-from pyrogram.types import *
-import os
 import json
 import requests
 import io
+import os
 import random
-from PIL import Image, PngImagePlugin
 import base64
+from PIL import Image, PngImagePlugin
+from pyrogram import Client, filters
+from pyrogram.types import *
+
+# Done! Congratulations on your new bot. You will find it at 
+# t.me/gootmornbot
+# You can now add a description, about section and profile picture for your bot, see /help for a list of commands. By the way, when you've finished creating your cool bot, ping our Bot Support if you want a better username for it. Just make sure the bot is fully operational before you do this.
+
+# Use this token to access the HTTP API:
+# Keep your token secure and store it safely, it can be used by anyone to control your bot.
+
+# For a description of the Bot API, see this page: https://core.telegram.org/bots/api
 
 API_ID = os.environ.get("API_ID", None) 
 API_HASH = os.environ.get("API_HASH", None) 
 TOKEN = os.environ.get("TOKEN", None) 
-SD_URL = os.environ.get("TOKEN", None) 
+SD_URL = os.environ.get("SD_URL", None) 
+print(SD_URL)
 
-bot = Client(
+app = Client(
     "stable",
     api_id=API_ID,
     api_hash=API_HASH,
     bot_token=TOKEN
 )
 
+
 @app.on_message(filters.command(["draw"]))
 def draw(client, message):
     msgs = message.text.split(' ', 1)
     if len(msgs) == 1:
-        message.reply_text("Format : /draw < text to anime image >")
+        message.reply_text("Format : /draw < text to image >")
         return
     msg = msgs[1]
 
@@ -34,6 +45,9 @@ def draw(client, message):
 
     r = requests.post(url=f'{SD_URL}/sdapi/v1/txt2img', json=payload).json()
 
+    # save the image to a file
+
+    def genr(): # generate random string
         chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
         chars1 = "1234564890"
         gen1 = random.choice(chars)
@@ -46,13 +60,14 @@ def draw(client, message):
         gen8 = random.choice(chars)
         gen9 = random.choice(chars)
         gen10 = random.choice(chars1)
-    word = f"{message.from_user.id}-MOE{gen1}{gen2}{gen3}{gen4}{gen5}{gen6}{gen7}{gen8}{gen9}{gen10}"
+        return f"{message.from_user.id}-MOE{gen1}{gen2}{gen3}{gen4}{gen5}{gen6}{gen7}{gen8}{gen9}{gen10}"
     
+    word=genr()
     for i in r['images']:
         image = Image.open(io.BytesIO(base64.b64decode(i.split(",", 1)[0])))
 
         png_payload = {"image": "data:image/png;base64," + i}
-        response2 = requests.post(url=f'{url}/sdapi/v1/png-info',
+        response2 = requests.post(url=f'{SD_URL}/sdapi/v1/png-info',
                                   json=png_payload)
 
         pnginfo = PngImagePlugin.PngInfo()
@@ -62,27 +77,66 @@ def draw(client, message):
         message.reply_photo(
             photo=f"{word}.png",
             caption=
-            f"Prompt - **{msg}**\n **[{message.from_user.first_name}-Kun](tg://user?id={message.from_user.id})**\n Join @WaifuAiSupport"
+            f"Prompt - **{msg}**\n **[{message.from_user.first_name}-Kun](tg://user?id={message.from_user.id})**"
         )
         os.remove(f"{word}.png")
         K.delete()
 
 
+# get the models, choose one, and set it as the model checkpoint
+@app.on_message(filters.command(["getmodels"]))
+async def get_models(client, message):
+    response = requests.get(url=f'{SD_URL}/sdapi/v1/sd-models')
+    if response.status_code == 200:
+        models_json = response.json()
+        # create buttons for each model name
+        buttons = []
+        for model in models_json:
+            buttons.append([InlineKeyboardButton(model['title'], callback_data=model['model_name'])])
+        # send the message
+        await message.reply_text(
+            text="Select a model [checkpoint] to use",
+            reply_markup=InlineKeyboardMarkup(buttons)
+        )
+
+@app.on_callback_query()
+async def process_callback(client, callback_query):
+    # if a model button is clicked, set sd_model_checkpoint to the selected model's title
+    sd_model_checkpoint = callback_query.data
+
+    # The sd_model_checkpoint needs to be set to the title from /sdapi/v1/sd-models
+    # post using /sdapi/v1/options
+
+    options = {
+        "sd_model_checkpoint": sd_model_checkpoint
+    }
+
+    # post the options
+    response = requests.post(url=f'{SD_URL}/sdapi/v1/options', json=options)
+    if response.status_code == 200:
+        # if the post was successful, send a message
+        await callback_query.message.reply_text("checpoint set to " + sd_model_checkpoint)
+    else:
+        # if the post was unsuccessful, send an error message
+        await callback_query.message.reply_text("Error setting options")
+
+
+# start the bot
 @app.on_message(filters.command(["start"], prefixes=["/", "!"]))
 async def start(client, message):
-    Photo = "https://media.discordapp.net/attachments/1028156834944655380/1062018608022171788/3aac7aaf-0065-40aa-9e4d-430c717b3d87.jpg"
+    # Photo = "https://i.imgur.com/79hHVX6.png"
 
     buttons = [[
         InlineKeyboardButton("Add to your group",
-                             url="http://t.me/botname?startgroup=true"),
+                             url="https://t.me/gootmornbot?startgroup=true")
-        InlineKeyboardButton("Channel", url="https://t.me/otakatsu"),
-        InlineKeyboardButton("Support", url="https://t.me/otakatsu_chat")
     ]]
-    await message.reply_photo(
+    await message.reply_text(
-        photo=Photo,
+
-        caption=
+        # photo=Photo,
-        f"Hello! I'm botname Ai and I can make an anime-styled picture!\n\n/generate - Reply to Image\n/draw text to anime image\n\nPowered by @Otakatsu",
+        text=
-        reply_markup=InlineKeyboardMarkup(buttons))
+        f"Hello!\nask me to imagine anything\n\n/draw text to image",
+        reply_markup=InlineKeyboardMarkup(buttons)
+        )
 
 
 app.run()