feat:put a4k to rt, all three outputs are working

This commit is contained in:
zpeng11 committed 2025-12-28 17:53:21 -05:00
1 parent c537839e6c
commit ebe0073eca
6 files changed
+29 -41

No files matched your search

+6 -6
View File
@@ -29,8 +29,8 @@ parser.add_argument('--alpha_split', action='store_true')
parser.add_argument('--bongo', action='store_true')
parser.add_argument('--extend_movement', action='store_true')
parser.add_argument('--filter_min_cutoff', type=float, default=1.0)
parser.add_argument('--filter_beta', type=float, default=0.0)
parser.add_argument('--filter_min_cutoff', type=float, default=0.3)
parser.add_argument('--filter_beta', type=float, default=1.0)
parser.add_argument('--simplify', type=int, default=1)
@@ -49,7 +49,7 @@ parser.add_argument('--interpolation_half', action='store_true')
parser.add_argument('--use_sr', action='store_true')
parser.add_argument('--sr_x4', action='store_true')
parser.add_argument('--sr_half', action='store_true')
parser.add_argument('--anime4k', action='store_true')
parser.add_argument('--sr_a4k', action='store_true')
parser.add_argument('--cache', type=str, default='256mb')
@@ -78,12 +78,12 @@ if not args.cam_input and args.mouse_input is None \
and args.ifm_input is None and args.osf_input is None:
args.debug_input = True # Default to debug input
if args.use_sr:
args.anime4k = False # Disable anime4k if SR is enabled
if args.sr_a4k:
args.use_sr = True
if args.output_spout2:
args.alpha_split = False # Disable alpha split for spout2 output
args.model_output_size = 1024 if args.use_sr or args.anime4k else 512
args.model_output_size = 1024 if args.use_sr else 512
print("Arguments:", args)
+2
View File
@@ -38,6 +38,7 @@ def get_core(
use_sr:bool = False,
sr_x4:bool = True,
sr_half:bool = True,
sr_a4k:bool = False,
):
if use_tensorrt:
try:
@@ -60,6 +61,7 @@ def get_core(
sr_model_enable=use_sr,
sr_model_scale=4 if sr_x4 else 2,
sr_model_fp16=sr_half,
sr_a4k=sr_a4k,
vram_cache_size=model_cache_size,
cache_max_giga=cacher_ram_size,
+10 -3
View File
@@ -10,6 +10,7 @@ from .utils.timer_wait import wait_until
from PIL import Image
from .utils.fps import FPS
import pyvirtualcam
from OpenGL.GL import GL_RGBA
def main():
@@ -87,6 +88,10 @@ def main():
backend='obs',
fmt=pyvirtualcam.PixelFormat.RGB)
print(f'Using virtual camera: {virtual_cam.device}')
elif args.output_spout2:
from PySpout import SpoutSender
spout_sender = SpoutSender("EasyVtuber", cam_width_scale * args.model_output_size,
args.model_output_size, GL_RGBA)
else:
print("Using OpenCV windows for output display.")
@@ -99,14 +104,16 @@ def main():
for i in range(args.interpolation_scale):
ret_batch_shm_channels[i].acquire()
for i in range(args.interpolation_scale):
wait_until(last_time + interval)
last_time += interval
if args.output_virtual_cam:
virtual_cam.send(np_ret_shms[i])
elif args.output_spout2:
spout_sender.send_image(np_ret_shms[i], False)
else:
cv2.imshow("Output Frame Batch {}".format(i), np_ret_shms[i])
cv2.waitKey(1)
wait_until(last_time + interval)
last_time += interval
ret_batch_shm_channels[i].release()
print("pipeline FPS: {:.2f}, Input FPS: {:.2f}, Model Avg Interval: {:.2f} ms, Cache Hit Ratio: {:.2f}%, GPU Cache Hit Ratio: {:.2f}%, current time {:.5f}".format(
infer_process.pipeline_fps_number.value,
+9 -30
View File
@@ -9,7 +9,6 @@ from .utils.channel_shared_mem import SharedMemoryExclusiveChannel
from .utils.pose_simplify import pose_simplify
from .utils.fps import FPS, Interval
from typing import List
import pyanime4k
class ModelClientProcess(Process):
def __init__(self, input_image, pose_position_shm: shared_memory.SharedMemory , input_fps):
@@ -48,14 +47,6 @@ class ModelClientProcess(Process):
for i in range(args.interpolation_scale)
]
if args.anime4k:
self.a4k_processor = pyanime4k.Processor(
processor_type="opencl",
device=0,
model="acnet-gan"
)
print("Anime4K Loaded")
model_infer_average_interval: Interval = Interval()
pipeline_fps = FPS()
@@ -77,7 +68,9 @@ class ModelClientProcess(Process):
use_sr=args.use_sr,
sr_half=args.sr_half,
sr_x4=args.sr_x4)
sr_x4=args.sr_x4,
sr_a4k=args.sr_a4k,
)
model.setImage(self.input_image)
model_infer_average_interval.start()
model.inference([np.zeros((1, 45), dtype=np.float32)]) # Warm up
@@ -172,20 +165,10 @@ class ModelClientProcess(Process):
(0, 80),
cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
rgba_image = None
alpha_channel = None
rgb_channels = None
if args.anime4k or args.alpha_split:
rgba_image = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2RGBA)
alpha_channel = np.ascontiguousarray(rgba_image[:, :, 3])
rgb_channels = np.ascontiguousarray(rgba_image[:,:,:3])
if args.anime4k:
rgb_channels = self.a4k_processor(rgb_channels)
alpha_channel = cv2.resize(alpha_channel, None, fx=2, fy=2)
rgba_image = cv2.merge((rgb_channels, alpha_channel))
if args.alpha_split:
rgba_image = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2RGBA)
alpha_channel = rgba_image[:, :, 3]
rgb_channels = rgba_image[:,:,:3]
alpha_image = cv2.cvtColor(alpha_channel, cv2.COLOR_GRAY2RGB)
rgb_channels = cv2.hconcat([rgb_channels, alpha_image])
@@ -193,18 +176,14 @@ class ModelClientProcess(Process):
if args.alpha_split:
bgr_channels = cv2.cvtColor(rgb_channels, cv2.COLOR_RGB2BGR)
else:
if args.anime4k:
bgr_channels = cv2.cvtColor(rgba_image, cv2.COLOR_RGBA2BGR)
else:
bgr_channels = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2BGR)
bgr_channels = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2BGR)
ret.append(bgr_channels)
elif args.output_virtual_cam:
if rgb_channels is None:
if not args.alpha_split:
rgb_channels = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2RGB)
if args.anime4k and not args.alpha_split:
rgb_channels = cv2.cvtColor(rgba_image, cv2.COLOR_RGBA2RGB)
ret.append(rgb_channels)
else:
rgba_image = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2RGBA)
ret.append(rgba_image)
return ret
+1 -1
View File
@@ -135,7 +135,7 @@ class OSFClientProcess(Process):
mouth_eye_vector[2] = 1 - data['leftEyeOpen']
mouth_eye_vector[3] = 1 - data['rightEyeOpen']
mouth_eye_vector[14] = max(data['MouthOpen'], 0) * 1.5 # Open larger mouth
mouth_eye_vector[14] = max(data['MouthOpen'], 0) * 2 # Open larger mouth
# print(mouth_eye_vector[14])
mouth_eye_vector[25] = iris_x_filter(-data['eyeRotationY'] * 3 - (data['rotationX']) / 57.3 * 1.5, timestamp=time.perf_counter())