| """TemporalFrameBuffer — ring buffer for video latents with HCA compression. |
| |
| Stores the last N video latents (local) and maintains a compressed long-range |
| cache via TernaryScaleTensor projection. Used for conditioning video generation |
| on previous time steps. |
| |
| Latent shape: [B, C, H', W'] where C=OPEN_SORA_LATENT_CHANNELS=4, |
| H'=VIDEO_HEIGHT=32, W'=VIDEO_WIDTH=32. Each "latent" is one 4-frame chunk. |
| """ |
| import torch |
| import torch.nn as nn |
| from ..kernel.ternary_scale import TernaryScaleTensor, TScaleType |
| from .ring_buffer import GPURingBuffer |
| from ..config import FRAME_BUFFER_LOCAL_SIZE, FRAME_BUFFER_CACHE_STRIDE, \ |
| OPEN_SORA_LATENT_CHANNELS, VIDEO_HEIGHT, VIDEO_WIDTH |
|
|
|
|
| class TemporalFrameBuffer(nn.Module): |
| def __init__(self, local_size=FRAME_BUFFER_LOCAL_SIZE, |
| cache_stride=FRAME_BUFFER_CACHE_STRIDE, |
| latent_channels=OPEN_SORA_LATENT_CHANNELS, |
| height=VIDEO_HEIGHT, width=VIDEO_WIDTH, |
| tscale_type=TScaleType.T32): |
| super().__init__() |
| self.latent_channels = latent_channels |
| self.spatial_dim = height * width |
| self.latent_flat_dim = latent_channels * self.spatial_dim |
|
|
| self.local = GPURingBuffer( |
| max_size=local_size, |
| dtype=torch.float32, |
| dim=self.latent_flat_dim, |
| ) |
|
|
| self.compress_proj = TernaryScaleTensor( |
| self.latent_flat_dim, |
| self.latent_flat_dim // 4, |
| tscale_type=tscale_type, |
| ) |
| self.compressed_cache = [] |
| self.cache_stride = cache_stride |
| self._frames_since_compress = 0 |
|
|
| def append(self, latent): |
| B = latent.shape[0] |
| flat = latent.reshape(B, -1) |
| self.local.append(flat) |
|
|
| self._frames_since_compress += 1 |
| if self._frames_since_compress >= self.cache_stride: |
| compressed = self.compress_proj(flat) |
| self.compressed_cache.append(compressed.detach()) |
| self._frames_since_compress = 0 |
|
|
| def get_local(self, n=None): |
| n = n or self.local.max_size |
| result = self.local.get_last_n(n) |
| if result.dim() == 0 or result.shape[0] == 0: |
| return torch.zeros(0, 1, self.latent_flat_dim) |
| if result.dim() == 1: |
| result = result.unsqueeze(0) |
| return result |
|
|
| def get_compressed_cache(self): |
| if not self.compressed_cache: |
| return torch.zeros(0, 1, self.latent_flat_dim // 4) |
| return torch.stack(self.compressed_cache, dim=0) |
|
|
| def get_conditioning(self, n_local=None): |
| return { |
| "local": self.get_local(n_local), |
| "compressed": self.get_compressed_cache(), |
| } |
|
|
| def reset(self): |
| self.local.reset() |
| self.compressed_cache = [] |
| self._frames_since_compress = 0 |
|
|
