import os
import random

import gradio as gr
import numpy as np
import torch

# import spaces #[uncomment to use ZeroGPU]
from diffusers import StableDiffusionPipeline
from peft import LoraConfig, PeftModel

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",
    "sd-legacy/stable-diffusion-v1-5",
]

model_dropdown = [
    "stabilityai/sdxl-turbo",
    "CompVis/stable-diffusion-v1-4",
    "sd-legacy/stable-diffusion-v1-5",
]


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
MODEL_NAME = "CompVis/stable-diffusion-v1-4"
CKPT_DIR = "sd-14-lora-1000"


def get_lora_sd_pipeline(
    ckpt_dir=CKPT_DIR,
    base_model_name_or_path=None,
    dtype=torch.float16,
    device="cuda",
    adapter_name="default",
):
    unet_sub_dir = os.path.join(ckpt_dir, "unet")
    text_encoder_sub_dir = os.path.join(ckpt_dir, "text_encoder")
    if os.path.exists(text_encoder_sub_dir) and base_model_name_or_path is None:
        config = LoraConfig.from_pretrained(text_encoder_sub_dir)
        base_model_name_or_path = config.base_model_name_or_path

    if base_model_name_or_path is None:
        raise ValueError("Please specify the base model name or path")

    pipe = StableDiffusionPipeline.from_pretrained(
        base_model_name_or_path, torch_dtype=dtype
    ).to(device)
    pipe.unet = PeftModel.from_pretrained(
        pipe.unet, unet_sub_dir, adapter_name=adapter_name
    )

    if os.path.exists(text_encoder_sub_dir):
        pipe.text_encoder = PeftModel.from_pretrained(
            pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name
        )

    if dtype in (torch.float16, torch.bfloat16):
        pipe.unet.half()
        pipe.text_encoder.half()

    return pipe


# @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)

    if model_id == "gstranger/kawaiicat-lora-1.4":
        # добавляем lora
        pipe = get_lora_sd_pipeline(
            os.path.join(CKPT_DIR, ""), adapter_name="sd-14-lora", dtype=torch_dtype
        ).to(device)
        pipe.safety_checker = None
        print(f"LoRA adapter loaded: {pipe.unet.active_adapters}")

    else:
        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()