import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
import subprocess

# ติดตั้ง flash-attn แม้จะไม่ได้ใช้โดยตรง (ข้าม build CUDA)
subprocess.run(
    'pip install flash-attn --no-build-isolation',
    env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
    shell=True
)

# ใช้ CPU
device = "cpu"

# โหลดโมเดลเบา: sd-turbo
pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/sd-turbo",
    torch_dtype=torch.float32
)
pipe = pipe.to(device)
pipe.safety_checker = None  # ปิด safety checker เพื่อความเร็ว

# เปิด attention_slicing (แม้บน CPU ก็ช่วยเรื่องหน่วยความจำ)
pipe.enable_attention_slicing()

# ฟังก์ชันสร้างภาพ
def generate_image(prompt):
    result = pipe(prompt, num_inference_steps=10, guidance_scale=3.0)
    image = result.images[0]
    return image

# Gradio UI
io = gr.Interface(
    fn=generate_image,
    inputs=[gr.Textbox(label="Enter your prompt")],
    outputs=[gr.Image(label="Generated Image")],
    theme="Yntec/HaleyCH_Theme_Orange",
    description="⚠ Running on CPU using sd-turbo. Optimized for speed with low inference steps."
)

# เปิด Gradio ด้วย queue ป้องกันค้างถ้ามีหลายคำสั่ง
io.queue().launch(debug=True)