#!/usr/bin/env python3 from typing import Dict, Any import torch from PIL import Image import base64 from io import BytesIO import numpy as np from diffusers import AutoencoderKL, DDIMScheduler from einops import repeat from omegaconf import OmegaConf from transformers import CLIPVisionModelWithProjection import cv2 import os import sys import skvideo.io from src.models.pose_guider import PoseGuider from src.models.unet_2d_condition import UNet2DConditionModel from src.models.unet_3d import UNet3DConditionModel from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline from src.utils.util import read_frames, get_fps, save_videos_grid import roop.globals from roop.core import start, decode_execution_providers, suggest_max_memory, suggest_execution_threads from roop.utilities import normalize_output_path from roop.processors.frame.core import get_frame_processors_modules # import onnxruntime as ort import gc import subprocess import requests import tempfile from rembg import remove import onnxruntime as ort import shutil device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device.type != 'cuda': raise ValueError("The model requires a GPU for inference.") class EndpointHandler(): def __init__(self, path=""): base_dir = os.path.dirname(os.path.abspath(__file__)) config_path = os.path.join(base_dir, 'configs', 'prompts', 'animation.yaml') if not os.path.exists(config_path): raise FileNotFoundError(f"The configuration file was not found at: {config_path}") self.config = OmegaConf.load(config_path) self.weight_dtype = torch.float16 self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.pipeline = None self._initialize_pipeline() def _initialize_pipeline(self): base_dir = os.path.dirname(os.path.abspath(__file__)) config_path = os.path.join(base_dir, 'pretrained_weights', 'sd-vae-ft-mse') if not os.path.exists(config_path): raise FileNotFoundError(f"The sd-vae-ft-mse folder was not found at: {config_path}") vae = AutoencoderKL.from_pretrained(config_path).to(self.device, dtype=self.weight_dtype) pretrained_base_model_path_unet = os.path.join(base_dir, 'pretrained_weights', 'stable-diffusion-v1-5', 'unet') print("model path is " + pretrained_base_model_path_unet) reference_unet = UNet2DConditionModel.from_pretrained( pretrained_base_model_path_unet ).to(dtype=self.weight_dtype, device=self.device) inference_config_path = os.path.join(base_dir, 'configs', 'inference', 'inference_v2.yaml') motion_module_path = os.path.join(base_dir, 'pretrained_weights', 'motion_module.pth') denoising_unet_path = os.path.join(base_dir, 'pretrained_weights', 'denoising_unet.pth') reference_unet_path = os.path.join(base_dir, 'pretrained_weights', 'reference_unet.pth') pose_guider_path = os.path.join(base_dir, 'pretrained_weights', 'pose_guider.pth') image_encoder_path = os.path.join(base_dir, 'pretrained_weights', 'image_encoder') infer_config = OmegaConf.load(inference_config_path) denoising_unet = UNet3DConditionModel.from_pretrained_2d( pretrained_base_model_path_unet, motion_module_path, unet_additional_kwargs=infer_config.unet_additional_kwargs, ).to(self.device, dtype=self.weight_dtype) pose_guider = PoseGuider(320, block_out_channels=(16, 32, 96, 256)).to(self.device, dtype=self.weight_dtype) image_enc = CLIPVisionModelWithProjection.from_pretrained(image_encoder_path).to(self.device, dtype=self.weight_dtype) sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs) scheduler = DDIMScheduler(**sched_kwargs) denoising_unet.load_state_dict(torch.load(denoising_unet_path, map_location="cpu"), strict=False) reference_unet.load_state_dict(torch.load(reference_unet_path, map_location="cpu")) pose_guider.load_state_dict(torch.load(pose_guider_path, map_location="cpu")) self.pipeline = Pose2VideoPipeline( vae=vae, image_encoder=image_enc, reference_unet=reference_unet, denoising_unet=denoising_unet, pose_guider=pose_guider, scheduler=scheduler ).to(self.device, dtype=self.weight_dtype) def _crop_face(self, image, save_path="cropped_face.jpg", margin=0.5): # Convert image to OpenCV format cv_image = np.array(image) cv_image = cv_image[:, :, ::-1].copy() # Load OpenCV face detector face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml') # Detect faces gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY) faces = face_cascade.detectMultiScale(gray, 1.1, 4) if len(faces) == 0: raise ValueError("No faces detected in the reference image.") # Crop the first face found with a margin x, y, w, h = faces[0] x_margin = int(margin * w) y_margin = int(margin * h) x1 = max(0, x - x_margin) y1 = max(0, y - y_margin // 2) # Less margin at the top x2 = min(cv_image.shape[1], x + w + x_margin) y2 = min(cv_image.shape[0], y + h + y_margin) # More margin at the bottom cropped_face = cv_image[y1:y2, x1:x2] # Convert back to PIL format cropped_face = Image.fromarray(cropped_face[:, :, ::-1]).convert("RGB") # Save the cropped face cropped_face.save(save_path, format="JPEG", quality=95) return cropped_face def _swap_face(self, source_path, target_video_path, output_path): # source_path = "input.jpg" # source_image.save(source_path, format="JPEG", quality=95) roop.globals.source_path = source_path roop.globals.target_path = target_video_path roop.globals.output_path = output_path roop.globals.frame_processors = ["face_swapper", "face_enhancer"] roop.globals.headless = True roop.globals.keep_fps = True roop.globals.keep_audio = True roop.globals.keep_frames = False roop.globals.many_faces = False # roop.globals.video_encoder = "libx264" roop.globals.video_quality = 50 roop.globals.max_memory = suggest_max_memory() # Set GPU execution provider roop.globals.execution_providers = decode_execution_providers(["CUDAExecutionProvider"]) roop.globals.execution_threads = suggest_execution_threads() # Ensure onnxruntime is using the GPU ort.set_default_logger_severity(3) # Suppress verbose logging providers = ['CUDAExecutionProvider'] options = ort.SessionOptions() options.intra_op_num_threads = 1 for frame_processor in get_frame_processors_modules(roop.globals.frame_processors): if hasattr(frame_processor, 'onnx_session'): frame_processor.onnx_session.set_providers(providers, options) # Clear CUDA cache before starting the face swapping process torch.cuda.empty_cache() start() # Clear CUDA cache after the face swapping process for frame_processor in roop.globals.frame_processors: del frame_processor torch.cuda.empty_cache() return os.path.join(os.getcwd(), output_path) def print_memory_stat_for_stuff(self, phase, log_file="memory_stats.log"): with open(log_file, "a") as f: f.write(f"Memory Stats - {phase}:\n") f.write(f"Allocated memory: {torch.cuda.memory_allocated() / 1024**2:.2f} MB\n") f.write(f"Reserved memory: {torch.cuda.memory_reserved() / 1024**2:.2f} MB\n") f.write(f"Max allocated memory: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB\n") f.write(f"Max reserved memory: {torch.cuda.max_memory_reserved() / 1024**2:.2f} MB\n") f.write("="*30 + "\n") def convert_to_playable_format(self, input_path, output_path): with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_file: temp_output_path = tmp_file.name command = f"ffmpeg -i {input_path} -c:v libx264 -preset fast -crf 18 -y {temp_output_path}" # Run the command with shell=True result = subprocess.run(command, shell=True, capture_output=True, text=True) print("Conversion STDOUT:", result.stdout) print("Conversion STDERR:", result.stderr) if result.returncode != 0: raise RuntimeError(f"FFmpeg conversion failed with exit code {result.returncode}") shutil.move(temp_output_path, output_path) def run_rife_interpolation(self, video_path, output_path, multi=2, scale=1.0): base_dir = os.path.dirname(os.path.abspath(__file__)) directory = os.path.join(base_dir, "Practical-RIFE", "inference_video.py") model_directory = os.path.join(base_dir, "Practical-RIFE", "train_log") command = f"python3 {directory} --video={video_path} --output={output_path} --multi={multi} --scale={scale} --model={model_directory}" # Run the command with shell=True result = subprocess.run(command, shell=True, capture_output=True, text=True) print(result) print(result.stdout) print(result.stderr) if result.returncode != 0: raise RuntimeError(f"RIFE interpolation failed with exit code {result.returncode}") # Overwrite the RIFE output with the converted playable format self.convert_to_playable_format(output_path, output_path) def speed_up_video(self, input_path, output_path, factor=4): command = f"ffmpeg -i {input_path} -filter:v setpts=PTS/{factor} -an {output_path}" # Run the command with shell=True result = subprocess.run(command, shell=True, capture_output=True, text=True) print("Speed Up Video STDOUT:", result.stdout) print("Speed Up Video STDERR:", result.stderr) if result.returncode != 0: raise RuntimeError(f"FFmpeg speed up failed with exit code {result.returncode}") def slow_down_video(self, input_path, output_path, factor=4): command = f"ffmpeg -i {input_path} -filter:v setpts={factor}*PTS -an {output_path}" # Run the command with shell=True result = subprocess.run(command, shell=True, capture_output=True, text=True) print("Slow Down Video STDOUT:", result.stdout) print("Slow Down Video STDERR:", result.stderr) if result.returncode != 0: raise RuntimeError(f"FFmpeg slow down failed with exit code {result.returncode}") def download_file(self, url: str, save_path: str): response = requests.get(url, stream=True) if response.status_code == 200: with open(save_path, 'wb') as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) else: raise ValueError(f"Failed to download file from {url}") def print_directory_contents(self, directory): for root, dirs, files in os.walk(directory): level = root.replace(directory, '').count(os.sep) indent = ' ' * 4 * (level) print(f"{indent}{os.path.basename(root)}/") subindent = ' ' * 4 * (level + 1) for f in files: print(f"{subindent}{f}") def __call__(self, data: Any) -> Dict[str, str]: inputs = data.get("inputs", {}) ref_image_url = inputs.get("ref_image_url", "") video_url = inputs.get("video_url", "") width = inputs.get("width", 512) height = inputs.get("height", 768) length = inputs.get("length", 24) num_inference_steps = inputs.get("num_inference_steps", 25) cfg = inputs.get("cfg", 3.5) seed = inputs.get("seed", 123) base_dir = os.path.dirname(os.path.abspath(__file__)) # Create a unique temporary directory for this request with tempfile.TemporaryDirectory() as temp_dir: print(f"Temporary directory created at {temp_dir}") # Debug statement video_root = os.path.join(temp_dir, "dw_poses_videos") os.makedirs(video_root, exist_ok=True) downloaded_video_path = os.path.join(video_root, "downloaded_video.mp4") downloaded_image_path = os.path.join(video_root, "downloaded_image.jpg") # Download the video from the URL self.download_file(video_url, downloaded_video_path) # Download the reference image from the URL self.download_file(ref_image_url, downloaded_image_path) ref_image = Image.open(downloaded_image_path) # Calculate new dimensions original_width, original_height = ref_image.size max_dimension = max(original_width, original_height) if max_dimension > 600: ratio = max_dimension / 600 width = int(original_width / ratio) height = int(original_height / ratio) else: width = original_width height = original_height # Remove the background from the reference image ref_image_no_bg = remove(ref_image) ref_image_no_bg_path = os.path.join(video_root, "ref_image_no_bg.png") ref_image_no_bg.save(ref_image_no_bg_path) pose_output_path = os.path.join(temp_dir, "pose_videos") # Run the extract_dwpose_from_vid.py script extract_pose_path = os.path.join(base_dir, 'extract_dwpose_from_vid.py') command = f'python3 {extract_pose_path} --video_root {video_root}' # Run the command with shell=True result = subprocess.run(command, shell=True, capture_output=True, text=True) if result.returncode != 0: raise RuntimeError(f"Error running extract_dwpose_from_vid.py: {result.stderr}") # Locate the extracted pose video save_dir = video_root + "_dwpose" print(f"Expected save directory: {save_dir}") # Debug statement pose_video_path = os.path.join(save_dir, "downloaded_video.mp4") if not os.path.exists(pose_video_path): print("Contents of the temporary directory:") self.print_directory_contents(temp_dir) raise FileNotFoundError(f"The pose video was not found at: {pose_video_path}") # Speed up the pose video by 4x sped_up_pose_video_path = os.path.join(temp_dir, "sped_up_pose_video.mp4") self.speed_up_video(pose_video_path, sped_up_pose_video_path, factor=4) torch.manual_seed(seed) pose_images = read_frames(sped_up_pose_video_path) src_fps = get_fps(sped_up_pose_video_path) pose_list = [] total_length = min(length, len(pose_images)) for pose_image_pil in pose_images[:total_length]: pose_list.append(pose_image_pil) video = self.pipeline( ref_image_no_bg, pose_list, width=width, height=height, video_length=total_length, num_inference_steps=num_inference_steps, guidance_scale=cfg ).videos save_dir = os.path.join(temp_dir, "output") if not os.path.exists(save_dir): os.makedirs(save_dir, exist_ok=True) animation_path = os.path.join(save_dir, "animation_output.mp4") save_videos_grid(video, animation_path, n_rows=1, fps=src_fps) # Crop the face from the reference image and save it cropped_face_path = os.path.join(save_dir, "cropped_face.jpg") cropped_face = self._crop_face(ref_image_no_bg, save_path=cropped_face_path) # Delete the pipeline and clear CUDA cache to free up memory del self.pipeline torch.cuda.empty_cache() # Perform face swapping # self.print_directory_contents(temp_dir) swapped_face_video_path = os.path.join(save_dir, "swapped_face_output.mp4") self._swap_face('./good_face.jpeg', animation_path, swapped_face_video_path) # Slow down the produced video by 4x self.print_directory_contents(temp_dir) slowed_down_animation_path = os.path.join(save_dir, "slowed_down_animation_output.mp4") self.slow_down_video(swapped_face_video_path, slowed_down_animation_path, factor=4) # Clear CUDA cache before RIFE interpolation torch.cuda.empty_cache() # Perform RIFE interpolation # self.print_directory_contents(temp_dir) rife_output_path = os.path.join(save_dir, "completed_result.mp4") self.run_rife_interpolation(slowed_down_animation_path, rife_output_path, multi=2, scale=0.5) # Encode the final video in base64 with open(rife_output_path, "rb") as video_file: video_base64 = base64.b64encode(video_file.read()).decode("utf-8") torch.cuda.empty_cache() return {"video": video_base64}