Update app.py (#2)

- Update app.py (a57ed4682be746df61032f8ddafc3cac7917a29a)


Co-authored-by: Scott Copeland <[email protected]>

app.py

@@ -5,163 +5,204 @@ import uuid
 import time
 import shutil
 
-os.makedirs("./output", exist_ok=True)
-
-def run(*args):
-    source, target, unique_id, *rest_args = args
-    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("./output", num_minutes=60)
-    filename = os.path.basename(target)
-    os.makedirs(f"./output/{unique_id}",exist_ok=True)
-    output = f"./output/{unique_id}/{filename}"
-    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]
-
-    cmd = (
-        f"python run.py --execution-providers cpu -s {source} -t {target} -o {output} "
-        f"--frame-processors {selected_frame_processors} "
-        f"--face-analyser-direction {face_analyser_direction} "
-    )
-    if face_recognition != 'none':
-        cmd += f"--face-recognition {face_recognition} "
-    if face_analyser_gender != 'none':
-        cmd += f"--face-analyser-gender {face_analyser_gender} "
-
-    if len(rest_args) > 4:
-        skip_audio = rest_args[4]
-        keep_fps = rest_args[5]
-        keep_temp = rest_args[6]
-        if skip_audio:
-            cmd += "--skip-audio "
-        if keep_fps:
-            cmd += "--keep-fps "
-        if keep_temp:
-            cmd += "--keep-temp "
-
-    try:
-        print("Started...", cmd)
-        output_text = sp.run(cmd, shell=True, 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(unique_id):
-    try:
-        output_path = f"./output/{unique_id}"
-        if os.path.exists(output_path):
-            print("Trying to delete ")
-            for filename in os.listdir(output_path):
-                file_path = os.path.join(output_path, filename)
-                if os.path.isfile(file_path):
-                    os.remove(file_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(directory, num_minutes=60):
-    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
-                    
+# 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'
+        )
 
-def get_theme() -> 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()
 
-with gr.Blocks(theme=get_theme(), title="DeepFakeAI 1.0.0") as ui:
+# 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):
-           frame_processor_checkbox = gr.CheckboxGroup(
-            choices = ['face_swapper', 'face_enhancer', 'frame_enhancer'],
-            label = 'FRAME PROCESSORS',
-            value = ['face_swapper']  # Default value
-        )
-
+            gr.Label("FRAME PROCESSORS")
+            app.frame_processor_checkbox.render()
     
     with gr.Group():
         with gr.Column(scale=3):
-            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'
-            )
-            face_analyser_age_dropdown = gr.Dropdown(
-                label = 'FACE RECOGNITION',
-                choices = ['none'] + ['reference', 'many'],
-                value = 'reference'
-            )
-            face_analyser_gender_dropdown = gr.Dropdown(
-                label = 'FACE ANALYSER GENDER',
-                choices = ['none'] + ['male', 'female'],
-                value = 'none'
-            )
-            unique_id = gr.Textbox(value=str(uuid.uuid4()), visible=False)
-    with gr.Tab("Image:               "):
+            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.Image(type="filepath", label="TARGET IMAGE")
+        target_image = gr.Files(label="TARGET IMAGE(S)")
         image_button = gr.Button("START")
         clear_button = gr.ClearButton(value="CLEAR")
-        image_output = gr.Image(label="OUTPUT")
+        image_output = gr.Files(label="OUTPUT")
         clear_button.add(image_output)
-
+        
         image_button.click(
-            run,
-            inputs=[source_image, target_image, unique_id, frame_processor_checkbox, face_analyser_direction_dropdown, face_analyser_age_dropdown, face_analyser_gender_dropdown],
+            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=clear_output, inputs=unique_id)
-
-    with gr.Tab("Video:               "):
+        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.Video(label="TARGET VIDEO")
+        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.Video(label="OUTPUT")
+        video_output = gr.Files(label="OUTPUT")
         clear_video_button.add(video_output)
         video_button.click(
-            run,
-            inputs=[source_image_video, target_video, unique_id, frame_processor_checkbox, face_analyser_direction_dropdown, face_analyser_age_dropdown, face_analyser_gender_dropdown, skip_audio, keep_fps, keep_temp],
+            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=clear_output, inputs=unique_id)
+        clear_video_button.click(fn=app.clear_output, inputs=app.unique_id)
 
 ui.launch(debug=True)