import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

# Load pre-trained model từ Hugging Face Hub
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
)

pipe.to("cpu")
pipe.enable_attention_slicing()

# Hàm sinh ảnh từ prompt
def generate_image(prompt):
    image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
    return image

# Giao diện Gradio
demo = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(lines=2, placeholder="Nhập mô tả bằng tiếng Anh...", label="Prompt"),
    outputs=gr.Image(label="Ảnh tạo ra"),
    title="Tạo ảnh từ văn bản với Stable Diffusion",
    description="Ví dụ: 'handwritten Vietnamese text on white paper' hoặc 'bold black Vietnamese text on wooden background'"
)

# Chạy app
if __name__ == "__main__":
    demo.launch()