Update app.py

app.py

@@ -1,89 +1,153 @@
 import gradio as gr
 import subprocess
 import os
-from datetime import datetime
-from huggingface_hub import login, hf_hub_download, upload_file, HfApi, upload_folder
-import tempfile
 import shutil
-# --- Config ---
-REPO_ID = os.environ.get("HFPATH")
-HF_TOKEN = os.environ.get("MAGIC")
+import tempfile
+from datetime import datetime
+from huggingface_hub import (
+    login,
+    hf_hub_download,
+    upload_file,
+    upload_folder,
+    HfApi,
+)
+
+# -------------------------------------------------
+# Cấu hình (đặt các biến môi trường trước khi chạy)
+# -------------------------------------------------
+REPO_ID = os.getenv("HFPATH")          # ví dụ: "username/dataset-name"
+HF_TOKEN = os.getenv("MAGIC")          # token Hugging Face của bạn
 login(HF_TOKEN)
 
-# --- Download function ---
 api = HfApi()
 
-# --- Download function ---
-def download_from_hf(subfolder):
+
+# -------------------------------------------------
+# Hàm phụ: tải xuống tất cả file trong một subfolder
+# -------------------------------------------------
+def download_from_hf(subfolder: str):
+    """Trả về danh sách đường dẫn cục bộ của các file trong `subfolder`."""
     downloaded_files = []
-    
-    # List all files in the repository
-    all_files = api.list_repo_files(repo_id=REPO_ID,repo_type="dataset")
-    
-    # Filter files that start with the specified subfolder
+
+    # Lấy danh sách mọi file trong repo dataset
+    all_files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset")
+
+    # Chỉ giữ lại những file bắt đầu bằng subfolder yêu cầu
     for file in all_files:
-        if file.startswith(subfolder + "/"):
-            # Download the file
-            downloaded_file = hf_hub_download(repo_id=REPO_ID, filename=file,repo_type="dataset" )
-            downloaded_files.append(downloaded_file)
-    
+        if file.startswith(f"{subfolder}/"):
+            local_path = hf_hub_download(
+                repo_id=REPO_ID,
+                filename=file,
+                repo_type="dataset",
+            )
+            downloaded_files.append(local_path)
+
     return downloaded_files
-# --- Upload function ---
-def upload_to_hf(filepath):
+
+
+# -------------------------------------------------
+# Hàm phụ: tải lên một file đơn lẻ vào thư mục theo ngày
+# -------------------------------------------------
+def upload_to_hf(filepath: str):
     filename = os.path.basename(filepath)
-    SUBFOLDER = datetime.now().strftime("%Y%m%d")
-    path_in_repo = f"{SUBFOLDER}/{filename}"
+    subfolder = datetime.now().strftime("%Y%m%d")
+    path_in_repo = f"{subfolder}/{filename}"
+
     upload_file(
         path_or_fileobj=filepath,
-        path_in_repo=path_in_repo,
+        path_in_repo=api,
         repo_id=REPO_ID,
-        repo_type="dataset"
+        repo_type="dataset",
     )
     return f"https://huggingface.co/datasets/{REPO_ID}/blob/main/{path_in_repo}"
-def upload_to_fd(path):
-    SUBFOLDER = datetime.now().strftime("%Y%m%d")
-    upload_folder(folder_path=path,
-                  path_in_repo=SUBFOLDER,
-                  repo_id=REPO_ID,
-                  repo_type="dataset")
-    return f"https://huggingface.co/datasets/{REPO_ID}/blob/main/{SUBFOLDER}"
-
-def run_scripts(subfolder, source):
-    outputfile = []
+
+
+# -------------------------------------------------
+# Hàm phụ: tải lên toàn bộ thư mục vào thư mục theo ngày
+# -------------------------------------------------
+def upload_folder_to_hf(folder_path: str):
+    subfolder = datetime.now().strftime("%Y%m%d")
+    upload_folder(
+        folder_path=folder_path,
+        path_in_repo=subfolder,
+        repo_id=REPO_ID,
+        repo_type="dataset",
+    )
+    return f"https://huggingface.co/datasets/{REPO_ID}/blob/main/{subfolder}"
+
+
+# -------------------------------------------------
+# Hàm chính được Gradio gọi
+# -------------------------------------------------
+def run_scripts(subfolder: str, source):
+    """
+    1. Tải về các file mục tiêu từ `subfolder`.
+    2. Chạy `run.py` cho mỗi file, tạo file output.
+    3. Sao chép các output vào thư mục tạm.
+    4. Tải lên thư mục tạm lên Hugging Face.
+    Trả về danh sách tên file output.
+    """
+    output_files = []
     target_files = download_from_hf(subfolder)
-    
-    
-    # Tạo thư mục tạm và lưu file đầu ra vào đó
+
+    # Thư mục tạm để gom kết quả
     with tempfile.TemporaryDirectory() as temp_dir:
-        print(f"Thư mục tạm thời được tạo: {temp_dir}")
+        print(f"Thư mục tạm được tạo: {temp_dir}")
 
         for target_file in target_files:
-        target_extension = os.path.splitext(target_file)[-1]
-        #filename = os.path.splitext(target_file)[0]
-        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-        output_path = f"output_{timestamp}{target_extension}"
-        cmd1 = ["python3", "run.py", "-s", source.name, "-t", target_file, "-o", output_path, "--frame-processor", "face_swapper", '--many-faces']
-        subprocess.run(cmd1)
-        outputfile.append(output_path)
-        print(output_path)
-    # Sao chép file vào thư mục tạm
-       destination_file = os.path.join(temp_dir, os.path.basename(output_path))
-       shutil.copy(output_path, destination_file)
-       print(f"Đã sao chép {output_path} vào {destination_file}")
-
-        # Tải lên thư mục tạm lên Hugging Face
-        upload_to_fd(temp_dir)
-    return outputfile
+            # Giữ nguyên phần mở rộng của file gốc
+            _, ext = os.path.splitext(target_file)
+            timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+            output_path = f"output_{timestamp}{ext}"
+
+            # Lệnh gọi script swap mặt
+            cmd = [
+                "python3",
+                "run.py",
+                "-s",
+                source.name,
+                "-t",
+                target_file,
+                "-o",
+                output_path,
+                "--frame-processor",
+                "face_swapper",
+                "--many-faces",
+            ]
+
+            # Thực thi, sẽ ném lỗi nếu có vấn đề
+            subprocess.run(cmd, check=True)
+
+            output_files.append(output_path)
+            print(f"Tạo: {output_path}")
+
+            # Sao chép vào thư mục tạm
+            dest_path = os.path.join(temp_dir, os.path.basename(output_path))
+            shutil.copy(output_path, dest_path)
+            print(f"Đã sao chép {output_path} → {dest_path}")
+
+        # Tải lên toàn bộ thư mục tạm
+        upload_folder_to_hf(temp_dir)
+
+    return output_files
+
+
+# -------------------------------------------------
+# Giao diện Gradio
+# -------------------------------------------------
 iface = gr.Interface(
     fn=run_scripts,
     inputs=[
-        "text",
-        "file"
+        "text",   # tên subfolder
+        "file",   # ảnh nguồn
     ],
     outputs="files",
     title="Face Swapper",
-    description="Enter a subfolder name from the dataset and upload a source image to swap faces.",
-    live=False
+    description=(
+        "Nhập tên subfolder trong dataset và tải lên ảnh nguồn để thực hiện "
+        "đổi mặt."
+    ),
+    live=False,
 )
 
 iface.launch()