import gradio as gr
from PIL import Image
import torch
import warnings

# Tắt cảnh báo về mô hình
warnings.filterwarnings("ignore")

# Cấu hình thiết bị và Mô hình
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = None
MODEL_ID = "runwayml/stable-diffusion-v1-5" 

try:
    from diffusers import StableDiffusionImg2ImgPipeline
    
    print(f"Đang tải mô hình Img2Img: {MODEL_ID} trên {device}...")
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
        MODEL_ID, 
        torch_dtype=torch.float16 if device == "cuda" else torch.float32
    ).to(device)
    print("Tải mô hình hoàn tất.")
    
except ImportError:
    # Lỗi khi không cài đặt 'diffusers'
    pipe = None
except Exception as e:
    # Lỗi khi tải mô hình (ví dụ: lỗi mạng, lỗi bộ nhớ)
    print(f"Lỗi khi tải mô hình Stable Diffusion: {e}")
    pipe = None


# --- HÀM TẠO ẢNH CHÍNH ---
def generate_image(
    reference_image: Image.Image, 
    prompt: str, 
    style: str, 
    guidance_scale: float, 
    strength: float
):
    """Thực hiện việc tạo ảnh bằng Img2Img Pipeline."""
    if pipe is None:
        raise gr.Error("Mô hình AI chưa được tải thành công. Vui lòng kiểm tra nhật ký lỗi.")

    # Chuẩn bị Prompt (Mô tả)
    style_keywords = {
        "photorealistic": "photorealistic, 4k, ultra detailed, cinematic lighting",
        "anime": "anime style, vibrant colors, trending on pixiv",
        "oilpainting": "oil painting, thick impasto, masterpiece on canvas",
        "conceptart": "concept art, matte painting, dramatic lighting, highly detailed"
    }
    full_prompt = f"{prompt}, {style_keywords.get(style, '')}"
    
    # Xử lý Kích thước Ảnh Tham Chiếu
    # Resizing ảnh đầu vào (thường là 512x512)
    reference_image = reference_image.resize((512, 512))
    
    # Gọi Pipeline Tạo Ảnh
    with torch.no_grad():
        result = pipe(
            prompt=full_prompt,
            image=reference_image,
            strength=strength,         
            guidance_scale=guidance_scale,
            num_inference_steps=50     
        ).images[0]
        
    return result


# --- THIẾT LẬP GIAO DIỆN GRADIO ---

gr.Interface(
    fn=generate_image,
    inputs=[
        # Đã XÓA 'source="upload"' để khắc phục lỗi TypeError
        gr.Image(type="pil", label="🖼️ Ảnh Tham Chiếu"), 
        gr.Textbox(
            label="📝 Mô Tả và Từ Khóa", 
            placeholder="Nhập mô tả chi tiết về ảnh bạn muốn tạo..."
        ),
        gr.Dropdown(
            label="🎨 Phong cách Nghệ thuật", 
            choices=["photorealistic", "anime", "oilpainting", "conceptart"], 
            value="photorealistic"
        ),
        gr.Slider(
            label="Mức độ Tuân theo Prompt (Guidance Scale)", 
            minimum=1.0, maximum=15.0, step=0.5, value=7.5
        ),
        gr.Slider(
            label="Cường độ Ảnh Tham Chiếu (Strength)", 
            minimum=0.0, maximum=1.0, step=0.05, value=0.8,
            info="Strength cao (0.8-1.0) giữ lại ít chi tiết ảnh gốc, Strength thấp (0.3-0.6) giữ lại nhiều hơn."
        )
    ],
    outputs=gr.Image(label="✅ Ảnh Kết Quả", type="pil"),
    title="✨ Ứng Dụng Tạo Ảnh Từ Ảnh Tham Chiếu (Stable Diffusion Img2Img)",
    description="Sử dụng Stable Diffusion Img2Img để tạo ảnh mới dựa trên ảnh tham chiếu và mô tả văn bản.",
    allow_flagging="never",
).launch(debug=True)