import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

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

# Avatar 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"
]

# Function to generate avatar
def generate_avatar(image, style):
    prompt = f"{style} style portrait of a person"
    result = pipe(prompt=prompt, image=image).images[0]
    return result

# Build Gradio UI
gr.Interface(
    fn=generate_avatar,
    inputs=[
        gr.Image(type="pil", label="Upload Your Selfie"),
        gr.Dropdown(choices=styles, label="Choose a Style")
    ],
    outputs=gr.Image(label="Your Avatar"),
    title="AvatarifyAI",
    description="Turn your selfie into an AI-generated avatar in awesome styles!"
)
demo.launch()