""" AI Video 4K Upscaler Module for Historical Series Optimized for NVIDIA GTX 1660 Super (6GB VRAM) & Real-ESRGAN NCNN Vulkan """ import os import sys import json import time import shutil import argparse import subprocess from pathlib import Path # Force UTF-8 encoding on Windows console if sys.platform.startswith("win"): try: sys.stdout.reconfigure(encoding="utf-8") sys.stderr.reconfigure(encoding="utf-8") except Exception: pass FFMPEG_EXE = r"C:\Projects\Historical\venv\Scripts\ffmpeg.exe" FFPROBE_EXE = r"C:\Projects\Historical\venv\Scripts\ffmpeg.exe" # Can probe via ffmpeg REAL_ESRGAN_EXE = r"C:\tools\realesrgan\realesrgan-ncnn-vulkan.exe" REAL_ESRGAN_DIR = r"C:\tools\realesrgan" def check_environment(): if not os.path.exists(FFMPEG_EXE): raise FileNotFoundError(f"FFmpeg not found at {FFMPEG_EXE}") if not os.path.exists(REAL_ESRGAN_EXE): raise FileNotFoundError(f"Real-ESRGAN not found at {REAL_ESRGAN_EXE}") def get_video_meta(video_path: str): cmd = [ FFMPEG_EXE, "-i", video_path ] res = subprocess.run(cmd, stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True) fps = 24.0 has_audio = False for line in res.stderr.splitlines(): if "Stream #" in line and "Video:" in line: if "fps" in line: parts = line.split(",") for p in parts: if "fps" in p: try: fps = float(p.strip().split()[0]) except ValueError: pass if "Stream #" in line and "Audio:" in line: has_audio = True return fps, has_audio def upscale_video( input_video: str, output_video: str, start_sec: float = 0, duration_sec: float = None, tile_size: int = 256, model_name: str = "realesrgan-x4plus", target_res: str = "3840x2160", cq: int = 19 ): check_environment() input_path = Path(input_video).resolve() output_path = Path(output_video).resolve() if not input_path.exists(): raise FileNotFoundError(f"Input file not found: {input_path}") start_time = time.time() work_dir = input_path.parent / f"_temp_upscale_{int(time.time())}" frames_in = work_dir / "in_frames" frames_out = work_dir / "out_frames" audio_file = work_dir / "audio.aac" if work_dir.exists(): shutil.rmtree(work_dir) frames_in.mkdir(parents=True, exist_ok=True) frames_out.mkdir(parents=True, exist_ok=True) print(f"=== BẮT ĐẦU UPSCALE VIDEO 4K ===") print(f"Nguồn: {input_path.name}") print(f"Đích: {output_path.name}") print(f"Mô hình: {model_name} (Tile size: {tile_size})") try: fps, has_audio = get_video_meta(str(input_path)) print(f"FPS: {fps} | Có Audio: {has_audio}") # 1. Trích xuất audio (nếu có) if has_audio: print("[1/4] Đang tách track âm thanh...") audio_cmd = [FFMPEG_EXE, "-y"] if start_sec > 0: audio_cmd.extend(["-ss", str(start_sec)]) audio_cmd.extend(["-i", str(input_path)]) if duration_sec: audio_cmd.extend(["-t", str(duration_sec)]) audio_cmd.extend(["-vn", "-c:a", "aac", "-b:a", "192k", str(audio_file)]) subprocess.run(audio_cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) # 2. Trích xuất frames print("[2/4] Đang trích xuất frames ảnh...") extract_cmd = [FFMPEG_EXE, "-y"] if start_sec > 0: extract_cmd.extend(["-ss", str(start_sec)]) extract_cmd.extend(["-i", str(input_path)]) if duration_sec: extract_cmd.extend(["-t", str(duration_sec)]) extract_cmd.extend(["-qscale:v", "1", str(frames_in / "frame_%08d.png")]) subprocess.run(extract_cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) frame_count = len(list(frames_in.glob("*.png"))) print(f"Tổng số frames cần xử lý: {frame_count}") if frame_count == 0: raise RuntimeError("Không có frame nào được trích xuất từ video nguồn!") # 3. Chạy Real-ESRGAN NCNN Vulkan trên GPU (Tối ưu hóa VRAM 6GB với -t 256 và -j 1:1:2) print(f"[3/4] Đang AI Super-Resolution trên NVIDIA GeForce GTX 1660 Super...") upscale_cmd = [ REAL_ESRGAN_EXE, "-i", str(frames_in), "-o", str(frames_out), "-n", model_name, "-s", "4", "-t", str(tile_size), "-j", "1:1:2", "-f", "png" ] proc = subprocess.Popen( upscale_cmd, cwd=REAL_ESRGAN_DIR, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1 ) for line in proc.stdout: line_str = line.strip() if "%" in line_str: print(f"\rTiến độ: {line_str}", end="", flush=True) elif "NVIDIA" in line_str: print(f"\nKhởi tạo GPU: {line_str}") proc.wait() if proc.returncode != 0: raise RuntimeError(f"Real-ESRGAN exited with error code {proc.returncode}") print("\nHoàn tất xử lý AI cho tất cả các frames!") # 4. Ghép lại thành video 4K bằng NVENC print("[4/4] Đang render video 4K (H.265 / HEVC NVENC)...") render_cmd = [ FFMPEG_EXE, "-y", "-framerate", str(fps), "-i", str(frames_out / "frame_%08d.png") ] if has_audio and audio_file.exists(): render_cmd.extend(["-i", str(audio_file)]) render_cmd.extend([ "-c:v", "hevc_nvenc", "-preset", "p6", "-rc:v", "vbr", "-cq:v", str(cq), "-b:v", "0", "-s", target_res, "-pix_fmt", "yuv420p" ]) if has_audio and audio_file.exists(): render_cmd.extend(["-c:a", "aac", "-b:a", "192k", "-shortest"]) render_cmd.append(str(output_path)) subprocess.run(render_cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) elapsed = round(time.time() - start_time, 1) out_size_mb = round(output_path.stat().st_size / (1024 * 1024), 2) print(f"\n==========================================") print(f"THÀNH CÔNG!") print(f"Video đầu ra: {output_path}") print(f"Dung lượng: {out_size_mb} MB") print(f"Thời gian: {elapsed} giây") print(f"==========================================") return str(output_path) finally: if work_dir.exists(): shutil.rmtree(work_dir, ignore_errors=True) print("Đã dọn dẹp các thư mục ảnh tạm.") if __name__ == "__main__": parser = argparse.ArgumentParser(description="AI 4K Video Upscaler for GTX 1660 Super") parser.add_argument("input", help="Đường dẫn video đầu vào") parser.add_argument("output", help="Đường dẫn video 4K đầu ra") parser.add_argument("--start", type=float, default=0, help="Giây bắt đầu (mặc định 0)") parser.add_argument("--duration", type=float, default=None, help="Thời lượng cắt (giây) để test nhanh") parser.add_argument("--tile", type=int, default=400, help="Tile size (400 tối ưu cho 6GB VRAM)") parser.add_argument("--model", default="realesrgan-x4plus", help="Mô hình Real-ESRGAN") parser.add_argument("--res", default="3840x2160", help="Độ phân giải mục tiêu") args = parser.parse_args() upscale_video( input_video=args.input, output_video=args.output, start_sec=args.start, duration_sec=args.duration, tile_size=args.tile, model_name=args.model, target_res=args.res )