app.py

import discord
import os
import threading
import gradio as gr
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



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')
dfif = Client("huggingface-projects/IF", DFIF_TOKEN)
#sdlu = Client("huggingface-projects/stable-diffusion-latent-upscaler", DFIF_TOKEN)




#---------------------------------------------------------------------------------------------------------------------------------------------- 

# Set up discord bot 
class MyClient(discord.Client):
    async def on_ready(self):
        print('Logged on as', self.user)

	#------------------------------------------
    async def on_message(self, message):
		ctx = await self.get_context(message)
        
        if message.content.startswith('!dfif'): #ex: !dfif sleepy cat
			prompt = string[5:]                 #ex: sleepy cat
			await dfif (ctx, prompt)            # pass prompt to dfif (for generation) and ctx (for keeping track of command msg)
			

    		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 = dfif.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)
    
                img1 = Image.open(first_png_path)
                img2 = Image.open(second_png_path)
                img3 = Image.open(third_png_path)
                img4 = Image.open(fourth_png_path)
    
                combined_image = Image.new('RGB', (img1.width * 2, img1.height * 2))
    
                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))
    
                combined_image_path = os.path.join(stage_1_results, 'combined_image.png')
                combined_image.save(combined_image_path)


            sent_message = await message.reply('Here is the combined image. React with the image number you want to upscale!', file=discord.File(f, 'combined_image.png'))


#---------------------------------------------------------------------------------------------------------------------------------------------- 




    
DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)
intents = discord.Intents.default()
intents.message_content = True
client = MyClient(intents=intents)

def run_bot():
    client.run(DISCORD_TOKEN)

threading.Thread(target=run_bot).start()

def greet(name):
    return "Hello " + name + "!"

demo = gr.Interface(fn=greet, inputs="text", outputs="text")
demo.launch()