940 lines
31 KiB
Python
940 lines
31 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")
|
|
|
|
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", "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
|
|
|
|
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
|
|
)
|
|
|
|
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.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 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..."
|
|
|
|
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"}
|
|
|
|
# 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}
|
|
|
|
@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)
|
|
|
|
# Mount static folder
|
|
app.mount("/", StaticFiles(directory=os.path.join(upscaler.BASE_DIR, "static"), html=True), name="static")
|