import gradio as gr
import torch
import torchvision.transforms as transforms
from pathlib import Path
import os

# Lấy Hugging Face token từ biến môi trường
HF_TOKEN = os.getenv("HF_TOKEN")  # Token từ biến môi trường
MODEL_REPO = "Arrcttacsrks/FaceStyleGan"
TARGET_EPOCH = 47

# Hàm tải mô hình từ thư mục local
def load_model(model_name, size):
    model_directory = Path(f"./models/{model_name}/{size}")
    if not model_directory.exists():
        raise FileNotFoundError(f"Không tìm thấy thư mục mô hình: {model_directory}")

    # Tải trọng số mô hình
    mapping_network = torch.load(model_directory / f"mapping_network_{TARGET_EPOCH}.pth", map_location="cpu")
    generator = torch.load(model_directory / f"generator_{TARGET_EPOCH}.pth", map_location="cpu")
    return generator, mapping_network

# Hàm sinh ảnh
def generate_image(model_name, size, alpha=0.5):
    generator, mapping_network = load_model(model_name, size)
    
    # Tạo noise
    noise = torch.randn(1, 512)  # Kích thước noise có thể thay đổi tùy theo mô hình
    with torch.no_grad():
        # Sinh ảnh từ generator và mapping network
        sample_w = mapping_network(noise)
        generated_image = generator(sample_w)
        
        # Chuyển đổi ảnh về định dạng numpy
        transform = transforms.ToPILImage()
        image = transform((generated_image + 1) / 2)  # Chuyển đổi ảnh về khoảng [0, 1]
        
    return image

# Lấy danh sách mô hình và kích thước
def get_model_names():
    return [folder.name for folder in Path("./models").glob("*") if folder.is_dir()]

def get_sizes(model_name):
    sizes = [folder.name for folder in (Path("./models") / model_name).iterdir() if folder.is_dir()]
    return sizes

# Tạo giao diện Gradio
with gr.Blocks() as demo:
    gr.Markdown("# FaceStyleGan Image Generator")
    
    model_name = gr.Dropdown(choices=get_model_names(), label="Chọn model", interactive=True)
    size = gr.Dropdown(label="Chọn kích thước", interactive=True)
    model_name.change(fn=get_sizes, inputs=model_name, outputs=size)
    
    alpha = gr.Slider(0, 1, 0.5, label="Hệ số điều chỉnh (alpha)")
    generate_button = gr.Button("Tạo Ảnh")
    
    output = gr.Image(label="Ảnh Đã Tạo")
    
    generate_button.click(
        fn=generate_image,
        inputs=[model_name, size, alpha],
        outputs=output
    )

demo.launch()