import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
import PIL
import io

# --- Configuration du modèle ---
MODEL_ID = "Dipl0/pepe-diffuser"

# Détection de l'appareil (GPU si disponible, sinon CPU)
device = "cuda" if torch.cuda.is_available() else "cpu"

# Chargement du pipeline Stable Diffusion
# Utilisation de torch_dtype=torch.float16 pour la VRAM si vous avez une carte GPU compatible
# Sinon, retirez-le ou changez-le pour torch.float32 pour le CPU ou les GPU moins performantes
try:
    pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float16)
    pipe = pipe.to(device)
except Exception as e:
    print(f"Avertissement : Chargement en float16 a échoué ({e}). Tentative en float32.")
    pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
    pipe = pipe.to(device)

# --- Fonction de génération d'image ---
def generate_image(prompt: str, negative_prompt: str, num_inference_steps: int, guidance_scale: float, seed: int) -> PIL.Image.Image:
    """
    Génère une image à partir du prompt en utilisant le modèle Dipl0/pepe-diffuser.
    """
    # Configuration du générateur pour la reproductibilité si un seed est fourni
    generator = None
    if seed > 0:
        generator = torch.Generator(device=device).manual_seed(seed)

    # Exécution du pipeline
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator,
    ).images[0]

    return image

# --- Création de l'interface Gradio ---

# Définition des composants d'entrée
prompt_input = gr.Textbox(
    label="Prompt (Description de l'image de Pepe)",
    placeholder="pepe the frog smoking a cigar, digital art, high resolution",
    lines=2
)
negative_prompt_input = gr.Textbox(
    label="Negative Prompt (Ce que vous ne voulez PAS dans l'image)",
    placeholder="blurry, low quality, deformed, text, watermark",
    lines=2
)
steps_slider = gr.Slider(
    minimum=1,
    maximum=10,
    step=5,
    value=1,
    label="Nombre d'étapes d'inférence (plus = meilleure qualité mais plus long)"
)
guidance_slider = gr.Slider(
    minimum=1.0,
    maximum=15.0,
    step=0.5,
    value=7.5,
    label="Échelle de Guidance (plus = l'image suit plus le prompt)"
)
seed_input = gr.Number(
    label="Seed (Graine) pour la reproductibilité (0 pour aléatoire)",
    value=0,
    minimum=0,
    step=1
)

# Définition du composant de sortie
image_output = gr.Image(
    label="Image Générée",
    type="pil",
    height=512,
    width=512
)

# Création de l'interface complète
gr.Interface(
    fn=generate_image,
    inputs=[
        prompt_input,
        negative_prompt_input,
        steps_slider,
        guidance_slider,
        seed_input
    ],
    outputs=image_output,
    title="🐸 Pepe Diffuser - Text-to-Image",
    description=f"Générez des images personnalisées de Pepe la grenouille en utilisant le modèle Hugging Face **{MODEL_ID}**.",
    allow_flagging="never",
).launch()