Skip to content

Use 256-pixel VideoVAE tiles by default #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ArthurOstapenko wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Use 256-pixel VideoVAE tiles by default #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -496,10 +496,11 @@ aggressive point. The latter produced a coherent walking fox and repeated at
8.02 seconds of DiT versus about 15.82 seconds natively. Native 256 uses the
same-cost spatial-RoPE adaptation described above; it remains a fast composition
preview rather than a substitute for a 512- or 768-class final render.
The video VAE automatically chooses a 256-320 pixel spatial tile from the
requested canvas geometry, minimizing repeated overlap work while keeping peak
storage bounded. `H3_VAE_TILE_PIXELS=256` restores the original conservative
tile plan for close-reference diagnosis.
The video VAE defaults to a 256-pixel spatial tile and 64-pixel minimum
overlap, matching the released checkpoint config. Set `H3_VAE_TILE_PIXELS`
to a multiple of 16 from 256 through 512 to override the tile size for
profiling. A value of 320 reproduces the previous automatic plan for the
measured 576x1024 case.

### Weight residency and streamed prompt encoding

Expand Down
26 changes: 5 additions & 21 deletions h3_video_vae.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -663,30 +663,16 @@ static int tile_count_for_extent(int extent, int tile_pixels) {
return count;
}

static int configured_tile_pixels(int pixel_height, int pixel_width) {
static int configured_tile_pixels(void) {
const char *value = getenv("H3_VAE_TILE_PIXELS");
if (value && *value) {
char *end = NULL;
long pixels = strtol(value, &end, 10);
if (end && !*end && pixels >= TILE_PIXELS && pixels <= 512 &&
pixels % SPATIAL_RATIO == 0) return (int)pixels;
}
int best = TILE_PIXELS;
uint64_t best_score = UINT64_MAX;
for (int pixels = TILE_PIXELS; pixels <= 320;
pixels += SPATIAL_RATIO) {
uint64_t tiles = (uint64_t)tile_count_for_extent(pixel_height, pixels) *
(uint64_t)tile_count_for_extent(pixel_width, pixels);
/* The resident VAE is dominated by linears and activation traffic;
* measured tile cost follows area more closely than cubic sequence
* growth at these shapes. Attention is still included in each tile. */
uint64_t score = tiles * (uint64_t)pixels * (uint64_t)pixels;
if (score < best_score) {
best = pixels;
best_score = score;
}
}
return best;
/* Match vae_tile_size from the released checkpoint config. */
return TILE_PIXELS;
}

static int tile_axis_build(int extent, int tile_pixels, tile_axis *axis,
Expand Down Expand Up @@ -910,8 +896,7 @@ h3_video_vae_decoder *h3_video_vae_decoder_load(
}
decoder->latent_h = latent_height;
decoder->latent_w = latent_width;
int tile_pixels = configured_tile_pixels(
latent_height * SPATIAL_RATIO, latent_width * SPATIAL_RATIO);
int tile_pixels = configured_tile_pixels();
int ok = load_latent_normalization(
weight_directory, decoder->latent_mean, decoder->latent_std,
error, error_size) &&
Expand Down Expand Up @@ -1213,8 +1198,7 @@ int h3_video_vae_decode(const char *weight_directory,
float latent_mean[LATENT_CHANNELS], latent_std[LATENT_CHANNELS];
if (!load_latent_normalization(weight_directory, latent_mean, latent_std,
error, error_size)) return 0;
int tile_pixels = configured_tile_pixels(
latent_height * SPATIAL_RATIO, latent_width * SPATIAL_RATIO);
int tile_pixels = configured_tile_pixels();
if (latent_time == 2 &&
(latent_height > TILE_PIXELS / SPATIAL_RATIO ||
latent_width > TILE_PIXELS / SPATIAL_RATIO)) {
Expand Down