revert to stable MyClient class, debug Error: None

app.py

@@ -1,183 +1,160 @@
 import discord
+import gradio_client
+from gradio_client import Client
+import gradio as gr
 import os
 import threading
-import gradio as gr
+
+#for deepfloydif
 import requests
 import json
 import random
-import time
-import re
-from discord import Embed, Color
-from discord.ext import commands
-# test
-from gradio_client import Client
 from PIL import Image
-#from ratelimiter import RateLimiter
-#
-from datetime import datetime
-from pytz import timezone
-#
-import asyncio
-
-breaking temporarily
-
-
-zurich_tz = timezone("Europe/Zurich")
-
-def convert_to_timezone(dt, tz):
-    return dt.astimezone(tz).strftime("%Y-%m-%d %H:%M:%S %Z")
-
-DFIF_TOKEN = os.getenv('HF_TOKEN')
-df = Client("huggingface-projects/IF", DFIF_TOKEN)
-sdlu = Client("huggingface-projects/stable-diffusion-latent-upscaler", DFIF_TOKEN)
-
-DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)
-intents = discord.Intents.default()
-intents.message_content = True
-
-bot = commands.Bot(command_prefix='!', intents=intents)
-
+import matplotlib.pyplot as plt
+import matplotlib.image as mpimg
+import time
 
-#----------------------------------------------------------------------------------------------------------------------------------------------
+# random + small llama #
 
-@bot.event
-async def on_ready():
-    print('Logged on as', bot.user)
-    bot.log_channel = bot.get_channel(1100458786826747945) # 1100458786826747945 = bot-test, 1107006391547342910 = lunarbot server
 
-#---------------------------------------------------------------------------------------------------------------------------------------------- 
-# Stage 1
-@bot.command()
-async def deepfloydif(ctx, *, prompt: str):
-    try:
-        prompt = prompt.strip()[:100]  # Limit the prompt length to 100 characters
-        prompt = re.sub(r'[^\w\s]', '', prompt)  # Remove special characters
+#todos
+#alert
+#fix error on first command on bot startup
+#stable diffusion upscale
+#buttons for deepfloydIF (1,2,3,4)
+#application commands instead of message content checks (more user-friendly)
 
-        def check_reaction(reaction, user):
-            return user == ctx.author and str(reaction.emoji) in ['1️⃣', '2️⃣', '3️⃣', '4️⃣']   
 
-        await ctx.message.add_reaction('👍')
-        thread = await ctx.message.create_thread(name=f'{ctx.author} Image Upscaling Thread ')
-        # create thread -> send new message inside thread + combined_image -> add reactions -> dfif2
-        await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...')
-        
-        number_of_images = 4
-        current_time = int(time.time())
-        random.seed(current_time)
-        seed = random.randint(0, 2**32 - 1)
-        stage_1_results, stage_1_param_path, stage_1_result_path = df.predict(prompt, "blur", seed, number_of_images, 7.0, 'smart100', 50, api_name="/generate64")
-        png_files = [f for f in os.listdir(stage_1_results) if f.endswith('.png')]
 
-        if png_files:
-            first_png = png_files[0]
-            second_png = png_files[1]
-            third_png = png_files[2]
-            fourth_png = png_files[3]
 
-            first_png_path = os.path.join(stage_1_results, first_png)
-            second_png_path = os.path.join(stage_1_results, second_png)
-            third_png_path = os.path.join(stage_1_results, third_png)
-            fourth_png_path = os.path.join(stage_1_results, fourth_png)
+DFIF_TOKEN = os.getenv('DFIF_TOKEN')
 
-            img1 = Image.open(first_png_path)
-            img2 = Image.open(second_png_path)
-            img3 = Image.open(third_png_path)
-            img4 = Image.open(fourth_png_path)
+#deepfloydIF
+#df = Client("DeepFloyd/IF", DFIF_TOKEN) #not reliable at the moment
+df = Client("huggingface-projects/IF", DFIF_TOKEN)
 
-            combined_image = Image.new('RGB', (img1.width * 2, img1.height * 2))
+#stable diffusion upscaler
+sdlu = Client("huggingface-projects/stable-diffusion-latent-upscaler", DFIF_TOKEN)
 
-            combined_image.paste(img1, (0, 0))
-            combined_image.paste(img2, (img1.width, 0))
-            combined_image.paste(img3, (0, img1.height))
-            combined_image.paste(img4, (img1.width, img1.height))
+# Set up discord bot 
+class MyClient(discord.Client):
+    async def on_ready(self):
+        print('Logged on as', self.user)
 
-            combined_image_path = os.path.join(stage_1_results, 'combined_image.png')
-            combined_image.save(combined_image_path)
 
-            # Trigger the second stage prediction
-            #await dfif2(ctx, stage_1_result_path)
+    async def on_message(self, message):
 
+        #safety checks----------------------------------------------------------------------------------------------------
 
-        #1
-        #sent_message = None
+        #   tldr, bot should run if
+        #1) it does not have @offline role
+        #2) user has @verified role
+        #3) bot is in #bot-test channel
+        
+        # bot won't respond to itself, prevents feedback loop + API spam
+        if message.author == self.user:
+            return       
+
+        # if the bot has this role, it won't run
+        OFFLINE_ROLE_ID = 1103676632667017266  # 1103676632667017266 = @offline / under maintenance
+        guild = message.guild
+        bot_member = guild.get_member(self.user.id)
+        if any(role.id == OFFLINE_ROLE_ID for role in bot_member.roles):
+            return
+        
+        # the message author needs this role in order to use the bot
+        REQUIRED_ROLE_ID = 897376942817419265 # 900063512829755413 = @verified, 897376942817419265 = @huggingfolks
+        if not any(role.id == REQUIRED_ROLE_ID for role in message.author.roles):
+            return            
 
-        #2
-        # create thread -> send new message inside thread + combined_image -> add reactions -> dfif2
-        #sent_message = await ctx.reply('Creating thread...')
+        # channels where bot will accept commands
+        ALLOWED_CHANNEL_IDS = [1100458786826747945] # 1100458786826747945 = #bot-test
+        if message.channel.id not in ALLOWED_CHANNEL_IDS:
+            return
 
-        #3
-        #ctx.message
-        # could replace sent_message with initial user command message? try later
-        #thread = await ctx.message.create_thread(name='Image Upscaling Thread')
+        #deepfloydif----------------------------------------------------------------------------------------------------
+            
+        if message.content.startswith('!deepfloydif'): # change to application commands, more intuitive
 
-        
-        #thread = await sent_message.create_thread(name='Image Upscaling Thread')
+            #(prompt, negative_prompt, seed, number_of_images, guidance_scale,custom_timesteps_1, number_of_inference_steps, api_name="/generate64") 
+            #-> (stage_1_results, stage_1_param_path, stage_1_result_path)
 
-        #4
-        with open(combined_image_path, 'rb') as f:
-            threadmsg = await thread.send(f'{ctx.author.mention}React with the image number you want to upscale!', file=discord.File(f, 'combined_image.png'))          
+            # input prompt
+            prompt = message.content[12:].strip()  
             
-        # bot reacts with appropriate emojis to the post, both showing the user what options they have,
-        # as well as showing which post to react to.
-        for emoji in ['1️⃣', '2️⃣', '3️⃣', '4️⃣']:
-            await threadmsg.add_reaction(emoji)   
-
-        reaction, user = await bot.wait_for('reaction_add', check=check_reaction)
-        if str(reaction.emoji) == '1️⃣':
-            await thread.send(f"{ctx.author.mention}Upscaling the first image...")
-            index = 0
-            await dfif2(ctx, index, stage_1_result_path, thread)
-        elif str(reaction.emoji) == '2️⃣':
-            await thread.send(f"{ctx.author.mention}Upscaling the second image...")
-            index = 1
-            await dfif2(ctx, index, stage_1_result_path, thread)            
-        elif str(reaction.emoji) == '3️⃣':
-            await thread.send(f"{ctx.author.mention}Upscaling the third image...")
-            index = 2
-            await dfif2(ctx, index, stage_1_result_path, thread)            
-        elif str(reaction.emoji) == '4️⃣':
-            await thread.send(f"{ctx.author.mention}Upscaling the fourth image...")
-            index = 3
-            await dfif2(ctx, index, stage_1_result_path, thread)            
-
-    #deepfloydif try/except
-    except Exception as e:
-        print(f"Error: {e}")
-        await ctx.reply('An error occurred while processing your request. Please wait 5 seconds before retrying.')
-        await ctx.message.add_reaction('❌')
-
-#----------------------------------------------------------------------------------------------------------------------------
-# Stage 2
-async def dfif2(ctx, index: int, stage_1_result_path, thread):
-    try:
-        selected_index_for_stage_2 = index
-        seed_2 = 0
-        guidance_scale_2 = 4
-        custom_timesteps_2 = 'smart50'
-        number_of_inference_steps_2 = 50
-        result_path = df.predict(stage_1_result_path, selected_index_for_stage_2, seed_2, 
-                                 guidance_scale_2, custom_timesteps_2, number_of_inference_steps_2, api_name='/upscale256')
-    
-        
-        with open(result_path, 'rb') as f:
-            await thread.send(f'{ctx.author.mention}Here is the upscaled image! :) ', file=discord.File(f, 'result.png'))
-            #await ctx.reply('Here is the result of the second stage', file=discord.File(f, 'result.png'))
-            await ctx.message.add_reaction('✔️')
+            negative_prompt = ''
+            seed = 0
+            number_of_images = 4
+            guidance_scale = 7
+            custom_timesteps_1 = 'smart50'
+            number_of_inference_steps = 50
             
+            stage_1_results, stage_1_param_path, stage_1_result_path = df.predict(
+                prompt,
+                negative_prompt,
+                seed,
+                number_of_images,
+                guidance_scale,
+                custom_timesteps_1,
+                number_of_inference_steps,
+                api_name='/generate64')
+
+            #stage_1_results, stage_1_param_path, stage_1_result_path = df.predict("gradio written on a wall", "blur", 1,1,7.0, 'smart100',50, api_name="/generate64")
+            
+            # stage_1_results -> path to directory with png files, so we isolate those
+            png_files = [f for f in os.listdir(stage_1_results) if f.endswith('.png')]
+
+            # merge images into larger, 2x2 image the way midjourney does it
+            if png_files:
+                first_png = png_files[0]
+                second_png = png_files[1]
+                third_png = png_files[2]
+                fourth_png = png_files[3]
+
+                '''
+                [],[],[],[] -> [][]
+                               [][]
+                
+                '''
+            
+                first_png_path = os.path.join(stage_1_results, first_png)
+                second_png_path = os.path.join(stage_1_results, second_png)
+                third_png_path = os.path.join(stage_1_results, third_png)
+                fourth_png_path = os.path.join(stage_1_results, fourth_png)
+            
+                img1 = Image.open(first_png_path)
+                img2 = Image.open(second_png_path)
+                img3 = Image.open(third_png_path)
+                img4 = Image.open(fourth_png_path)
+            
+                # create a new blank image with the size of the combined images (2x2)
+                combined_image = Image.new('RGB', (img1.width * 2, img1.height * 2))
+            
+                # paste the individual images into the combined image
+                combined_image.paste(img1, (0, 0))
+                combined_image.paste(img2, (img1.width, 0))
+                combined_image.paste(img3, (0, img1.height))
+                combined_image.paste(img4, (img1.width, img1.height))
+            
+                # save the combined image
+                combined_image_path = os.path.join(stage_1_results, 'combined_image.png')
+                combined_image.save(combined_image_path)
+            
+            # send the combined image file as a discord attachment
+            with open(combined_image_path, 'rb') as f:
+                await message.reply('Here is the combined image', file=discord.File(f, 'combined_image.png'))
 
-        
-
-    except Exception as e:
-        print(f"Error: {e}")
-        await ctx.reply('An error occurred while processing stage 2 upscaling. Please try again later.')
-        await ctx.message.add_reaction('❌')
-#----------------------------------------------------------------------------------------------------------------------------
 
+       
+            
 
-        
+DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)
+intents = discord.Intents.default()
+intents.message_content = True
+client = MyClient(intents=intents)
 
 def run_bot():
-    bot.run(DISCORD_TOKEN)
+  client.run(DISCORD_TOKEN)
 
 threading.Thread(target=run_bot).start()
 
@@ -185,4 +162,4 @@ def greet(name):
     return "Hello " + name + "!"
 
 demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+demo.launch()