""" AI Video 4K Upscaler Module for Historical Series Optimized for Dual Intel Xeon (56 Threads) & NVIDIA GTX 1660 Super (6GB VRAM) Using Real-ESRGAN NCNN Vulkan with Fast-I/O & Multi-threaded Post-Processing """ 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" 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 get_cpu_thread_allocation(): """Tự động phân bổ luồng CPU cho kiến trúc đa nhân (hỗ trợ tối đa Dual Xeon 56 luồng)""" total = os.cpu_count() or 8 # load : proc : save # load: giải nén frame đầu vào (CPU) # proc: hàng đợi gửi lệnh GPU Vulkan (GPU queue, 2 là tối ưu cho 1 GPU) # save: nén và ghi frame 4K (CPU) load_threads = min(16, max(4, total // 4)) proc_threads = 2 save_threads = min(28, max(8, total // 2)) return f"{load_threads}:{proc_threads}:{save_threads}", total def upscale_video( input_video: str, output_video: str, start_sec: float = 0, duration_sec: float = None, tile_size: int = 1280, model_name: str = "realesr-animevideov3", 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, ignore_errors=True) frames_in.mkdir(parents=True, exist_ok=True) frames_out.mkdir(parents=True, exist_ok=True) thread_param, total_cpu_threads = get_cpu_thread_allocation() print(f"=== BẮT ĐẦU UPSCALE VIDEO 4K (HYBRID XEON 56-THREAD + GTX 1660 SUPER) ===") print(f"Nguồn: {input_path.name}") print(f"Đích: {output_path.name}") print(f"Mô hình: {model_name}") print(f"Tile Size: {tile_size} (Tối ưu hóa VRAM 6GB & triệt tiêu padding thừa)") print(f"Luồng NCNN: -j {thread_param} (Tổng số luồng hệ thống: {total_cpu_threads})") print(f"Định dạng I/O: JPG Ultra-Quality (Khử hoàn toàn nghẽn ghi đĩa PNG)") try: fps, has_audio = get_video_meta(str(input_path)) print(f"FPS nguồn: {fps} | Có Audio: {has_audio}") # 1. Trích xuất audio (nếu có) if has_audio: print("\n[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 (dùng JPG qscale:2 + đa luồng CPU Xeon cực nhanh) print("\n[2/4] Đang trích xuất frames ảnh (Đa luồng CPU AVX2)...") extract_cmd = [ FFMPEG_EXE, "-y", "-threads", str(min(32, total_cpu_threads)) ] 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", "2", "-pix_fmt", "yuvj420p", str(frames_in / "frame_%08d.jpg") ]) t_ext = time.time() subprocess.run(extract_cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) ext_time = time.time() - t_ext frame_count = len(list(frames_in.glob("*.jpg"))) print(f"Đã trích xuất {frame_count} frames trong {ext_time:.2f}s ({frame_count / max(0.01, ext_time):.1f} FPS)") 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 + CPU Đa luồng print(f"\n[3/4] Đang chạy AI Super-Resolution (GTX 1660 Super + {thread_param} Threads)...") upscale_cmd = [ REAL_ESRGAN_EXE, "-i", str(frames_in), "-o", str(frames_out), "-n", model_name, "-s", "4", "-t", str(tile_size), "-j", thread_param, "-f", "jpg" ] t_ai = time.time() 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 độ AI: {line_str}", end="", flush=True) elif "NVIDIA" in line_str: print(f"Thiết bị GPU: {line_str}") proc.wait() if proc.returncode != 0: raise RuntimeError(f"Real-ESRGAN exited with error code {proc.returncode}") ai_time = time.time() - t_ai avg_ai_fps = frame_count / max(0.01, ai_time) print(f"\nHoàn tất xử lý AI! Thời gian: {ai_time:.2f}s -> Tốc độ: {avg_ai_fps:.2f} FPS") # 4. Ghép lại thành video 4K bằng NVENC + CPU Lanczos Resampling siêu nét print("\n[4/4] Đang render video 4K (CPU Lanczos Filter + Turing NVENC HEVC)...") render_cmd = [ FFMPEG_EXE, "-y", "-threads", str(min(24, total_cpu_threads)), "-framerate", str(fps), "-i", str(frames_out / "frame_%08d.jpg") ] if has_audio and audio_file.exists(): render_cmd.extend(["-i", str(audio_file)]) # CPU Lanczos downsampling 5K -> 4K + Unsharp mask làm nổi bật chi tiết render_cmd.extend([ "-vf", "scale=3840:2160:flags=lanczos,unsharp=5:5:0.6:3:3:0.3", "-c:v", "hevc_nvenc", "-preset", "p4", "-rc:v", "vbr", "-cq:v", str(cq), "-b:v", "0", "-spatial-aq", "1", "-b_ref_mode", "middle", "-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)) t_ren = time.time() subprocess.run(render_cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) ren_time = time.time() - t_ren total_elapsed = round(time.time() - start_time, 1) out_size_mb = round(output_path.stat().st_size / (1024 * 1024), 2) total_fps = frame_count / max(0.01, total_elapsed) print(f"\n==========================================") print(f"UPSCALE THÀNH CÔNG!") print(f"Video đầu ra: {output_path.name}") print(f"Dung lượng: {out_size_mb} MB") print(f"Thời gian tổng: {total_elapsed}s (Tách: {ext_time:.1f}s | AI: {ai_time:.1f}s | Render: {ren_time:.1f}s)") print(f"Tốc độ trung bình: {total_fps:.2f} FPS") 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 tệp tạm thời.") if __name__ == "__main__": parser = argparse.ArgumentParser(description="AI 4K Video Upscaler - Dual Xeon 56-Thread & 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=1280, help="Tile size (1280 tối ưu cho 6GB VRAM & giảm padding lặp)") parser.add_argument("--model", default="realesr-animevideov3", help="Mô hình Real-ESRGAN (mặc định realesr-animevideov3 SRVGGNetCompact)") parser.add_argument("--res", default="3840x2160", help="Độ phân giải mục tiêu") parser.add_argument("--cq", type=int, default=19, help="Chất lượng nén NVENC CQ (mặc định 19)") 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, cq=args.cq )