Files
video_upscaler/app/main.py
T
david 006eece417 Add persistent default temp directory setting with web UI
- Add settings.json persistence (load/save) in backend
- Add GET/POST /api/settings endpoints for default_temp_dir
- Add Default Temp Directory config UI in Temporary Storage section
- Auto-fallback: jobs use saved default if no per-job temp_dir set
- Per-job Advanced Options placeholder syncs with saved default
- Add settings.json to .gitignore for per-machine config
2026-06-29 09:51:12 -04:00

1260 lines
42 KiB
Python

import os
import uuid
import queue
import threading
import asyncio
import subprocess
import urllib.request
import json
from typing import Dict, List, Any
from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import app.upscaler as upscaler
app = FastAPI(title="AI Video Upscaler Dashboard")
# CORS middleware for local development
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# In-memory databases
class CustomJobQueue:
def __init__(self):
self.queue = []
self.lock = threading.Lock()
self.condition = threading.Condition(self.lock)
def put(self, job_id: str):
with self.lock:
if job_id not in self.queue:
self.queue.append(job_id)
self.condition.notify()
def get(self) -> str:
with self.lock:
while not self.queue:
self.condition.wait()
return self.queue.pop(0)
def remove(self, job_id: str) -> bool:
with self.lock:
if job_id in self.queue:
self.queue.remove(job_id)
return True
return False
def get_all(self) -> List[str]:
with self.lock:
return list(self.queue)
def reorder(self, job_ids: List[str]):
with self.lock:
valid_ids = [jid for jid in job_ids if jid in self.queue]
missing_ids = [jid for jid in self.queue if jid not in valid_ids]
self.queue = valid_ids + missing_ids
def task_done(self):
pass
def empty(self) -> bool:
with self.lock:
return len(self.queue) == 0
def qsize(self) -> int:
with self.lock:
return len(self.queue)
jobs_db: Dict[str, upscaler.UpscaleJob] = {}
ws_connections: Dict[str, List[WebSocket]] = {}
preview_db: Dict[str, Dict[str, str]] = {} # preview_id -> {orig, upscaled}
# Custom thread-safe queue for upscaling jobs to support reordering & cancellation
job_queue = CustomJobQueue()
queue_lock = threading.Lock()
current_running_job_id = None
main_loop = None
JOBS_FILE = os.path.join(upscaler.BASE_DIR, "jobs.json")
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": ""}
if os.path.exists(SETTINGS_FILE):
try:
with open(SETTINGS_FILE, "r") as f:
saved = json.load(f)
defaults.update(saved)
except Exception as e:
print(f"Error loading settings: {e}")
return defaults
def save_settings(settings: dict):
"""Save persistent app settings to settings.json"""
try:
with open(SETTINGS_FILE, "w") as f:
json.dump(settings, f, indent=4)
except Exception as e:
print(f"Error saving settings: {e}")
def load_jobs_db():
global jobs_db
if os.path.exists(JOBS_FILE):
try:
with open(JOBS_FILE, "r") as f:
data = json.load(f)
for job_id, job_data in data.items():
job = upscaler.UpscaleJob.from_dict(job_data)
# Automatically put queued items back in the queue
if job.status == "queued":
job_queue.put(job_id)
# Mark active items as interrupted so they can be resumed
elif job.status in ["analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling"]:
job.status = "interrupted"
job.eta = "Interrupted"
jobs_db[job_id] = job
except Exception as e:
print(f"Error loading jobs database: {e}")
def save_jobs_db():
try:
with open(JOBS_FILE, "w") as f:
data = {job_id: job.to_dict() for job_id, job in jobs_db.items()}
json.dump(data, f, indent=4)
except Exception as e:
print(f"Error saving jobs database: {e}")
@app.on_event("startup")
def startup_event():
global main_loop
main_loop = asyncio.get_event_loop()
load_jobs_db()
global_webhook_url = None
class WebhookSetupRequest(BaseModel):
url: str
def send_webhook_notification(url: str, payload: dict):
def post():
try:
req = urllib.request.Request(
url,
data=json.dumps(payload).encode('utf-8'),
headers={'Content-Type': 'application/json'}
)
with urllib.request.urlopen(req, timeout=5) as response:
response.read()
except Exception as e:
print(f"Webhook notification failed to {url}: {e}")
threading.Thread(target=post, daemon=True).start()
# Broadcast updates to websockets and webhooks
def broadcast_progress(job_id: str, data: dict):
save_jobs_db()
job = jobs_db.get(job_id)
if job:
data["is_preview"] = getattr(job, "is_preview", False)
if getattr(job, "is_preview", False):
data["original_preview_file"] = f"original_{job_id}.{getattr(job, 'transcode_format', 'mp4')}"
if job_id in ws_connections and main_loop is not None:
active_sockets = list(ws_connections[job_id])
for websocket in active_sockets:
try:
# Schedule the coroutine in the main event loop
asyncio.run_coroutine_threadsafe(websocket.send_json(data), main_loop)
except Exception as e:
print(f"Error sending websocket message: {e}")
# Process webhooks
job = jobs_db.get(job_id)
webhook_url = None
if job and getattr(job, "webhook_url", None):
webhook_url = job.webhook_url
elif global_webhook_url:
webhook_url = global_webhook_url
if webhook_url:
payload = {
"job_id": job_id,
"status": data.get("status"),
"progress": data.get("progress"),
"current_frame": data.get("current_frame"),
"total_frames": data.get("total_frames"),
"eta": data.get("eta"),
"error": data.get("error"),
"output_file": data.get("output_file")
}
send_webhook_notification(webhook_url, payload)
# Sequential queue processor thread
def queue_worker():
global current_running_job_id
while True:
try:
job_id = job_queue.get()
if job_id is None:
break
job = jobs_db.get(job_id)
if not job or job.status == "cancelled":
job_queue.task_done()
continue
with queue_lock:
current_running_job_id = job_id
print(f"Starting job {job_id}...")
upscaler.run_upscale_pipeline(job, broadcast_progress)
with queue_lock:
current_running_job_id = None
job_queue.task_done()
except Exception as e:
print(f"Error in queue worker: {e}")
# Start the background queue worker thread
worker_thread = threading.Thread(target=queue_worker, daemon=True)
worker_thread.start()
class StartUpscaleRequest(BaseModel):
file_id: 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 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
timestamp_sec: float = 1.0
model: str = "realesr-animevideov3"
scale: int = 4
tile_size: int = 256
gpu_ids: str | None = None
# Endpoints
UPLOAD_METADATA_FILE = os.path.join(upscaler.UPLOAD_DIR, "metadata.json")
def load_upload_metadata():
if os.path.exists(UPLOAD_METADATA_FILE):
try:
with open(UPLOAD_METADATA_FILE, "r") as f:
return json.load(f)
except Exception:
pass
return {}
def save_upload_metadata(metadata):
try:
with open(UPLOAD_METADATA_FILE, "w") as f:
json.dump(metadata, f, indent=4)
except Exception:
pass
@app.post("/api/upload")
async def upload_video(file: UploadFile = File(...)):
"""Upload video to workspace directory and parse metadata"""
import time
file_id = str(uuid.uuid4())
ext = os.path.splitext(file.filename)[1].lower()
if ext not in [".mp4", ".mkv", ".avi", ".mov", ".webm"]:
raise HTTPException(status_code=400, detail="Invalid video format. Supported: .mp4, .mkv, .avi, .mov, .webm")
save_path = os.path.join(upscaler.UPLOAD_DIR, f"{file_id}{ext}")
with open(save_path, "wb") as buffer:
content = await file.read()
buffer.write(content)
info = upscaler.get_video_info(save_path)
if not info:
# cleanup if reading failed
os.remove(save_path)
raise HTTPException(status_code=400, detail="Could not read video file metadata. File may be corrupted.")
# Save upload metadata
metadata = load_upload_metadata()
metadata[file_id] = {
"file_id": file_id,
"original_filename": file.filename,
"ext": ext,
"size_bytes": os.path.getsize(save_path),
"upload_time": time.time(),
"metadata": info
}
save_upload_metadata(metadata)
return {
"file_id": file_id,
"filename": file.filename,
"path": save_path,
"metadata": info
}
@app.get("/api/uploads")
def list_uploads():
"""List all uploaded video source files"""
metadata = load_upload_metadata()
valid_uploads = []
metadata_updated = False
if os.path.exists(upscaler.UPLOAD_DIR):
all_files = os.listdir(upscaler.UPLOAD_DIR)
for filename in all_files:
if filename == "metadata.json":
continue
file_path = os.path.join(upscaler.UPLOAD_DIR, filename)
file_id, ext = os.path.splitext(filename)
if file_id in metadata:
valid_uploads.append(metadata[file_id])
else:
info = upscaler.get_video_info(file_path)
if info:
entry = {
"file_id": file_id,
"original_filename": filename,
"ext": ext,
"size_bytes": os.path.getsize(file_path),
"upload_time": os.path.getmtime(file_path),
"metadata": info
}
metadata[file_id] = entry
valid_uploads.append(entry)
metadata_updated = True
# Clean up missing files from metadata
for fid in list(metadata.keys()):
ext = metadata[fid].get("ext", ".mp4")
expected_file = os.path.join(upscaler.UPLOAD_DIR, f"{fid}{ext}")
if not os.path.exists(expected_file):
del metadata[fid]
metadata_updated = True
if metadata_updated:
save_upload_metadata(metadata)
# Sort by upload time desc
valid_uploads.sort(key=lambda x: x.get("upload_time", 0), reverse=True)
return valid_uploads
@app.delete("/api/uploads/{file_id}")
@app.post("/api/uploads/delete/{file_id}")
def delete_upload(file_id: str):
"""Delete an uploaded source file"""
metadata = load_upload_metadata()
if file_id not in metadata:
exts = [".mp4", ".mkv", ".avi", ".mov", ".webm"]
file_path = None
for ext in exts:
p = os.path.join(upscaler.UPLOAD_DIR, f"{file_id}{ext}")
if os.path.exists(p):
file_path = p
break
if not file_path:
raise HTTPException(status_code=404, detail="Upload file not found.")
os.remove(file_path)
return {"file_id": file_id, "status": "deleted"}
ext = metadata[file_id].get("ext", ".mp4")
file_path = os.path.join(upscaler.UPLOAD_DIR, f"{file_id}{ext}")
if os.path.exists(file_path):
os.remove(file_path)
del metadata[file_id]
save_upload_metadata(metadata)
return {"file_id": file_id, "status": "deleted"}
@app.post("/api/upscale/start")
def start_upscale(req: StartUpscaleRequest):
"""Queue upscaling task"""
# Find uploaded file
file_path = None
exts = [".mp4", ".mkv", ".avi", ".mov", ".webm"]
for ext in exts:
test_path = os.path.join(upscaler.UPLOAD_DIR, f"{req.file_id}{ext}")
if os.path.exists(test_path):
file_path = test_path
break
if not file_path:
raise HTTPException(status_code=404, detail="Uploaded file not found.")
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 or load_settings().get("default_temp_dir") or None
)
jobs_db[job_id] = job
job_queue.put(job_id)
save_jobs_db()
# Broadcast initial queued progress
broadcast_progress(job_id, {
"status": "queued",
"progress": 0.0,
"current_frame": 0,
"total_frames": 0,
"eta": "Calculating..."
})
return {
"job_id": job_id,
"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 or load_settings().get("default_temp_dir") or None
)
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"""
job = jobs_db.get(job_id)
if not job:
raise HTTPException(status_code=404, detail="Job not found.")
return {
"job_id": job.job_id,
"status": job.status,
"progress": job.progress,
"current_frame": job.current_frame,
"total_frames": job.total_frames,
"eta": job.eta,
"error": job.error,
"output_file": os.path.basename(job.output_file) if job.output_file else None,
"is_preview": getattr(job, "is_preview", False),
"original_preview_file": f"original_{job_id}.{getattr(job, 'transcode_format', 'mp4')}" if getattr(job, "is_preview", False) else None
}
@app.post("/api/upscale/cancel/{job_id}")
def cancel_job(job_id: str):
"""Cancel a running or queued job"""
job = jobs_db.get(job_id)
if not job:
raise HTTPException(status_code=404, detail="Job not found.")
# Remove from queue if it was queued
job_queue.remove(job_id)
job.cancel()
# Broadcast cancellation status
broadcast_progress(job_id, {
"status": "cancelled",
"progress": job.progress,
"current_frame": job.current_frame,
"total_frames": job.total_frames,
"eta": "N/A"
})
save_jobs_db()
return {"job_id": job_id, "status": "cancelled"}
@app.post("/api/preview/generate")
async def generate_preview(req: PreviewRequest):
"""Generate high quality upscaled single frame preview"""
file_path = None
exts = [".mp4", ".mkv", ".avi", ".mov", ".webm"]
for ext in exts:
test_path = os.path.join(upscaler.UPLOAD_DIR, f"{req.file_id}{ext}")
if os.path.exists(test_path):
file_path = test_path
break
if not file_path:
raise HTTPException(status_code=404, detail="Video file not found.")
preview_id = str(uuid.uuid4())
preview_temp_dir = os.path.join(upscaler.TEMP_DIR, f"preview_{preview_id}")
os.makedirs(preview_temp_dir, exist_ok=True)
orig_path = os.path.join(preview_temp_dir, "orig.jpg")
upscaled_path = os.path.join(preview_temp_dir, "upscaled.jpg")
# Extract frame
extracted = upscaler.extract_single_frame(file_path, req.timestamp_sec, orig_path)
if not extracted:
raise HTTPException(status_code=500, detail="Failed to extract preview frame from video.")
# Upscale frame
upscaled = upscaler.upscale_image_file(orig_path, upscaled_path, req.model, req.scale, req.tile_size, req.gpu_ids)
if not upscaled:
raise HTTPException(status_code=500, detail="Failed to upscale preview frame.")
preview_db[preview_id] = {
"orig": orig_path,
"upscaled": upscaled_path
}
return {"preview_id": preview_id}
@app.get("/api/preview/original/{preview_id}")
def get_preview_original(preview_id: str):
paths = preview_db.get(preview_id)
if not paths or not os.path.exists(paths["orig"]):
raise HTTPException(status_code=404, detail="Preview frame not found.")
return FileResponse(paths["orig"])
@app.get("/api/preview/upscaled/{preview_id}")
def get_preview_upscaled(preview_id: str):
paths = preview_db.get(preview_id)
if not paths or not os.path.exists(paths["upscaled"]):
raise HTTPException(status_code=404, detail="Preview frame not found.")
return FileResponse(paths["upscaled"])
@app.get("/api/models")
def get_models():
"""Scan models directory and return list of available models"""
models_dir = os.path.join(upscaler.BASE_DIR, "realesrgan-bin", "models")
if not os.path.exists(models_dir):
return []
models = set()
for filename in os.listdir(models_dir):
if filename.endswith(".param"):
name = filename[:-6] # strip '.param'
# Strip scale suffix if present
for suffix in ["-x2", "-x3", "-x4"]:
if name.endswith(suffix):
name = name[:-len(suffix)]
break
models.add(name)
return sorted(list(models))
@app.get("/api/download/{filename}")
def download_file(filename: str):
file_path = os.path.join(upscaler.OUTPUT_DIR, filename)
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)
def safe_delete_file(file_path: str):
if file_path and os.path.exists(file_path):
try:
os.remove(file_path)
print(f"Deleted file: {file_path}")
except Exception as e:
print(f"Failed to delete {file_path}: {e}")
@app.get("/api/jobs")
def list_jobs():
"""List details of all submitted jobs in queue-sorted order"""
active_id = current_running_job_id
queued_ids = job_queue.get_all()
# Sort active first, then queued in order, then history by start time descending
def get_sort_key(job):
if job.job_id == active_id:
return (0, 0)
elif job.job_id in queued_ids:
return (1, queued_ids.index(job.job_id))
else:
t = job.start_time if job.start_time is not None else 0
return (2, -t)
sorted_jobs = sorted(jobs_db.values(), key=get_sort_key)
return [
{
"job_id": job.job_id,
"status": job.status,
"progress": job.progress,
"current_frame": job.current_frame,
"total_frames": job.total_frames,
"eta": job.eta,
"error": job.error,
"model": job.model,
"scale": job.scale,
"output_file": os.path.basename(job.output_file) if job.output_file else None,
"video_path": job.video_path,
"is_preview": getattr(job, "is_preview", False),
"queue_position": queued_ids.index(job.job_id) if job.job_id in queued_ids else -1 if job.job_id == active_id else None
}
for job in sorted_jobs
]
@app.delete("/api/jobs/{job_id}")
@app.post("/api/jobs/delete/{job_id}")
def delete_job(job_id: str):
job = jobs_db.get(job_id)
if not job:
raise HTTPException(status_code=404, detail="Job not found.")
# Remove from queue if it is queued
job_queue.remove(job_id)
# Safely delete original preview video if present
for ext in [".mp4", ".mkv", ".avi", ".mov", ".webm"]:
orig_prev_path = os.path.join(upscaler.OUTPUT_DIR, f"original_{job_id}{ext}")
if os.path.exists(orig_prev_path):
safe_delete_file(orig_prev_path)
# Safely delete upscaled output file
if job.output_file:
safe_delete_file(job.output_file)
# Clean up any residual temp folder
job_temp_dir = os.path.join(upscaler.TEMP_DIR, job_id)
if os.path.exists(job_temp_dir):
try:
import shutil
shutil.rmtree(job_temp_dir)
except Exception:
pass
# Delete from in-memory db
if job_id in jobs_db:
del jobs_db[job_id]
save_jobs_db()
return {"job_id": job_id, "status": "purged"}
@app.post("/api/jobs/purge-all")
def purge_all_jobs():
# Purge uploads folder
for f in os.listdir(upscaler.UPLOAD_DIR):
safe_delete_file(os.path.join(upscaler.UPLOAD_DIR, f))
# Purge outputs folder
for f in os.listdir(upscaler.OUTPUT_DIR):
safe_delete_file(os.path.join(upscaler.OUTPUT_DIR, f))
# Purge temp folder
for f in os.listdir(upscaler.TEMP_DIR):
path = os.path.join(upscaler.TEMP_DIR, f)
try:
if os.path.isdir(path):
import shutil
shutil.rmtree(path)
else:
os.remove(path)
except Exception:
pass
# Reset in-memory database
jobs_db.clear()
# Re-initialize custom queue
global job_queue
job_queue = CustomJobQueue()
# Reset upload metadata file
save_upload_metadata({})
save_jobs_db()
return {"status": "all purged"}
class ReorderQueueRequest(BaseModel):
job_ids: List[str]
@app.post("/api/queue/reorder")
def reorder_queue(req: ReorderQueueRequest):
"""Reorder the job queue"""
job_queue.reorder(req.job_ids)
save_jobs_db()
return {"status": "success", "queue": job_queue.get_all()}
@app.get("/api/queue")
def get_queue():
"""Get the current job queue order"""
return {"queue": job_queue.get_all()}
@app.post("/api/upscale/resume/{job_id}")
def resume_job(job_id: str):
"""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.")
# Re-queue the job
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()
broadcast_progress(job_id, {
"status": "queued",
"progress": job.progress,
"current_frame": job.current_frame,
"total_frames": job.total_frames,
"eta": "Queued for resume..."
})
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):
await websocket.accept()
if job_id not in ws_connections:
ws_connections[job_id] = []
ws_connections[job_id].append(websocket)
# Send current state immediately
job = jobs_db.get(job_id)
if job:
await websocket.send_json({
"status": job.status,
"progress": job.progress,
"current_frame": job.current_frame,
"total_frames": job.total_frames,
"eta": job.eta,
"error": job.error,
"output_file": os.path.basename(job.output_file) if job.output_file else None
})
try:
while True:
# Just keep the connection alive
await websocket.receive_text()
except WebSocketDisconnect:
if job_id in ws_connections:
ws_connections[job_id].remove(websocket)
if not ws_connections[job_id]:
del ws_connections[job_id]
# Clean up previews on exit
@app.on_event("shutdown")
def cleanup_previews():
for preview_id, paths in preview_db.items():
try:
parent_dir = os.path.dirname(paths["orig"])
if os.path.exists(parent_dir):
import shutil
shutil.rmtree(parent_dir)
except Exception:
pass
# System Diagnostics Endpoint
@app.get("/api/diagnostics")
def get_diagnostics():
# RAM Diagnostics
ram = {"total": 0, "available": 0, "used": 0, "percent": 0.0}
try:
if os.path.exists('/proc/meminfo'):
mem_info = {}
with open('/proc/meminfo', 'r') as f:
for line in f:
parts = line.split(':')
if len(parts) == 2:
name = parts[0].strip()
val = parts[1].replace('kB', '').strip()
mem_info[name] = int(val)
total = mem_info.get('MemTotal', 0) * 1024
free = mem_info.get('MemFree', 0) * 1024
available = mem_info.get('MemAvailable', total - free) * 1024
used = total - available
percent = round((used / total) * 100, 1) if total > 0 else 0.0
ram = {
"total": total,
"available": available,
"used": used,
"percent": percent
}
except Exception as e:
print(f"Error getting RAM diagnostics: {e}")
# CPU Diagnostics
cpu = {"count": os.cpu_count(), "load_avg": [], "percent": 0.0}
try:
load_1, load_5, load_15 = os.getloadavg()
cpu["load_avg"] = [load_1, load_5, load_15]
cpu["percent"] = min(100.0, round((load_1 / cpu["count"]) * 100, 1)) if cpu["count"] else 0.0
except Exception as e:
print(f"Error getting CPU diagnostics: {e}")
# GPU Diagnostics
gpu = {"available": False, "gpus": []}
try:
cmd = [
"nvidia-smi",
"--query-gpu=utilization.gpu,utilization.memory,memory.total,memory.free,memory.used,name,temperature.gpu",
"--format=csv,noheader,nounits"
]
res = subprocess.run(cmd, capture_output=True, text=True, check=True)
lines = res.stdout.strip().split('\n')
gpu_list = []
for line in lines:
if not line.strip():
continue
parts = [p.strip() for p in line.split(',')]
if len(parts) >= 7:
gpu_list.append({
"name": parts[5],
"gpu_utilization_percent": float(parts[0]),
"memory_utilization_percent": float(parts[1]),
"memory_total_mb": float(parts[2]),
"memory_free_mb": float(parts[3]),
"memory_used_mb": float(parts[4]),
"temperature_c": float(parts[6])
})
if gpu_list:
gpu["available"] = True
gpu["gpus"] = gpu_list
except Exception as e:
gpu["error"] = str(e)
return {
"ram": ram,
"cpu": cpu,
"gpu": gpu
}
# Outputs Gallery Endpoints
@app.get("/api/outputs")
def list_outputs():
output_dir = upscaler.OUTPUT_DIR
if not os.path.exists(output_dir):
return []
files = []
for filename in os.listdir(output_dir):
if filename.startswith("original_"):
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}"
})
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)):
raise HTTPException(status_code=403, detail="Access denied.")
if not os.path.exists(file_path):
raise HTTPException(status_code=404, detail="File not found.")
try:
os.remove(file_path)
return {"filename": filename, "status": "deleted"}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to delete file: {e}")
# Custom Model Uploading Endpoint
@app.post("/api/models/upload")
async def upload_model(
param_file: UploadFile = File(None),
bin_file: UploadFile = File(None),
files: List[UploadFile] = File(None)
):
models_dir = os.path.join(upscaler.BASE_DIR, "realesrgan-bin", "models")
os.makedirs(models_dir, exist_ok=True)
uploaded_files = []
all_files = []
if files:
all_files.extend(files)
if param_file:
all_files.append(param_file)
if bin_file:
all_files.append(bin_file)
if not all_files:
raise HTTPException(status_code=400, detail="No files uploaded. Please upload .param or .bin files.")
for f in all_files:
filename = f.filename
ext = os.path.splitext(filename)[1].lower()
if ext not in [".param", ".bin"]:
raise HTTPException(
status_code=400,
detail=f"Invalid file extension: {ext}. Only .param and .bin files are allowed for models."
)
safe_filename = os.path.basename(filename)
save_path = os.path.join(models_dir, safe_filename)
with open(save_path, "wb") as buffer:
content = await f.read()
buffer.write(content)
uploaded_files.append(safe_filename)
return {
"status": "success",
"uploaded": uploaded_files
}
# Webhook Setup/Retrieval Endpoints
@app.post("/api/webhook/setup")
def setup_webhook(req: WebhookSetupRequest):
global global_webhook_url
global_webhook_url = req.url
return {"status": "success", "webhook_url": global_webhook_url}
# ─── Settings Endpoints ─────────────────────────────────────────────
class UpdateSettingsRequest(BaseModel):
default_temp_dir: str = ""
@app.get("/api/settings")
def get_settings():
"""Get persistent app settings"""
settings = load_settings()
return settings
@app.post("/api/settings")
def update_settings(req: UpdateSettingsRequest):
"""Update persistent app settings"""
settings = load_settings()
# Validate the directory path if provided
new_dir = req.default_temp_dir.strip()
if new_dir:
# Ensure the parent directory exists or can be created
try:
os.makedirs(new_dir, exist_ok=True)
except OSError as e:
raise HTTPException(status_code=400, detail=f"Invalid directory path: {e}")
settings["default_temp_dir"] = new_dir
save_settings(settings)
return {"status": "saved", "settings": settings}
@app.get("/api/uploads/{filename}")
def get_upload_file(filename: str):
file_path = os.path.join(upscaler.UPLOAD_DIR, filename)
resolved_path = os.path.abspath(file_path)
if not resolved_path.startswith(os.path.abspath(upscaler.UPLOAD_DIR)):
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)
class CleanupTempRequest(BaseModel):
paths: List[str] | None = None
def get_dir_stats(path: str) -> tuple[int, int]:
total_size = 0
total_files = 0
if not os.path.exists(path):
return 0, 0
for root, dirs, files in os.walk(path):
for f in files:
fp = os.path.join(root, f)
try:
if os.path.exists(fp):
total_size += os.path.getsize(fp)
total_files += 1
except Exception:
pass
return total_size, total_files
@app.get("/api/system/temp-info")
def get_temp_info():
temp_folders = []
# 1. Collect all potential temp root directories
# Default temp root
temp_roots = {upscaler.TEMP_DIR}
# Custom temp roots from jobs db
for job in jobs_db.values():
if getattr(job, "temp_dir", None):
temp_roots.add(job.temp_dir)
# 2. Scan folders inside the temp roots
seen_paths = set()
for root_dir in temp_roots:
if not os.path.exists(root_dir):
continue
try:
for item in os.listdir(root_dir):
item_path = os.path.join(root_dir, item)
if not os.path.isdir(item_path):
continue
if item_path in seen_paths:
continue
seen_paths.add(item_path)
# Check if it corresponds to a job_id (which is usually a UUID)
job_id = item
job = jobs_db.get(job_id)
size_bytes, file_count = get_dir_stats(item_path)
if job:
job_status = job.status
video_name = os.path.basename(job.video_path)
else:
job_status = "orphaned"
video_name = "Unknown Video (Orphaned)"
temp_folders.append({
"job_id": job_id,
"path": item_path,
"size_bytes": size_bytes,
"file_count": file_count,
"status": job_status,
"video_name": video_name
})
except Exception as e:
print(f"Error scanning temp root {root_dir}: {e}")
return {"temp_folders": temp_folders}
@app.post("/api/system/temp-cleanup")
def cleanup_temp(req: CleanupTempRequest):
import shutil
cleaned = []
errors = []
# Identify active job IDs
active_statuses = ["analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling"]
active_job_ids = {job.job_id for job in jobs_db.values() if job.status in active_statuses}
# Get current temp folder status
temp_info = get_temp_info()
folders = temp_info["temp_folders"]
target_paths = req.paths if req.paths is not None else [f["path"] for f in folders]
for f in folders:
path = f["path"]
if path not in target_paths:
continue
# Safety checks:
if f["job_id"] in active_job_ids:
errors.append(f"Cannot clean up active job {f['job_id']}")
continue
resolved_path = os.path.realpath(path)
is_sub_of_temp = resolved_path.startswith(os.path.realpath(upscaler.TEMP_DIR))
if not is_sub_of_temp:
for job in jobs_db.values():
if getattr(job, "temp_dir", None):
if resolved_path.startswith(os.path.realpath(job.temp_dir)):
is_sub_of_temp = True
break
if not is_sub_of_temp:
errors.append(f"Safety constraint: Path {path} is not in a valid temp root.")
continue
try:
if os.path.exists(path):
shutil.rmtree(path)
cleaned.append(path)
# Update DB state if needed
job = jobs_db.get(f["job_id"])
if job and job.status in ["paused", "interrupted", "failed"]:
job.status = "interrupted"
job.progress = 0.0
job.current_frame = 0
job.eta = "Temp files cleaned. Will restart on resume."
else:
cleaned.append(path)
except Exception as e:
errors.append(f"Error removing {path}: {str(e)}")
save_jobs_db()
return {"cleaned": cleaned, "errors": errors}
# Mount static folder
app.mount("/", StaticFiles(directory=os.path.join(upscaler.BASE_DIR, "static"), html=True), name="static")