Spaces:
Running on Zero
Running on Zero
File size: 4,799 Bytes
294e8f2 | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | from omegaconf import OmegaConf, DictConfig
from typing import List, Tuple, Union
from PIL import Image
import numpy as np
from collections import OrderedDict
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder
from pathlib import Path
from copy import deepcopy
from torch.utils.data.distributed import DistributedSampler
def parse_configs(config: Union[DictConfig, str]) -> Tuple[DictConfig, DictConfig, DictConfig, DictConfig, DictConfig, DictConfig, DictConfig]:
"""Load a config file and return component sections as DictConfigs."""
if isinstance(config, str):
config = OmegaConf.load(config)
pr_config = config.get("stage_1", None)
stage2_config = config.get("stage_2", None)
transport_config = config.get("transport", None)
sampler_config = config.get("sampler", None)
guidance_config = config.get("guidance", None)
misc = config.get("misc", None)
training_config = config.get("training", None)
eval_config = config.get("eval", None)
return pr_config, stage2_config, transport_config, sampler_config, guidance_config, misc, training_config, eval_config
def none_or_str(value):
if value == 'None':
return None
return value
def center_crop_arr(pil_image, image_size):
"""
Center cropping implementation from ADM.
https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
"""
while min(*pil_image.size) >= 2 * image_size:
pil_image = pil_image.resize(
tuple(x // 2 for x in pil_image.size), resample=Image.BOX
)
scale = image_size / min(*pil_image.size)
pil_image = pil_image.resize(
tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
)
arr = np.array(pil_image)
crop_y = (arr.shape[0] - image_size) // 2
crop_x = (arr.shape[1] - image_size) // 2
return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
#################################################################################
# Training Helper Functions #
#################################################################################
def requires_grad(model, flag=True):
"""
Set requires_grad flag for all parameters in a model.
"""
for p in model.parameters():
p.requires_grad = flag
@torch.no_grad()
def update_ema(ema_model, model, decay=0.9999):
"""
Step the EMA model towards the current model.
"""
ema_params = OrderedDict(ema_model.named_parameters())
model_params = OrderedDict(model.named_parameters())
for name, param in model_params.items():
# TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
def prepare_dataloader(
data_path: Path,
batch_size: int,
workers: int,
transform: List= None,
) -> Tuple[DataLoader, None]:
"""
Prepare dataloader for training. Accelerator will handle distributed sampling automatically.
Returns loader and None (sampler is handled by Accelerator).
"""
dataset = ImageFolder(str(data_path), transform=transform)
loader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=True, # Accelerator will handle distributed sampling
num_workers=workers,
pin_memory=True,
drop_last=True,
)
return loader, None
def prepare_dataloader_distributed(
data_path: Path,
batch_size: int,
workers: int,
rank: int,
world_size: int,
transform: List= None,
worker_init_fn=None,
) -> Tuple[DataLoader, DistributedSampler]:
dataset = ImageFolder(str(data_path), transform=transform)
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
loader = DataLoader(
dataset,
batch_size=batch_size,
sampler=sampler,
num_workers=workers,
pin_memory=True,
drop_last=True,
worker_init_fn=worker_init_fn,
)
return loader, sampler
# Note: get_autocast_scaler is no longer needed with Accelerate
# Accelerator handles mixed precision automatically
def get_autocast_scaler(args) -> Tuple[dict, torch.cuda.amp.GradScaler | None]:
if args.precision == "fp16":
scaler = GradScaler()
autocast_kwargs = dict(enabled=True, dtype=torch.float16)
elif args.precision == "bf16":
scaler = None
autocast_kwargs = dict(enabled=True, dtype=torch.bfloat16)
else:
scaler = None
autocast_kwargs = dict(enabled=False)
return scaler, autocast_kwargs |