import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

# Danh sách các model nhẹ, phù hợp cho CPU
MODELS = {
    "SD v1.5 (Cơ bản)": "runwayml/stable-diffusion-v1-5",
    "Tiny SD (Siêu nhanh - Khuyên dùng)": "segmind/tiny-sd",
    "Dreamlike Photoreal": "dreamlike-art/dreamlike-photoreal-2.0"
}

current_model_id = ""
pipe = None

def load_model(model_id, progress=gr.Progress()):
    global pipe, current_model_id
    if model_id != current_model_id:
        progress(0, desc="Đang tải model mới... (Vui lòng đợi)")
        pipe = StableDiffusionPipeline.from_pretrained(
            model_id, 
            torch_dtype=torch.float32,
            safety_checker=None # Tắt để tiết kiệm RAM
        )
        # Tối ưu scheduler để chạy nhanh hơn
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        current_model_id = model_id
    return pipe

def generate(prompt, model_name, steps, progress=gr.Progress(track_tqdm=True)):
    if not prompt:
        return None
    
    # 1. Load model đã chọn
    model_id = MODELS[model_name]
    pipeline = load_model(model_id)
    
    # 2. Tạo ảnh với Progress Bar
    # track_tqdm=True sẽ tự động bắt các bước loop của diffusers để hiện %
    image = pipeline(
        prompt, 
        num_inference_steps=int(steps),
        guidance_scale=7.0
    ).images[0]
    
    return image

# Giao diện UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"# 🎨 Stable Diffusion Free CPU - Minh Đức AI")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Mô tả ảnh", placeholder="Một con mèo đi thám hiểm mặt trăng...")
            model_name = gr.Dropdown(choices=list(MODELS.keys()), value="Tiny SD (Siêu nhanh - Khuyên dùng)", label="Chọn Model")
            steps = gr.Slider(minimum=10, maximum=30, value=15, step=1, label="Số bước chạy (Steps - Càng cao càng chậm)")
            btn = gr.Button("Bắt đầu tạo ảnh", variant="primary")
        
        with gr.Column():
            output = gr.Image(label="Kết quả")

    btn.click(fn=generate, inputs=[prompt, model_name, steps], outputs=output)

demo.launch()