import gradio as gr
import numpy as np
import random
from diffusers import DiffusionPipeline
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "stabilityai/sdxl-turbo"  # Replace with the model you want to use

torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
pipe = pipe.to(device)

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

def generate_cat_image():
    seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    
    image = pipe(
        prompt="A cute cat, highly detailed, 8k",
        negative_prompt="",
        guidance_scale=7.5,
        num_inference_steps=20,
        width=512,
        height=512,
        generator=generator,
    ).images[0]
    
    return image

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(" # Cat Image Generator 🐱")
        result = gr.Image(label="Generated Cat Image", show_label=False)
        run_button = gr.Button("Generate New Cat")
    
    run_button.click(fn=generate_cat_image, inputs=[], outputs=[result])

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