import gradio as gr
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"

# Charger le modèle avec scheduler
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to(device)
pipe.enable_attention_slicing()

def generate_image(prompt):
    if prompt.strip() == "":
        return None
    image = pipe(prompt, num_inference_steps=20, height=512, width=512).images[0]
    return image

description = """
# 🌄 Générateur de Paysages IA
Tape un mot-clé ou une idée de paysage et clique sur **Générer** !
"""

demo = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Description du paysage", placeholder="ex : forêt mystique au coucher du soleil"),
    outputs=gr.Image(label="Image générée"),
    title="Paysage IA",
    description=description,
)

demo.launch()