| import json |
| import os |
| import re |
| from flask import Flask, render_template_string, send_from_directory |
|
|
| |
|
|
| |
| FOLDERS = { |
| 'baseline': 'baseline', |
| 'teacache': 'teacache', |
| 'magcache': 'magcache', |
| 'ours': 'ours' |
| } |
|
|
| PROMPT_FILE = 'prompts.json' |
|
|
| |
|
|
| app = Flask(__name__) |
| BASE_DIR = os.path.abspath(os.path.dirname(__file__)) |
|
|
| def get_prompts(): |
| path = os.path.join(BASE_DIR, PROMPT_FILE) |
| if not os.path.exists(path): |
| return [] |
| |
| with open(path, 'r', encoding='utf-8') as f: |
| try: |
| content = json.load(f) |
| if isinstance(content, list): return content |
| except: |
| f.seek(0) |
| lines = [line.strip() for line in f if line.strip()] |
| return lines |
|
|
| def find_file_by_id(folder, video_id): |
| dir_path = os.path.join(BASE_DIR, folder) |
| if not os.path.exists(dir_path): return None |
| |
| for fname in os.listdir(dir_path): |
| if fname.startswith(video_id) and fname.lower().endswith('.mp4'): |
| return fname |
| return None |
|
|
| def build_tasks(): |
| tasks = [] |
| prompts = get_prompts() |
| |
| baseline_path = os.path.join(BASE_DIR, FOLDERS['baseline']) |
| if not os.path.exists(baseline_path): |
| print("Baseline folder missing!") |
| return [] |
|
|
| baseline_files = sorted([f for f in os.listdir(baseline_path) if f.lower().endswith('.mp4')]) |
| id_regex = re.compile(r'^(\d{4})') |
|
|
| for i, b_file in enumerate(baseline_files): |
| match = id_regex.match(b_file) |
| if not match: continue |
| |
| vid_id = match.group(1) |
| prompt_text = prompts[i] if i < len(prompts) else "(No prompt)" |
|
|
| tea_file = find_file_by_id(FOLDERS['teacache'], vid_id) |
| mag_file = find_file_by_id(FOLDERS['magcache'], vid_id) |
| our_file = find_file_by_id(FOLDERS['ours'], vid_id) |
|
|
| tasks.append({ |
| "index": i + 1, |
| "id": vid_id, |
| "prompt": prompt_text, |
| "baseline": f"videos/{FOLDERS['baseline']}/{b_file}", |
| "teacache": f"videos/{FOLDERS['teacache']}/{tea_file}" if tea_file else "", |
| "magcache": f"videos/{FOLDERS['magcache']}/{mag_file}" if mag_file else "", |
| "ours": f"videos/{FOLDERS['ours']}/{our_file}" if our_file else "" |
| }) |
|
|
| print(f"Loaded {len(tasks)} videos.") |
| return tasks |
|
|
| TASKS = build_tasks() |
|
|
| |
|
|
| HTML = """ |
| <!DOCTYPE html> |
| <html lang="en"> |
| <head> |
| <meta charset="UTF-8"> |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> |
| <title>2x2 Visualizer</title> |
| <script src="https://cdn.tailwindcss.com"></script> |
| <style> |
| body { background-color: #0f172a; color: #e2e8f0; font-family: sans-serif; } |
| |
| /* Darker theme to make videos pop */ |
| .video-box { |
| background: #1e293b; |
| padding: 0.5rem; |
| border-radius: 0.5rem; |
| height: 100%; |
| display: flex; |
| flex-direction: column; |
| } |
| |
| video { |
| width: 100%; |
| height: auto; |
| border-radius: 0.25rem; |
| background: #000; |
| flex-grow: 1; |
| } |
| |
| .label { |
| text-align: center; |
| font-weight: bold; |
| font-size: 1rem; |
| margin-bottom: 0.5rem; |
| color: #94a3b8; |
| text-transform: uppercase; |
| letter-spacing: 0.05em; |
| } |
| |
| /* Highlight 'Ours' */ |
| .highlight { border: 2px solid #3b82f6; background: #1e3a8a; } |
| .highlight .label { color: #60a5fa; } |
| |
| .prompt-box { |
| background: #1e293b; |
| border-bottom: 1px solid #334155; |
| padding: 1rem; |
| position: sticky; |
| top: 0; |
| z-index: 50; |
| } |
| </style> |
| </head> |
| <body class="h-screen flex flex-col overflow-hidden"> |
| |
| <div class="prompt-box flex flex-col md:flex-row gap-4 items-center flex-none shadow-lg"> |
| <div class="flex-none w-full md:w-48 space-y-2"> |
| <div class="flex justify-between items-center text-xs text-gray-400 font-mono"> |
| <span>ID: <span id="vid-id" class="text-white">--</span></span> |
| <span><span id="idx" class="text-white">0</span> / <span id="total">0</span></span> |
| </div> |
| <div class="flex gap-2"> |
| <button onclick="move(-1)" class="flex-1 bg-gray-700 hover:bg-gray-600 text-white py-2 rounded font-bold text-sm transition">Prev</button> |
| <button onclick="move(1)" class="flex-1 bg-blue-600 hover:bg-blue-500 text-white py-2 rounded font-bold text-sm transition">Next</button> |
| </div> |
| </div> |
| |
| <div class="flex-grow overflow-y-auto max-h-20 w-full"> |
| <p id="prompt-text" class="text-sm md:text-lg text-gray-200 leading-snug font-light">Loading...</p> |
| </div> |
| </div> |
| |
| <div class="flex-grow overflow-y-auto p-4"> |
| <div class="w-full h-full max-w-[1920px] mx-auto"> |
| |
| <div class="grid grid-cols-1 md:grid-cols-2 gap-4 h-full"> |
| |
| <div class="video-box"> |
| <div class="label">Baseline</div> |
| <video id="v-baseline" controls loop muted playsinline></video> |
| </div> |
| |
| <div class="video-box"> |
| <div class="label">TeaCache</div> |
| <video id="v-teacache" controls loop muted playsinline></video> |
| </div> |
| |
| <div class="video-box"> |
| <div class="label">MagCache</div> |
| <video id="v-magcache" controls loop muted playsinline></video> |
| </div> |
| |
| <div class="video-box highlight"> |
| <div class="label">Ours</div> |
| <video id="v-ours" controls loop muted playsinline></video> |
| </div> |
| |
| </div> |
| </div> |
| </div> |
| |
| <script> |
| const tasks = {{ tasks_json | safe }}; |
| let current = 0; |
| |
| const els = { |
| prompt: document.getElementById('prompt-text'), |
| id: document.getElementById('vid-id'), |
| idx: document.getElementById('idx'), |
| total: document.getElementById('total'), |
| v1: document.getElementById('v-baseline'), |
| v2: document.getElementById('v-teacache'), |
| v3: document.getElementById('v-magcache'), |
| v4: document.getElementById('v-ours') |
| }; |
| |
| els.total.innerText = tasks.length; |
| |
| function loadTask(index) { |
| if (tasks.length === 0) return; |
| |
| if (index < 0) index = tasks.length - 1; |
| if (index >= tasks.length) index = 0; |
| current = index; |
| |
| const t = tasks[current]; |
| |
| els.prompt.innerText = t.prompt; |
| els.id.innerText = t.id; |
| els.idx.innerText = t.index; |
| |
| els.v1.src = t.baseline; |
| els.v2.src = t.teacache; |
| els.v3.src = t.magcache; |
| els.v4.src = t.ours; |
| |
| [els.v1, els.v2, els.v3, els.v4].forEach(v => { |
| if(v.src) { |
| v.currentTime = 0; |
| v.play().catch(e => {}); |
| } |
| }); |
| } |
| |
| function move(dir) { loadTask(current + dir); } |
| |
| document.addEventListener('keydown', e => { |
| if (e.key === "ArrowLeft") move(-1); |
| if (e.key === "ArrowRight") move(1); |
| }); |
| |
| if (tasks.length > 0) loadTask(0); |
| else els.prompt.innerText = "No tasks found. Check console."; |
| </script> |
| </body> |
| </html> |
| """ |
|
|
| @app.route('/') |
| def home(): |
| return render_template_string(HTML, tasks_json=json.dumps(TASKS)) |
|
|
| @app.route('/videos/<folder>/<path:filename>') |
| def serve_file(folder, filename): |
| if folder in FOLDERS: |
| return send_from_directory(os.path.join(BASE_DIR, FOLDERS[folder]), filename) |
| return "Error", 404 |
|
|
| if __name__ == '__main__': |
| print("Starting server...") |
| app.run(port=5001, debug=True) |
