import os
import gradio as gr
import torch
from diffusers import AutoPipelineForText2Image

# ---- Model choice (a fast, lightweight VLM for text->image) ----
MODEL_ID = os.environ.get("MODEL_ID", "stabilityai/sdxl-turbo")

# ---- Load pipeline  ----
def load_pipeline():
    
    use_cuda = torch.cuda.is_available()
    dtype = torch.float16 if use_cuda else torch.float32

    
    kwargs = {"torch_dtype": dtype}
    if use_cuda:
        kwargs["variant"] = "fp16"

    pipe = AutoPipelineForText2Image.from_pretrained(MODEL_ID, **kwargs)
    if use_cuda:
        pipe = pipe.to("cuda")
    else:
        pipe = pipe.to("cpu")
    

    return pipe

PIPE = load_pipeline()

# ---- Generation function  ----
def generate_image(prompt, steps, guidance, width, height):
    
    
    gen = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(0)

    # SDXL-Turbo is designed for very few steps (1–6).
    result = PIPE(
        prompt=prompt,
        negative_prompt= None,
        num_inference_steps=int(steps),
        guidance_scale=int(guidance),
        width=int(width),
        height=int(height),
        generator=gen,
    )
    image = result.images[0]
    return image

# ---- Gradio UI (Blocks) ----
with gr.Blocks(title="Text→Image (Diffusers + Gradio)") as interface:
    gr.Markdown(
        "# Text → Image\n"
        f"**Model:** `{MODEL_ID}` "
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt", placeholder="a mountain landscape with a warm sunlight"
            )
            
            with gr.Row():
                steps = gr.Slider(1, 6, value=4, step=1, label="Steps")
                guidance = gr.Slider(0, 15, value=1, step=1, label="Guidance")

            with gr.Row():
                width = gr.Dropdown(
                    choices=[384, 448, 512, 640, 768, 1024], value=384, label="Width"
                )
                height = gr.Dropdown(
                    choices=[384, 448, 512, 640, 768, 1024], value=384, label="Height"
                )
            
            run_btn = gr.Button("Generate", variant="primary")

        with gr.Column(scale=1):
            out = gr.Image(label="Result", type="pil")

    run_btn.click(
        fn=generate_image,
        inputs=[prompt, steps, guidance, width, height],
        outputs=[out],
        queue=True,
        api_name="generate",
    )

if __name__ == "__main__":
    
    interface.queue(max_size=32).launch()