File size: 2,016 Bytes
5dd27d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | """Streaming VAE timings at 960x544: encoder (1 frame, then 4-frame chunks) and decoder (1 latent -> 4 frames)."""
import json
import time
import torch
from sftrain import conditions
device = torch.device('cuda')
enc = conditions.vae_encoder(device)
dec = conditions.vae_decoder(device)
H, W = conditions.HEIGHT, conditions.WIDTH
print(json.dumps(dict(H=H, W=W, enc_dtype=str(getattr(enc.config, 'dtype', '?')), dec_dtype=str(dec.config.dtype))), flush=True)
def timed(fn, n):
out = []
for _ in range(n):
torch.cuda.synchronize()
t = time.perf_counter()
fn()
torch.cuda.synchronize()
out.append((time.perf_counter() - t) * 1000)
return out
with torch.no_grad():
for rep in range(2):
cache = enc.initialize_autoregressive_cache()
first = torch.rand(1, 1, 3, H, W, device=device) * 2 - 1
t0 = timed(lambda: enc(first, autoregressive_index=0, cache=cache), 1)
ts = []
for i in range(1, 11):
chunk = torch.rand(1, 4, 3, H, W, device=device) * 2 - 1
ts += timed(lambda: enc(chunk, autoregressive_index=i, cache=cache), 1)
full = torch.rand(1, 81, 3, H, W, device=device) * 2 - 1
tf = timed(lambda: enc(full, autoregressive_index=0, cache=enc.initialize_autoregressive_cache()), 1)
lat = torch.randn(1, 11, 48, conditions.LATENT_H, conditions.LATENT_W, device=device, dtype=dec.config.dtype)
dcache = dec.initialize_autoregressive_cache()
d0 = timed(lambda: dec(lat[:, :1], autoregressive_index=0, cache=dcache), 1)
ds = []
for i in range(1, 11):
ds += timed(lambda: dec(lat[:, i:i + 1], autoregressive_index=i, cache=dcache), 1)
print(json.dumps(dict(rep=rep, enc_first_ms=round(t0[0], 1), enc_4frame_ms=[round(x, 1) for x in ts],
enc_81frames_ms=round(tf[0], 1), dec_first_ms=round(d0[0], 1),
dec_1latent_4frames_ms=[round(x, 1) for x in ds])), flush=True)
|