Update app.py

app.py

@@ -4,19 +4,20 @@ import torchvision.transforms as transforms
 from pathlib import Path
 import os
 
-# Cấu hình
-MODEL_PATH = "./models"
+# 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
+# Hàm tải mô hình từ thư mục local
 def load_model(model_name, size):
-    model_directory = Path(MODEL_PATH) / model_name / str(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 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")
+        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
@@ -38,10 +39,10 @@ def generate_image(model_name, size, alpha=0.5):
 
 # 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()]
+    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(MODEL_PATH) / model_name).iterdir() if folder.is_dir()]
+    sizes = [folder.name for folder in (Path("./models") / model_name).iterdir() if folder.is_dir()]
     return sizes
 
 # Tạo giao diện Gradio