Update app.py

app.py

@@ -1,279 +1,285 @@
-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("Installing 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"Error executing command: {e}")
-        return None
-
-package_name = 'nvidia-cudnn-cu12'
-version = get_pip_version(package_name)
-print(f"The installed version of {package_name} is: {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("Search results:\n", process.stdout)
-else:
-    print("An error occurred while executing the command:", 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/'
-
-commands = [
-    ['mv', source_path, destination_path],
-    ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cublas/lib/libcublas.so.12", destination_path],
-    ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cufft/lib/libcufft.so.11", destination_path],
-    ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cufft/lib/libcufftw.so.11", destination_path],
-    ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cuda_runtime/lib/libcudart.so.12", destination_path],
-    ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cuda_cupti/lib/libcupti.so.12", destination_path],
-    ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/curand/lib/libcurand.so.10", destination_path],
-    ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/cusolver/lib/libcusolver.so.11", destination_path],
-    ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/cusolver/lib/libcusolverMg.so.11", destination_path],
-    ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/cusparse/lib/libcusparse.so.12", destination_path],
-]
-
-for command in commands:
-    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("Search results:\n", process.stdout)
-else:
-    print("An error occurred while executing the command:", 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"Downloading {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} downloaded successfully.")
-            else:
-                print(f"Failed to download {filename}. Status code: {response.status_code}")
-        else:
-            print(f"{filename} already exists. Skipping 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"Error removing temporary directory: {e}")
-
-    def extract_frames(video_path: str):
-        video_capture = cv2.VideoCapture(video_path)
-        if not video_capture.isOpened():
-            print("Error opening video.")
-            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("Error opening video.")
-        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"Error removing file: {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)
-    remove_temp_directory(temp_dir)
-    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()],
-    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"
-)
-
+
+import os
+
+# Set the LD_LIBRARY_PATH
+os.environ['LD_LIBRARY_PATH'] = '/usr/local/lib/python3.10/site-packages/nvidia/cudnn/lib:' + os.environ.get('LD_LIBRARY_PATH', '')
+
+# Now you can import your libraries that depend on this path
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
+import requests
+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("Installing 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"Error executing command: {e}")
+        return None
+
+package_name = 'nvidia-cudnn-cu12'
+version = get_pip_version(package_name)
+print(f"The installed version of {package_name} is: {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("Search results:\n", process.stdout)
+else:
+    print("An error occurred while executing the command:", 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/'
+
+commands = [
+    ['mv', source_path, destination_path],
+    ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cublas/lib/libcublas.so.12", destination_path],
+    ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cufft/lib/libcufft.so.11", destination_path],
+    ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cufft/lib/libcufftw.so.11", destination_path],
+    ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cuda_runtime/lib/libcudart.so.12", destination_path],
+    ['mv', "/usr/local/lib/python3.10/site-packages/nvidia/cuda_cupti/lib/libcupti.so.12", destination_path],
+    ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/curand/lib/libcurand.so.10", destination_path],
+    ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/cusolver/lib/libcusolver.so.11", destination_path],
+    ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/cusolver/lib/libcusolverMg.so.11", destination_path],
+    ['cp', "/usr/local/lib/python3.10/site-packages/nvidia/cusparse/lib/libcusparse.so.12", destination_path],
+]
+
+for command in commands:
+    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("Search results:\n", process.stdout)
+else:
+    print("An error occurred while executing the command:", 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"Downloading {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} downloaded successfully.")
+            else:
+                print(f"Failed to download {filename}. Status code: {response.status_code}")
+        else:
+            print(f"{filename} already exists. Skipping 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"Error removing temporary directory: {e}")
+
+    def extract_frames(video_path: str):
+        video_capture = cv2.VideoCapture(video_path)
+        if not video_capture.isOpened():
+            print("Error opening video.")
+            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("Error opening video.")
+        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"Error removing file: {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)
+    remove_temp_directory(temp_dir)
+    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()],
+    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()