|
|
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) |
