Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from collections.abc import Iterator | |
| import torch | |
| from .padding import target_frame_count | |
| class TemporalVideoStream: | |
| """Expand a source timeline on demand and release consumed LQ frames. | |
| Only EOF introduces temporal padding. Frame indices refer to the expanded | |
| model timeline, including its unique first frame. | |
| """ | |
| def __init__(self, chunks: Iterator[torch.Tensor], compression: int, model_compression: int = 4): | |
| if compression not in (1, 2, 4) or model_compression % compression: | |
| raise ValueError("invalid temporal compression") | |
| self.chunks = iter(chunks) | |
| self.compression = compression | |
| self.repeat_factor = model_compression // compression | |
| self.source_frames = 0 | |
| self.effective_frames = 0 | |
| self.expanded_frames = 0 | |
| self.eof = False | |
| self.buffer = None | |
| self.offset = 0 | |
| self.last_frame = None | |
| self.max_buffer_frames = 0 | |
| def _append(self, value: torch.Tensor) -> None: | |
| if not value.shape[2]: | |
| return | |
| self.buffer = value.clone() if self.buffer is None else torch.cat((self.buffer, value), dim=2) | |
| self.expanded_frames += value.shape[2] | |
| self.max_buffer_frames = max(self.max_buffer_frames, self.buffer.shape[2]) | |
| def ensure(self, stop: int) -> bool: | |
| while self.expanded_frames < stop and not self.eof: | |
| try: | |
| value = next(self.chunks) | |
| except StopIteration: | |
| self.eof = True | |
| if not self.source_frames: | |
| raise ValueError("No frames decoded from input video") | |
| self.effective_frames = target_frame_count(self.source_frames, self.compression) | |
| missing = (self.effective_frames - self.source_frames) * self.repeat_factor | |
| self._append(self.last_frame.repeat(1, 1, missing, 1, 1)) | |
| break | |
| if value.ndim != 5 or value.shape[:2] != (1, 3) or not value.shape[2]: | |
| raise ValueError("Input chunks must be nonempty [1,3,F,H,W] tensors") | |
| if self.last_frame is not None and value.shape[-2:] != self.last_frame.shape[-2:]: | |
| raise ValueError("Input chunk dimensions changed") | |
| first = self.source_frames == 0 | |
| self.source_frames += value.shape[2] | |
| self.last_frame = value[:, :, -1:].clone() | |
| if first: | |
| expanded = torch.cat((value[:, :, :1], value[:, :, 1:].repeat_interleave(self.repeat_factor, dim=2)), dim=2) | |
| else: | |
| expanded = value.repeat_interleave(self.repeat_factor, dim=2) | |
| self._append(expanded) | |
| return self.expanded_frames >= stop | |
| def window(self, start: int, stop: int) -> torch.Tensor: | |
| self.ensure(stop) | |
| if start < self.offset or stop <= start: | |
| raise ValueError("Requested frames have already been released or window is empty") | |
| value = self.buffer[:, :, start - self.offset:stop - self.offset] | |
| if not value.shape[2]: | |
| value = self.last_frame[:, :, :0] | |
| missing = stop - start - value.shape[2] | |
| if missing: | |
| value = torch.cat((value, self.last_frame.repeat(1, 1, missing, 1, 1)), dim=2) | |
| return value.clone() | |
| def release_before(self, frame: int) -> None: | |
| frame = min(frame, self.expanded_frames) | |
| if frame > self.offset: | |
| self.buffer = self.buffer[:, :, frame - self.offset:].clone() | |
| self.offset = frame | |
| def projection_chunks(self, chunk_latents: int, model_compression: int): | |
| index = 0 | |
| stride = chunk_latents * model_compression | |
| while self.ensure(1 + index * stride): | |
| start = 0 if index == 0 else 1 + index * stride | |
| stop = 1 + (index + 1) * stride | |
| value = self.window(start, stop) | |
| last = self.eof and self.expanded_frames < stop | |
| yield index, value, start, last | |
| index += 1 | |
| def close(self) -> None: | |
| close = getattr(self.chunks, "close", None) | |
| if close is not None: | |
| close() | |
| self.buffer = None | |
| self.last_frame = None | |