import torch
from diffusers import AutoPipelineForText2Image
import gradio as gr

# DÙNG PHIÊN BẢN NHẸ CỦA SDXL-TURBO (chạy được trên CPU Space miễn phí)
# Model này chỉ 2.5GB, không cần CUDA, vẫn đẹp
pipe = AutoPipelineForText2Image.from_pretrained(
    "segmind/Segmind-Vega",  # SDXL-Turbo nhẹ nhất 2025, chạy ngon CPU
    torch_dtype=torch.float32,  # CPU thì dùng float32
    safety_checker=None
)

# Load LoRA Á Đông (vẫn dùng được)
pipe.load_lora_weights("ntc-ai/SDXL-LoRA-slider.asian", weight_name="asian.safetensors")

def generate(prompt):
    # Prompt booster Việt Nam
    full_prompt = f"{prompt}, vietnamese girl, ao dai, realistic, detailed, 8k"
    image = pipe(
        prompt=full_prompt,
        num_inference_steps=6,     # CPU thì tăng lên 6 để đẹp hơn
        guidance_scale=0.0,
        height=768,
        width=768
    ).images[0]
    return image

# UI đẹp
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# SDXL-Turbo Việt Nam – Chạy Miễn Phí 24/7")
    gr.Markdown("Không cần GPU • Ảnh đẹp • Đúng áo dài, nón lá, phở")

    txt = gr.Textbox(
        label="Prompt tiếng Việt",
        placeholder="cô gái mặc áo dài đứng giữa cánh đồng lúa vàng",
        lines=2
    )
    btn = gr.Button("Tạo ảnh", variant="primary")
    out = gr.Image()

    gr.Examples([
        "cô gái mặc áo dài đứng giữa cánh đồng lúa vàng",
        "mèo đội nón lá ăn phở Hà Nội",
        "Hội An về đêm đèn lồng lung linh",
        "cà phê sữa đá Sài Gòn"
    ], inputs=txt)

    btn.click(generate, txt, out)
    txt.submit(generate, txt, out)

demo.launch()