Deep-Live-Cam/google-colab/DeepLive_Google_Colab.ipynb

1503 lines
49 KiB
Plaintext
Raw Normal View History

2024-08-16 11:47:09 +08:00
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "U3Fc3OXuDeqP"
},
"source": [
"# **Install InsightFace and Dependencies:**\n",
"\n",
" Run the following code to install the required packages.\n",
"\n",
" Note: This installation varies and depends on your cuda and cudnn version. Incase of face swapping issue refer to the link in the debug session to know the version to install."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true,
"id": "D-UgHjSFBE9d"
},
"outputs": [],
"source": [
"#!pip install onnx==1.16.0\n",
"!pip install onnxruntime-gpu==1.18.0 --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/\n",
"!pip install insightface #==0.7.3\n",
"#!pip install onnxruntime==1.18.0\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Q81Sf-CFarjq"
},
"outputs": [],
"source": [
"!apt-get install -y ffmpeg"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GYFbuuCdvTaO"
},
"source": [
"### DeepFakeLive\n",
"Follow the steps below to create the necessery directories for the project"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"id": "_MDoAtpyvV_l"
},
"outputs": [],
"source": [
"!cd /content\n",
"!rm -rf *"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"id": "RzPl9sPVgSoP"
},
"outputs": [],
"source": [
"!mkdir deepfakecollab"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "lJJbUDokN0bF",
"outputId": "a2ef8b6e-97a8-4ef2-f986-e34e6eb1252e"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"/content/deepfakecollab\n"
]
}
],
"source": [
"cd deepfakecollab"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"id": "ZUzAF6L2N6__"
},
"outputs": [],
"source": [
"!mkdir Scripts"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ziSLHxab_j0h"
},
"outputs": [],
"source": [
"!ffmpeg -version"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0rTCd2R1uw7A"
},
"source": [
"### [Optional] FRP\n",
"\n",
"Follow the steps below to setup FRP. You will also need to host the FRPS on your VPS (free).\n",
"\n",
"Note: You cannot run both FRP and ngrok together. Its ether FRP and ngrok"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ckp_dM52N9K3"
},
"outputs": [],
"source": [
"!touch Scripts/get_frs.sh\n",
"\n",
"getfrs = \"\"\"#!/usr/bin/env bash\n",
"\n",
"# Check if frpc is installed\n",
"command -v frpc >/dev/null 2>&1\n",
"if [[ $? -ne 0 ]]; then\n",
" echo \"frpc is not found, installing...\"\n",
" wget -q -nc https://github.com/fatedier/frp/releases/download/v0.59.0/frp_0.59.0_linux_amd64.tar.gz\n",
" tar -xzf frp_0.59.0_linux_amd64.tar.gz\n",
" echo \"Done!\"\n",
"fi\"\"\"\n",
"with open('Scripts/get_frs.sh', 'w') as f:\n",
" f.write(getfrs)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "w21kxVRmODFZ"
},
"outputs": [],
"source": [
"!touch Scripts/open_tunnel_frs.sh\n",
"\n",
"getfrs = \"\"\"#!/usr/bin/env bash\n",
"\n",
"cmd=\"frp_0.59.0_linux_amd64/frpc -c frp_0.59.0_linux_amd64/frpc.toml\"\n",
"\n",
"kill -9 $(ps aux | grep $cmd | awk '{print $2}') 2> /dev/null\n",
"\n",
"echo Opening tunnel\n",
"$cmd\"\"\"\n",
"with open('Scripts/open_tunnel_frs.sh', 'w') as f:\n",
" f.write(getfrs)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tmOY654Mia79"
},
"outputs": [],
"source": [
"!chmod +x Scripts/get_frs.sh\n",
"!chmod +x Scripts/open_tunnel_frs.sh"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Br-P5pkOwFg1",
"outputId": "663dae11-27c6-4329-a768-58faa1dd14e9"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"frpc is not found, installing...\n",
"Done!\n"
]
}
],
"source": [
"!Scripts/get_frs.sh"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "9n-qESinvM22"
},
"source": [
"### [Optional] Ngrok\n",
"\n",
"Follow to Setup Ngrok.\n",
"\n",
"Note: You need an API key from Ngrok to use tcp for free you would need to add a billing details to their platform"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"id": "S9U6x85DvTTV"
},
"outputs": [],
"source": [
"!touch Scripts/get_ngrok.sh\n",
"\n",
"getfrs = \"\"\"#!/usr/bin/env bash\n",
"\n",
"# Check if frpc is installed\n",
"command -v frpc >/dev/null 2>&1\n",
"if [[ $? -ne 0 ]]; then\n",
" echo \"ngrok is not found, installing...\"\n",
" wget -q -nc https://bin.equinox.io/c/bNyj1mQVY4c/ngrok-v3-stable-linux-amd64.tgz\n",
" tar -xzf ngrok-v3-stable-linux-amd64.tgz\n",
" echo \"Done!\"\n",
"fi\"\"\"\n",
"with open('Scripts/get_ngrok.sh', 'w') as f:\n",
" f.write(getfrs)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"id": "ZYKKKRZSwA82"
},
"outputs": [],
"source": [
"!touch Scripts/open_tunnel_ngrok.sh\n",
"\n",
"getfrs = \"\"\"#!/usr/bin/env bash\n",
"\n",
"cmd=\"./ngrok start --all --config ngrok.conf\"\n",
"\n",
"kill -9 $(ps aux | grep $cmd | awk '{print $2}') 2> /dev/null\n",
"\n",
"echo Opening tunnel\n",
"$cmd\"\"\"\n",
"with open('Scripts/open_tunnel_ngrok.sh', 'w') as f:\n",
" f.write(getfrs)"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {
"id": "A2FQrmDQwW__"
},
"outputs": [],
"source": [
"!chmod +x Scripts/get_ngrok.sh\n",
"!chmod +x Scripts/open_tunnel_ngrok.sh"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "WoNs1gSiweBm",
"outputId": "9e24a0a6-b50b-49f6-a928-2c004917927e"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"ngrok is not found, installing...\n",
"Done!\n"
]
}
],
"source": [
"!Scripts/get_ngrok.sh"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"id": "uCPALISZwo3K"
},
"outputs": [],
"source": [
"# Paste your authtoken here in quotes\n",
"authtoken = \"\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uVpDmuTmwtO9"
},
"source": [
"Set your region\n",
"\n",
"Code | Region\n",
"--- | ---\n",
"us | United States\n",
"eu | Europe\n",
"ap | Asia/Pacific\n",
"au | Australia\n",
"sa | South America\n",
"jp | Japan\n",
"in | India"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {
"id": "MM-UUANCwwqo"
},
"outputs": [],
"source": [
"# Set your region here in quotes\n",
"region = \"eu\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "JmNGrasvwnVX"
},
"source": [
"### Create Model Folder\n",
"\n",
"Create and download into the model folder"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {
"id": "rYAgrA75wv4U"
},
"outputs": [],
"source": [
"!mkdir -p Model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "EU1ZXGlGw8Rb"
},
"outputs": [],
"source": [
"!wget https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128_fp16.onnx -P /content/deepfakecollab/Model\n",
"!wget https://huggingface.co/hacksider/deep-live-cam/resolve/main/GFPGANv1.4.pth -P /content/deepfakecollab/Model"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "sNHGT2FjjZiK"
},
"source": [
"### [Ignore] DeBuggin (Run only when its necessery)\n",
"\n",
"Debugging to ensure cuda was used. if not check the version of the cudnn, edit the install of the onnxruntime package to install the right version. https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "ZQr-gWl6Ibw7",
"outputId": "ef468d22-a4c5-4f73-932b-dea535cc25ed"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"#define CUDNN_MAJOR 8\n",
"#define CUDNN_MINOR 9\n",
"#define CUDNN_PATCHLEVEL 6\n",
"--\n",
"#define CUDNN_VERSION (CUDNN_MAJOR * 1000 + CUDNN_MINOR * 100 + CUDNN_PATCHLEVEL)\n",
"\n",
"/* cannot use constexpr here since this is a C-only file */\n"
]
}
],
"source": [
"!cat /usr/include/cudnn_version.h | grep CUDNN_MAJOR -A 2"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "RZ9r-s_9Kdha",
"outputId": "b8778f0e-f10a-4725-cac7-67a435992100"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"1.18.0\n"
]
}
],
"source": [
"import onnxruntime as rt\n",
"print(rt.__version__)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Stwb4OHv_hs0",
"outputId": "41d0963a-31df-4fea-f80b-cbebef56d07c"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Active providers: ['CUDAExecutionProvider', 'CPUExecutionProvider']\n"
]
}
],
"source": [
"# Ensure the GPU providers are set explicitly\n",
"ort_session = rt.InferenceSession(\n",
" \"/content/deepfakecollab/Model/inswapper_128_fp16.onnx\",\n",
" providers=[\"CUDAExecutionProvider\", \"CPUExecutionProvider\"]\n",
")\n",
"\n",
"# Verify the active provider again\n",
"print(\"Active providers:\", ort_session.get_providers())\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "DTbmyO-CK-8b"
},
"source": [
"### Create the Colab Server\n",
"\n",
"In your Google Colab notebook, set up a TCP server that will receive frames, process them using the FACE_SWAPPER model, and send back the results."
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {
"id": "V_Luv36jjcmH"
},
"outputs": [],
"source": [
"import socket\n",
"import cv2\n",
"import numpy as np\n",
"import insightface\n",
"import threading\n",
"import torch\n",
"import onnxruntime\n",
"from typing import Any\n",
"from insightface.app.common import Face\n",
"import matplotlib.pyplot as plt\n"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "P4xE5bpeosP_",
"outputId": "7562a24e-ca5f-491e-884e-ce43b2ca0325"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"['CUDAExecutionProvider']\n"
]
}
],
"source": [
"# Check if CUDA is available\n",
"if 'CUDAExecutionProvider' in onnxruntime.get_available_providers():\n",
" providers = ['CUDAExecutionProvider']\n",
"elif 'TensorrtExecutionProvider' in onnxruntime.get_available_providers():\n",
" providers = ['TensorrtExecutionProvider']\n",
"else:\n",
" providers = ['CPUExecutionProvider']\n",
"print(providers)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "0OFrqUjgz7zb"
},
"outputs": [],
"source": [
"\n",
"FACE_SWAPPER = None\n",
"FACE_ANALYSER = insightface.app.FaceAnalysis(name='buffalo_l', providers=providers)\n",
"FACE_ANALYSER.prepare(ctx_id=0, det_size=(640, 640))\n",
"THREAD_LOCK = threading.Lock()\n",
"Frame = np.ndarray[Any, Any]"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {
"id": "ZbPSlDq_zJLj"
},
"outputs": [],
"source": [
"def get_face_swapper() -> Any:\n",
" global FACE_SWAPPER\n",
"\n",
" with THREAD_LOCK:\n",
" if FACE_SWAPPER is None:\n",
" model_path = \"/content/deepfakecollab/Model/inswapper_128_fp16.onnx\"\n",
" FACE_SWAPPER = insightface.model_zoo.get_model(model_path, providers=['CUDAExecutionProvider'])\n",
" return FACE_SWAPPER"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {
"id": "EnNJl5kDOjrL"
},
"outputs": [],
"source": [
"def swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:\n",
" return get_face_swapper().get(temp_frame, target_face, source_face, paste_back=True)"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {
"id": "y0VduuqyaA8R"
},
"outputs": [],
"source": [
"def get_face_analyser() -> Any:\n",
" #FACE_ANALYSER.prepare(ctx_id=0, det_size=(640, 640))\n",
" return FACE_ANALYSER"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {
"id": "ZTHjQ0p2aI9c"
},
"outputs": [],
"source": [
"def get_one_face(frame: Frame) -> Any:\n",
" face = get_face_analyser().get(frame)\n",
" try:\n",
" return min(face, key=lambda x: x.bbox[0])\n",
" except ValueError:\n",
" return None\n"
]
},
{
"cell_type": "code",
"execution_count": 23,
"metadata": {
"id": "ygR97Za2aTVO"
},
"outputs": [],
"source": [
"def get_many_faces(frame: Frame) -> Any:\n",
" try:\n",
" return get_face_analyser().get(frame)\n",
" except IndexError:\n",
" return None"
]
},
{
"cell_type": "code",
"execution_count": 24,
"metadata": {
"id": "05SG37Fbwip1"
},
"outputs": [],
"source": [
"def process_frame(source_face: Face, temp_frame: Frame,manyface: bool) -> Frame:\n",
"\n",
" if manyface:\n",
" many_faces = get_many_faces(temp_frame)\n",
" if many_faces:\n",
" for target_face in many_faces:\n",
" temp_frame = swap_face(source_face, target_face, temp_frame)\n",
" else:\n",
" target_face = get_one_face(temp_frame)\n",
" if target_face:\n",
" temp_frame = swap_face(source_face, target_face, temp_frame)\n",
" return temp_frame"
]
},
{
"cell_type": "code",
"execution_count": 25,
"metadata": {
"id": "znqUf8Gnq4dJ"
},
"outputs": [],
"source": [
"# Input and output ports for communication\n",
"local_in_source = 5555\n",
"local_in_temp = 5556\n",
"local_out_frame = 5557"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "FehUyVNvGGgZ"
},
"source": [
"### [Optional] For Live Streaming\n",
"For Live Streaming from webcam run this cell but for just image swap run the next cell. \n",
"\n",
"Note: Don't run both cells at same time. You should either run this or the one below\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 26,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "VYQRVPsdxizn",
"outputId": "def84ce6-46cf-401f-8e60-f1f0216b2975"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"PULL one socket bound to tcp://127.0.0.1:5555\n",
"OutputStream ffmpeg bound to tcp://127.0.0.1:5557?listen\n",
"InputStream ffmpeg bound from tcp://127.0.0.1:5556?listen\n"
]
}
],
"source": [
"import zmq\n",
"import threading\n",
"import cv2\n",
"import numpy as np\n",
"import msgpack\n",
"import queue\n",
"import time\n",
"import zlib\n",
"from tqdm import tqdm\n",
"import subprocess\n",
"from collections import deque\n",
"#import matplotlib.pyplot as plt\n",
"\n",
"def create_demo_image():\n",
" # Create a demo image (e.g., a solid color or pattern)\n",
" demo_image = np.zeros((540, 960, 3), dtype=np.uint8) # Black image\n",
" cv2.putText(demo_image, 'Demo Image', (50, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA)\n",
" return demo_image\n",
"\n",
"def pull_socket(local_in_port):\n",
" context = zmq.Context()\n",
" socket = context.socket(zmq.REP)\n",
" socket.setsockopt(zmq.RCVHWM, 100000)\n",
" socket.setsockopt(zmq.LINGER, 0)\n",
" address = f\"tcp://127.0.0.1:{local_in_port}\"\n",
" socket.bind(address) # Binding to a different local port\n",
" print(f\"PULL one socket bound to {address}\")\n",
" return socket\n",
"\n",
"\n",
"# Compress image\n",
"def compress_image(image, quality=95):\n",
" # Set the JPEG quality parameter\n",
" encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]\n",
" # Encode the image as a JPEG\n",
" result, encimg = cv2.imencode('.jpg', image, encode_param)\n",
"\n",
" if not result:\n",
" raise Exception(\"Image encoding failed\")\n",
"\n",
" # Decode the encoded image back to an image format\n",
" decimg = cv2.imdecode(encimg, 1)\n",
" return decimg\n",
"# Decompress image\n",
"def decompress_image(encimg):\n",
" image = cv2.imdecode(np.frombuffer(encimg, np.uint8), cv2.IMREAD_COLOR)\n",
" return image\n",
"\n",
"# Global variables\n",
"frames_array=deque(maxlen=2000)\n",
"source_frame = None\n",
"is_manyFace = None\n",
"frameSize = '960x540'\n",
"fps = None\n",
"\n",
"#functions\n",
"def pull_worker(pull_socket):\n",
" global source_frame,is_manyFace,frameSize,fps\n",
" while True:\n",
" try:\n",
" # Receive the JSON with total chunks\n",
" meta_data_json = pull_socket.recv_json()\n",
" #print(meta_data_json)\n",
" total_chunk = meta_data_json['total_chunk']\n",
" # Send acknowledgment for metadata\n",
" pull_socket.send_string(\"ACK\")\n",
" # Receive the array bytes\n",
" source_array_bytes =b''\n",
" for i in range(total_chunk):\n",
" chunk = pull_socket.recv()\n",
" source_array_bytes += chunk\n",
" pull_socket.send_string(f\"ACK {i + 1}/{total_chunk}\")\n",
"\n",
"\n",
" end_message = pull_socket.recv()\n",
" if end_message == b\"END\":\n",
" pull_socket.send_string(\"Final ACK\")\n",
"\n",
" # Deserialize the bytes back to an ndarray\n",
" source_array = np.frombuffer(source_array_bytes, dtype=np.dtype(meta_data_json['dtype_source'])).reshape(meta_data_json['shape_source'])\n",
"\n",
" #plt.imshow(source_array[:, :, ::-1])\n",
" #plt.show()\n",
" #frame_queue.append([\"source\", source_array])\n",
" source_frame = source_array\n",
" is_manyFace = meta_data_json['manyface']\n",
" frameSize = meta_data_json['size']\n",
" fps = meta_data_json[\"fps\"]\n",
"\n",
"\n",
" #process_queue.put((\"source\", source_array))\n",
" break\n",
" except zmq.Again:\n",
" # Sleep briefly to avoid busy-waiting\n",
" time.sleep(0.01)\n",
" except Exception as e:\n",
" print(f\"Error: {e}\")\n",
"def pull_worker_two(local_in_temp):\n",
" ffmpeg_receive_command = [\n",
" 'ffmpeg',\n",
" '-i',f'tcp://127.0.0.1:{local_in_temp}?listen',\n",
" '-f','rawvideo',\n",
" '-pix_fmt','bgr24',\n",
" '-s','960x540',\n",
" 'pipe:1'\n",
" ]\n",
" ffmpeg_receive_process = subprocess.Popen((ffmpeg_receive_command), stdout=subprocess.PIPE)\n",
" timefame =1/25\n",
" print(f\"InputStream ffmpeg bound from tcp://127.0.0.1:{local_in_temp}?listen\")\n",
" global source_frame,is_manyFace\n",
" while True:\n",
" try:\n",
"\n",
" # Receive the JSON with total chunks\n",
" # Read decoded frame from FFmpeg\n",
" raw_frame = ffmpeg_receive_process.stdout.read(960 * 540 * 3)\n",
" if not raw_frame:\n",
" break\n",
" framex = np.frombuffer(raw_frame, dtype=np.uint8).reshape((540, 960, 3))\n",
" #print(framex)\n",
" source_array = source_frame\n",
" is_many_face = is_manyFace\n",
"\n",
" if source_array is not None:\n",
" processed_array = process_frame(get_one_face(source_array),framex,is_many_face)\n",
" #plt.imshow(processed_array[:, :, ::-1])\n",
" #plt.show()\n",
" frames_array.append(processed_array)\n",
"\n",
" time.sleep(timefame)\n",
"\n",
"\n",
" except zmq.Again:\n",
" # Sleep briefly to avoid busy-waiting\n",
" time.sleep(0.01)\n",
" except Exception as e:\n",
" print(f\"Error: {e}\")\n",
"def push_worker(local_out_frame):\n",
" source_array = None\n",
" temp_array = None\n",
" is_many_face = None\n",
" global source_frame,is_manyFace,frameSize,fps\n",
" print(f\"OutputStream ffmpeg bound to tcp://127.0.0.1:{local_out_frame}?listen\")\n",
" ffmpeg_encode_command = [\n",
" 'ffmpeg',\n",
" '-f', 'rawvideo',\n",
" '-pix_fmt', 'bgr24',\n",
" '-s', '960x540',\n",
" '-r', '5',\n",
" '-i', 'pipe:',\n",
" #'-vf', 'fps=5',\n",
" #'-c:v', 'libx264',\n",
" #'-probesize', '32',\n",
" #'-analyzeduration', '0',\n",
" '-f', 'mpegts',\n",
" f'tcp://127.0.0.1:{local_out_frame}?listen'\n",
" ]\n",
" timefame =1/5\n",
"\n",
" ffmpeg_encode_process = subprocess.Popen((ffmpeg_encode_command), stdin=subprocess.PIPE)\n",
"\n",
"\n",
" demo_image = create_demo_image()\n",
" frames_to_skip = 200 # Number of frames to skip to reduce delay\n",
" frame_count =0\n",
" wait_frame = 100\n",
" try:\n",
" while True:\n",
"\n",
" # Get the processed array and metadata from the queue\n",
" if len(frames_array)>wait_frame:\n",
"\n",
" while len(frames_array)>7:#frame_count<frames_to_skip:\n",
" frames_array.popleft()\n",
" frame_count+=1\n",
" if len(frames_array)==7:\n",
" wait_frame = 5\n",
"\n",
" temp_array = frames_array.popleft()#process_queue.get()\n",
" source_array = source_frame\n",
" is_many_face = is_manyFace\n",
" #print(source_array,temp_array)\n",
" if source_array is not None and temp_array is not None:\n",
" #print(\"frames_array,\",len(frames_array))\n",
" #processed_array =process_frame(get_one_face(source_array),temp_array,is_many_face)\n",
" processed_array_bytes = temp_array.tobytes()#processed_array.tobytes()#temp_array.tobytes()\n",
"\n",
" ffmpeg_encode_process.stdin.write(processed_array_bytes)\n",
" #push_socket.send(zlib.compress(processed_array_bytes))\n",
" source_array = None\n",
" temp_array = None\n",
" is_many_face = None\n",
" else:\n",
" framex = demo_image\n",
" # Write the frame to FFmpeg for encoding and streaming\n",
" ffmpeg_encode_process.stdin.write(framex.tobytes())\n",
" time.sleep(timefame)\n",
" finally:\n",
" #ffmpeg_receive_process.terminate()\n",
" ffmpeg_encode_process.terminate()\n",
"\n",
"# Create sockets\n",
"pull_socket = pull_socket(local_in_source)\n",
"\n",
"# Run both workers in separate threads\n",
"# Start the pull worker thread\n",
"pull_thread = threading.Thread(target=pull_worker, args=(pull_socket,))\n",
"pull_thread.start()\n",
"# Start the push worker thread\n",
"pull_thread_two = threading.Thread(target=pull_worker_two, args=(local_in_temp,))\n",
"pull_thread_two.start()\n",
"# Start the push worker thread\n",
"push_thread = threading.Thread(target=push_worker,args=(local_out_frame,))\n",
"push_thread.start()\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Eb_cbyoaHKrT"
},
"source": [
"### [Optional] For Image Swap\n",
"For image swapping run this cell.\n",
"\n",
"Note: You can not run both cell at same time. Either This or The cell Above (For Live Streaming)"
]
},
{
"cell_type": "code",
"execution_count": 26,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "A_nrCfTpHWS4",
"outputId": "10038760-2787-44ed-83fb-35e74a757205"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"PULL one socket bound to tcp://127.0.0.1:5555\n",
"PULL one socket bound to tcp://127.0.0.1:5556\n",
"PUSH socket bound to tcp://127.0.0.1:5557\n"
]
}
],
"source": [
"import zmq\n",
"import threading\n",
"import cv2\n",
"import numpy as np\n",
"import msgpack\n",
"import queue\n",
"import time\n",
"import zlib\n",
"from tqdm import tqdm\n",
"import subprocess\n",
"from collections import deque\n",
"import matplotlib as plt\n",
"\n",
"def push_socket(local_out_port):\n",
" context = zmq.Context()\n",
" socket = context.socket(zmq.REQ)\n",
" socket.setsockopt(zmq.SNDHWM, 100000)\n",
" socket.setsockopt(zmq.LINGER, 0)\n",
" address = f\"tcp://127.0.0.1:{local_out_port}\"\n",
" socket.bind(address) # Binding to a local port\n",
" print(f\"PUSH socket bound to {address}\")\n",
" return socket\n",
"\n",
"def pull_socket(local_in_port):\n",
" context = zmq.Context()\n",
" socket = context.socket(zmq.REP)\n",
" socket.setsockopt(zmq.RCVHWM, 100000)\n",
" socket.setsockopt(zmq.LINGER, 0)\n",
" address = f\"tcp://127.0.0.1:{local_in_port}\"\n",
" socket.bind(address) # Binding to a different local port\n",
" print(f\"PULL one socket bound to {address}\")\n",
" return socket\n",
"\n",
"\n",
"# Compress image\n",
"def compress_image(image, quality=95):\n",
" # Set the JPEG quality parameter\n",
" encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]\n",
" # Encode the image as a JPEG\n",
" result, encimg = cv2.imencode('.jpg', image, encode_param)\n",
"\n",
" if not result:\n",
" raise Exception(\"Image encoding failed\")\n",
"\n",
" # Decode the encoded image back to an image format\n",
" decimg = cv2.imdecode(encimg, 1)\n",
" return decimg\n",
"# Decompress image\n",
"def decompress_image(encimg):\n",
" image = cv2.imdecode(np.frombuffer(encimg, np.uint8), cv2.IMREAD_COLOR)\n",
" return image\n",
"\n",
"# Global variables\n",
"frames_array=deque(maxlen=2000)\n",
"source_frame = None\n",
"is_manyFace = None\n",
"frameSize = '640x480'\n",
"fps = None\n",
"\n",
"#functions\n",
"def pull_worker(pull_socket):\n",
" global source_frame,is_manyFace,frameSize,fps\n",
" while True:\n",
" try:\n",
" # Receive the JSON with total chunks\n",
" meta_data_json = pull_socket.recv_json()\n",
" #print(meta_data_json)\n",
" total_chunk = meta_data_json['total_chunk']\n",
" # Send acknowledgment for metadata\n",
" pull_socket.send_string(\"ACK\")\n",
" # Receive the array bytes\n",
" source_array_bytes =b''\n",
" for i in range(total_chunk):\n",
" chunk = pull_socket.recv()\n",
" source_array_bytes += chunk\n",
" pull_socket.send_string(f\"ACK {i + 1}/{total_chunk}\")\n",
"\n",
"\n",
" end_message = pull_socket.recv()\n",
" if end_message == b\"END\":\n",
" pull_socket.send_string(\"Final ACK\")\n",
"\n",
" # Deserialize the bytes back to an ndarray\n",
" source_array = np.frombuffer(source_array_bytes, dtype=np.dtype(meta_data_json['dtype_source'])).reshape(meta_data_json['shape_source'])\n",
"\n",
" #plt.imshow(source_array[:, :, ::-1])\n",
" #plt.show()\n",
" #frame_queue.append([\"source\", source_array])\n",
" source_frame = source_array\n",
" is_manyFace = meta_data_json['manyface']\n",
" frames_array.append([\"source\",source_array,is_manyFace])\n",
"\n",
"\n",
" #process_queue.put((\"source\", source_array))\n",
" #break\n",
" except zmq.Again:\n",
" # Sleep briefly to avoid busy-waiting\n",
" time.sleep(0.01)\n",
" except Exception as e:\n",
" print(f\"Error: {e}\")\n",
"def pull_worker_two(pull_socket_two):\n",
"\n",
" while True:\n",
" try:\n",
"\n",
" # Receive the JSON with total chunks\n",
" meta_data_json = pull_socket_two.recv_json()\n",
" #print(meta_data_json)\n",
" total_chunk = meta_data_json['total_chunk']\n",
" # Send acknowledgment for metadata\n",
" pull_socket_two.send_string(\"ACK\")\n",
" # Receive the array bytes\n",
" temp_array_bytes =b''\n",
" for i in range(total_chunk):\n",
" chunk = pull_socket_two.recv()\n",
" temp_array_bytes += chunk\n",
" pull_socket_two.send_string(f\"ACK {i + 1}/{total_chunk}\")\n",
"\n",
"\n",
" end_message = pull_socket_two.recv()\n",
" if end_message == b\"END\":\n",
" pull_socket_two.send_string(\"Final ACK\")\n",
"\n",
" # Deserialize the bytes back to an ndarray\n",
" temp_array = np.frombuffer(temp_array_bytes, dtype=np.dtype(meta_data_json['dtype_temp'])).reshape(meta_data_json['shape_temp'])\n",
" #if source_frame is not None:\n",
" frames_array.append([\"temp\",temp_array])\n",
" print(\"added\",len(frames_array))\n",
" #break\n",
"\n",
"\n",
"\n",
"\n",
" except zmq.Again:\n",
" # Sleep briefly to avoid busy-waiting\n",
" time.sleep(0.01)\n",
" except Exception as e:\n",
" print(f\"Error: {e}\")\n",
"def push_worker(push_socket):\n",
"\n",
" global source_frame,is_manyFace,frameSize,fps\n",
" temp_frm = None\n",
" source_frm = None\n",
" is_manyface = None\n",
" try:\n",
" while True:\n",
" # Get the processed array and metadata from the queue\n",
" if len(frames_array)>0:\n",
" #print(frames_array)\n",
" item = frames_array.popleft()#process_queue.get()\n",
" if item[0]==\"source\":\n",
" source_frm = item[1]\n",
" is_manyface = item[2]\n",
" if item[0]==\"temp\":\n",
" temp_frm = item[1]\n",
" print(\"Recieved\")\n",
"\n",
" if temp_frm is not None and source_frm is not None:\n",
"\n",
" processed_frm =process_frame(get_one_face(source_frm),temp_frm,is_manyface)\n",
"\n",
" face_bytes = processed_frm.tobytes()\n",
" chunk_size = 1024*200\n",
" total_chunk = len(face_bytes) // chunk_size + 1\n",
" metadata ={\n",
"\n",
" 'dtype_source':str(processed_frm.dtype),\n",
" 'shape_source':processed_frm.shape,\n",
" 'size':'640x480',\n",
" 'fps':'60'\n",
" #'shape_temp':temp_frame.shape\n",
" }\n",
" new_metadata = {'total_chunk': total_chunk}\n",
" metadata.update(new_metadata)\n",
" # Send metadata first\n",
" push_socket.send_json(metadata)\n",
" # Wait for acknowledgment for metadata\n",
" ack = push_socket.recv_string()\n",
" with tqdm(total=total_chunk, desc=\"Sending chunks\", unit=\"chunk\") as pbar:\n",
" for i in range(total_chunk):\n",
" chunk = face_bytes[i * chunk_size:(i + 1) * chunk_size]\n",
" # Send the chunk\n",
" push_socket.send(chunk)\n",
" # Wait for acknowledgment after sending each chunk\n",
" ack = push_socket.recv_string()\n",
" pbar.set_postfix_str(f'Chunk {i + 1}/{total_chunk} ack: {ack}')\n",
" pbar.update(1)\n",
"\n",
" # Send a final message to indicate all chunks are sent\n",
" push_socket.send(b\"END\")\n",
" # Wait for the final reply\n",
" final_reply_message = push_socket.recv_string()\n",
" print(f\"Received final reply: {final_reply_message}\")\n",
"\n",
" temp_frm = None\n",
" source_frm = None\n",
" is_manyface = None\n",
"\n",
"\n",
" except Exception as e:\n",
" print (f\"Error in Push Wokr {e}\")\n",
"\n",
"local_in_source = 5555\n",
"local_in_temp = 5556\n",
"local_out_frame = 5557\n",
"\n",
"# Create sockets\n",
"pull_socket_ = pull_socket(local_in_source)\n",
"pull_socket_two = pull_socket(local_in_temp)\n",
"push_socket_ = push_socket(local_out_frame)\n",
"# Run both workers in separate threads\n",
"# Start the pull worker thread\n",
"pull_thread = threading.Thread(target=pull_worker, args=(pull_socket_,))\n",
"pull_thread.start()\n",
"# Start the push worker thread\n",
"pull_thread_two = threading.Thread(target=pull_worker_two, args=(pull_socket_two,))\n",
"pull_thread_two.start()\n",
"# Start the push worker thread\n",
"push_thread = threading.Thread(target=push_worker,args=(push_socket_,))\n",
"push_thread.start()\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "a9JW1rhH3H-T"
},
"source": [
"### [Optional] Open FRP tunnel"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "VNDmeHcx3Zrx"
},
"outputs": [],
"source": [
"\n",
"frpc_config = f\"\"\"serverAddr = \"194.113.64.71\"\n",
"serverPort = 7000\n",
"\n",
"[[proxies]]\n",
"name = \"Pull-tcp\"\n",
"type = \"tcp\"\n",
"localIP = \"127.0.0.1\"\n",
"localPort = {local_in_source}\n",
"remotePort = 6000\n",
"\n",
"[[proxies]]\n",
"name = \"Pull-tcp_two\"\n",
"type = \"tcp\"\n",
"localIP = \"127.0.0.1\"\n",
"localPort = {local_in_temp}\n",
"remotePort = 6001\n",
"\n",
"[[proxies]]\n",
"name = \"Push-tcp\"\n",
"type = \"tcp\"\n",
"localIP = \"127.0.0.1\"\n",
"localPort = {local_out_frame}\n",
"remotePort = 6002\"\"\"\n",
"with open('frp_0.59.0_linux_amd64/frpc.toml', 'w') as f:\n",
" f.write(frpc_config)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#Remote tcp\n",
"print(f'tcp://194.113.64.71:6000 ---> {local_in_source}')\n",
"print(f'tcp://194.113.64.71:6001 ---> {local_in_temp}')\n",
"print(f'tcp://194.113.64.71:6002 ---> {local_out_frame}')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Ucjrp4fM3mQ4"
},
"outputs": [],
"source": [
"from subprocess import Popen, PIPE\n",
"import time\n",
"ps = Popen('/content/deepfakecollab/Scripts/open_tunnel_frs.sh', stdout=PIPE, stderr=PIPE)\n",
"time.sleep(3)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "SDFT4XXJkVxT",
"outputId": "1738d849-a3e1-403d-95d0-d3a4e4f6cd44"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Opening tunnel\n",
"\u001b[1;34m2024-08-03 20:10:33.967 [I] [sub/root.go:142] start frpc service for config file [frp_0.59.0_linux_amd64/frpc.toml]\n",
"\u001b[0m\u001b[1;34m2024-08-03 20:10:33.967 [I] [client/service.go:294] try to connect to server...\n",
"\u001b[0m\u001b[1;34m2024-08-03 20:10:34.377 [I] [client/service.go:286] [b793872517fab479] login to server success, get run id [b793872517fab479]\n",
"\u001b[0m\u001b[1;34m2024-08-03 20:10:34.377 [I] [proxy/proxy_manager.go:173] [b793872517fab479] proxy added: [Pull-tcp Pull-tcp_two Push-tcp]\n",
"\u001b[0m\u001b[1;34m2024-08-03 20:10:34.514 [I] [client/control.go:168] [b793872517fab479] [Pull-tcp] start proxy success\n",
"\u001b[0m\u001b[1;34m2024-08-03 20:10:34.514 [I] [client/control.go:168] [b793872517fab479] [Pull-tcp_two] start proxy success\n",
"\u001b[0m\u001b[1;34m2024-08-03 20:10:34.514 [I] [client/control.go:168] [b793872517fab479] [Push-tcp] start proxy success\n",
"\u001b[0m^C\n"
]
}
],
"source": [
"!/content/deepfakecollab/Scripts/open_tunnel_frs.sh"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "buCnH_YpxKWa"
},
"source": [
"### [Optional] Open Ngrok tunnel"
]
},
{
"cell_type": "code",
"execution_count": 27,
"metadata": {
"id": "HyY2B5bmxlui"
},
"outputs": [],
"source": [
"from subprocess import Popen, PIPE\n",
"import shlex\n",
"import json\n",
"import time\n",
"\n",
"\n",
"def run_with_pipe(command):\n",
" commands = list(map(shlex.split,command.split(\"|\")))\n",
" ps = Popen(commands[0], stdout=PIPE, stderr=PIPE)\n",
" for command in commands[1:]:\n",
" ps = Popen(command, stdin=ps.stdout, stdout=PIPE, stderr=PIPE)\n",
" return ps.stdout.readlines()\n",
"\n",
"\n",
"def get_tunnel_adresses():\n",
" info = run_with_pipe(\"curl http://localhost:4040/api/tunnels\")\n",
" assert info\n",
"\n",
" info = json.loads(info[0])\n",
" for tunnel in info['tunnels']:\n",
" url = tunnel['public_url']\n",
" port = url.split(':')[-1]\n",
" local_port = tunnel['config']['addr'].split(':')[-1]\n",
" print(f'{url} -> {local_port} [{tunnel[\"name\"]}]')\n",
" if tunnel['name'] == 'input':\n",
" in_addr = url\n",
" elif tunnel['name'] == 'inputtwo':\n",
" in_addrtwo = url\n",
" elif tunnel['name'] == 'output':\n",
" out_addr = url\n",
" else:\n",
" print(f'unknown tunnel: {tunnel[\"name\"]}')\n",
"\n",
" return in_addr,in_addrtwo, out_addr"
]
},
{
"cell_type": "code",
"execution_count": 28,
"metadata": {
"id": "Lnqdi11sxGGR"
},
"outputs": [],
"source": [
"config =\\\n",
"f\"\"\"\n",
"version: 2\n",
"authtoken: {authtoken}\n",
"region: {region}\n",
"console_ui: False\n",
"tunnels:\n",
" input:\n",
" addr: {local_in_source}\n",
" proto: tcp\n",
" inputtwo:\n",
" addr: {local_in_temp}\n",
" proto: tcp\n",
" output:\n",
" addr: {local_out_frame}\n",
" proto: tcp\n",
"\"\"\"\n",
"\n",
"with open('ngrok.conf', 'w') as f:\n",
" f.write(config)"
]
},
{
"cell_type": "code",
"execution_count": 29,
"metadata": {
"id": "bQiD7DmexYvy"
},
"outputs": [],
"source": [
"# (Re)Open tunnel\n",
"ps = Popen('/content/deepfakecollab/Scripts/open_tunnel_ngrok.sh', stdout=PIPE, stderr=PIPE)\n",
"time.sleep(3)"
]
},
{
"cell_type": "code",
"execution_count": 30,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "4kC3axJEx5e8",
"outputId": "21867038-e5c4-4038-9876-c41a569b09e7"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tcp://0.tcp.eu.ngrok.io:11640 -> 5556 [inputtwo]\n",
"tcp://4.tcp.eu.ngrok.io:16503 -> 5555 [input]\n",
"tcp://0.tcp.eu.ngrok.io:13155 -> 5557 [output]\n",
"Tunnel opened\n"
]
}
],
"source": [
"# Get tunnel addresses\n",
"try:\n",
" in_addr,in_addr_two, out_addr = get_tunnel_adresses()\n",
" print(\"Tunnel opened\")\n",
"except Exception as e:\n",
" [print(l.decode(), end='') for l in ps.stdout.readlines()]\n",
" print(\"Something went wrong, reopen the tunnel\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "F9Qs1U8gu-l3"
},
"source": [
"### **DeBugging (Ignore)**"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "d_gTh_hfbstK"
},
"source": [
"### View Source Image"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "otrLL3NnHCt3",
"outputId": "60a822ea-b55c-4f31-cf19-ddc4390de81f"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"False\n"
]
}
],
"source": [
"print(is_manyFace)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "9Kqc-HRYwa08",
"outputId": "0ff9e514-b029-4df2-c52f-230ddd9dd2a6"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"0\n"
]
}
],
"source": [
"print(len(frames_array))"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [
"GYFbuuCdvTaO",
"0rTCd2R1uw7A",
"9n-qESinvM22",
"JmNGrasvwnVX",
"sNHGT2FjjZiK",
"DTbmyO-CK-8b",
"FehUyVNvGGgZ",
"Eb_cbyoaHKrT",
"a9JW1rhH3H-T",
"buCnH_YpxKWa",
"d_gTh_hfbstK"
],
"gpuType": "L4",
"machine_shape": "hm",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}