Download native-eiffel/native_capture.py from peih/matrix-shooter-train: direct link, hf CLI and curl.
- Browser
- Download file 14.4 kB
-
https://huggingface.co/peih/matrix-shooter-train/resolve/main/native-eiffel/native_capture.py
- Command line
-
hf download hf://peih/matrix-shooter-train/native-eiffel/native_capture.py
-
curl -L -o native_capture.py https://huggingface.co/peih/matrix-shooter-train/resolve/main/native-eiffel/native_capture.py
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 | |