import gradio as gr
import requests
from PIL import Image
import io

API_URL = "https://k2mar-mon-api-sd.hf.space/generate"

def generate_image(prompt, steps, guidance_scale):
    if not prompt.strip():
        raise gr.Error("Veuillez entrer un prompt.")

    try:
        response = requests.post(
            API_URL,
            json={
                "prompt": prompt,
                "steps": int(steps),
                "guidance_scale": float(guidance_scale)
            },
            timeout=600
        )

        if response.status_code == 503:
            raise gr.Error("Le serveur est occupé, réessayez dans quelques minutes.")

        if response.status_code != 200:
            raise gr.Error(f"Erreur API : {response.status_code}")

        image = Image.open(io.BytesIO(response.content))
        return image

    except requests.exceptions.Timeout:
        raise gr.Error("Timeout — la génération a pris trop de temps.")
    except requests.exceptions.ConnectionError:
        raise gr.Error("Impossible de joindre l'API. Vérifiez que le Space est démarré.")

with gr.Blocks(title="Générateur d'images IA") as demo:

    gr.Markdown("# Générateur d'images IA")
    gr.Markdown("Modèle Stable Diffusion fine-tuné · Solar Panel")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Ex: solar panel rooftop installation, technical diagram...",
                lines=3
            )
            with gr.Row():
                steps = gr.Slider(
                    minimum=5, maximum=30, value=20, step=1,
                    label="Steps"
                )
                guidance = gr.Slider(
                    minimum=1.0, maximum=15.0, value=7.5, step=0.5,
                    label="Guidance scale"
                )

            gr.Examples(
                examples=[
                    ["solar panel installation diagram, technical drawing", 20, 7.5],
                    ["electrical wiring schematic, blueprint style", 20, 7.5],
                    ["rooftop solar array, photorealistic, sunny day", 25, 8.0],
                    ["solar energy system overview, infographic style", 20, 7.0],
                ],
                inputs=[prompt, steps, guidance]
            )

            btn = gr.Button("Générer", variant="primary")

        with gr.Column():
            output = gr.Image(label="Image générée", type="pil")

    btn.click(
        fn=generate_image,
        inputs=[prompt, steps, guidance],
        outputs=output
    )

demo.launch()