853 lines
38 KiB
Python
853 lines
38 KiB
Python
import os
|
|
import sys
|
|
import subprocess
|
|
import json
|
|
import shutil
|
|
import time
|
|
import threading
|
|
from typing import Dict, Any, Callable
|
|
|
|
dist_chunks: Dict[str, list] = {}
|
|
dist_chunks_lock = threading.Lock()
|
|
|
|
_ffmpeg_filters_cache = {}
|
|
def has_ffmpeg_filter(filter_name: str) -> bool:
|
|
global _ffmpeg_filters_cache
|
|
if filter_name in _ffmpeg_filters_cache:
|
|
return _ffmpeg_filters_cache[filter_name]
|
|
try:
|
|
p = subprocess.Popen(["ffmpeg", "-filters"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
stdout, _ = p.communicate(timeout=5)
|
|
_ffmpeg_filters_cache[filter_name] = filter_name in stdout
|
|
except Exception:
|
|
_ffmpeg_filters_cache[filter_name] = False
|
|
return _ffmpeg_filters_cache[filter_name]
|
|
|
|
# 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,
|
|
ai_face_restoration: bool = False, ai_rife_interpolation: bool = False,
|
|
ai_audio_denoise: bool = False, temp_dir: str = None,
|
|
output_dir: str = None, source_filename: str = None):
|
|
self.job_id = job_id
|
|
self.temp_dir = temp_dir
|
|
self.output_dir = output_dir
|
|
self.source_filename = source_filename
|
|
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:
|
|
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/pause
|
|
self._processes = []
|
|
self._is_cancelled = False
|
|
self._is_paused = False
|
|
self._lock = threading.Lock()
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Serialize job attributes, excluding internal thread/process resources."""
|
|
return {k: v for k, v in self.__dict__.items() if not k.startswith('_')}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict) -> 'UpscaleJob':
|
|
"""Deserialize job from dictionary, reconstructing internal locks and processes."""
|
|
job = cls(
|
|
job_id=data.get('job_id'),
|
|
video_path=data.get('video_path'),
|
|
model=data.get('model'),
|
|
scale=data.get('scale', 4),
|
|
tile_size=data.get('tile_size', 256),
|
|
preserve_audio=data.get('preserve_audio', True),
|
|
webhook_url=data.get('webhook_url'),
|
|
transcode_format=data.get('transcode_format', 'mp4'),
|
|
is_preview=data.get('is_preview', False)
|
|
)
|
|
for k, v in data.items():
|
|
setattr(job, k, v)
|
|
job._processes = []
|
|
job._is_cancelled = False
|
|
job._is_paused = data.get('_is_paused', False) or (data.get('status') == 'paused')
|
|
job._lock = threading.Lock()
|
|
return job
|
|
|
|
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 pause(self):
|
|
with self._lock:
|
|
if self.status in ["queued", "pending"]:
|
|
self.status = "paused"
|
|
self.eta = "Paused"
|
|
elif self.status in ["analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling"]:
|
|
self._is_paused = True
|
|
self.status = "paused"
|
|
self.eta = "Paused"
|
|
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, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False) -> subprocess.Popen:
|
|
with self._lock:
|
|
if self._is_cancelled:
|
|
raise InterruptedError("Job was cancelled")
|
|
if getattr(self, "_is_paused", False):
|
|
raise InterruptedError("Job was paused")
|
|
|
|
p = subprocess.Popen(
|
|
cmd,
|
|
stdout=stdout,
|
|
stderr=stderr,
|
|
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]):
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] starting pipeline. Video path: {job.video_path}, Model: {job.model}, Scale: {job.scale}")
|
|
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
|
|
base_temp = job.temp_dir if getattr(job, "temp_dir", None) else TEMP_DIR
|
|
job_temp_dir = os.path.join(base_temp, 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 (Support Skipping on Resume)
|
|
skip_extraction = False
|
|
if os.path.exists(input_frames_dir):
|
|
extracted_files = sorted([f for f in os.listdir(input_frames_dir) if f.startswith("frame_")])
|
|
if len(extracted_files) > 0:
|
|
skip_extraction = True
|
|
print(f"Job {job.job_id}: Found existing input frames ({len(extracted_files)} frames). Skipping extraction step.")
|
|
job.total_frames = len(extracted_files)
|
|
|
|
if not skip_extraction:
|
|
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")
|
|
])
|
|
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] extracting frames. Command: {' '.join(extract_cmd)}")
|
|
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)
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] extraction completed. Extracted {actual_total} frames.")
|
|
if actual_total == 0:
|
|
raise RuntimeError("No frames extracted from video")
|
|
|
|
job.total_frames = actual_total
|
|
else:
|
|
actual_total = job.total_frames
|
|
|
|
# Step 2: Upscale Frames (Support Resuming by Skipping already upscaled frames)
|
|
if os.path.exists(output_frames_dir):
|
|
output_files = os.listdir(output_frames_dir)
|
|
skipped_frames = 0
|
|
for f in output_files:
|
|
if f.startswith("frame_") and f.endswith(".jpg"):
|
|
out_path = os.path.join(output_frames_dir, f)
|
|
if os.path.exists(out_path) and os.path.getsize(out_path) > 0:
|
|
in_path = os.path.join(input_frames_dir, f)
|
|
if os.path.exists(in_path):
|
|
try:
|
|
os.remove(in_path)
|
|
skipped_frames += 1
|
|
except Exception as ex:
|
|
print(f"Error removing resumed frame {in_path}: {ex}")
|
|
if skipped_frames > 0:
|
|
print(f"Job {job.job_id}: Skipping {skipped_frames} already upscaled frames.")
|
|
|
|
remaining_inputs = len(os.listdir(input_frames_dir)) if os.path.exists(input_frames_dir) else 0
|
|
is_coordinator = False
|
|
|
|
if remaining_inputs == 0:
|
|
print(f"Job {job.job_id}: All frames already upscaled. Skipping upscaling step.")
|
|
job.update_status("upscaling", progress=80.0, current_frame=actual_total)
|
|
on_progress_update(job.job_id, {"status": "upscaling", "progress": 80.0, "current_frame": actual_total, "total_frames": actual_total})
|
|
else:
|
|
job.update_status("upscaling", progress=20, current_frame=actual_total - remaining_inputs)
|
|
on_progress_update(job.job_id, {"status": "upscaling", "progress": 20, "current_frame": actual_total - remaining_inputs, "total_frames": actual_total})
|
|
|
|
# Check if role in settings is "coordinator"
|
|
settings = {}
|
|
main_mod = sys.modules.get("app.main")
|
|
if main_mod and hasattr(main_mod, "load_settings"):
|
|
try:
|
|
settings = main_mod.load_settings()
|
|
except Exception:
|
|
pass
|
|
if not settings:
|
|
settings_path = os.path.join(BASE_DIR, "settings.json")
|
|
if os.path.exists(settings_path):
|
|
try:
|
|
with open(settings_path, "r") as f:
|
|
settings = json.load(f)
|
|
except Exception:
|
|
pass
|
|
|
|
is_coordinator = settings.get("role") == "coordinator"
|
|
skip_local_upscale = False
|
|
if is_coordinator:
|
|
skip_local_upscale = True
|
|
files_to_upscale = sorted([f for f in os.listdir(input_frames_dir) if f.startswith("frame_")])
|
|
chunk_size = 50
|
|
chunks_list = []
|
|
for idx_chunk, i in enumerate(range(0, len(files_to_upscale), chunk_size)):
|
|
chunk_files = files_to_upscale[i : i + chunk_size]
|
|
chunks_list.append({
|
|
"chunk_id": f"{job.job_id}_{idx_chunk}",
|
|
"files": chunk_files,
|
|
"status": "pending",
|
|
"worker_url": None,
|
|
"updated_at": time.time()
|
|
})
|
|
with dist_chunks_lock:
|
|
dist_chunks[job.job_id] = chunks_list
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] Coordinator mode active. Initialized {len(chunks_list)} chunks.")
|
|
|
|
current_tile_size = job.tile_size
|
|
while not skip_local_upscale:
|
|
# 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()
|
|
upscale_stdout_path = os.path.join(job_temp_dir, "upscale_stdout.log")
|
|
upscale_stderr_path = os.path.join(job_temp_dir, "upscale_stderr.log")
|
|
|
|
with open(upscale_stdout_path, "w") as f_out, open(upscale_stderr_path, "w") as f_err:
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] starting realesrgan upscaling. Model: {job.model}, Tile size: {current_tile_size}. Command: {' '.join(upscale_cmd)}")
|
|
p_upscale = job.run_command(upscale_cmd, stdout=f_out, stderr=f_err)
|
|
|
|
# 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
|
|
this_run_processed = processed_files - (actual_total - remaining_inputs)
|
|
if this_run_processed > 0:
|
|
sec_per_frame = elapsed / this_run_processed
|
|
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)
|
|
|
|
# Read stdout/stderr from files
|
|
if os.path.exists(upscale_stdout_path):
|
|
with open(upscale_stdout_path, "r") as f_out:
|
|
stdout = f_out.read()
|
|
else:
|
|
stdout = ""
|
|
|
|
if os.path.exists(upscale_stderr_path):
|
|
with open(upscale_stderr_path, "r") as f_err:
|
|
stderr = f_err.read()
|
|
else:
|
|
stderr = ""
|
|
|
|
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 only output frames that we attempted to upscale in this run
|
|
for filename in os.listdir(input_frames_dir):
|
|
out_path = os.path.join(output_frames_dir, filename)
|
|
if os.path.exists(out_path):
|
|
try:
|
|
os.unlink(out_path)
|
|
except Exception:
|
|
pass
|
|
continue
|
|
|
|
raise RuntimeError(f"Real-ESRGAN failed with exit code {p_upscale.returncode}: {err_msg}")
|
|
else:
|
|
break
|
|
|
|
# If coordinator, wait for all chunks to be completed
|
|
if is_coordinator:
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] Coordinator waiting for all chunks to complete...")
|
|
while True:
|
|
if job._is_cancelled:
|
|
return
|
|
if getattr(job, "_is_paused", False) or job.status == "paused":
|
|
return
|
|
|
|
with dist_chunks_lock:
|
|
chunks = dist_chunks.get(job.job_id, [])
|
|
if not chunks:
|
|
break
|
|
|
|
all_done = all(c["status"] == "completed" for c in chunks)
|
|
completed_count = sum(1 for c in chunks if c["status"] == "completed")
|
|
total_chunks = len(chunks)
|
|
|
|
upscale_progress = 20.0
|
|
if total_chunks > 0:
|
|
upscale_progress += (completed_count / total_chunks) * 60.0
|
|
|
|
processed_files = len(os.listdir(output_frames_dir))
|
|
job.update_status(
|
|
"upscaling",
|
|
progress=upscale_progress,
|
|
current_frame=processed_files,
|
|
eta=f"Waiting for workers... Chunks: {completed_count}/{total_chunks}"
|
|
)
|
|
on_progress_update(job.job_id, {
|
|
"status": "upscaling",
|
|
"progress": upscale_progress,
|
|
"current_frame": processed_files,
|
|
"total_frames": actual_total,
|
|
"eta": f"Workers processing chunks: {completed_count}/{total_chunks}"
|
|
})
|
|
|
|
if all_done:
|
|
break
|
|
|
|
# Timeout check: reset chunk if assigned but no update in 60s
|
|
with dist_chunks_lock:
|
|
for c in chunks:
|
|
if c["status"] == "assigned" and time.time() - c["updated_at"] > 60:
|
|
print(f"Chunk {c['chunk_id']} timed out. Requeuing.")
|
|
c["status"] = "pending"
|
|
c["worker_url"] = None
|
|
c["updated_at"] = time.time()
|
|
|
|
time.sleep(1.0)
|
|
|
|
# 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 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"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] starting GFPGAN 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):
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] starting RIFE frame interpolation.")
|
|
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})
|
|
|
|
# Get source filename basename, append with _upscaled_YYYY-MM-DD
|
|
source_file = job.source_filename if getattr(job, "source_filename", None) else job.video_path
|
|
base_name = os.path.basename(source_file)
|
|
name_without_ext, _ = os.path.splitext(base_name)
|
|
|
|
date_str = time.strftime("%Y-%m-%d")
|
|
transcode_fmt = getattr(job, "transcode_format", "mp4")
|
|
out_filename = f"{name_without_ext}_upscaled_{date_str}.{transcode_fmt}"
|
|
|
|
output_dir = job.output_dir if getattr(job, "output_dir", None) else OUTPUT_DIR
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
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"
|
|
|
|
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
|
|
|
|
# RIFE outputs files as %08d.jpg, upscaler outputs as frame_%08d.jpg
|
|
frame_pattern = "%08d.jpg" if use_rife else "frame_%08d.jpg"
|
|
|
|
# Construct ffmpeg reassembly command
|
|
assemble_cmd = [
|
|
"ffmpeg", "-y",
|
|
"-framerate", str(assemble_fps),
|
|
"-i", os.path.join(assemble_frames_dir, frame_pattern)
|
|
]
|
|
|
|
# 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?"])
|
|
if getattr(job, "ai_audio_denoise", False) and has_ffmpeg_filter("arnnoise"):
|
|
acodec_denoise = "libvorbis" if transcode_fmt == "webm" else "aac"
|
|
assemble_cmd.extend([
|
|
"-af", "arnnoise",
|
|
"-c:a", acodec_denoise
|
|
])
|
|
else:
|
|
if getattr(job, "ai_audio_denoise", False):
|
|
print(f"Warning: ai_audio_denoise was requested but 'arnnoise' filter is not available in ffmpeg. Copying audio without denoise.")
|
|
assemble_cmd.extend(["-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, "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
|
|
# Use 'framerate' filter instead of 'minterpolate' to avoid memory exhaustion on high resolutions
|
|
vf_filters.append(f"framerate=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
|
|
])
|
|
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] starting video assembly. Command: {' '.join(assemble_cmd)}")
|
|
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
|
|
})
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] pipeline completed successfully. Output file: {job.output_file}")
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
if not job._is_cancelled and not getattr(job, "_is_paused", False) and job.status not in ["paused", "queued"] and "paused" not in str(e).lower():
|
|
job.update_status("failed", error=str(e))
|
|
on_progress_update(job.job_id, {"status": "failed", "error": str(e)})
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] pipeline failed. Error: {e}")
|
|
else:
|
|
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] pipeline halted. Status: {job.status}")
|
|
finally:
|
|
# Clean up temp frames to save space only if completed or cancelled
|
|
try:
|
|
if job.status in ["completed", "cancelled"]:
|
|
if os.path.exists(job_temp_dir):
|
|
shutil.rmtree(job_temp_dir)
|
|
else:
|
|
print(f"Job {job.job_id} finished with status {job.status}. Preserving temp directory {job_temp_dir} for potential resume.")
|
|
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()
|