Update app.py

app.py

@@ -9,6 +9,12 @@ import logging
 import io
 import time
 from PIL import Image
+import torch
+from diffusers import (
+    StableDiffusionXLPipeline, 
+    KDPM2AncestralDiscreteScheduler,
+    AutoencoderKL
+)
 
 # Load environment variables from the .env file
 load_dotenv()
@@ -16,27 +22,15 @@ load_dotenv()
 DISCORD_TOKEN = os.getenv('dsTOK')
 HF_API_KEY = os.getenv('HFREAD')
 
-API_URL = "https://api-inference.huggingface.co/models/Corcelio/mobius"
-headers = {"Authorization": f"Bearer {HF_API_KEY}"}
-
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Function to query the Hugging Face model with a structured prompt
-def query(payload, max_retries=5, delay=5):
-    retries = 0
-    while retries < max_retries:
-        try:
-            response = requests.post(API_URL, headers=headers, json=payload)
-            response.raise_for_status()
-            return response.content
-        except requests.exceptions.RequestException as e:
-            logger.error(f"Error querying the API: {e}. Retrying in {delay} seconds...")
-            time.sleep(delay)
-            retries += 1
-    logger.error(f"Failed to query the API after {max_retries} attempts.")
-    return None
+# Load VAE component and configure the pipeline
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+pipe = StableDiffusionXLPipeline.from_pretrained("Corcelio/mobius", vae=vae, torch_dtype=torch.float16)
+pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+pipe.to('cuda')
 
 # Initialize the Discord bot with a default prefix
 intents = discord.Intents.default()
@@ -49,32 +43,38 @@ 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='generate')
+@bot.command(name='gen')
 @commands.cooldown(rate=1, per=10, type=commands.BucketType.user)
 async def generate(ctx, *description):
     """
-    Command to generate an image based on a description using the Hugging Face model.
+    Command to generate an image based on a description using the Stable Diffusion XL model.
     """
     if not description:
         await ctx.send("You need to provide a description for the image.")
         return
 
-    description = " ".join(description)
+    prompt = " ".join(description)
+    negative_prompt = "bad quality, low resolution, ugly, worse quality, grayscale, blurry"
+
     try:
-        payload = {"inputs": description}
-        image_bytes = query(payload)
+        # Generate the image
+        image = pipe(
+            prompt, 
+            negative_prompt=negative_prompt, 
+            width=512,
+            height=512,
+            guidance_scale=5,
+            num_inference_steps=25,
+            clip_skip=3
+        ).images[0]
 
-        if image_bytes:
-            image = Image.open(io.BytesIO(image_bytes))
-            with io.BytesIO() as image_binary:
-                image.save(image_binary, 'PNG')
-                image_binary.seek(0)
-                await ctx.send(file=discord.File(fp=image_binary, filename='generated_image.png'))
-        else:
-            await ctx.send("An error occurred while generating the image. Please try again later.")
+        with io.BytesIO() as image_binary:
+            image.save(image_binary, 'PNG')
+            image_binary.seek(0)
+            await ctx.send(file=discord.File(fp=image_binary, filename='generated_image.png'))
     except Exception as e:
         logger.error(f"Error processing the generate command: {e}")
-        await ctx.send("An error occurred while processing your request. Please try again later.")
+        await ctx.send("An error occurred while processing your request.")
 
 @generate.error
 async def generate_error(ctx, error):