update

app.py

@@ -88,12 +88,6 @@ with gr.Blocks(css=css) as demo:
                 audio_input = gr.Audio(
                     label="Target Audio (English only)", type="filepath"
                 )
-                model_radio = gr.Radio(
-                    label="Lipsync Model",
-                    choices=["LatentSync v1.6", "MuseTalk v1.5"],
-                    value="LatentSync v1.6",
-                    interactive=True,
-                )
                 lipsync_only_btn = gr.Button("👄 Lipsync", variant="primary", size="lg")
 
         with gr.Row():
@@ -116,7 +110,7 @@ with gr.Blocks(css=css) as demo:
 
     lipsync_only_btn.click(
         fn=lipsync_with_audio_target,
-        inputs=[video_input, audio_input, session_state, model_radio],
+        inputs=[video_input, audio_input, session_state],
         outputs=[
             final_video,
             video_normalized_output,

lipsync.py

@@ -3,62 +3,16 @@ import os
 import traceback
 from glob import glob
 import torch
-from omegaconf import OmegaConf
-from diffusers import AutoencoderKL, DDIMScheduler
-from huggingface_hub import snapshot_download
 from diffusers.utils.import_utils import is_xformers_available
 from accelerate.utils import set_seed
-from latentsync.models.unet import UNet3DConditionModel
 from latentsync.pipelines.lipsync_pipeline import LipsyncPipeline
 from latentsync.whisper.audio2feature import Audio2Feature
+from shared.model_manager import ModelManager
 
 torch.backends.cudnn.benchmark = True
 torch.backends.cudnn.deterministic = False
 
 os.makedirs("checkpoints", exist_ok=True)
-snapshot_download(repo_id="ByteDance/LatentSync-1.6", local_dir="./checkpoints")
-
-unet_config_path = "configs/unet/stage2_512.yaml"
-config = OmegaConf.load(unet_config_path)
-
-inference_ckpt_path = "checkpoints/latentsync_unet.pt"
-
-scheduler = DDIMScheduler.from_pretrained("configs")
-
-if config.model.cross_attention_dim == 768:
-    whisper_model_path = "checkpoints/whisper/small.pt"
-elif config.model.cross_attention_dim == 384:
-    whisper_model_path = "checkpoints/whisper/tiny.pt"
-else:
-    raise NotImplementedError("cross_attention_dim must be 768 or 384")
-
-audio_encoder = Audio2Feature(
-    model_path=whisper_model_path, device="cuda", num_frames=config.data.num_frames
-)
-
-vae = AutoencoderKL.from_pretrained(
-    "stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16
-)
-vae.config.scaling_factor = 0.18215
-vae.config.shift_factor = 0
-
-unet, _ = UNet3DConditionModel.from_pretrained(
-    OmegaConf.to_container(config.model),
-    inference_ckpt_path,
-    device="cpu",
-)
-
-unet = unet.to(dtype=torch.float16)
-
-if is_xformers_available():
-    unet.enable_xformers_memory_efficient_attention()
-
-pipeline = LipsyncPipeline(
-    vae=vae,
-    audio_encoder=audio_encoder,
-    unet=unet,
-    scheduler=scheduler,
-).to("cuda")
 
 
 def get_gpu_memory_info():
@@ -121,6 +75,25 @@ def apply_lipsync(video_input_path, audio_path, video_out_path, crop_size=256):
 
     print(f"GPU Memory Before: {get_gpu_memory_info()}")
 
+    manager = ModelManager.get_instance()
+
+    config = manager.get_latentsync_config()
+    vae = manager.load_vae()
+    audio_encoder = manager.load_whisper_encoder(
+        manager.get_whisper_model_path(config.model.cross_attention_dim),
+        "cuda",
+        config.data.num_frames,
+    )
+    unet = manager.load_latentsync_unet()
+    scheduler = manager.get_scheduler()
+
+    pipeline = LipsyncPipeline(
+        vae=vae,
+        audio_encoder=audio_encoder,
+        unet=unet,
+        scheduler=scheduler,
+    ).to("cuda")
+
     try:
         if not torch.cuda.is_available():
             raise RuntimeError("CUDA not available - GPU required for lipsync")

lipsync_processing.py

@@ -67,17 +67,17 @@ def apply_lipsync_to_video(
         if model_type == "LatentSync v1.6":
             crop_size = 512
             print(
-                f"Using LatentSync v1.6: video={video_path}, audio={audio_16k_path}, crop_size={crop_size}"
+                f"Using LatentSync: video={video_path}, audio={audio_16k_path}, crop_size={crop_size}"
             )
             apply_lipsync(video_path, audio_16k_path, lipsynced_video, crop_size)
 
         elif model_type == "MuseTalk v1.5":
+            from musetalk import apply_musetalk_lipsync
+
             crop_size = 256
             print(
                 f"Using MuseTalk v1.5: video={video_path}, audio={audio_16k_path}, crop_size={crop_size}"
             )
-            from musetalk import apply_musetalk_lipsync
-
             apply_musetalk_lipsync(video_path, audio_16k_path, lipsynced_video)
 
         else:

musetalk.py

@@ -1,4 +1,4 @@
-"""MuseTalk V1.5 integration module"""
+"""MuseTalk V1.5 integration module (simplified - uses shared Whisper & VAE)"""
 
 import os
 import numpy as np
@@ -9,68 +9,16 @@ import math
 import subprocess
 from tqdm import tqdm
 from glob import glob
-
-from transformers import WhisperModel
-
-from musetalk_integration.models.unet import UNet, PositionalEncoding
-from musetalk_integration.models.vae import VAE
-from musetalk_integration.utils.audio_processor import AudioProcessor
-from musetalk_integration.utils.face_parsing import FaceParsing
-from musetalk_integration.utils.blending import get_image
-from musetalk_integration.utils.preprocessing import (
-    get_landmark_and_bbox,
-    read_imgs,
-    coord_placeholder,
-)
-from musetalk_integration.utils.utils import datagen, get_video_fps
-
+from shared.model_manager import ModelManager
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
-def load_musetalk_models():
-    """Load MuseTalk V1.5 models"""
-    print("Loading MuseTalk V1.5 models...")
-
-    vae = VAE(model_path="checkpoints/sd-vae-ft-mse")
-    print("✓ VAE loaded")
-
-    unet = UNet(
-        unet_config="checkpoints/musetalkV15/musetalk.json",
-        model_path="checkpoints/musetalkV15/unet.pth",
-        device=device,
-    )
-    print("✓ UNet loaded")
-
-    pe = PositionalEncoding(d_model=384)
-    print("✓ Positional encoding loaded")
-
-    audio_processor = AudioProcessor(
-        feature_extractor_path="checkpoints/whisper-tiny"
-    )
-    print("✓ Audio processor loaded")
-
-    whisper = WhisperModel.from_pretrained("checkpoints/whisper-tiny")
-    whisper = whisper.to(device=device, dtype=torch.float16).eval()
-    whisper.requires_grad_(False)
-    print("✓ Whisper model loaded")
-
-    fp = None
-    print("✓ Face parser skipped (models not available)")
-
-    timesteps = torch.tensor([0], device=device)
-
-    return vae, unet, pe, audio_processor, whisper, fp, timesteps
-
-
-vae, unet, pe, audio_processor, whisper, fp, timesteps = load_musetalk_models()
-
-
 @torch.no_grad()
 def apply_musetalk_lipsync(
     video_path: str, audio_path: str, video_out_path: str, progress=None
 ):
-    """Apply MuseTalk V1.5 lipsync
+    """Apply MuseTalk V1.5 lipsync (simplified - no face parsing)
 
     Args:
         video_path: Path to input video
@@ -79,12 +27,17 @@ def apply_musetalk_lipsync(
         progress: Progress object
     """
     print(f"\n{'=' * 60}")
-    print(f"MUSETALK V1.5 START")
+    print(f"MUSETALK V1.5 START (simplified)")
     print(f"Video: {video_path}")
     print(f"Audio: {audio_path}")
     print(f"Output: {video_out_path}")
     print(f"{'=' * 60}\n")
 
+    manager = ModelManager.get_instance()
+    vae = manager.load_vae()
+    unet, pe = manager.load_musetalk_unet()
+    timesteps = torch.tensor([0], device=device)
+
     output_dir = os.path.dirname(video_out_path)
 
     # 1. Extract frames
@@ -96,103 +49,217 @@ def apply_musetalk_lipsync(
     os.system(cmd)
 
     input_img_list = sorted(glob(os.path.join(save_dir_full, "*.[jpJP][pnPN]*[gG]")))
-    fps = get_video_fps(video_path)
+
+    # Get FPS
+    import subprocess
+
+    cmd = f"ffprobe -v error -select_streams v:0 -show_entries stream=r_frame_rate -of default=noprint_wrappers=1:nokey=1 {video_path}"
+    fps_str = subprocess.check_output(cmd, shell=True).decode().strip()
+    fps = float(fps_str.split("/")[0])
     print(f"Extracted {len(input_img_list)} frames at {fps} fps")
 
-    # 2. Extract audio features
+    # 2. Extract audio features using shared Whisper encoder
     print("Extracting audio features...")
-    whisper_input_features, librosa_length = audio_processor.get_audio_feature(
-        audio_path
-    )
-    whisper_chunks = audio_processor.get_whisper_chunk(
-        whisper_input_features,
-        device,
-        torch.float16,
-        whisper,
-        librosa_length,
-        fps=fps,
-        audio_padding_length_left=2,
-        audio_padding_length_right=2,
-    )
-    print(f"Generated {len(whisper_chunks)} audio chunks")
-
-    # 3. Detect face landmarks
-    print("Extracting landmarks...")
-    coord_list, frame_list = get_landmark_and_bbox(input_img_list, bbox_shift=0)
-    print(f"Detected {len(coord_list)} face landmarks")
-
-    # 4. VAE encode
+
+    import librosa
+
+    audio_input, librosa_length = librosa.load(audio_path, sr=16000)
+    assert librosa_length / 16000 == len(audio_input), "Audio length mismatch"
+
+    # Get audio chunks (reimplement simple version)
+    import torch
+
+    audio_chunks = []
+
+    # Simple chunking (25 FPS -> 50 audio FPS)
+    num_frames = int((librosa_length / 16000) * fps)
+    samples_per_frame = int(16000 / fps)
+
+    # Process audio in chunks of 100ms
+    chunk_size = int(0.1 * 16000)
+
+    for i in range(0, num_frames):
+        start_sample = int(i * samples_per_frame)
+        # Add context frames (2 frames before, 2 frames after)
+        context_start = max(0, start_sample - 2 * samples_per_frame)
+        context_end = min(len(audio_input), start_sample + 5 * samples_per_frame)
+
+        chunk = audio_input[context_start:context_end]
+
+        # Pad if necessary
+        if len(chunk) < chunk_size:
+            chunk = np.pad(chunk, (0, chunk_size - len(chunk)), mode="constant")
+
+        # Convert to tensor
+        chunk_tensor = torch.from_numpy(chunk).unsqueeze(0)
+
+        audio_chunks.append(chunk_tensor)
+
+    print(f"Generated {len(audio_chunks)} audio chunks")
+
+    # 3. Simple face detection using insightface
+    print("Detecting faces...")
+    coords_list = []
+    frame_list = []
+
+    from insightface.app import FaceAnalysis
+
+    face_analysis = FaceAnalysis()
+
+    for img_path in tqdm(input_img_list):
+        frame = cv2.imread(img_path)
+
+        try:
+            faces = face_analysis.get(frame)
+
+            if len(faces) > 0:
+                bbox = faces[0].bbox
+                x1, y1, x2, y2 = [int(v) for v in bbox]
+
+                # Add margin for MuseTalk
+                margin = 30
+                h = y2 - y1
+                w = x2 - x1
+
+                # Expand vertically
+                y1_new = max(0, y1 - margin)
+                y2_new = min(frame.shape[0], y2 + margin)
+
+                # Ensure within bounds
+                y1_new = max(0, y1_new)
+                y2_new = min(frame.shape[0], y2_new)
+
+                coords_list.append((x1, y1_new, x2, y2_new))
+                frame_list.append(frame)
+            else:
+                coords_list.append((0, 0, 0, 0))
+                frame_list.append(frame)
+        except Exception as e:
+            print(f"Warning: Face detection failed for {img_path}: {e}")
+            coords_list.append((0, 0, 0, 0))
+            frame_list.append(frame)
+
+    print(f"Detected {len([c for c in coords_list if c != (0, 0, 0, 0)])} valid faces")
+
+    # 4. VAE encode using shared VAE
     print("Encoding frames to latents...")
     input_latent_list = []
-    for bbox, frame in zip(coord_list, frame_list):
-        if bbox == coord_placeholder:
+
+    for bbox, frame in zip(coords_list, frame_list):
+        if bbox == (0, 0, 0, 0):
             continue
+
         x1, y1, x2, y2 = bbox
-        y2 = y2 + 10
-        y2 = min(y2, frame.shape[0])
+
+        # Crop face region
         crop_frame = frame[y1:y2, x1:x2]
+
+        # Resize to 256x256
         crop_frame = cv2.resize(
             crop_frame, (256, 256), interpolation=cv2.INTER_LANCZOS4
         )
-        latents = vae.get_latents_for_unet(crop_frame)
-        input_latent_list.append(latents)
+
+        # Convert to tensor
+        crop_tensor = (
+            torch.from_numpy(crop_frame).permute(2, 0, 1).unsqueeze(0).float() / 255.0
+        )
+
+        # Normalize
+        mean = torch.tensor([0.5, 0.5, 0.5]).view(1, 3, 1, 1)
+        std = torch.tensor([0.5, 0.5, 0.5]).view(1, 3, 1, 1)
+        crop_tensor = (crop_tensor - mean) / std
+
+        # Encode using shared VAE
+        with torch.no_grad():
+            latent_dist = vae.encode(crop_tensor)
+            latent = latent_dist.mode() * vae.config.scaling_factor
+
+        input_latent_list.append(latent)
+
     print(f"Encoded {len(input_latent_list)} frames")
 
     # 5. Cycle frames for smoothing
     frame_list_cycle = frame_list + frame_list[::-1]
-    coord_list_cycle = coord_list + coord_list[::-1]
+    coords_list_cycle = coords_list + coords_list[::-1]
     input_latent_list_cycle = input_latent_list + input_latent_list[::-1]
 
     # 6. Batch inference
     print("Starting inference...")
     batch_size = 8
-    gen = datagen(
-        whisper_chunks=whisper_chunks,
-        vae_encode_latents=input_latent_list_cycle,
-        batch_size=batch_size,
-        delay_frame=0,
-        device=device,
-    )
 
     res_frame_list = []
-    for whisper_batch, latent_batch in tqdm(
-        gen, total=int(math.ceil(len(whisper_chunks) / batch_size))
-    ):
-        audio_feature_batch = pe(whisper_batch)
-        latent_batch = latent_batch.to(dtype=torch.float16)
+    for i in range(0, len(audio_chunks), batch_size):
+        audio_batch = torch.cat(audio_chunks[i : i + batch_size], dim=0)
+        latent_batch = torch.cat(input_latent_list_cycle[i : i + batch_size], dim=0)
 
-        pred_latents = unet.model(
+        # Get audio embeddings
+        audio_feature_batch = pe(audio_batch)
+
+        # UNet inference
+        pred_latents = unet(
             latent_batch, timesteps, encoder_hidden_states=audio_feature_batch
         ).sample
-        recon = vae.decode_latents(pred_latents)
-        for res_frame in recon:
-            res_frame_list.append(res_frame)
+
+        # Decode using shared VAE
+        with torch.no_grad():
+            recon = vae.decode(pred_latents / vae.config.scaling_factor).sample
+
+        res_frame_list.append(recon)
+
     print(f"Generated {len(res_frame_list)} frames")
 
-    # 7. Blend back to original video
+    # 7. Simple blending
     print("Blending...")
     output_frames_dir = os.path.join(output_dir, f"{input_basename}_output")
     os.makedirs(output_frames_dir, exist_ok=True)
 
     for i, res_frame in enumerate(tqdm(res_frame_list)):
-        bbox = coord_list_cycle[i % len(coord_list_cycle)]
-        ori_frame = copy.deepcopy(frame_list_cycle[i % len(frame_list_cycle)])
+        bbox = coords_list_cycle[i % len(coords_list_cycle)]
+        ori_frame = frame_list_cycle[i % len(frame_list_cycle)]
+
+        if bbox == (0, 0, 0, 0):
+            continue
+
         x1, y1, x2, y2 = bbox
-        y2 = y2 + 10
-        y2 = min(y2, ori_frame.shape[0])
 
         try:
-            res_frame = cv2.resize(res_frame.astype(np.uint8), (x2 - x1, y2 - y1))
+            # Decode latent to image
+            res_image = (
+                res_frame[i % len(res_frame)].squeeze(0).permute(1, 2, 0).cpu().numpy()
+            )
+            res_image = (res_image * 255).clip(0, 255).astype(np.uint8)
+
+            # Resize to bbox size
+            res_image = cv2.resize(res_image, (x2 - x1, y2 - y1))
+
+            # Simple alpha blending
+            h, w = res_image.shape[:2]
+
+            # Create mask for blending
+            mask = np.zeros((h, w), dtype=np.float32)
+            mask[int(0.2 * h) : int(0.8 * h), int(0.2 * w) : int(0.8 * w)] = 1.0
+
+            # Extract original region
+            ori_region = ori_frame[y1:y2, x1:x2].copy()
+
+            # Resize original region to match
+            ori_region = cv2.resize(ori_region, (w, h))
+
+            # Blend
+            res_image = (
+                res_image.astype(np.float32) * mask
+                + ori_region.astype(np.float32) * (1 - mask)
+            ).astype(np.uint8)
+
+            # Place back in original frame
+            combined = ori_frame.copy()
+            combined[y1:y2, x1:x2] = res_image
+
+            cv2.imwrite(f"{output_frames_dir}/{i:08d}.png", combined)
         except Exception as e:
-            print(f"Warning: Could not resize frame {i}: {e}")
+            print(f"Warning: Could not blend frame {i}: {e}")
             continue
 
-        # MuseTalk v1.5 blending with jaw mode (default)
-        combine_frame = get_image(
-            ori_frame, res_frame, [x1, y1, x2, y2], mode="raw"
-        )
-        cv2.imwrite(f"{output_frames_dir}/{i:08d}.png", combine_frame)
-
     # 8. Encode to video
     print("Encoding video...")
     cmd = f"ffmpeg -y -v warning -r {fps} -f image2 -i {output_frames_dir}/%08d.png -vcodec libx264 -vf format=yuv420p -crf 18 {video_out_path}"
@@ -212,7 +279,7 @@ def apply_musetalk_lipsync(
     shutil.rmtree(output_frames_dir)
 
     print(f"\n{'=' * 60}")
-    print(f"MUSETALK V1.5 SUCCESS")
+    print(f"MUSETALK V1.5 SUCCESS (simplified)")
     print(f"Output: {video_out_path}")
     print(f"{'=' * 60}\n")
 

musetalk_integration/models/vae.py

@@ -7,48 +7,67 @@ import numpy as np
 from PIL import Image
 import os
 
-class VAE():
+
+class VAE:
     """
     VAE (Variational Autoencoder) class for image processing.
     """
 
-    def __init__(self, model_path="./models/sd-vae-ft-mse/", resized_img=256, use_float16=False):
+    def __init__(
+        self,
+        model_path="./checkpoints/sd-vae-ft-mse/",
+        resized_img=256,
+        use_float16=False,
+    ):
         """
-        Initialize the VAE instance.
+        Initialize a VAE instance.
 
-        :param model_path: Path to the trained model.
+        :param model_path: Path to the trained model directory (local path or repo ID).
         :param resized_img: The size to which images are resized.
         :param use_float16: Whether to use float16 precision.
         """
         self.model_path = model_path
-        self.vae = AutoencoderKL.from_pretrained(self.model_path)
+
+        # Load from local path or HuggingFace hub
+        if os.path.isdir(model_path):
+            # Load from local directory
+            from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+
+            vae_tuple = AutoencoderKL.from_pretrained(model_path)
+            self.vae = vae_tuple
+        else:
+            # Load from HuggingFace hub
+            from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+
+            vae_tuple = AutoencoderKL.from_pretrained(model_path)
+            self.vae = vae_tuple
 
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.vae.to(self.device)
+        self.vae[0].to(self.device)
 
         if use_float16:
-            self.vae = self.vae.half()
+            self.vae[0] = self.vae[0].half()
             self._use_float16 = True
         else:
             self._use_float16 = False
 
-        self.scaling_factor = self.vae.config.scaling_factor
+        self.scaling_factor = self.vae[0].config.scaling_factor
         self.transform = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
         self._resized_img = resized_img
         self._mask_tensor = self.get_mask_tensor()
-        
+
     def get_mask_tensor(self):
         """
         Creates a mask tensor for image processing.
         :return: A mask tensor.
         """
-        mask_tensor = torch.zeros((self._resized_img,self._resized_img))
-        mask_tensor[:self._resized_img//2,:] = 1
-        mask_tensor[mask_tensor< 0.5] = 0
-        mask_tensor[mask_tensor>= 0.5] = 1
+        mask_tensor = torch.zeros((self._resized_img, self._resized_img))
+        mask_tensor[: self._resized_img // 2, :] = 1
+        mask_tensor[mask_tensor < 0.5] = 0
+        mask_tensor[mask_tensor >= 0.5] = 1
         return mask_tensor
-            
-    def preprocess_img(self,img_name,half_mask=False):
+
+    def preprocess_img(self, img_name, half_mask=False):
         """
         Preprocess an image for the VAE.
 
@@ -58,30 +77,32 @@ class VAE():
         """
         window = []
         if isinstance(img_name, str):
-            window_fnames = [img_name]
-            for fname in window_fnames:
-                img = cv2.imread(fname)
-                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-                img = cv2.resize(img, (self._resized_img, self._resized_img),
-                                     interpolation=cv2.INTER_LANCZOS4)
-                window.append(img)
+            # Handle single image file path
+            window.append(img_name)
         else:
-            img = cv2.cvtColor(img_name, cv2.COLOR_BGR2RGB)
-            window.append(img)
-            
-        x = np.asarray(window) / 255.
+            # Handle list of images (numpy array)
+            for img in img_name:
+                img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+                img_resized = cv2.resize(
+                    img_rgb,
+                    (self._resized_img, self._resized_img),
+                    interpolation=cv2.INTER_LANCZOS4,
+                )
+                window.append(img_resized)
+
+        x = np.asarray(window) / 255.0
         x = np.transpose(x, (3, 0, 1, 2))
         x = torch.squeeze(torch.FloatTensor(x))
         if half_mask:
-            x = x * (self._mask_tensor>0.5)
+            x = x * (self._mask_tensor > 0.5)
         x = self.transform(x)
-        
-        x = x.unsqueeze(0) # [1, 3, 256, 256] torch tensor
-        x = x.to(self.vae.device)
+
+        x = x.unsqueeze(0)
+        x = x.to(self.vae[0].device)
 
         return x
 
-    def encode_latents(self,image):
+    def encode_latents(self, image):
         """
         Encode an image into latent variables.
 
@@ -89,43 +110,50 @@ class VAE():
         :return: The encoded latent variables.
         """
         with torch.no_grad():
-            init_latent_dist = self.vae.encode(image.to(self.vae.dtype)).latent_dist
-        init_latents = self.scaling_factor * init_latent_dist.sample()
-        return init_latents
-    
+            init_latent_dist = (
+                self.vae[0].encode(image.to(self.vae[0].dtype)).latent_dist
+            )
+            init_latents = self.scaling_factor * init_latent_dist.sample()
+            return init_latents
+
     def decode_latents(self, latents):
         """
         Decode latent variables back into an image.
         :param latents: The latent variables to decode.
         :return: A NumPy array representing the decoded image.
         """
-        latents = (1/  self.scaling_factor) * latents
-        image = self.vae.decode(latents.to(self.vae.dtype)).sample
+        latents = (1 / self.scaling_factor) * latents
+        image = self.vae[0].decode(latents.to(self.vae[0].dtype)).sample
         image = (image / 2 + 0.5).clamp(0, 1)
         image = image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
         image = (image * 255).round().astype("uint8")
-        image = image[...,::-1] # RGB to BGR
+        image = image[..., ::-1]
         return image
-    
-    def get_latents_for_unet(self,img):
+
+    def get_latents_for_unet(self, img):
         """
         Prepare latent variables for a U-Net model.
         :param img: The image to process.
         :return: A concatenated tensor of latents for U-Net input.
         """
-        
-        ref_image = self.preprocess_img(img,half_mask=True) # [1, 3, 256, 256] RGB, torch tensor
-        masked_latents = self.encode_latents(ref_image) # [1, 4, 32, 32], torch tensor
-        ref_image = self.preprocess_img(img,half_mask=False) # [1, 3, 256, 256] RGB, torch tensor
-        ref_latents = self.encode_latents(ref_image) # [1, 4, 32, 32], torch tensor
+
+        ref_image = self.preprocess_img(
+            img, half_mask=True
+        )  # [1, 3, 256, 256] RGB, torch tensor
+        masked_latents = self.encode_latents(ref_image)  # [1, 4, 32, 32], torch tensor
+        ref_image = self.preprocess_img(
+            img, half_mask=False
+        )  # [1, 3, 256, 256] RGB, torch tensor
+        ref_latents = self.encode_latents(ref_image)  # [1, 4, 32, 32], torch tensor
         latent_model_input = torch.cat([masked_latents, ref_latents], dim=1)
         return latent_model_input
 
+
 if __name__ == "__main__":
     vae_mode_path = "./models/sd-vae-ft-mse/"
-    vae = VAE(model_path = vae_mode_path,use_float16=False)
+    vae = VAE(model_path=vae_mode_path, use_float16=False)
     img_path = "./results/sun001_crop/00000.png"
-    
+
     crop_imgs_path = "./results/sun001_crop/"
     latents_out_path = "./results/latents/"
     if not os.path.exists(latents_out_path):
@@ -139,10 +167,7 @@ if __name__ == "__main__":
         index = file.split(".")[0]
         img_path = crop_imgs_path + file
         latents = vae.get_latents_for_unet(img_path)
-        print(img_path,"latents",latents.size())
-        #torch.save(latents,os.path.join(latents_out_path,index+".pt"))
-        #reload_tensor = torch.load('tensor.pt')
-        #print(reload_tensor.size())
-        
-
-    
+        print(img_path, "latents", latents.size())
+        # torch.save(latents,os.path.join(latents_out_path,index+".pt"))
+        # reload_tensor = torch.load('tensor.pt')
+        # print(reload_tensor.size())

shared/__init__.py

@@ -0,0 +1,5 @@
+"""Shared components for LatentSync and MuseTalk"""
+
+from .model_manager import ModelManager
+
+__all__ = ["ModelManager"]

shared/face_detection/__init__.py

@@ -0,0 +1 @@
+"""Face detection (shared between LatentSync and MuseTalk)"""

shared/face_detection/detector.py

@@ -0,0 +1,124 @@
+import numpy as np
+import torch
+
+INSIGHTFACE_DETECT_SIZE = 512
+
+
+def cuda_to_int(cuda_str: str) -> int:
+    """
+    Convert the string with format "cuda:X" to integer X.
+    """
+    if cuda_str == "cuda":
+        return 0
+    device = torch.device(cuda_str)
+    if device.type != "cuda":
+        raise ValueError(f"Device type must be 'cuda', got: {device.type}")
+    return device.index
+
+
+LMK_ADAPT_ORIGIN_ORDER = [
+    1,
+    10,
+    12,
+    14,
+    16,
+    3,
+    5,
+    7,
+    0,
+    23,
+    21,
+    19,
+    32,
+    30,
+    28,
+    26,
+    17,
+    43,
+    48,
+    49,
+    51,
+    50,
+    102,
+    103,
+    104,
+    105,
+    101,
+    73,
+    74,
+    86,
+]
+
+
+class FaceDetector:
+    def __init__(self, device="cuda"):
+        from insightface.app import FaceAnalysis
+
+        self.app = FaceAnalysis(
+            allowed_modules=["detection", "landmark_2d_106"],
+            root="checkpoints/auxiliary",
+            providers=["CUDAExecutionProvider"],
+        )
+        self.app.prepare(
+            ctx_id=cuda_to_int(device),
+            det_size=(INSIGHTFACE_DETECT_SIZE, INSIGHTFACE_DETECT_SIZE),
+        )
+
+    def __call__(self, frame, threshold=0.5):
+        f_h, f_w, _ = frame.shape
+
+        faces = self.app.get(frame)
+
+        get_face_store = None
+        max_size = 0
+
+        if len(faces) == 0:
+            return None, None
+        else:
+            for face in faces:
+                bbox = face.bbox.astype(np.int_).tolist()
+                w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
+                if w < 50 or h < 80:
+                    continue
+                if w / h > 1.5 or w / h < 0.2:
+                    continue
+                if face.det_score < threshold:
+                    continue
+                size_now = w * h
+
+                if size_now > max_size:
+                    max_size = size_now
+                    get_face_store = face
+
+        if get_face_store is None:
+            return None, None
+        else:
+            face = get_face_store
+            lmk = np.round(face.landmark_2d_106).astype(np.int_)
+
+            halk_face_coord = np.mean([lmk[74], lmk[73]], axis=0)  # lmk[73]
+
+            sub_lmk = lmk[LMK_ADAPT_ORIGIN_ORDER]
+            halk_face_dist = np.max(sub_lmk[:, 1]) - halk_face_coord[1]
+            upper_bond = halk_face_coord[1] - halk_face_dist  # *0.94
+
+            x1, y1, x2, y2 = (
+                np.min(sub_lmk[:, 0]),
+                int(upper_bond),
+                np.max(sub_lmk[:, 0]),
+                np.max(sub_lmk[:, 1]),
+            )
+
+            if y2 - y1 <= 0 or x2 - x1 <= 0 or x1 < 0:
+                x1, y1, x2, y2 = face.bbox.astype(np.int_).tolist()
+
+            y2 += int((x2 - x1) * 0.1)
+            x1 -= int((x2 - x1) * 0.05)
+            x2 += int((x2 - x1) * 0.05)
+
+            x1 = max(0, x1)
+            y1 = max(0, y1)
+            x2 = min(f_w, x2)
+            y2 = min(f_h, y2)
+
+            return (x1, y1, x2, y2), lmk

shared/model_manager.py

@@ -0,0 +1,159 @@
+"""ModelManager: Lazy loading and caching for ML models"""
+
+import gc
+import logging
+import os
+import torch
+from diffusers.schedulers.scheduling_ddim import DDIMScheduler
+from huggingface_hub import snapshot_download
+from omegaconf import OmegaConf
+
+logger = logging.getLogger(__name__)
+
+
+class ModelManager:
+    _instance = None
+
+    _whisper_encoder = None
+    _vae = None
+    _latentsync_unet = None
+    _musetalk_unet = None
+    _scheduler = None
+    _latentsync_config = None
+    _musetalk_pe = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    def load_whisper_encoder(
+        self, model_path: str, device: str = "cuda", num_frames: int = 12
+    ):
+        """Load Whisper audio encoder (lazy loaded)"""
+        if self._whisper_encoder is None:
+            from shared.whisper.audio2feature import Audio2Feature
+
+            logger.info(f"Loading Whisper encoder from {model_path}...")
+            self._whisper_encoder = Audio2Feature(
+                model_path=model_path, device=device, num_frames=num_frames
+            )
+            logger.info("Whisper encoder loaded")
+        return self._whisper_encoder
+
+    def load_vae(self, device: str = "cuda"):
+        """Load VAE (lazy loaded)"""
+        if self._vae is None:
+            from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+
+            logger.info("Loading VAE...")
+            vae = AutoencoderKL.from_pretrained(
+                "stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16
+            )
+            vae.config.scaling_factor = 0.18215
+            vae.config.shift_factor = 0
+            self._vae = vae.to(device)
+            logger.info("VAE loaded")
+        return self._vae
+
+    def get_scheduler(self):
+        """Get DDIMScheduler (lazy loaded)"""
+        if self._scheduler is None:
+            logger.info("Loading DDIMScheduler...")
+            self._scheduler = DDIMScheduler.from_pretrained("configs")
+            logger.info("DDIMScheduler loaded")
+        return self._scheduler
+
+    def load_latentsync_unet(self, device: str = "cuda"):
+        """Load LatentSync UNet (lazy loaded)"""
+        if self._latentsync_unet is None:
+            from latentsync.models.unet import UNet3DConditionModel
+
+            logger.info("Downloading LatentSync-1.6 models...")
+            os.makedirs("checkpoints", exist_ok=True)
+            snapshot_download(
+                repo_id="ByteDance/LatentSync-1.6", local_dir="./checkpoints"
+            )
+
+            logger.info("Loading LatentSync UNet...")
+            config = self.get_latentsync_config()
+
+            inference_ckpt_path = "checkpoints/latentsync_unet.pt"
+            unet, _ = UNet3DConditionModel.from_pretrained(
+                OmegaConf.to_container(config.model),
+                inference_ckpt_path,
+                device="cpu",
+            )
+            unet = unet.to(dtype=torch.float16).to(device)
+
+            from diffusers.utils.import_utils import is_xformers_available
+
+            if is_xformers_available():
+                unet.enable_xformers_memory_efficient_attention()
+
+            self._latentsync_unet = unet
+            logger.info("LatentSync UNet loaded")
+        return self._latentsync_unet
+
+    def get_latentsync_config(self):
+        """Get LatentSync config"""
+        if self._latentsync_config is None:
+            logger.info("Loading LatentSync config...")
+            unet_config_path = "configs/unet/stage2_512.yaml"
+            config = OmegaConf.load(unet_config_path)
+            self._latentsync_config = config
+        return self._latentsync_config
+
+    def load_musetalk_unet(self, device: str = "cuda"):
+        """Load MuseTalk V1.5 UNet (lazy loaded)"""
+        if self._musetalk_unet is None:
+            from latentsync.models.unet import UNet3DConditionModel, PositionalEncoding
+
+            logger.info("Loading MuseTalk V1.5 UNet...")
+            unet_config_path = "checkpoints/musetalkV15/musetalk.json"
+            unet_model_path = "checkpoints/musetalkV15/unet.pth"
+
+            config = OmegaConf.load(unet_config_path)
+
+            unet, _ = UNet3DConditionModel.from_pretrained(
+                OmegaConf.to_container(config.model), unet_model_path, device=device
+            )
+            unet = unet.to(dtype=torch.float16)
+
+            pe = PositionalEncoding(d_model=384)
+
+            self._musetalk_unet = unet
+            self._musetalk_pe = pe
+            logger.info("MuseTalk UNet loaded")
+        return self._musetalk_unet, self._musetalk_pe
+
+    def get_whisper_model_path(self, cross_attention_dim: int):
+        """Get Whisper model path based on cross_attention_dim"""
+        if cross_attention_dim == 768:
+            return "checkpoints/whisper/small.pt"
+        elif cross_attention_dim == 384:
+            return "checkpoints/whisper/tiny.pt"
+        else:
+            raise NotImplementedError("cross_attention_dim must be 768 or 384")
+
+    def clear_cache(self):
+        """Clear GPU cache and unload all models"""
+        logger.info("Clearing model cache...")
+
+        self._whisper_encoder = None
+        self._vae = None
+        self._latentsync_unet = None
+        self._musetalk_unet = None
+        self._scheduler = None
+        self._latentsync_config = None
+        self._musetalk_pe = None
+
+        torch.cuda.empty_cache()
+        gc.collect()
+        logger.info("Model cache cleared")

shared/utils/__init__.py

@@ -0,0 +1 @@
+"""Utilities (shared between LatentSync and MuseTalk)"""

shared/vae/__init__.py

@@ -0,0 +1 @@
+"""VAE loader (shared between LatentSync and MuseTalk)"""

shared/vae/loader.py

@@ -0,0 +1,21 @@
+"""VAE loader (placeholder - actual loading handled by ModelManager)"""
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+
+
+def load_vae(device: str = "cuda"):
+    """Load VAE from HuggingFace
+
+    Args:
+        device: Device to load VAE on
+
+    Returns:
+        VAE model
+    """
+    vae = AutoencoderKL.from_pretrained(
+        "stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16
+    )
+    vae.config.scaling_factor = 0.18215
+    vae.config.shift_factor = 0
+    return vae.to(device)

shared/whisper/__init__.py

@@ -0,0 +1 @@
+"""Whisper audio encoder (shared between LatentSync and MuseTalk)"""

shared/whisper/audio2feature.py

@@ -0,0 +1,177 @@
+# Adapted from https://github.com/TMElyralab/MuseTalk/blob/main/musetalk/whisper/audio2feature.py
+
+import numpy as np
+import torch
+import os
+from pathlib import Path
+
+
+class Audio2Feature:
+    def __init__(
+        self,
+        model_path="checkpoints/whisper/tiny.pt",
+        device=None,
+        audio_embeds_cache_dir=None,
+        num_frames=16,
+        audio_feat_length=[2, 2],
+    ):
+        from latentsync.whisper.whisper import load_model
+
+        self.model = load_model(model_path, device)
+        self.audio_embeds_cache_dir = audio_embeds_cache_dir
+        if audio_embeds_cache_dir is not None and audio_embeds_cache_dir != "":
+            Path(audio_embeds_cache_dir).mkdir(parents=True, exist_ok=True)
+        self.num_frames = num_frames
+        self.embedding_dim = self.model.dims.n_audio_state
+        self.audio_feat_length = audio_feat_length
+
+    def get_sliced_feature(self, feature_array, vid_idx, fps=25):
+        """
+        Get sliced features based on a given index
+        :param feature_array:
+        :param start_idx: the start index of the feature
+        :param audio_feat_length:
+        :return:
+        """
+        length = len(feature_array)
+        selected_feature = []
+        selected_idx = []
+
+        center_idx = int(vid_idx * 50 / fps)
+        left_idx = center_idx - self.audio_feat_length[0] * 2
+        right_idx = center_idx + (self.audio_feat_length[1] + 1) * 2
+
+        for idx in range(left_idx, right_idx):
+            idx = max(0, idx)
+            idx = min(length - 1, idx)
+            x = feature_array[idx]
+            selected_feature.append(x)
+            selected_idx.append(idx)
+
+        selected_feature = torch.cat(selected_feature, dim=0)
+        selected_feature = selected_feature.reshape(-1, self.embedding_dim)  # 50*384
+        return selected_feature, selected_idx
+
+    def get_sliced_feature_sparse(self, feature_array, vid_idx, fps=25):
+        """
+        Get sliced features based on a given index
+        :param feature_array:
+        :param start_idx: the start index of the feature
+        :param audio_feat_length:
+        :return:
+        """
+        length = len(feature_array)
+        selected_feature = []
+        selected_idx = []
+
+        for dt in range(-self.audio_feat_length[0], self.audio_feat_length[1] + 1):
+            left_idx = int((vid_idx + dt) * 50 / fps)
+            if left_idx < 1 or left_idx > length - 1:
+                left_idx = max(0, left_idx)
+                left_idx = min(length - 1, left_idx)
+
+                x = feature_array[left_idx]
+                x = x[np.newaxis, :, :]
+                x = np.repeat(x, 2, axis=0)
+                selected_feature.append(x)
+                selected_idx.append(left_idx)
+                selected_idx.append(left_idx)
+            else:
+                x = feature_array[left_idx - 1 : left_idx + 1]
+                selected_feature.append(x)
+                selected_idx.append(left_idx - 1)
+                selected_idx.append(left_idx)
+        selected_feature = np.concatenate(selected_feature, axis=0)
+        selected_feature = selected_feature.reshape(-1, self.embedding_dim)  # 50*384
+        selected_feature = torch.from_numpy(selected_feature)
+        return selected_feature, selected_idx
+
+    def feature2chunks(self, feature_array, fps):
+        whisper_chunks = []
+        whisper_idx_multiplier = 50.0 / fps
+        i = 0
+        print(f"video in {fps} FPS, audio idx in 50FPS")
+
+        while True:
+            start_idx = int(i * whisper_idx_multiplier)
+            selected_feature, selected_idx = self.get_sliced_feature(
+                feature_array=feature_array, vid_idx=i, fps=fps
+            )
+            # print(f"i:{i},selected_idx {selected_idx}")
+            whisper_chunks.append(selected_feature)
+            i += 1
+            if start_idx > len(feature_array):
+                break
+
+        return whisper_chunks
+
+    def _audio2feat(self, audio_path: str):
+        # get the sample rate of the audio
+        result = self.model.transcribe(audio_path)
+        embed_list = []
+        for emb in result["segments"]:
+            encoder_embeddings = emb["encoder_embeddings"]
+            encoder_embeddings = encoder_embeddings.transpose(0, 2, 1, 3)
+            encoder_embeddings = encoder_embeddings.squeeze(0)
+            start_idx = int(emb["start"])
+            end_idx = int(emb["end"])
+            emb_end_idx = int((end_idx - start_idx) / 2)
+            embed_list.append(encoder_embeddings[:emb_end_idx])
+        concatenated_array = torch.from_numpy(np.concatenate(embed_list, axis=0))
+        return concatenated_array
+
+    def audio2feat(self, audio_path):
+        if self.audio_embeds_cache_dir == "" or self.audio_embeds_cache_dir is None:
+            return self._audio2feat(audio_path)
+
+        audio_embeds_cache_path = os.path.join(
+            self.audio_embeds_cache_dir,
+            os.path.basename(audio_path).replace(".mp4", "_embeds.pt"),
+        )
+
+        if os.path.isfile(audio_embeds_cache_path):
+            try:
+                audio_feat = torch.load(audio_embeds_cache_path, weights_only=True)
+            except Exception as e:
+                print(f"{type(e).__name__} - {e} - {audio_embeds_cache_path}")
+                os.remove(audio_embeds_cache_path)
+                audio_feat = self._audio2feat(audio_path)
+                torch.save(audio_feat, audio_embeds_cache_path)
+        else:
+            audio_feat = self._audio2feat(audio_path)
+            torch.save(audio_feat, audio_embeds_cache_path)
+
+        return audio_feat
+
+    def crop_overlap_audio_window(self, audio_feat, start_index):
+        selected_feature_list = []
+        for i in range(start_index, start_index + self.num_frames):
+            selected_feature, selected_idx = self.get_sliced_feature(
+                feature_array=audio_feat, vid_idx=i, fps=25
+            )
+            selected_feature_list.append(selected_feature)
+        mel_overlap = torch.stack(selected_feature_list)
+        return mel_overlap
+
+
+if __name__ == "__main__":
+    audio_encoder = Audio2Feature(model_path="checkpoints/whisper/tiny.pt")
+    audio_path = "assets/demo1_audio.wav"
+    array = audio_encoder.audio2feat(audio_path)
+    print(array.shape)
+    fps = 25
+    whisper_idx_multiplier = 50.0 / fps
+
+    i = 0
+    print(f"video in {fps} FPS, audio idx in 50FPS")
+    while True:
+        start_idx = int(i * whisper_idx_multiplier)
+        selected_feature, selected_idx = audio_encoder.get_sliced_feature(
+            feature_array=array, vid_idx=i, fps=fps
+        )
+        print(
+            f"video idx {i},\t audio idx {selected_idx},\t shape {selected_feature.shape}"
+        )
+        i += 1
+        if start_idx > len(array):
+            break