Implement advanced AI enhancements UI tracking, styling, requirements, and documentation modal
This commit is contained in:
+8
-2
@@ -98,7 +98,7 @@ def load_jobs_db():
|
||||
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"]:
|
||||
elif job.status in ["analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling"]:
|
||||
job.status = "interrupted"
|
||||
job.eta = "Interrupted"
|
||||
jobs_db[job_id] = job
|
||||
@@ -232,6 +232,9 @@ class StartUpscaleRequest(BaseModel):
|
||||
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
|
||||
|
||||
class PreviewRequest(BaseModel):
|
||||
file_id: str
|
||||
@@ -416,7 +419,10 @@ def start_upscale(req: StartUpscaleRequest):
|
||||
interpolation=req.interpolation,
|
||||
webhook_url=req.webhook_url,
|
||||
transcode_format=req.transcode_format,
|
||||
is_preview=req.is_preview
|
||||
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
|
||||
)
|
||||
|
||||
jobs_db[job_id] = job
|
||||
|
||||
+102
-8
@@ -23,13 +23,18 @@ class UpscaleJob:
|
||||
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):
|
||||
webhook_url: str = None, transcode_format: str = "mp4", is_preview: bool = False,
|
||||
ai_face_restoration: bool = False, ai_rife_interpolation: bool = False,
|
||||
ai_audio_denoise: bool = False):
|
||||
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
|
||||
self.ai_face_restoration = ai_face_restoration
|
||||
self.ai_rife_interpolation = ai_rife_interpolation
|
||||
self.ai_audio_denoise = ai_audio_denoise
|
||||
|
||||
# Trim mapping
|
||||
if ss is not None:
|
||||
@@ -448,6 +453,87 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
|
||||
processed_files = len(os.listdir(output_frames_dir))
|
||||
job.update_status("upscaling", progress=80.0, current_frame=processed_files)
|
||||
|
||||
# Step 2.6: AI Face Restoration (GFPGAN)
|
||||
if getattr(job, "ai_face_restoration", False):
|
||||
job.update_status("restoring_faces", progress=81.0)
|
||||
on_progress_update(job.job_id, {"status": "restoring_faces", "progress": 81.0})
|
||||
|
||||
import importlib.util
|
||||
gfpgan_installed = importlib.util.find_spec("gfpgan") is not None
|
||||
|
||||
if gfpgan_installed:
|
||||
print(f"Job {job.job_id}: GFPGAN detected. Running Face Restoration...")
|
||||
restored_dir = os.path.join(job_temp_dir, "restored_frames")
|
||||
os.makedirs(restored_dir, exist_ok=True)
|
||||
|
||||
gfpgan_cmd = [
|
||||
sys.executable, "-m", "gfpgan.inference_gfpgan",
|
||||
"-i", output_frames_dir,
|
||||
"-o", restored_dir,
|
||||
"-v", "1.4",
|
||||
"-s", "1",
|
||||
"--bg_upsampler", "None"
|
||||
]
|
||||
|
||||
p_gfp = job.run_command(gfpgan_cmd)
|
||||
stdout, stderr = p_gfp.communicate()
|
||||
job.cleanup_process(p_gfp)
|
||||
|
||||
if p_gfp.returncode == 0:
|
||||
gfp_output_path = os.path.join(restored_dir, "restored_imgs")
|
||||
if os.path.exists(gfp_output_path) and len(os.listdir(gfp_output_path)) > 0:
|
||||
for filename in os.listdir(gfp_output_path):
|
||||
src_f = os.path.join(gfp_output_path, filename)
|
||||
dst_f = os.path.join(output_frames_dir, filename)
|
||||
try:
|
||||
shutil.copy2(src_f, dst_f)
|
||||
except Exception as e:
|
||||
print(f"Error copying restored face frame: {e}")
|
||||
print(f"Job {job.job_id}: Face Restoration completed successfully.")
|
||||
else:
|
||||
print(f"Job {job.job_id}: GFPGAN did not generate outputs in restored_imgs.")
|
||||
else:
|
||||
print(f"Job {job.job_id}: GFPGAN failed (exit code {p_gfp.returncode}). Continuing with normal upscale.")
|
||||
else:
|
||||
print(f"Job {job.job_id}: 'gfpgan' package is not installed in the virtual environment. Skipping face restoration. To enable, run: pip install gfpgan realesrgan")
|
||||
|
||||
# Step 2.7: AI Frame Interpolation (RIFE)
|
||||
rife_frames_dir = os.path.join(job_temp_dir, "rife_frames")
|
||||
use_rife = False
|
||||
|
||||
if getattr(job, "ai_rife_interpolation", False):
|
||||
rife_bin = os.path.join(BASE_DIR, "rife-bin", "rife-ncnn-vulkan")
|
||||
if os.path.isfile(rife_bin):
|
||||
job.update_status("interpolating", progress=83.0)
|
||||
on_progress_update(job.job_id, {"status": "interpolating", "progress": 83.0})
|
||||
os.makedirs(rife_frames_dir, exist_ok=True)
|
||||
|
||||
try:
|
||||
os.chmod(rife_bin, 0o755)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
rife_cmd = [
|
||||
rife_bin,
|
||||
"-i", output_frames_dir,
|
||||
"-o", rife_frames_dir,
|
||||
"-f", "jpg"
|
||||
]
|
||||
if getattr(job, "gpu_ids", None) is not None:
|
||||
rife_cmd.extend(["-g", str(job.gpu_ids)])
|
||||
|
||||
p_rife = job.run_command(rife_cmd)
|
||||
stdout, stderr = p_rife.communicate()
|
||||
job.cleanup_process(p_rife)
|
||||
|
||||
if p_rife.returncode == 0:
|
||||
use_rife = True
|
||||
print(f"Job {job.job_id}: Successfully ran RIFE frame interpolation.")
|
||||
else:
|
||||
print(f"Job {job.job_id}: RIFE failed (exit code {p_rife.returncode}). Falling back to FFmpeg interpolation.")
|
||||
else:
|
||||
print(f"Job {job.job_id}: RIFE binary not found at {rife_bin}. Falling back to FFmpeg interpolation.")
|
||||
|
||||
# Step 3: Reassemble video
|
||||
job.update_status("assembling", progress=85.0)
|
||||
on_progress_update(job.job_id, {"status": "assembling", "progress": 85.0})
|
||||
@@ -463,12 +549,15 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
|
||||
if transcode_fmt == "webm":
|
||||
vcodec = "libvpx-vp9"
|
||||
acodec = "libvorbis"
|
||||
|
||||
assemble_frames_dir = rife_frames_dir if use_rife else output_frames_dir
|
||||
assemble_fps = fps * 2 if (use_rife or getattr(job, "double_fps", False) or getattr(job, "interpolation", False)) else fps
|
||||
|
||||
# Construct ffmpeg reassembly command
|
||||
assemble_cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-framerate", str(fps),
|
||||
"-i", os.path.join(output_frames_dir, "frame_%08d.jpg")
|
||||
"-framerate", str(assemble_fps),
|
||||
"-i", os.path.join(assemble_frames_dir, "frame_%08d.jpg")
|
||||
]
|
||||
|
||||
# We need the original video as the second input (index 1) if we preserve audio or subtitles
|
||||
@@ -485,10 +574,15 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
|
||||
assemble_cmd.extend(["-map", "0:v:0"])
|
||||
|
||||
if job.preserve_audio:
|
||||
assemble_cmd.extend([
|
||||
"-map", "1:a:0?",
|
||||
"-c:a", acodec
|
||||
])
|
||||
assemble_cmd.extend(["-map", "1:a:0?"])
|
||||
if getattr(job, "ai_audio_denoise", False):
|
||||
acodec_denoise = "libvorbis" if transcode_fmt == "webm" else "aac"
|
||||
assemble_cmd.extend([
|
||||
"-af", "arnnoise",
|
||||
"-c:a", acodec_denoise
|
||||
])
|
||||
else:
|
||||
assemble_cmd.extend(["-c:a", acodec])
|
||||
|
||||
if getattr(job, "preserve_subtitles", True):
|
||||
assemble_cmd.extend([
|
||||
@@ -502,7 +596,7 @@ def run_upscale_pipeline(job: UpscaleJob, on_progress_update: Callable[[str, Dic
|
||||
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):
|
||||
if (getattr(job, "double_fps", False) or getattr(job, "interpolation", False) or getattr(job, "ai_rife_interpolation", False)) and not use_rife:
|
||||
target_fps = fps * 2 if getattr(job, "double_fps", False) else 60
|
||||
if target_fps < fps:
|
||||
target_fps = fps
|
||||
|
||||
Reference in New Issue
Block a user