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

# Charger le pipeline de diffusion depuis Hugging Face
model_name = "Yaquv/rickthenpc"
device = "cuda" if torch.cuda.is_available() else "cpu"

try:
    pipe = StableDiffusionPipeline.from_pretrained(model_name)
    pipe = pipe.to(device)
except Exception as e:
    print(f"Erreur lors du chargement du modèle : {e}")
    pipe = None

# Fonction de génération et de post-traitement
def generate_image(prompt):
    """
    Génère une image à partir du prompt en utilisant le modèle Hugging Face.
    """
    if pipe is None:
        return "Le modèle n'a pas pu être chargé."

    try:
        # Générer l'image
        image = pipe(prompt).images[0]

        return image

    except Exception as e:
        # Gestion des erreurs avec un retour informatif
        return f"Erreur lors de la génération : {str(e)}"

# Interface Gradio
iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.Image(label="Image générée"),
    title="Générateur d'images RickTheNPC",
    description="Entrez un prompt et générez une image avec le modèle RickTheNPC."
)

# Lancer l'application
if __name__ == "__main__":
    iface.launch()