david/video_upscaler
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement batch processing, pause/resume, custom temp directory, and detailed pipeline logs

Browse Source
This commit is contained in:
david
2026-06-29 08:04:35 -04:00
parent d1ea029afb
commit d35a48bf64
5 changed files with 507 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/main.py
+135 -2
View File
@@ -235,6 +235,35 @@ class StartUpscaleRequest(BaseModel):
ai_face_restoration: bool = False
ai_rife_interpolation: bool = False
ai_audio_denoise: bool = False
temp_dir: str | None = None
class StartBatchUpscaleRequest(BaseModel):
file_ids: List[str]
model: str = "realesr-animevideov3"
scale: int = 4
tile_size: int = 256
preserve_audio: bool = True
ss: str | None = None
t: str | None = None
gpu_ids: str | None = None
tta: bool = False
unsharp: bool = False
double_fps: bool = False
preserve_subtitles: bool = True
start_sec: float | None = None
end_sec: float | None = None
crf: int = 18
preset: str = "medium"
denoise: bool = False
sharpen: bool = False
interpolation: bool = False
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
temp_dir: str | None = None
class PreviewRequest(BaseModel):
file_id: str
@@ -422,7 +451,8 @@ def start_upscale(req: StartUpscaleRequest):
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
ai_audio_denoise=req.ai_audio_denoise,
temp_dir=req.temp_dir
)
jobs_db[job_id] = job
@@ -443,6 +473,80 @@ def start_upscale(req: StartUpscaleRequest):
"status": "queued"
}
@app.post("/api/upscale/start/batch")
def start_upscale_batch(req: StartBatchUpscaleRequest):
"""Queue multiple upscaling tasks with one set of settings"""
queued_jobs = []
# 1. Validation loop
for file_id in req.file_ids:
file_path = None
exts = [".mp4", ".mkv", ".avi", ".mov", ".webm"]
for ext in exts:
test_path = os.path.join(upscaler.UPLOAD_DIR, f"{file_id}{ext}")
if os.path.exists(test_path):
file_path = test_path
break
if not file_path:
raise HTTPException(status_code=404, detail=f"Uploaded file {file_id} not found.")
# 2. Queueing loop
for file_id in req.file_ids:
file_path = None
for ext in [".mp4", ".mkv", ".avi", ".mov", ".webm"]:
test_path = os.path.join(upscaler.UPLOAD_DIR, f"{file_id}{ext}")
if os.path.exists(test_path):
file_path = test_path
break
job_id = str(uuid.uuid4())
job = upscaler.UpscaleJob(
job_id=job_id,
video_path=file_path,
model=req.model,
scale=req.scale,
tile_size=req.tile_size,
preserve_audio=req.preserve_audio,
ss=req.ss,
t=req.t,
gpu_ids=req.gpu_ids,
tta=req.tta,
unsharp=req.unsharp,
double_fps=req.double_fps,
preserve_subtitles=req.preserve_subtitles,
start_sec=req.start_sec,
end_sec=req.end_sec,
crf=req.crf,
preset=req.preset,
denoise=req.denoise,
sharpen=req.sharpen,
interpolation=req.interpolation,
webhook_url=req.webhook_url,
transcode_format=req.transcode_format,
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,
temp_dir=req.temp_dir
)
jobs_db[job_id] = job
job_queue.put(job_id)
queued_jobs.append({"job_id": job_id, "status": "queued"})
save_jobs_db()
# Broadcast initial queued progress for all queued jobs
for qj in queued_jobs:
broadcast_progress(qj["job_id"], {
"status": "queued",
"progress": 0.0,
"current_frame": 0,
"total_frames": 0,
"eta": "Calculating..."
})
return {"jobs": queued_jobs}
@app.get("/api/upscale/status/{job_id}")
def get_status(job_id: str):
"""Get status of upscale job"""
@@ -699,7 +803,7 @@ def get_queue():
@app.post("/api/upscale/resume/{job_id}")
def resume_job(job_id: str):
"""Resume an interrupted/failed upscale job"""
"""Resume an interrupted/failed/paused upscale job"""
job = jobs_db.get(job_id)
if not job:
raise HTTPException(status_code=404, detail="Job not found.")
@@ -708,6 +812,8 @@ def resume_job(job_id: str):
job.status = "queued"
job.error = None
job.eta = "Queued for resume..."
if hasattr(job, "_is_paused"):
job._is_paused = False
job_queue.put(job_id)
save_jobs_db()
@@ -722,6 +828,33 @@ def resume_job(job_id: str):
return {"job_id": job_id, "status": "queued"}
@app.post("/api/upscale/pause/{job_id}")
def pause_job(job_id: str):
"""Pause a running or queued job"""
job = jobs_db.get(job_id)
if not job:
raise HTTPException(status_code=404, detail="Job not found.")
if job.status not in ["queued", "pending", "analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling"]:
raise HTTPException(status_code=400, detail=f"Job in status {job.status} cannot be paused.")
# Remove from queue if it is in queue
job_queue.remove(job_id)
# Call pause logic on the job
job.pause()
save_jobs_db()
broadcast_progress(job_id, {
"status": "paused",
"progress": job.progress,
"current_frame": job.current_frame,
"total_frames": job.total_frames,
"eta": "Paused"
})
return {"job_id": job_id, "status": "paused"}
# Websocket endpoint for real-time progress updates
@app.websocket("/ws/progress/{job_id}")
async def websocket_progress(websocket: WebSocket, job_id: str):
app/upscaler.py
+41 -5
View File
@@ -26,8 +26,9 @@ class UpscaleJob:
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):
ai_audio_denoise: bool = False, temp_dir: str = None):
self.job_id = job_id
self.temp_dir = temp_dir
self.video_path = video_path
self.model = model
self.scale = scale
@@ -79,9 +80,10 @@ class UpscaleJob:
self.start_time = None
self.output_file = None
# Track processes to allow cancellation
# 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:
@@ -106,6 +108,7 @@ class UpscaleJob:
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
@@ -137,10 +140,32 @@ class UpscaleJob:
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,
@@ -240,6 +265,7 @@ def upscale_image_file(input_path: str, output_path: str, model: str, scale: int
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)
@@ -254,7 +280,8 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
fps = info["fps"]
# Create job temp directories
job_temp_dir = os.path.join(TEMP_DIR, job.job_id)
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")
@@ -299,6 +326,7 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
if not skip_extraction:
job.update_status("extracting", progress=10)
on_progress_update(job.job_id, {"status": "extracting", "progress": 10})
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [Job {job.job_id}] extracting frames. Command: {' '.join(extract_cmd)}")
# High quality JPG frames to balance disk usage and speed
extract_cmd = ["ffmpeg", "-y"]
@@ -327,6 +355,7 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
# 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")
@@ -384,6 +413,7 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
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
@@ -479,7 +509,7 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
gfpgan_installed = importlib.util.find_spec("gfpgan") is not None
if gfpgan_installed:
print(f"Job {job.job_id}: GFPGAN detected. Running Face Restoration...")
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)
@@ -521,6 +551,7 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
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)
@@ -630,6 +661,7 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
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)
@@ -645,13 +677,17 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
"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:
if not job._is_cancelled and not getattr(job, "_is_paused", False) and job.status != "paused":
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: