app.py

import gradio as gr
import subprocess as sp
import os
import uuid
import time
import shutil

# Constants
OUTPUT_DIR = './output'
NUM_MINUTES = 60

# Create output directory if it doesn't exist
os.makedirs(OUTPUT_DIR, exist_ok=True)

class DeepFakeAI:
    def __init__(self):
        # Initialize the frame processor checkbox
        self.frame_processor_checkbox = gr.CheckboxGroup(
            choices=['face_swapper', 'face_enhancer', 'frame_enhancer'],
            label='FRAME PROCESSORS',
            value=['face_swapper']  # Default value
        )
        
        # Initialize the face analyser direction dropdown
        self.face_analyser_direction_dropdown = gr.Dropdown(
            label='FACE ANALYSER DIRECTION',
            choices=['left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small'],
            value='left-right'
        )
        
        # Initialize the face analyser age dropdown
        self.face_analyser_age_dropdown = gr.Dropdown(
            label='FACE RECOGNITION',
            choices=['none'] + ['reference', 'many'],
            value='reference'
        )
        
        # Initialize the face analyser gender dropdown
        self.face_analyser_gender_dropdown = gr.Dropdown(
            label='FACE ANALYSER GENDER',
            choices=['none'] + ['male', 'female'],
            value='none'
        )
        
        # Initialize the unique id textbox
        self.unique_id = gr.Textbox(value=str(uuid.uuid4()), visible=False)
    
    def run(self, *args):
        # Unpack input arguments
        source, target, unique_id, *rest_args = args
        
        # Check if source and target files exist
        if not os.path.exists(source):
            return "Source file does not exist"
        if not os.path.exists(target):
            return "Target file does not exist"
        
        # Remove old directories
        self.remove_old_directories(OUTPUT_DIR, NUM_MINUTES)
        
        # Get filename and create output directory
        filename = os.path.basename(target)
        output_dir = os.path.join(OUTPUT_DIR, unique_id)
        os.makedirs(output_dir, exist_ok=True)
        output = os.path.join(output_dir, filename)
        
        # Get frame processor and face analyser settings
        frame_processor = rest_args[0]
        selected_frame_processors = ' '.join(frame_processor)
        face_analyser_direction = rest_args[1]
        face_recognition = rest_args[2]
        face_analyser_gender = rest_args[3]
        
        # Construct command to run
        cmd = [
            'python', 'run.py',
            '--execution-providers', 'cpu',
            '-s', source,
            '-t', target,
            '-o', output,
            '--frame-processors', selected_frame_processors,
            '--face-analyser-direction', face_analyser_direction
        ]
        if face_recognition != 'none':
            cmd.extend(['--face-recognition', face_recognition])
        if face_analyser_gender != 'none':
            cmd.extend(['--face-analyser-gender', face_analyser_gender])
        
        # Add additional flags if present
        if len(rest_args) > 4:
            skip_audio = rest_args[4]
            keep_fps = rest_args[5]
            keep_temp = rest_args[6]
            if skip_audio:
                cmd.append('--skip-audio')
            if keep_fps:
                cmd.append('--keep-fps')
            if keep_temp:
                cmd.append('--keep-temp')
        
        try:
            print("Started...", cmd)
            output_text = sp.run(cmd, capture_output=True, text=True).stdout
            print(output_text)
            return output
        except Exception as e:
            return f"An error occurred: {str(e)}"
    
    def clear_output(self, unique_id):
        try:
            output_path = os.path.join(OUTPUT_DIR, unique_id)
            if os.path.exists(output_path):
                print("Trying to delete")
                shutil.rmtree(output_path)
                print(f"Output files in {output_path} are deleted")
                return "Output files for unique_id deleted"
            else:
                print(f"Output files in {output_path} does not exist")
                return "Output directory for (output_path} does not exist"
        except Exception as e:
            return f"An error occurred: {str(e)}"
    
    def remove_old_directories(self, directory, num_minutes):
        now = time.time()
        
        for r, d, f in os.walk(directory):
            for dir_name in d:
                dir_path = os.path.join(r, dir_name)
                timestamp = os.path.getmtime(dir_path)
                age_minutes = (now - timestamp) / 60 # Convert to minutes
                
                if age_minutes >= num_minutes:
                    try:
                        print("Removing", dir_path)
                        shutil.rmtree(dir_path)
                        print("Directory removed:", dir_path)
                    except Exception as e:
                        print(e)
                        pass
    
    def get_theme(self) -> gr.Theme:
        return gr.themes.Soft(
            primary_hue=gr.themes.colors.red,
            secondary_hue=gr.themes.colors.gray,
            font=gr.themes.GoogleFont('Inter')
        ).set(
            background_fill_primary='*neutral_50',
            block_label_text_size='*text_sm',
            block_title_text_size='*text_sm'
        )

# Create an instance of the DeepFakeAI class
app = DeepFakeAI()

# Create the UI
with gr.Blocks(theme=app.get_theme(), title="DeepFakeAI 1.0.0") as ui:
    with gr.Group():
        gr.HTML('<center><a href="https://codegenius.me">DeepFakeAI 1.0.1</a></center>')
    
    with gr.Group():
        with gr.Column(scale=3):
            gr.Label("FRAME PROCESSORS")
            app.frame_processor_checkbox.render()
    
    with gr.Group():
        with gr.Column(scale=3):
            gr.Label("FACE ANALYSER DIRECTION")
            app.face_analyser_direction_dropdown.render()
            gr.Label("FACE RECOGNITION")
            app.face_analyser_age_dropdown.render()
            gr.Label("FACE ANALYSER GENDER")
            app.face_analyser_gender_dropdown.render()
            app.unique_id.render()
    
    with gr.Tab("Image:"):
        source_image = gr.Image(type="filepath", label="SOURCE IMAGE")
        target_image = gr.Files(label="TARGET IMAGE(S)")
        image_button = gr.Button("START")
        clear_button = gr.ClearButton(value="CLEAR")
        image_output = gr.Files(label="OUTPUT")
        clear_button.add(image_output)
        
        image_button.click(
            app.run,
            inputs=[source_image, target_image, app.unique_id, app.frame_processor_checkbox, app.face_analyser_direction_dropdown, app.face_analyser_age_dropdown, app.face_analyser_gender_dropdown],
            outputs=image_output
        )
        clear_button.click(fn=app.clear_output, inputs=app.unique_id)
    
    with gr.Tab("Video:"):
        source_image_video = gr.Image(type="filepath", label="SOURCE IMAGE")
        target_video = gr.Files(label="TARGET VIDEO(S)")
        with gr.Group():
            skip_audio = gr.Checkbox(label="SKIP AUDIO")
            keep_fps = gr.Checkbox(label="KEEP FPS")
            keep_temp = gr.Checkbox(label="KEEP TEMP")
        video_button = gr.Button("START")
        clear_video_button = gr.ClearButton(value="CLEAR")
        video_output = gr.Files(label="OUTPUT")
        clear_video_button.add(video_output)
        video_button.click(
            app.run,
            inputs=[source_image_video, target_video, app.unique_id, app.frame_processor_checkbox, app.face_analyser_direction_dropdown, app.face_analyser_age_dropdown, app.face_analyser_gender_dropdown, skip_audio, keep_fps, keep_temp],
            outputs=video_output
        )
        clear_video_button.click(fn=app.clear_output, inputs=app.unique_id)

ui.launch(debug=True)