| import importlib |
| import os |
| import os.path as osp |
| import shutil |
| import sys |
| from pathlib import Path |
| import time |
| import av |
| import numpy as np |
| import torch |
| import torchvision |
| from einops import rearrange |
| from PIL import Image |
| import cv2 |
|
|
| def ensure_file_written(file_path, retries=10, delay=0.1): |
| for i in range(retries): |
| if os.path.exists(file_path): |
| image = Image.open(file_path) |
| if image is not None: |
| return True |
| else: |
| print(i) |
| time.sleep(delay) |
| return False |
|
|
| def compute_snr(noise_scheduler, timesteps): |
| """ |
| Computes SNR as per |
| https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849 |
| """ |
| alphas_cumprod = noise_scheduler.alphas_cumprod |
| sqrt_alphas_cumprod = alphas_cumprod**0.5 |
| sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5 |
|
|
| |
| |
| sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[ |
| timesteps |
| ].float() |
| while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape): |
| sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None] |
| alpha = sqrt_alphas_cumprod.expand(timesteps.shape) |
|
|
| sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to( |
| device=timesteps.device |
| )[timesteps].float() |
| while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape): |
| sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None] |
| sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape) |
|
|
| |
| snr = (alpha / sigma) ** 2 |
| return snr |
|
|
| def seed_everything(seed): |
| import random |
|
|
| import numpy as np |
|
|
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| np.random.seed(seed % (2**32)) |
| random.seed(seed) |
|
|
|
|
| def import_filename(filename): |
| spec = importlib.util.spec_from_file_location("mymodule", filename) |
| module = importlib.util.module_from_spec(spec) |
| sys.modules[spec.name] = module |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| def delete_additional_ckpt(base_path, num_keep): |
| dirs = [] |
| for d in os.listdir(base_path): |
| if d.startswith("checkpoint-"): |
| dirs.append(d) |
| num_tot = len(dirs) |
| if num_tot <= num_keep: |
| return |
| |
| del_dirs = sorted(dirs, key=lambda x: int(x.split("-")[-1]))[: num_tot - num_keep] |
| for d in del_dirs: |
| path_to_dir = osp.join(base_path, d) |
| if osp.exists(path_to_dir): |
| shutil.rmtree(path_to_dir) |
|
|
|
|
| def save_videos_from_pil(pil_images, path, fps=8): |
| import av |
|
|
| save_fmt = Path(path).suffix |
| os.makedirs(os.path.dirname(path), exist_ok=True) |
| width, height = pil_images[0].size |
|
|
| if save_fmt == ".mp4": |
| codec = "libx264" |
| container = av.open(path, "w") |
| stream = container.add_stream(codec, rate=fps) |
|
|
| stream.width = width |
| stream.height = height |
|
|
| for pil_image in pil_images: |
| |
| av_frame = av.VideoFrame.from_image(pil_image) |
| container.mux(stream.encode(av_frame)) |
| container.mux(stream.encode()) |
| container.close() |
|
|
| elif save_fmt == ".gif": |
| pil_images[0].save( |
| fp=path, |
| format="GIF", |
| append_images=pil_images[1:], |
| save_all=True, |
| duration=(1 / fps * 1000), |
| loop=0, |
| ) |
| else: |
| raise ValueError("Unsupported file type. Use .mp4 or .gif.") |
|
|
|
|
| def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8): |
| videos = rearrange(videos, "b c t h w -> t b c h w") |
| height, width = videos.shape[-2:] |
| outputs = [] |
|
|
| for x in videos: |
| x = torchvision.utils.make_grid(x, nrow=n_rows) |
| x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) |
| if rescale: |
| x = (x + 1.0) / 2.0 |
| x = (x * 255).numpy().astype(np.uint8) |
| x = Image.fromarray(x) |
|
|
| outputs.append(x) |
|
|
| os.makedirs(os.path.dirname(path), exist_ok=True) |
|
|
| save_videos_from_pil(outputs, path, fps) |
|
|
|
|
| def read_frames(video_path): |
| container = av.open(video_path) |
|
|
| video_stream = next(s for s in container.streams if s.type == "video") |
| frames = [] |
| for packet in container.demux(video_stream): |
| for frame in packet.decode(): |
| image = Image.frombytes( |
| "RGB", |
| (frame.width, frame.height), |
| frame.to_rgb().to_ndarray(), |
| ) |
| frames.append(image) |
|
|
| return frames |
|
|
|
|
| def get_fps(video_path): |
| container = av.open(video_path) |
| video_stream = next(s for s in container.streams if s.type == "video") |
| fps = video_stream.average_rate |
| container.close() |
| return fps |
|
|
