#!/usr/bin/env python

import os
import random
import gradio as gr
import numpy as np
import PIL.Image
import torch
from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler

DESCRIPTION = 'This space is an API service meant to be used by frontend applications.'
if not torch.cuda.is_available():
    DESCRIPTION += '\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>'

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = int(os.getenv('MAX_IMAGE_SIZE', '1024'))
SECRET_TOKEN = os.getenv('SECRET_TOKEN', 'default_secret')

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
if torch.cuda.is_available():
    unet = UNet2DConditionModel.from_pretrained(
        "latent-consistency/lcm-ssd-1b",
        torch_dtype=torch.float16,
        variant="fp16"
    )

    pipe = DiffusionPipeline.from_pretrained(
        "segmind/SSD-1B",
        unet=unet,
        torch_dtype=torch.float16,
        variant="fp16"
    )

    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    pipe.to(device)
else:
    pipe = None
    
def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed


def generate(prompt: str,
             negative_prompt: str = '',
             prompt_2: str = '',
             negative_prompt_2: str = '',
             use_negative_prompt: bool = False,
             use_prompt_2: bool = False,
             use_negative_prompt_2: bool = False,
             seed: int = 0,
             width: int = 1024,
             height: int = 1024,
             guidance_scale: float = 1.0,
             num_inference_steps: int = 4,
             secret_token: str = '') -> PIL.Image.Image:
    if secret_token != SECRET_TOKEN:
        raise gr.Error(
            f'Invalid secret token. Please fork the original space if you want to use it for yourself.')
        
    generator = torch.Generator().manual_seed(seed)

    if not use_negative_prompt:
        negative_prompt = None  # type: ignore
    if not use_prompt_2:
        prompt_2 = None  # type: ignore
    if not use_negative_prompt_2:
        negative_prompt_2 = None  # type: ignore

    return pipe(prompt=prompt,
                negative_prompt=negative_prompt,
                prompt_2=prompt_2,
                negative_prompt_2=negative_prompt_2,
                width=width,
                height=height,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                generator=generator,
                output_type='pil').images[0]

with gr.Blocks(css='style.css') as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Box():
        with gr.Row():
            secret_token = gr.Text(
                label='Secret Token',
                max_lines=1,
                placeholder='Enter your secret token',
            )
            prompt = gr.Text(
                label='Prompt',
                show_label=False,
                max_lines=1,
                placeholder='Enter your prompt',
                container=False,
            )
            run_button = gr.Button('Run', scale=0)
        result = gr.Image(label='Result', show_label=False)
        with gr.Accordion('Advanced options', open=False):
            with gr.Row():
                use_negative_prompt = gr.Checkbox(label='Use negative prompt',
                                                  value=False)
                use_prompt_2 = gr.Checkbox(label='Use prompt 2', value=False)
                use_negative_prompt_2 = gr.Checkbox(
                    label='Use negative prompt 2', value=False)
            negative_prompt = gr.Text(
                label='Negative prompt',
                max_lines=1,
                placeholder='Enter a negative prompt',
                visible=False,
            )
            prompt_2 = gr.Text(
                label='Prompt 2',
                max_lines=1,
                placeholder='Enter your prompt',
                visible=False,
            )
            negative_prompt_2 = gr.Text(
                label='Negative prompt 2',
                max_lines=1,
                placeholder='Enter a negative prompt',
                visible=False,
            )

            seed = gr.Slider(label='Seed',
                             minimum=0,
                             maximum=MAX_SEED,
                             step=1,
                             value=0)
            randomize_seed = gr.Checkbox(label='Randomize seed', value=True)
            with gr.Row():
                width = gr.Slider(
                    label='Width',
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                height = gr.Slider(
                    label='Height',
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            guidance_scale = gr.Slider(
                label='Guidance scale',
                minimum=1,
                maximum=20,
                step=0.1,
                value=1.0)
            num_inference_steps = gr.Slider(
                label='Number of inference steps',
                minimum=2,
                maximum=8,
                step=1,
                value=4)

    use_negative_prompt.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt,
        outputs=negative_prompt
    )
    use_prompt_2.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_prompt_2,
        outputs=prompt_2
    )
    use_negative_prompt_2.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt_2,
        outputs=negative_prompt_2
    )

    inputs = [
        prompt,
        negative_prompt,
        prompt_2,
        negative_prompt_2,
        use_negative_prompt,
        use_prompt_2,
        use_negative_prompt_2,
        seed,
        width,
        height,
        guidance_scale,
        num_inference_steps,
        secret_token,
    ]
    prompt.submit(
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed
    ).then(
        fn=generate,
        inputs=inputs,
        outputs=result,
        api_name='run',
    )
    negative_prompt.submit(
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed
    ).then(
        fn=generate,
        inputs=inputs,
        outputs=result
    )
    run_button.click(
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed
    ).then(
        fn=generate,
        inputs=inputs,
        outputs=result
    )
demo.queue(max_size=6).launch()