import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

def generate_image(prompt):
    try:
        # Charger le modèle en précision float32 pour CPU
        model_id = "runwayml/stable-diffusion-v1-5"
        pipe = StableDiffusionPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float32,  # Float32 pour compatibilité CPU
            low_cpu_mem_usage=True,     # Réduit l'usage mémoire
        )
        pipe = pipe.to("cpu")  # Forcer explicitement CPU

        # Générer l'image (réduire la taille pour économiser mémoire)
        image = pipe(prompt, num_inference_steps=20, height=256, width=256).images[0]
        return image
    except Exception as e:
        return f"Erreur : {str(e)}"

# Interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Image Generator with Stable Diffusion (CPU)")
    with gr.Row():
        prompt = gr.Textbox(label="Enter your prompt")
        btn = gr.Button("Generate")
    output = gr.Image()
    btn.click(generate_image, inputs=prompt, outputs=output)

demo.launch()