Update app.py

app.py

@@ -1,3 +1,301 @@
+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
 
-exec(os.environ.get('APP'))
+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, enhance = False) -> None:
+        source_face = get_one_face(cv2.imread(source_path))
+        target_frame = cv2.imread(target_path)
+        result = process_frame(source_face, target_frame, enhance)
+        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)
+    
+    if not target_path.endswith('.mp4') and not target_path.endswith('.mov') and not target_path.endswith('.avi'):
+        process_image(source_path, target_path, "./output.jpg", enhance)
+        return "./output.jpg"
+
+    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)
+    remove_temp_directory(temp_dir)
+    return output_path
+
+app = gr.Interface(
+    fn=process_video, 
+    inputs=[gr.Image(type='filepath'), gr.File(label="Upload Image or Video", file_types=["image", "video"]), gr.Checkbox(label="Use Face GAN(increase render time)", value=False)], 
+    outputs=[gr.File()],
+    description="Videos get downsampled to 720p and 24fps. To modify the code or purchase a modification, send an email to [email protected] to donate to the owner of the space: https://donate.stripe.com/3csg0D0tadXU4mYcMM"
+)
+
+app.launch()