import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load model
pipe = StableDiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo")
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# Available styles
styles = [
    "Anime", "Cartoon", "Pixel Art", "Cyberpunk", "Disney",
    "GTA", "Barbie", "Pixar", "Watercolor", "3D Render",
    "Futuristic Soldier", "Zombie Apocalypse", "Royal Portrait",
    "Oil Painting", "Astronaut", "Cyber Angel", "Superhero",
    "Line Art Sketch", "Pop Art", "Cute Chibi"
]

# Avatar generation function
def generate_avatar(name, style):
    prompt = f"{style} style portrait of {name}"
    image = pipe(prompt=prompt).images[0]
    return image

# Gradio UI
gr.Interface(
    fn=generate_avatar,
    inputs=[
        gr.Textbox(label="Who do you want to see? (e.g. a young woman, a boy, etc.)"),
        gr.Dropdown(choices=styles, label="Choose a Style")
    ],
    outputs=gr.Image(label="Generated Avatar"),
    title="AvatarifyAI",
    description="Generate stylish avatars using SDXL Turbo. Enter who you want to see and select a style!"
).launch()