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

Implement distributed cluster upscaling, custom output dirs, date-based naming, active preemption, and SMB mappings

Browse Source 
- Add Output Destination folder selector with file browser and fallback to outputs/
- Apply YYYY-MM-DD start-date suffix naming rule to upscaled outputs
- Implement active job preemption (demoting active running job back to queue) via reorder
- Implement Coordinator / Worker roles with automatic pull-based chunk processing
- Support Shared Storage mode and HTTP-based frame upload/download fallback
- Add SMB mapped share configuration and background CIFS mount support
- Update index.html and app.js with the new Settings sections and modals
This commit is contained in:
david
2026-06-29 13:13:08 -04:00
parent f662c5806f
commit f610c86f8b
4 changed files with 1228 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/main.py
+544 -29
View File
@@ -4,7 +4,9 @@ import queue
import threading
import asyncio
import subprocess
import time
import urllib.request
import urllib.parse
import json
from typing import Dict, List, Any
from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks, HTTPException, WebSocket, WebSocketDisconnect
@@ -90,7 +92,16 @@ SETTINGS_FILE = os.path.join(upscaler.BASE_DIR, "settings.json")
def load_settings() -> dict:
"""Load persistent app settings from settings.json"""
defaults = {"default_temp_dir": ""}
defaults = {
"default_temp_dir": "",
"default_output_dir": "",
"role": "standalone",
"coordinator_url": "",
"worker_node_url": "http://localhost:8000",
"use_shared_storage": False,
"shared_storage_path": "",
"smb_shares": []
}
if os.path.exists(SETTINGS_FILE):
try:
with open(SETTINGS_FILE, "r") as f:
@@ -141,6 +152,7 @@ def startup_event():
global main_loop
main_loop = asyncio.get_event_loop()
load_jobs_db()
restart_dist_worker()
global_webhook_url = None
@@ -259,6 +271,7 @@ class StartUpscaleRequest(BaseModel):
ai_rife_interpolation: bool = False
ai_audio_denoise: bool = False
temp_dir: str | None = None
output_dir: str | None = None
class StartBatchUpscaleRequest(BaseModel):
file_ids: List[str]
@@ -287,6 +300,7 @@ class StartBatchUpscaleRequest(BaseModel):
ai_rife_interpolation: bool = False
ai_audio_denoise: bool = False
temp_dir: str | None = None
output_dir: str | None = None
class PreviewRequest(BaseModel):
file_id: str
@@ -447,6 +461,18 @@ def start_upscale(req: StartUpscaleRequest):
if not file_path:
raise HTTPException(status_code=404, detail="Uploaded file not found.")
metadata = load_upload_metadata()
source_filename = None
if req.file_id in metadata:
source_filename = metadata[req.file_id].get("original_filename")
if not source_filename and file_path:
source_filename = os.path.basename(file_path)
settings = load_settings()
output_dir = req.output_dir or settings.get("default_output_dir") or None
if output_dir:
output_dir = output_dir.strip()
job_id = str(uuid.uuid4())
job = upscaler.UpscaleJob(
job_id=job_id,
@@ -475,7 +501,9 @@ def start_upscale(req: StartUpscaleRequest):
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 or load_settings().get("default_temp_dir") or None
temp_dir=req.temp_dir or settings.get("default_temp_dir") or None,
output_dir=output_dir,
source_filename=source_filename
)
jobs_db[job_id] = job
@@ -513,6 +541,11 @@ def start_upscale_batch(req: StartBatchUpscaleRequest):
if not file_path:
raise HTTPException(status_code=404, detail=f"Uploaded file {file_id} not found.")
settings = load_settings()
output_dir = req.output_dir or settings.get("default_output_dir") or None
if output_dir:
output_dir = output_dir.strip()
# 2. Queueing loop
for file_id in req.file_ids:
file_path = None
@@ -522,6 +555,13 @@ def start_upscale_batch(req: StartBatchUpscaleRequest):
file_path = test_path
break
metadata = load_upload_metadata()
source_filename = None
if file_id in metadata:
source_filename = metadata[file_id].get("original_filename")
if not source_filename and file_path:
source_filename = os.path.basename(file_path)
job_id = str(uuid.uuid4())
job = upscaler.UpscaleJob(
job_id=job_id,
@@ -550,7 +590,9 @@ def start_upscale_batch(req: StartBatchUpscaleRequest):
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 or load_settings().get("default_temp_dir") or None
temp_dir=req.temp_dir or settings.get("default_temp_dir") or None,
output_dir=output_dir,
source_filename=source_filename
)
jobs_db[job_id] = job
job_queue.put(job_id)
@@ -685,11 +727,40 @@ def get_models():
return sorted(list(models))
@app.get("/api/download/{filename}")
def download_file(filename: str):
file_path = os.path.join(upscaler.OUTPUT_DIR, filename)
def download_file(filename: str, path: str | None = None):
allowed_dirs = {os.path.abspath(upscaler.OUTPUT_DIR)}
settings = load_settings()
default_out = settings.get("default_output_dir", "").strip()
if default_out:
allowed_dirs.add(os.path.abspath(default_out))
for job in jobs_db.values():
if getattr(job, "output_dir", None):
allowed_dirs.add(os.path.abspath(job.output_dir))
if path:
resolved_path = os.path.abspath(path)
if os.path.isdir(resolved_path):
file_path = os.path.join(resolved_path, filename)
else:
file_path = resolved_path
else:
file_path = os.path.join(upscaler.OUTPUT_DIR, filename)
resolved_file_path = os.path.abspath(file_path)
allowed = False
for allowed_dir in allowed_dirs:
if resolved_file_path.startswith(allowed_dir):
allowed = True
break
if not allowed:
raise HTTPException(status_code=403, detail="Access denied.")
if not os.path.exists(file_path):
raise HTTPException(status_code=404, detail="File not found.")
return FileResponse(file_path, media_type="application/octet-stream", filename=filename)
return FileResponse(file_path, media_type="application/octet-stream", filename=os.path.basename(file_path))
def safe_delete_file(file_path: str):
if file_path and os.path.exists(file_path):
@@ -814,15 +885,53 @@ class ReorderQueueRequest(BaseModel):
@app.post("/api/queue/reorder")
def reorder_queue(req: ReorderQueueRequest):
"""Reorder the job queue"""
"""Reorder the job queue with preemption support"""
global current_running_job_id
running_id = current_running_job_id
if running_id and running_id in req.job_ids:
idx = req.job_ids.index(running_id)
if idx > 0:
# Preempt the running job!
job = jobs_db.get(running_id)
if job:
print(f"Preempting running job {running_id} (moved to index {idx} in queue)")
try:
job.pause()
except Exception as e:
print(f"Error pausing job {running_id} for preemption: {e}")
# Re-queue the job
job.status = "queued"
job.error = None
job.eta = "Queued (preempted)..."
if hasattr(job, "_is_paused"):
job._is_paused = False
job_queue.put(running_id)
broadcast_progress(running_id, {
"status": "queued",
"progress": job.progress,
"current_frame": job.current_frame,
"total_frames": job.total_frames,
"eta": "Queued (preempted)..."
})
job_queue.reorder(req.job_ids)
save_jobs_db()
return {"status": "success", "queue": job_queue.get_all()}
q = job_queue.get_all()
if current_running_job_id:
q = [current_running_job_id] + q
return {"status": "success", "queue": q}
@app.get("/api/queue")
def get_queue():
"""Get the current job queue order"""
return {"queue": job_queue.get_all()}
q = job_queue.get_all()
if current_running_job_id:
q = [current_running_job_id] + q
return {"queue": q}
@app.post("/api/upscale/resume/{job_id}")
def resume_job(job_id: str):
@@ -999,32 +1108,83 @@ def get_diagnostics():
# Outputs Gallery Endpoints
@app.get("/api/outputs")
def list_outputs():
output_dir = upscaler.OUTPUT_DIR
if not os.path.exists(output_dir):
return []
output_dirs = {upscaler.OUTPUT_DIR}
settings = load_settings()
default_out = settings.get("default_output_dir", "").strip()
if default_out:
output_dirs.add(default_out)
for job in jobs_db.values():
if job.status == "completed" and job.output_file:
output_dirs.add(os.path.dirname(job.output_file))
elif getattr(job, "output_dir", None):
output_dirs.add(job.output_dir)
files = []
for filename in os.listdir(output_dir):
if filename.startswith("original_"):
seen_files = set()
for out_dir in output_dirs:
if not os.path.exists(out_dir):
continue
file_path = os.path.join(output_dir, filename)
if os.path.isfile(file_path):
stat = os.stat(file_path)
files.append({
"filename": filename,
"size": stat.st_size,
"modified": stat.st_mtime,
"url": f"/api/download/{filename}"
})
try:
for filename in os.listdir(out_dir):
if filename.startswith("original_"):
continue
file_path = os.path.join(out_dir, filename)
abs_file_path = os.path.abspath(file_path)
if abs_file_path in seen_files:
continue
if os.path.isfile(file_path):
seen_files.add(abs_file_path)
stat = os.stat(file_path)
if os.path.abspath(out_dir) == os.path.abspath(upscaler.OUTPUT_DIR):
url = f"/api/download/{filename}"
else:
url = f"/api/download/{filename}?path={urllib.parse.quote(out_dir)}"
files.append({
"filename": filename,
"size": stat.st_size,
"modified": stat.st_mtime,
"url": url,
"path": file_path
})
except Exception as e:
print(f"Error scanning output dir {out_dir}: {e}")
files.sort(key=lambda x: x["modified"], reverse=True)
return files
@app.delete("/api/outputs/delete/{filename}")
@app.post("/api/outputs/delete/{filename}")
def delete_output(filename: str):
file_path = os.path.join(upscaler.OUTPUT_DIR, filename)
resolved_path = os.path.abspath(file_path)
if not resolved_path.startswith(os.path.abspath(upscaler.OUTPUT_DIR)):
def delete_output(filename: str, path: str | None = None):
allowed_dirs = {os.path.abspath(upscaler.OUTPUT_DIR)}
settings = load_settings()
default_out = settings.get("default_output_dir", "").strip()
if default_out:
allowed_dirs.add(os.path.abspath(default_out))
for job in jobs_db.values():
if getattr(job, "output_dir", None):
allowed_dirs.add(os.path.abspath(job.output_dir))
if path:
resolved_path = os.path.abspath(path)
if os.path.isdir(resolved_path):
file_path = os.path.join(resolved_path, filename)
else:
file_path = resolved_path
else:
file_path = os.path.join(upscaler.OUTPUT_DIR, filename)
resolved_file_path = os.path.abspath(file_path)
allowed = False
for allowed_dir in allowed_dirs:
if resolved_file_path.startswith(allowed_dir):
allowed = True
break
if not allowed:
raise HTTPException(status_code=403, detail="Access denied.")
if not os.path.exists(file_path):
@@ -1090,6 +1250,13 @@ def setup_webhook(req: WebhookSetupRequest):
# ─── Settings Endpoints ─────────────────────────────────────────────
class UpdateSettingsRequest(BaseModel):
default_temp_dir: str = ""
default_output_dir: str = ""
role: str = "standalone"
coordinator_url: str = ""
worker_node_url: str = "http://localhost:8000"
use_shared_storage: bool = False
shared_storage_path: str = ""
smb_shares: List[Any] = []
@app.get("/api/settings")
def get_settings():
@@ -1165,12 +1332,38 @@ def update_settings(req: UpdateSettingsRequest, background_tasks: BackgroundTask
try:
os.makedirs(new_dir, exist_ok=True)
except OSError as e:
raise HTTPException(status_code=400, detail=f"Invalid directory path: {e}")
raise HTTPException(status_code=400, detail=f"Invalid default_temp_dir path: {e}")
new_out_dir = req.default_output_dir.strip()
if new_out_dir:
try:
os.makedirs(new_out_dir, exist_ok=True)
except OSError as e:
raise HTTPException(status_code=400, detail=f"Invalid default_output_dir path: {e}")
new_shared_path = req.shared_storage_path.strip()
if req.use_shared_storage and new_shared_path:
try:
os.makedirs(new_shared_path, exist_ok=True)
except OSError as e:
raise HTTPException(status_code=400, detail=f"Invalid shared_storage_path path: {e}")
role_changed = (settings.get("role") != req.role)
# Save settings
settings["default_temp_dir"] = new_dir
settings["default_output_dir"] = new_out_dir
settings["role"] = req.role
settings["coordinator_url"] = req.coordinator_url
settings["worker_node_url"] = req.worker_node_url
settings["use_shared_storage"] = req.use_shared_storage
settings["shared_storage_path"] = req.shared_storage_path
settings["smb_shares"] = req.smb_shares
save_settings(settings)
if role_changed:
restart_dist_worker()
# Identify and pause active jobs
active_jobs = []
old_real = os.path.realpath(old_dir)
@@ -1400,5 +1593,327 @@ def cleanup_temp(req: CleanupTempRequest):
save_jobs_db()
return {"cleaned": cleaned, "errors": errors}
# ─── Distributed Processing Coordinator & Worker Endpoints ────────────────
class ChunkDoneRequest(BaseModel):
job_id: str
chunk_id: str
@app.get("/api/dist/request-chunk")
def request_chunk(worker_url: str = ""):
settings = load_settings()
if settings.get("role") != "coordinator":
raise HTTPException(status_code=400, detail="Node is not configured as a coordinator.")
with upscaler.dist_chunks_lock:
job_id = current_running_job_id
if not job_id:
return {"status": "idle"}
job = jobs_db.get(job_id)
if not job or job.status != "upscaling":
return {"status": "idle"}
chunks = upscaler.dist_chunks.get(job_id, [])
for chunk in chunks:
if chunk["status"] == "pending":
chunk["status"] = "assigned"
chunk["worker_url"] = worker_url
chunk["updated_at"] = time.time()
return {
"status": "assigned",
"job_id": job_id,
"chunk_id": chunk["chunk_id"],
"model": job.model,
"scale": job.scale,
"tile_size": job.tile_size,
"gpu_ids": job.gpu_ids,
"tta": job.tta,
"files": chunk["files"],
"use_shared_storage": settings.get("use_shared_storage", False),
"shared_storage_path": settings.get("shared_storage_path", ""),
"job_temp_dir": os.path.join(settings.get("shared_storage_path") if settings.get("use_shared_storage") else upscaler.TEMP_DIR, job_id)
}
return {"status": "no_chunks"}
@app.get("/api/dist/download-frame/{job_id}/{filename}")
def download_frame(job_id: str, filename: str):
settings = load_settings()
if settings.get("role") != "coordinator":
raise HTTPException(status_code=400, detail="Node is not configured as a coordinator.")
if "/" in filename or "\\" in filename or filename in [".", ".."]:
raise HTTPException(status_code=400, detail="Invalid filename")
base_temp = settings.get("shared_storage_path") if settings.get("use_shared_storage") else upscaler.TEMP_DIR
file_path = os.path.join(base_temp, job_id, "input_frames", filename)
if not os.path.exists(file_path):
raise HTTPException(status_code=404, detail="Frame not found")
return FileResponse(file_path)
@app.post("/api/dist/upload-chunk/{job_id}")
async def upload_chunk(job_id: str, files: List[UploadFile] = File(...)):
settings = load_settings()
if settings.get("role") != "coordinator":
raise HTTPException(status_code=400, detail="Node is not configured as a coordinator.")
base_temp = settings.get("shared_storage_path") if settings.get("use_shared_storage") else upscaler.TEMP_DIR
output_frames_dir = os.path.join(base_temp, job_id, "output_frames")
os.makedirs(output_frames_dir, exist_ok=True)
for file in files:
filename = os.path.basename(file.filename)
if "/" in filename or "\\" in filename or filename in [".", ".."]:
continue
save_path = os.path.join(output_frames_dir, filename)
with open(save_path, "wb") as buffer:
content = await file.read()
buffer.write(content)
return {"status": "success"}
@app.post("/api/dist/chunk-done")
def chunk_done(req: ChunkDoneRequest):
settings = load_settings()
if settings.get("role") != "coordinator":
raise HTTPException(status_code=400, detail="Node is not configured as a coordinator.")
with upscaler.dist_chunks_lock:
chunks = upscaler.dist_chunks.get(req.job_id, [])
for chunk in chunks:
if chunk["chunk_id"] == req.chunk_id:
chunk["status"] = "completed"
chunk["updated_at"] = time.time()
return {"status": "completed"}
raise HTTPException(status_code=404, detail="Chunk not found.")
# ─── SMB Mounting & Background Worker Thread ───────────────────────────────
def mount_smb_share(share: dict) -> bool:
"""Mount an SMB share if not already mounted"""
local_path = share.get("local_path", "").strip()
share_path = share.get("share_path", "").strip()
username = share.get("username", "").strip()
password = share.get("password", "").strip()
if not local_path or not share_path:
print("SMB Share: missing local_path or share_path")
return False
try:
os.makedirs(local_path, exist_ok=True)
except Exception as e:
print(f"SMB Share: failed to create mount point {local_path}: {e}")
return False
try:
with open("/proc/mounts", "r") as f:
mounts = f.read()
if os.path.realpath(local_path) in mounts or share_path in mounts:
print(f"SMB Share: {share_path} is already mounted at {local_path}")
return True
except Exception:
pass
options = []
if username:
options.append(f"username={username}")
if password:
options.append(f"password={password}")
cmd = ["mount", "-t", "cifs", share_path, local_path]
if options:
cmd.extend(["-o", ",".join(options)])
try:
print(f"Mounting SMB share: {' '.join(cmd)}")
res = subprocess.run(cmd, capture_output=True, text=True)
if res.returncode == 0:
print(f"SMB Share: successfully mounted {share_path} to {local_path}")
return True
else:
print(f"SMB Share: failed to mount (exit {res.returncode}): {res.stderr}")
return False
except Exception as e:
print(f"SMB Share: exception mounting: {e}")
return False
def upload_files_multipart(url: str, files: List[tuple[str, str]]):
"""Upload files using multipart/form-data via urllib"""
import uuid
import mimetypes
boundary = f"----WebKitFormBoundary{uuid.uuid4().hex}"
parts = []
for field_name, file_path in files:
filename = os.path.basename(file_path)
mime_type = mimetypes.guess_type(file_path)[0] or "image/jpeg"
parts.append(f"--{boundary}".encode('utf-8'))
parts.append(f'Content-Disposition: form-data; name="{field_name}"; filename="{filename}"'.encode('utf-8'))
parts.append(f"Content-Type: {mime_type}".encode('utf-8'))
parts.append(b"")
with open(file_path, "rb") as f:
parts.append(f.read())
parts.append(f"--{boundary}--".encode('utf-8'))
body = b"\r\n".join(parts)
headers = {
"Content-Type": f"multipart/form-data; boundary={boundary}",
"Content-Length": str(len(body))
}
req = urllib.request.Request(url, data=body, headers=headers, method="POST")
with urllib.request.urlopen(req) as res:
return res.read()
def process_dist_chunk(chunk: dict) -> bool:
job_id = chunk["job_id"]
chunk_id = chunk["chunk_id"]
model = chunk["model"]
scale = chunk["scale"]
tile_size = chunk["tile_size"]
gpu_ids = chunk.get("gpu_ids")
files = chunk["files"]
use_shared_storage = chunk.get("use_shared_storage", False)
coordinator_url = load_settings().get("coordinator_url", "").rstrip("/")
print(f"Worker: Processing chunk {chunk_id} for job {job_id} ({len(files)} files)")
local_job_temp = os.path.join(upscaler.TEMP_DIR, f"worker_{job_id}_{chunk_id}")
local_input_dir = os.path.join(local_job_temp, "input_frames")
local_output_dir = os.path.join(local_job_temp, "output_frames")
if not use_shared_storage:
os.makedirs(local_input_dir, exist_ok=True)
os.makedirs(local_output_dir, exist_ok=True)
try:
input_paths = []
output_paths = []
for filename in files:
if use_shared_storage:
job_temp_dir = chunk["job_temp_dir"]
in_path = os.path.join(job_temp_dir, "input_frames", filename)
out_path = os.path.join(job_temp_dir, "output_frames", filename)
else:
in_path = os.path.join(local_input_dir, filename)
out_path = os.path.join(local_output_dir, filename)
download_url = f"{coordinator_url}/api/dist/download-frame/{job_id}/{filename}"
try:
urllib.request.urlretrieve(download_url, in_path)
except Exception as dl_err:
print(f"Worker: Failed to download frame {filename}: {dl_err}")
return False
input_paths.append(in_path)
output_paths.append(out_path)
for in_p, out_p in zip(input_paths, output_paths):
success = upscaler.upscale_image_file(in_p, out_p, model, scale, tile_size, gpu_ids)
if not success:
print(f"Worker: Failed to upscale frame {in_p}")
return False
if not use_shared_storage:
upload_url = f"{coordinator_url}/api/dist/upload-chunk/{job_id}"
file_tuples = [("files", out_p) for out_p in output_paths]
try:
upload_files_multipart(upload_url, file_tuples)
except Exception as up_err:
print(f"Worker: Failed to upload output frames: {up_err}")
return False
done_url = f"{coordinator_url}/api/dist/chunk-done"
try:
req = urllib.request.Request(
done_url,
data=json.dumps({"job_id": job_id, "chunk_id": chunk_id}).encode('utf-8'),
headers={"Content-Type": "application/json"},
method="POST"
)
with urllib.request.urlopen(req) as res:
res.read()
except Exception as done_err:
print(f"Worker: Failed to notify done: {done_err}")
return False
return True
finally:
if not use_shared_storage and os.path.exists(local_job_temp):
try:
shutil.rmtree(local_job_temp)
except Exception:
pass
dist_worker_running = False
dist_worker_thread = None
dist_worker_stop_event = threading.Event()
def dist_worker_loop():
global dist_worker_running
dist_worker_running = True
print("Background worker loop started.")
settings = load_settings()
for share in settings.get("smb_shares", []):
mount_smb_share(share)
while not dist_worker_stop_event.is_set():
settings = load_settings()
role = settings.get("role", "standalone")
if role != "worker":
time.sleep(5.0)
continue
coordinator_url = settings.get("coordinator_url", "").strip().rstrip("/")
if not coordinator_url:
time.sleep(10.0)
continue
worker_node_url = settings.get("worker_node_url", "http://localhost:8000")
try:
req_url = f"{coordinator_url}/api/dist/request-chunk?worker_url={urllib.parse.quote(worker_node_url)}"
req = urllib.request.Request(req_url, method="GET")
with urllib.request.urlopen(req, timeout=10) as res:
chunk_data = json.loads(res.read().decode('utf-8'))
if chunk_data.get("status") == "assigned":
success = process_dist_chunk(chunk_data)
if not success:
print(f"Worker: failed to process chunk {chunk_data.get('chunk_id')}. Will retry later.")
else:
print(f"Worker: successfully processed chunk {chunk_data.get('chunk_id')}")
else:
time.sleep(3.0)
except Exception as e:
print(f"Worker loop error: {e}")
time.sleep(5.0)
dist_worker_running = False
print("Background worker loop stopped.")
def restart_dist_worker():
global dist_worker_thread, dist_worker_stop_event
dist_worker_stop_event.set()
if dist_worker_thread and dist_worker_thread.is_alive():
dist_worker_thread.join(timeout=2.0)
dist_worker_stop_event.clear()
dist_worker_thread = threading.Thread(target=dist_worker_loop, daemon=True)
dist_worker_thread.start()
# Mount static folder
app.mount("/", StaticFiles(directory=os.path.join(upscaler.BASE_DIR, "static"), html=True), name="static")
app/upscaler.py
+108 -5
View File
@@ -7,6 +7,9 @@ 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
@@ -39,9 +42,12 @@ 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, temp_dir: str = None):
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
@@ -404,8 +410,45 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
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 True:
while not skip_local_upscale:
# Launch Real-ESRGAN on directory
upscale_cmd = [
BIN_PATH,
@@ -509,6 +552,57 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
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)
@@ -599,9 +693,18 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
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"upscaled_{job.job_id}.{transcode_fmt}"
out_filepath = os.path.join(OUTPUT_DIR, out_filename)
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
@@ -697,7 +800,7 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
except Exception as e:
import traceback
traceback.print_exc()
if not job._is_cancelled and not getattr(job, "_is_paused", False) and job.status != "paused":
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}")