import gradio as gr
import numpy as np
import datasets, os, time, random
from PIL import Image

dataset = datasets.load_dataset('parsee-mizuhashi/realrealrealreal', use_auth_token=os.environ["HF_TOKEN"], num_proc=8) #train dataset
def apply_noise_layer(noise: np.ndarray, img: np.ndarray, percent):
    if percent <= 0:
        return noise
    elif percent >= 1:
        return img
    else:
        #apply the noise with transperancy equal to 1-percent
        t = img * percent + noise * (1-percent)
        t = Image.fromarray(t.astype(np.uint8))
        return t
    
def generate(sampler, steps, use_thunder=False):
    t = 0.25
    if "++" in sampler:
        t = 0.4
    if use_thunder:
        t = 0.1
    # get random img from dataset
    imgid: Image = random.choice(list(dataset["train"]["image"]))
    basimg = np.array(imgid)
    for _ in range(steps):
        image = np.random.random((1024, 1024, 3))
        image = image * 255
        image = image.astype(np.uint8)
        yield apply_noise_layer(image, basimg, (_+1) / steps)
        time.sleep(t)

    yield basimg


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=2):
            pass
        with gr.Column(scale=3):
            positive = gr.Textbox(interactive=True, show_label=False)
            negative = gr.Textbox(placeholder="negative", interactive=True, show_label=False)
            generate_button = gr.Button("Generate")
            img = gr.Image(interactive=False, show_label=False)
            with gr.Accordion("Advanced", open=False):
                
                sampler = gr.Dropdown(
                choices=["Euler A", "DPM++ 2M SDE", "DPM++ 2S Ancestral", "DDPM"],
                value="DPM++ 2M SDE",
                )
                steps = gr.Slider(minimum=1, maximum=50, value=10, step=1)
                use_thunder = gr.Checkbox(label="Use WD Thunder to generate", value=False)

        generate_button.click(generate, inputs=[sampler, steps, use_thunder], outputs=img)
        with gr.Column(scale=2):
            pass
        
demo.queue(max_size=100, default_concurrency_limit=10)
demo.launch()