ui/tabs/extras_tab.py

import os
import gradio as gr
import shutil
import roop.utilities as util
import roop.util_ffmpeg as ffmpeg
import roop.globals
from roop.utilities import clean_dir

frame_filters_map = {
    "Colorize B/W Images (Deoldify Artistic)": {
        "colorizer": {"subtype": "deoldify_artistic"}
    },
    "Colorize B/W Images (Deoldify Stable)": {
        "colorizer": {"subtype": "deoldify_stable"}
    },
    "Background remove": {"removebg": {"subtype": ""}},
    "Filter Stylize": {"filter_generic": {"subtype": "stylize"}},
    "Filter Detail Enhance": {"filter_generic": {"subtype": "detailenhance"}},
    "Filter Pencil Sketch": {"filter_generic": {"subtype": "pencil"}},
    "Filter Cartoon": {"filter_generic": {"subtype": "cartoon"}},
    "Filter C64": {"filter_generic": {"subtype": "C64"}},
}

frame_upscalers_map = {
    "ESRGAN x2": {"upscale": {"subtype": "esrganx2"}},
    "ESRGAN x4": {"upscale": {"subtype": "esrganx4"}},
    "LSDIR x4": {"upscale": {"subtype": "lsdirx4"}},
}


def extras_tab():
    filternames = ["None"]
    for f in frame_filters_map.keys():
        filternames.append(f)
    upscalernames = ["None"]
    for f in frame_upscalers_map.keys():
        upscalernames.append(f)

    with gr.Tab("🎉 Extras"):
        with gr.Row():
            files_to_process = gr.Files(
                label="File(s) to process",
                file_count="multiple",
                file_types=["image", "video"],
            )
        with gr.Row(variant="panel"):
            with gr.Accordion(label="Video/GIF", open=False):
                with gr.Row(variant="panel"):
                    with gr.Column():
                        gr.Markdown("""
                                    # Poor man's video editor
                                    Re-encoding uses your configuration from the Settings Tab.
    """)
                    with gr.Column():
                        cut_start_time = gr.Slider(
                            0,
                            1000000,
                            value=0,
                            label="Start Frame",
                            step=1.0,
                            interactive=True,
                        )
                    with gr.Column():
                        cut_end_time = gr.Slider(
                            1,
                            1000000,
                            value=1,
                            label="End Frame",
                            step=1.0,
                            interactive=True,
                        )
                    with gr.Column():
                        extras_chk_encode = gr.Checkbox(
                            label="Re-encode videos (necessary for videos with different codecs)",
                            value=False,
                        )
                        start_cut_video = gr.Button("Cut video")
                        start_extract_frames = gr.Button("Extract frames")
                        start_join_videos = gr.Button("Join videos")

                with gr.Row(variant="panel"):
                    with gr.Column():
                        gr.Markdown("""
                                    # Create video/gif from images
    """)
                    with gr.Column():
                        extras_fps = gr.Slider(
                            minimum=0,
                            maximum=120,
                            value=30,
                            label="Video FPS",
                            step=1.0,
                            interactive=True,
                        )
                        extras_images_folder = gr.Textbox(
                            show_label=False, placeholder="/content/", interactive=True
                        )
                    with gr.Column():
                        extras_chk_creategif = gr.Checkbox(
                            label="Create GIF from video", value=False
                        )
                        extras_create_video = gr.Button("Create")
                with gr.Row(variant="panel"):
                    with gr.Column():
                        gr.Markdown("""
                                    # Create video from gif
    """)
                    with gr.Column():
                        extras_video_fps = gr.Slider(
                            minimum=0,
                            maximum=120,
                            value=0,
                            label="Video FPS",
                            step=1.0,
                            interactive=True,
                        )
                    with gr.Column():
                        extras_create_video_from_gif = gr.Button("Create")
                with gr.Row(variant="panel"):
                    with gr.Column(scale=2):
                        gr.Markdown("""
                                    # Repair video

                                    Uses FFMpeg to fix corrupt videos. 
    """)
                    with gr.Column():
                        extras_repair_video = gr.Button("Repair")

        with gr.Row(variant="panel"):
            with gr.Accordion(label="Full frame processing", open=True):
                with gr.Row(variant="panel"):
                    filterselection = gr.Dropdown(
                        filternames,
                        value="None",
                        label="Colorizer/FilterFX",
                        interactive=True,
                    )
                    upscalerselection = gr.Dropdown(
                        upscalernames, value="None", label="Enhancer", interactive=True
                    )
                with gr.Row(variant="panel"):
                    start_frame_process = gr.Button("Start processing")

        with gr.Row():
            gr.Button("👀 Open Output Folder", size="sm").click(
                fn=lambda: util.open_folder(roop.globals.output_path)
            )
        with gr.Row():
            extra_files_output = gr.Files(
                label="Resulting output files", file_count="multiple"
            )

    start_cut_video.click(
        fn=on_cut_video,
        inputs=[files_to_process, cut_start_time, cut_end_time, extras_chk_encode],
        outputs=[extra_files_output],
    )
    start_extract_frames.click(
        fn=on_extras_extract_frames,
        inputs=[files_to_process],
        outputs=[extra_files_output],
    )
    start_join_videos.click(
        fn=on_join_videos,
        inputs=[files_to_process, extras_chk_encode],
        outputs=[extra_files_output],
    )
    extras_create_video.click(
        fn=on_extras_create_video,
        inputs=[
            files_to_process,
            extras_images_folder,
            extras_fps,
            extras_chk_creategif,
        ],
        outputs=[extra_files_output],
    )
    extras_create_video_from_gif.click(
        fn=on_extras_create_video_from_gif,
        inputs=[files_to_process, extras_video_fps],
        outputs=[extra_files_output],
    )
    extras_repair_video.click(
        fn=on_extras_repair_video,
        inputs=[files_to_process],
        outputs=[extra_files_output],
    )
    start_frame_process.click(
        fn=on_frame_process,
        inputs=[files_to_process, filterselection, upscalerselection],
        outputs=[extra_files_output],
    )


def on_cut_video(files, cut_start_frame, cut_end_frame, reencode):
    if files is None:
        return None

    resultfiles = []
    for tf in files:
        f = tf.name
        destfile = util.get_destfilename_from_path(f, roop.globals.output_path, "_cut")
        ffmpeg.cut_video(f, destfile, cut_start_frame, cut_end_frame, reencode)
        if os.path.isfile(destfile):
            resultfiles.append(destfile)
        else:
            gr.Error("Cutting video failed!")
    return resultfiles


def on_join_videos(files, chk_encode):
    if files is None:
        return None

    filenames = []
    for f in files:
        filenames.append(f.name)
    destfile = util.get_destfilename_from_path(
        filenames[0], roop.globals.output_path, "_join"
    )
    sorted_filenames = util.sort_filenames_ignore_path(filenames)
    ffmpeg.join_videos(sorted_filenames, destfile, not chk_encode)
    resultfiles = []
    if os.path.isfile(destfile):
        resultfiles.append(destfile)
    else:
        gr.Error("Joining videos failed!")
    return resultfiles


def on_extras_create_video_from_gif(files, fps):
    if files is None:
        return None

    filenames = []
    resultfiles = []
    for f in files:
        filenames.append(f.name)

    destfilename = os.path.join(
        roop.globals.output_path, "img2video." + roop.globals.CFG.output_video_format
    )
    ffmpeg.create_video_from_gif(filenames[0], destfilename)
    if os.path.isfile(destfilename):
        resultfiles.append(destfilename)
    return resultfiles


def on_extras_repair_video(files):
    if files is None:
        return None

    resultfiles = []
    for tf in files:
        f = tf.name
        destfile = util.get_destfilename_from_path(
            f, roop.globals.output_path, "_repair"
        )
        ffmpeg.repair_video(f, destfile)
        if os.path.isfile(destfile):
            resultfiles.append(destfile)
        else:
            gr.Error("Repairing video failed!")
    return resultfiles


def on_extras_create_video(files, images_path, fps, create_gif):
    if images_path is None:
        return None
    resultfiles = []
    if len(files) > 0 and util.is_video(files[0]) and create_gif:
        destfilename = files[0]
    else:
        util.sort_rename_frames(os.path.dirname(images_path))
        destfilename = os.path.join(
            roop.globals.output_path,
            "img2video." + roop.globals.CFG.output_video_format,
        )
        ffmpeg.create_video("", destfilename, fps, images_path)
        if os.path.isfile(destfilename):
            resultfiles.append(destfilename)
        else:
            return None
    if create_gif:
        gifname = util.get_destfilename_from_path(destfilename, "./output", ".gif")
        ffmpeg.create_gif_from_video(destfilename, gifname)
        if os.path.isfile(destfilename):
            resultfiles.append(gifname)
    return resultfiles


def on_extras_extract_frames(files):
    if files is None:
        return None

    resultfiles = []
    for tf in files:
        f = tf.name
        resfolder = ffmpeg.extract_frames(f)
        for file in os.listdir(resfolder):
            outfile = os.path.join(resfolder, file)
            if os.path.isfile(outfile):
                resultfiles.append(outfile)
    return resultfiles


def on_frame_process(files, filterselection, upscaleselection):
    import pathlib
    from roop.core import batch_process_with_options
    from roop.ProcessEntry import ProcessEntry
    from roop.ProcessOptions import ProcessOptions
    from ui.main import prepare_environment

    if files is None:
        return None

    if roop.globals.CFG.clear_output:
        clean_dir(roop.globals.output_path)
    prepare_environment()
    list_files_process: list[ProcessEntry] = []

    for tf in files:
        list_files_process.append(ProcessEntry(tf.name, 0, 0, 0))

    processoroptions = {}
    filter = next((x for x in frame_filters_map.keys() if x == filterselection), None)
    if filter is not None:
        processoroptions.update(frame_filters_map[filter])
    filter = next(
        (x for x in frame_upscalers_map.keys() if x == upscaleselection), None
    )
    if filter is not None:
        processoroptions.update(frame_upscalers_map[filter])
    options = ProcessOptions(
        None, processoroptions, 0, 0, "all", 0, None, None, 0, 128, False, False
    )
    batch_process_with_options(list_files_process, options, None)
    outdir = pathlib.Path(roop.globals.output_path)
    outfiles = [str(item) for item in outdir.rglob("*") if item.is_file()]
    return outfiles