import torch
import gradio as gr
from diffusers import DiffusionPipeline


MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"


# --------------------------------------------------
# GPU configuration
# --------------------------------------------------

assert torch.cuda.is_available(), "CUDA GPU is required"

device = "cuda"

print("GPU:", torch.cuda.get_device_name(0))
print(
    "VRAM:",
    round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2),
    "GB"
)


# Enable TF32 for supported NVIDIA GPUs
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True


# --------------------------------------------------
# Load SDXL in FP16
# --------------------------------------------------

pipe = DiffusionPipeline.from_pretrained(
    MODEL_ID,
    dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
)

pipe = pipe.to(device)


# --------------------------------------------------
# Optimizations
# --------------------------------------------------

# Use channels_last memory layout
pipe.unet.to(memory_format=torch.channels_last)
pipe.vae.to(memory_format=torch.channels_last)


# Compile the UNet for faster repeated inference
print("Compiling UNet...")

pipe.unet = torch.compile(
    pipe.unet,
    mode="reduce-overhead",
    fullgraph=True
)

print("UNet compilation enabled.")


# --------------------------------------------------
# Generate
# --------------------------------------------------

def generate_image(prompt):

    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    print("Generating image...")
    print("Prompt:", prompt)

    with torch.inference_mode():

        image = pipe(
            prompt=prompt,
            num_inference_steps=20,
            guidance_scale=7.0,
            height=768,
            width=768,
        ).images[0]

    return image


# --------------------------------------------------
# Gradio UI
# --------------------------------------------------

demo = gr.Interface(
    fn=generate_image,

    inputs=gr.Textbox(
        label="Prompt",
        placeholder="Describe the image you want...",
        value=(
            "a professional portrait of a confident woman "
            "in a modern office, natural lighting, high quality"
        ),
        lines=3,
    ),

    outputs=gr.Image(
        label="Generated Image",
        type="pil",
    ),

    title="Stable Diffusion XL Image Generator",

    description=(
        "Fast SDXL generation using NVIDIA T4 GPU"
    ),

    examples=[
        [
            "a professional portrait of a confident woman "
            "in a modern office, natural lighting, high quality"
        ],
        [
            "a futuristic city at night, neon lights, "
            "cinematic photography"
        ],
        [
            "a beautiful mountain landscape at sunrise, "
            "highly detailed"
        ],
        [
            "a cute golden retriever sitting in a park, "
            "realistic photography"
        ],
    ],
)


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