Update app.py

app.py

@@ -1,107 +1,91 @@
-#!/usr/bin/env python
-
-from setuptools import find_packages, setup
-
+# -*- 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,
+)
+from roop.processors.frame.core import get_frame_processors_modules
+from roop.utilities import normalize_output_path
 import os
-import subprocess
-import time
-
-version_file = 'realesrgan/version.py'
-
-
-def readme():
-    with open('README.md', encoding='utf-8') as f:
-        content = f.read()
-    return content
-
-
-def get_git_hash():
-
-    def _minimal_ext_cmd(cmd):
-        # construct minimal environment
-        env = {}
-        for k in ['SYSTEMROOT', 'PATH', 'HOME']:
-            v = os.environ.get(k)
-            if v is not None:
-                env[k] = v
-        # LANGUAGE is used on win32
-        env['LANGUAGE'] = 'C'
-        env['LANG'] = 'C'
-        env['LC_ALL'] = 'C'
-        out = subprocess.Popen(cmd, stdout=subprocess.PIPE, env=env).communicate()[0]
-        return out
-
-    try:
-        out = _minimal_ext_cmd(['git', 'rev-parse', 'HEAD'])
-        sha = out.strip().decode('ascii')
-    except OSError:
-        sha = 'unknown'
-
-    return sha
-
-
-def get_hash():
-    if os.path.exists('.git'):
-        sha = get_git_hash()[:7]
-    else:
-        sha = 'unknown'
-
-    return sha
-
-
-def write_version_py():
-    content = """# GENERATED VERSION FILE
-# TIME: {}
-__version__ = '{}'
-__gitsha__ = '{}'
-version_info = ({})
-"""
-    sha = get_hash()
-    with open('VERSION', 'r') as f:
-        SHORT_VERSION = f.read().strip()
-    VERSION_INFO = ', '.join([x if x.isdigit() else f'"{x}"' for x in SHORT_VERSION.split('.')])
-
-    version_file_str = content.format(time.asctime(), SHORT_VERSION, sha, VERSION_INFO)
-    with open(version_file, 'w') as f:
-        f.write(version_file_str)
-
-
-def get_version():
-    with open(version_file, 'r') as f:
-        exec(compile(f.read(), version_file, 'exec'))
-    return locals()['__version__']
-
-
-def get_requirements(filename='requirements.txt'):
-    here = os.path.dirname(os.path.realpath(__file__))
-    with open(os.path.join(here, filename), 'r') as f:
-        requires = [line.replace('\n', '') for line in f.readlines()]
-    return requires
-
-
-if __name__ == '__main__':
-    write_version_py()
-    setup(
-        name='realesrgan',
-        version=get_version(),
-        description='Real-ESRGAN aims at developing Practical Algorithms for General Image Restoration',
-        long_description=readme(),
-        long_description_content_type='text/markdown',
-        author='Xintao Wang',
-        author_email='xintao.wang@outlook.com',
-        keywords='computer vision, pytorch, image restoration, super-resolution, esrgan, real-esrgan',
-        url='https://github.com/xinntao/Real-ESRGAN',
-        include_package_data=True,
-        packages=find_packages(exclude=('options', 'datasets', 'experiments', 'results', 'tb_logger', 'wandb')),
-        classifiers=[
-            'Development Status :: 4 - Beta',
-            'License :: OSI Approved :: Apache Software License',
-            'Operating System :: OS Independent',
-            'Programming Language :: Python :: 3',
-            'Programming Language :: Python :: 3.7',
-            'Programming Language :: Python :: 3.8',
-        ],
-        license='BSD-3-Clause License',
-        setup_requires=['cython', 'numpy'],
-        install_requires=get_requirements(),
-        zip_safe=False)
+import random
+from PIL import Image
+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_dir = "temp"  # Sử dụng thư mục cố định
+    os.makedirs(session_dir, exist_ok=True)
+
+    # Tạo tên file ngẫu nhiên
+    source_filename = f"source_{random.randint(1000, 9999)}.jpg"
+    target_filename = f"target_{random.randint(1000, 9999)}.jpg"
+    output_filename = f"output_{random.randint(1000, 9999)}.jpg"
+
+    source_path = os.path.join(session_dir, source_filename)
+    target_path = os.path.join(session_dir, target_filename)
+
+    source_image = Image.fromarray(source_file)
+    source_image.save(source_path)
+    target_image = Image.fromarray(target_file)
+    target_image.save(target_path)
+
+    print("source_path: ", source_path)
+    print("target_path: ", target_path)
+
+    # 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.")
+
+    # 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.")
+
+    output_path = os.path.join(session_dir, output_filename)
+    normalized_output_path = normalize_output_path(source_path, target_path, output_path)
+
+    frame_processors = ["face_swapper", "face_enhancer"] if doFaceEnhancer else ["face_swapper"]
+
+    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 = 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.execution_providers = decode_execution_providers(['cpu'])
+    roop.globals.reference_face_position = 0
+    roop.globals.similar_face_distance = 0.6
+    roop.globals.max_memory = 60
+    roop.globals.execution_threads = 8
+    
+    start()
+    return normalized_output_path
+
+app = gr.Interface(
+    fn=swap_face, 
+    inputs=[
+        gr.Image(), 
+        gr.Image(), 
+        gr.Checkbox(label="Face Enhancer?", info="Do face enhancement?")
+    ], 
+    outputs="image"
+)
+app.launch()