Общая оптимизация: больше мультипроцессинга и асинка

Author: Olexasha

main.py

@@ -1,17 +1,22 @@
 import asyncio
+import glob
 import os
 import re
 from datetime import datetime
+from math import ceil
 
 from src.audio.audio_handling import AudioHandler
 from src.config.comp_params import ComputerParams
 from src.config.settings import (
+    BATCH_VIDEO_PATH,
     END_BATCH_TO_UPSCALE,
     FINAL_VIDEO,
     INPUT_BATCHES_DIR,
     ORIGINAL_VIDEO,
     START_BATCH_TO_UPSCALE,
+    TMP_VIDEO_PATH,
 )
+from src.files.file_actions import delete_file
 from src.frames.frames_helpers import extract_frames_to_batches, get_fps_accurate
 from src.frames.upscale import delete_frames, upscale_batches
 from src.utils.logger import logger
@@ -26,7 +31,7 @@ def print_header(title: str) -> None:
 
 def print_bottom(title: str) -> None:
     logger.info(f"✅ {title.upper()}".center(50))
-    logger.info(f"{'=' * 50}")
+    logger.info(f"{'=' * 50}\n")
 
 
 async def clean_up(audio: AudioHandler) -> None:
@@ -36,6 +41,12 @@ async def clean_up(audio: AudioHandler) -> None:
         asyncio.to_thread(audio.delete_audio_if_exists),
         asyncio.to_thread(delete_frames, del_upscaled=False),
         asyncio.to_thread(delete_frames, del_upscaled=True),
+        asyncio.to_thread(
+            map, delete_file, glob.glob(os.path.join(BATCH_VIDEO_PATH, "*.mp4"))
+        ),
+        asyncio.to_thread(
+            map, delete_file, glob.glob(os.path.join(TMP_VIDEO_PATH, "*.mp4"))
+        ),
     )
     logger.debug("Временные файлы успешно удалены")
 
@@ -80,13 +91,7 @@ async def process_batches(
         logger.success(f"Батчи {start_batch}-{end_batch} успешно апскейлены")
 
         batches_to_perform = [f"batch_{i}" for i in range(start_batch, end_batch + 1)]
-        short_video = await video.build_short_video(batches_to_perform)
-        if short_video:
-            logger.info(f"Видео собрано: {short_video}")
-            await asyncio.to_thread(delete_frames, del_upscaled=True)
-            logger.debug(
-                f"Обработанные кадры удалены из батчей {start_batch}-{end_batch}"
-            )
+        video.build_short_video(batches_to_perform)
         start_batch += threads
 
 
@@ -104,21 +109,21 @@ async def main():
         logger.info(
             "Параметры системы:"
             f"\n\tОС: {my_computer.cpu_name}"
-            f"\n\tЯдра: {my_computer.cpu_threads}"
+            f"\n\tCPU потоки: {my_computer.cpu_threads}"
             f"\n\tБезопасные потоки: {my_computer.safe_cpu_threads}"
-            f"\n\tСкорость SSD: {my_computer.ssd_speed} MB/s"
-            f"\n\tRAM: {my_computer.ram_total} GB"
+            f"\n\tСкорость SSD: ~{my_computer.ssd_speed} MB/s"
+            f"\n\tRAM: ~{my_computer.ram_total} GB"
             f"\n\tПараметры нейронок: -j {ai_threads}"
             f"\n\tПуть к нейронке апскейла: {ai_realesrgan_path}"
         )
 
         print_header("извлекаем 'сырьё' из видео...")
-        audio = AudioHandler(threads=process_threads)
+        audio = AudioHandler(threads=my_computer.safe_cpu_threads // 2)
         await clean_up(audio)
 
         # запускаем извлечение аудио из видео в фоне
         asyncio.create_task(audio.extract_audio())
-        extract_frames_to_batches(process_threads)
+        extract_frames_to_batches(my_computer.safe_cpu_threads // 2)
         fps = await asyncio.to_thread(get_fps_accurate, ORIGINAL_VIDEO)
         video = VideoHandler(fps=fps)
         print_bottom("`сырьё` из видео извлечено")
@@ -135,11 +140,11 @@ async def main():
             START_BATCH_TO_UPSCALE,
             end_batch_to_upscale,
         )
-        await asyncio.to_thread(delete_frames, del_upscaled=False)
         print_bottom("апскейленные короткие видео сгенерированы")
 
         print_header("начало финальной сборки видео...")
-        final_merge = video.build_final_video()
+        total_short_videos = ceil(end_batch_to_upscale / process_threads)
+        final_merge = await video.build_final_video(total_short_videos)
         logger.success(f"Общее видео собрано: {final_merge}")
 
         logger.info("Добавление аудиодорожки к финальному видео")

poetry.lock

@@ -12,8 +12,6 @@ numpy = "^2.2.6"
 imageio = "^2.37.0"
 python-dotenv = "^1.1.1"
 opencv-python = "^4.12.0.88"
-tqdm = "^4.67.1"
-colorama = "^0.4.6"
 psutil = "^7.0.0"
 logging = "==0.4.9.6"
 

pyproject.toml

@@ -68,7 +68,6 @@ def extract_audio_sync(self) -> Optional[str]:
             "-b:a", self.BITRATE, "-threads", str(self.threads),
             "-loglevel", "error", audio_file,
         ]
-
         try:
             result = subprocess.run(
                 ["ffmpeg", *ffmpeg_args],
@@ -82,6 +81,7 @@ def extract_audio_sync(self) -> Optional[str]:
                 )
             self.__check_audio_extracted(audio_file)
             self.audio_path = audio_file
+            logger.success(f"Аудио успешно извлечено {audio_file}")
         except subprocess.CalledProcessError as e:
             logger.error(f"Ошибка извлечения аудио: {e.stderr.decode()}")
             raise RuntimeError(f"Ошибка извлечения аудио: {e.stderr.decode()}")
@@ -124,7 +124,7 @@ def insert_audio(self) -> None:
         ]
         try:
             run_ffmpeg_command_with_progress(
-                cmd, duration, desc="Добавление аудио к видео", unit="сек"
+                cmd, duration, desc="Импортирование аудиоряда в видео", unit="сек"
             )
         except subprocess.CalledProcessError as e:
             logger.error(f"Ошибка при добавлении аудио: {str(e)}")

src/audio/audio_handling.py

@@ -44,8 +44,8 @@ def run_ffmpeg_command_with_progress(
                         or progress_percent >= 100
                     ):
                         logger.info(
-                            f"{desc}: {progress_percent:.1f}% "
-                            f"({current_time:.1f}/{duration:.1f} {unit})"
+                            f"{desc}: {current_time:.1f}/{duration:.1f}{unit} "
+                            f"({progress_percent:.1f}%)"
                         )
                         last_logged_progress = int(progress_percent)
         process.wait()

src/audio/audio_helpers.py

@@ -64,9 +64,9 @@ def get_optimal_threads(self) -> Tuple[str, int]:
         Рассчитывает оптимальные параметры для -j load:proc:save и количество процессов/тредов.
         """
         processes = self._calculate_processing_threads()
-        load_threads = self._calculate_load_threads()
+        save_threads = self._calculate_save_threads()
         proc_threads = self._calculate_proc_threads(processes)
-        save_threads = max(1, load_threads - 1)
+        load_threads = max(1, save_threads - 1)
 
         # подбираем load:proc:save
         j_params = f"{load_threads}:{proc_threads}:{save_threads}"
@@ -105,15 +105,15 @@ def _estimate_ssd_speed(self) -> float:
         finally:
             delete_file(test_file)
 
-    def _calculate_load_threads(self) -> int:
+    def _calculate_save_threads(self) -> int:
         """Потоки для загрузки (зависит от SSD)."""
         if self.ssd_speed >= 2000:  # NVMe
-            return 3
+            return 2
         return 2 if self.ssd_speed >= 500 else 1  # SATA SSD или HDD
 
     def _calculate_proc_threads(self, processes: int) -> int:
         """Потоки для обработки с учётом количества процессов."""
-        safe_threads = max(2, self.safe_cpu_threads // max(1, processes))
+        safe_threads = max(2, self.safe_cpu_threads // max(1, processes) - 1)
         return min(8, safe_threads)  # Не больше 8 даже для мощных CPU
 
     @staticmethod

src/config/comp_params.py

@@ -9,7 +9,12 @@
 ORIGINAL_VIDEO = os.path.join(ROOT_DIR, "data", "input_video", "naruto_test2.mkv")
 AUDIO_PATH = os.path.join(ROOT_DIR, "data", "audio")
 TMP_VIDEO_PATH = os.path.join(ROOT_DIR, "data", "tmp_video")
-FINAL_VIDEO = os.path.join(ROOT_DIR, "data", "output_video", "final_video.mp4")
+FINAL_VIDEO = os.path.join(
+    ROOT_DIR,
+    "data",
+    "output_video",
+    f"{os.path.splitext(os.path.basename(ORIGINAL_VIDEO))[0]}_enhanced.mp4",
+)
 BATCH_VIDEO_PATH = os.path.join(ROOT_DIR, "data", "video_batches")
 INPUT_BATCHES_DIR = os.path.join(ROOT_DIR, "data", "default_frame_batches")
 OUTPUT_BATCHES_DIR = os.path.join(ROOT_DIR, "data", "upscaled_frame_batches")

src/config/settings.py

@@ -1,9 +1,7 @@
 import os
-import sys
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
 import cv2
-from tqdm import tqdm
 
 from src.config.settings import (
     FRAMES_PER_BATCH,
@@ -67,16 +65,14 @@ def extract_frames_to_batches(
     logger.info(f"Начало извлечения {total_frames} кадров из видео")
     logger.debug(f"Параметры: потоки={threads}, batch_size={batch_size}")
 
-    with tqdm(
-        total=total_frames,
-        desc="Извлечение фреймов",
-        ncols=150,
-        file=sys.stdout,
-    ) as pbar, ThreadPoolExecutor(max_workers=threads) as executor:
+    with ThreadPoolExecutor(max_workers=threads) as executor:
         futures = []
         current_batch_dir = make_default_batch_dir(output_dir)
         logger.debug(f"Создан первый батч: {current_batch_dir}")
 
+        processed_frames = 0
+        last_logged_percent = 0
+
         for frame_num in range(1, total_frames + 1):
             ret, frame = cap.read()
             if not ret:
@@ -97,10 +93,20 @@ def extract_frames_to_batches(
             if len(futures) >= threads * 2:
                 for future in as_completed(futures[:threads]):
                     futures.remove(future)
-                pbar.update(threads)
+                processed_frames += threads
+
+            current_percent = (processed_frames * 100) // total_frames
+            if current_percent >= last_logged_percent + 5:
+                logger.info(
+                    f"Извлечение фреймов: {processed_frames}/{total_frames} "
+                    f"({current_percent}%)"
+                )
+                last_logged_percent = current_percent
 
         for _ in as_completed(futures):
-            pbar.update(1)
+            processed_frames += 1
+
+        logger.info(f"Извлечение фреймов: {total_frames}/{total_frames} " f"(100%)")
 
     cap.release()
     logger.success(f"Успешно извлечено {total_frames} кадров")

src/frames/frames_helpers.py

@@ -2,10 +2,9 @@
 import glob
 import os
 import subprocess
-import sys
+import time
 from concurrent.futures import ProcessPoolExecutor
 from pathlib import Path
-from tqdm import tqdm
 
 from src.config.settings import (
     INPUT_BATCHES_DIR,
@@ -46,32 +45,54 @@ def delete_frames(del_upscaled: bool, del_only_dirs: bool = True):
 
 
 async def monitor_progress(
-    total_frames: int,
-    is_processing: list,
-    batch_numbers: range
+    total_frames: int, is_processing: list, batch_numbers: range
 ) -> None:
     """Мониторинг прогресса с частым обновлением и визуализацией"""
     processed_frames = 0
-    with tqdm(
-        total=total_frames,
-        desc="Обработка батчей "
-        f"{batch_numbers[0]}-{batch_numbers[-1]}",
-        unit=f"фрейм",
-        ncols=150,
-        file=sys.stdout,
-    ) as pbar:
-        while is_processing[0]:
-            if processed_frames >= total_frames:
-                break
-            processed_frames = count_frames_in_certain_batches(
-                OUTPUT_BATCHES_DIR, batch_numbers
-            )
-            pbar.n = processed_frames
-            pbar.refresh()
-            await asyncio.sleep(5)  # Пауза для периодического обновления
-        pbar.n = processed_frames
-        pbar.refresh()
-        pbar.close()
+    start_time = time.time()
+
+    last_logged = 0
+    log_every_sec = 15  # логировать каждые N секунд
+
+    while is_processing[0]:
+        current_frames = count_frames_in_certain_batches(
+            OUTPUT_BATCHES_DIR, batch_numbers
+        )
+
+        # Обновляем только если количество изменилось
+        if current_frames > processed_frames:
+            processed_frames = current_frames
+            progress_percent = (processed_frames / total_frames) * 100
+
+            # Логируем регулярно или при значительном прогрессе
+            current_time = time.time()
+            if (
+                current_time - last_logged >= log_every_sec
+                or progress_percent >= 100
+                or processed_frames == total_frames
+            ):
+                elapsed = current_time - start_time
+                fps = processed_frames / elapsed if elapsed > 0 else 0
+                remaining = (total_frames - processed_frames) / fps if fps > 0 else 0
+
+                logger.info(
+                    f"Апскейл фреймов: {processed_frames}/{total_frames} "
+                    f"({progress_percent:.1f}%) | "
+                    f"Прошло: {elapsed:.1f}сек | "
+                    f"Осталось: {remaining:.1f}сек"
+                )
+                last_logged = current_time
+        if processed_frames >= total_frames:
+            break
+        await asyncio.sleep(5)  # Проверяем прогресс каждые 5 секунд
+    total_time = time.time() - start_time
+    logger.info(
+        f"Апскейл фреймов: {total_frames}/{total_frames} "
+        f"(100%) | Прошло: {total_time:.1f}сек"
+    )
+    logger.success(
+        f"Обработка завершена. (Средняя скорость: {total_frames / total_time:.1f} FPS)"
+    )
 
 
 def _upscale(ai_threads: str, ai_realesrgan_path: str, batch_num: int):
@@ -116,7 +137,7 @@ async def upscale_batches(
 
     logger.info(
         f"Начало обработки батчей {start_batch}-{end_batch} "
-        f"(всего фреймов: {total_frames}, потоков: {process_threads})"
+        f"(всего фреймов: {total_frames}, процессов: {process_threads})"
     )
 
     monitor_task = asyncio.create_task(