1

app.py

@@ -17,20 +17,55 @@ import gradio as gr
 import spaces
 from PIL import Image
 import cv2
+import torch  # used for cuda sync & empty_cache
 
 # ----------------- BASIC INFO -----------------
 CHECKPOINT_URL = "https://github.com/yyang181/colormnet/releases/download/v0.1/DINOv2FeatureV6_LocalAtten_s2_154000.pth"
 CHECKPOINT_LOCAL = "DINOv2FeatureV6_LocalAtten_s2_154000.pth"
 
-TITLE = "ColorMNet — ZeroGPU (CUDA-only) Video Colorization with Reference Image"
+TITLE = "ColorMNet — 视频着色 / Video Colorization (ZeroGPU, CUDA-only)"
 DESC = """
-上传**黑白视频**与**参考图像**，点击“开始着色”。  
+**中文**  
+上传**黑白视频**与**参考图像**，点击「开始着色 / Start Coloring」。  
 此版本在 **app.py 中调度 ZeroGPU**，并**在同一进程**调用 `test.py` 的入口函数。  
-临时工作目录结构：
-- 抽帧：`_colormnet_tmp/input_video/<视频名>/00000.png ...`
-- 参考：`_colormnet_tmp/ref/<视频名>/ref.png`
-- 输出：`_colormnet_tmp/output/<视频名>/*.png`
+临时工作目录结构：  
+- 抽帧：`_colormnet_tmp/input_video/<视频名>/00000.png ...`  
+- 参考：`_colormnet_tmp/ref/<视频名>/ref.png`  
+- 输出：`_colormnet_tmp/output/<视频名>/*.png`  
 - 合成视频：`_colormnet_tmp/<视频名>.mp4`
+
+**English**  
+Upload a **B&W video** and a **reference image**, then click “Start Coloring”.  
+This app runs **ZeroGPU scheduling in `app.py`** and calls `test.py` **in-process**.  
+Temp workspace layout:  
+- Frames: `_colormnet_tmp/input_video/<stem>/00000.png ...`  
+- Reference: `_colormnet_tmp/ref/<stem>/ref.png`  
+- Output frames: `_colormnet_tmp/output/<stem>/*.png`  
+- Final video: `_colormnet_tmp/<stem>.mp4`
+"""
+
+PAPER = """
+### 论文 / Paper
+**ECCV 2024 — ColorMNet: A Memory-based Deep Spatial-Temporal Feature Propagation Network for Video Colorization**  
+
+如果你喜欢这个项目，欢迎到 GitHub 点个 ⭐ Star：  
+**GitHub**: https://github.com/yyang181/colormnet
+"""
+
+BADGES_HTML = """
+<div style="display:flex;gap:12px;align-items:center;flex-wrap:wrap;">
+  <a href="https://github.com/yyang181/colormnet" target="_blank" title="Open GitHub Repo">
+    <img alt="GitHub Repo"
+         src="https://img.shields.io/badge/GitHub-colormnet-181717?logo=github" />
+  </a>
+  <a href="https://github.com/yyang181/colormnet/stargazers" target="_blank" title="Star on GitHub">
+    <img alt="GitHub Repo stars"
+         src="https://img.shields.io/github/stars/yyang181/colormnet?style=social" />
+  </a>
+  <span style="opacity:0.85">
+    喜欢这个项目就点个 ⭐ Star / If you like it, please ⭐ Star!
+  </span>
+</div>
 """
 
 # ----------------- TEMP WORKDIR -----------------
@@ -83,11 +118,11 @@ def video_to_frames_dir(video_path: str, frames_dir: str):
         out_path = path.join(frames_dir, f"{idx:05d}.png")
         ok = cv2.imwrite(out_path, frame)
         if not ok:
-            raise RuntimeError(f"写入抽帧失败: {out_path}")
+            raise RuntimeError(f"写入抽帧失败 / Failed to write: {out_path}")
         idx += 1
     cap.release()
     if idx == 0:
-        raise RuntimeError("Input video has no readable frames.")
+        raise RuntimeError("视频无可读帧 / Input video has no readable frames.")
     return w, h, fps, idx
 
 def encode_frames_to_video(frames_dir: str, out_path: str, fps: float):
@@ -166,9 +201,9 @@ def gradio_infer(
 ):
     # 1) 基本校验与临时目录
     if bw_video is None:
-        return None, "请上传黑白视频。"
+        return None, "请上传黑白视频 / Please upload a B&W video."
     if ref_image is None:
-        return None, "请上传参考图像。"
+        return None, "请上传参考图像 / Please upload a reference image."
     reset_temp_root()
 
     # 2) 解析视频源路径 & 目标 <video_stem>
@@ -177,7 +212,7 @@ def gradio_infer(
     elif isinstance(bw_video, str):
         src_video_path = bw_video
     else:
-        return None, "无法读取视频输入。"
+        return None, "无法读取视频输入 / Failed to read video input."
 
     video_stem = path.splitext(path.basename(src_video_path))[0]
 
@@ -195,7 +230,7 @@ def gradio_infer(
     try:
         _w, _h, fps, _n = video_to_frames_dir(src_video_path, input_frames_dir)
     except Exception as e:
-        return None, f"抽帧失败：\n{e}"
+        return None, f"抽帧失败 / Frame extraction failed:\n{e}"
 
     # 5) 参考帧 -> ref/<stem>/ref.png
     ref_png_path = path.join(ref_dir, "ref.png")
@@ -203,14 +238,14 @@ def gradio_infer(
         try:
             ref_image.save(ref_png_path)
         except Exception as e:
-            return None, f"保存参考图像失败：\n{e}"
+            return None, f"保存参考图像失败 / Failed to save reference image:\n{e}"
     elif isinstance(ref_image, str):
         try:
             shutil.copy2(ref_image, ref_png_path)
         except Exception as e:
-            return None, f"复制参考图像失败：\n{e}"
+            return None, f"复制参考图像失败 / Failed to copy reference image:\n{e}"
     else:
-        return None, "无法读取参考图像输入。"
+        return None, "无法读取参考图像输入 / Failed to read reference image."
 
     # 6) 收集 UI 配置
     default_config = {
@@ -259,7 +294,7 @@ def gradio_infer(
     try:
         import test  # 确保 test.py 同目录且有 run_cli 函数
     except Exception as e:
-        return None, f"导入 test.py 失败：\n{e}"
+        return None, f"导入 test.py 失败 / Failed to import test.py：\n{e}"
 
     args_list = build_args_list_for_test(
         d16_batch_path=input_root,   # 指向 input_video 根
@@ -280,44 +315,52 @@ def gradio_infer(
         log = f"Args: {' '.join(args_list)}\n\n{buf.getvalue()}\n\nERROR: {e}"
         return None, log
 
-    # 在合成 mp4 之前加：
-    torch.cuda.synchronize()
+    # 在合成 mp4 之前：清空 CUDA
+    try:
+        torch.cuda.synchronize()
+    except Exception:
+        pass
     try:
         del network, processor, loader, vid_reader, data, rgb, msk, prob
     except Exception:
         pass
-    torch.cuda.empty_cache()
+    try:
+        torch.cuda.empty_cache()
+    except Exception:
+        pass
 
     # 9) 合成 mp4：从 output/<stem>/ 帧合成 -> TEMP_ROOT/<stem>.mp4
     out_frames = path.join(output_root, video_stem)
     if not path.isdir(out_frames):
-        return None, f"未找到输出帧目录：{out_frames}\n\n{log}"
+        return None, f"未找到输出帧目录 / Output frame dir not found：{out_frames}\n\n{log}"
     final_mp4 = path.join(TEMP_ROOT, f"{video_stem}.mp4")
     try:
         encode_frames_to_video(out_frames, final_mp4, fps=fps)
     except Exception as e:
-        return None, f"合成视频失败：\n{e}\n\n{log}"
+        return None, f"合成视频失败 / Video mux failed：\n{e}\n\n{log}"
 
-    return final_mp4, f"完成 ✅\n\n{log}"
+    return final_mp4, f"完成 ✅ / Done ✅\n\n{log}"
 
 # ----------------- UI -----------------
 with gr.Blocks() as demo:
     gr.Markdown(f"# {TITLE}")
+    gr.HTML(BADGES_HTML)     # 页头徽章
+    gr.Markdown(PAPER)
     gr.Markdown(DESC)
 
-    debug_shapes = gr.Checkbox(label="调试日志（仅用于显示更完整日志）", value=False)
+    debug_shapes = gr.Checkbox(label="调试日志 / Debug Logs（仅用于显示更完整日志 / show verbose logs）", value=False)
 
     with gr.Row():
-        inp_video = gr.Video(label="黑白视频（mp4/webm/avi）", interactive=True)
-        inp_ref = gr.Image(label="参考图像（RGB）", type="pil")
+        inp_video = gr.Video(label="黑白视频（mp4/webm/avi） / B&W Video", interactive=True)
+        inp_ref = gr.Image(label="参考图像（RGB） / Reference Image (RGB)", type="pil")
         gr.Examples(
-            label="示例输入",
+            label="示例 / Examples",
             examples=[["./example/4.mp4", "./example/4.png"]],
             inputs=[inp_video, inp_ref],
             cache_examples=False,
         )
 
-    with gr.Accordion("高级参数设置（传给 test.py）", open=False):
+    with gr.Accordion("高级参数设置 / Advanced Settings（传给 test.py / passed to test.py）", open=False):
         with gr.Row():
             first_not_exemplar = gr.Checkbox(label="FirstFrameIsNotExemplar (--FirstFrameIsNotExemplar)", value=True)
             reverse = gr.Checkbox(label="reverse (--reverse)", value=False)
@@ -339,10 +382,10 @@ with gr.Blocks() as demo:
             flip = gr.Checkbox(label="flip (--flip)", value=False)
             size = gr.Number(label="size (--size)", value=-1, precision=0)
 
-    run_btn = gr.Button("开始着色（同进程调用 test.py）")
+    run_btn = gr.Button("开始着色 / Start Coloring（同进程调用 test.py / in-process）")
     with gr.Row():
-        out_video = gr.Video(label="输出视频（着色结果）")
-        status = gr.Textbox(label="状态 / 日志输出（test.py stdout/stderr）", interactive=False, lines=16)
+        out_video = gr.Video(label="输出视频（着色结果） / Output (Colorized)")
+        status = gr.Textbox(label="状态 / 日志输出（test.py stdout/stderr） / Status & Logs", interactive=False, lines=16)
 
     run_btn.click(
         fn=gradio_infer,
@@ -357,6 +400,10 @@ with gr.Blocks() as demo:
         outputs=[out_video, status]
     )
 
+    # 页脚徽章
+    gr.HTML("<hr/>")
+    gr.HTML(BADGES_HTML)
+
 if __name__ == "__main__":
     try:
         ensure_checkpoint()