from typing import List

import time
import gradio as gr
from modules.sprites import generate_sides, build_gifs, build_spritesheet


def generate(prompt, thresh):
		timestamp = int(time.time())

		sides = generate_sides(prompt, 3)[0]	
		spritesheet = build_spritesheet(sides, prompt, timestamp=timestamp, thresh=thresh)[0]

		filepaths = build_gifs(sides, prompt, save=True, timestamp=timestamp, thresh=thresh)[1]

		return spritesheet, filepaths[0], filepaths[1], filepaths[2], filepaths[3]


demo = gr.Blocks()

with demo:
	gr.Markdown("""
		# Stable Diffusion Sprite Sheets

		Generate a sprite sheet of pixel art character sides and their corresponding walk animations! Checkpoint by [Onodofthenorth](https://huggingface.co/Onodofthenorth/SD_PixelArt_SpriteSheet_Generator). Sprites are 32x32 pixels scaled up to 96x96. NSFW content replaced with blank sprites.
	""")

	with gr.Row():
		with gr.Column():
			prompt = gr.Textbox(label="Prompt", placeholder="Enter text prompt")
			threshold = gr.Slider(label="Background removal threshold", placeholder="Tweak how strong the background removal is", minimum=0, maximum=255, value=128)

		button = gr.Button("Generate")

	with gr.Box():
		with gr.Row():
			spritesheet = gr.Image(label="Sprite Sheet")

			with gr.Row():
				front = gr.Image(label="Front")
				back = gr.Image(label="Back")
				left = gr.Image(label="Left")
				right = gr.Image(label="Right")

		button.click(fn=generate, inputs=[prompt, threshold], outputs=[spritesheet, front, back, left, right])

demo.queue()
demo.launch(show_api=False)