Update app.py

app.py

@@ -3,15 +3,29 @@ import subprocess
 import shutil
 import os
 from datetime import datetime
-from huggingface_hub import login, upload_file
+from huggingface_hub import login, download_file, list_datasets, upload_file
 
 # --- Config ---
 REPO_ID = "mr2along/Picture"
 HF_TOKEN = os.environ.get("MAGIC")
-#print(HF_TOKEN)
-#if not HF_TOKEN:
-#    print("⚠️ Không tìm thấy biến môi trường 'MAGIC'")
 login(HF_TOKEN)
+
+# --- Download function ---
+def download_from_hf(subfolder):
+    downloaded_files = []
+    files = list_datasets(REPO_ID)
+    for file in files:
+        if file.filename.startswith(subfolder + "/"):
+            filename = file.filename.split("/")[-1]
+            download_file(
+                path_or_url=file.filename,
+                repo_id=REPO_ID,
+                repo_type="dataset",
+                local_path=filename
+            )
+            downloaded_files.append(filename)
+    return downloaded_files
+
 # --- Upload function ---
 def upload_to_hf(filepath):
     filename = os.path.basename(filepath)
@@ -25,14 +39,15 @@ def upload_to_hf(filepath):
     )
     return f"https://huggingface.co/datasets/{REPO_ID}/blob/main/{filename}"
 
-def run_scripts(target, source):
-    outputfile=[]
-    for target_file in target :
-        target_extension = os.path.splitext(target_file.name)[-1]
-        target_name = os.path.splitext(target_file.name)[1]
+def run_scripts(subfolder, source):
+    outputfile = []
+    target_files = download_from_hf(subfolder)
+    for target_file in target_files:
+        target_extension = os.path.splitext(target_file)[-1]
+        target_name = os.path.splitext(target_file)[1]
         timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
         output_path1 = "output" + timestamp + target_extension
-        cmd1 = ["python3", "run.py", "-s", source.name, "-t", target_file.name, "-o", output_path1, "--frame-processor", "face_swapper","face_enhancer",'--many-faces']
+        cmd1 = ["python3", "run.py", "-s", source.name, "-t", target_file, "-o", output_path1, "--frame-processor", "face_swapper","face_enhancer",'--many-faces']
         subprocess.run(cmd1)
         outputfile.append(output_path1)
         print(output_path1)
@@ -42,13 +57,13 @@ def run_scripts(target, source):
 iface = gr.Interface(
     fn=run_scripts,
     inputs=[
-        "files",
+        gr.Textbox(label="Target subfolder"),
         "file"
     ],
     outputs="files",
     title="Face swapper",
-    description="Upload a target image/video and a source image to swap faces.",
+    description="Enter a subfolder name from the dataset and upload a source image to swap faces.",
     live=False
 )
 
-iface.launch()
+iface.launch()