app.py

import onnxruntime as ort
from huggingface_hub import hf_hub_download
import requests
import os
import gradio as gr
import spaces
from typing import Any, List, Callable
import cv2
import insightface
import time
import tempfile
import subprocess
import gfpgan

print("instaling cudnn 9")
import subprocess
import sys

def get_pip_version(package_name):
    try:
        result = subprocess.run(
            [sys.executable, '-m', 'pip', 'show', package_name],
            capture_output=True,
            text=True,
            check=True
        )
        output = result.stdout
        version_line = next(line for line in output.split('\n') if line.startswith('Version:'))
        return version_line.split(': ')[1]
    except subprocess.CalledProcessError as e:
        print(f"Erro ao executar o comando: {e}")
        return None

package_name = 'nvidia-cudnn-cu12' 
version = get_pip_version(package_name)
print(f"A versão instalada de {package_name} é: {version}")

command = "find / -path /proc -prune -o -path /sys -prune -o -name 'libcudnn*' -print"

process = subprocess.run(command, shell=True, text=True, capture_output=True)

if process.returncode == 0:
    print("Resultados da busca:\n", process.stdout)
else:
    print("Houve um erro na execução do comando:", process.stderr)

source_path = '/usr/local/lib/python3.10/site-packages/nvidia/cublas/lib/libcublasLt.so.12'
destination_path = '/usr/local/lib/python3.10/site-packages/nvidia/cudnn/lib/'

command = ['mv', source_path, destination_path]
subprocess.run(command, check=True)

command = ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cublas/lib/libcublas.so.12", destination_path]
subprocess.run(command, check=True)

command = ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cufft/lib/libcufft.so.11", destination_path]
subprocess.run(command, check=True)

command = ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cufft/lib/libcufftw.so.11", destination_path]
subprocess.run(command, check=True)

command = ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cuda_runtime/lib/libcudart.so.12", destination_path]
subprocess.run(command, check=True)

command = ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cuda_cupti/lib/libcupti.so.12", destination_path]
subprocess.run(command, check=True)

command = ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/curand/lib/libcurand.so.10", destination_path]
subprocess.run(command, check=True)

command = ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/cusolver/lib/libcusolver.so.11", destination_path]
subprocess.run(command, check=True)

command = ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/cusolver/lib/libcusolverMg.so.11", destination_path]
subprocess.run(command, check=True)

command = ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/cusparse/lib/libcusparse.so.12", destination_path]
subprocess.run(command, check=True)

command = "find / -path /proc -prune -o -path /sys -prune -o -name 'libcu*' -print"
process = subprocess.run(command, shell=True, text=True, capture_output=True)

if process.returncode == 0:
    print("Resultados da busca:\n", process.stdout)
else:
    print("Houve um erro na execução do comando:", process.stderr)


print("done")
print("---------------------")

print(ort.get_available_providers())

def conditional_download(download_directory_path, urls):
    if not os.path.exists(download_directory_path):
        os.makedirs(download_directory_path)

    for url in urls:
        filename = url.split('/')[-1]
        file_path = os.path.join(download_directory_path, filename)

        if not os.path.exists(file_path):
            print(f"Baixando {filename}...")

            response = requests.get(url, stream=True)
            if response.status_code == 200:
                with open(file_path, 'wb') as file:
                    for chunk in response.iter_content(chunk_size=8192):
                        file.write(chunk)
                print(f"{filename} baixado com sucesso.")
            else:
                print(f"Falha ao baixar {filename}. Status code: {response.status_code}")
        else:
            print(f"{filename} já existe. Pulando o download.")

model_path = hf_hub_download(repo_id="countfloyd/deepfake", filename="inswapper_128.onnx")
conditional_download("./", ['https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth'])

FACE_SWAPPER = None
FACE_ANALYSER = None
FACE_ENHANCER = None

@spaces.GPU(duration=300, enable_queue=True)
def process_video(source_path: str, target_path: str, enhance = False, progress=gr.Progress(), output_path='result.mp4') -> None:
    def get_face_analyser():
        global FACE_ANALYSER

        if FACE_ANALYSER is None:
            FACE_ANALYSER = insightface.app.FaceAnalysis(name='buffalo_l', providers=["CUDAExecutionProvider"])
            FACE_ANALYSER.prepare(ctx_id=0, det_size=(640, 640))

        return FACE_ANALYSER

    def get_face_enhancer() -> Any:
        global FACE_ENHANCER

        if FACE_ENHANCER is None:
            FACE_ENHANCER = gfpgan.GFPGANer(model_path="./GFPGANv1.4.pth", upscale=2, ) # type: ignore[attr-defined]

        return FACE_ENHANCER

    def get_one_face(frame):
        face = get_face_analyser().get(frame)
        try:
            return min(face, key=lambda x: x.bbox[0])
        except ValueError:
            return None

    def get_face_swapper():
        global FACE_SWAPPER

        if FACE_SWAPPER is None:
            FACE_SWAPPER = insightface.model_zoo.get_model(model_path, providers=["CUDAExecutionProvider"])

        return FACE_SWAPPER



    def swap_face(source_face, target_face, temp_frame):
        return get_face_swapper().get(temp_frame, target_face, source_face, paste_back=True)


    def process_frame(source_face, temp_frame, enhance):
        target_face = get_one_face(temp_frame)
        if target_face:
            temp_frame = swap_face(source_face, target_face, temp_frame)
            if enhance:
                temp_frame = enhance_face(temp_frame)
        return temp_frame


    def process_image(source_path: str, target_path: str, output_path: str) -> None:
        source_face = get_one_face(cv2.imread(source_path))
        target_frame = cv2.imread(target_path)
        result = process_frame(source_face, target_frame)
        cv2.imwrite(output_path, result)

    def create_temp_directory():
        temp_dir = tempfile.mkdtemp()
        return temp_dir

    def enhance_face(temp_frame):
        _, _, temp_frame = get_face_enhancer().enhance(
            temp_frame,
            paste_back=True
        )
        return temp_frame


    def remove_temp_directory(temp_dir):
        try:
            for filename in os.listdir(temp_dir):
                file_path = os.path.join(temp_dir, filename)
                if os.path.isfile(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    os.rmdir(file_path)
            os.rmdir(temp_dir)
        except Exception as e:
            print(f"Erro ao remover a pasta temporária: {e}")

    def extract_frames(video_path: str):
        video_capture = cv2.VideoCapture(video_path)
        if not video_capture.isOpened():
            print("Erro ao abrir o vídeo.")
            return []

        frames = []
        
        while True:
            ret, frame = video_capture.read()
            if not ret:
                break
            
            frames.append(frame)
        
        video_capture.release()
        
        return frames


    def get_video_fps(video_path: str) -> float:
        video_capture = cv2.VideoCapture(video_path)
        if not video_capture.isOpened():
            raise ValueError("Erro ao abrir o vídeo.")
        fps = video_capture.get(cv2.CAP_PROP_FPS)
        video_capture.release()
        return fps

    def create_video_from_frames(temp_dir: str, output_video_path: str, fps: float) -> None:
        temp_frames_pattern = os.path.join(temp_dir, "frame_%04d.jpg")
        
        ffmpeg_command = [
            'ffmpeg',
            '-y',
            '-framerate', str(fps),
            '-i', temp_frames_pattern,
            '-c:v', 'libx264',
            '-pix_fmt', 'yuv420p',
            '-preset', 'ultrafast',
            output_video_path
        ]
        
        subprocess.run(ffmpeg_command, check=True)


    def extract_audio(video_path: str, audio_path: str) -> None:
        ffmpeg_command = [
            'ffmpeg',
            '-y',
            '-i', video_path,
            '-q:a', '0',
            '-map', 'a',
            '-preset', 'ultrafast',
            audio_path
        ]
        
        subprocess.run(ffmpeg_command, check=True)

    def add_audio_to_video(video_path: str, audio_path: str, output_video_path: str) -> None:
        ffmpeg_command = [
            'ffmpeg',
            '-y',
            '-i', video_path,
            '-i', audio_path,
            '-c:v', 'copy',
            '-c:a', 'aac',
            '-strict', 'experimental',
            '-preset', 'ultrafast',
            output_video_path
        ]
        
        subprocess.run(ffmpeg_command, check=True)


    def delete_file(file_path: str) -> None:
        try:
            os.remove(file_path)
        except Exception as e:
            print(f"Erro ao remover o arquivo: {e}")

    def reduce_video(video_path: str, output_video_path: str) -> None:
        ffmpeg_command = [
            'ffmpeg',
                '-y',
            '-i', video_path,
            '-vf', "scale='if(gte(iw,ih),720,-1)':'if(gte(iw,ih),-1,720)',pad=ceil(iw/2)*2:ceil(ih/2)*2",
                '-preset', 'ultrafast',
            '-r','24',
            output_video_path
        ]
        subprocess.run(ffmpeg_command, check=True)
        
    temp_dir = create_temp_directory()
    video_input = temp_dir + "/input.mp4"
    reduce_video(target_path , video_input)

    
    source_face = get_one_face(cv2.imread(source_path))
    frames = extract_frames(video_input)     
    
    for index, frame in progress.tqdm(enumerate(frames), total=len(frames)):
        processed_frame = process_frame(source_face, frame, enhance)
        frame_filename = os.path.join(temp_dir, f"frame_{index:04d}.jpg")
        cv2.imwrite(frame_filename, processed_frame)

    video_path = temp_dir + "/out.mp4"
    create_video_from_frames(temp_dir, video_path, get_video_fps(video_input))
    audio_path = temp_dir + "/audio.wav"
    extract_audio(video_input, audio_path)

    add_audio_to_video(video_path, audio_path, output_path)

    return output_path

app = gr.Interface(
    fn=process_video, 
    inputs=[gr.Image(type='filepath'), gr.Video(), gr.Checkbox(label="Use Face GAN(increase render time)", value=False)], 
    outputs=[gr.Video()],
   )

app.launch(share='True')