Initial commit of video upscaler
This commit is contained in:
+519
@@ -0,0 +1,519 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user