Update app.py

app.py

@@ -1,6 +1,7 @@
 import discord
 from discord.ext import commands
 from gradio_client import Client
+import asyncio
 import multiprocessing
 import subprocess
 import sys
@@ -28,23 +29,35 @@ async def on_message(message):
     if message.channel.id != CHANNEL_ID and not isinstance(message.channel, discord.Thread):
         return
 
-    # Generate image using Gradio client
     try:
-        result = gradio_client.predict(
-            prompt=message.content,
-            seed=123,
-            randomize_seed=False,
-            width=1024,
-            height=576,
-            guidance_scale=5,
-            num_inference_steps=30,
-            api_name="/infer_t2i"
+        result = await asyncio.wait_for(
+            asyncio.to_thread(
+                gradio_client.predict,
+                prompt=message.content,
+                seed=123,
+                randomize_seed=False,
+                width=1024,
+                height=576,
+                guidance_scale=5,
+                num_inference_steps=30,
+                api_name="/infer_t2i"
+            ),
+            timeout=60  # 60 seconds timeout
         )
-        
-        # Send the generated image back to the user
         await message.channel.send(file=discord.File(result))
+    except asyncio.TimeoutError:
+        await message.channel.send("Image generation timed out. Please try again.")
     except Exception as e:
         await message.channel.send(f"An error occurred: {str(e)}")
+        print(f"Error in on_message: {str(e)}")
+
+@bot.event
+async def on_disconnect():
+    print("Bot disconnected. Attempting to reconnect...")
+
+@bot.event
+async def on_resume():
+    print("Bot connection resumed.")
 
 def run_web():
     subprocess.run([sys.executable, "web.py"])