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