diff --git a/assets/gpt1.mp3 b/assets/gpt1.mp3 new file mode 100644 index 0000000..8ef3aac Binary files /dev/null and b/assets/gpt1.mp3 differ diff --git a/assets/gpt1.mp4 b/assets/gpt1.mp4 new file mode 100644 index 0000000..c1c2367 Binary files /dev/null and b/assets/gpt1.mp4 differ diff --git a/latentsync/pipelines/lipsync_pipeline.py b/latentsync/pipelines/lipsync_pipeline.py index 693155e..29999a0 100644 --- a/latentsync/pipelines/lipsync_pipeline.py +++ b/latentsync/pipelines/lipsync_pipeline.py @@ -1,4 +1,16 @@ -# Adapted from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/pipelines/pipeline_animation.py +# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. import inspect import math @@ -32,7 +44,7 @@ from ..models.unet import UNet3DConditionModel from ..utils.util import read_video, read_audio, write_video, check_ffmpeg_installed -from ..utils.image_processor import ImageProcessor, load_fixed_mask +from ..utils.image_processor import ImageProcessor, load_fixed_mask # load_fixed_mask is still external from ..whisper.audio2feature import Audio2Feature import tqdm import soundfile as sf @@ -233,12 +245,19 @@ def set_progress_bar_config(self, **kwargs): self._progress_bar_config = {} self._progress_bar_config.update(kwargs) + # THIS STATIC METHOD IS CORRECTLY LOCATED AND MODIFIED HERE @staticmethod def paste_surrounding_pixels_back(decoded_latents, pixel_values, masks, device, weight_dtype): - # Paste the surrounding pixels back, because we only want to change the mouth region + # `masks` tensor (passed as 3rd arg) is 0.0 for lips (black in your PNG), 1.0 for face (white in your PNG). + # We want to apply `decoded_latents` where the mask is black (lips). + # We want to preserve `pixel_values` where the mask is white (face). + pixel_values = pixel_values.to(device=device, dtype=weight_dtype) - masks = masks.to(device=device, dtype=weight_dtype) - combined_pixel_values = decoded_latents * masks + pixel_values * (1 - masks) + masks = masks.to(device=device, dtype=weight_dtype) + + # CORRECTED LOGIC: + combined_pixel_values = decoded_latents * (1 - masks) + pixel_values * masks + return combined_pixel_values @staticmethod @@ -249,65 +268,98 @@ def pixel_values_to_images(pixel_values: torch.Tensor): images = images.cpu().numpy() return images - def affine_transform_video(self, video_frames: np.ndarray): + def affine_transform_video(self, video_frames: np.ndarray) -> (List[Optional[torch.Tensor]], List[Optional[list]], List[Optional[np.ndarray]], List[bool]): + """ + Processes video frames to detect faces. If no face is detected in a frame, + it appends None for face, box, and affine_matrix, and False to has_face_map. + Returns: + - faces: List of preprocessed face tensors (or None) + - boxes: List of bounding boxes (or None) + - affine_matrices: List of affine transformation matrices (or None) + - has_face_map: Boolean list, True if face detected, False otherwise + """ faces = [] boxes = [] affine_matrices = [] - print(f"Affine transforming {len(video_frames)} faces...") - for frame in tqdm.tqdm(video_frames): - face, box, affine_matrix = self.image_processor.affine_transform(frame) - faces.append(face) - boxes.append(box) - affine_matrices.append(affine_matrix) - - faces = torch.stack(faces) - return faces, boxes, affine_matrices + has_face_map = [] # True if face detected, False otherwise + + print(f"Detecting faces in {len(video_frames)} frames...") + for i, frame in enumerate(tqdm.tqdm(video_frames)): + try: + face, box, affine_matrix = self.image_processor.affine_transform(frame) + faces.append(face) # face is already a torch.Tensor here + boxes.append(box) + affine_matrices.append(affine_matrix) + has_face_map.append(True) + except RuntimeError as e: + if "Face not detected" in str(e): + faces.append(None) + boxes.append(None) + affine_matrices.append(None) + has_face_map.append(False) + else: + raise e # Re-raise any other unexpected RuntimeErrors + return faces, boxes, affine_matrices, has_face_map def restore_video(self, faces: torch.Tensor, video_frames: np.ndarray, boxes: list, affine_matrices: list): - video_frames = video_frames[: len(faces)] + """ + This is the original restore_video function. It should only be called with + faces that have successfully been lipsynced (i.e., had faces detected and processed). + `faces` here is a batch of (C, resolution, resolution) lipsynced cropped faces. + `video_frames` here is a batch of corresponding original full frames. + `boxes` and `affine_matrices` are corresponding lists of original detected face data. + """ + # Ensure input lengths match - crucial for correct mapping + if faces.shape[0] != video_frames.shape[0] or \ + faces.shape[0] != len(boxes) or \ + faces.shape[0] != len(affine_matrices): + print("Warning: Mismatch in input lengths for restore_video. Proceeding with min length.") + min_len = min(faces.shape[0], video_frames.shape[0], len(boxes), len(affine_matrices)) + faces = faces[:min_len] + video_frames = video_frames[:min_len] + boxes = boxes[:min_len] + affine_matrices = affine_matrices[:min_len] + out_frames = [] - print(f"Restoring {len(faces)} faces...") - for index, face in enumerate(tqdm.tqdm(faces)): + for index, lipsynced_cropped_face_tensor in enumerate(tqdm.tqdm(faces, desc="Restoring lipsynced faces...")): # Renamed `face` to clarify content x1, y1, x2, y2 = boxes[index] - height = int(y2 - y1) - width = int(x2 - x1) - face = torchvision.transforms.functional.resize( - face, size=(height, width), interpolation=transforms.InterpolationMode.BICUBIC, antialias=True + height_orig_face = int(y2 - y1) + width_orig_face = int(x2 - x1) + + # This step resizes the (C, resolution, resolution) lipsynced cropped face + # to the actual bounding box dimensions (C, height_orig_face, width_orig_face). + processed_face_tensor_for_restore_img = torchvision.transforms.functional.resize( + lipsynced_cropped_face_tensor, + size=(height_orig_face, width_orig_face), + interpolation=transforms.InterpolationMode.BICUBIC, + antialias=True + ) + + # self.image_processor.restorer.restore_img blends the processed face region + # back onto the original full frame using the affine matrix. + out_frame = self.image_processor.restorer.restore_img( + video_frames[index], # Original full frame for this index + processed_face_tensor_for_restore_img, # The resized lipsynced face region + affine_matrices[index] # Affine matrix for this face ) - out_frame = self.image_processor.restorer.restore_img(video_frames[index], face, affine_matrices[index]) out_frames.append(out_frame) return np.stack(out_frames, axis=0) def loop_video(self, whisper_chunks: list, video_frames: np.ndarray): - # If the audio is longer than the video, we need to loop the video + # This function remains unchanged from your original. + # It handles looping/trimming video frames to match audio length. if len(whisper_chunks) > len(video_frames): - faces, boxes, affine_matrices = self.affine_transform_video(video_frames) num_loops = math.ceil(len(whisper_chunks) / len(video_frames)) loop_video_frames = [] - loop_faces = [] - loop_boxes = [] - loop_affine_matrices = [] for i in range(num_loops): if i % 2 == 0: loop_video_frames.append(video_frames) - loop_faces.append(faces) - loop_boxes += boxes - loop_affine_matrices += affine_matrices else: loop_video_frames.append(video_frames[::-1]) - loop_faces.append(faces.flip(0)) - loop_boxes += boxes[::-1] - loop_affine_matrices += affine_matrices[::-1] - video_frames = np.concatenate(loop_video_frames, axis=0)[: len(whisper_chunks)] - faces = torch.cat(loop_faces, dim=0)[: len(whisper_chunks)] - boxes = loop_boxes[: len(whisper_chunks)] - affine_matrices = loop_affine_matrices[: len(whisper_chunks)] else: video_frames = video_frames[: len(whisper_chunks)] - faces, boxes, affine_matrices = self.affine_transform_video(video_frames) - - return video_frames, faces, boxes, affine_matrices + return video_frames # Return the effective video frames @torch.no_grad() def __call__( @@ -315,7 +367,7 @@ def __call__( video_path: str, audio_path: str, video_out_path: str, - num_frames: int = 16, + num_frames: int = 16, # This is the batch size for UNet, not number of frames to process video_fps: int = 25, audio_sample_rate: int = 16000, height: Optional[int] = None, @@ -324,7 +376,7 @@ def __call__( guidance_scale: float = 1.5, weight_dtype: Optional[torch.dtype] = torch.float16, eta: float = 0.0, - mask_image_path: str = "latentsync/utils/mask.png", + mask_image_path: str = "latentsync/utils/mask3.png", # Updated default mask path temp_dir: str = "temp", generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, @@ -338,9 +390,10 @@ def __call__( # 0. Define call parameters device = self._execution_device - mask_image = load_fixed_mask(height, mask_image_path) + # The mask_image is loaded here and passed to ImageProcessor + mask_image = load_fixed_mask(height, mask_image_path) # Call global function self.image_processor = ImageProcessor(height, device="cuda", mask_image=mask_image) - self.set_progress_bar_config(desc=f"Sample frames: {num_frames}") + self.set_progress_bar_config(desc=f"Processing video frames...") # 1. Default height and width to unet height = height or self.unet.config.sample_size * self.vae_scale_factor @@ -349,9 +402,6 @@ def __call__( # 2. Check inputs self.check_inputs(height, width, callback_steps) - # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) - # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` - # corresponds to doing no classifier free guidance. do_classifier_free_guidance = guidance_scale > 1.0 # 3. set timesteps @@ -361,105 +411,156 @@ def __call__( # 4. Prepare extra step kwargs. extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + # 5. Read audio and video whisper_feature = self.audio_encoder.audio2feat(audio_path) whisper_chunks = self.audio_encoder.feature2chunks(feature_array=whisper_feature, fps=video_fps) + original_video_frames_raw = read_video(video_path, use_decord=False) - audio_samples = read_audio(audio_path) - video_frames = read_video(video_path, use_decord=False) - - video_frames, faces, boxes, affine_matrices = self.loop_video(whisper_chunks, video_frames) + # Get the effective video frames after looping/trimming based on audio length + effective_video_frames = self.loop_video(whisper_chunks, original_video_frames_raw) + + # Initialize the list for final output frames with copies of the effective_video_frames. + # Frames will be replaced in this list ONLY IF a face is detected and processed. + final_synced_frames_list = [f for f in effective_video_frames] # Make a mutable list copy - synced_video_frames = [] + # Perform face detection for all effective video frames. + # This will return three lists for faces, boxes, and affine matrices, + # where elements are None if no face was found, and also a boolean map. + all_detected_faces, all_detected_boxes, all_detected_affine_matrices, has_face_map = \ + self.affine_transform_video(effective_video_frames) num_channels_latents = self.vae.config.latent_channels - - # Prepare latent variables - all_latents = self.prepare_latents( - len(whisper_chunks), - num_channels_latents, - height, - width, - weight_dtype, - device, - generator, - ) - - num_inferences = math.ceil(len(whisper_chunks) / num_frames) - for i in tqdm.tqdm(range(num_inferences), desc="Doing inference..."): - if self.unet.add_audio_layer: - audio_embeds = torch.stack(whisper_chunks[i * num_frames : (i + 1) * num_frames]) - audio_embeds = audio_embeds.to(device, dtype=weight_dtype) + + # Collect data for frames that actually need processing (i.e., have faces) + # Store tuples of (original_idx, face_tensor, box, affine_matrix, audio_chunk) + data_for_unet_processing = [] + for i in range(len(effective_video_frames)): + if has_face_map[i]: + data_for_unet_processing.append({ + "original_idx": i, + "face": all_detected_faces[i], + "box": all_detected_boxes[i], + "affine_matrix": all_detected_affine_matrices[i], + "audio_chunk": whisper_chunks[i] # Ensure audio chunk aligns + }) + + if not data_for_unet_processing: # Check if the list is empty (no faces detected at all) + print("No faces detected in any frame. Returning original video.") + # Skip all UNet processing, proceed to video/audio writing + synced_video_frames_np = np.stack(final_synced_frames_list, axis=0) # Convert back to numpy array + else: + # Iterate through data_for_unet_processing in batches of `num_frames` + num_batches = math.ceil(len(data_for_unet_processing) / num_frames) + + print(f"Total batches with faces for UNet inference: {num_batches}") + + for batch_num in range(num_batches): + start_idx_batch = batch_num * num_frames + end_idx_batch = min((batch_num + 1) * num_frames, len(data_for_unet_processing)) + current_batch_data = data_for_unet_processing[start_idx_batch:end_idx_batch] + + batch_size_current = len(current_batch_data) # Actual size of this specific batch + + # Collect data for the current batch + batch_original_indices = [item["original_idx"] for item in current_batch_data] + batch_input_faces = torch.stack([item["face"] for item in current_batch_data]) # Stack faces for UNet input + batch_audio_embeds = torch.stack([item["audio_chunk"] for item in current_batch_data]) + batch_boxes = [item["box"] for item in current_batch_data] + batch_affine_matrices = [item["affine_matrix"] for item in current_batch_data] + + # --- START ORIGINAL UNET PROCESSING FLOW FOR A BATCH (Restored) --- + # This section should be as close as possible to how it worked before "no-face" handling. + + # Prepare latents for the current batch + batch_latents = self.prepare_latents( + batch_size_current, # Use actual batch size + num_channels_latents, + height, + width, + weight_dtype, + device, + generator, + ) + + # Audio embeds for UNet + audio_embeds_for_unet = batch_audio_embeds.to(device, dtype=weight_dtype) if do_classifier_free_guidance: - null_audio_embeds = torch.zeros_like(audio_embeds) - audio_embeds = torch.cat([null_audio_embeds, audio_embeds]) - else: - audio_embeds = None - inference_faces = faces[i * num_frames : (i + 1) * num_frames] - latents = all_latents[:, :, i * num_frames : (i + 1) * num_frames] - ref_pixel_values, masked_pixel_values, masks = self.image_processor.prepare_masks_and_masked_images( - inference_faces, affine_transform=False - ) - - # 7. Prepare mask latent variables - mask_latents, masked_image_latents = self.prepare_mask_latents( - masks, - masked_pixel_values, - height, - width, - weight_dtype, - device, - generator, - do_classifier_free_guidance, - ) - - # 8. Prepare image latents - ref_latents = self.prepare_image_latents( - ref_pixel_values, - device, - weight_dtype, - generator, - do_classifier_free_guidance, - ) - - # 9. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for j, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - unet_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents - - unet_input = self.scheduler.scale_model_input(unet_input, t) - - # concat latents, mask, masked_image_latents in the channel dimension - unet_input = torch.cat([unet_input, mask_latents, masked_image_latents, ref_latents], dim=1) - - # predict the noise residual - noise_pred = self.unet(unet_input, t, encoder_hidden_states=audio_embeds).sample - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_audio = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_audio - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - - # call the callback, if provided - if j == len(timesteps) - 1 or ((j + 1) > num_warmup_steps and (j + 1) % self.scheduler.order == 0): - progress_bar.update() - if callback is not None and j % callback_steps == 0: - callback(j, t, latents) - - # Recover the pixel values - decoded_latents = self.decode_latents(latents) - decoded_latents = self.paste_surrounding_pixels_back( - decoded_latents, ref_pixel_values, 1 - masks, device, weight_dtype - ) - synced_video_frames.append(decoded_latents) - - synced_video_frames = self.restore_video(torch.cat(synced_video_frames), video_frames, boxes, affine_matrices) - - audio_samples_remain_length = int(synced_video_frames.shape[0] / video_fps * audio_sample_rate) + null_audio_embeds = torch.zeros_like(audio_embeds_for_unet) + audio_embeds_for_unet = torch.cat([null_audio_embeds, audio_embeds_for_unet]) + + # Mask and Masked Image Latent preparation + # `masks` here will be 0 for lips, 1 for face. + ref_pixel_values, masked_pixel_values, masks = self.image_processor.prepare_masks_and_masked_images( + batch_input_faces, affine_transform=False # batch_input_faces is list of tensors + ) + mask_latents, masked_image_latents = self.prepare_mask_latents( + masks, masked_pixel_values, height, width, weight_dtype, device, generator, do_classifier_free_guidance + ) + + # Reference Image Latent preparation + ref_latents = self.prepare_image_latents( + ref_pixel_values, device, weight_dtype, generator, do_classifier_free_guidance + ) + + # Denoising loop + current_latents = batch_latents + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + with self.progress_bar(total=num_inference_steps) as progress_bar: + for j, t in enumerate(timesteps): + unet_input = torch.cat([current_latents] * 2) if do_classifier_free_guidance else current_latents + unet_input = self.scheduler.scale_model_input(unet_input, t) + unet_input = torch.cat([unet_input, mask_latents, masked_image_latents, ref_latents], dim=1) + + noise_pred = self.unet(unet_input, t, encoder_hidden_states=audio_embeds_for_unet).sample + + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_audio = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_audio - noise_pred_uncond) + + current_latents = self.scheduler.step(noise_pred, t, current_latents, **extra_step_kwargs).prev_sample + + if j == len(timesteps) - 1 or ((j + 1) > num_warmup_steps and (j + 1) % self.scheduler.order == 0): + progress_bar.update() + if callback is not None and j % callback_steps == 0: + callback(j, t, current_latents) + + # Recover the pixel values for the current batch of generated faces + decoded_latents = self.decode_latents(current_latents) + + # Apply the staticmethod `paste_surrounding_pixels_back` from THIS class. + # It now correctly handles your mask convention (0=lips, 1=face) + # This returns the lipsynced cropped face (C, resolution, resolution) + generated_lipsynced_cropped_faces_batch = self.paste_surrounding_pixels_back( + decoded_latents, ref_pixel_values, masks, device, weight_dtype # Pass original `masks` (0=lips, 1=face) + ) + + # --- END ORIGINAL UNET PROCESSING FLOW FOR A BATCH --- + + # Now, use the original `restore_video` method to blend these processed faces + # back onto their respective original full frames. + + # Gather original full frames for this batch + batch_original_full_frames = np.stack([effective_video_frames[idx] for idx in batch_original_indices], axis=0) + + # Call the original restore_video function + restored_full_frames_for_batch = self.restore_video( + generated_lipsynced_cropped_faces_batch, # Faces from UNet (C, resolution, resolution) + batch_original_full_frames, # Corresponding original full frames (N, H, W, C) + batch_boxes, # Corresponding original boxes + batch_affine_matrices # Corresponding original affine matrices + ) + + # Update the final list with the restored full frames + for k, original_idx in enumerate(batch_original_indices): + final_synced_frames_list[original_idx] = restored_full_frames_for_batch[k] + + + # Convert list of frames back to numpy array for video writing + synced_video_frames_np = np.stack(final_synced_frames_list, axis=0) + + # Audio processing (remains unchanged as it's continuous) + audio_samples = read_audio(audio_path) + audio_samples_remain_length = int(synced_video_frames_np.shape[0] / video_fps * audio_sample_rate) audio_samples = audio_samples[:audio_samples_remain_length].cpu().numpy() if is_train: @@ -469,9 +570,9 @@ def __call__( shutil.rmtree(temp_dir) os.makedirs(temp_dir, exist_ok=True) - write_video(os.path.join(temp_dir, "video.mp4"), synced_video_frames, fps=video_fps) + write_video(os.path.join(temp_dir, "video.mp4"), synced_video_frames_np, fps=video_fps) sf.write(os.path.join(temp_dir, "audio.wav"), audio_samples, audio_sample_rate) command = f"ffmpeg -y -loglevel error -nostdin -i {os.path.join(temp_dir, 'video.mp4')} -i {os.path.join(temp_dir, 'audio.wav')} -c:v libx264 -crf 18 -c:a aac -q:v 0 -q:a 0 {video_out_path}" - subprocess.run(command, shell=True) + subprocess.run(command, shell=True) \ No newline at end of file diff --git a/latentsync/whisper/audio2feature.py b/latentsync/whisper/audio2feature.py index c8dc609..6e87b70 100644 --- a/latentsync/whisper/audio2feature.py +++ b/latentsync/whisper/audio2feature.py @@ -10,13 +10,24 @@ class Audio2Feature: def __init__( self, - model_path="checkpoints/whisper/tiny.pt", + # The model_path argument passed here is conceptually for where the model *should* be, + # but the load_model function needs a specific way to find it. + # We'll use this argument to construct the download_root if needed. + model_path="checkpoints/whisper/tiny.pt", # This still acts as a default/reference device=None, audio_embeds_cache_dir=None, num_frames=16, audio_feat_length=[2, 2], ): - self.model = load_model(model_path, device) + # Determine the absolute path to the whisper model directory + # CHECKPOINT_PATH is /checkpoints from modal_app.py + # So, the whisper models are in /checkpoints/whisper/ + whisper_download_root = Path("/checkpoints") / "whisper" + print(f"Loading Whisper model from download_root: {whisper_download_root}") + + # Use the 'download_root' argument of load_model to specify where to find 'tiny.pt' + self.model = load_model("tiny", device=device, download_root=str(whisper_download_root)) + 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) @@ -164,4 +175,4 @@ def crop_overlap_audio_window(self, audio_feat, start_index): print(f"video idx {i},\t audio idx {selected_idx},\t shape {selected_feature.shape}") i += 1 if start_idx > len(array): - break + break \ No newline at end of file diff --git a/modal_app.py b/modal_app.py index 8e7e65a..642deca 100644 --- a/modal_app.py +++ b/modal_app.py @@ -10,56 +10,46 @@ from fastapi import FastAPI, UploadFile, File, Form, Response, HTTPException -# --- App and Volume Setup --- +# --- Modal Setup --- app = modal.App("latentsync-api") -# This is a persistent shared volume to store model checkpoints volume = modal.NetworkFileSystem.from_name("latentsync-checkpoints-vol", create_if_missing=True) # --- Paths --- -# We define remote paths for the code and checkpoints REMOTE_CODE_PATH = Path("/app") CHECKPOINT_PATH = Path("/checkpoints") -# --- Model Download Logic --- -# This function runs once during the image build to download and cache the models +# --- Download Required Models --- def _download_models(): from huggingface_hub import hf_hub_download - # Ensure the target directory exists CHECKPOINT_PATH.mkdir(parents=True, exist_ok=True) (CHECKPOINT_PATH / "whisper").mkdir(exist_ok=True) - # Download the main UNet model hf_hub_download( repo_id="ByteDance/LatentSync-1.6", filename="latentsync_unet.pt", local_dir=CHECKPOINT_PATH, local_dir_use_symlinks=False, ) - # Download the Whisper model hf_hub_download( repo_id="ByteDance/LatentSync-1.6", - filename="whisper/tiny.pt", + filename="whisper/tiny.pt", # Correct path within the repo local_dir=CHECKPOINT_PATH, local_dir_use_symlinks=False, ) - # Download the SyncNet model hf_hub_download( repo_id="ByteDance/LatentSync-1.6", filename="stable_syncnet.pt", local_dir=CHECKPOINT_PATH, local_dir_use_symlinks=False, ) - print("All models downloaded successfully.") + print("✅ All models downloaded.") -# --- Modal Image Definition --- -# This defines the container environment, installing all necessary dependencies +# --- Modal Image Configuration --- latentsync_image = ( modal.Image.debian_slim(python_version="3.10") - .apt_install("ffmpeg") # ffmpeg is essential for video processing - # First, install PyTorch with the specific CUDA version via a direct command + .apt_install("ffmpeg") .run_commands("pip install torch==2.5.1 torchvision==0.20.1 --extra-index-url https://download.pytorch.org/whl/cu121") - # Then, install the rest of the packages .pip_install( "diffusers==0.32.2", "transformers==4.48.0", @@ -87,44 +77,39 @@ def _download_models(): "fastapi", "uvicorn", "python-multipart", - "requests", # Added for downloading files from URLs + "requests" ) - # Run the model download function after dependencies are installed .run_function( _download_models, - network_file_systems={str(CHECKPOINT_PATH): volume}, # Mount the volume and run the download + network_file_systems={str(CHECKPOINT_PATH): volume}, ) - # Set the PYTHONPATH to include the app directory, making local modules importable - .env({"PYTHONPATH": str(REMOTE_CODE_PATH)}) - # Add local code last to optimize build speed on file changes - .add_local_dir(".", remote_path=str(REMOTE_CODE_PATH)) + .env({"PYTHONPATH": "/app"}) + .add_local_dir(".", remote_path="/app") + .add_local_dir("scripts", remote_path="/app/scripts") + .add_local_dir("configs", remote_path="/app/configs") ) -# --- Inference Class --- -# This class encapsulates the model and the inference logic. -# It's decorated with @app.cls to run on a GPU-equipped container on Modal. +# --- GPU Class for Inference --- @app.cls( - gpu="A10G", # A10G has 24GB VRAM, suitable for the 18GB requirement + gpu="A10G", image=latentsync_image, network_file_systems={str(CHECKPOINT_PATH): volume}, - timeout=600, # Set a 10-minute timeout for inference + timeout=1800, ) class LatentSync: @modal.enter() def setup(self): - """ - This method runs once when the container for the class starts. - We change the directory and load the model configuration. - """ + import sys from omegaconf import OmegaConf - + + sys.path.append("/app") # Ensure scripts import works os.chdir(REMOTE_CODE_PATH) - print(f"Current working directory: {os.getcwd()}") - - # Load the configuration file for the model - config_path = REMOTE_CODE_PATH / "configs" / "unet" / "stage2_512.yaml" + + config_path = REMOTE_CODE_PATH / "configs/unet/stage2_512.yaml" self.config = OmegaConf.load(config_path) - print("Model configuration loaded.") + + print("✅ Setup complete. Current working directory:", os.getcwd()) + print("📂 Scripts directory:", os.listdir("/app/scripts")) @modal.method() def generate( @@ -137,19 +122,14 @@ def generate( inference_steps: int, seed: int, ) -> bytes: - """ - The main inference method. It takes paths to video/audio and generation - parameters, runs the lipsync pipeline, and returns the output video as bytes. - """ from scripts.inference import main as inference_main - # Create temporary directories for input and output temp_input_dir = Path("/tmp/input") temp_input_dir.mkdir(parents=True, exist_ok=True) + video_path = temp_input_dir / video_filename audio_path = temp_input_dir / audio_filename - # Write the received bytes to temporary files in this container with open(video_path, "wb") as f: f.write(video_bytes) with open(audio_path, "wb") as f: @@ -157,121 +137,81 @@ def generate( output_dir = Path("/tmp/output") output_dir.mkdir(parents=True, exist_ok=True) - current_time = datetime.now().strftime("%Y%m%d_%H%M%S") - output_video_path = str(output_dir / f"result_{current_time}.mp4") + output_video_path = output_dir / f"result_{current_time}.mp4" - # Create an arguments namespace object, using the new temp file paths args = argparse.Namespace( inference_ckpt_path=str(CHECKPOINT_PATH / "latentsync_unet.pt"), video_path=str(video_path), audio_path=str(audio_path), - video_out_path=output_video_path, + video_out_path=str(output_video_path), inference_steps=inference_steps, guidance_scale=guidance_scale, seed=seed, temp_dir=str(output_dir), - enable_deepcache=True, # Enable for better performance + enable_deepcache=True, ) - - # It's good practice to work with a copy of the config for each run - # to avoid stateful issues if a container is reused. + run_config = deepcopy(self.config) - - # Update the config copy with runtime arguments run_config["run"].update({ "guidance_scale": guidance_scale, "inference_steps": inference_steps, }) - - print("Running inference...") + try: - # Call the main inference function from the original script inference_main(config=run_config, args=args) - print("Inference complete.") - - # Read the generated video file and return its content with open(output_video_path, "rb") as f: - content = f.read() - return content - + return f.read() except Exception as e: - print(f"Error during inference: {e}") + print("❌ Inference failed:", e) raise # --- FastAPI Web Server --- -# We define a FastAPI app to handle web requests. fastapi_app = FastAPI() def _get_filename_from_url(url: str) -> str: - """Helper function to extract a filename from a URL.""" return os.path.basename(urlparse(url).path) @fastapi_app.post("/lipsync", response_class=Response) async def lipsync( - # --- Input Options --- - # User can provide either a file upload OR a URL for video and audio - video: Optional[UploadFile] = File(None, description="Video file to be lip-synced."), - audio: Optional[UploadFile] = File(None, description="Audio file to sync with the video."), - video_url: Optional[str] = Form(None, description="URL of the video file."), - audio_url: Optional[str] = Form(None, description="URL of the audio file."), - - # --- Generation Parameters --- - guidance_scale: float = Form(1.5, description="Classifier-free guidance scale."), - inference_steps: int = Form(20, description="Number of DDIM inference steps."), - seed: int = Form(1247, description="Random seed for generation."), + video: Optional[UploadFile] = File(None), + audio: Optional[UploadFile] = File(None), + video_url: Optional[str] = Form(None), + audio_url: Optional[str] = Form(None), + guidance_scale: float = Form(1.5), + inference_steps: int = Form(20), + seed: int = Form(1247), ): - """ - This endpoint performs lip-syncing on a video using a target audio. - It accepts either direct file uploads or URLs for the video and audio sources. - """ - video_bytes: Optional[bytes] = None - video_filename: Optional[str] = None - audio_bytes: Optional[bytes] = None - audio_filename: Optional[str] = None + video_bytes = audio_bytes = None + video_filename = audio_filename = None - # --- Step 1: Process Video Input (File or URL) --- if video and video_url: - raise HTTPException(status_code=400, detail="Provide either a video file or a video_url, not both.") - + raise HTTPException(status_code=400, detail="Only one of video or video_url allowed.") if video: video_bytes = await video.read() video_filename = video.filename elif video_url: - try: - print(f"Downloading video from: {video_url}") - response = requests.get(video_url) - response.raise_for_status() # Raise an exception for bad status codes (4xx or 5xx) - video_bytes = response.content - video_filename = _get_filename_from_url(video_url) or "video_from_url.mp4" - except requests.exceptions.RequestException as e: - raise HTTPException(status_code=400, detail=f"Failed to download video from URL: {e}") - else: - raise HTTPException(status_code=400, detail="Either a video file or video_url must be provided.") + r = requests.get(video_url) + r.raise_for_status() + video_bytes = r.content + video_filename = _get_filename_from_url(video_url) - # --- Step 2: Process Audio Input (File or URL) --- if audio and audio_url: - raise HTTPException(status_code=400, detail="Provide either an audio file or an audio_url, not both.") - + raise HTTPException(status_code=400, detail="Only one of audio or audio_url allowed.") if audio: audio_bytes = await audio.read() audio_filename = audio.filename elif audio_url: - try: - print(f"Downloading audio from: {audio_url}") - response = requests.get(audio_url) - response.raise_for_status() - audio_bytes = response.content - audio_filename = _get_filename_from_url(audio_url) or "audio_from_url.mp3" - except requests.exceptions.RequestException as e: - raise HTTPException(status_code=400, detail=f"Failed to download audio from URL: {e}") - else: - raise HTTPException(status_code=400, detail="Either an audio file or audio_url must be provided.") - - # --- Step 3: Run Inference --- - # Instantiate the Modal class and call the generation method remotely + r = requests.get(audio_url) + r.raise_for_status() + audio_bytes = r.content + audio_filename = _get_filename_from_url(audio_url) + + if not video_bytes or not audio_bytes: + raise HTTPException(status_code=400, detail="Video and audio inputs are required.") + model = LatentSync() - output_video_bytes = model.generate.remote( + output = model.generate.remote( video_bytes, audio_bytes, video_filename, @@ -280,13 +220,9 @@ async def lipsync( inference_steps, seed ) - - # --- Step 4: Return Result --- - # Return the generated video as a response - return Response(content=output_video_bytes, media_type="video/mp4") + return Response(content=output, media_type="video/mp4") -# --- Modal ASGI App --- -# This serves the FastAPI application using Modal's web hosting capabilities. +# --- ASGI App Entry Point --- @app.function( image=latentsync_image, network_file_systems={str(CHECKPOINT_PATH): volume}, @@ -294,4 +230,4 @@ async def lipsync( ) @modal.asgi_app() def web_server(): - return fastapi_app + return fastapi_app \ No newline at end of file