woutermans
11a3511120
The dumps are one harness frame per preset (11/13/12) with D4R_CUDA_REPLAY_DUMP_DIR set. The K/M/L models built from them are byte-identical to k-def/m-def/l-def.
96 lines
5.3 KiB
Python
96 lines
5.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Prepare a full-frame native K replay from a capture. Does not translate kernels.
|
|
|
|
prepare_k.py CAPTURE_DIR OUTPUT_DIR [--frames N] [--start FRAME]
|
|
prepare_k.py CAPTURE_DIR OUTPUT_DIR --lut-only
|
|
The model's reconstruction tables can be packaged using OUTPUT_DIR/lut.bin. --lut-only writes only that, from the
|
|
first output launch's replay (a capture without the launch dumps: install.sh's weight dump).
|
|
"""
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
import struct
|
|
from capture import INPUT, OUTPUT, ORDER, launches, input_params, output_params
|
|
|
|
|
|
def quote(path):
|
|
return json.dumps(str(Path(path).resolve()))
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('capture', type=Path)
|
|
parser.add_argument('output', type=Path)
|
|
parser.add_argument('--frames', type=int, default=0, help='0 = all captured frames')
|
|
parser.add_argument('--start', type=int, default=0)
|
|
parser.add_argument('--lut-only', action='store_true', help='write only lut.bin')
|
|
a = parser.parse_args()
|
|
if a.lut_only:
|
|
outputs = launches(a.capture, OUTPUT)
|
|
if not outputs:
|
|
raise ValueError('capture holds no K output launch')
|
|
a.output.mkdir(parents=True, exist_ok=True)
|
|
(a.output / 'lut.bin').write_bytes(outputs[0].allocation(280, 16 * 32768))
|
|
return
|
|
inputs, outputs = launches(a.capture, INPUT), launches(a.capture, OUTPUT)
|
|
if not inputs or not outputs or len(outputs) > len(inputs):
|
|
raise ValueError('capture must contain complete K input/output frame pairs')
|
|
pairs = list(zip(inputs, outputs))
|
|
for i, (ip, op) in enumerate(pairs):
|
|
if not ip.sequence < op.sequence or (i+1 < len(inputs) and op.sequence >= inputs[i+1].sequence):
|
|
raise ValueError('capture has a missing or out-of-order input/output launch')
|
|
if a.start < 0 or a.start >= len(pairs) or a.frames < 0:
|
|
raise ValueError('invalid frame selection')
|
|
end = len(pairs) if a.frames == 0 else min(len(pairs), a.start + a.frames)
|
|
selected = pairs[a.start:end]
|
|
# Unselected later frames may be incomplete (for example a capture stopped while writing them).
|
|
for ip, op in selected:
|
|
for resource in (*ip.resources.values(), *op.resources.values()): resource.validate()
|
|
layer_launches = {name: launches(a.capture, 'dltss_pwin_'+name+'_layer') for name in ORDER}
|
|
first_in, first_out = selected[0]
|
|
geometry = []
|
|
for name in ORDER:
|
|
matches = [p for p in layer_launches[name] if first_in.sequence < p.sequence < first_out.sequence]
|
|
if len(matches) != 1:
|
|
raise ValueError(f'frame {a.start}: expected one {name} launch, found {len(matches)}')
|
|
geometry.append((*matches[0].i2(0), *matches[0].i2(56)))
|
|
final = first_out.resources[312]
|
|
a.output.mkdir(parents=True, exist_ok=True)
|
|
(a.output / 'shape.bin').write_bytes(struct.pack('<2I', final.width, final.height) + b''.join(struct.pack('<4i', *s) for s in geometry))
|
|
# NGX's whole allocation: sixteen reconstruction tables, of which a launch uses the one for its upscaling ratio
|
|
(a.output / 'lut.bin').write_bytes(first_out.allocation(280, 16 * 32768))
|
|
lines = [f'shape {quote(a.output / "shape.bin")}',
|
|
f'history {quote(first_in.resources[176].path)} {quote(first_out.resources[232].path)} {quote(first_in.resources[224].path)}']
|
|
resource_ids = {}
|
|
references = []
|
|
for number, (ip, op) in enumerate(selected):
|
|
if not ip.sequence < op.sequence or (number and ip.sequence <= selected[number-1][1].sequence):
|
|
raise ValueError('capture frames are not in network order')
|
|
if op.resources[312].width != final.width or op.resources[312].height != final.height:
|
|
raise ValueError('resolution changes require another Engine')
|
|
frame_geometry = []
|
|
for name, expected in zip(ORDER, geometry):
|
|
matches = [p for p in layer_launches[name] if ip.sequence < p.sequence < op.sequence]
|
|
if len(matches) != 1 or matches[0].i2(0) != expected[:2]:
|
|
raise ValueError(f'frame {number}: missing or changed {name} geometry')
|
|
frame_geometry.append((*matches[0].i2(0), *matches[0].i2(56)))
|
|
ids = []
|
|
for off in (168, 184, 192):
|
|
r = ip.resources[off]
|
|
if r.path not in resource_ids:
|
|
i = len(resource_ids); resource_ids[r.path] = i
|
|
lines.append(f'image {i} {quote(r.path)} {r.width} {r.height} {r.channels} {r.bits}')
|
|
ids.append(resource_ids[r.path])
|
|
inp, outp = a.output / f'frame-{number}.input', a.output / f'frame-{number}.output'
|
|
inp.write_bytes(input_params(ip)); outp.write_bytes(output_params(op))
|
|
windows = a.output / f'frame-{number}.windows'
|
|
windows.write_bytes(b''.join(struct.pack('<4i', *s) for s in frame_geometry))
|
|
lines.append(f'frame {quote(inp)} {quote(outp)} {quote(windows)} ' + ' '.join(map(str, ids)))
|
|
references.append({str(off): str(op.resources[off].path) for off in (288,296,304,312)})
|
|
(a.output / 'case.txt').write_text('\n'.join(lines) + '\n')
|
|
(a.output / 'references.json').write_text(json.dumps(references, indent=2) + '\n')
|
|
print(f'Prepared {len(selected)} frames, {final.width}x{final.height}, token grid {geometry[0][0]}x{geometry[0][1]}')
|
|
|
|
|
|
if __name__ == '__main__': main()
|