import os import subprocess import json import shutil import time import threading from typing import Dict, Any, Callable # Paths config BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) BIN_PATH = os.path.join(BASE_DIR, "realesrgan-bin", "realesrgan-ncnn-vulkan") UPLOAD_DIR = os.path.join(BASE_DIR, "uploads") OUTPUT_DIR = os.path.join(BASE_DIR, "outputs") TEMP_DIR = os.path.join(BASE_DIR, "temp") # Ensure directories exist for d in [UPLOAD_DIR, OUTPUT_DIR, TEMP_DIR]: os.makedirs(d, exist_ok=True) class UpscaleJob: def __init__(self, job_id: str, video_path: str, model: str, scale: int, tile_size: int, preserve_audio: bool, ss: str = None, t: str = None, gpu_ids: str = None, tta: bool = False, 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): 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 # Trim mapping if ss is not None: self.ss = ss elif start_sec is not None: self.ss = str(start_sec) else: self.ss = None if t is not None: self.t = t elif end_sec is not None: start = float(self.ss) if self.ss is not None else 0.0 self.t = str(max(0.0, end_sec - start)) else: self.t = None self.gpu_ids = gpu_ids self.tta = tta self.unsharp = unsharp self.double_fps = double_fps self.preserve_subtitles = preserve_subtitles self.start_sec = start_sec self.end_sec = end_sec self.crf = crf self.preset = preset self.denoise = denoise self.sharpen = sharpen self.interpolation = interpolation self.webhook_url = webhook_url self.transcode_format = transcode_format self.is_preview = is_preview self.original_preview_file = None self.status = "pending" self.progress = 0.0 self.current_frame = 0 self.total_frames = 0 self.eta = "Calculating..." self.error = None self.start_time = None self.output_file = None # Track processes to allow cancellation self._processes = [] self._is_cancelled = False self._lock = threading.Lock() def update_status(self, status: str, progress: float = None, current_frame: int = None, eta: str = None, error: str = None): with self._lock: self.status = status if progress is not None: self.progress = round(progress, 2) if current_frame is not None: self.current_frame = current_frame if eta is not None: self.eta = eta if error is not None: self.error = error def cancel(self): with self._lock: self._is_cancelled = True self.status = "cancelled" self.eta = "N/A" for p in self._processes: try: p.terminate() p.wait(timeout=2) except Exception: try: p.kill() except Exception: pass self._processes.clear() def run_command(self, cmd: list, shell=False) -> subprocess.Popen: with self._lock: if self._is_cancelled: raise InterruptedError("Job was cancelled") p = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=shell ) self._processes.append(p) return p def cleanup_process(self, p: subprocess.Popen): with self._lock: if p in self._processes: self._processes.remove(p) def get_video_info(video_path: str) -> dict: """Extract metadata using ffprobe""" cmd = [ "ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=avg_frame_rate,duration,width,height,codec_name,nb_frames", "-of", "json", video_path ] try: result = subprocess.run(cmd, capture_output=True, text=True, check=True) data = json.loads(result.stdout) if not data.get("streams"): return {} stream = data["streams"][0] # Calculate FPS fps_str = stream.get("avg_frame_rate", "30/1") if "/" in fps_str: num, den = map(float, fps_str.split("/")) fps = num / den if den != 0 else 30.0 else: fps = float(fps_str) if fps_str else 30.0 # Get total frames nb_frames = stream.get("nb_frames") if nb_frames and nb_frames.isdigit(): total_frames = int(nb_frames) else: duration = float(stream.get("duration", 0)) total_frames = int(duration * fps) return { "width": int(stream.get("width", 0)), "height": int(stream.get("height", 0)), "fps": round(fps, 3), "duration": round(float(stream.get("duration", 0)), 2), "codec": stream.get("codec_name", "unknown"), "total_frames": total_frames } except Exception as e: print(f"Error reading video info: {e}") return {} def extract_single_frame(video_path: str, timestamp_sec: float, output_path: str) -> bool: """Extract a single frame at timestamp for preview""" cmd = [ "ffmpeg", "-y", "-ss", str(timestamp_sec), "-i", video_path, "-vframes", "1", "-f", "image2", output_path ] try: subprocess.run(cmd, capture_output=True, check=True) return os.path.exists(output_path) except Exception as e: print(f"Error extracting single frame: {e}") return False def upscale_image_file(input_path: str, output_path: str, model: str, scale: int, tile_size: int, gpu_ids: str = None) -> bool: """Run Real-ESRGAN on a single image file""" cmd = [ BIN_PATH, "-i", input_path, "-o", output_path, "-n", model, "-s", str(scale), "-t", str(tile_size) ] if gpu_ids is not None: cmd.extend(["-g", str(gpu_ids)]) try: subprocess.run(cmd, capture_output=True, check=True) return os.path.exists(output_path) except Exception as e: print(f"Error upscaling single image: {e}") return False def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dict[str, Any]], None]): job.start_time = time.time() job.update_status("analyzing", progress=5) # Get video info info = get_video_info(job.video_path) if not info: job.update_status("failed", error="Could not read video metadata.") on_progress_update(job.job_id, {"status": "failed", "error": "Could not read video metadata"}) return job.total_frames = info["total_frames"] fps = info["fps"] # Create job temp directories job_temp_dir = os.path.join(TEMP_DIR, job.job_id) input_frames_dir = os.path.join(job_temp_dir, "input_frames") output_frames_dir = os.path.join(job_temp_dir, "output_frames") os.makedirs(input_frames_dir, exist_ok=True) os.makedirs(output_frames_dir, exist_ok=True) try: # If preview, extract the original 5s clip first transcode_fmt = getattr(job, "transcode_format", "mp4") if getattr(job, "is_preview", False): orig_preview_filename = f"original_{job.job_id}.{transcode_fmt}" orig_preview_filepath = os.path.join(OUTPUT_DIR, orig_preview_filename) job.original_preview_file = orig_preview_filepath cut_cmd = ["ffmpeg", "-y"] if job.ss is not None: cut_cmd.extend(["-ss", str(job.ss)]) if job.t is not None: cut_cmd.extend(["-t", str(job.t)]) cut_cmd.extend([ "-i", job.video_path, "-map", "0:v:0", "-map", "0:a:0?", "-c:v", "libx264", "-c:a", "aac", orig_preview_filepath ]) try: subprocess.run(cut_cmd, capture_output=True, check=True) except Exception as cut_err: print(f"Error cutting original preview video: {cut_err}") # Step 1: Extract Frames job.update_status("extracting", progress=10) on_progress_update(job.job_id, {"status": "extracting", "progress": 10}) # High quality JPG frames to balance disk usage and speed extract_cmd = ["ffmpeg", "-y"] if job.ss is not None: extract_cmd.extend(["-ss", str(job.ss)]) if job.t is not None: extract_cmd.extend(["-t", str(job.t)]) extract_cmd.extend(["-i", job.video_path]) # Apply unsharp pre-filter if enabled if getattr(job, "unsharp", False): extract_cmd.extend(["-vf", "unsharp"]) extract_cmd.extend([ "-q:v", "2", os.path.join(input_frames_dir, "frame_%08d.jpg") ]) p_extract = job.run_command(extract_cmd) stdout, stderr = p_extract.communicate() job.cleanup_process(p_extract) if p_extract.returncode != 0: raise RuntimeError(f"FFmpeg frame extraction failed: {stderr}") # Count actual frames extracted extracted_files = sorted([f for f in os.listdir(input_frames_dir) if f.startswith("frame_")]) actual_total = len(extracted_files) if actual_total == 0: raise RuntimeError("No frames extracted from video") job.total_frames = actual_total # Step 2: Upscale Frames job.update_status("upscaling", progress=20, current_frame=0) on_progress_update(job.job_id, {"status": "upscaling", "progress": 20, "current_frame": 0, "total_frames": actual_total}) current_tile_size = job.tile_size while True: # Launch Real-ESRGAN on directory upscale_cmd = [ BIN_PATH, "-i", input_frames_dir, "-o", output_frames_dir, "-n", job.model, "-s", str(job.scale), "-t", str(current_tile_size), "-f", "jpg" ] if getattr(job, "gpu_ids", None) is not None: upscale_cmd.extend(["-g", str(job.gpu_ids)]) if getattr(job, "tta", False): upscale_cmd.append("-x") upscale_start_time = time.time() p_upscale = job.run_command(upscale_cmd) # Monitor thread for output files while p_upscale.poll() is None: if job._is_cancelled: return processed_files = len(os.listdir(output_frames_dir)) progress_pct = 20.0 + (float(processed_files) / actual_total) * 60.0 # upscaling is 20% to 80% # Estimate ETA elapsed = time.time() - upscale_start_time if processed_files > 0: sec_per_frame = elapsed / processed_files rem_frames = actual_total - processed_files eta_sec = rem_frames * sec_per_frame # Format ETA if eta_sec > 60: eta_str = f"{int(eta_sec // 60)}m {int(eta_sec % 60)}s" else: eta_str = f"{int(eta_sec)}s" else: eta_str = "Calculating..." job.update_status("upscaling", progress=progress_pct, current_frame=processed_files, eta=eta_str) on_progress_update(job.job_id, { "status": "upscaling", "progress": progress_pct, "current_frame": processed_files, "total_frames": actual_total, "eta": eta_str }) time.sleep(0.5) stdout, stderr = p_upscale.communicate() job.cleanup_process(p_upscale) if job._is_cancelled: return if p_upscale.returncode != 0: err_msg = (stdout or "") + "\n" + (stderr or "") is_alloc_error = any(x in err_msg.lower() for x in ["vkallocatememory", "out of memory", "allocation", "vram", "failed to allocate"]) if is_alloc_error: if current_tile_size <= 0: next_tile_size = 256 else: next_tile_size = current_tile_size // 2 if next_tile_size >= 32: print(f"Job {job.job_id}: Real-ESRGAN failed with VRAM allocation error. Retrying with tile size halved from {current_tile_size} to {next_tile_size}.") current_tile_size = next_tile_size # Clean up output frames directory before retrying for filename in os.listdir(output_frames_dir): file_path = os.path.join(output_frames_dir, filename) try: if os.path.isfile(file_path) or os.path.islink(file_path): os.unlink(file_path) elif os.path.isdir(file_path): shutil.rmtree(file_path) except Exception as cleanup_err: print(f"Error cleaning file {file_path}: {cleanup_err}") continue raise RuntimeError(f"Real-ESRGAN failed with exit code {p_upscale.returncode}: {err_msg}") else: break # Final validation of upscale output processed_files = len(os.listdir(output_frames_dir)) job.update_status("upscaling", progress=80.0, current_frame=processed_files) # Step 3: Reassemble video job.update_status("assembling", progress=85.0) on_progress_update(job.job_id, {"status": "assembling", "progress": 85.0}) transcode_fmt = getattr(job, "transcode_format", "mp4") out_filename = f"upscaled_{job.job_id}.{transcode_fmt}" out_filepath = os.path.join(OUTPUT_DIR, out_filename) job.output_file = out_filepath # Choose codecs based on format vcodec = "libx264" acodec = "copy" if transcode_fmt == "webm": vcodec = "libvpx-vp9" acodec = "libvorbis" # Construct ffmpeg reassembly command assemble_cmd = [ "ffmpeg", "-y", "-framerate", str(fps), "-i", os.path.join(output_frames_dir, "frame_%08d.jpg") ] # We need the original video as the second input (index 1) if we preserve audio or subtitles need_orig_input = job.preserve_audio or getattr(job, "preserve_subtitles", True) if need_orig_input: orig_input_cmd = [] if getattr(job, "ss", None) is not None: orig_input_cmd.extend(["-ss", str(job.ss)]) if getattr(job, "t", None) is not None: orig_input_cmd.extend(["-t", str(job.t)]) orig_input_cmd.extend(["-i", job.video_path]) assemble_cmd.extend(orig_input_cmd) assemble_cmd.extend(["-map", "0:v:0"]) if job.preserve_audio: assemble_cmd.extend([ "-map", "1:a:0?", "-c:a", acodec ]) if getattr(job, "preserve_subtitles", True): assemble_cmd.extend([ "-map", "1:s?", "-c:s", "mov_text" ]) # Apply filters vf_filters = [] if getattr(job, "denoise", False): 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): target_fps = fps * 2 if getattr(job, "double_fps", False) else 60 if target_fps < fps: target_fps = fps vf_filters.append(f"minterpolate=fps={target_fps}") if vf_filters: assemble_cmd.extend(["-vf", ",".join(vf_filters)]) assemble_cmd.extend([ "-c:v", vcodec, "-pix_fmt", "yuv420p", "-crf", str(getattr(job, "crf", 18)), "-preset", getattr(job, "preset", "medium"), out_filepath ]) p_assemble = job.run_command(assemble_cmd) stdout, stderr = p_assemble.communicate() job.cleanup_process(p_assemble) if p_assemble.returncode != 0: raise RuntimeError(f"FFmpeg video assembly failed: {stderr}") # Step 4: Complete job.update_status("completed", progress=100.0, eta="Done") on_progress_update(job.job_id, { "status": "completed", "progress": 100.0, "eta": "Done", "output_file": out_filename }) except Exception as e: import traceback traceback.print_exc() if not job._is_cancelled: job.update_status("failed", error=str(e)) on_progress_update(job.job_id, {"status": "failed", "error": str(e)}) finally: # Clean up temp frames to save space try: if os.path.exists(job_temp_dir): shutil.rmtree(job_temp_dir) except Exception as cleanup_err: print(f"Error during temp cleanup: {cleanup_err}") # Trigger webhook if URL is provided if getattr(job, "webhook_url", None): def trigger_webhook_task(): import urllib.request import json try: payload = { "job_id": job.job_id, "status": job.status, "progress": job.progress, "error": job.error, "output_file": os.path.basename(job.output_file) if job.output_file else None, "model": job.model, "scale": job.scale, "duration": round(time.time() - job.start_time, 2) if job.start_time else 0 } req = urllib.request.Request( job.webhook_url, data=json.dumps(payload).encode('utf-8'), headers={'Content-Type': 'application/json'}, method='POST' ) with urllib.request.urlopen(req, timeout=5) as response: pass except Exception as ex: print(f"Error triggering webhook: {ex}") threading.Thread(target=trigger_webhook_task, daemon=True).start()