import gradio as gr
import torch
from diffusers import AutoPipelineForText2Image

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32
model_id = "stabilityai/sd-turbo"

print("[IMAGE] Loading", model_id, "on", device)
pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=dtype)
pipe = pipe.to(device)
print("[IMAGE] Pipeline ready.")


def generate_image(prompt):
    full_prompt = prompt + ", high quality, realistic photo"
    return pipe(prompt=full_prompt, num_inference_steps=2, guidance_scale=0.0).images[0]


with gr.Blocks(title="Tachelhit Image Generator") as demo:
    prompt_input = gr.Textbox(label="Prompt")
    generate_btn = gr.Button("Generate")
    image_output = gr.Image(label="Result", type="filepath")
    generate_btn.click(generate_image, prompt_input, image_output, api_name="/generate")

demo.launch(server_name="0.0.0.0", server_port=7860)