run.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import configparser
import datetime as dt
import logging
import sys
import time
from enum import Enum
from pathlib import Path
from threading import Lock, Thread
from typing import List, Optional, Tuple

import imageio
import numpy as np
import torch
import typer
from humanize.time import precisedelta
from imageio.plugins.ffmpeg import FfmpegFormat
from rich import get_console, print
from rich.logging import RichHandler
from rich.markdown import Markdown
from rich.progress import (
    BarColumn,
    Progress,
    ProgressColumn,
    Task,
    TaskID,
    TimeRemainingColumn,
)
from rich.text import Text

from upscale import AlphaOptions, SeamlessOptions, Upscale
from utils.video import FrameDiff, are_same_imgs, find_scenes, get_diff_frame, get_frame

# from rich.traceback import install as install_traceback

# install_traceback()

app = typer.Typer()


class DeinterpaintOptions(str, Enum):
    even = "even"
    odd = "odd"


class FpsSpeedColumn(ProgressColumn):
    """Renders human readable FPS speed."""

    def render(self, task: Task) -> Text:
        speed = task.finished_speed or task.speed
        if speed is None:
            return Text("? FPS", style="progress.data.speed")
        return Text(f"{speed:.2f} FPS", style="progress.data.speed")


@app.command()
def image(
    models: List[str] = typer.Argument(...),
    input: Path = typer.Option(Path("input"), "--input", "-i", help="Input folder"),
    reverse: bool = typer.Option(False, "--reverse", "-r", help="Reverse Order"),
    output: Path = typer.Option(Path("output"), "--output", "-o", help="Output folder"),
    skip_existing: bool = typer.Option(
        False,
        "--skip-existing",
        "-se",
        help="Skip existing output files",
    ),
    delete_input: bool = typer.Option(
        False,
        "--delete-input",
        "-di",
        help="Delete input files after upscaling",
    ),
    seamless: SeamlessOptions = typer.Option(
        None,
        "--seamless",
        "-s",
        case_sensitive=False,
        help="Helps seamlessly upscale an image. tile = repeating along edges. mirror = reflected along edges. replicate = extended pixels along edges. alpha_pad = extended alpha border.",
    ),
    cpu: bool = typer.Option(False, "--cpu", "-c", help="Use CPU instead of CUDA"),
    fp16: bool = typer.Option(
        False,
        "--floating-point-16",
        "-fp16",
        help="Use FloatingPoint16/Halftensor type for images.",
    ),
    device_id: int = typer.Option(
        0, "--device-id", "-did", help="The numerical ID of the GPU you want to use."
    ),
    multi_gpu: bool = typer.Option(False, "--multi-gpu", "-mg", help="Multi GPU"),
    cache_max_split_depth: bool = typer.Option(
        False,
        "--cache-max-split-depth",
        "-cmsd",
        help="Caches the maximum recursion depth used by the split/merge function. Useful only when upscaling images of the same size.",
    ),
    binary_alpha: bool = typer.Option(
        False,
        "--binary-alpha",
        "-ba",
        help="Whether to use a 1 bit alpha transparency channel, Useful for PSX upscaling",
    ),
    ternary_alpha: bool = typer.Option(
        False,
        "--ternary-alpha",
        "-ta",
        help="Whether to use a 2 bit alpha transparency channel, Useful for PSX upscaling",
    ),
    alpha_threshold: float = typer.Option(
        0.5,
        "--alpha-threshold",
        "-at",
        help="Only used when binary_alpha is supplied. Defines the alpha threshold for binary transparency",
    ),
    alpha_boundary_offset: float = typer.Option(
        0.2,
        "--alpha-boundary-offset",
        "-abo",
        help="Only used when binary_alpha is supplied. Determines the offset boundary from the alpha threshold for half transparency.",
    ),
    alpha_mode: AlphaOptions = typer.Option(
        "alpha_separately",
        "--alpha-mode",
        "-am",
        help="Type of alpha processing to use. no_alpha = is no alpha processing. bas = is BA's difference method. alpha_separately = is upscaling the alpha channel separately (like IEU). swapping = is swapping an existing channel with the alpha channel.",
    ),
    imagemagick: bool = typer.Option(
        False,
        "--imagemagick",
        "-im",
        help="Use ImageMagick to save the upscaled image (higher quality but slower). Disabled when using multi_gpu mode.",
    ),
    jpg: bool = typer.Option(
        False,
        "--jpg",
        "-j",
        help="Convert the image to jpg",
    ),
    resize: int = typer.Option(
        100,
        "--resize",
        "-r",
        help="Resize percent",
    ),
    zip: bool = typer.Option(
        False,
        "--zip",
        "-z",
        help="Compress the output to zip file",
    ),
    verbose: bool = typer.Option(
        False,
        "--verbose",
        "-v",
        help="Verbose mode",
    ),
):
    logging.basicConfig(
        level=logging.DEBUG if verbose else logging.WARNING,
        format="%(message)s",
        datefmt="[%X]",
        handlers=[RichHandler(markup=True)],
    )
    start_time = time.process_time()
    for model in models:
        upscale = Upscale(
            model=model,
            seamless=seamless,
            cpu=cpu,
            fp16=fp16,
            device_id=device_id,
            multi_gpu=multi_gpu,
            cache_max_split_depth=cache_max_split_depth,
            binary_alpha=binary_alpha,
            ternary_alpha=ternary_alpha,
            alpha_threshold=alpha_threshold,
            alpha_boundary_offset=alpha_boundary_offset,
            alpha_mode=alpha_mode,
            imagemagick=imagemagick,
            jpg=jpg,
            resize=resize,
            zip=zip,
        )
        models_str = model.split("+") if "+" in model else model.split(">")
        upscale.folder(
            input=input,
            output=output
            if len(models) == 1 or zip
            else output.joinpath("_".join([Path(x).stem for x in models_str])),
            skip_existing=skip_existing,
            reverse=reverse,
            delete_input=delete_input,
        )
    log = logging.getLogger()
    log.info(
        f"Images upscaled in {precisedelta(dt.timedelta(seconds=time.process_time() - start_time))}"
    )


def video_thread_func(
    device: torch.device,
    num_lock: int,
    multi_gpu: bool,
    input: Path,
    start_frame: int,
    end_frame: int,
    num_frames: int,
    progress: Progress,
    task_upscaled_id: TaskID,
    ai_upscaled_path: Path,
    fps: int,
    quality: float,
    ffmpeg_params: str,
    deinterpaint: DeinterpaintOptions,
    diff_mode: bool,
    ssim: bool,
    min_ssim: float,
    chunk_size: int,
    padding_size: int,
    scale: int,
    upscale: Upscale,
    config: configparser.ConfigParser,
    scenes_ini: Path,
):
    log = logging.getLogger()
    video_reader: FfmpegFormat.Reader = imageio.get_reader(str(input.absolute()))
    start_time = time.process_time()
    last_frame = None
    last_frame_ai = None
    current_frame = None
    frames_diff: List[Optional[FrameDiff]] = []
    video_reader.set_image_index(start_frame - 1)
    start_frame_str = str(start_frame).zfill(len(str(num_frames)))
    end_frame_str = str(end_frame).zfill(len(str(num_frames)))
    task_scene_desc = f'Scene [green]"{start_frame_str}_{end_frame_str}"[/]'
    if multi_gpu and len(upscale.devices) > 1:
        if device.type == "cuda":
            device_name = torch.cuda.get_device_name(device.index)
        else:
            device_name = "CPU"
        task_scene_desc += f" ({device_name})"
    task_scene_id = progress.add_task(
        description=task_scene_desc,
        total=end_frame - start_frame + 1,
        completed=0,
        refresh=True,
    )
    video_writer_params = {"quality": quality, "macro_block_size": None}
    if ffmpeg_params:
        if "-crf" in ffmpeg_params:
            del video_writer_params["quality"]
        video_writer_params["output_params"] = ffmpeg_params.split()
    video_writer: FfmpegFormat.Writer = imageio.get_writer(
        str(
            ai_upscaled_path.joinpath(
                f"{start_frame_str}_{end_frame_str}.mp4"
            ).absolute()
        ),
        fps=fps,
        **video_writer_params,
    )
    duplicated_frames = 0
    total_duplicated_frames = 0
    for current_frame_idx in range(start_frame, end_frame + 1):
        frame = video_reader.get_next_data()

        if deinterpaint is not None:
            for i in range(
                0 if deinterpaint == DeinterpaintOptions.even else 1, frame.shape[0], 2
            ):
                frame[i : i + 1] = (0, 255, 0)  # (B, G, R)

        if not diff_mode:
            if last_frame is not None and are_same_imgs(
                last_frame, frame, ssim, min_ssim
            ):
                frame_ai = last_frame_ai
                if duplicated_frames == 0:
                    start_duplicated_frame = current_frame_idx - 1
                duplicated_frames += 1
            else:
                frame_ai = upscale.image(frame, device, multi_gpu_release_device=False)
                if duplicated_frames != 0:
                    start_duplicated_frame_str = str(start_duplicated_frame).zfill(
                        len(str(num_frames))
                    )
                    current_frame_idx_str = str(current_frame_idx - 1).zfill(
                        len(str(num_frames))
                    )
                    log.info(
                        f"Detected {duplicated_frames} duplicated frame{'' if duplicated_frames==1 else 's'} ({start_duplicated_frame_str}-{current_frame_idx_str})"
                    )
                    total_duplicated_frames += duplicated_frames
                    duplicated_frames = 0
            video_writer.append_data(frame_ai)
            last_frame = frame
            last_frame_ai = frame_ai
            progress.advance(task_upscaled_id)
            progress.advance(task_scene_id)
        else:
            if current_frame is None:
                current_frame = frame
            else:
                frame_diff = get_diff_frame(
                    current_frame, frame, chunk_size, padding_size, ssim, min_ssim
                )
                if (
                    frame_diff is None
                ):  # the frame is equal to current_frame, the best scenario!!!
                    frames_diff.append(frame_diff)
                else:
                    h_diff, w_diff, c_diff = frame_diff.frame.shape
                    h, w, c = current_frame.shape
                    if w * h > w_diff * h_diff:  # TODO difference of size > 20%
                        frames_diff.append(frame_diff)
                    else:
                        current_frame_ai = upscale.image(
                            current_frame, device, multi_gpu_release_device=False
                        )
                        video_writer.append_data(current_frame_ai)
                        progress.advance(task_upscaled_id)
                        progress.advance(task_scene_id)
                        current_frame = frame
                        for frame_diff in frames_diff:
                            if frame_diff is None:
                                frame_ai = current_frame_ai
                            else:
                                diff_ai = upscale.image(
                                    frame_diff.frame,
                                    device,
                                    multi_gpu_release_device=False,
                                )
                                frame_diff_ai = frame_diff
                                frame_diff_ai.frame = diff_ai
                                frame_ai = get_frame(
                                    current_frame_ai,
                                    frame_diff_ai,
                                    scale,
                                    chunk_size,
                                    padding_size,
                                )
                            video_writer.append_data(frame_ai)
                            progress.advance(task_upscaled_id)
                            progress.advance(task_scene_id)
                        frames_diff = []
    if diff_mode:
        if len(frames_diff) > 0:
            current_frame_ai = upscale.image(
                current_frame, device, multi_gpu_release_device=False
            )
            video_writer.append_data(current_frame_ai)
            progress.advance(task_upscaled_id)
            progress.advance(task_scene_id)
            for frame_diff in frames_diff:
                if frame_diff is None:
                    frame_ai = current_frame
                else:
                    diff_ai = upscale.image(
                        frame_diff.frame, device, multi_gpu_release_device=False
                    )
                    frame_diff_ai = frame_diff
                    frame_diff_ai.frame = diff_ai
                    frame_ai = get_frame(
                        current_frame_ai,
                        frame_diff_ai,
                        scale,
                        chunk_size,
                        padding_size,
                    )
                video_writer.append_data(frame_ai)
                progress.advance(task_upscaled_id)
                progress.advance(task_scene_id)
            current_frame = None
            frames_diff = []
        elif current_frame is not None:
            current_frame_ai = upscale.image(
                current_frame, device, multi_gpu_release_device=False
            )
            video_writer.append_data(current_frame_ai)
            progress.advance(task_upscaled_id)
            progress.advance(task_scene_id)
    if duplicated_frames != 0:
        start_duplicated_frame_str = str(start_duplicated_frame).zfill(
            len(str(num_frames))
        )
        current_frame_idx_str = str(current_frame_idx - 1).zfill(len(str(num_frames)))
        log.info(
            f"Detected {duplicated_frames} duplicated frame{'' if duplicated_frames==1 else 's'} ({start_duplicated_frame_str}-{current_frame_idx_str})"
        )
        total_duplicated_frames += duplicated_frames
        duplicated_frames = 0

    video_writer.close()
    task_scene = next(task for task in progress.tasks if task.id == task_scene_id)

    config.set(f"{start_frame_str}_{end_frame_str}", "upscaled", "True")
    config.set(
        f"{start_frame_str}_{end_frame_str}",
        "duplicated_frames",
        f"{total_duplicated_frames}",
    )
    finished_speed = task_scene.finished_speed or task_scene.speed or 0.01
    config.set(
        f"{start_frame_str}_{end_frame_str}",
        "average_fps",
        f"{finished_speed:.2f}",
    )
    with open(scenes_ini, "w") as configfile:
        config.write(configfile)
    log.info(
        f"Frames from {str(start_frame).zfill(len(str(num_frames)))} to {str(end_frame).zfill(len(str(num_frames)))} upscaled in {precisedelta(dt.timedelta(seconds=time.process_time() - start_time))}"
    )
    if total_duplicated_frames > 0:
        total_frames = end_frame - (start_frame - 1)
        seconds_saved = (
            (
                (1 / finished_speed * total_frames)
                - (total_duplicated_frames * 0.04)  # 0.04 seconds per duplicate frame
            )
            / (total_frames - total_duplicated_frames)
            * total_duplicated_frames
        )
        log.info(
            f"Total number of duplicated frames from {str(start_frame).zfill(len(str(num_frames)))} to {str(end_frame).zfill(len(str(num_frames)))}: {total_duplicated_frames} (saved ≈ {precisedelta(dt.timedelta(seconds=seconds_saved))})"
        )
    progress.remove_task(task_scene_id)
    if multi_gpu:
        upscale.devices[device][num_lock].release()


@app.command()
def video(
    model: str = typer.Argument(...),
    input: Path = typer.Option(
        Path("input/video.mp4"), "--input", "-i", help="Input video"
    ),
    output: Path = typer.Option(
        Path("output/video.mp4"), "--output", "-o", help="Output video"
    ),
    seamless: SeamlessOptions = typer.Option(
        None,
        "--seamless",
        "-s",
        case_sensitive=False,
        help="Helps seamlessly upscale an image. tile = repeating along edges. mirror = reflected along edges. replicate = extended pixels along edges. alpha_pad = extended alpha border.",
    ),
    # cpu: bool = typer.Option(False, "--cpu", "-c", help="Use CPU instead of CUDA"),
    fp16: bool = typer.Option(
        False,
        "--floating-point-16",
        "-fp16",
        help="Use FloatingPoint16/Halftensor type for images.",
    ),
    device_id: int = typer.Option(
        0, "--device-id", "-did", help="The numerical ID of the GPU you want to use."
    ),
    multi_gpu: bool = typer.Option(False, "--multi-gpu", "-mg", help="Multi GPU"),
    scenes_per_gpu: int = typer.Option(
        1,
        "--scenes-per-gpu",
        "-spg",
        help="Number of scenes to be upscaled at the same time using the same GPU. 0 for automatic mode",
    ),
    cache_max_split_depth: bool = typer.Option(
        False,
        "--cache-max-split-depth",
        "-cmsd",
        help="Caches the maximum recursion depth used by the split/merge function. Useful only when upscaling images of the same size.",
    ),
    ssim: bool = typer.Option(
        False,
        "--ssim",
        "-ssim",
        help="True to enable duplication frame removal using ssim. False to use np.all().",
    ),
    min_ssim: float = typer.Option(0.9987, "--min-ssim", "-ms", help="Min SSIM value."),
    diff_mode: bool = typer.Option(
        False, "--diff", "-d", help="Enable diff mode (beta)."
    ),
    chunk_size: int = typer.Option(
        16,
        "--chunk-size",
        "-cs",
        help="Only used with diff mode. Chunk size to be able to generate the frame difference (beta).",
    ),
    padding_size: int = typer.Option(
        2,
        "--padding-size",
        "-ps",
        help="Only used with diff mode. Padding size between each chunk (beta).",
    ),
    quality: float = typer.Option(
        10,
        "--quality",
        "-q",
        min=0,
        max=10,
        help="Video quality.",
    ),
    ffmpeg_params: str = typer.Option(
        None,
        "--ffmpeg-params",
        "--ffmpeg",
        help='FFmpeg parameters to save the scenes. If -crf is present, the quality parameter will be ignored. Example: "-c:v libx265 -crf 5 -pix_fmt yuv444p10le -preset medium -x265-params pools=none -threads 8".',
    ),
    # deduplication: bool = typer.Option(
    #     False,
    #     "--deduplication",
    #     "-d",
    #     help="True to enable duplication frame removal",
    # ),
    deinterpaint: DeinterpaintOptions = typer.Option(
        None,
        "--deinterpaint",
        "-dp",
        case_sensitive=False,
        help="De-interlacing by in-painting. Fills odd or even rows with green (#00FF00). Useful for models like Joey's 1x_DeInterPaint.",
    ),
    verbose: bool = typer.Option(
        False,
        "--verbose",
        "-v",
        help="Verbose mode",
    ),
):
    logging.basicConfig(
        level=logging.DEBUG if verbose else logging.WARNING,
        format="%(message)s",
        datefmt="[%X]",
        handlers=[RichHandler(markup=True)],
    )
    log = logging.getLogger()

    input = input.resolve()
    output = output.resolve()
    if not input.exists():
        log.error(f'Video "{input}" does not exist.')
        sys.exit(1)
    elif input.is_dir():
        log.error(f'Video "{input}" is a directory.')
        sys.exit(1)
    elif output.is_dir():
        log.error(f'Video "{output}" is a directory.')
        sys.exit(1)
    # elif not output.exists():
    #     output=input.with_name(f"{input.stem}_ai.mp4")

    upscale = Upscale(
        model=model,
        seamless=seamless,
        # cpu=cpu,
        fp16=fp16,
        device_id=device_id,
        cache_max_split_depth=cache_max_split_depth,
        alpha_mode=AlphaOptions.no_alpha,
        multi_gpu=multi_gpu,
    )
    if len(upscale.model_chain) > 1 and deinterpaint is not None:
        log.error("Model Chain and DeInterPaint cannot be used at the same time.")
        exit(1)

    project_path = output.parent.joinpath(f"{output.stem}").absolute()
    ai_upscaled_path = project_path.joinpath("scenes")
    scenes_ini = project_path.joinpath("scenes.ini")
    frames_todo: List[Tuple[int, int]] = []
    frames_upscaled: List[Tuple[int, int]] = []
    config = configparser.ConfigParser()
    if project_path.is_dir():
        resume_mode = True
        log.info(f'Resuming project "{project_path}"')
        config.read(scenes_ini)
        for scene in config.sections():
            start_frame, end_frame = scene.split("_")
            start_frame = int(start_frame)
            end_frame = int(end_frame)
            if config.getboolean(scene, "upscaled") == True:
                frames_upscaled.append((start_frame, end_frame))
            else:
                frames_todo.append((start_frame, end_frame))
    else:
        resume_mode = False
        with get_console().status("Detecting scenes..."):
            scenes = find_scenes(str(input.absolute()))
        log.info(f"Detected {len(scenes)} scene{'' if len(scenes)==1 else 's'}")

        ai_upscaled_path.mkdir(parents=True, exist_ok=True)
        num_frames = scenes[-1][1].get_frames()
        for scene in scenes:
            start_frame = str(scene[0].get_frames() + 1).zfill(len(str(num_frames)))
            end_frame = str(scene[1].get_frames()).zfill(len(str(num_frames)))
            config[f"{start_frame}_{end_frame}"] = {
                "upscaled": "False",
                "duplicated_frames": "None",
                "average_fps": "None",
            }
            frames_todo.append((int(start_frame), int(end_frame)))

        with open(scenes_ini, "w") as configfile:
            config.write(configfile)

    video_reader: FfmpegFormat.Reader = imageio.get_reader(str(input.absolute()))
    fps = video_reader.get_meta_data()["fps"]
    num_frames = video_reader.count_frames()
    scale = 1

    if diff_mode:
        with get_console().status("Detecting the model's scale..."):
            img = video_reader.get_data(0)
            h, w, c = img.shape
            img = np.resize(img, (h // 4, w // 4, c))  # resize for fast upscaling
            height, width, channels = img.shape
            height_ai, width_ai, channels_ai = upscale.image(img).shape
            scale = int(width_ai / width)
        log.info(f"Model's scale: x{scale}")

    if scenes_per_gpu < 1:
        with get_console().status(
            f"Detecting how many scenes can be upscaled at the same time..."
        ):
            img = video_reader.get_data(0)
            upscale.image(img)
        reserved = torch.cuda.memory_reserved(device_id)
        if multi_gpu:
            for i in range(torch.cuda.device_count()):
                device = torch.device(f"cuda:{i}")
                device_name = torch.cuda.get_device_name(i)
                total = torch.cuda.get_device_properties(i).total_memory
                # TODO upscale using the device i
                num_scenes_same_time = 0
                reserved_temp = 0
                while reserved_temp < total:
                    reserved_temp += reserved
                    num_scenes_same_time += 1
                if reserved_temp >= total:
                    num_scenes_same_time -= 1
                log.info(
                    f'Number of scenes to upscale at the same time on "{device_name}": {num_scenes_same_time}'
                )
                upscale.devices[device] = [Lock() for _ in range(num_scenes_same_time)]
        else:
            device = torch.device(f"cuda:{device_id}")
            device_name = torch.cuda.get_device_name(device_id)
            total = torch.cuda.get_device_properties(device_id).total_memory
            num_scenes_same_time = 0
            reserved_temp = 0
            while reserved_temp < total:
                reserved_temp += reserved
                num_scenes_same_time += 1
            if reserved_temp >= total:
                num_scenes_same_time -= 1
            log.info(
                f'Number of scenes to upscale at the same time on "{device_name}": {num_scenes_same_time}'
            )
            upscale.devices[device] = [Lock() for _ in range(num_scenes_same_time)]
            if num_scenes_same_time > 1:
                multi_gpu = True
                upscale.multi_gpu = True
    else:
        for device in upscale.devices.keys():
            upscale.devices[device] = [Lock() for _ in range(scenes_per_gpu)]
        if scenes_per_gpu > 1:
            multi_gpu = True
            upscale.multi_gpu = True

    with Progress(
        # SpinnerColumn(),
        "[progress.description]{task.description}",
        "[progress.percentage]{task.percentage:>3.0f}%",
        BarColumn(),
        TimeRemainingColumn(),
        FpsSpeedColumn(),
    ) as progress:
        num_frames_upscaled = 0
        for start_frame, end_frame in frames_upscaled:
            num_frames_upscaled += end_frame - start_frame + 1
        task_upscaled_id = progress.add_task(
            f'Upscaling [green]"{input.name}"[/]', total=num_frames
        )
        if num_frames_upscaled > 0:
            log.info(f"Skipped {num_frames_upscaled} frames already upscaled")
            progress.update(
                task_upscaled_id, completed=num_frames_upscaled, refresh=True
            )

        if len(upscale.model_chain) > 1:
            # Fix model chain (because the models were not loaded if the threads start at the same time)
            upscale.image(255 * np.zeros([10, 10, 3], dtype=np.uint8))

        threads = []
        for start_frame, end_frame in frames_todo:
            num_lock = 0
            if multi_gpu:
                device, num_lock = upscale.get_available_device(first_lock=False)
            else:
                device = list(upscale.devices.keys())[0]
            video_thread_func_args = {
                "device": device,
                "num_lock": num_lock,
                "multi_gpu": multi_gpu,
                "input": input,
                "start_frame": start_frame,
                "end_frame": end_frame,
                "num_frames": num_frames,
                "progress": progress,
                "task_upscaled_id": task_upscaled_id,
                "ai_upscaled_path": ai_upscaled_path,
                "fps": fps,
                "quality": quality,
                "ffmpeg_params": ffmpeg_params,
                "deinterpaint": deinterpaint,
                "diff_mode": diff_mode,
                "ssim": ssim,
                "min_ssim": min_ssim,
                "chunk_size": chunk_size,
                "padding_size": padding_size,
                "scale": scale,
                "upscale": upscale,
                "config": config,
                "scenes_ini": scenes_ini,
            }
            if multi_gpu:
                x = Thread(target=video_thread_func, kwargs=video_thread_func_args)
                threads.append(x)
                x.daemon = True
                x.start()
            else:
                video_thread_func(**video_thread_func_args)
        for thread in threads:
            thread.join()

    with open(
        project_path.joinpath("ffmpeg_list.txt"), "w", encoding="utf-8"
    ) as outfile:
        for mp4_path in ai_upscaled_path.glob("*.mp4"):
            outfile.write(f"file '{mp4_path.relative_to(project_path).as_posix()}'\n")
    total_duplicated_frames = 0
    total_average_fps = 0
    for section in config.sections():
        total_duplicated_frames += config.getint(section, "duplicated_frames")
        total_average_fps += config.getfloat(section, "average_fps")
    total_average_fps = total_average_fps / len(config.sections())
    if not resume_mode:
        task_upscaled = next(
            task for task in progress.tasks if task.id == task_upscaled_id
        )
        total_average_fps = task_upscaled.finished_speed or task_upscaled.speed or 0.01
    if total_duplicated_frames > 0:
        seconds_saved = (
            (
                (1 / total_average_fps * num_frames)
                - (total_duplicated_frames * 0.04)  # 0.04 seconds per duplicate frame
            )
            / (num_frames - total_duplicated_frames)
            * total_duplicated_frames
        )
        log.info(
            f"Total number of duplicated frames: {total_duplicated_frames} (saved ≈ {precisedelta(dt.timedelta(seconds=seconds_saved))})"
        )
    log.info(f"Total FPS: {total_average_fps:.2f}")
    print("\nUpscale completed!\n")

    bad_scenes = []
    with get_console().status(
        "Checking the correct number of frames of the mp4 files..."
    ):
        for mp4_path in ai_upscaled_path.glob("*.mp4"):
            start_frame, end_frame = mp4_path.stem.split("_")
            num_frames = int(end_frame) - int(start_frame) + 1
            with imageio.get_reader(str(mp4_path.absolute())) as video_reader:
                frames_mp4 = video_reader.count_frames()
            if num_frames != frames_mp4:
                bad_scenes.append(f"{mp4_path.stem}")

    if len(bad_scenes) > 0:
        for scene in bad_scenes:
            config.set(scene, "upscaled", "False")
        with open(scenes_ini, "w") as configfile:
            config.write(configfile)
        if len(bad_scenes) == 1:
            bad_scenes_str = f"[green]{bad_scenes[0]}[/]"
        else:
            bad_scenes_str = f'[green]{"[/], [green]".join(bad_scenes[:-1])}[/] and [green]{bad_scenes[-1]}[/]'
        print(f"The following scenes were incorrectly upscaled: {bad_scenes_str}.")
        print(f"Please re-run the script to finish upscaling them.")
    else:
        print(
            f'Go to the "{project_path}" directory and run the following command to concatenate the scenes.'
        )
        print(
            Markdown(
                f"`ffmpeg -f concat -safe 0 -i ffmpeg_list.txt -i {input.absolute()} -map 0:v -map 1:a -c copy {output.name}`"
            )
        )


if __name__ == "__main__":
    app()