SECourses.py

import os
import subprocess
import gradio as gr
from retinaface import RetinaFace
from PIL import Image
import filetype
from datetime import datetime
import re
import sys
import torch
import argparse

import platform, os

def open_folder():
    open_folder_path = os.path.abspath("outputs")
    if platform.system() == "Windows":
        os.startfile(open_folder_path)
    elif platform.system() == "Linux":
        os.system(f'xdg-open "{open_folder_path}"')


# Get the path to the currently activated Python executable
python_executable = sys.executable

def display_media(file):
    # Determine the type of the uploaded file using filetype
    if file is None:
       return gr.update(visible=False), gr.update(visible=False)
    kind = filetype.guess(file.name)

    if kind is None:
        return gr.update(visible=False), gr.update(visible=False)

    if kind.mime.startswith('video'):
        return gr.update(value=file.name, visible=True), gr.update(visible=False)
    elif kind.mime.startswith('audio'):
        return gr.update(visible=False), gr.update(value=file.name, visible=True)
    else:
        return gr.update(visible=False), gr.update(visible=False)


parser = argparse.ArgumentParser()
parser.add_argument("--share", type=str, default=False, help="Set to True to share the app publicly.")
args = parser.parse_args()


# Function to extract audio from video using FFmpeg
def extract_audio(video_path, audio_path):
    command = [python_executable, "-m", "ffmpeg", "-i", video_path, "-vn", "-acodec", "libmp3lame", "-q:a", "2", audio_path]
    subprocess.call(command)

# Function to convert audio to MP3 using FFmpeg
def convert_audio_to_mp3(audio_path, mp3_path):
    command = ["ffmpeg", "-i", audio_path, "-acodec", "libmp3lame", "-q:a", "2", mp3_path]
    subprocess.call(command)

def crop_and_save_image(image_path, auto_crop, crop_width, crop_height, crop_expansion):
    cropped_image = auto_crop_image(image_path, crop_expansion, crop_size=(crop_width, crop_height))
    if cropped_image is not None:
        cropped_folder = os.path.join("outputs", "cropped_images")
        os.makedirs(cropped_folder, exist_ok=True)
        
        # Get the base name and extension of the image file
        base_name, extension = os.path.splitext(os.path.basename(image_path))
        
        # Initialize the counter for the image number
        counter = 1
        
        # Generate the new image name with the incremented number
        new_image_name = f"{base_name}_{counter:04d}{extension}"
        cropped_image_path = os.path.join(cropped_folder, new_image_name)
        
        # Check if the image already exists and increment the counter until a unique name is found
        while os.path.exists(cropped_image_path):
            counter += 1
            new_image_name = f"{base_name}_{counter:04d}{extension}"
            cropped_image_path = os.path.join(cropped_folder, new_image_name)
        
        # Save the cropped image with the new name
        cropped_image.save(cropped_image_path, format='PNG')
        return cropped_image_path
    return None

# Function to generate kps sequence and audio from video
def generate_kps_sequence_and_audio(video_path, kps_sequence_save_path, audio_save_path):
    command = [python_executable, "scripts/extract_kps_sequence_and_audio.py", "--video_path", video_path, "--kps_sequence_save_path", kps_sequence_save_path, "--audio_save_path", audio_save_path]
    subprocess.call(command)

def auto_crop_image(image_path, expand_percent, crop_size=(512, 512)):
    # Check if CUDA is available
    if torch.cuda.is_available():
        device = 'cuda'
        print("Using GPU for RetinaFace detection.")
    else:
        device = 'cpu'
        print("Using CPU for RetinaFace detection.")

    # Load image
    img = Image.open(image_path)

    # Perform face detection
    faces = RetinaFace.detect_faces(image_path)

    if not faces:
        print("No faces detected.")
        return None

    # Assuming 'faces' is a dictionary of detected faces
    # Pick the first face detected
    face = list(faces.values())[0]
    landmarks = face['landmarks']

    # Extract the landmarks
    right_eye = landmarks['right_eye']
    left_eye = landmarks['left_eye']
    right_mouth = landmarks['mouth_right']
    left_mouth = landmarks['mouth_left']

    # Calculate the distance between the eyes
    eye_distance = abs(right_eye[0] - left_eye[0])

    # Estimate the head width and height
    head_width = eye_distance * 4.5  # Increase the width multiplier
    head_height = eye_distance * 6.5  # Increase the height multiplier

    # Calculate the center point between the eyes
    eye_center_x = (right_eye[0] + left_eye[0]) // 2
    eye_center_y = (right_eye[1] + left_eye[1]) // 2

    # Calculate the top-left and bottom-right coordinates of the assumed head region
    head_left = max(0, int(eye_center_x - head_width // 2))
    head_top = max(0, int(eye_center_y - head_height // 2))  # Adjust the top coordinate
    head_right = min(img.width, int(eye_center_x + head_width // 2))
    head_bottom = min(img.height, int(eye_center_y + head_height // 2))  # Adjust the bottom coordinate

    # Save the assumed head image
    assumed_head_img = img.crop((head_left, head_top, head_right, head_bottom))
    assumed_head_img.save("assumed_head.png", format='PNG')

    # Calculate the expansion in pixels and the new dimensions
    expanded_w = int(head_width * (1 + expand_percent))
    expanded_h = int(head_height * (1 + expand_percent))

    # Calculate the top-left and bottom-right points of the expanded box
    center_x, center_y = head_left + head_width // 2, head_top + head_height // 2
    left = max(0, center_x - expanded_w // 2)
    right = min(img.width, center_x + expanded_w // 2)
    top = max(0, center_y - expanded_h // 2)
    bottom = min(img.height, center_y + expanded_h // 2)

    # Crop the image with the expanded boundaries
    cropped_img = img.crop((left, top, right, bottom))
    cropped_img.save("expanded_face.png", format='PNG')

    # Calculate the aspect ratio of the cropped image
    cropped_width, cropped_height = cropped_img.size
    aspect_ratio = cropped_width / cropped_height

    # Calculate the target dimensions based on the desired crop size
    target_width = crop_size[0]
    target_height = crop_size[1]

    # Adjust the crop to match the desired aspect ratio
    if aspect_ratio > target_width / target_height:
        # Crop from left and right
        new_width = int(cropped_height * target_width / target_height)
        left_crop = (cropped_width - new_width) // 2
        right_crop = left_crop + new_width
        top_crop = 0
        bottom_crop = cropped_height
    else:
        # Crop from top and bottom
        new_height = int(cropped_width * target_height / target_width)
        top_crop = (cropped_height - new_height) // 2
        bottom_crop = top_crop + new_height
        left_crop = 0
        right_crop = cropped_width

    # Crop the image with the adjusted boundaries
    final_cropped_img = cropped_img.crop((left_crop, top_crop, right_crop, bottom_crop))
    final_cropped_img.save("final_cropped_img.png", format='PNG')

    # Resize the cropped image to the desired size (512x512 by default) with best quality
    resized_img = final_cropped_img.resize(crop_size, resample=Image.LANCZOS)

    # Save the resized image as PNG
    resized_img.save(image_path, format='PNG')
    return resized_img
     

def generate_output_video(reference_image_path, audio_path, kps_path, output_path, retarget_strategy, num_inference_steps, reference_attention_weight, audio_attention_weight, auto_crop, crop_width, crop_height, crop_expansion,image_width,image_height, low_vram):
    print("auto cropping...")
    if auto_crop:
        auto_crop_image(reference_image_path,crop_expansion, crop_size=(crop_width, crop_height))
    
    print("starting inference...")
    command = [
        python_executable, "inference.py",
        "--reference_image_path", reference_image_path,
        "--audio_path", audio_path,
        "--kps_path", kps_path,
        "--output_path", output_path,
        "--retarget_strategy", retarget_strategy,
        "--num_inference_steps", str(num_inference_steps),
        "--reference_attention_weight", str(reference_attention_weight),
        "--audio_attention_weight", str(audio_attention_weight),     
        "--image_width", str(image_width),
        "--image_height", str(image_height)
    ]

    if low_vram:  # Add the --save_gpu_memory flag if Low VRAM is checked
        command.append("--save_gpu_memory")
    
    with open("executed_command.txt", "w") as file:
        file.write(" ".join(command))
    
    subprocess.call(command)
    return output_path, reference_image_path

def sanitize_folder_name(name):
    # Define a regex pattern to match invalid characters for both Linux and Windows
    invalid_chars = r'[<>:"/\\|?*\x00-\x1F]'
    # Replace invalid characters with an underscore
    sanitized_name = re.sub(invalid_chars, '_', name)
    return sanitized_name

# Function to handle the input and generate the output
def process_input(reference_image, target_input, retarget_strategy, num_inference_steps, reference_attention_weight, audio_attention_weight, auto_crop, crop_width, crop_height, crop_expansion,image_width,image_height,low_vram):
    # Create temp_process directory for intermediate files
    temp_process_dir = "temp_process"
    os.makedirs(temp_process_dir, exist_ok=True)
    
    input_file_name = os.path.splitext(os.path.basename(reference_image))[0]
    input_file_name=sanitize_folder_name(input_file_name)
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    temp_dir = os.path.join(temp_process_dir, f"{input_file_name}_{timestamp}")
    os.makedirs(temp_dir, exist_ok=True)
    
    kind = filetype.guess(target_input)
    if not kind:
        raise ValueError("Cannot determine file type. Please provide a valid video or audio file.")
    
    mime_type = kind.mime
    
    if mime_type.startswith("video/"):  # Video input
        audio_path = os.path.join(temp_dir, "target_audio.mp3")
        kps_path = os.path.join(temp_dir, "kps.pth")
        print("generating generate_kps_sequence_and_audio...")
        generate_kps_sequence_and_audio(target_input, kps_path, audio_path)
    elif mime_type.startswith("audio/"):  # Audio input
        audio_path = target_input
        if mime_type != "audio/mpeg":
            mp3_path = os.path.join(temp_dir, "target_audio_converted.mp3")
            convert_audio_to_mp3(target_input, mp3_path)
            audio_path = mp3_path
        kps_path = ""
    else:
        raise ValueError("Unsupported file type. Please provide a video or audio file.")
    
    output_dir = "outputs"
    os.makedirs(output_dir, exist_ok=True)
    output_file_name = f"{input_file_name}_result_"
    output_file_name=sanitize_folder_name(output_file_name)
    output_file_ext = ".mp4"
    output_file_count = 1
    while os.path.exists(os.path.join(output_dir, f"{output_file_name}{output_file_count:04d}{output_file_ext}")):
        output_file_count += 1
    output_path = os.path.join(output_dir, f"{output_file_name}{output_file_count:04d}{output_file_ext}")

    
    output_video_path, cropped_image_path = generate_output_video(reference_image, audio_path, kps_path, output_path, retarget_strategy, num_inference_steps, reference_attention_weight, audio_attention_weight, auto_crop,crop_width,crop_height, crop_expansion,image_width,image_height,low_vram)
    
    return output_video_path, cropped_image_path

def launch_interface():
    retarget_strategies = ["fix_face", "no_retarget", "offset_retarget", "naive_retarget"]
   
    with gr.Blocks() as demo:
        gr.Markdown("# Tencent AI Lab - V-Express Image to Animation V4 : https://www.patreon.com/posts/105251204")
        with gr.Row():          
            with gr.Column():
                input_image = gr.Image(label="Reference Image", format="png", type="filepath", height=512)
                generate_button = gr.Button("Generate Talking Video")
                low_vram = gr.Checkbox(label="Low VRAM - Greatly reduces VRAM usage but takes longer", value=False,visible=False)
                crop_button = gr.Button("Crop Image")
                with gr.Row():

                    with gr.Column(min_width=0):
                        image_width = gr.Number(label="Target Video Width", value=512)

                    with gr.Column(min_width=0):
                        image_height = gr.Number(label="Target Video Height", value=512)

                with gr.Row():
                    with gr.Column(min_width=0):
                        retarget_strategy = gr.Dropdown(retarget_strategies, label="Retarget Strategy", value="fix_face")
                    with gr.Column(min_width=0):
                        inference_steps = gr.Slider(10, 90, step=1, label="Number of Inference Steps", value=30)

                with gr.Row():
                    with gr.Column(min_width=0):
                        reference_attention = gr.Slider(0.80, 1.1, step=0.01, label="Reference Attention Weight", value=0.95)
                    with gr.Column(min_width=0):
                        audio_attention = gr.Slider(1.0, 5.0, step=0.1, label="Audio Attention Weight", value=3.0)             

                with gr.Row(visible=True) as crop_size_row:
                    with gr.Column(min_width=0):
                        auto_crop = gr.Checkbox(label="Auto Crop Image", value=True)
                    with gr.Column(min_width=0):
                        crop_expansion = gr.Slider(0.0, 1.0, step=0.01, label="Face Focus Expansion Percent", value=0.15)
                with gr.Row():
                    with gr.Column(min_width=0):
                        crop_width = gr.Number(label="Crop Width", value=512)
                    with gr.Column(min_width=0):
                        crop_height = gr.Number(label="Crop Height", value=512)
        
            with gr.Column():
                input_video = gr.File(
                    label="Target Input (Image or Video)",
                    type="filepath",
                    file_count="single",
                    file_types=[
                        ".mp4", ".avi", ".mov", ".wmv", ".flv", ".mkv", ".webm",  # Video extensions
                        ".3gp", ".m4v", ".mpg", ".mpeg", ".m2v", ".m4v", ".mts",  # More video extensions
                        ".mp3", ".wav", ".aac", ".flac", ".m4a", ".wma", ".ogg"   # Audio extensions
                    ],
                    height=512        )
                video_output = gr.Video(visible=False)
                audio_output = gr.Audio(visible=False)
    
                input_video.change(display_media, inputs=input_video, outputs=[video_output, audio_output])
                btn_open_outputs = gr.Button("Open Outputs Folder")
                btn_open_outputs.click(fn=open_folder)
                gr.Markdown("""



                            Retarget Strategies



                            Only target audio : fix_face



                            Input picture and target video (same person - best practice) select : no_retarget



                            Input picture and target video (different person) select : offset_retarget or naive_retarget



                            Please look examples in Tests folder to see which settings you like most. I feel like offset_retarget is best



                            You can turn up reference_attention_weight to make the model maintain higher character consistency, and turn down audio_attention_weight to reduce mouth artifacts. E.g. setting both values to 1.0

                            """)



            with gr.Column():
                output_video = gr.Video(label="Generated Video", height=512)
                output_image = gr.Image(label="Cropped Image")
        
        
        generate_button.click(
            fn=process_input,
            inputs=[
                input_image,
                input_video,
                retarget_strategy,
                inference_steps,
                reference_attention,
                audio_attention,
                auto_crop,
                crop_width,
                crop_height,
                crop_expansion,
                image_width,
                image_height,
                low_vram
            ],
            outputs=[output_video, output_image]
        )

        crop_button.click(
            fn=crop_and_save_image,
            inputs=[
                input_image,
                auto_crop,
                crop_width,
                crop_height,
                crop_expansion
            ],
            outputs=output_image
        )
    
    demo.queue()
    demo.launch(inbrowser=True,share=args.share)

# Run the Gradio interface
launch_interface()