From 7b7cb2866a2e0fc75bc620579355136b2538c3d9 Mon Sep 17 00:00:00 2001 From: David Kifer Date: Wed, 24 Jun 2026 13:45:34 -0400 Subject: [PATCH] Implement advanced AI enhancements UI tracking, styling, requirements, and documentation modal --- app/main.py | 10 ++++- app/upscaler.py | 110 ++++++++++++++++++++++++++++++++++++++++++---- context/README.md | 6 +++ requirements.txt | 3 ++ static/app.js | 60 ++++++++++++++++++++++--- static/index.html | 60 +++++++++++++++++++++++++ static/styles.css | 3 ++ 7 files changed, 236 insertions(+), 16 deletions(-) diff --git a/app/main.py b/app/main.py index 8f2e652..342f699 100644 --- a/app/main.py +++ b/app/main.py @@ -98,7 +98,7 @@ def load_jobs_db(): if job.status == "queued": job_queue.put(job_id) # Mark active items as interrupted so they can be resumed - elif job.status in ["analyzing", "extracting", "upscaling", "assembling"]: + elif job.status in ["analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling"]: job.status = "interrupted" job.eta = "Interrupted" jobs_db[job_id] = job @@ -232,6 +232,9 @@ class StartUpscaleRequest(BaseModel): webhook_url: str | None = None transcode_format: str = "mp4" is_preview: bool = False + ai_face_restoration: bool = False + ai_rife_interpolation: bool = False + ai_audio_denoise: bool = False class PreviewRequest(BaseModel): file_id: str @@ -416,7 +419,10 @@ def start_upscale(req: StartUpscaleRequest): interpolation=req.interpolation, webhook_url=req.webhook_url, transcode_format=req.transcode_format, - is_preview=req.is_preview + is_preview=req.is_preview, + ai_face_restoration=req.ai_face_restoration, + ai_rife_interpolation=req.ai_rife_interpolation, + ai_audio_denoise=req.ai_audio_denoise ) jobs_db[job_id] = job diff --git a/app/upscaler.py b/app/upscaler.py index 950d6d9..1ca1746 100644 --- a/app/upscaler.py +++ b/app/upscaler.py @@ -23,13 +23,18 @@ class UpscaleJob: unsharp: bool = False, double_fps: bool = False, preserve_subtitles: bool = True, start_sec: float = None, end_sec: float = None, crf: int = 18, preset: str = "medium", denoise: bool = False, sharpen: bool = False, interpolation: bool = False, - webhook_url: str = None, transcode_format: str = "mp4", is_preview: bool = False): + webhook_url: str = None, transcode_format: str = "mp4", is_preview: bool = False, + ai_face_restoration: bool = False, ai_rife_interpolation: bool = False, + ai_audio_denoise: bool = False): self.job_id = job_id self.video_path = video_path self.model = model self.scale = scale self.tile_size = tile_size self.preserve_audio = preserve_audio + self.ai_face_restoration = ai_face_restoration + self.ai_rife_interpolation = ai_rife_interpolation + self.ai_audio_denoise = ai_audio_denoise # Trim mapping if ss is not None: @@ -448,6 +453,87 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic processed_files = len(os.listdir(output_frames_dir)) job.update_status("upscaling", progress=80.0, current_frame=processed_files) + # Step 2.6: AI Face Restoration (GFPGAN) + if getattr(job, "ai_face_restoration", False): + job.update_status("restoring_faces", progress=81.0) + on_progress_update(job.job_id, {"status": "restoring_faces", "progress": 81.0}) + + import importlib.util + gfpgan_installed = importlib.util.find_spec("gfpgan") is not None + + if gfpgan_installed: + print(f"Job {job.job_id}: GFPGAN detected. Running Face Restoration...") + restored_dir = os.path.join(job_temp_dir, "restored_frames") + os.makedirs(restored_dir, exist_ok=True) + + gfpgan_cmd = [ + sys.executable, "-m", "gfpgan.inference_gfpgan", + "-i", output_frames_dir, + "-o", restored_dir, + "-v", "1.4", + "-s", "1", + "--bg_upsampler", "None" + ] + + p_gfp = job.run_command(gfpgan_cmd) + stdout, stderr = p_gfp.communicate() + job.cleanup_process(p_gfp) + + if p_gfp.returncode == 0: + gfp_output_path = os.path.join(restored_dir, "restored_imgs") + if os.path.exists(gfp_output_path) and len(os.listdir(gfp_output_path)) > 0: + for filename in os.listdir(gfp_output_path): + src_f = os.path.join(gfp_output_path, filename) + dst_f = os.path.join(output_frames_dir, filename) + try: + shutil.copy2(src_f, dst_f) + except Exception as e: + print(f"Error copying restored face frame: {e}") + print(f"Job {job.job_id}: Face Restoration completed successfully.") + else: + print(f"Job {job.job_id}: GFPGAN did not generate outputs in restored_imgs.") + else: + print(f"Job {job.job_id}: GFPGAN failed (exit code {p_gfp.returncode}). Continuing with normal upscale.") + else: + print(f"Job {job.job_id}: 'gfpgan' package is not installed in the virtual environment. Skipping face restoration. To enable, run: pip install gfpgan realesrgan") + + # Step 2.7: AI Frame Interpolation (RIFE) + rife_frames_dir = os.path.join(job_temp_dir, "rife_frames") + use_rife = False + + if getattr(job, "ai_rife_interpolation", False): + rife_bin = os.path.join(BASE_DIR, "rife-bin", "rife-ncnn-vulkan") + if os.path.isfile(rife_bin): + job.update_status("interpolating", progress=83.0) + on_progress_update(job.job_id, {"status": "interpolating", "progress": 83.0}) + os.makedirs(rife_frames_dir, exist_ok=True) + + try: + os.chmod(rife_bin, 0o755) + except Exception: + pass + + rife_cmd = [ + rife_bin, + "-i", output_frames_dir, + "-o", rife_frames_dir, + "-f", "jpg" + ] + if getattr(job, "gpu_ids", None) is not None: + rife_cmd.extend(["-g", str(job.gpu_ids)]) + + p_rife = job.run_command(rife_cmd) + stdout, stderr = p_rife.communicate() + job.cleanup_process(p_rife) + + if p_rife.returncode == 0: + use_rife = True + print(f"Job {job.job_id}: Successfully ran RIFE frame interpolation.") + else: + print(f"Job {job.job_id}: RIFE failed (exit code {p_rife.returncode}). Falling back to FFmpeg interpolation.") + else: + print(f"Job {job.job_id}: RIFE binary not found at {rife_bin}. Falling back to FFmpeg interpolation.") + # Step 3: Reassemble video job.update_status("assembling", progress=85.0) on_progress_update(job.job_id, {"status": "assembling", "progress": 85.0}) @@ -463,12 +549,15 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic if transcode_fmt == "webm": vcodec = "libvpx-vp9" acodec = "libvorbis" + + assemble_frames_dir = rife_frames_dir if use_rife else output_frames_dir + assemble_fps = fps * 2 if (use_rife or getattr(job, "double_fps", False) or getattr(job, "interpolation", False)) else fps # Construct ffmpeg reassembly command assemble_cmd = [ "ffmpeg", "-y", - "-framerate", str(fps), - "-i", os.path.join(output_frames_dir, "frame_%08d.jpg") + "-framerate", str(assemble_fps), + "-i", os.path.join(assemble_frames_dir, "frame_%08d.jpg") ] # We need the original video as the second input (index 1) if we preserve audio or subtitles @@ -485,10 +574,15 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic assemble_cmd.extend(["-map", "0:v:0"]) if job.preserve_audio: - assemble_cmd.extend([ - "-map", "1:a:0?", - "-c:a", acodec - ]) + assemble_cmd.extend(["-map", "1:a:0?"]) + if getattr(job, "ai_audio_denoise", False): + acodec_denoise = "libvorbis" if transcode_fmt == "webm" else "aac" + assemble_cmd.extend([ + "-af", "arnnoise", + "-c:a", acodec_denoise + ]) + else: + assemble_cmd.extend(["-c:a", acodec]) if getattr(job, "preserve_subtitles", True): assemble_cmd.extend([ @@ -502,7 +596,7 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic vf_filters.append("hqdn3d") if getattr(job, "sharpen", False) or getattr(job, "unsharp", False): vf_filters.append("unsharp=3:3:0.5:3:3:0.5") - if getattr(job, "interpolation", False) or getattr(job, "double_fps", False): + if (getattr(job, "double_fps", False) or getattr(job, "interpolation", False) or getattr(job, "ai_rife_interpolation", False)) and not use_rife: target_fps = fps * 2 if getattr(job, "double_fps", False) else 60 if target_fps < fps: target_fps = fps diff --git a/context/README.md b/context/README.md index 7cf3b0f..8f97861 100644 --- a/context/README.md +++ b/context/README.md @@ -56,3 +56,9 @@ A typical upscale job follows these sequential steps: - **Self-Sufficiency**: Running `python3 start.py` automatically checks for a local virtual environment (`venv/` or `.venv/`). - **Auto-Provisioning**: If no virtual environment is found, `start.py` will initialize one in `venv/`, upgrade `pip`, install all dependencies listed in `requirements.txt`, mark the Real-ESRGAN binary as executable (`chmod +x`), and create necessary folders (`uploads/`, `outputs/`, `temp/`). - **Seamless Launch**: It then automatically launches the server process using the newly created environment interpreter. + +### 4. Advanced AI Enhancements +- **AI Face Restoration (GFPGAN)**: Runs as a python subprocess invoking `gfpgan.inference_gfpgan` to restore and clear up human faces in low-resolution video frames. Results are copied directly back into the frame output folder before motion interpolation and final video assembly. +- **AI Frame Interpolation (RIFE)**: Runs using the `rife-ncnn-vulkan` binary (expected in `rife-bin/`). Smooths motion by generating and inserting intermediate frames, doubling the framerate. Falls back to FFmpeg's `minterpolate` optical flow filter if the Vulkan binary is not present. +- **AI Audio Denoising (RNNoise)**: Transports and filters audio using the deep-learning-based `arnnoise` FFmpeg filter, eliminating background noise from output tracks during assembly. + diff --git a/requirements.txt b/requirements.txt index 32ef7e9..8defef8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,6 @@ python-multipart>=0.0.9 websockets>=13.0 pydantic>=2.0 torch>=2.0 +gfpgan>=1.3.8 +realesrgan>=0.3.0 + diff --git a/static/app.js b/static/app.js index 21d1e2a..45adcb6 100644 --- a/static/app.js +++ b/static/app.js @@ -600,6 +600,10 @@ generateVideoPreviewBtn.addEventListener("click", async () => { const filterSharpen = document.getElementById("filter-sharpen"); const filterFps = document.getElementById("filter-fps"); + const aiFaceRestoration = document.getElementById("ai-face-restoration"); + const aiRifeInterpolation = document.getElementById("ai-rife-interpolation"); + const aiAudioDenoise = document.getElementById("ai-audio-denoise"); + const transcodeFormat = transcodeSelect.value; try { @@ -626,7 +630,10 @@ generateVideoPreviewBtn.addEventListener("click", async () => { interpolation: filterFps.checked, webhook_url: null, transcode_format: transcodeFormat, - is_preview: true + is_preview: true, + ai_face_restoration: aiFaceRestoration ? aiFaceRestoration.checked : false, + ai_rife_interpolation: aiRifeInterpolation ? aiRifeInterpolation.checked : false, + ai_audio_denoise: aiAudioDenoise ? aiAudioDenoise.checked : false }) }); @@ -711,6 +718,10 @@ startUpscaleBtn.addEventListener("click", async () => { const filterSharpen = document.getElementById("filter-sharpen"); const filterFps = document.getElementById("filter-fps"); const webhookUrl = document.getElementById("webhook-url"); + + const aiFaceRestoration = document.getElementById("ai-face-restoration"); + const aiRifeInterpolation = document.getElementById("ai-rife-interpolation"); + const aiAudioDenoise = document.getElementById("ai-audio-denoise"); try { const res = await fetch("/api/upscale/start", { @@ -733,7 +744,10 @@ startUpscaleBtn.addEventListener("click", async () => { denoise: filterDenoise.checked, sharpen: filterSharpen.checked, interpolation: filterFps.checked, - webhook_url: webhookUrl.value.trim() || null + webhook_url: webhookUrl.value.trim() || null, + ai_face_restoration: aiFaceRestoration ? aiFaceRestoration.checked : false, + ai_rife_interpolation: aiRifeInterpolation ? aiRifeInterpolation.checked : false, + ai_audio_denoise: aiAudioDenoise ? aiAudioDenoise.checked : false }) }); @@ -816,6 +830,10 @@ function updateProgressUI(data) { statusText = "Queued in pipeline. Waiting for GPU lock..."; } else if (data.status === "extracting") { statusText = "Extracting video frames..."; + } else if (data.status === "restoring_faces") { + statusText = "GFPGAN: Running Face Restoration on upscaled frames..."; + } else if (data.status === "interpolating") { + statusText = "RIFE: Interpolating frames for smooth motion..."; } else if (data.status === "assembling") { statusText = "Assembling video preview and audio tracks..."; } @@ -891,6 +909,10 @@ function updateProgressUI(data) { appendLogLine("FFmpeg: Splitting video stream to high-fidelity JPG frames...", "info"); } else if (data.status === "upscaling") { appendLogLine(`Vulkan GPU: Processing frame ${data.current_frame} of ${data.total_frames} (ETA: ${data.eta})...`, "info"); + } else if (data.status === "restoring_faces") { + appendLogLine("GFPGAN: Running Face Restoration on upscaled frames...", "info"); + } else if (data.status === "interpolating") { + appendLogLine("RIFE: Interpolating frames for smooth motion...", "info"); } else if (data.status === "assembling") { appendLogLine("FFmpeg: Reassembling enhanced frames & merging audio tracks...", "info"); } else if (data.status === "completed") { @@ -1021,7 +1043,7 @@ async function checkActiveJobs() { const jobs = await res.json(); const activeJob = jobs.find(job => - ["queued", "analyzing", "extracting", "upscaling", "assembling"].includes(job.status) + ["queued", "analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling"].includes(job.status) ); if (activeJob) { @@ -1318,7 +1340,7 @@ async function loadQueue() { const jobs = await res.json(); const displayJobs = jobs.filter(job => - ["queued", "analyzing", "extracting", "upscaling", "assembling", "completed", "failed", "cancelled", "interrupted"].includes(job.status) + ["queued", "analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling", "completed", "failed", "cancelled", "interrupted"].includes(job.status) ); if (displayJobs.length === 0) { @@ -1380,7 +1402,7 @@ async function loadQueue() { Clear `; - } else if (["queued", "analyzing", "extracting", "upscaling", "assembling"].includes(job.status)) { + } else if (["queued", "analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling"].includes(job.status)) { actionsHtml = ` ${reorderHtml}
AI Face Restoration (GFPGAN) + + + +
+ +
@@ -590,6 +609,47 @@
+ + + diff --git a/static/styles.css b/static/styles.css index a014897..bbe40c9 100644 --- a/static/styles.css +++ b/static/styles.css @@ -653,11 +653,14 @@ input:checked + .slider:before { .status-badge.queued { background: rgba(255, 184, 0, 0.1); border: 1px solid var(--warning); color: var(--warning); } .status-badge.extracting { background: rgba(189, 0, 255, 0.15); border: 1px solid var(--secondary); color: #e499ff; } .status-badge.upscaling { background: rgba(0, 240, 255, 0.1); border: 1px solid var(--primary); color: var(--primary); } +.status-badge.restoring_faces { background: rgba(255, 0, 255, 0.1); border: 1px solid var(--secondary); color: var(--secondary); } +.status-badge.interpolating { background: rgba(0, 255, 184, 0.1); border: 1px solid var(--primary); color: var(--primary); } .status-badge.assembling { background: rgba(0, 240, 255, 0.1); border: 1px solid var(--primary); color: var(--primary); } .status-badge.completed { background: rgba(0, 255, 135, 0.1); border: 1px solid var(--success); color: var(--success); } .status-badge.failed { background: rgba(255, 0, 85, 0.1); border: 1px solid var(--danger); color: var(--danger); } .status-badge.interrupted { background: rgba(255, 255, 255, 0.08); border: 1px solid var(--text-muted); color: var(--text-muted); } + .bar-container { width: 100%; height: 16px;