import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Tải mô hình Stable Diffusion
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    use_auth_token=False
)
#pipe = pipe.to("cuda")  # Giả sử bạn dùng GPU
pipe.enable_attention_slicing()  # Tối ưu RAM

# Hàm tạo hình ảnh
def generate_image(prompt, negative_prompt="", num_inference_steps=50, guidance_scale=7.5):
    image = pipe(
        prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=int(num_inference_steps),
        guidance_scale=guidance_scale
    ).images[0]
    return image

# Tạo giao diện với Blocks
with gr.Blocks() as demo:
    gr.Markdown("# Text-to-Image with Stable Diffusion")
    gr.Markdown("Enter a prompt to generate an image.")
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="E.g., 'A futuristic city at sunset'")
        negative_prompt = gr.Textbox(label="Negative Prompt (optional)", placeholder="E.g., 'blurry, low quality'")
    with gr.Row():
        steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=50, step=1)
        guidance = gr.Slider(label="Guidance Scale", minimum=1, maximum=20, value=7.5, step=0.5)
    btn = gr.Button("Generate")
    output = gr.Image(label="Generated Image")
    btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance],
        outputs=output
    )

# Khởi chạy
demo.launch()