matrix-shooter-train / native-eiffel /native_capture.py
peih's picture
backup native-eiffel
5dd27d0 verified
Raw History Blame Contribute Delete
14.4 kB
"""Game capture without a browser: the Node engine steps the game, the GPU draws it.
``NativeSceneRenderer`` has the interface ``GameCaptureWorker`` expects from a
renderer (the same one ``NativeMatrixRenderer`` implements), so the session,
its capture thread and the model loop are unchanged. Condition frames stay on
the GPU as (H, W, 3) uint8 tensors and go to the proxy encoder as they are;
the game-window frames come back to the host for the paired video, as before.
``install()`` is called from the deployment adapter. Switches:
EDEN_NATIVE_CAPTURE=1 use it for the slugs in EDEN_NATIVE_SLUGS
EDEN_NATIVE_SLUGS=eiffel-assault
EDEN_NATIVE_AUX=full render every frame for the game-window video;
'condition' renders only the model's frames and
repeats the latest one in that video. Defaults
to 'full' only when that video is sent to the
player (FLASHDREAMS_PAIRED_VIDEO=1).
EDEN_LATE_CAPTURE=1 run each chunk as one burst at the end of its
paced window, so the model sees fresh input
EDEN_CHUNK_PERIOD_MS=437 pace a 12-frame chunk to this wall time. Faster
capture would otherwise speed the game up, since
presentation follows production; 437 ms is the
live median before this change.
EDEN_NATIVE_FEATURES=1 encode each condition frame with the proxy encoder
as soon as it is rendered, on the capture stream,
and hand the model the finished features; the
model loop then skips its own proxy pass.
In a burst the chunk is one engine command: the engine streams each condition
frame's scene back the moment it is simulated, the GPU renders and encodes it
while the engine keeps stepping, and the chunk is ready one render and one
single-frame encode after the last condition frame.
"""
import json
import os
import threading
import time
from collections import deque
from dataclasses import replace
from pathlib import Path
import torch
from loguru import logger
from engine_client import Engine, standalone_config
from native_render import NativeRenderer
# Proxy features by chunk id, from the capture thread to RenderControlStream.take().
FEATURES = {}
_FEATURES_KEPT = 4
class ProxyFeatures:
"""Model input that already went through the proxy encoder."""
__slots__ = ('proxy_features',)
def __init__(self, proxy_features):
self.proxy_features = proxy_features
def enabled_for(slug):
if os.environ.get('EDEN_NATIVE_CAPTURE', '0') != '1':
return False
return slug in os.environ.get('EDEN_NATIVE_SLUGS', 'eiffel-assault').split(',')
def _late_capture():
return os.environ.get('EDEN_LATE_CAPTURE', '1') == '1'
def _features_enabled():
return os.environ.get('EDEN_NATIVE_FEATURES', '1') == '1'
def _proxy_encoder():
"""A second copy of the live proxy encoder, built from the same profile as the runner's."""
from lightx2v.models.input_encoders.edenlive.proxy_dino import ProxyDINO
root = Path(os.environ.get('EDEN_DEPLOY_ROOT', '/root/autodl-tmp/edenlive'))
profile = json.loads(Path(os.environ.get('EDEN_5090_PROFILE', root / 'runtime-5090.json')).read_text())
# As EdenLiveRunner.load_proxy_encoder() chooses for the cached Sage kernels.
attention = 'sage_attn2' if profile.get('self_attn_1_type') in ('sage_attn2_sm90_fused', 'sage_attn2_sm120_cached') else 'torch_sdpa'
return ProxyDINO(str(root / 'models/EdenLive-x2v'), profile.get('proxy_resize_mode', 'bilinear'),
use_cuda_graph=profile.get('use_cuda_graph', False), attention_type=attention,
quantized_ckpt=profile.get('proxy_quantized_ckpt'))
def _aux_full():
default = 'full' if os.environ.get('FLASHDREAMS_PAIRED_VIDEO') == '1' else 'condition'
return os.environ.get('EDEN_NATIVE_AUX', default) == 'full'
def chunk_period_s(frame_count):
return float(os.environ.get('EDEN_CHUNK_PERIOD_MS', '437')) / 1000 * frame_count / 12
def split_release(events):
"""Hold back from the first release of a key pressed earlier in the same batch.
Delivered together, a press and its release would land before one
simulation step and a game that polls held keys would never see the tap.
The remainder is delivered with the next frame.
"""
pressed = set()
for index, event in enumerate(events):
action = event.get('action')
if action == 'pointer':
continue
if event.get('phase') == 'down':
pressed.add(action)
elif event.get('phase') == 'up' and action in pressed:
return events[:index], events[index:]
return events, []
class NativeSceneRenderer:
def __init__(self, slug, *, width=960, height=544, games_root=None):
root = Path(games_root or os.environ.get('EDEN_GAMES_ROOT', '/root/autodl-tmp/edenlive/EdenLive/realtime/games'))
entry = next(g for g in json.loads((root / 'registry.json').read_text())['games'] if g['slug'] == slug)
self.slug, self.width, self.height = slug, width, height
self.engine = Engine(slug, width, height, games_root=root, config=standalone_config(entry))
# Its own high-priority stream: a render never queues behind model kernels.
self.stream = torch.cuda.Stream(priority=-1)
with torch.cuda.stream(self.stream):
self.renderer = NativeRenderer(self.engine.scene, self.engine.scene_arrays)
self.webgl_renderer = f'native nvdiffrast on {torch.cuda.get_device_name()}'
self.last_evaluate_s = self.last_decode_s = 0.0
self.last_burst_s = 0.0
self._bursts = deque(maxlen=8)
self._target_end = None
self._carry = []
self._last_host = None
self._last_device = None
self.proxy = None
self.measure_frames = False
self.frame_to_features_s = []
if _features_enabled():
with torch.cuda.stream(self.stream):
self.proxy = _proxy_encoder()
logger.info('Game2V capture: {} runs browser-free ({} draws)', slug, len(self.engine.scene['draws']))
def render_chunk(self, chunk, *, duration_s=0.0, drain_events=None):
started = time.perf_counter()
late = _late_capture() and duration_s > 0
if late:
# Chunks end on a fixed schedule, each one period after the last
# one's scheduled end, so finishing a burst early never compounds
# into a faster game. After a pause the schedule restarts.
previous = self._target_end
target = started + duration_s if previous is None or previous < started - duration_s else previous + duration_s
self._target_end = target
if self._bursts:
time.sleep(max(0.0, target - (max(self._bursts) + 0.003) - time.perf_counter()))
else:
self._target_end = None
burst_started = time.perf_counter()
if late or duration_s <= 0:
return self._burst(chunk, drain_events, burst_started)
aux_full = _aux_full()
frames, images = [], []
engine_s = 0.0
packets = chunk.frames
draws = []
def send(index):
# Input is read when a step is queued, as the browser reads it per frame.
events = self._carry + (drain_events() if drain_events is not None else [])
events, self._carry = split_release(events)
draw = packets[index].capture_condition or aux_full or (self._last_device is None and not draws)
draws.append(draw)
self.engine.send_advance(packets[index].simulation_tick, events, capture=draw)
with torch.cuda.stream(self.stream):
# Paced without a burst (EDEN_LATE_CAPTURE=0): one step at a time,
# input read before each frame, spread over the chunk's window.
for index, packet in enumerate(packets):
send(index)
step_started = time.perf_counter()
header, scene = self.engine.receive()
engine_s += time.perf_counter() - step_started
if draws[index]:
self._last_device = self.renderer.render(scene)
images.append(self._last_device)
frames.append((packet, header.get('state'), self._last_device, len(images) - 1))
remaining = started + duration_s * (index + 1) / len(packets) - time.perf_counter()
if remaining > 0:
time.sleep(remaining)
# Game-window frames for the paired video, in one copy.
host = torch.stack(images).cpu().numpy() if images else None
# Everything the model will read from this chunk is finished on the
# render stream before the chunk is handed over.
self.stream.synchronize()
out = []
for packet, state, device, slot in frames:
ui = host[slot] if slot >= 0 else self._last_host
out.append(replace(packet, ui_frame=ui, ui_state=state,
condition_frame=device if packet.capture_condition else None))
if host is not None:
self._last_host = host[-1]
burst = time.perf_counter() - burst_started
self.last_burst_s = burst
self._bursts.append(burst)
self.last_evaluate_s = engine_s
self.last_decode_s = burst - engine_s
return replace(chunk, frames=tuple(out))
def _burst(self, chunk, drain_events, burst_started):
"""The whole chunk as one engine command, condition frames rendered and encoded as they arrive."""
packets = chunk.frames
aux_full = _aux_full()
# Input that piled up before the burst opens it; a release that would
# cancel a press in the same frame moves to the next frame.
pending = self._carry + (drain_events() if drain_events is not None else [])
plan = []
for index, packet in enumerate(packets):
now, pending = split_release(pending)
draw = packet.capture_condition or aux_full or (index == 0 and self._last_device is None)
plan.append(dict(targetTick=packet.simulation_tick, events=now, capture=draw))
self._carry = pending
self.engine.send_chunk(plan)
rendered, features = {}, []
encode = self.proxy is not None
with torch.cuda.stream(self.stream):
while True:
header, scene = self.engine.receive()
if header['type'] != 'scene':
break
index = header['index']
arrived = time.perf_counter()
image = self.renderer.render(scene)
rendered[index] = image
if encode and packets[index].capture_condition:
# The encoder's CUDA graph reuses its output buffer.
features.append(self.proxy.infer(image.permute(2, 0, 1)[None].contiguous()).clone())
if self.measure_frames:
# Measurement only: wait for this frame's features here.
torch.cuda.current_stream().synchronize()
self.frame_to_features_s.append(time.perf_counter() - arrived)
states = header['states']
engine_done = time.perf_counter()
order = sorted(rendered)
host = torch.stack([rendered[i] for i in order]).cpu().numpy() if order else None
self.stream.synchronize()
ready = time.perf_counter()
# Engine-to-features for the chunk's last rendered frame, and for the
# whole chunk from the moment the engine command went out.
self.last_frame_to_features_s = ready - arrived if rendered else 0.0
self.last_engine_s = engine_done - burst_started
out, slot = [], -1
for index, packet in enumerate(packets):
if index in rendered:
self._last_device = rendered[index]
slot = order.index(index)
ui = host[slot] if slot >= 0 else self._last_host
out.append(replace(packet, ui_frame=ui, ui_state=states[index],
condition_frame=rendered.get(index) if packet.capture_condition else None))
if host is not None:
self._last_host = host[-1]
if encode and len(features) == sum(p.capture_condition for p in packets):
FEATURES[chunk.chunk_id] = torch.cat(features, dim=2)
for stale in sorted(FEATURES)[:-_FEATURES_KEPT]:
FEATURES.pop(stale, None)
burst = time.perf_counter() - burst_started
self.last_burst_s = burst
self._bursts.append(burst)
self.last_evaluate_s = header['timing']['chunk_ms'] / 1000
self.last_decode_s = burst - self.last_evaluate_s
return replace(chunk, frames=tuple(out))
def probe(self, *, drain_events=None):
events = self._carry + (drain_events() if drain_events is not None else [])
self._carry = []
return {'state': self.engine.probe(events).get('state')}
def close(self):
engine, self.engine = self.engine, None
if engine is not None:
engine.close()
self.renderer = self.proxy = None
self._last_device = self._last_host = None
FEATURES.clear()
def install(game_renderer, session):
"""Route enabled slugs to NativeSceneRenderer and pace their chunks to a fixed period."""
if getattr(game_renderer, '_native_capture_installed', False):
return
original_type = game_renderer._renderer_type
def renderer_type(slug, games_root=None):
if enabled_for(slug):
return NativeSceneRenderer
return original_type(slug, games_root)
game_renderer._renderer_type = renderer_type
original_pace = session._pace_duration_s
def pace_duration_s(state, step_index):
if not enabled_for(state.config.game_slug):
return original_pace(state, step_index)
return chunk_period_s(int(state.pipeline.get_num_output_frames(step_index)))
session._pace_duration_s = pace_duration_s
game_renderer._native_capture_installed = True