Update app.py

app.py

@@ -10,58 +10,33 @@ from face_swapper import Inswapper, paste_to_whole
 from face_analyser import analyse_face
 from face_enhancer import load_face_enhancer_model, get_available_enhancer_names
 import tempfile
-from huggingface_hub import login, hf_hub_download, upload_file, HfApi,upload_folder
+from huggingface_hub import login, hf_hub_download, upload_file, HfApi, upload_folder
 from datetime import datetime
 import shutil
 
-
 # --- Config ---
 REPO_ID = os.environ.get("HFPATH")
 HF_TOKEN = os.environ.get("MAGIC")
 login(HF_TOKEN)
+
 # --- Download function ---
 api = HfApi()
 
-# --- Download function ---
 def download_from_hf(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
+    all_files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset")
     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_file = hf_hub_download(repo_id=REPO_ID, filename=file, repo_type="dataset")
             downloaded_files.append(downloaded_file)
-    
     return downloaded_files
+
 # --- Upload function ---
-def upload_to_hf(filepath):
-    filename = os.path.basename(filepath)
-    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,
-        repo_id=REPO_ID,
-        repo_type="dataset"
-    )
-    return f"https://huggingface.co/datasets/{REPO_ID}/blob/main/{path_in_repo}"
-    # --- Upload function ---
 def upload_to_fd(path):
-    #filename = os.path.basename(filepath)
     SUBFOLDER = datetime.now().strftime("%Y%m%d")
-    path_in_repo = f"{SUBFOLDER}"
-    upload_folder(
-        folder_path=path,
-        path_in_repo=path_in_repo,
-        repo_id=REPO_ID,
-        repo_type="dataset"
-    )
-    return f"https://huggingface.co/datasets/{REPO_ID}/blob/main/{path_in_repo}"
-# ------------------------------
+    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}"
+
 # ------------------------------ ARGS ------------------------------
 parser = argparse.ArgumentParser(description="Face Swapper (Multi-target, Male+Female Sources)")
 parser.add_argument("--out_dir", default=os.getcwd())
@@ -109,10 +84,10 @@ def swap_on_frame(frame_bgr, analysed_source_male, analysed_source_female, enhan
 
         if analysed_source_male is not None and analysed_source_female is not None:
             src_face = analysed_source_male if gender == 1 else analysed_source_female
-        elif analysed_source_male is not None and analysed_source_female is None:
+        elif analysed_source_male is not None:
             if gender == 1:
                 src_face = analysed_source_male
-        elif analysed_source_female is not None and analysed_source_male is None:
+        elif analysed_source_female is not None:
             if gender == 0:
                 src_face = analysed_source_female
 
@@ -133,7 +108,7 @@ def swap_on_frame(frame_bgr, analysed_source_male, analysed_source_female, enhan
             crop_mask=(0, 0, 0, 0),
             blur_amount=0.1,
             erode_amount=0.15,
-            blend_method="laplacian",
+            blend_method="laplacian
         )
 
     if enhancer_name != "NONE":
@@ -166,41 +141,33 @@ def swap_faces(target_files, male_file, female_file, enhancer_name="NONE"):
 
     output_files = []
 
-    for f in target_files:
-        target_path = f.name
-        ext = os.path.splitext(target_path)[-1].lower()
-
-        # -------------------- IMAGE --------------------
-        if ext in [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"]:
-            frame_bgr = cv2.imread(target_path)
-            out_frame = swap_on_frame(frame_bgr, analysed_source_male, analysed_source_female, enhancer_name)
+    # Tạo thư mục tạm và lưu file đầu ra vào đó
+    with tempfile.TemporaryDirectory() as temp_dir:
+        print(f"Thư mục tạm thời được tạo: {temp_dir}")
 
-            # Save the output image to a temporary file
-            with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as temp_file:
-                #timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-                output_path = temp_file.name
-                cv2.imwrite(output_path, out_frame)  # Convert BGR to RGB before saving
-                output_files.append(output_path)
-                #upload_to_hf(output_path)
+        for f in target_files:
+            target_path = f.name
+            ext = os.path.splitext(target_path)[-1].lower()
 
+            # -------------------- IMAGE --------------------
+            if ext in [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"]:
+                frame_bgr = cv2.imread(target_path)
+                out_frame = swap_on_frame(frame_bgr, analysed_source_male, analysed_source_female, enhancer_name)
 
+                # Tạo file tạm và ghi bức ảnh vào đó
+                with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as temp_file:
+                    output_path = temp_file.name
+                    cv2.imwrite(output_path, out_frame)  # Ghi bức ảnh vào file tạm
+                    output_files.append(output_path)
+                    print(f"File tạm được tạo tại: {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ạo thư mục tạm thời
-    with tempfile.TemporaryDirectory() as temp_dir:
-        print(f"Thư mục tạm thời được tạo: {temp_dir}")
-    
-    # Lặp qua từng đường dẫn file trong danh sách
-    for source_file in output_files:
-        # Tạo đường dẫn đầy đủ cho file đích trong thư mục tạm
-        destination_file = os.path.join(temp_dir, os.path.basename(source_file))
-        
-        # Sao chép file vào thư mục tạm
-        shutil.copy(source_file, destination_file)
-        print(f"Đã sao chép {source_file} vào {destination_file}")
-
-    # Tải lên thư mục tạm lên Hugging Face
-    upload_to_fd(temp_dir)
+        # Tải lên thư mục tạm lên Hugging Face
+        upload_to_fd(temp_dir)
 
     print(f"✔ Hoàn tất tất cả trong {time.time() - start_time:.2f}s")
     return output_files
@@ -231,4 +198,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()