import gradio as gr
import numpy as np
import os, random, json, spaces, torch, time, subprocess

import torch
# from transformers import AutoProcessor, AutoTokenizer
# from diffusers import DiffusionPipeline
from diffusers import NewbiePipeline
from transformers import AutoModel

"""
for anyone having problem with flash_attn, you need to build it from source. 
This compiles the library against your specific environment:

git clone https://github.com/Dao-AILab/flash-attention.git
pip install --no-build-isolation flash-attention/.
"""

from utils import prompt_utils

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

device = "cuda"
MODEL_REPO = "Disty0/NewBie-image-Exp0.1-Diffusers"
text_encoder_2 = AutoModel.from_pretrained(
    MODEL_REPO, 
    subfolder="text_encoder_2", 
    trust_remote_code=True, 
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)
pipe = NewbiePipeline.from_pretrained(
    MODEL_REPO, 
    text_encoder_2=text_encoder_2, 
    torch_dtype=torch.bfloat16
).to("cuda")
del text_encoder_2

def read_file(path: str) -> str:
    with open(path, 'r', encoding='utf-8') as f:
        content = f.read()
    return content

def prepare(prompt, is_polish_prompt):
    if not is_polish_prompt: return prompt, False
    system_prompt = read_file('system_prompt.md')
    polished_prompt = prompt_utils.polish_prompt(prompt, system_prompt)
    return polished_prompt, True

@spaces.GPU
def inference(
    prompt,
    negative_prompt="blurry ugly bad",
    width=1024,
    height=1024,
    seed=42,
    randomize_seed=True,
    guidance_scale=3.5,
    num_inference_steps=8,
    progress=gr.Progress(track_tqdm=True),
):
    timestamp = time.time()
    print(f"timestamp: {timestamp}")

    # generation
    if randomize_seed: seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

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

    return image, seed

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

with open('examples/0_examples.json', 'r') as file: examples = json.load(file)

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        with gr.Column():
            gr.HTML(read_file("static/header.html"))
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="Prompt",
                    show_label=False,
                    lines=2,
                    placeholder="Enter your prompt",
                    # container=False,
                )
                is_polish_prompt = gr.Checkbox(label="Polish prompt", value=True)
                run_button = gr.Button("Generate", variant="primary")
                with gr.Accordion("Advanced Settings", open=False):
                    
                    negative_prompt = gr.Textbox(
                        label="Negative prompt",
                        lines=2,
                        container=False,
                        placeholder="Enter your negative prompt",
                        value="blurry ugly bad"
                    )
                    num_inference_steps = gr.Slider(
                        label="Steps",
                        minimum=1,
                        maximum=50,
                        step=1,
                        value=20,
                    )
                    with gr.Row():
                        width = gr.Slider(
                            label="Width",
                            minimum=512,
                            maximum=1280,
                            step=32,
                            value=768, 
                        )

                        height = gr.Slider(
                            label="Height",
                            minimum=512,
                            maximum=1280,
                            step=32,
                            value=1024,
                        )
                    with gr.Row():
                        seed = gr.Slider(
                            label="Seed",
                            minimum=0,
                            maximum=MAX_SEED,
                            step=1,
                            value=42,
                        )
                        guidance_scale = gr.Slider(
                            label="Guidance scale",
                            minimum=0.0,
                            maximum=10.0,
                            step=0.1,
                            value=3.5,
                        )

                    
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Column():
                output_image = gr.Image(label="Generated image", show_label=False)
                polished_prompt = gr.Textbox(label="Final prompt",lines=2, interactive=False)

        gr.Examples(examples=examples, inputs=[prompt])
        gr.Markdown(read_file("static/footer.md"))

    run_button.click(
        fn=prepare,
        inputs=[prompt, is_polish_prompt],
        outputs=[polished_prompt, is_polish_prompt]
    ).then(
        fn=inference,
        inputs=[
            polished_prompt,
            negative_prompt,
            width,
            height,
            seed,
            randomize_seed,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[output_image, seed],
    )


if __name__ == "__main__":
    demo.launch(mcp_server=True, css=css)