File size: 2,786 Bytes
b9fc164 | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | import random
import re
from pathlib import Path
import numpy as np
from einops import rearrange
import torch
import torchvision.io
from src.config import VAE_TEMPORAL_FACTOR
def seed_everything(seed: int | None = None) -> None:
if seed is not None:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def save_video(tensor: torch.Tensor, save_path: str, fps: int = 30) -> None:
"""Save (T, 3, H, W) pixels as video, including one-frame clips."""
if tensor.ndim != 4 or tensor.shape[1] != 3 or min(tensor.shape) <= 0:
raise ValueError(f"Video must have shape (T, 3, H, W), got {tuple(tensor.shape)}.")
if fps <= 0:
raise ValueError("Video FPS must be positive.")
if tensor.is_floating_point():
tensor = tensor.mul(255).clamp(0, 255).to(torch.uint8)
elif tensor.dtype != torch.uint8:
raise ValueError("Video pixels must be uint8 or floating-point values in [0, 1].")
video_tensor = rearrange(tensor.cpu(), "t c h w -> t h w c")
torchvision.io.write_video(save_path, video_tensor, fps=fps, options={"crf": "8"})
def get_from_dict(data: dict, path: str):
"""Look up a dot-separated key path."""
if not isinstance(data, dict):
raise TypeError(f"Expected a dictionary, got {type(data).__name__}")
keys = path.split(".")
current_level = data
for key in keys:
current_level = current_level[key]
return current_level
def find_latest_checkpoint(ckpt_dir: Path) -> Path | None:
if not ckpt_dir.exists():
return None
step_pattern = re.compile(r"global_step(\d+)")
latest_step = -1
latest_ckpt = None
for subdir in ckpt_dir.iterdir():
if subdir.is_dir():
match = step_pattern.fullmatch(subdir.name)
if match:
step = int(match.group(1))
checkpoint = subdir / f"step_{step}.pth"
if step > latest_step and checkpoint.is_file() and checkpoint.stat().st_size > 0:
latest_step = step
latest_ckpt = subdir
return latest_ckpt.resolve() if latest_ckpt else None
def _resolve_full_length_target_frames(
total_frames: int,
source_fps: float | None,
target_fps: float | None,
) -> int:
"""Resample the source duration and floor to a valid 8n+1 frame count."""
if total_frames <= 0:
raise ValueError("A source video must contain at least one frame.")
if source_fps and target_fps and source_fps > 0 and target_fps > 0:
resampled = int(round((total_frames - 1) * target_fps / source_fps)) + 1
else:
resampled = total_frames
return max(1, ((resampled - 1) // VAE_TEMPORAL_FACTOR) * VAE_TEMPORAL_FACTOR + 1)
|