import os
import gradio as gr
import numpy as np
import random
import spaces
from diffusers import DiffusionPipeline
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

try:
    from dotenv import load_dotenv

    load_dotenv()
except:
    print("failed to import dotenv (this is not a problem on the production)")

device = "cuda" if torch.cuda.is_available() else "cpu"

HF_TOKEN = os.environ.get("HF_TOKEN")
assert HF_TOKEN is not None

IMAGE_MODEL_REPO_ID = os.environ.get(
    "IMAGE_MODEL_REPO_ID", "OnomaAIResearch/Illustrious-xl-early-release-v0"
)
DART_V3_REPO_ID = os.environ.get("DART_V3_REPO_ID", None)
assert DART_V3_REPO_ID is not None

torch_dtype = torch.bfloat16

dart = AutoModelForCausalLM.from_pretrained(
    DART_V3_REPO_ID,
    torch_dtype=torch_dtype,
    token=HF_TOKEN,
)
tokenizer = AutoTokenizer.from_pretrained(DART_V3_REPO_ID)

pipe = DiffusionPipeline.from_pretrained(IMAGE_MODEL_REPO_ID, torch_dtype=torch_dtype)
pipe = pipe.to(device)
pipe = torch.compile(pipe)


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

TEMPLATE = (
    "<|bos|>"
    #
    "<|rating:general|>"
    "{aspect_ratio}"
    "<|length:medium|>"
    #
    "<copyright>original</copyright>"
    #
    "<character></character>"
    #
    "<general>"
)


@torch.inference_mode
def generate_prompt(aspect_ratio: str):
    input_ids = tokenizer.encode_plus(
        TEMPLATE.format(aspect_ratio=aspect_ratio)
    ).input_ids

    output_ids = dart.generate(
        input_ids,
        max_new_tokens=256,
        temperature=1.0,
        top_p=1.0,
        top_k=100,
        num_beams=1,
    )[0]

    generated = output_ids[len(input_ids) :]
    decoded = ", ".join(tokenizer.batch_decode(generated))

    return decoded


@spaces.GPU  # [uncomment to use ZeroGPU]
def infer(
    negative_prompt: str,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

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

    prompt = generate_prompt("<|aspect_ratio:square|>")
    print(prompt)

    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
    ).images[0]

    return image, prompt, seed


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(f"""
        # Random IllustriousXL
        """)

        with gr.Row():
            run_button = gr.Button("Generate random", scale=0)

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

        with gr.Accordion("Generation details", open=False):
            prompt_txt = gr.Textbox("Generated prompt", interactive=False)

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter a negative prompt",
                visible=False,
                value=" worst quality, comic, multiple views, bad quality, low quality, lowres, displeasing, very displeasing, bad anatomy, bad hands, scan artifacts, monochrome, greyscale, signature, twitter username, jpeg artifacts, 2koma, 4koma, guro, extra digits, fewer digits",
            )

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

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

            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=1.0,
                    maximum=10.0,
                    step=0.5,
                    value=6.5,
                )

                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=20,
                )

    gr.on(
        triggers=[run_button.click],
        fn=infer,
        inputs=[
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, prompt_txt, seed],
    )

demo.queue().launch()