import gradio as gr
import torch
import numpy as np
import random
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image

dtype = torch.float16
device = "cuda" if torch.cuda.is_available() else "cpu"
MAX_SEED = np.iinfo(np.int32).max

MODEL_ID = "runwayml/stable-diffusion-v1-5"

STYLES = {
    "Anime": {
        "prompt": "anime style portrait, detailed face, vibrant colors, high quality, studio ghibli"
    },
    "Pixel Art": {
        "prompt": "pixel art avatar, 16-bit style, retro game character, sharp pixels"
    },
    "Ghibli": {
        "prompt": "ghibli studio style illustration, soft colors, dreamy atmosphere, watercolor"
    },
}

print("Chargement du modèle...")
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=dtype,
    safety_checker=None
)
pipe.to(device)
print("Modèle chargé !")

def generate_avatar(
    input_image,
    style_name,
    strength,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True)
):
    if input_image is None:
        raise gr.Error("Veuillez uploader une photo !")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    prompt = STYLES[style_name]["prompt"]

    input_image = input_image.resize((512, 512))

    result = pipe(
        prompt=prompt,
        image=input_image,
        strength=strength,
        num_inference_steps=20,
        guidance_scale=7.5,
        generator=generator,
    ).images[0]

    return (input_image, result), seed


with gr.Blocks() as demo:
    gr.Markdown("# Avatar Generator — Transformez votre photo en personnage")

    with gr.Row():
        with gr.Column():
            input_img = gr.Image(label="Votre photo", type="pil")
            style = gr.Radio(
                choices=list(STYLES.keys()),
                label="Style d'avatar",
                value="Anime"
            )
            strength = gr.Slider(
                label="Intensité du style",
                minimum=0.3, maximum=0.95,
                step=0.05, value=0.7,
                info="Bas = garde votre visage, Haut = stylisation forte"
            )
            with gr.Accordion("Paramètres avancés", open=False):
                seed = gr.Slider(0, MAX_SEED, step=1, value=42, label="Seed")
                randomize_seed = gr.Checkbox(label="Seed aléatoire", value=True)

            btn = gr.Button("Générer mon avatar", variant="primary")

        result_slider = gr.ImageSlider(label="Avant / Après", type="pil")

    seed_out = gr.Number(label="Seed utilisé", visible=False)

    btn.click(
        fn=generate_avatar,
        inputs=[input_img, style, strength, seed, randomize_seed],
        outputs=[result_slider, seed_out]
    )

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