From f610c86f8ba95a8184c165db257d9cc1f27e6785 Mon Sep 17 00:00:00 2001 From: david kifer Date: Mon, 29 Jun 2026 13:13:08 -0400 Subject: [PATCH] Implement distributed cluster upscaling, custom output dirs, date-based naming, active preemption, and SMB mappings - 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 --- app/main.py | 573 +++++++++++++++++++++++++++++++++++++++++++--- app/upscaler.py | 113 ++++++++- static/app.js | 435 ++++++++++++++++++++++++++++++++++- static/index.html | 153 +++++++++++++ 4 files changed, 1228 insertions(+), 46 deletions(-) diff --git a/app/main.py b/app/main.py index 7dc4efa..830f246 100644 --- a/app/main.py +++ b/app/main.py @@ -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") diff --git a/app/upscaler.py b/app/upscaler.py index a5d98e4..4310cc3 100644 --- a/app/upscaler.py +++ b/app/upscaler.py @@ -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}") diff --git a/static/app.js b/static/app.js index a5a055d..d2bfc78 100644 --- a/static/app.js +++ b/static/app.js @@ -6,6 +6,16 @@ let socket = null; let sliderInitialized = false; // Error formatting helper +function getDownloadUrl(filePath) { + if (!filePath) return ""; + if (filePath.startsWith("http://") || filePath.startsWith("https://")) return filePath; + if (filePath.startsWith("/")) { + const filename = filePath.split("/").pop(); + return `/api/download/${encodeURIComponent(filename)}?path=${encodeURIComponent(filePath)}`; + } + return `/api/download/${encodeURIComponent(filePath)}`; +} + function formatFetchError(err, defaultMsg) { if (!err) return defaultMsg; if (typeof err === "string") return err; @@ -1000,8 +1010,8 @@ function updateProgressUI(data) { vidOriginal.style.display = "block"; vidUpscaledWrapper.style.display = "block"; - vidOriginal.src = `/api/download/${data.original_preview_file}`; - vidUpscaled.src = `/api/download/${data.output_file}`; + vidOriginal.src = getDownloadUrl(data.original_preview_file); + vidUpscaled.src = getDownloadUrl(data.output_file); vidOriginal.load(); vidUpscaled.load(); @@ -1060,7 +1070,7 @@ function updateProgressUI(data) { if (socket) socket.close(); stepProgress.classList.remove("active"); stepFinished.classList.add("active"); - downloadLink.href = `/api/download/${data.output_file}`; + downloadLink.href = getDownloadUrl(data.output_file); loadGallery(); }, 1500); @@ -1302,7 +1312,7 @@ if (finishedPlayComparisonBtn) { if (job) { const originalFilename = job.video_path.split("/").pop(); const originalUrl = `/api/uploads/${originalFilename}`; - const upscaledUrl = `/api/download/${job.output_file.split("/").pop()}`; + const upscaledUrl = getDownloadUrl(job.output_file); openVideoComparisonModal(originalUrl, upscaledUrl); } } catch (err) { @@ -1521,10 +1531,16 @@ async function loadQueue() { let actionsHtml = ""; let reorderHtml = ""; - if (job.status === "queued" && job.queue_position !== null && job.queue_position !== undefined) { - const queuedJobs = displayJobs.filter(j => j.status === "queued"); - const isFirst = job.queue_position === 0; - const isLast = job.queue_position === queuedJobs.length - 1; + const isActiveOrQueued = ["queued", "analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling"].includes(job.status); + if (isActiveOrQueued && job.queue_position !== null && job.queue_position !== undefined) { + const activeQueuedJobs = displayJobs.filter(j => + ["queued", "analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling"].includes(j.status) + ).filter(j => j.queue_position !== null && j.queue_position !== undefined); + + const minPos = Math.min(...activeQueuedJobs.map(j => j.queue_position)); + const maxPos = Math.max(...activeQueuedJobs.map(j => j.queue_position)); + const isFirst = job.queue_position === minPos; + const isLast = job.queue_position === maxPos; reorderHtml = ` + + `; + + const deleteBtn = item.querySelector(".btn-delete-smb"); + if (deleteBtn) { + deleteBtn.addEventListener("click", async (e) => { + e.stopPropagation(); + if (confirm(`Are you sure you want to unmap SMB share "${share.name}"?`)) { + await deleteSmbShare(share.name); + } + }); + } + + smbSharesList.appendChild(item); + }); + } + } + } catch (err) { + console.error("Error loading cluster settings:", err); + } +} + +async function deleteSmbShare(shareName) { + try { + const getRes = await fetch("/api/settings"); + if (!getRes.ok) throw new Error("Failed to load settings."); + const settings = await getRes.json(); + + if (settings.smb_shares) { + settings.smb_shares = settings.smb_shares.filter(s => s.name !== shareName); + } + + const saveRes = await fetch("/api/settings", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(settings) + }); + + if (saveRes.ok) { + loadClusterSettings(); + } else { + const err = await saveRes.json().catch(() => ({})); + alert(`Failed to delete SMB share: ${err.detail || "Unknown error"}`); + } + } catch (err) { + console.error("Error deleting SMB share:", err); + alert("Network error deleting SMB share."); + } +} + async function loadDefaultTempDir() { try { const res = await fetch("/api/settings"); if (!res.ok) return; const settings = await res.json(); + + // Temp dir settings const savedDir = settings.default_temp_dir || ""; const defaultInput = document.getElementById("default-temp-dir-input"); if (defaultInput) defaultInput.value = savedDir; - // Update per-job input placeholder const perJobInput = document.getElementById("temp-dir-input"); if (perJobInput && savedDir && !perJobInput.value.trim()) { perJobInput.placeholder = `Default: ${savedDir} (override here)`; } + + // Output dir settings + const savedOutputDir = settings.default_output_dir || ""; + const defaultOutputDirInput = document.getElementById("default-output-dir-input"); + if (defaultOutputDirInput) defaultOutputDirInput.value = savedOutputDir; + const perJobOutputDirInput = document.getElementById("output-dir-input"); + if (perJobOutputDirInput && savedOutputDir && !perJobOutputDirInput.value.trim()) { + perJobOutputDirInput.placeholder = `Default: ${savedOutputDir} (override here)`; + } } catch (err) { console.error("Failed to load settings:", err); } diff --git a/static/index.html b/static/index.html index ca16c67..f2dda27 100644 --- a/static/index.html +++ b/static/index.html @@ -290,6 +290,20 @@

+ +
+ +
+ + +
+

+ Specify a custom location to save the final upscaled video file. +

+
+
Pre-Filters & Enhancing @@ -609,6 +623,95 @@
+ +
+
+ + Default Output Directory +
+

+ Set a persistent default directory where completed upscale videos will be saved. This is applied automatically unless overridden per-job in Advanced Options. +

+
+ + + +
+
+
+ + +
+
+ + Distributed Processing & Node Roles +
+

+ Link multiple video upscaler servers together to accelerate upscaling tasks by distributing frame processing workload. +

+ +
+
+ + +
+ + + + +
+ + +
+ + +
+ + +
+

+ SMB Network Shares Mapping + +

+
+ +
+ No mapped network shares. +
+
+
+ + +
+
+
Total Temp Size @@ -768,6 +871,56 @@
+ + +