import time import torch import torch.nn as nn import tha2.poser.modes.mode_20 import tha3.poser.modes.standard_float import tha3.poser.modes.separable_float import tha3.poser.modes.standard_half import tha3.poser.modes.separable_half from torch.nn.functional import interpolate from args import args from collections import OrderedDict # THA4 适配器导入 try: from tha4_adapter import THA4Wrapper THA4_AVAILABLE = True except ImportError: THA4_AVAILABLE = False print("Warning: THA4 adapter not available") class TalkingAnimeLight(nn.Module): def __init__(self): super(TalkingAnimeLight, self).__init__() self.face_morpher = tha2.poser.modes.mode_20.load_face_morpher('pretrained/face_morpher.pt') self.two_algo_face_rotator = tha2.poser.modes.mode_20.load_face_rotater('pretrained/two_algo_face_rotator.pt') self.combiner = tha2.poser.modes.mode_20.load_combiner('pretrained/combiner.pt') self.face_cache = OrderedDict() self.tot = 0 self.hit = 0 def forward(self, image, mouth_eye_vector, pose_vector, mouth_eye_vector_c, ratio=None): x = image.clone() if args.perf == 'model': tic = time.perf_counter() input_hash = hash(tuple(mouth_eye_vector_c)) cached = self.face_cache.get(input_hash) self.tot += 1 if cached is None: mouth_eye_morp_image = self.face_morpher(image[:, :, 32:224, 32:224], mouth_eye_vector) self.face_cache[input_hash] = mouth_eye_morp_image.detach() if len(self.face_cache) > args.max_gpu_cache_len: self.face_cache.popitem(last=False) else: self.hit += 1 mouth_eye_morp_image = cached self.face_cache.move_to_end(input_hash) if args.debug and ratio is not None: ratio.value = self.hit / self.tot if args.perf == 'model': print(" - face_morpher", (time.perf_counter() - tic) * 1000) tic = time.perf_counter() x[:, :, 32:224, 32:224] = mouth_eye_morp_image rotate_image = self.two_algo_face_rotator(x, pose_vector)[:2] if args.perf == 'model': print(" - rotator", (time.perf_counter() - tic) * 1000) tic = time.perf_counter() output_image = self.combiner(rotate_image[0], rotate_image[1], pose_vector) if args.perf == 'model': print(" - combiner", (time.perf_counter() - tic) * 1000) tic = time.perf_counter() return output_image class TalkingAnime3(nn.Module): def __init__(self): super(TalkingAnime3, self).__init__() if args.model == "standard_float": if args.eyebrow: self.eyebrow_decomposer = tha3.poser.modes.standard_float.load_eyebrow_decomposer( 'data/models/standard_float/eyebrow_decomposer.pt') self.eyebrow_morphing_combiner = tha3.poser.modes.standard_float.load_eyebrow_morphing_combiner( 'data/models/standard_float/eyebrow_morphing_combiner.pt') self.face_morpher = tha3.poser.modes.standard_float.load_face_morpher( 'data/models/standard_float/face_morpher.pt') self.two_algo_face_body_rotator = tha3.poser.modes.standard_float.load_two_algo_generator( 'data/models/standard_float/two_algo_face_body_rotator.pt') self.editor = tha3.poser.modes.standard_float.load_editor('data/models/standard_float/editor.pt') elif args.model == "standard_half": if args.eyebrow: self.eyebrow_decomposer = tha3.poser.modes.standard_half.load_eyebrow_decomposer( 'data/models/standard_half/eyebrow_decomposer.pt') self.eyebrow_morphing_combiner = tha3.poser.modes.standard_half.load_eyebrow_morphing_combiner( 'data/models/standard_half/eyebrow_morphing_combiner.pt') self.face_morpher = tha3.poser.modes.standard_half.load_face_morpher( 'data/models/standard_half/face_morpher.pt') self.two_algo_face_body_rotator = tha3.poser.modes.standard_half.load_two_algo_generator( 'data/models/standard_half/two_algo_face_body_rotator.pt') self.editor = tha3.poser.modes.standard_half.load_editor('data/models/standard_half/editor.pt') elif args.model == "separable_float": if args.eyebrow: self.eyebrow_decomposer = tha3.poser.modes.separable_float.load_eyebrow_decomposer( 'data/models/separable_float/eyebrow_decomposer.pt') self.eyebrow_morphing_combiner = tha3.poser.modes.separable_float.load_eyebrow_morphing_combiner( 'data/models/separable_float/eyebrow_morphing_combiner.pt') self.face_morpher = tha3.poser.modes.separable_float.load_face_morpher( 'data/models/separable_float/face_morpher.pt') self.two_algo_face_body_rotator = tha3.poser.modes.separable_float.load_two_algo_generator( 'data/models/separable_float/two_algo_face_body_rotator.pt') self.editor = tha3.poser.modes.separable_float.load_editor('data/models/separable_float/editor.pt') elif args.model == "separable_half": if args.eyebrow: self.eyebrow_decomposer = tha3.poser.modes.separable_half.load_eyebrow_decomposer( 'data/models/separable_half/eyebrow_decomposer.pt') self.eyebrow_morphing_combiner = tha3.poser.modes.separable_half.load_eyebrow_morphing_combiner( 'data/models/separable_half/eyebrow_morphing_combiner.pt') self.face_morpher = tha3.poser.modes.separable_half.load_face_morpher( 'data/models/separable_half/face_morpher.pt') self.two_algo_face_body_rotator = tha3.poser.modes.separable_half.load_two_algo_generator( 'data/models/separable_half/two_algo_face_body_rotator.pt') self.editor = tha3.poser.modes.separable_half.load_editor('data/models/separable_half/editor.pt') else: raise RuntimeError("Invalid model: '%s'" % args.model) self.face_cache = OrderedDict() self.tot = 0 self.hit = 0 def forward(self, image, mouth_eye_vector, pose_vector, eyebrow_vector, mouth_eye_vector_c, eyebrow_vector_c, ratio=None): if args.perf == 'model': tic = time.perf_counter() x = image.clone() if args.eyebrow: input_hash = hash(tuple(eyebrow_vector_c + mouth_eye_vector_c)) else: input_hash = hash(tuple(mouth_eye_vector_c)) cached = self.face_cache.get(input_hash) self.tot += 1 if cached is None: face_image = x[:, :, 32:32 + 192, (32 + 128):(32 + 192 + 128)].clone() if args.eyebrow: eyebrow_morp_image = self.eyebrow_decomposer(x[:, :, 64:192, 64 + 128:192 + 128].clone()) eyebrow_morp_image = \ self.eyebrow_morphing_combiner(eyebrow_morp_image[3], eyebrow_morp_image[0], eyebrow_vector)[2] face_image[:, :, 32:32 + 128, 32:32 + 128] = eyebrow_morp_image mouth_eye_morp_image = self.face_morpher(face_image, mouth_eye_vector)[0] self.face_cache[input_hash] = mouth_eye_morp_image.detach() if len(self.face_cache) > args.max_gpu_cache_len: self.face_cache.popitem(last=False) else: self.hit += 1 mouth_eye_morp_image = cached self.face_cache.move_to_end(input_hash) if args.debug and ratio is not None: ratio.value = self.hit / self.tot if args.perf == 'model': print(" - face_morpher", (time.perf_counter() - tic) * 1000) tic = time.perf_counter() x[:, :, 32:32 + 192, (32 + 128):(32 + 192 + 128)] = mouth_eye_morp_image x_half = interpolate(x, size=(256, 256), mode='bilinear', align_corners=False) rotate_image = self.two_algo_face_body_rotator(x_half, pose_vector) if args.perf == 'model': print(" - rotator", (time.perf_counter() - tic) * 1000) tic = time.perf_counter() output_image = self.editor(x, interpolate(rotate_image[1], size=(512, 512), mode='bilinear', align_corners=False), interpolate(rotate_image[2], size=(512, 512), mode='bilinear', align_corners=False), pose_vector)[0] if args.perf == 'model': print(" - editor", (time.perf_counter() - tic) * 1000) tic = time.perf_counter() return output_image class TalkingAnime(nn.Module): def __init__(self): super(TalkingAnime, self).__init__() def forward(self, image, mouth_eye_vector, pose_vector): x = image.clone() mouth_eye_morp_image = self.face_morpher(image[:, :, 32:224, 32:224], mouth_eye_vector) x[:, :, 32:224, 32:224] = mouth_eye_morp_image rotate_image = self.two_algo_face_rotator(x, pose_vector)[:2] output_image = self.combiner(rotate_image[0], rotate_image[1], pose_vector) return output_image class TalkingAnime4(nn.Module): """ THA4 model class, compatible with THA3 TalkingAnime3 interface Uses THA4's Siren architecture for inference PyTorch only, float32 precision """ def __init__(self, device): super(TalkingAnime4, self).__init__() if not THA4_AVAILABLE: raise RuntimeError( "THA4 adapter not available. Check tha4_adapter.py" ) # Create THA4 wrapper self.wrapper = THA4Wrapper(device=device) print("THA4 model initialized (float32)") def forward(self, image, mouth_eye_vector, pose_vector, eyebrow_vector, mouth_eye_vector_c, eyebrow_vector_c, ratio=None): """ Forward inference, compatible with THA3 interface Args: image: [batch, 4, 512, 512] mouth_eye_vector: [batch, 27] pose_vector: [batch, 6] eyebrow_vector: [batch, 12] mouth_eye_vector_c: compressed (for caching) eyebrow_vector_c: compressed (for caching) ratio: GPU cache hit ratio Returns: output_image: [batch, 4, 512, 512] """ if args.perf == 'model': tic = time.perf_counter() # Call THA4 wrapper output_image = self.wrapper.forward( image, mouth_eye_vector, pose_vector, eyebrow_vector, mouth_eye_vector_c, eyebrow_vector_c, ratio ) if args.perf == 'model': print(" - tha4_inference", (time.perf_counter() - tic) * 1000) return output_image def to(self, device): """Move model to specified device""" self.wrapper.to(device) return super().to(device)