Spaces:

huggingface-projects
/

huggingbots

Paused

huggingbots / deepfloydif.py

lunarflu

Synced repo using 'sync_with_huggingface' Github Action

4decbdc about 3 years ago

10.4 kB

import asyncio
import glob
import os
import pathlib
import random

import discord
from gradio_client import Client
from PIL import Image

from discord.ui import Button, View

HF_TOKEN = os.getenv("HF_TOKEN")
deepfloydif_client = Client("huggingface-projects/IF", HF_TOKEN)

BOT_USER_ID = 1086256910572986469 if os.getenv("TEST_ENV", False) else 1102236653545861151
DEEPFLOYDIF_CHANNEL_ID = 1121834257959092234 if os.getenv("TEST_ENV", False) else 1119313215675973714


def deepfloydif_generate64_inference(prompt):
    """Generates four images based on a prompt"""
    negative_prompt = ""
    seed = random.randint(0, 1000)
    number_of_images = 4
    guidance_scale = 7
    custom_timesteps_1 = "smart50"
    number_of_inference_steps = 50
    (
        stage_1_images,
        stage_1_param_path,
        path_for_upscale256_upscaling,
    ) = deepfloydif_client.predict(
        prompt,
        negative_prompt,
        seed,
        number_of_images,
        guidance_scale,
        custom_timesteps_1,
        number_of_inference_steps,
        api_name="/generate64",
    )
    return [stage_1_images, stage_1_param_path, path_for_upscale256_upscaling]


def deepfloydif_upscale256_inference(index, path_for_upscale256_upscaling):
    """Upscales one of the images from deepfloydif_generate64_inference based on the chosen index"""
    selected_index_for_upscale256 = index
    seed_2 = 0
    guidance_scale_2 = 4
    custom_timesteps_2 = "smart50"
    number_of_inference_steps_2 = 50
    result_path = deepfloydif_client.predict(
        path_for_upscale256_upscaling,
        selected_index_for_upscale256,
        seed_2,
        guidance_scale_2,
        custom_timesteps_2,
        number_of_inference_steps_2,
        api_name="/upscale256",
    )
    return result_path


def deepfloydif_upscale1024_inference(index, path_for_upscale256_upscaling, prompt):
    """Upscales to stage 2, then stage 3"""
    selected_index_for_upscale256 = index
    seed_2 = 0  # default seed for stage 2 256 upscaling
    guidance_scale_2 = 4  # default for stage 2
    custom_timesteps_2 = "smart50"  # default for stage 2
    number_of_inference_steps_2 = 50  # default for stage 2
    negative_prompt = ""  # empty (not used, could add in the future)

    seed_3 = 0  # default for stage 3 1024 upscaling
    guidance_scale_3 = 9  # default for stage 3
    number_of_inference_steps_3 = 40  # default for stage 3

    result_path = deepfloydif_client.predict(
        path_for_upscale256_upscaling,
        selected_index_for_upscale256,
        seed_2,
        guidance_scale_2,
        custom_timesteps_2,
        number_of_inference_steps_2,
        prompt,
        negative_prompt,
        seed_3,
        guidance_scale_3,
        number_of_inference_steps_3,
        api_name="/upscale1024",
    )
    return result_path


def load_image(png_files, stage_1_images):
    """Opens images as variables so we can combine them later"""
    results = []
    for file in png_files:
        png_path = os.path.join(stage_1_images, file)
        results.append(Image.open(png_path))
    return results


def combine_images(png_files, stage_1_images, partial_path):
    if os.environ.get("TEST_ENV") == "True":
        print("Combining images for deepfloydif_generate64")
    images = load_image(png_files, stage_1_images)
    combined_image = Image.new("RGB", (images[0].width * 2, images[0].height * 2))
    combined_image.paste(images[0], (0, 0))
    combined_image.paste(images[1], (images[0].width, 0))
    combined_image.paste(images[2], (0, images[0].height))
    combined_image.paste(images[3], (images[0].width, images[0].height))
    combined_image_path = os.path.join(stage_1_images, f"{partial_path}.png")
    combined_image.save(combined_image_path)
    return combined_image_path


async def deepfloydif_generate64(ctx, prompt, client):
    """DeepfloydIF command (generate images with realistic text using slash commands)"""
    try:
        if ctx.author.id != BOT_USER_ID:
            if ctx.channel.id == DEEPFLOYDIF_CHANNEL_ID:
                if os.environ.get("TEST_ENV") == "True":
                    print("Safety checks passed for deepfloydif_generate64")
                channel = client.get_channel(DEEPFLOYDIF_CHANNEL_ID)
                # interaction.response message can't be used to create a thread, so we create another message
                message = await ctx.send(f"**{prompt}** - {ctx.author.mention} <a:loading:1114111677990981692>")

                loop = asyncio.get_running_loop()
                result = await loop.run_in_executor(None, deepfloydif_generate64_inference, prompt)
                stage_1_images = result[0]
                path_for_upscale256_upscaling = result[2]

                partial_path = pathlib.Path(path_for_upscale256_upscaling).name
                png_files = list(glob.glob(f"{stage_1_images}/**/*.png"))

                if png_files:
                    await message.delete()
                    combined_image_path = combine_images(png_files, stage_1_images, partial_path)
                    if os.environ.get("TEST_ENV") == "True":
                        print("Images combined for deepfloydif_generate64")

                    with Image.open(combined_image_path) as img:
                        width, height = img.size
                        new_width = width * 3
                        new_height = height * 3
                        resized_img = img.resize((new_width, new_height))
                        x2_combined_image_path = combined_image_path
                        resized_img.save(x2_combined_image_path)

                    # making image bigger, more readable
                    with open(x2_combined_image_path, "rb") as f:  # was combined_image_path
                        button1 = Button(custom_id="0", emoji="↖")
                        button2 = Button(custom_id="1", emoji="↗")
                        button3 = Button(custom_id="2", emoji="↙")
                        button4 = Button(custom_id="3", emoji="↘")

                        async def button_callback(interaction):
                            index = int(interaction.data["custom_id"])  # 0,1,2,3

                            await interaction.response.send_message(
                                f"{interaction.user.mention} <a:loading:1114111677990981692>", ephemeral=True
                            )
                            result_path = await deepfloydif_upscale256(index, path_for_upscale256_upscaling)

                            # create and use upscale 1024 button
                            with open(result_path, "rb") as f:
                                upscale1024 = Button(
                                    label="High-quality upscale (x4)", custom_id=str(index)
                                )  # "0", "1" etc
                                upscale1024.callback = upscale1024_callback
                                view = View(timeout=None)
                                view.add_item(upscale1024)

                                await interaction.delete_original_response()
                                await channel.send(
                                    content=(
                                        f"{interaction.user.mention} Here is the upscaled image! Click the button"
                                        " to upscale even more!"
                                    ),
                                    file=discord.File(f, f"{prompt}.png"),
                                    view=view,
                                )

                        async def upscale1024_callback(interaction):
                            index = int(interaction.data["custom_id"])

                            await interaction.response.send_message(
                                f"{interaction.user.mention} <a:loading:1114111677990981692>", ephemeral=True
                            )
                            result_path = await deepfloydif_upscale1024(index, path_for_upscale256_upscaling, prompt)

                            with open(result_path, "rb") as f:
                                await interaction.delete_original_response()
                                await channel.send(
                                    content=f"{interaction.user.mention} Here's your high-quality x16 image!",
                                    file=discord.File(f, f"{prompt}.png"),
                                )

                        button1.callback = button_callback
                        button2.callback = button_callback
                        button3.callback = button_callback
                        button4.callback = button_callback

                        view = View(timeout=None)
                        view.add_item(button1)
                        view.add_item(button2)
                        view.add_item(button3)
                        view.add_item(button4)

                        # could store this message as combined_image_dfif in case it's useful for future testing
                        await channel.send(
                            f"**{prompt}** - {ctx.author.mention} Click a button to upscale! (make larger + enhance"
                            " quality)",
                            file=discord.File(f, f"{partial_path}.png"),
                            view=view,
                        )
                else:
                    await ctx.send(f"{ctx.author.mention} No PNG files were found, cannot post them!")

    except Exception as e:
        print(f"Error: {e}")


async def deepfloydif_upscale256(index: int, path_for_upscale256_upscaling):
    """upscaling function for images generated using /deepfloydif"""
    try:
        loop = asyncio.get_running_loop()
        result_path = await loop.run_in_executor(
            None, deepfloydif_upscale256_inference, index, path_for_upscale256_upscaling
        )
        return result_path

    except Exception as e:
        print(f"Error: {e}")


async def deepfloydif_upscale1024(index: int, path_for_upscale256_upscaling, prompt):
    """upscaling function for images generated using /deepfloydif"""
    try:
        loop = asyncio.get_running_loop()
        result_path = await loop.run_in_executor(
            None, deepfloydif_upscale1024_inference, index, path_for_upscale256_upscaling, prompt
        )
        return result_path

    except Exception as e:
        print(f"Error: {e}")

	import asyncio
	import glob
	import os
	import pathlib
	import random

	import discord
	from gradio_client import Client
	from PIL import Image

	from discord.ui import Button, View

	HF_TOKEN = os.getenv("HF_TOKEN")
	deepfloydif_client = Client("huggingface-projects/IF", HF_TOKEN)

	BOT_USER_ID = 1086256910572986469 if os.getenv("TEST_ENV", False) else 1102236653545861151
	DEEPFLOYDIF_CHANNEL_ID = 1121834257959092234 if os.getenv("TEST_ENV", False) else 1119313215675973714


	def deepfloydif_generate64_inference(prompt):
	"""Generates four images based on a prompt"""
	negative_prompt = ""
	seed = random.randint(0, 1000)
	number_of_images = 4
	guidance_scale = 7
	custom_timesteps_1 = "smart50"
	number_of_inference_steps = 50
	(
	stage_1_images,
	stage_1_param_path,
	path_for_upscale256_upscaling,
	) = deepfloydif_client.predict(
	prompt,
	negative_prompt,
	seed,
	number_of_images,
	guidance_scale,
	custom_timesteps_1,
	number_of_inference_steps,
	api_name="/generate64",
	)
	return [stage_1_images, stage_1_param_path, path_for_upscale256_upscaling]


	def deepfloydif_upscale256_inference(index, path_for_upscale256_upscaling):
	"""Upscales one of the images from deepfloydif_generate64_inference based on the chosen index"""
	selected_index_for_upscale256 = index
	seed_2 = 0
	guidance_scale_2 = 4
	custom_timesteps_2 = "smart50"
	number_of_inference_steps_2 = 50
	result_path = deepfloydif_client.predict(
	path_for_upscale256_upscaling,
	selected_index_for_upscale256,
	seed_2,
	guidance_scale_2,
	custom_timesteps_2,
	number_of_inference_steps_2,
	api_name="/upscale256",
	)
	return result_path


	def deepfloydif_upscale1024_inference(index, path_for_upscale256_upscaling, prompt):
	"""Upscales to stage 2, then stage 3"""
	selected_index_for_upscale256 = index
	seed_2 = 0 # default seed for stage 2 256 upscaling
	guidance_scale_2 = 4 # default for stage 2
	custom_timesteps_2 = "smart50" # default for stage 2
	number_of_inference_steps_2 = 50 # default for stage 2
	negative_prompt = "" # empty (not used, could add in the future)

	seed_3 = 0 # default for stage 3 1024 upscaling
	guidance_scale_3 = 9 # default for stage 3
	number_of_inference_steps_3 = 40 # default for stage 3

	result_path = deepfloydif_client.predict(
	path_for_upscale256_upscaling,
	selected_index_for_upscale256,
	seed_2,
	guidance_scale_2,
	custom_timesteps_2,
	number_of_inference_steps_2,
	prompt,
	negative_prompt,
	seed_3,
	guidance_scale_3,
	number_of_inference_steps_3,
	api_name="/upscale1024",
	)
	return result_path


	def load_image(png_files, stage_1_images):
	"""Opens images as variables so we can combine them later"""
	results = []
	for file in png_files:
	png_path = os.path.join(stage_1_images, file)
	results.append(Image.open(png_path))
	return results


	def combine_images(png_files, stage_1_images, partial_path):
	if os.environ.get("TEST_ENV") == "True":
	print("Combining images for deepfloydif_generate64")
	images = load_image(png_files, stage_1_images)
	combined_image = Image.new("RGB", (images[0].width * 2, images[0].height * 2))
	combined_image.paste(images[0], (0, 0))
	combined_image.paste(images[1], (images[0].width, 0))
	combined_image.paste(images[2], (0, images[0].height))
	combined_image.paste(images[3], (images[0].width, images[0].height))
	combined_image_path = os.path.join(stage_1_images, f"{partial_path}.png")
	combined_image.save(combined_image_path)
	return combined_image_path


	async def deepfloydif_generate64(ctx, prompt, client):
	"""DeepfloydIF command (generate images with realistic text using slash commands)"""
	try:
	if ctx.author.id != BOT_USER_ID:
	if ctx.channel.id == DEEPFLOYDIF_CHANNEL_ID:
	if os.environ.get("TEST_ENV") == "True":
	print("Safety checks passed for deepfloydif_generate64")
	channel = client.get_channel(DEEPFLOYDIF_CHANNEL_ID)
	# interaction.response message can't be used to create a thread, so we create another message
	message = await ctx.send(f"{prompt} - {ctx.author.mention} <a:loading:1114111677990981692>")

	loop = asyncio.get_running_loop()
	result = await loop.run_in_executor(None, deepfloydif_generate64_inference, prompt)
	stage_1_images = result[0]
	path_for_upscale256_upscaling = result[2]

	partial_path = pathlib.Path(path_for_upscale256_upscaling).name
	png_files = list(glob.glob(f"{stage_1_images}/*/.png"))

	if png_files:
	await message.delete()
	combined_image_path = combine_images(png_files, stage_1_images, partial_path)
	if os.environ.get("TEST_ENV") == "True":
	print("Images combined for deepfloydif_generate64")

	with Image.open(combined_image_path) as img:
	width, height = img.size
	new_width = width * 3
	new_height = height * 3
	resized_img = img.resize((new_width, new_height))
	x2_combined_image_path = combined_image_path
	resized_img.save(x2_combined_image_path)

	# making image bigger, more readable
	with open(x2_combined_image_path, "rb") as f: # was combined_image_path
	button1 = Button(custom_id="0", emoji="↖")
	button2 = Button(custom_id="1", emoji="↗")
	button3 = Button(custom_id="2", emoji="↙")
	button4 = Button(custom_id="3", emoji="↘")

	async def button_callback(interaction):
	index = int(interaction.data["custom_id"]) # 0,1,2,3

	await interaction.response.send_message(
	f"{interaction.user.mention} <a:loading:1114111677990981692>", ephemeral=True
	)
	result_path = await deepfloydif_upscale256(index, path_for_upscale256_upscaling)

	# create and use upscale 1024 button
	with open(result_path, "rb") as f:
	upscale1024 = Button(
	label="High-quality upscale (x4)", custom_id=str(index)
	) # "0", "1" etc
	upscale1024.callback = upscale1024_callback
	view = View(timeout=None)
	view.add_item(upscale1024)

	await interaction.delete_original_response()
	await channel.send(
	content=(
	f"{interaction.user.mention} Here is the upscaled image! Click the button"
	" to upscale even more!"
	),
	file=discord.File(f, f"{prompt}.png"),
	view=view,
	)

	async def upscale1024_callback(interaction):
	index = int(interaction.data["custom_id"])

	await interaction.response.send_message(
	f"{interaction.user.mention} <a:loading:1114111677990981692>", ephemeral=True
	)
	result_path = await deepfloydif_upscale1024(index, path_for_upscale256_upscaling, prompt)

	with open(result_path, "rb") as f:
	await interaction.delete_original_response()
	await channel.send(
	content=f"{interaction.user.mention} Here's your high-quality x16 image!",
	file=discord.File(f, f"{prompt}.png"),
	)

	button1.callback = button_callback
	button2.callback = button_callback
	button3.callback = button_callback
	button4.callback = button_callback

	view = View(timeout=None)
	view.add_item(button1)
	view.add_item(button2)
	view.add_item(button3)
	view.add_item(button4)

	# could store this message as combined_image_dfif in case it's useful for future testing
	await channel.send(
	f"{prompt} - {ctx.author.mention} Click a button to upscale! (make larger + enhance"
	" quality)",
	file=discord.File(f, f"{partial_path}.png"),
	view=view,
	)
	else:
	await ctx.send(f"{ctx.author.mention} No PNG files were found, cannot post them!")

	except Exception as e:
	print(f"Error: {e}")


	async def deepfloydif_upscale256(index: int, path_for_upscale256_upscaling):
	"""upscaling function for images generated using /deepfloydif"""
	try:
	loop = asyncio.get_running_loop()
	result_path = await loop.run_in_executor(
	None, deepfloydif_upscale256_inference, index, path_for_upscale256_upscaling
	)
	return result_path

	except Exception as e:
	print(f"Error: {e}")


	async def deepfloydif_upscale1024(index: int, path_for_upscale256_upscaling, prompt):
	"""upscaling function for images generated using /deepfloydif"""
	try:
	loop = asyncio.get_running_loop()
	result_path = await loop.run_in_executor(
	None, deepfloydif_upscale1024_inference, index, path_for_upscale256_upscaling, prompt
	)
	return result_path

	except Exception as e:
	print(f"Error: {e}")