import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
import os # Pour gérer les variables d'environnement

# --- Configuration du modèle ---
# Choisissez un modèle Stable Diffusion.
# "runwayml/stable-diffusion-v1-5" est un bon point de départ.
# Pour des modèles plus petits et plus rapides, vous pourriez envisager "stabilityai/stable-diffusion-xl-base-1.0"
# (attention, XL est plus lourd) ou des versions optimisées.
MODEL_ID = "runwayml/stable-diffusion-v1-5" # Plus petit et plus rapide pour le Free Tier

# Vérifiez si CUDA est disponible (GPU) et utilisez-le si oui.
# Sinon, utilisez le CPU (ce sera beaucoup plus lent).
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Chargement du modèle {MODEL_ID} sur {DEVICE}...")

# Charger le pipeline de Stable Diffusion.
# Utilisez torch_dtype=torch.float16 pour économiser de la mémoire GPU si disponible.
# Cela est crucial pour le free tier de Hugging Face Spaces.
pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32)
pipe.to(DEVICE)

print("Modèle chargé.")

# --- Fonction de génération d'images ---
def generate_image(prompt, negative_prompt="", num_inference_steps=30, guidance_scale=7.5):
    """
    Génère une image à partir d'un prompt texte.
    """
    with torch.no_grad(): # Désactive le calcul des gradients pour l'inférence
        image = pipe(
            prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale
        ).images[0]
    return image

# --- Création de l'interface Gradio ---
print("Création de l'interface Gradio...")

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Votre prompt (description de l'image)", placeholder="Un astronaute chevauchant un cheval sur Mars, style réaliste"),
        gr.Textbox(label="Prompt négatif (éléments à éviter)", placeholder="mauvaise qualité, moche, flou"),
        gr.Slider(minimum=10, maximum=100, step=1, value=30, label="Étapes d'inférence (qualité/temps)"),
        gr.Slider(minimum=1.0, maximum=20.0, step=0.1, value=7.5, label="Échelle de guidance (fidélité au prompt)"),
    ],
    outputs=gr.Image(type="pil", label="Image générée"),
    title="Mon Fooocus Gratuit",
    description="Générez des images étonnantes avec Stable Diffusion, hébergé gratuitement sur Hugging Face Spaces !",
    article="""
    <p>Cette application est une démo simplifiée de génération d'images, similaire à Fooocus, utilisant Hugging Face Diffusers et Gradio.</p>
    <p>Elle est hébergée gratuitement sur Hugging Face Spaces. Les performances peuvent varier en fonction de la charge et du modèle choisi.</p>
    """
)

# Lancer l'interface Gradio
# Pour le développement local, vous pouvez utiliser share=True pour obtenir un lien public temporaire.
# Pour Hugging Face Spaces, vous n'avez pas besoin de share=True.
if __name__ == "__main__":
    iface.launch(debug=True) # debug=True affichera plus d'informations utiles en cas de problème