diff --git a/shizuchat/modules/Draw.py b/shizuchat/modules/Draw.py new file mode 100644 index 0000000..9f8e436 --- /dev/null +++ b/shizuchat/modules/Draw.py @@ -0,0 +1,43 @@ +from pyrogram import Client, filters +import requests +from shizuchat import shizuchat as app + +# Store user state to track if they are awaiting a query for the draw command +user_states = {} + +# Define the handler for /draw Command. +@app.on_message(filters.command("draw")) +async def ask_for_query(client, message): + # Check if a query is provided directly with the /draw command + if len(message.command) == 1: # If no query is provided + # Set the user state to 'awaiting query' + user_states[message.from_user.id] = 'awaiting_query' + await message.reply("**Please provide a query to generate an image.**") + else: + # If the query is provided, process it directly + query = " ".join(message.command[1:]) + await generate_image(client, message, query) + + +# Function to generate image based on the query +async def generate_image(client, message, query): + # Make a request to the text2img API + url = f"https://text2img.codesearch.workers.dev/prompt={query}" + try: + response = requests.get(url) + + # Check if the response is successful + if response.status_code == 200: + data = response.json() + image_url = data.get("image") + if image_url: + # Send the image with a caption mentioning the user + caption = f"**Image generated by {message.from_user.mention}**" + await message.reply_photo(photo=image_url, caption=caption) + else: + await message.reply("**Sorry, I couldn't generate an image for your query.**") + else: + await message.reply(f"**Error: {response.status_code} - {response.text}. Please try again later.**") + except Exception as e: + await message.reply(f"**An error occurred: {str(e)}. Please try again later.**") +