Update app.py

app.py

@@ -1,71 +1,152 @@
-import gradio as gr
+import sys
 import os
 import tempfile
 from pathlib import Path
-# 根据原项目实际结构调整导入路径（示例假设核心类为 VideoWatermarkRemover）
-from kling.core import VideoWatermarkRemover  # 需替换为实际模块路径
+import subprocess
+import gradio as gr
 
-def process_video(input_video):
+# -------------------------- 关键修复：添加项目根目录到 Python 路径 --------------------------
+# 获取当前文件（app.py）的父目录（即 Hugging Face Space 的工作目录 /app）
+project_root = Path(__file__).parent.resolve()
+sys.path.append(str(project_root))  # 确保能导入 kling 模块
+
+# -------------------------- 导入项目核心模块（修复子模块问题） --------------------------
+try:
+    from kling.core import VideoWatermarkRemover  # 替换为实际模块路径
+except ImportError as e:
+    raise RuntimeError(
+        f"导入 kling 模块失败：{str(e)}\n"
+        "请检查：\n"
+        "1. 是否通过 `git submodule update --init --recursive` 初始化子模块？\n"
+        "2. 项目根目录是否正确添加到 sys.path？"
+    )
+
+# -------------------------- 视频处理函数（含储存检视逻辑） --------------------------
+def process_video(input_video, progress=gr.Progress()):
     """
-    处理输入视频（去水印）的核心函数
+    处理输入视频（去水印）并生成可下载的输出视频
     input_video: Gradio 上传的视频文件（路径为临时文件路径）
+    返回：(输出视频路径, 下载链接/None)
     """
     try:
-        # 临时保存输入视频（Hugging Face Space 会自动清理临时文件）
+        # -------------------- 步骤 1：保存输入视频到临时文件 --------------------
         with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as f:
-            f.write(input_video.read())
+            f.write(input_video.read())  # 写入输入视频数据
             input_path = f.name
+        progress(0.1, desc="✅ 输入视频已保存到临时文件")
 
-        # 输出路径（临时文件，Gradio 会自动展示）
-        output_path = os.path.join(tempfile.gettempdir(), "output.mp4")
+        # -------------------- 步骤 2：检查 FFmpeg 是否可用 --------------------
+        try:
+            subprocess.run(
+                ["ffmpeg", "-version"],
+                check=True,
+                capture_output=True,
+                text=True
+            )
+        except Exception as e:
+            raise RuntimeError(f"❌ FFmpeg 未安装或不可用：{str(e)}")
 
-        # 初始化去水印模型（根据原项目实际参数调整）
+        # -------------------- 步骤 3：初始化去水印模型 --------------------
+        # 模型路径：假设模型文件位于项目的 kling/models 目录下（需上传至仓库）
+        model_path = str(project_root / "kling" / "models" / "watermark_remover.pth")
         remover = VideoWatermarkRemover(
-            device="cpu",  # 强制使用 CPU（关键！）
-            model_path="path/to/pretrained_model"  # 替换为实际模型路径（本地或 Hugging Face Hub）
+            device="cpu",  # 强制使用 CPU
+            model_path=model_path
         )
+        progress(0.3, desc="✅ 模型加载完成")
+
+        # -------------------- 步骤 4：执行去水印推理 --------------------
+        # 生成唯一输出文件名（避免覆盖）
+        input_filename = os.path.splitext(os.path.basename(input_path))[0]  # 输入文件名（不含后缀）
+        output_filename = f"{input_filename}_去水印后.mp4"  # 自定义输出文件名
+        output_path = os.path.join(tempfile.gettempdir(), output_filename)  # 临时输出路径
 
-        # 执行去水印（根据原项目实际函数名调整，如 remove_watermark）
+        # 执行去水印（根据原项目实际函数名调整）
         remover.remove_watermark(
             input_path=input_path,
-            output_path=output_path,
-            # 其他参数（如水印位置、强度等，根据原项目功能补充）
-            # watermark_region=(100, 100, 300, 200),  # 示例：(x1, y1, x2, y2)
-            # strength=0.8
+            output_path=output_path
+            # 可选参数（根据原项目功能调整）：
+            # watermark_region=(100, 100, 300, 200),  # 水印区域 (x1, y1, x2, y2)
+            # strength=0.8                           # 去水印强度
         )
+        progress(0.9, desc="✅ 去水印完成")
+
+        # -------------------- 步骤 5：生成下载链接（Hugging Face Space 特性） --------------------
+        # Space 会自动将临时文件暴露为可下载链接（格式：/files/文件名）
+        # 注意：临时文件仅在当前会话有效（约 72 小时未访问会被清理）
+        download_link = f"/files/{output_filename}"  # 下载链接路径
 
-        return output_path  # Gradio 会自动读取视频并展示
+        return output_path, download_link  # 返回视���路径和下载链接
 
     except Exception as e:
-        return f"处理失败：{str(e)}"
+        # 捕获所有异常并返回友好提示
+        error_msg = f"❌ 处理失败：{str(e)}"
+        return error_msg, None  # 错误时返回 None 隐藏下载链接
 
-# 创建 Gradio 界面
-with gr.Blocks(title="KLing 视频去水印工具") as demo:
-    gr.Markdown("# 🚀 KLing 视频去水印工具（Hugging Face Space CPU 版）")
-    gr.Markdown("上传含水印视频，自动去除水印（仅支持 CPU 推理，耗时可能较长）")
+# -------------------------- 创建 Gradio 界面（含储存检视组件） --------------------------
+with gr.Blocks(title="KLing 视频去水印工具（CPU 版）") as demo:
+    gr.Markdown("# 🚀 KLing 视频去水印工具")
+    gr.Markdown("上传含水印视频，自动去除水印（支持下载处理结果）")
 
+    # -------------------- 输入组件 --------------------
     with gr.Row():
         input_video = gr.Video(
-            label="输入视频（支持 MP4、MOV 等格式）",
+            label="输入视频（支持 MP4、MOV）",
             type="filepath",  # 获取文件路径而非二进制数据
-            interactive=True
+            interactive=True,
+            max_file_size=1024 * 1024 * 200,  # 限制 200MB
+            file_types=["video/mp4", "video/quicktime"]
         )
+
+    # -------------------- 输出组件 --------------------
+    with gr.Row():
         output_video = gr.Video(label="输出视频（去水印后）", interactive=False)
+        download_btn = gr.Button("⬇️ 下载视频", variant="secondary", interactive=False)  # 初始不可用
 
+    # -------------------- 操作按钮 --------------------
     submit_btn = gr.Button("开始去水印", variant="primary")
+    spinner = gr.Spinner(visible=False, label="处理中...")  # 加载提示
+
+    # -------------------- 下载链接提示（可选） --------------------
+    download_link = gr.HTML(visible=False)  # 隐藏的 HTML 组件（用于显示长链接）
+
+    # -------------------- 绑定事件（核心逻辑） --------------------
+    def process_wrapper(input_video):
+        """包装处理函数，控制界面组件状态"""
+        spinner.visible = True  # 显示加载提示
+        download_btn.interactive = False  # 初始禁用下载按钮
+
+        # 执行去水印处理
+        output_path, download_path = process_video(input_video)
+
+        # 处理结果反馈
+        if download_path:
+            # 处理成功：显示输出视频、启用下载按钮、显示下载链接（可选）
+            output_video.value = output_path  # 更新输出视频组件
+            download_btn.interactive = True   # 启用下载按钮
+            # 可选：显示完整下载链接（适合需要复制链接的场景）
+            # download_link.value = f"🔗 下载链接：<a href='{download_path}'>点击下载</a>"
+            # download_link.visible = True
+            return output_video, download_btn, None  # 隐藏错误提示
+        else:
+            # 处理失败：清空输出视频、禁用下载按钮、显示错误提示
+            output_video.value = None
+            download_btn.interactive = False
+            return output_video, download_btn, output_path  # 显示错误信息
 
-    # 绑定按钮点击事件（输入→处理函数→输出）
+    # 绑定提交按钮点击事件
     submit_btn.click(
-        fn=process_video,
+        fn=process_wrapper,
         inputs=input_video,
-        outputs=output_video
+        outputs=[output_video, download_btn, gr.Markdown()]  # 关联输出组件和错误提示
     )
 
-    # 示例说明（可选）
-    gr.Markdown("### 注意事项")
-    gr.Markdown("- 仅支持 CPU 推理，大视频（>100MB）处理时间较长")
-    gr.Markdown("- 模型路径需提前配置（可上传至 Hugging Face Hub 或使用绝对路径）")
-    gr.Markdown("- 若报错，检查 `requirements.txt` 依赖是否完整安装")
+    # -------------------- 注意事项 --------------------------
+    gr.Markdown("### ⚠️ 注意事项")
+    gr.Markdown("- 处理后的视频仅保留 72 小时（Hugging Face Space 临时文件策略）")
+    gr.Markdown("- 大视频（>500MB）处理时间较长（CPU 推理较慢），请耐心等待")
+    gr.Markdown("- 若下载按钮不可用，说明处理失败，请检查上方错误提示")
+    gr.Markdown("- 模型文件需提前上传至 `kling/models` 目录（或 Hugging Face Hub）")
 
 if __name__ == "__main__":
     # 本地调试（可选，部署到 Space 时无需运行）