Update bot.py

bot.py

@@ -1,5 +1,5 @@
-import discord, aiohttp, asyncio, datetime
-from discord import app_commands
+import discord, aiohttp, asyncio, datetime, re
+from discord import ui, app_commands
 from discord.ext import commands
 from discord.ext.commands import bot
 
@@ -8,7 +8,131 @@
 #chunk_guilds_at_startup is for not saving guilds datas in the cache (less ram consumption)
 #case_insensitive is for using command with full caps for exemple
 
-tree = client.tree #for slash commands
+tree = client.tree
+
+start_time = datetime.datetime.now(datetime.timezone.utc)
+
+async def generate_image(prompt, model, timeout=120):
+    counter = 0
+    api_url = "https://api.sitius.ir/"
+    data = {
+        "model": model,
+        "prompt": prompt,
+        "steps": 30,
+        "cfg_scale": 7,
+        "sampler" : "Euler",
+        "negative_prompt": "canvas frame, cartoon, 3d, ((disfigured)), ((bad art)), ((deformed)), ((extra limbs)), ((close up)), ((b&w)), weird colors, blurry, (((duplicate))), ((morbid)), ((mutilated)), [out of frame], extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), ((ugly)), (((bad proportions))), (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), (fused fingers), (too many fingers), (((long neck))), Photoshop, video game, tiling, poorly drawn feet, body out of frame, nsfw"
+    }
+    headers = {"auth": "test"}
+
+    async with aiohttp.ClientSession() as session:
+        async with session.post(f"{api_url}v1/generate/", json=data, headers=headers) as response:
+            job_id = await response.json()
+        while counter < timeout :
+            async with session.get(api_url + f"v1/image/{job_id}/") as r:
+                if r.status == 200:
+                    url = await r.json()
+                    return url
+                await asyncio.sleep(0.5)
+                counter += 0.5
+
+
+
+
+
+
+
+
+generating_embed = discord.Embed(
+        title="Generating Image",
+        description="Image Is Generating, It takes 30 ~ 50 seconds to generate",
+        color=discord.Color.orange()
+    )
+
+
+
+
+
+class RetryButton(ui.View):
+    def __init__(self, prompt, model, user, p1) -> None:
+        super().__init__()
+        self.prompt = prompt
+        self.model = model
+        self.user = user
+        self.p1 = p1
+
+
+    @discord.ui.button(label='retry', emoji="♻️", style=discord.ButtonStyle.green, custom_id='retry')
+    async def retry(self, interaction: discord.Interaction, button: ui.Button):
+        button.disabled = True
+        await interaction.message.edit(view=self)
+        await interaction.response.edit_message(content=None, embed=generating_embed)
+
+        image = await generate_image(self.prompt, self.model)
+
+
+        if image:
+            print(f"Prompt: {self.prompt}\nImage URL: {image}") 
+            image_embed = discord.Embed(title="Generated Image", color=discord.Color.blue())
+            image_embed.set_image(url=image)
+            image_embed.add_field(name="Prompt", value=self.prompt, inline=False)
+            image_embed.add_field(name="Model", value=self.model, inline=False)
+            button.disabled = False
+            await interaction.message.edit(view=self)
+            await interaction.edit_original_response(content=None, embed=image_embed)
+
+
+
+        else:
+            await interaction.response.edit_message(content="Failed to upload image.", embed=None)
+
+    async def interaction_check(self, interaction):
+        if interaction.user.id != self.p1:
+            await interaction.response.send_message(f"You do not own this command {interaction.user.mention}", ephemeral= True)
+        return interaction.user.id == self.p1
+
+
+
+
+async def check_words(prompt: str) -> bool:
+    banned_patterns = [
+        r"\bloli\b",
+        r"\bbaby\b",
+        r"\bshota\b",
+        r"\bunderage\b",
+        r"\bkid\b",
+        r"\bchild\b",
+        r"\blittle girl\b",
+        r"\byoung girl\b",
+        r"\bpetite\b",
+        r"\blittle boy\b",
+        r"\byoung boy\b",
+        r"\bteen\b",
+        r"\btween\b",
+        r"\bminor\b",
+        r"\badolescent\b",
+        r"\bpreteen\b",
+        r"\bsmall girl\b",
+        r"\bsmall boy\b",
+        # Match age patterns only if the age is below 18
+        r"\b(1[0-7]|[1-9])\s?yo\b",                # e.g., 2yo, 15yo (but not 18yo or older)
+        r"\b(1[0-7]|[1-9])\s?years?\s?old\b",       # e.g., 3 years old, 10 yrs old (but not 18 or older)
+        r"\b(1[0-7]|[1-9])\s?-year-old\b",          # e.g., 4-year-old, 7-year-old (but not 18 or older)
+    ]
+
+    prompt_lower = prompt.lower()
+    for pattern in banned_patterns:
+        if re.search(pattern, prompt_lower):
+            return True
+    return False
+
+
+
+
+
+
+
+
 
 @client.event
 async def on_ready():
@@ -21,119 +145,153 @@ async def update_stats():
         total_members = sum(guild.member_count for guild in client.guilds)
         activity_text = f"In {len(client.guilds)} servers | Serving {total_members} members"
         await client.change_presence(activity=discord.Game(name=activity_text))
-        await asyncio.sleep(360)
+        await asyncio.sleep(1800)
+
+@client.event
+async def on_message(message):
+    if message.author == client.user:
+        return
+
+    if message.channel.type == discord.ChannelType.private:
+        title = 'The Bot Can Only Be Used In Servers'
+    else:
+        title='Bot Is Now Based On / Commands Only'
 
+    response_embed = discord.Embed(title=title,
+            description=f"[Invite](https://discord.com/oauth2/authorize?client_id=1180105597874610206&permissions=277025515584&scope=bot) the bot for your server!\n[Support](https://discord.com/invite/bVUa2En9) our community server!",
+            color=discord.Color.blue())
+
+    if message.content == client.user.mention:
+        await message.channel.send(embed=response_embed)
 
-#generating embed
-generating_embed = discord.Embed(
-        title="Generating Image",
-        description="Image Is Generating, It takes 30 ~ 50 seconds to generate",
-        color=discord.Color.orange())
+    await client.process_commands(message)
 
+
+@tree.command(name = "stats", description = "Get bot statistics")
+async def stats(interaction: discord.Interaction):
+    total_servers = len(client.guilds)
+    total_members = sum(guild.member_count for guild in client.guilds)
+    uptime = datetime.datetime.now(datetime.timezone.utc) - start_time
+    total_seconds, microseconds = divmod(uptime.microseconds + (uptime.seconds + uptime.days * 86400) * 10 ** 6, 10 ** 6)
+    days = int(total_seconds // 86400)
+    uptime_string = f"{days} days, {int((total_seconds % 86400) // 3600)} hours, {int((total_seconds % 3600) // 60)} minutes, {int(total_seconds % 60)} seconds, and {int(microseconds)} microseconds"
+    embed = discord.Embed(
+        title = "Bot Statistics",
+        description = f"Here are the current statistics of the bot:\n\nTotal Servers: {total_servers}\nTotal Members: {total_members}\nUptime: {uptime_string}",
+        color = discord.Color.blue()
+    )
+    await interaction.response.send_message(embed=embed)
+
+
+@client.command(name="machine", description="Get bot's machine statistics")
+async def machine(ctx):
+    try:
+        ram = psutil.virtual_memory()
+        cpu = psutil.cpu_percent()
+        disk = psutil.disk_usage('/')
 
-#define the api header
-async def head(prompt, image_format, model):
-    width , height = map(int, image_format.split("x"))
-    api_key = ""#change by your api key on the bot from visioncraft : https://t.me/VisionCraft_bot (use the command /key)
-    data = {
-        "model": model,
-        "prompt": prompt,
-        "token": api_key,
-        "width": width,
-        "height": height
-    }
-    return data
+        performance_rate = int((ram.percent + cpu + disk.percent) / 3)
 
-#make the request to the api
-async def generate_image(data) -> bytes:
-    api_url = "https://visioncraft.top"
-    async with aiohttp.ClientSession() as session:
-        async with session.post(f"{api_url}/sd", json=data) as response:
-            image = await response.read()
-            return image
+        embed = discord.Embed(title="Machine Stats", color=discord.Color.blue())
+        embed.add_field(name="Free RAM", value=f"{ram.available / (1024 ** 3):.2f} GB / {ram.total / (1024 ** 3):.2f} GB", inline=True)
+        embed.add_field(name="Free CPU", value=f"{100 - cpu:.2f}%", inline=True)
+        embed.add_field(name="Free Disk Space", value=f"{disk.free / (1024 ** 3):.2f} GB / {disk.total / (1024 ** 3):.2f} GB", inline=True)
+        embed.add_field(name="Performance Rate", value=f"{performance_rate}% out of 100%", inline=False)
 
-#upload the image on a file host
-async def telegraph_file_upload(file_bytes):
-    url = 'https://telegra.ph/upload'
-    try:
-        data = aiohttp.FormData()
-        data.add_field('file', file_bytes, filename='image.png', content_type='image/png')
-
-        async with aiohttp.ClientSession() as session:
-            async with session.post(url, data=data) as response:
-                response_data = await response.json()
-                if response.status == 200 and response_data and isinstance(response_data, list) and 'src' in response_data[0]:
-                    telegraph_url = response_data[0]['src']
-                    full_url = f'https://telegra.ph{telegraph_url}'
-                    return full_url
-                else:
-                    print(f'Unexpected response data or status: {response_data}, Status: {response.status}')
-                    return None
+        await ctx.send(embed=embed)
     except Exception as e:
-        print(f'Error uploading file to Telegraph: {str(e)}')
-        return None
+        await ctx.send(f"An error occurred: {e}")
+
+@client.event
+async def on_guild_join(guild):
+    # Details about the server
+    owner = guild.owner
+    server_name = guild.name
+    if owner:  # Ensure there is an owner object
+        owner_name = f"{owner.name}#{owner.discriminator}"  # The owner's username and discriminator
+        owner_id = owner.id  # The owner's ID
+
+        # Construct the message
+        message_content = f"{owner_name} (ID: {owner_id}) added the bot to their server {server_name}."
+
+        # Fetch your user object by your Discord user ID
+        user = await client.fetch_user(761886210120744990)  # Your Discord ID
+
+        # Send the message
+        if user:
+            await user.send(message_content)
+        else:
+            print("Failed to fetch user for DM.")
+    else:
+        print("Failed to retrieve owner information.")
+
+
+
 
-#simple command
 @client.command()
 async def imagine(ctx, *, prompt: str):
     message = await ctx.send(embed=generating_embed)
+    if await check_words(prompt):
+        return await message.edit(content="your prompt contain banned word", embed=None)
+    image = await generate_image(prompt, model="Realistic_Vision_v5.0")
 
-    data = await head(prompt, image_format="1024x1024", model="RealVisXLV40Turbo-40")
-    image = await generate_image(data)
-    telegraph_url = await telegraph_file_upload(image)
 
-    if telegraph_url:
+    if image:
+        view=RetryButton(prompt, model, ctx.user.name, ctx.user.id)
+
         image_embed = discord.Embed(title="Generated Image", color=discord.Color.blue())
-        image_embed.set_image(url=telegraph_url)
+        image_embed.set_image(url=image)
+        image_embed.add_field(name="Prompt", value=prompt, inline=False)
+        image_embed.add_field(name="Model", value="Realistic_Vision_v5.0", inline=False)
 
-        await message.edit(content=None, embed=image_embed)
-        print(f"Image URL: {telegraph_url}")  # Log the image URL to the console
+        await message.edit(content=None, embed=image_embed, view = view)
+        print(f"Image URL: {image}")  # Log the image URL to the console
     else:
         await message.edit(content="Failed to upload image.", embed=None)
 
 
 
-#slash command
+#how to use
 @tree.command(description="generate images by ai")
-@app_commands.describe(prompt="description of the image", private="make the image visble for you or everyone", image_format="the format of the image", image_type="type of the image like anime, 3d, etc...")
+@app_commands.describe(prompt="description of the image", private="make the image visble for you or everyone", image_type="type of the image like anime, 3d, etc...")
 @app_commands.choices(
-    image_format = [
-    app_commands.Choice(name="default・Square", value="1024x1024"),
-    app_commands.Choice(name="Portrait", value="768x1024"),
-    app_commands.Choice(name="Landscape", value="1024x768"),
-],
     private = [
     app_commands.Choice(name="True", value="True"),
     app_commands.Choice(name="False", value="False"),
     ],
     image_type = [
-        app_commands.Choice(name="default・realistic", value="RealVisXLV40Turbo-40"),
-        app_commands.Choice(name="3D", value="UnrealXL-v1"),
-        app_commands.Choice(name="anime", value="AnimefromHaDeS-v16HQ"),
-    ]
+        app_commands.Choice(name="default・realistic", value="Realistic_Vision_v5.0"),
+        app_commands.Choice(name="3D", value="lyriel_v1.6"),
+        app_commands.Choice(name="anime", value="anything_V5"),
+
+        ]
+    
 )
-async def imagine(ctx : discord.Interaction, prompt: str, image_format: app_commands.Choice[str] = None, image_type: app_commands.Choice[str] = None, private: app_commands.Choice[str] = None):
+async def imagine(ctx : discord.Interaction, prompt: str, image_type: app_commands.Choice[str] = None, private: app_commands.Choice[str] = None):
     private_value = True if private is not None and private.value == "True" else False 
-    img_format = image_format.value if image_format is not None else "1024x1024"
-    model = image_type.value if image_type is not None else "RealVisXLV40Turbo-40"
+    model = image_type.value if image_type is not None else "Realistic_Vision_v5.0"
 
     await ctx.response.send_message(embed=generating_embed, ephemeral=bool(private_value))
-
+    if await check_words(prompt):
+        return await ctx.edit_original_response(content="your prompt contain banned word", embed=None)
 
-    data = await head(prompt, img_format, model)
-    image = await generate_image(data)
-    telegraph_url = await telegraph_file_upload(image)
+
+    image = await generate_image(prompt, model)
+
 
-    if telegraph_url:
+    if image:
+        view=RetryButton(prompt, model, ctx.user.name, ctx.user.id)
+        print(f"Prompt: {prompt}\nImage URL: {image}") 
         image_embed = discord.Embed(title="Generated Image", color=discord.Color.blue())
-        image_embed.set_image(url=telegraph_url)
-        print(f"Image URL: {telegraph_url}")
-        await ctx.edit_original_response(content=None, embed=image_embed)
+        image_embed.set_image(url=image)
+        image_embed.add_field(name="Prompt", value=prompt, inline=False)
+        image_embed.add_field(name="Model", value=model, inline=False)
+        await ctx.edit_original_response(content=None, embed=image_embed, view=view)
 
 
     else:
+
         await ctx.edit_original_response(content="Failed to upload image.", embed=None)
 
-
 
 client.run("")#use your bot token