import random

import gradio as gr
import numpy as np
import torch

# import spaces #[uncomment to use ZeroGPU]
from diffusers import StableDiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"
# model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
# model_repo_id = "CompVis/stable-diffusion-v1-4"
# model_dropdown = ["stabilityai/sdxl-turbo", "CompVis/stable-diffusion-v1-4"]

models = [
    "gstranger/kawaiicat-lora-1.4",
    "CompVis/stable-diffusion-v1-4",
    "stabilityai/sdxl-turbo",
]


if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

# pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
# pipe = pipe.to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024


# @spaces.GPU #[uncomment to use ZeroGPU]
def infer(
    model_id,
    prompt,
    negative_prompt,
    randomize_seed,
    width,
    height,
    # model_repo_id=model_repo_id,
    seed=42,
    guidance_scale=7,
    num_inference_steps=50,
    progress=gr.Progress(track_tqdm=True),
    lora_scale=1,
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)

    pipe = StableDiffusionPipeline.from_pretrained(
        model_id,
        torch_dtype=torch_dtype,
        requires_safety_checker=False,
        safety_checker=None,
    )
    pipe = pipe.to(device)

    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
        cross_attention_kwargs={"scale": lora_scale},
    ).images[0]

    return image, seed


examples = [
    "kawaiicat. The cat is sitting. The cat's tail is curled up at the end. The cat is pleased and is enjoying its time.",
    "kawaiicat. The cat is sitting upright. The cat is eating some noodles with the chopsticks from a green bowl, which it's holding in his hands.",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # Text-to-Image kawaiicat Stickers")
        with gr.Row():
            # Dropdown to select the model from Hugging Face
            model_id = gr.Dropdown(
                label="Model",
                choices=models,
                value=models[0],  # Default model
            )

        lora_scale = gr.Slider(
            label="LORA Scale",
            minimum=0,
            maximum=1,
            step=0.01,
            value=1,
        )

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )

            run_button = gr.Button("Run", scale=0, variant="primary")

        result = gr.Image(label="Result", show_label=False)

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter a negative prompt",
                visible=True,
            )

            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=42,
            )

            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=512,  # Replace with defaults that work for your model
                )

                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=512,  # Replace with defaults that work for your model
                )

            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=10.0,  # Replace with defaults that work for your model
                )

                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=50,  # Replace with defaults that work for your model
                )

        gr.Examples(examples=examples, inputs=[prompt])
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            model_id,
            prompt,
            negative_prompt,
            randomize_seed,
            width,
            height,
            seed,
            guidance_scale,
            num_inference_steps,
            lora_scale,
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.launch()