Update app.py

app.py

@@ -6,6 +6,8 @@ import discord
 from discord.ext import commands
 from dotenv import load_dotenv
 import logging
+import io
+from PIL import Image
 
 # Load environment variables from the .env file
 load_dotenv()
@@ -13,7 +15,7 @@ load_dotenv()
 DISCORD_TOKEN = os.getenv('dsTOK')
 HF_API_KEY = os.getenv('HFREAD')
 
-API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
+API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
 headers = {"Authorization": f"Bearer {HF_API_KEY}"}
 
 # Set up logging
@@ -21,29 +23,14 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Function to query the Hugging Face model with a structured prompt
-def query_huggingface(prompt, temperature=0.5, max_tokens=100, top_k=50, top_p=0.9):
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "temperature": temperature,
-            "max_new_tokens": max_tokens,
-            "top_k": top_k,
-            "top_p": top_p
-        }
-    }
+def query(payload):
     try:
         response = requests.post(API_URL, headers=headers, json=payload)
         response.raise_for_status()
-        result = response.json()
-        if isinstance(result, list) and len(result) > 0 and 'generated_text' in result[0]:
-            return result[0]['generated_text']
-        elif 'generated_text' in result:
-            return result['generated_text']
-        else:
-            return "No response generated."
+        return response.content
     except requests.exceptions.RequestException as e:
         logger.error(f"Error querying the API: {e}")
-        return f"Error: {str(e)}"
+        return None
 
 # Initialize the Discord bot with a default prefix
 intents = discord.Intents.default()
@@ -56,36 +43,37 @@ bot = commands.Bot(command_prefix=commands.when_mentioned_or(default_prefix), in
 async def on_ready():
     logger.info(f'Bot is ready. Logged in as {bot.user}')
 
-@bot.command(name='ask')
+@bot.command(name='generate')
 @commands.cooldown(rate=1, per=10, type=commands.BucketType.user)
-async def ask(ctx, *question):
+async def generate(ctx, *description):
     """
-    Command to ask a question to the Hugging Face model with an instructive prompt.
+    Command to generate an image based on a description using the Hugging Face model.
     """
-    if not question:
-        await ctx.send("You need to ask a question.")
+    if not description:
+        await ctx.send("You need to provide a description for the image.")
         return
 
-    question = " ".join(question)
+    description = " ".join(description)
     try:
-        # Create a structured prompt
-        prompt = f"Answer the following question accurately and concisely: {question}"
-        response = query_huggingface(prompt, temperature=0.8, max_tokens=1000, top_k=70, top_p=0.4)
-        
-        # Extract and clean the response
-        clean_response = response.replace(prompt, '').strip()
-        
-        await ctx.send(clean_response)
+        payload = {"inputs": description}
+        image_bytes = query(payload)
+
+        if image_bytes:
+            image = Image.open(io.BytesIO(image_bytes))
+            image.save("generated_image.png")
+            await ctx.send(file=discord.File("generated_image.png"))
+        else:
+            await ctx.send("An error occurred while generating the image.")
     except Exception as e:
-        logger.error(f"Error processing the ask command: {e}")
+        logger.error(f"Error processing the generate command: {e}")
         await ctx.send("An error occurred while processing your request.")
 
-@ask.error
-async def ask_error(ctx, error):
+@generate.error
+async def generate_error(ctx, error):
     if isinstance(error, commands.CommandOnCooldown):
         await ctx.send(f"This command is on cooldown. Please try again after {int(error.retry_after)} seconds.")
     else:
-        logger.error(f"Error in ask command: {error}")
+        logger.error(f"Error in generate command: {error}")
         await ctx.send("An error occurred. Please try again later.")
 
 @bot.command(name='setprefix')
@@ -98,7 +86,7 @@ async def set_prefix(ctx, prefix: str):
 async def help_custom(ctx):
     help_text = (
         "Here are the commands you can use:\n"
-        "!ask <question> - Ask a question to the AI model.\n"
+        "!generate <description> - Generate an image based on the description.\n"
         "!setprefix <prefix> - Change the command prefix (admin only).\n"
         "!help_custom - Display this help message."
     )
@@ -107,3 +95,4 @@ async def help_custom(ctx):
 # Run the bot
 bot.run(DISCORD_TOKEN)
 
+