Initial commit of video upscaler
This commit is contained in:
+684
@@ -0,0 +1,684 @@
|
||||
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
|
||||
jobs_db: Dict[str, upscaler.UpscaleJob] = {}
|
||||
ws_connections: Dict[str, List[WebSocket]] = {}
|
||||
preview_db: Dict[str, Dict[str, str]] = {} # preview_id -> {orig, upscaled}
|
||||
|
||||
# FIFO queue for upscaling jobs to prevent GPU memory overload
|
||||
job_queue = queue.Queue()
|
||||
queue_lock = threading.Lock()
|
||||
current_running_job_id = None
|
||||
|
||||
main_loop = None
|
||||
|
||||
@app.on_event("startup")
|
||||
def startup_event():
|
||||
global main_loop
|
||||
main_loop = asyncio.get_event_loop()
|
||||
|
||||
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):
|
||||
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
|
||||
@app.post("/api/upload")
|
||||
async def upload_video(file: UploadFile = File(...)):
|
||||
"""Upload video to workspace directory and parse metadata"""
|
||||
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.")
|
||||
|
||||
return {
|
||||
"file_id": file_id,
|
||||
"filename": file.filename,
|
||||
"path": save_path,
|
||||
"metadata": info
|
||||
}
|
||||
|
||||
@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)
|
||||
|
||||
# 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.")
|
||||
|
||||
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"
|
||||
})
|
||||
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"""
|
||||
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)
|
||||
}
|
||||
for job in jobs_db.values()
|
||||
]
|
||||
|
||||
@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.")
|
||||
|
||||
# Safely delete uploaded input file
|
||||
if job.video_path:
|
||||
safe_delete_file(job.video_path)
|
||||
|
||||
# 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]
|
||||
|
||||
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()
|
||||
|
||||
return {"status": "all purged"}
|
||||
|
||||
# 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")
|
||||
+519
@@ -0,0 +1,519 @@
|
||||
import os
|
||||
import subprocess
|
||||
import json
|
||||
import shutil
|
||||
import time
|
||||
import threading
|
||||
from typing import Dict, Any, Callable
|
||||
|
||||
# Paths config
|
||||
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
BIN_PATH = os.path.join(BASE_DIR, "realesrgan-bin", "realesrgan-ncnn-vulkan")
|
||||
UPLOAD_DIR = os.path.join(BASE_DIR, "uploads")
|
||||
OUTPUT_DIR = os.path.join(BASE_DIR, "outputs")
|
||||
TEMP_DIR = os.path.join(BASE_DIR, "temp")
|
||||
|
||||
# Ensure directories exist
|
||||
for d in [UPLOAD_DIR, OUTPUT_DIR, TEMP_DIR]:
|
||||
os.makedirs(d, exist_ok=True)
|
||||
|
||||
class UpscaleJob:
|
||||
def __init__(self, job_id: str, video_path: str, model: str, scale: int, tile_size: int, preserve_audio: bool,
|
||||
ss: str = None, t: str = None, gpu_ids: str = None, tta: bool = False,
|
||||
unsharp: bool = False, double_fps: bool = False, preserve_subtitles: bool = True,
|
||||
start_sec: float = None, end_sec: float = None, crf: int = 18, preset: str = "medium",
|
||||
denoise: bool = False, sharpen: bool = False, interpolation: bool = False,
|
||||
webhook_url: str = None, transcode_format: str = "mp4", is_preview: bool = False):
|
||||
self.job_id = job_id
|
||||
self.video_path = video_path
|
||||
self.model = model
|
||||
self.scale = scale
|
||||
self.tile_size = tile_size
|
||||
self.preserve_audio = preserve_audio
|
||||
|
||||
# Trim mapping
|
||||
if ss is not None:
|
||||
self.ss = ss
|
||||
elif start_sec is not None:
|
||||
self.ss = str(start_sec)
|
||||
else:
|
||||
self.ss = None
|
||||
|
||||
if t is not None:
|
||||
self.t = t
|
||||
elif end_sec is not None:
|
||||
start = float(self.ss) if self.ss is not None else 0.0
|
||||
self.t = str(max(0.0, end_sec - start))
|
||||
else:
|
||||
self.t = None
|
||||
|
||||
self.gpu_ids = gpu_ids
|
||||
self.tta = tta
|
||||
self.unsharp = unsharp
|
||||
self.double_fps = double_fps
|
||||
self.preserve_subtitles = preserve_subtitles
|
||||
self.start_sec = start_sec
|
||||
self.end_sec = end_sec
|
||||
self.crf = crf
|
||||
self.preset = preset
|
||||
self.denoise = denoise
|
||||
self.sharpen = sharpen
|
||||
self.interpolation = interpolation
|
||||
self.webhook_url = webhook_url
|
||||
self.transcode_format = transcode_format
|
||||
self.is_preview = is_preview
|
||||
self.original_preview_file = None
|
||||
|
||||
self.status = "pending"
|
||||
self.progress = 0.0
|
||||
self.current_frame = 0
|
||||
self.total_frames = 0
|
||||
self.eta = "Calculating..."
|
||||
self.error = None
|
||||
self.start_time = None
|
||||
self.output_file = None
|
||||
|
||||
# Track processes to allow cancellation
|
||||
self._processes = []
|
||||
self._is_cancelled = False
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def update_status(self, status: str, progress: float = None, current_frame: int = None, eta: str = None, error: str = None):
|
||||
with self._lock:
|
||||
self.status = status
|
||||
if progress is not None:
|
||||
self.progress = round(progress, 2)
|
||||
if current_frame is not None:
|
||||
self.current_frame = current_frame
|
||||
if eta is not None:
|
||||
self.eta = eta
|
||||
if error is not None:
|
||||
self.error = error
|
||||
|
||||
def cancel(self):
|
||||
with self._lock:
|
||||
self._is_cancelled = True
|
||||
self.status = "cancelled"
|
||||
self.eta = "N/A"
|
||||
for p in self._processes:
|
||||
try:
|
||||
p.terminate()
|
||||
p.wait(timeout=2)
|
||||
except Exception:
|
||||
try:
|
||||
p.kill()
|
||||
except Exception:
|
||||
pass
|
||||
self._processes.clear()
|
||||
|
||||
def run_command(self, cmd: list, shell=False) -> subprocess.Popen:
|
||||
with self._lock:
|
||||
if self._is_cancelled:
|
||||
raise InterruptedError("Job was cancelled")
|
||||
|
||||
p = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
shell=shell
|
||||
)
|
||||
self._processes.append(p)
|
||||
return p
|
||||
|
||||
def cleanup_process(self, p: subprocess.Popen):
|
||||
with self._lock:
|
||||
if p in self._processes:
|
||||
self._processes.remove(p)
|
||||
|
||||
def get_video_info(video_path: str) -> dict:
|
||||
"""Extract metadata using ffprobe"""
|
||||
cmd = [
|
||||
"ffprobe",
|
||||
"-v", "error",
|
||||
"-select_streams", "v:0",
|
||||
"-show_entries", "stream=avg_frame_rate,duration,width,height,codec_name,nb_frames",
|
||||
"-of", "json",
|
||||
video_path
|
||||
]
|
||||
try:
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, check=True)
|
||||
data = json.loads(result.stdout)
|
||||
if not data.get("streams"):
|
||||
return {}
|
||||
|
||||
stream = data["streams"][0]
|
||||
|
||||
# Calculate FPS
|
||||
fps_str = stream.get("avg_frame_rate", "30/1")
|
||||
if "/" in fps_str:
|
||||
num, den = map(float, fps_str.split("/"))
|
||||
fps = num / den if den != 0 else 30.0
|
||||
else:
|
||||
fps = float(fps_str) if fps_str else 30.0
|
||||
|
||||
# Get total frames
|
||||
nb_frames = stream.get("nb_frames")
|
||||
if nb_frames and nb_frames.isdigit():
|
||||
total_frames = int(nb_frames)
|
||||
else:
|
||||
duration = float(stream.get("duration", 0))
|
||||
total_frames = int(duration * fps)
|
||||
|
||||
return {
|
||||
"width": int(stream.get("width", 0)),
|
||||
"height": int(stream.get("height", 0)),
|
||||
"fps": round(fps, 3),
|
||||
"duration": round(float(stream.get("duration", 0)), 2),
|
||||
"codec": stream.get("codec_name", "unknown"),
|
||||
"total_frames": total_frames
|
||||
}
|
||||
except Exception as e:
|
||||
print(f"Error reading video info: {e}")
|
||||
return {}
|
||||
|
||||
def extract_single_frame(video_path: str, timestamp_sec: float, output_path: str) -> bool:
|
||||
"""Extract a single frame at timestamp for preview"""
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-ss", str(timestamp_sec),
|
||||
"-i", video_path,
|
||||
"-vframes", "1",
|
||||
"-f", "image2",
|
||||
output_path
|
||||
]
|
||||
try:
|
||||
subprocess.run(cmd, capture_output=True, check=True)
|
||||
return os.path.exists(output_path)
|
||||
except Exception as e:
|
||||
print(f"Error extracting single frame: {e}")
|
||||
return False
|
||||
|
||||
def upscale_image_file(input_path: str, output_path: str, model: str, scale: int, tile_size: int, gpu_ids: str = None) -> bool:
|
||||
"""Run Real-ESRGAN on a single image file"""
|
||||
cmd = [
|
||||
BIN_PATH,
|
||||
"-i", input_path,
|
||||
"-o", output_path,
|
||||
"-n", model,
|
||||
"-s", str(scale),
|
||||
"-t", str(tile_size)
|
||||
]
|
||||
if gpu_ids is not None:
|
||||
cmd.extend(["-g", str(gpu_ids)])
|
||||
try:
|
||||
subprocess.run(cmd, capture_output=True, check=True)
|
||||
return os.path.exists(output_path)
|
||||
except Exception as e:
|
||||
print(f"Error upscaling single image: {e}")
|
||||
return False
|
||||
|
||||
def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dict[str, Any]], None]):
|
||||
job.start_time = time.time()
|
||||
job.update_status("analyzing", progress=5)
|
||||
|
||||
# Get video info
|
||||
info = get_video_info(job.video_path)
|
||||
if not info:
|
||||
job.update_status("failed", error="Could not read video metadata.")
|
||||
on_progress_update(job.job_id, {"status": "failed", "error": "Could not read video metadata"})
|
||||
return
|
||||
|
||||
job.total_frames = info["total_frames"]
|
||||
fps = info["fps"]
|
||||
|
||||
# Create job temp directories
|
||||
job_temp_dir = os.path.join(TEMP_DIR, job.job_id)
|
||||
input_frames_dir = os.path.join(job_temp_dir, "input_frames")
|
||||
output_frames_dir = os.path.join(job_temp_dir, "output_frames")
|
||||
|
||||
os.makedirs(input_frames_dir, exist_ok=True)
|
||||
os.makedirs(output_frames_dir, exist_ok=True)
|
||||
|
||||
try:
|
||||
# If preview, extract the original 5s clip first
|
||||
transcode_fmt = getattr(job, "transcode_format", "mp4")
|
||||
if getattr(job, "is_preview", False):
|
||||
orig_preview_filename = f"original_{job.job_id}.{transcode_fmt}"
|
||||
orig_preview_filepath = os.path.join(OUTPUT_DIR, orig_preview_filename)
|
||||
job.original_preview_file = orig_preview_filepath
|
||||
|
||||
cut_cmd = ["ffmpeg", "-y"]
|
||||
if job.ss is not None:
|
||||
cut_cmd.extend(["-ss", str(job.ss)])
|
||||
if job.t is not None:
|
||||
cut_cmd.extend(["-t", str(job.t)])
|
||||
cut_cmd.extend([
|
||||
"-i", job.video_path,
|
||||
"-map", "0:v:0",
|
||||
"-map", "0:a:0?",
|
||||
"-c:v", "libx264",
|
||||
"-c:a", "aac",
|
||||
orig_preview_filepath
|
||||
])
|
||||
try:
|
||||
subprocess.run(cut_cmd, capture_output=True, check=True)
|
||||
except Exception as cut_err:
|
||||
print(f"Error cutting original preview video: {cut_err}")
|
||||
|
||||
# Step 1: Extract Frames
|
||||
job.update_status("extracting", progress=10)
|
||||
on_progress_update(job.job_id, {"status": "extracting", "progress": 10})
|
||||
|
||||
# High quality JPG frames to balance disk usage and speed
|
||||
extract_cmd = ["ffmpeg", "-y"]
|
||||
if job.ss is not None:
|
||||
extract_cmd.extend(["-ss", str(job.ss)])
|
||||
if job.t is not None:
|
||||
extract_cmd.extend(["-t", str(job.t)])
|
||||
extract_cmd.extend(["-i", job.video_path])
|
||||
|
||||
# Apply unsharp pre-filter if enabled
|
||||
if getattr(job, "unsharp", False):
|
||||
extract_cmd.extend(["-vf", "unsharp"])
|
||||
|
||||
extract_cmd.extend([
|
||||
"-q:v", "2",
|
||||
os.path.join(input_frames_dir, "frame_%08d.jpg")
|
||||
])
|
||||
|
||||
p_extract = job.run_command(extract_cmd)
|
||||
stdout, stderr = p_extract.communicate()
|
||||
job.cleanup_process(p_extract)
|
||||
|
||||
if p_extract.returncode != 0:
|
||||
raise RuntimeError(f"FFmpeg frame extraction failed: {stderr}")
|
||||
|
||||
# Count actual frames extracted
|
||||
extracted_files = sorted([f for f in os.listdir(input_frames_dir) if f.startswith("frame_")])
|
||||
actual_total = len(extracted_files)
|
||||
if actual_total == 0:
|
||||
raise RuntimeError("No frames extracted from video")
|
||||
|
||||
job.total_frames = actual_total
|
||||
|
||||
# Step 2: Upscale Frames
|
||||
job.update_status("upscaling", progress=20, current_frame=0)
|
||||
on_progress_update(job.job_id, {"status": "upscaling", "progress": 20, "current_frame": 0, "total_frames": actual_total})
|
||||
|
||||
current_tile_size = job.tile_size
|
||||
while True:
|
||||
# Launch Real-ESRGAN on directory
|
||||
upscale_cmd = [
|
||||
BIN_PATH,
|
||||
"-i", input_frames_dir,
|
||||
"-o", output_frames_dir,
|
||||
"-n", job.model,
|
||||
"-s", str(job.scale),
|
||||
"-t", str(current_tile_size),
|
||||
"-f", "jpg"
|
||||
]
|
||||
if getattr(job, "gpu_ids", None) is not None:
|
||||
upscale_cmd.extend(["-g", str(job.gpu_ids)])
|
||||
if getattr(job, "tta", False):
|
||||
upscale_cmd.append("-x")
|
||||
|
||||
upscale_start_time = time.time()
|
||||
p_upscale = job.run_command(upscale_cmd)
|
||||
|
||||
# Monitor thread for output files
|
||||
while p_upscale.poll() is None:
|
||||
if job._is_cancelled:
|
||||
return
|
||||
|
||||
processed_files = len(os.listdir(output_frames_dir))
|
||||
progress_pct = 20.0 + (float(processed_files) / actual_total) * 60.0 # upscaling is 20% to 80%
|
||||
|
||||
# Estimate ETA
|
||||
elapsed = time.time() - upscale_start_time
|
||||
if processed_files > 0:
|
||||
sec_per_frame = elapsed / processed_files
|
||||
rem_frames = actual_total - processed_files
|
||||
eta_sec = rem_frames * sec_per_frame
|
||||
|
||||
# Format ETA
|
||||
if eta_sec > 60:
|
||||
eta_str = f"{int(eta_sec // 60)}m {int(eta_sec % 60)}s"
|
||||
else:
|
||||
eta_str = f"{int(eta_sec)}s"
|
||||
else:
|
||||
eta_str = "Calculating..."
|
||||
|
||||
job.update_status("upscaling", progress=progress_pct, current_frame=processed_files, eta=eta_str)
|
||||
on_progress_update(job.job_id, {
|
||||
"status": "upscaling",
|
||||
"progress": progress_pct,
|
||||
"current_frame": processed_files,
|
||||
"total_frames": actual_total,
|
||||
"eta": eta_str
|
||||
})
|
||||
time.sleep(0.5)
|
||||
|
||||
stdout, stderr = p_upscale.communicate()
|
||||
job.cleanup_process(p_upscale)
|
||||
|
||||
if job._is_cancelled:
|
||||
return
|
||||
|
||||
if p_upscale.returncode != 0:
|
||||
err_msg = (stdout or "") + "\n" + (stderr or "")
|
||||
is_alloc_error = any(x in err_msg.lower() for x in ["vkallocatememory", "out of memory", "allocation", "vram", "failed to allocate"])
|
||||
|
||||
if is_alloc_error:
|
||||
if current_tile_size <= 0:
|
||||
next_tile_size = 256
|
||||
else:
|
||||
next_tile_size = current_tile_size // 2
|
||||
|
||||
if next_tile_size >= 32:
|
||||
print(f"Job {job.job_id}: Real-ESRGAN failed with VRAM allocation error. Retrying with tile size halved from {current_tile_size} to {next_tile_size}.")
|
||||
current_tile_size = next_tile_size
|
||||
# Clean up output frames directory before retrying
|
||||
for filename in os.listdir(output_frames_dir):
|
||||
file_path = os.path.join(output_frames_dir, filename)
|
||||
try:
|
||||
if os.path.isfile(file_path) or os.path.islink(file_path):
|
||||
os.unlink(file_path)
|
||||
elif os.path.isdir(file_path):
|
||||
shutil.rmtree(file_path)
|
||||
except Exception as cleanup_err:
|
||||
print(f"Error cleaning file {file_path}: {cleanup_err}")
|
||||
continue
|
||||
|
||||
raise RuntimeError(f"Real-ESRGAN failed with exit code {p_upscale.returncode}: {err_msg}")
|
||||
else:
|
||||
break
|
||||
|
||||
# Final validation of upscale output
|
||||
processed_files = len(os.listdir(output_frames_dir))
|
||||
job.update_status("upscaling", progress=80.0, current_frame=processed_files)
|
||||
|
||||
# Step 3: Reassemble video
|
||||
job.update_status("assembling", progress=85.0)
|
||||
on_progress_update(job.job_id, {"status": "assembling", "progress": 85.0})
|
||||
|
||||
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)
|
||||
job.output_file = out_filepath
|
||||
|
||||
# Choose codecs based on format
|
||||
vcodec = "libx264"
|
||||
acodec = "copy"
|
||||
if transcode_fmt == "webm":
|
||||
vcodec = "libvpx-vp9"
|
||||
acodec = "libvorbis"
|
||||
|
||||
# Construct ffmpeg reassembly command
|
||||
assemble_cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-framerate", str(fps),
|
||||
"-i", os.path.join(output_frames_dir, "frame_%08d.jpg")
|
||||
]
|
||||
|
||||
# We need the original video as the second input (index 1) if we preserve audio or subtitles
|
||||
need_orig_input = job.preserve_audio or getattr(job, "preserve_subtitles", True)
|
||||
if need_orig_input:
|
||||
orig_input_cmd = []
|
||||
if getattr(job, "ss", None) is not None:
|
||||
orig_input_cmd.extend(["-ss", str(job.ss)])
|
||||
if getattr(job, "t", None) is not None:
|
||||
orig_input_cmd.extend(["-t", str(job.t)])
|
||||
orig_input_cmd.extend(["-i", job.video_path])
|
||||
assemble_cmd.extend(orig_input_cmd)
|
||||
|
||||
assemble_cmd.extend(["-map", "0:v:0"])
|
||||
|
||||
if job.preserve_audio:
|
||||
assemble_cmd.extend([
|
||||
"-map", "1:a:0?",
|
||||
"-c:a", acodec
|
||||
])
|
||||
|
||||
if getattr(job, "preserve_subtitles", True):
|
||||
assemble_cmd.extend([
|
||||
"-map", "1:s?",
|
||||
"-c:s", "mov_text"
|
||||
])
|
||||
|
||||
# Apply filters
|
||||
vf_filters = []
|
||||
if getattr(job, "denoise", False):
|
||||
vf_filters.append("hqdn3d")
|
||||
if getattr(job, "sharpen", False) or getattr(job, "unsharp", False):
|
||||
vf_filters.append("unsharp=3:3:0.5:3:3:0.5")
|
||||
if getattr(job, "interpolation", False) or getattr(job, "double_fps", False):
|
||||
target_fps = fps * 2 if getattr(job, "double_fps", False) else 60
|
||||
if target_fps < fps:
|
||||
target_fps = fps
|
||||
vf_filters.append(f"minterpolate=fps={target_fps}")
|
||||
|
||||
if vf_filters:
|
||||
assemble_cmd.extend(["-vf", ",".join(vf_filters)])
|
||||
|
||||
assemble_cmd.extend([
|
||||
"-c:v", vcodec,
|
||||
"-pix_fmt", "yuv420p",
|
||||
"-crf", str(getattr(job, "crf", 18)),
|
||||
"-preset", getattr(job, "preset", "medium"),
|
||||
out_filepath
|
||||
])
|
||||
|
||||
p_assemble = job.run_command(assemble_cmd)
|
||||
stdout, stderr = p_assemble.communicate()
|
||||
job.cleanup_process(p_assemble)
|
||||
|
||||
if p_assemble.returncode != 0:
|
||||
raise RuntimeError(f"FFmpeg video assembly failed: {stderr}")
|
||||
|
||||
# Step 4: Complete
|
||||
job.update_status("completed", progress=100.0, eta="Done")
|
||||
on_progress_update(job.job_id, {
|
||||
"status": "completed",
|
||||
"progress": 100.0,
|
||||
"eta": "Done",
|
||||
"output_file": out_filename
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
if not job._is_cancelled:
|
||||
job.update_status("failed", error=str(e))
|
||||
on_progress_update(job.job_id, {"status": "failed", "error": str(e)})
|
||||
finally:
|
||||
# Clean up temp frames to save space
|
||||
try:
|
||||
if os.path.exists(job_temp_dir):
|
||||
shutil.rmtree(job_temp_dir)
|
||||
except Exception as cleanup_err:
|
||||
print(f"Error during temp cleanup: {cleanup_err}")
|
||||
|
||||
# Trigger webhook if URL is provided
|
||||
if getattr(job, "webhook_url", None):
|
||||
def trigger_webhook_task():
|
||||
import urllib.request
|
||||
import json
|
||||
try:
|
||||
payload = {
|
||||
"job_id": job.job_id,
|
||||
"status": job.status,
|
||||
"progress": job.progress,
|
||||
"error": job.error,
|
||||
"output_file": os.path.basename(job.output_file) if job.output_file else None,
|
||||
"model": job.model,
|
||||
"scale": job.scale,
|
||||
"duration": round(time.time() - job.start_time, 2) if job.start_time else 0
|
||||
}
|
||||
req = urllib.request.Request(
|
||||
job.webhook_url,
|
||||
data=json.dumps(payload).encode('utf-8'),
|
||||
headers={'Content-Type': 'application/json'},
|
||||
method='POST'
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=5) as response:
|
||||
pass
|
||||
except Exception as ex:
|
||||
print(f"Error triggering webhook: {ex}")
|
||||
|
||||
threading.Thread(target=trigger_webhook_task, daemon=True).start()
|
||||
Reference in New Issue
Block a user