Create app.py

app.py

@@ -0,0 +1,66 @@
+import gradio as gr
+import torch
+import torchvision.transforms as transforms
+from pathlib import Path
+import os
+
+# Cấu hình
+MODEL_PATH = "./models"
+TARGET_EPOCH = 47
+
+# Hàm tải mô hình
+def load_model(model_name, size):
+    model_directory = Path(MODEL_PATH) / model_name / str(size)
+    if not model_directory.exists():
+        raise FileNotFoundError(f"Không tìm thấy thư mục model: {model_directory}")
+    
+    # Tải mô hình với trọng số
+    generator = torch.load(model_directory / f"generator_{TARGET_EPOCH}.pth")
+    mapping_network = torch.load(model_directory / f"mapping_network_{TARGET_EPOCH}.pth")
+    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(MODEL_PATH).glob("*") if folder.is_dir()]
+
+def get_sizes(model_name):
+    sizes = [folder.name for folder in (Path(MODEL_PATH) / 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()