import gradio as gr
import numpy as np
import random
import uuid
from PIL import Image
import spaces
from diffusers import DiffusionPipeline
import torch

DESCRIPTIONx = """## SD-3.5 LARGE TURBO """

device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "stabilityai/stable-diffusion-3.5-large-turbo"

if torch.cuda.is_available():
    torch_dtype = torch.bfloat16
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

# Define styles
style_list = [
    {
        "name": "3840 x 2160",
        "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
    },
    {
        "name": "2560 x 1440",
        "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
    },
    {
        "name": "HD+",
        "prompt": "hyper-realistic 2K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
    },
    {
        "name": "Style Zero",
        "prompt": "{prompt}",
        "negative_prompt": "",
    },
]

STYLE_NAMES = [style["name"] for style in style_list]
DEFAULT_STYLE_NAME = STYLE_NAMES[0]

grid_sizes = {
    "2x1": (2, 1),
    "1x2": (1, 2),
    "2x2": (2, 2),
    "2x3": (2, 3),
    "3x2": (3, 2),
    "1x1": (1, 1)
}

@spaces.GPU(duration=60, enable_queue=True)
def infer(
    prompt,
    negative_prompt="",
    seed=42,
    randomize_seed=False,
    width=1024,
    height=1024,
    guidance_scale=0.0,
    num_inference_steps=4,
    style="Style Zero",
    grid_size="1x1",
    progress=gr.Progress(track_tqdm=True),
):
    
    selected_style = next(s for s in style_list if s["name"] == style)
    styled_prompt = selected_style["prompt"].format(prompt=prompt)
    styled_negative_prompt = selected_style["negative_prompt"]

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

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

    grid_size_x, grid_size_y = grid_sizes.get(grid_size, (2, 2))
    num_images = grid_size_x * grid_size_y

    images = []
    for _ in range(num_images):
        image = pipe(
            prompt=styled_prompt,
            negative_prompt=styled_negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
        ).images[0]
        images.append(image)

    # Create a grid image
    grid_img = Image.new('RGB', (width * grid_size_x, height * grid_size_y))

    for i, img in enumerate(images[:num_images]):
        grid_img.paste(img, (i % grid_size_x * width, i // grid_size_x * height))

    # Save the grid image
    unique_name = str(uuid.uuid4()) + ".png"
    grid_img.save(unique_name)
    
    return unique_name, seed

examples = [
    "A capybara wearing a suit holding a sign that reads Hello World",
]

css = '''
.gradio-container{max-width: 585px !important}
h1{text-align:center}
footer {
    visibility: hidden
}
'''

with gr.Blocks(css=css, theme="prithivMLmods/Minecraft-Theme") as demo:
    gr.Markdown(DESCRIPTIONx)  
    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.Row(visible=True):
        grid_size_selection = gr.Dropdown(
            choices=["2x1", "1x2", "2x2", "2x3", "3x2", "1x1"],
            value="1x1",
            label="Grid Size"
        )

    with gr.Row(visible=True):     
        style_selection = gr.Radio(
            show_label=True,
            container=True,
            interactive=True,
            choices=STYLE_NAMES,
            value=DEFAULT_STYLE_NAME,
            label="Quality Style",
        )

    with gr.Accordion("Advanced Settings", open=False):
        negative_prompt = gr.Text(
            label="Negative prompt",
            max_lines=1,
            placeholder="Enter a negative prompt",
            visible=False,
        )

        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=512,
                maximum=MAX_IMAGE_SIZE,
                step=32,
                value=1024,
            )

            height = gr.Slider(
                label="Height",
                minimum=512,
                maximum=MAX_IMAGE_SIZE,
                step=32,
                value=1024,
            )

        with gr.Row():
            guidance_scale = gr.Slider(
                label="Guidance scale",
                minimum=0.0,
                maximum=7.5,
                step=0.1,
                value=0.0,
            )

            num_inference_steps = gr.Slider(
                label="Number of inference steps",
                minimum=1,
                maximum=50,
                step=1,
                value=4,
            )
                
    gr.Examples(
        examples=examples, 
        inputs=[prompt], 
        outputs=[result, seed], 
        fn=infer, 
        cache_examples=True, 
        cache_mode="lazy"
    )
        
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            style_selection,
            grid_size_selection,
        ],
        outputs=[result, seed],
    )

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