import gradio as gr
import numpy as np
import random
# import spaces #[uncomment to use ZeroGPU]
from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, AutoencoderKL
import torch

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

# Load VAE
vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix",
    torch_dtype=torch.float16
)

# Load base pipeline with LoRA
base = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    vae=vae,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
)
base.scheduler = EulerAncestralDiscreteScheduler.from_config(base.scheduler.config)
base.to(device)
base.load_lora_weights(
    "Darebal/sdxl_lora",
    subfolder = "dataset2000_8r_5epochs_flip_text_512_500ws",
    weight_name ="pytorch_lora_weights.safetensors"
)

# Load refiner pipeline
refiner = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    text_encoder_2=base.text_encoder_2,
    vae=base.vae,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
)
refiner.scheduler = EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config)
refiner.to(device)

# Функція генерації зображення
def infer(
    prompt,
    negative_prompt,
    guidance_scale,
    num_inference_steps,
    seed=None,
    progress=gr.Progress(track_tqdm=True),
):
    high_noise_frac = 0.8
    height = 512
    width = 512

    generator = torch.manual_seed(seed) if seed is not None else None

    # Stage 1: base latent generation
    latent_output = base(
        prompt=prompt + " on a white background",
        negative_prompt=negative_prompt if negative_prompt else None,
        height=height,
        width=width,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        denoising_end=high_noise_frac,
        output_type="latent",
        generator=generator,  # <- Pass generator here
    )

    # Stage 2: refine the image
    refined_output = refiner(
        prompt=prompt + " on a white background",
        negative_prompt=negative_prompt if negative_prompt else None,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        denoising_start=high_noise_frac,
        image=latent_output.images,
        generator=generator,  # <- Pass generator here too
    )

    return refined_output.images[0]

examples = [
    "A crazy clown tattoo",
    "A green dragon tattoo",
    "A skull tattoo",
]

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(" # Tattoo image from text generator")

        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",
            )

            guidance_scale = gr.Slider(
                label="Guidance scale",
                minimum=0.0,
                maximum=15.0,
                step=0.1,
                value=8.0,
            )

            num_inference_steps = gr.Slider(
                label="Number of inference steps",
                minimum=1,
                maximum=100,
                step=1,
                value=32,
            )
            
            seed = gr.Number(
                label="Seed (optional)",
                value=None,
                precision=0,
                interactive=True,
                placeholder="Random seed",
            )
            
            gr.Markdown("""
            **Tips for better results:**
            - **Prompts** Better use short prompts. Include word "tattoo" in your prompt. It can be at the start of sentence e.g. "A tattoo of dragon" or in the end e.g. "A rose tattoo". This way you can generate a better image and choose the right one.
            - **Guidance Scale** (Recommended: 6-8)
            - **Inference Steps** (Recommended: 32)
            - **Negative prompt** (Recommended: detailed background, blurry, abstract, painting, low quality, hyper-detailed)
            - **Random seed** (Default: None) If random seed is provided then the model generates the same image with the same parameters every time. If not, then each time a random image.
            """)
            

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

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