Update app.py

app.py

@@ -17,118 +17,17 @@ from datetime import datetime
 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
-    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)
-    
-    return downloaded_files
+
 # --- Upload function ---
-def upload_to_hf(filepath):
-    #filename = os.path.basename(filepath)
+def upload_to_hf(folder_path):
     SUBFOLDER = datetime.now().strftime("%Y%m%d")
-    path_in_repo = f"{SUBFOLDER}"
     api.upload_folder(
-        folder_path=filepath,
-        path_in_repo=path_in_repo,
+        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/{path_in_repo}"
-# ------------------------------ ARGS ------------------------------
-parser = argparse.ArgumentParser(description="Face Swapper (Multi-target, Male+Female Sources)")
-parser.add_argument("--out_dir", default=os.getcwd())
-parser.add_argument("--batch_size", default=32)
-parser.add_argument("--cuda", action="store_true", default=False)
-user_args = parser.parse_args()
-
-USE_CUDA = user_args.cuda
-DEF_OUTPUT_PATH = user_args.out_dir
-BATCH_SIZE = int(user_args.batch_size)
-
-# ------------------------------ DEVICE ------------------------------
-PROVIDER = ["CPUExecutionProvider"]
-if USE_CUDA and "CUDAExecutionProvider" in onnxruntime.get_available_providers():
-    PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"]
-    print(">>> Running on CUDA")
-else:
-    USE_CUDA = False
-    print(">>> Running on CPU")
-
-device = "cuda" if USE_CUDA else "cpu"
-EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None
-
-# ------------------------------ MODELS ------------------------------
-FACE_ANALYSER = insightface.app.FaceAnalysis(name="buffalo_l", providers=PROVIDER)
-FACE_ANALYSER.prepare(ctx_id=0, det_size=(640, 640), det_thresh=0.6)
-
-FACE_SWAPPER = Inswapper(
-    model_file="./assets/pretrained_models/inswapper_128.onnx",
-    batch_size=(BATCH_SIZE if USE_CUDA else 1),
-    providers=PROVIDER,
-)
-
-# ------------------------------ ENHANCERS ------------------------------
-ENHANCER_CHOICES = ["NONE"] + get_available_enhancer_names()
-
-# ------------------------------ CORE SWAP FUNC ------------------------------
-def swap_on_frame(frame_bgr, analysed_source_male, analysed_source_female, enhancer_name="NONE"):
-    analysed_faces = FACE_ANALYSER.get(frame_bgr)
-    preds, matrs = [], []
-
-    for analysed_face in analysed_faces:
-        gender = analysed_face.get("gender", 1)  # 1 = male, 0 = female
-        src_face = None
-
-        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:
-            if gender == 1:
-                src_face = analysed_source_male
-        elif analysed_source_female is not None and analysed_source_male is None:
-            if gender == 0:
-                src_face = analysed_source_female
-
-        if src_face is None:
-            continue
-
-        batch_pred, batch_matr = FACE_SWAPPER.get([frame_bgr], [analysed_face], [src_face])
-        preds.extend(batch_pred)
-        matrs.extend(batch_matr)
-        EMPTY_CACHE()
-
-    for p, m in zip(preds, matrs):
-        frame_bgr = paste_to_whole(
-            foreground=p,
-            background=frame_bgr,
-            matrix=m,
-            mask=None,
-            crop_mask=(0, 0, 0, 0),
-            blur_amount=0.1,
-            erode_amount=0.15,
-            blend_method="laplacian",
-        )
-
-    if enhancer_name != "NONE":
-        try:
-            model, runner = load_face_enhancer_model(name=enhancer_name, device=device)
-            frame_bgr = runner(frame_bgr, model)
-        except Exception as e:
-            print(f"[Enhancer] Error while running {enhancer_name}: {e}")
-
-    return frame_bgr
+    return f"https://huggingface.co/datasets/{REPO_ID}/blob/main/{SUBFOLDER}"
 
 # ------------------------------ PROCESS ------------------------------
 def swap_faces(target_files, male_file, female_file, enhancer_name="NONE"):
@@ -148,40 +47,37 @@ def swap_faces(target_files, male_file, female_file, enhancer_name="NONE"):
 
     if analysed_source_male is None and analysed_source_female is None:
         raise ValueError("❌ Cần ít nhất 1 khuôn mặt nguồn (Nam hoặc Nữ).")
-# Tạo thư mục tạm để lưu ảnh đầu ra
+
+    # Tạo thư mục tạm để lưu ảnh đầu ra
     with tempfile.TemporaryDirectory() as temp_dir:
+        print(f"Temporary directory created at: {temp_dir}")
         output_files = []
 
         for f in target_files:
             target_path = f.name
             ext = os.path.splitext(target_path)[-1].lower()
-    #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)
-
-            # 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)
-                # Lưu ảnh đầu ra vào thư mục tạm
-                output_path = os.path.join(temp_dir,temp_file.name )
-                cv2.imwrite(output_path, out_frame)  # Lưu ảnh vào thư mục tạm
+
+            # -------------------- 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)
+
+                # Save the output image to a temporary file
+                output_path = os.path.join(temp_dir, f"output{len(output_files)}{ext}")
+                cv2.imwrite(output_path, out_frame)  # Save image to temp directory
                 output_files.append(output_path)
 
+        # Upload the temporary directory to Hugging Face
+        if os.path.isdir(temp_dir):
+            upload_to_hf(temp_dir)
+        else:
+            print(f"Error: {temp_dir} is not a valid directory.")
 
-    upload_to_hf(temp_dir)
     print(f"✔ Hoàn tất tất cả trong {time.time() - start_time:.2f}s")
     return output_files
 
+
+
 # ------------------------------ UI ------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("## 🧑‍🦱➡👩 Face Swapper (Upload nhiều file target + nguồn nam/nữ) + Enhancer")