mirror of
https://github.com/facefusion/facefusion.git
synced 2026-09-17 04:55:26 +02:00
95 lines
3.6 KiB
Python
95 lines
3.6 KiB
Python
import ctypes
|
|
import struct
|
|
from typing import Optional
|
|
|
|
from facefusion.libraries import aom as aom_module
|
|
from facefusion.types import AomEncoder, BitRate, Buffer, Resolution
|
|
|
|
|
|
def create(frame_resolution : Resolution, bitrate : BitRate, thread_count : int, cpu_count : int) -> Optional[AomEncoder]:
|
|
aom_library = aom_module.create_static_library()
|
|
|
|
if aom_library:
|
|
aom_encoder = ctypes.create_string_buffer(1152)
|
|
aom_codec = ctypes.c_void_p.in_dll(aom_library, 'aom_codec_av1_cx_algo')
|
|
|
|
config_buffer = ctypes.create_string_buffer(1024)
|
|
|
|
if aom_library.aom_codec_enc_config_default(ctypes.byref(aom_codec), config_buffer, 1) == 0:
|
|
struct.pack_into('I', config_buffer, 4, thread_count)
|
|
struct.pack_into('I', config_buffer, 12, frame_resolution[0])
|
|
struct.pack_into('I', config_buffer, 16, frame_resolution[1])
|
|
struct.pack_into('I', config_buffer, 136, bitrate)
|
|
struct.pack_into('I', config_buffer, 192, 30)
|
|
|
|
if aom_library.aom_codec_enc_init_ver(aom_encoder, ctypes.byref(aom_codec), config_buffer, 0, 25) == 0:
|
|
aom_library.aom_codec_control(aom_encoder, 13, ctypes.c_int(cpu_count))
|
|
aom_library.aom_codec_control(aom_encoder, 75, ctypes.c_int(1))
|
|
aom_library.aom_codec_control(aom_encoder, 106, ctypes.c_int(0))
|
|
aom_library.aom_codec_control(aom_encoder, 122, ctypes.c_int(0))
|
|
aom_library.aom_codec_control(aom_encoder, 123, ctypes.c_int(0))
|
|
ctypes.memmove(ctypes.addressof(aom_encoder) + 128, config_buffer, 1024)
|
|
return aom_encoder
|
|
|
|
return None
|
|
|
|
|
|
def encode(aom_encoder : AomEncoder, input_buffer : Buffer, frame_resolution : Resolution, frame_index : int) -> Buffer:
|
|
aom_library = aom_module.create_static_library()
|
|
output_buffer = bytes()
|
|
|
|
if aom_library:
|
|
temp_buffer = ctypes.create_string_buffer(256)
|
|
encode_buffer = ctypes.create_string_buffer(input_buffer)
|
|
|
|
if aom_library.aom_img_wrap(temp_buffer, 0x102, frame_resolution[0], frame_resolution[1], 1, encode_buffer) and aom_library.aom_codec_encode(aom_encoder, temp_buffer, frame_index, 1, 0, 1) == 0:
|
|
output_buffer = collect(aom_encoder)
|
|
|
|
return output_buffer
|
|
|
|
|
|
def collect(aom_encoder : AomEncoder) -> Buffer:
|
|
aom_library = aom_module.create_static_library()
|
|
output_parts = []
|
|
|
|
packet_cursor = ctypes.c_void_p(0)
|
|
packet = aom_library.aom_codec_get_cx_data(aom_encoder, ctypes.byref(packet_cursor))
|
|
|
|
while packet:
|
|
if ctypes.c_int.from_address(packet).value == 0:
|
|
buffer_pointer = ctypes.c_void_p.from_address(packet + 8).value
|
|
buffer_size = ctypes.c_size_t.from_address(packet + 16).value
|
|
output_parts.append(ctypes.string_at(buffer_pointer, buffer_size))
|
|
|
|
packet = aom_library.aom_codec_get_cx_data(aom_encoder, ctypes.byref(packet_cursor))
|
|
|
|
return bytes().join(output_parts)
|
|
|
|
|
|
def update_resolution(aom_encoder : AomEncoder, frame_resolution : Resolution) -> bool:
|
|
aom_library = aom_module.create_static_library()
|
|
|
|
if aom_library:
|
|
struct.pack_into('I', aom_encoder, 128 + 12, frame_resolution[0])
|
|
struct.pack_into('I', aom_encoder, 128 + 16, frame_resolution[1])
|
|
return aom_library.aom_codec_enc_config_set(aom_encoder, ctypes.cast(ctypes.addressof(aom_encoder) + 128, ctypes.c_void_p)) == 0
|
|
|
|
return False
|
|
|
|
|
|
def update_bitrate(aom_encoder : AomEncoder, bitrate : BitRate) -> bool:
|
|
aom_library = aom_module.create_static_library()
|
|
|
|
if aom_library:
|
|
struct.pack_into('I', aom_encoder, 128 + 136, bitrate)
|
|
return aom_library.aom_codec_enc_config_set(aom_encoder, ctypes.cast(ctypes.addressof(aom_encoder) + 128, ctypes.c_void_p)) == 0
|
|
|
|
return False
|
|
|
|
|
|
def destroy(aom_encoder : AomEncoder) -> None:
|
|
aom_library = aom_module.create_static_library()
|
|
|
|
if aom_library:
|
|
aom_library.aom_codec_destroy(aom_encoder)
|