import os
import gradio as gr
from dotenv import load_dotenv
from huggingface_hub import InferenceClient

load_dotenv()

HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise RuntimeError(
        "No se encontró HF_TOKEN"
    )

MODEL_ID = "FireRedTeam/FireRed-Image-Edit-1.0"

client = InferenceClient(
    provider="auto",
    api_key=HF_TOKEN,
    timeout=300,
)

def edit_image(input_image, prompt):
    if input_image is None:
        raise gr.Error("Sube una imagen primero.")

    if not prompt or not prompt.strip():
        raise gr.Error("Escribe un prompt de edición.")

    try:
        edited_image = client.image_to_image(
            image=input_image,
            prompt=prompt.strip(),
            model=MODEL_ID,
        )
        return edited_image

    except Exception as e:
        raise gr.Error(f"Error durante la inferencia: {e}")

with gr.Blocks(title="FireRed Image Editor") as demo:
    gr.Markdown(
        """
# FireRed Image Editor
Sube una imagen, escribe una instrucción y genera una versión editada.
"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="Upload Images",
                type="pil",
                height=420
            )

            prompt = gr.Textbox(
                label="Edit Prompt",
                placeholder="e.g., transform into anime, upscale, change lighting...",
                lines=2
            )

            edit_btn = gr.Button("Edit Image", variant="primary", size="lg")

        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Output Image",
                type="pil",
                height=420
            )

    edit_btn.click(
        fn=edit_image,
        inputs=[input_image, prompt],
        outputs=output_image,
        show_progress="full"
    )

    prompt.submit(
        fn=edit_image,
        inputs=[input_image, prompt],
        outputs=output_image
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        theme=gr.themes.Soft(),
        ssr_mode=False,
    )