adding gradio

app.py

@@ -1,82 +1,122 @@
-import gradio as gr
 import os
-import zipfile
-from pathlib import Path
-
-# The base directory that will be explored
-BASE_DIR = Path(__file__).resolve().parent / "your_target_folder"
-
-# Ensure BASE_DIR exists
-BASE_DIR.mkdir(parents=True, exist_ok=True)
-
-def resolve_path(current_rel_path: str) -> Path:
-    """Secure path resolution within BASE_DIR"""
-    resolved_path = (BASE_DIR / current_rel_path).resolve()
-    if BASE_DIR not in resolved_path.parents and resolved_path != BASE_DIR:
-        raise ValueError("Access outside base directory is not allowed.")
-    return resolved_path
-
-def list_dir(current_rel_path: str = ""):
-    current_path = resolve_path(current_rel_path)
-
-    # Parent folder navigation
-    parent_rel = str(Path(current_rel_path).parent) if current_rel_path else ""
-    entries = []
-
-    if current_path != BASE_DIR:
-        entries.append(("..", "⬆️ Parent Folder"))
-
-    # List directories and files
-    for item in sorted(current_path.iterdir()):
-        rel_item = os.path.relpath(item, BASE_DIR)
-        label = f"📁 {item.name}" if item.is_dir() else f"📄 {item.name}"
-        entries.append((rel_item, label))
-
-    return gr.update(choices=entries, value=None), f"Currently in: /{current_rel_path}"
-
-def download_entry(selected_rel_path: str):
-    selected_path = resolve_path(selected_rel_path)
-
-    if selected_path.is_file():
-        return selected_path
-    elif selected_path.is_dir():
-        zip_path = f"/tmp/{selected_path.name}.zip"
-        with zipfile.ZipFile(zip_path, "w") as zipf:
-            for root, dirs, files in os.walk(selected_path):
-                for file in files:
-                    abs_file = os.path.join(root, file)
-                    arc_file = os.path.relpath(abs_file, selected_path)
-                    zipf.write(abs_file, arc_file)
-        return zip_path
+import gradio as gr
+import subprocess
+import uuid
+import shutil
+from huggingface_hub import snapshot_download
+
+# ----------------------------------------
+# Step 1: Download Model Weights
+# ----------------------------------------
+MODEL_REPO = "roll-ai/DOVE"
+MODEL_PATH = "pretrained_models/"
+
+if not os.path.exists(MODEL_PATH) or len(os.listdir(MODEL_PATH)) == 0:
+    print("🔽 Downloading model weights from Hugging Face Hub...")
+    snapshot_download(
+        repo_id=MODEL_REPO,
+        repo_type="dataset",
+        local_dir=MODEL_PATH,
+        local_dir_use_symlinks=False
+    )
+    print("✅ Download complete.")
+
+# ----------------------------------------
+# Step 2: Setup Directories
+# ----------------------------------------
+INFERENCE_SCRIPT = "inference_script.py"
+OUTPUT_DIR = "results/DOVE/demo"
+UPLOAD_DIR = "input_videos"
+
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+# ----------------------------------------
+# Step 3: Inference Function
+# ----------------------------------------
+def run_inference(video_path, save_format):
+    input_name = f"{uuid.uuid4()}.mp4"
+    input_path = os.path.join(UPLOAD_DIR, input_name)
+    shutil.copy(video_path, input_path)
+
+    # --- Run inference script ---
+    cmd = [
+        "python", INFERENCE_SCRIPT,
+        "--input_dir", UPLOAD_DIR,
+        "--model_path", MODEL_PATH,
+        "--output_path", OUTPUT_DIR,
+        "--is_vae_st",
+        "--save_format", save_format
+    ]
+
+    try:
+        inference_result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            check=True
+        )
+        print("📄 Inference stdout:\n", inference_result.stdout)
+        print("⚠️ Inference stderr:\n", inference_result.stderr)
+    except subprocess.CalledProcessError as e:
+        print("❌ Inference failed.")
+        print("⚠️ STDOUT:\n", e.stdout)
+        print("⚠️ STDERR:\n", e.stderr)
+        return f"Inference failed:\n{e.stderr}", None
+
+    # --- Convert .mkv to .mp4 ---
+    mkv_path = os.path.join(OUTPUT_DIR, input_name).replace(".mp4", ".mkv")
+    mp4_path = os.path.join(OUTPUT_DIR, input_name)
+
+    if os.path.exists(mkv_path):
+        convert_cmd = [
+            "ffmpeg", "-y", "-i", mkv_path, "-c:v", "copy", "-c:a", "aac", mp4_path
+        ]
+        try:
+            convert_result = subprocess.run(
+                convert_cmd,
+                capture_output=True,
+                text=True,
+                check=True
+            )
+            print("🔄 FFmpeg stdout:\n", convert_result.stdout)
+            print("⚠️ FFmpeg stderr:\n", convert_result.stderr)
+        except subprocess.CalledProcessError as e:
+            print("❌ FFmpeg conversion failed.")
+            print("⚠️ STDOUT:\n", e.stdout)
+            print("⚠️ STDERR:\n", e.stderr)
+            return f"Inference OK, but conversion failed:\n{e.stderr}", None
+
+    if os.path.exists(mp4_path):
+        return "Inference successful!", mp4_path
     else:
-        return None
+        return "Output video not found.", None
 
+# ----------------------------------------
+# Step 4: Gradio Interface
+# ----------------------------------------
 with gr.Blocks() as demo:
-    current_path = gr.State("")
+    gr.Markdown("# 🎥 DOVE Video SR + Restoration Inference Demo")
+    gr.Markdown("⚙️ **Note:** Default `save_format` is `yuv444p`. If playback fails, try `yuv420p` for compatibility.")
 
-    gr.Markdown("# 📁 Folder Browser")
-    
     with gr.Row():
-        folder_dropdown = gr.Dropdown(label="Select File or Folder", choices=[])
-        refresh_btn = gr.Button("🔄 Refresh")
-    
-    status_text = gr.Textbox(label="Current Path", interactive=False)
-    download_btn = gr.Button("⬇️ Download Selected")
-    file_output = gr.File(label="Download Result")
-
-    # Events
-    refresh_btn.click(fn=list_dir, inputs=current_path, outputs=[folder_dropdown, status_text])
-    
-    folder_dropdown.change(
-        fn=lambda x: (x, *list_dir(x)),  # update path, refresh list
-        inputs=folder_dropdown,
-        outputs=[current_path, folder_dropdown, status_text],
-    )
+        input_video = gr.Video(label="Upload input video", type="filepath")
+        output_video = gr.Video(label="Output video")
 
-    download_btn.click(fn=download_entry, inputs=folder_dropdown, outputs=file_output)
-
-    # Initial trigger
-    demo.load(fn=list_dir, inputs=current_path, outputs=[folder_dropdown, status_text])
+    with gr.Row():
+        save_format = gr.Dropdown(
+            choices=["yuv444p", "yuv420p"],
+            value="yuv444p",
+            label="Save format (for video playback compatibility)"
+        )
+
+    run_button = gr.Button("Run Inference")
+    status = gr.Textbox(label="Status")
+
+    run_button.click(
+        fn=run_inference,
+        inputs=[input_video, save_format],
+        outputs=[status, output_video],
+    )
 
 demo.launch()
-

inference_script.py

@@ -751,4 +751,4 @@ if __name__ == "__main__":
         with open(out_path, 'w') as f:
             json.dump(output, f, indent=2)
 
-    print("All videos processed.")
+    print("All videos processed.")

requirements.txt

@@ -1,4 +1,3 @@
-gradio
 accelerate>=1.1.1
 transformers>=4.46.2
 numpy==1.26.0
@@ -19,3 +18,7 @@ opencv-python
 decord
 av
 torchdiffeq
+diffusers["torch"]
+transformers
+pyiqa
+huggingface_hub