import gradio as gr

import os

from PIL import Image
import random

import torch
from diffusers import StableDiffusionPipeline, AutoencoderKL

def gen_seed():
    random_data = os.urandom(3)
    seed = int.from_bytes(random_data, byteorder="big")
    return seed

repo = "IDKiro/sdxs-512-0.9"
weight_type = torch.float32     # or float16

# Load model.
pipe = StableDiffusionPipeline.from_pretrained(repo, torch_dtype=weight_type)

# use original VAE
# pipe.vae = AutoencoderKL.from_pretrained("IDKiro/sdxs-512-0.9/vae_large")
 
#pipe.to("cuda")

prompt = "portrait photo of a girl, photograph, highly detailed face, depth of field, moody light, golden hour"

def sdxs_run(prompt, steps, guidance, seed):
    # Ensure using 1 inference step and CFG set to 0.
    image = pipe(
        prompt, 
        num_inference_steps=steps, 
        guidance_scale=guidance,
        generator=torch.Generator(device="cpu").manual_seed(seed)
    ).images[0]
    return image

#image.save("output.png")

def update_seed(rand, seed):
    if rand:
        return gen_seed()
    else:
        return seed

desc = """# SDXS CPU Test Space
Just a quick test. Model is `sdxs-512-0.9` for txt2img.
"""

with gr.Blocks() as demo:
    gr.Markdown(desc)
    with gr.Group():
        with gr.Row():
            img = gr.Image(label='SDXS Generated Image')
        with gr.Row():
            prompt = gr.Textbox(label='Enter your prompt (English)', scale=8, value="portrait photo of a girl, photograph, highly detailed face, depth of field, moody light, golden hour")
        with gr.Row():
            with gr.Accordion("More options", open=False):
                steps = gr.Slider(label="Number of steps", value=1, minimum=1, maximum=20, step=1)
                guidance = gr.Slider(label="Guidance", value=0, minimum=0, maximum=2, step=0.1)
                seed = gr.Slider(label="Seed", minimum=20, maximum=100000000, step=1, randomize=True)
                rand = gr.Checkbox(label="Randomize Seed After Generation?", value=True)
        with gr.Row():
            submit = gr.Button(scale=1, variant='primary')
            #clear = gr.ClearButton(components=[])
        submit.click(fn=sdxs_run, inputs=[prompt, steps, guidance, seed], outputs=img).then(fn=update_seed, inputs=[rand, seed], outputs=seed)

demo.queue(max_size=20).launch()