Update app.py

app.py

@@ -1,24 +1,31 @@
-# -* coding:UTF-8 -*
+# -*- coding:UTF-8 -*-
 # !/usr/bin/env python
+import spaces
 import numpy as np
 import gradio as gr
+import gradio.exceptions
 import roop.globals
 from roop.core import (
     start,
     decode_execution_providers,
-    suggest_max_memory,
-    suggest_execution_threads,
 )
 from roop.processors.frame.core import get_frame_processors_modules
 from roop.utilities import normalize_output_path
 import os
 from PIL import Image
+import uuid
+import onnxruntime as ort
+import cv2
+from roop.face_analyser import get_one_face
 
+@spaces.GPU
+def swap_face(source_file, target_file, doFaceEnhancer):
+    session_id = str(uuid.uuid4())  # Tạo một UUID duy nhất cho mỗi phiên làm việc
+    session_dir = f"temp/{session_id}"
+    os.makedirs(session_dir, exist_ok=True)
 
-def swap_face(source_file, target_file,doFaceEnhancer):
-
-    source_path = "input.jpg"
-    target_path = "target.jpg"
+    source_path = os.path.join(session_dir, "input.jpg")
+    target_path = os.path.join(session_dir, "target.jpg")
 
     source_image = Image.fromarray(source_file)
     source_image.save(source_path)
@@ -28,68 +35,68 @@ def swap_face(source_file, target_file,doFaceEnhancer):
     print("source_path: ", source_path)
     print("target_path: ", target_path)
 
-    roop.globals.source_path = source_path
-    roop.globals.target_path = target_path
-    output_path = "output.jpg"
-    roop.globals.output_path = normalize_output_path(
-        roop.globals.source_path, roop.globals.target_path, output_path
-    )
-    if doFaceEnhancer == True:
-        roop.globals.frame_processors = ["face_swapper","face_enhancer"]
-    else:
-        roop.globals.frame_processors = ["face_swapper"]
-    roop.globals.headless = True
-    roop.globals.keep_fps = True
-    roop.globals.keep_audio = True
-    roop.globals.keep_frames = False
-    roop.globals.many_faces = False
-    roop.globals.video_encoder = "libx264"
-    roop.globals.video_quality = 18
-    roop.globals.max_memory = suggest_max_memory()
-    roop.globals.execution_providers = decode_execution_providers(["cuda"])
-    roop.globals.execution_threads = suggest_execution_threads()
+    # Check if a face is detected in the source image
+    source_face = get_one_face(cv2.imread(source_path))
+    if source_face is None:
+        raise gradio.exceptions.Error("No face in source path detected.")
 
-    print(
-        "start process",
-        roop.globals.source_path,
-        roop.globals.target_path,
-        roop.globals.output_path,
-    )
+    # Check if a face is detected in the target image
+    target_face = get_one_face(cv2.imread(target_path))
+    if target_face is None:
+        raise gradio.exceptions.Error("No face in target path detected.")
 
-    for frame_processor in get_frame_processors_modules(
-        roop.globals.frame_processors
-    ):
-        if not frame_processor.pre_check():
-            return
+    output_path = os.path.join(session_dir, "output.jpg")
+    normalized_output_path = normalize_output_path(source_path, target_path, output_path)
 
-    start()
-    return output_path
+    frame_processors = ["face_swapper", "face_enhancer"] if doFaceEnhancer else ["face_swapper"]
+    headless = True
+    keep_fps = True
+    keep_audio = True
+    keep_frames = False
+    many_faces = False
+    video_encoder = "libx264"
+    video_quality = 18
+    max_memory = "12G"
+    execution_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]  # Ưu tiên GPU, sau đó là CPU
+    execution_threads = 4
+    reference_face_position = 0
+    similar_face_distance = 0.6
 
-# HTML code for the ad banner
-ad_code = """
-<div style="text-align: center;">
-<iframe src="https://www.fiverr.com/gig_widgets?id=U2FsdGVkX1+9vQaosK0mq7xGXzIb/5hdFYIQsmHlpeMwFm5gingveDZKXBz2oJwV7ZwvToxseOx/KzKtr77ryB3qTb83/+GvCyI8OKpZcSgM0sDHzEPlZnqNeu8Db4y/IXlAceA1UcYT5Z3cW5vBLrgJ/WOpyMWSiyfwwMmGngjltQBexcivW7ukRd+/0yos7GCbfGjGdnLPHm2LB/CB9w2mA1i8cVv9LU+UUdj/O5KLRXtnychl9wAMkPTpmwMiyuAfUZy0nbK/Xa5O2UecmCvn7wicTr0TCTpLmIxaEh3YftXOVr9e36OGVYTmy5nlsgujkqJPI7wL0dXfon0ru4kvoycC7UD6m/whMoxrWZOf386qV2eEeXzk3vKO/emXUUfSavtofzSBtjLcyAOVhnHJg27PvVPgvFTQdy0o3F1M0DaHYTVW3Ln45MPYvMSH&affiliate_id=36184&strip_google_tagmanager=true" loading="lazy" data-with-title="false" class="fiverr_nga_frame" frameborder="0" height="350" width="100%" referrerpolicy="no-referrer-when-downgrade" data-mode="random_gigs" onload=" var frame = this; var script = document.createElement('script'); script.addEventListener('load', function() { window.FW_SDK.register(frame); }); script.setAttribute('src', 'https://www.fiverr.com/gig_widgets/sdk'); document.body.appendChild(script); " ></iframe>
-</div>
-	"""
+    print("Available providers:", ort.get_available_providers())  # Kiểm tra các provider hiện có
+    print("Configured execution providers:", execution_providers)
 
-# HTML code for the ad banner
-ad_code2 = """
-<div style="text-align: center;">
-    <a href="https://beta.publishers.adsterra.com/referral/UNXJYTziBP" target="_blank" style="display: inline-block;">
-        <img decoding="async" alt="banner" src="https://landings-cdn.adsterratech.com/referralBanners/gif/468x120_adsterra_reff.gif">
-    </a>
-    <a href="https://go.fiverr.com/visit/?bta=36184&brand=fiverrcpa&landingPage=https%253A%252F%252Fwww.fiverr.com%252Fcategories%252Fprogramming-tech%252Fai-coding%252Fai-applications%253Fsource%253Dcategory_tree" target="_blank" style="display: inline-block;">
-        <img fetchpriority="high" decoding="async" width="468" height="120" src="https://ziverr.xyz/wp-content/uploads/2024/06/PASSIVE-1.gif" class="attachment-large size-large wp-image-1266" alt="">
-    </a>
-</div>
-	"""
+    for frame_processor in get_frame_processors_modules(frame_processors):
+        if not frame_processor.pre_check():
+            print(f"Pre-check failed for {frame_processor}")
+            raise gradio.exceptions.Error(f"Pre-check failed for {frame_processor}")
 
+    roop.globals.source_path = source_path
+    roop.globals.target_path = target_path
+    roop.globals.output_path = normalized_output_path
+    roop.globals.frame_processors = frame_processors
+    roop.globals.headless = headless
+    roop.globals.keep_fps = keep_fps
+    roop.globals.keep_audio = keep_audio
+    roop.globals.keep_frames = keep_frames
+    roop.globals.many_faces = many_faces
+    roop.globals.video_encoder = video_encoder
+    roop.globals.video_quality = video_quality
+    roop.globals.max_memory = max_memory
+    roop.globals.execution_providers = execution_providers
+    roop.globals.execution_threads = execution_threads
+    roop.globals.reference_face_position = reference_face_position
+    roop.globals.similar_face_distance = similar_face_distance
+
+    start()
+    return normalized_output_path
 
 app = gr.Interface(
-    fn=swap_face,
-    # ???=gr.HTML(ad_code2),  # Remove this line or replace '???'
-    inputs=[gr.Image(), gr.Image(), gr.Checkbox(label="face_enhancer?", info="do face enhancer?")],
-    outputs="image",
-    # Directly pass the HTML code as a string
-    description=ad_code2 
+    fn=swap_face, 
+    inputs=[
+        gr.Image(), 
+        gr.Image(), 
+        gr.Checkbox(label="Face Enhancer?", info="Do face enhancement?")
+    ], 
+    outputs="image"
 )
+app.launch()