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

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

pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
pipe = pipe.to(device)
pipe.load_lora_weights("strangerzonehf/SD3.5-Turbo-Portrait-LoRA", weight_name="SD3.5-Turbo-Portrait.safetensors")
pipe.fuse_lora(lora_scale=1.0)

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

# Style presets
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 = [s["name"] for s in style_list]

def randomize_seed_fn(seed, randomize):
    return random.randint(0, MAX_SEED) if randomize else seed

def save_image(img):
    filename = str(uuid.uuid4()) + ".png"
    img.save(filename)
    return filename

@spaces.GPU
def generate_images(
    prompt,
    style,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    num_images,
    progress=gr.Progress(track_tqdm=True)
):
    seed = randomize_seed_fn(seed, randomize_seed)
    generator = torch.Generator(device=device).manual_seed(seed)

    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 not negative_prompt else negative_prompt

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

    image_paths = [save_image(img) for img in images]
    return image_paths, seed

# CSS & Interface
css = '''
.gradio-container {
    max-width: 150%;
    margin: 0 auto;
}
h1 { text-align: center; }
footer { visibility: hidden; }
'''

examples = [
    "portrait photo of a futuristic astronaut",
    "macro shot of a water droplet on a leaf",
    "hyper-realistic food photography of a burger",
    "cyberpunk city at night, rain, neon lights",
    "ultra detailed fantasy landscape with dragons",
]

with gr.Blocks(css=css, theme="YTheme/GMaterial") as demo:
    gr.Markdown("## SD3.5 Turbo Portrait")

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Row():
                prompt = gr.Text(
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    container=False,
                )
                run_button = gr.Button("Run", scale=0, variant="primary")

            result_gallery = gr.Gallery(show_label=False, format="png", columns=2, object_fit="contain")

            with gr.Accordion("Advanced Settings", open=False):
                num_images = gr.Slider(
                    label="Number of Images",
                    minimum=1,
                    maximum=10,
                    value=5,
                    step=1,
                )
                style = gr.Dropdown(label="Select Style", choices=STYLE_NAMES, value=STYLE_NAMES[0])

                negative_prompt = gr.Text(
                    label="Negative Prompt",
                    max_lines=4,
                    lines=3,
                    value="cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly"
                )
                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=64, value=1024)
                    height = gr.Slider(label="Height", minimum=512, maximum=MAX_IMAGE_SIZE, step=64, value=1024)
                with gr.Row():
                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=15, step=0.5, value=0.0)
                    num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=30, step=1, value=4)

        with gr.Column(scale=1):
            gr.Examples(
                examples=examples,
                inputs=prompt,
                cache_examples=False,
            )

    gr.on(
        triggers=[prompt.submit, run_button.click],
        fn=generate_images,
        inputs=[
            prompt,
            style,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            num_images
        ],
        outputs=[result_gallery, seed],
        api_name="generate"
    )

if __name__ == "__main__":
    demo.queue(max_size=40).launch(ssr_mode=False)