import os
import spaces
import torch
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
import gradio as gr
import random
import tqdm

# Enable TQDM progress tracking
tqdm.monitor_interval = 0

# Load the diffusion pipeline
pipe = StableDiffusionXLPipeline.from_pretrained(
    "kayfahaarukku/UrangDiffusion-1.0", 
    torch_dtype=torch.float16, 
    custom_pipeline="lpw_stable_diffusion_xl",
)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)

# Function to generate an image
@spaces.GPU
def generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
    pipe.to('cuda')
    
    if randomize_seed:
        seed = random.randint(0, 99999999)
    
    original_prompt = prompt
    original_negative_prompt = negative_prompt
    
    if use_defaults:
        prompt = f"{prompt}, masterpiece, best quality"
        negative_prompt = f"nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, {negative_prompt}"
    generator = torch.manual_seed(seed)
    
    def callback(step, timestep, latents):
        progress(step / num_inference_steps)
        return
    
    width, height = map(int, resolution.split('x'))
    image = pipe(
        prompt, 
        negative_prompt=negative_prompt,
        width=width,
        height=height, 
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator,
        callback=callback,
        callback_steps=1
    ).images[0]

    torch.cuda.empty_cache()

    return image, seed, original_prompt, original_negative_prompt

# Define Gradio interface
def interface_fn(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
    image, seed, original_prompt, original_negative_prompt = generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress)
    
    details = f"""Prompt: {original_prompt}

Negative prompt: {original_negative_prompt}

Steps: {num_inference_steps}, CFG scale: {guidance_scale}, Seed: {seed}, Size: {resolution}

Default quality tags: {"Enabled" if use_defaults else "Disabled"}"""

    if use_defaults:
        details += f"""

Default prompt addition: , masterpiece, best quality

Default negative prompt addition: nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name"""
    
    return image, seed, gr.update(value=seed), details

def reset_inputs():
    return gr.update(value=''), gr.update(value=''), gr.update(value=True), gr.update(value='832x1216'), gr.update(value=7), gr.update(value=28), gr.update(value=0), gr.update(value=True)

with gr.Blocks(title="UrangDiffusion 1.0 Demo", theme="NoCrypt/miku@1.2.1") as demo:
    gr.HTML(
        "<h1>UrangDiffusion 1.0 Demo</h1>"
        "This demo is intended to showcase what the model is capable of and is not intended to be the main generation platform. Results produced with Diffusers are not the best, and it's highly recommended for you to get the model running inside Stable Diffusion WebUI or ComfyUI."
        )
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(lines=2, placeholder="Enter prompt here", label="Prompt")
            negative_prompt_input = gr.Textbox(lines=2, placeholder="Enter negative prompt here", label="Negative Prompt")
            use_defaults_input = gr.Checkbox(label="Use Default Quality Tags and Negative Prompt", value=True)
            resolution_input = gr.Radio(
                choices=[
                    "1024x1024", "1152x896", "896x1152", "1216x832", "832x1216",
                    "1344x768", "768x1344", "1536x640", "640x1536"
                ],
                label="Resolution",
                value="832x1216"
            )
            guidance_scale_input = gr.Slider(minimum=1, maximum=20, step=0.5, label="Guidance Scale", value=7)
            num_inference_steps_input = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=28)
            seed_input = gr.Slider(minimum=0, maximum=99999999, step=1, label="Seed", value=0, interactive=True)
            randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=True)
            generate_button = gr.Button("Generate")
            reset_button = gr.Button("Reset")

        with gr.Column():
            output_image = gr.Image(type="pil", label="Generated Image")
        with gr.Accordion("Generation Details", open=False):
            generation_info = gr.Textbox(
                label="", 
                max_lines=15, 
                interactive=False,
                elem_id="generation_info",
                show_copy_button=True
            )
            gr.Markdown(
                """
                ### Recommended prompt formatting:
                `1girl/1boy, character name, from what series, everything else in any order, masterpiece, best quality`

                **PS:** `masterpiece, best quality` is automatically added when "Use Default Quality Tags and Negative Prompt" is enabled

                ### Recommended settings:
                - Steps: 25-30
                - CFG: 5-7
                """
            )

    generate_button.click(
        interface_fn,
        inputs=[
            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
        ],
        outputs=[output_image, seed_input, seed_input, generation_info]
    )
    
    reset_button.click(
        reset_inputs,
        inputs=[],
        outputs=[
            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
        ]
    )

demo.queue(max_size=20).launch(share=False)