1341 lines
47 KiB
Python
1341 lines
47 KiB
Python
import torch
|
|
import json
|
|
import yaml
|
|
import torchvision
|
|
from torch import nn, optim
|
|
from transformers import (
|
|
AutoTokenizer,
|
|
CLIPTextModelWithProjection,
|
|
CLIPVisionModelWithProjection,
|
|
)
|
|
from warmup_scheduler import GradualWarmupScheduler
|
|
import torch.multiprocessing as mp
|
|
import os
|
|
import numpy as np
|
|
import re
|
|
import sys
|
|
|
|
sys.path.append(os.path.abspath("./"))
|
|
|
|
from dataclasses import dataclass
|
|
from torch.distributed import init_process_group, destroy_process_group, barrier
|
|
from ..gdf import GDF_dual_fixlrt as GDF
|
|
from ..gdf import EpsilonTarget, CosineSchedule
|
|
from ..gdf import (
|
|
VPScaler,
|
|
CosineTNoiseCond,
|
|
DDPMSampler,
|
|
P2LossWeight,
|
|
AdaptiveLossWeight,
|
|
)
|
|
from torchtools.transforms import SmartCrop
|
|
from fractions import Fraction
|
|
from ..modules.effnet import EfficientNetEncoder
|
|
from ..modules.model_4stage_lite import (
|
|
StageC,
|
|
ResBlock,
|
|
AttnBlock,
|
|
TimestepBlock,
|
|
FeedForwardBlock,
|
|
)
|
|
from ..modules.common_ckpt import GlobalResponseNorm
|
|
from ..modules.previewer import Previewer
|
|
from ..core.data import Bucketeer
|
|
from .base import DataCore, TrainingCore
|
|
from tqdm import tqdm
|
|
from ..core import WarpCore
|
|
from ..core.utils import EXPECTED, EXPECTED_TRAIN, load_or_fail
|
|
|
|
from accelerate import init_empty_weights
|
|
from accelerate.utils import set_module_tensor_to_device
|
|
from contextlib import contextmanager
|
|
from .dist_core import *
|
|
import glob
|
|
from torch.utils.data import DataLoader, Dataset
|
|
from torch.nn.parallel import DistributedDataParallel as DDP
|
|
from torch.utils.data.distributed import DistributedSampler
|
|
from PIL import Image
|
|
from ..core.utils import (
|
|
EXPECTED,
|
|
EXPECTED_TRAIN,
|
|
update_weights_ema,
|
|
create_folder_if_necessary,
|
|
)
|
|
from ..core.utils import Base
|
|
import torch.nn.functional as F
|
|
import functools
|
|
import math
|
|
import copy
|
|
import random
|
|
from ..modules.lora import apply_lora, apply_retoken, LoRA, ReToken
|
|
|
|
Image.MAX_IMAGE_PIXELS = None
|
|
torch.manual_seed(23)
|
|
random.seed(23)
|
|
np.random.seed(23)
|
|
# 7978026
|
|
|
|
|
|
class Null_Model(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x):
|
|
pass
|
|
|
|
|
|
def identity(x):
|
|
if isinstance(x, bytes):
|
|
x = x.decode("utf-8")
|
|
return x
|
|
|
|
|
|
def check_nan_inmodel(model, meta=""):
|
|
for name, param in model.named_parameters():
|
|
if torch.isnan(param).any():
|
|
print(f"nan detected in {name}", meta)
|
|
return True
|
|
# print("no nan", meta)
|
|
return False
|
|
|
|
|
|
class mydist_dataset(Dataset):
|
|
def __init__(self, rootpath, tmp_prompt, img_processor=None):
|
|
|
|
self.img_pathlist = glob.glob(os.path.join(rootpath, "*.jpg"))
|
|
self.img_pathlist = self.img_pathlist * 100000
|
|
self.img_processor = img_processor
|
|
self.length = len(self.img_pathlist)
|
|
self.caption = tmp_prompt
|
|
|
|
def __getitem__(self, idx):
|
|
|
|
imgpath = self.img_pathlist[idx]
|
|
txt = self.caption
|
|
|
|
try:
|
|
img = Image.open(imgpath).convert("RGB")
|
|
w, h = img.size
|
|
if self.img_processor is not None:
|
|
img = self.img_processor(img)
|
|
|
|
except:
|
|
print("exception", imgpath)
|
|
return self.__getitem__(random.randint(0, self.length - 1))
|
|
return dict(captions=txt, images=img)
|
|
|
|
def __len__(self):
|
|
return self.length
|
|
|
|
|
|
class WurstCore(TrainingCore, DataCore, WarpCore):
|
|
@dataclass(frozen=True)
|
|
class Config(TrainingCore.Config, DataCore.Config, WarpCore.Config):
|
|
# TRAINING PARAMS
|
|
lr: float = EXPECTED_TRAIN
|
|
warmup_updates: int = EXPECTED_TRAIN
|
|
dtype: str = None
|
|
|
|
# MODEL VERSION
|
|
model_version: str = EXPECTED # 3.6B or 1B
|
|
clip_image_model_name: str = "openai/clip-vit-large-patch14"
|
|
clip_text_model_name: str = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
|
|
|
|
# CHECKPOINT PATHS
|
|
effnet_checkpoint_path: str = EXPECTED
|
|
previewer_checkpoint_path: str = EXPECTED
|
|
generator_checkpoint_path: str = None
|
|
ultrapixel_path: str = EXPECTED
|
|
|
|
# gdf customization
|
|
adaptive_loss_weight: str = None
|
|
|
|
# LoRA STUFF
|
|
module_filters: list = EXPECTED
|
|
rank: int = EXPECTED
|
|
train_tokens: list = EXPECTED
|
|
use_ddp: bool = EXPECTED
|
|
tmp_prompt: str = EXPECTED
|
|
|
|
@dataclass(frozen=True)
|
|
class Data(Base):
|
|
dataset: Dataset = EXPECTED
|
|
dataloader: DataLoader = EXPECTED
|
|
iterator: any = EXPECTED
|
|
sampler: DistributedSampler = EXPECTED
|
|
|
|
@dataclass(frozen=True)
|
|
class Models(TrainingCore.Models, DataCore.Models, WarpCore.Models):
|
|
effnet: nn.Module = EXPECTED
|
|
previewer: nn.Module = EXPECTED
|
|
train_norm: nn.Module = EXPECTED
|
|
train_lora: nn.Module = EXPECTED
|
|
|
|
@dataclass(frozen=True)
|
|
class Schedulers(WarpCore.Schedulers):
|
|
generator: any = None
|
|
|
|
@dataclass(frozen=True)
|
|
class Extras(TrainingCore.Extras, DataCore.Extras, WarpCore.Extras):
|
|
gdf: GDF = EXPECTED
|
|
sampling_configs: dict = EXPECTED
|
|
effnet_preprocess: torchvision.transforms.Compose = EXPECTED
|
|
|
|
info: TrainingCore.Info
|
|
config: Config
|
|
|
|
def setup_extras_pre(self) -> Extras:
|
|
gdf = GDF(
|
|
schedule=CosineSchedule(clamp_range=[0.0001, 0.9999]),
|
|
input_scaler=VPScaler(),
|
|
target=EpsilonTarget(),
|
|
noise_cond=CosineTNoiseCond(),
|
|
loss_weight=(
|
|
AdaptiveLossWeight()
|
|
if self.config.adaptive_loss_weight is True
|
|
else P2LossWeight()
|
|
),
|
|
)
|
|
sampling_configs = {
|
|
"cfg": 5,
|
|
"sampler": DDPMSampler(gdf),
|
|
"shift": 1,
|
|
"timesteps": 20,
|
|
}
|
|
|
|
if self.info.adaptive_loss is not None:
|
|
gdf.loss_weight.bucket_ranges = torch.tensor(
|
|
self.info.adaptive_loss["bucket_ranges"]
|
|
)
|
|
gdf.loss_weight.bucket_losses = torch.tensor(
|
|
self.info.adaptive_loss["bucket_losses"]
|
|
)
|
|
|
|
effnet_preprocess = torchvision.transforms.Compose(
|
|
[
|
|
torchvision.transforms.Normalize(
|
|
mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)
|
|
)
|
|
]
|
|
)
|
|
|
|
clip_preprocess = torchvision.transforms.Compose(
|
|
[
|
|
torchvision.transforms.Resize(
|
|
224, interpolation=torchvision.transforms.InterpolationMode.BICUBIC
|
|
),
|
|
torchvision.transforms.CenterCrop(224),
|
|
torchvision.transforms.Normalize(
|
|
mean=(0.48145466, 0.4578275, 0.40821073),
|
|
std=(0.26862954, 0.26130258, 0.27577711),
|
|
),
|
|
]
|
|
)
|
|
|
|
if self.config.training:
|
|
transforms = torchvision.transforms.Compose(
|
|
[
|
|
torchvision.transforms.ToTensor(),
|
|
torchvision.transforms.Resize(
|
|
self.config.image_size[-1],
|
|
interpolation=torchvision.transforms.InterpolationMode.BILINEAR,
|
|
antialias=True,
|
|
),
|
|
SmartCrop(self.config.image_size, randomize_p=0.3, randomize_q=0.2),
|
|
]
|
|
)
|
|
else:
|
|
transforms = None
|
|
|
|
return self.Extras(
|
|
gdf=gdf,
|
|
sampling_configs=sampling_configs,
|
|
transforms=transforms,
|
|
effnet_preprocess=effnet_preprocess,
|
|
clip_preprocess=clip_preprocess,
|
|
)
|
|
|
|
def get_conditions(
|
|
self,
|
|
batch: dict,
|
|
models: Models,
|
|
extras: Extras,
|
|
is_eval=False,
|
|
is_unconditional=False,
|
|
eval_image_embeds=False,
|
|
return_fields=None,
|
|
):
|
|
conditions = super().get_conditions(
|
|
batch,
|
|
models,
|
|
extras,
|
|
is_eval,
|
|
is_unconditional,
|
|
eval_image_embeds,
|
|
return_fields=return_fields
|
|
or ["clip_text", "clip_text_pooled", "clip_img"],
|
|
)
|
|
return conditions
|
|
|
|
def setup_models(self, extras: Extras) -> Models: # configure model
|
|
|
|
dtype = (
|
|
getattr(torch, self.config.dtype) if self.config.dtype else torch.bfloat16
|
|
)
|
|
|
|
# EfficientNet encoderin
|
|
effnet = EfficientNetEncoder()
|
|
effnet_checkpoint = load_or_fail(self.config.effnet_checkpoint_path)
|
|
effnet.load_state_dict(
|
|
effnet_checkpoint
|
|
if "state_dict" not in effnet_checkpoint
|
|
else effnet_checkpoint["state_dict"]
|
|
)
|
|
effnet.eval().requires_grad_(False).to(self.device)
|
|
del effnet_checkpoint
|
|
|
|
# Previewer
|
|
previewer = Previewer()
|
|
previewer_checkpoint = load_or_fail(self.config.previewer_checkpoint_path)
|
|
previewer.load_state_dict(
|
|
previewer_checkpoint
|
|
if "state_dict" not in previewer_checkpoint
|
|
else previewer_checkpoint["state_dict"]
|
|
)
|
|
previewer.eval().requires_grad_(False).to(self.device)
|
|
del previewer_checkpoint
|
|
|
|
@contextmanager
|
|
def dummy_context():
|
|
yield None
|
|
|
|
loading_context = dummy_context if self.config.training else init_empty_weights
|
|
|
|
# Diffusion models
|
|
with loading_context():
|
|
generator_ema = None
|
|
if self.config.model_version == "3.6B":
|
|
generator = StageC()
|
|
if self.config.ema_start_iters is not None: # default setting
|
|
generator_ema = StageC()
|
|
elif self.config.model_version == "1B":
|
|
print("in line 155 1b light model", self.config.model_version)
|
|
generator = StageC(
|
|
c_cond=1536,
|
|
c_hidden=[1536, 1536],
|
|
nhead=[24, 24],
|
|
blocks=[[4, 12], [12, 4]],
|
|
)
|
|
|
|
if self.config.ema_start_iters is not None and self.config.training:
|
|
generator_ema = StageC(
|
|
c_cond=1536,
|
|
c_hidden=[1536, 1536],
|
|
nhead=[24, 24],
|
|
blocks=[[4, 12], [12, 4]],
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown model version {self.config.model_version}")
|
|
|
|
if loading_context is dummy_context:
|
|
generator.load_state_dict(
|
|
load_or_fail(self.config.generator_checkpoint_path)
|
|
)
|
|
else:
|
|
for param_name, param in load_or_fail(
|
|
self.config.generator_checkpoint_path
|
|
).items():
|
|
set_module_tensor_to_device(generator, param_name, "cpu", value=param)
|
|
|
|
generator._init_extra_parameter()
|
|
generator = generator.to(torch.bfloat16).to(self.device)
|
|
|
|
train_norm = nn.ModuleList()
|
|
|
|
cnt_norm = 0
|
|
for mm in generator.modules():
|
|
if isinstance(mm, GlobalResponseNorm):
|
|
|
|
train_norm.append(Null_Model())
|
|
cnt_norm += 1
|
|
|
|
train_norm.append(generator.agg_net)
|
|
train_norm.append(generator.agg_net_up)
|
|
sdd = torch.load(self.config.ultrapixel_path, map_location="cpu")
|
|
collect_sd = {}
|
|
for k, v in sdd.items():
|
|
collect_sd[k[7:]] = v
|
|
train_norm.load_state_dict(collect_sd)
|
|
|
|
# CLIP encoders
|
|
tokenizer = AutoTokenizer.from_pretrained(self.config.clip_text_model_name)
|
|
text_model = (
|
|
CLIPTextModelWithProjection.from_pretrained(
|
|
self.config.clip_text_model_name
|
|
)
|
|
.requires_grad_(False)
|
|
.to(dtype)
|
|
.to(self.device)
|
|
)
|
|
image_model = (
|
|
CLIPVisionModelWithProjection.from_pretrained(
|
|
self.config.clip_image_model_name
|
|
)
|
|
.requires_grad_(False)
|
|
.to(dtype)
|
|
.to(self.device)
|
|
)
|
|
|
|
# PREPARE LORA
|
|
train_lora = nn.ModuleList()
|
|
update_tokens = []
|
|
for tkn_regex, aggr_regex in self.config.train_tokens:
|
|
if (tkn_regex.startswith("[") and tkn_regex.endswith("]")) or (
|
|
tkn_regex.startswith("<") and tkn_regex.endswith(">")
|
|
):
|
|
# Insert new token
|
|
tokenizer.add_tokens([tkn_regex])
|
|
# add new zeros embedding
|
|
new_embedding = torch.zeros_like(
|
|
text_model.text_model.embeddings.token_embedding.weight.data
|
|
)[:1]
|
|
if (
|
|
aggr_regex is not None
|
|
): # aggregate embeddings to provide an interesting baseline
|
|
aggr_tokens = [
|
|
v
|
|
for k, v in tokenizer.vocab.items()
|
|
if re.search(aggr_regex, k) is not None
|
|
]
|
|
if len(aggr_tokens) > 0:
|
|
new_embedding = text_model.text_model.embeddings.token_embedding.weight.data[
|
|
aggr_tokens
|
|
].mean(
|
|
dim=0, keepdim=True
|
|
)
|
|
elif self.is_main_node:
|
|
print(
|
|
f"WARNING: No tokens found for aggregation regex {aggr_regex}. It will be initialized as zeros."
|
|
)
|
|
text_model.text_model.embeddings.token_embedding.weight.data = torch.cat(
|
|
[
|
|
text_model.text_model.embeddings.token_embedding.weight.data,
|
|
new_embedding,
|
|
],
|
|
dim=0,
|
|
)
|
|
selected_tokens = [len(tokenizer.vocab) - 1]
|
|
else:
|
|
selected_tokens = [
|
|
v
|
|
for k, v in tokenizer.vocab.items()
|
|
if re.search(tkn_regex, k) is not None
|
|
]
|
|
update_tokens += selected_tokens
|
|
update_tokens = list(set(update_tokens)) # remove duplicates
|
|
|
|
apply_retoken(text_model.text_model.embeddings.token_embedding, update_tokens)
|
|
|
|
apply_lora(generator, filters=self.config.module_filters, rank=self.config.rank)
|
|
for module in generator.modules():
|
|
if isinstance(module, LoRA) or (
|
|
hasattr(module, "_fsdp_wrapped_module")
|
|
and isinstance(module._fsdp_wrapped_module, LoRA)
|
|
):
|
|
train_lora.append(module)
|
|
|
|
train_lora.append(
|
|
text_model.text_model.embeddings.token_embedding.parametrizations.weight[0]
|
|
)
|
|
|
|
if os.path.exists(
|
|
os.path.join(
|
|
self.config.output_path,
|
|
self.config.experiment_id,
|
|
"train_lora.safetensors",
|
|
)
|
|
):
|
|
sdd = torch.load(
|
|
os.path.join(
|
|
self.config.output_path,
|
|
self.config.experiment_id,
|
|
"train_lora.safetensors",
|
|
),
|
|
map_location="cpu",
|
|
)
|
|
collect_sd = {}
|
|
for k, v in sdd.items():
|
|
collect_sd[k[7:]] = v
|
|
train_lora.load_state_dict(collect_sd, strict=True)
|
|
|
|
train_norm.to(self.device).train().requires_grad_(True)
|
|
|
|
if generator_ema is not None:
|
|
|
|
generator_ema.load_state_dict(
|
|
load_or_fail(self.config.generator_checkpoint_path)
|
|
)
|
|
generator_ema._init_extra_parameter()
|
|
pretrained_pth = os.path.join(
|
|
self.config.output_path,
|
|
self.config.experiment_id,
|
|
"generator.safetensors",
|
|
)
|
|
if os.path.exists(pretrained_pth):
|
|
generator_ema.load_state_dict(
|
|
torch.load(pretrained_pth, map_location="cpu")
|
|
)
|
|
|
|
generator_ema.eval().requires_grad_(False)
|
|
|
|
check_nan_inmodel(generator, "generator")
|
|
|
|
if self.config.use_ddp and self.config.training:
|
|
|
|
train_lora = DDP(
|
|
train_lora, device_ids=[self.device], find_unused_parameters=True
|
|
)
|
|
|
|
return self.Models(
|
|
effnet=effnet,
|
|
previewer=previewer,
|
|
train_norm=train_norm,
|
|
generator=generator,
|
|
generator_ema=generator_ema,
|
|
tokenizer=tokenizer,
|
|
text_model=text_model,
|
|
image_model=image_model,
|
|
train_lora=train_lora,
|
|
)
|
|
|
|
def setup_optimizers(
|
|
self, extras: Extras, models: Models
|
|
) -> TrainingCore.Optimizers:
|
|
|
|
params = []
|
|
params += list(models.train_lora.module.parameters())
|
|
optimizer = optim.AdamW(params, lr=self.config.lr)
|
|
|
|
return self.Optimizers(generator=optimizer)
|
|
|
|
def ema_update(self, ema_model, source_model, beta):
|
|
for param_src, param_ema in zip(
|
|
source_model.parameters(), ema_model.parameters()
|
|
):
|
|
param_ema.data.mul_(beta).add_(param_src.data, alpha=1 - beta)
|
|
|
|
def sync_ema(self, ema_model):
|
|
print("sync ema", torch.distributed.get_world_size())
|
|
for param in ema_model.parameters():
|
|
torch.distributed.all_reduce(param.data, op=torch.distributed.ReduceOp.SUM)
|
|
param.data /= torch.distributed.get_world_size()
|
|
|
|
def setup_optimizers_backup(
|
|
self, extras: Extras, models: Models
|
|
) -> TrainingCore.Optimizers:
|
|
|
|
optimizer = optim.AdamW(
|
|
models.generator.up_blocks.parameters(), lr=self.config.lr
|
|
)
|
|
optimizer = self.load_optimizer(
|
|
optimizer,
|
|
"generator_optim",
|
|
fsdp_model=models.generator if self.config.use_fsdp else None,
|
|
)
|
|
return self.Optimizers(generator=optimizer)
|
|
|
|
def setup_schedulers(
|
|
self, extras: Extras, models: Models, optimizers: TrainingCore.Optimizers
|
|
) -> Schedulers:
|
|
scheduler = GradualWarmupScheduler(
|
|
optimizers.generator, multiplier=1, total_epoch=self.config.warmup_updates
|
|
)
|
|
scheduler.last_epoch = self.info.total_steps
|
|
return self.Schedulers(generator=scheduler)
|
|
|
|
def setup_data(self, extras: Extras) -> WarpCore.Data:
|
|
# SETUP DATASET
|
|
dataset_path = self.config.webdataset_path
|
|
|
|
dataset = mydist_dataset(
|
|
dataset_path,
|
|
self.config.tmp_prompt,
|
|
(
|
|
torchvision.transforms.ToTensor()
|
|
if self.config.multi_aspect_ratio is not None
|
|
else extras.transforms
|
|
),
|
|
)
|
|
|
|
# SETUP DATALOADER
|
|
real_batch_size = self.config.batch_size // (
|
|
self.world_size * self.config.grad_accum_steps
|
|
)
|
|
|
|
sampler = DistributedSampler(
|
|
dataset, rank=self.process_id, num_replicas=self.world_size, shuffle=True
|
|
)
|
|
dataloader = DataLoader(
|
|
dataset,
|
|
batch_size=real_batch_size,
|
|
num_workers=4,
|
|
pin_memory=True,
|
|
collate_fn=identity if self.config.multi_aspect_ratio is not None else None,
|
|
sampler=sampler,
|
|
)
|
|
if self.is_main_node:
|
|
print(
|
|
f"Training with batch size {self.config.batch_size} ({real_batch_size}/GPU)"
|
|
)
|
|
|
|
if self.config.multi_aspect_ratio is not None:
|
|
aspect_ratios = [float(Fraction(f)) for f in self.config.multi_aspect_ratio]
|
|
dataloader_iterator = Bucketeer(
|
|
dataloader,
|
|
density=[ss * ss for ss in self.config.image_size],
|
|
factor=32,
|
|
ratios=aspect_ratios,
|
|
p_random_ratio=self.config.bucketeer_random_ratio,
|
|
interpolate_nearest=False,
|
|
) # , use_smartcrop=True)
|
|
else:
|
|
|
|
dataloader_iterator = iter(dataloader)
|
|
|
|
return self.Data(
|
|
dataset=dataset,
|
|
dataloader=dataloader,
|
|
iterator=dataloader_iterator,
|
|
sampler=sampler,
|
|
)
|
|
|
|
def setup_ddp(self, experiment_id, single_gpu=False, rank=0):
|
|
|
|
if not single_gpu:
|
|
local_rank = rank
|
|
process_id = rank
|
|
world_size = get_world_size()
|
|
|
|
self.process_id = process_id
|
|
self.is_main_node = process_id == 0
|
|
self.device = torch.device(local_rank)
|
|
self.world_size = world_size
|
|
|
|
os.environ["MASTER_ADDR"] = "localhost"
|
|
os.environ["MASTER_PORT"] = "14443"
|
|
torch.cuda.set_device(local_rank)
|
|
init_process_group(
|
|
backend="nccl",
|
|
rank=local_rank,
|
|
world_size=world_size,
|
|
# init_method=init_method,
|
|
)
|
|
print(f"[GPU {process_id}] READY")
|
|
else:
|
|
self.is_main_node = rank == 0
|
|
self.process_id = rank
|
|
self.device = torch.device("cuda:0")
|
|
self.world_size = 1
|
|
# print("Running in single thread, DDP not enabled.")
|
|
|
|
# Training loop --------------------------------
|
|
def get_target_lr_size(self, ratio, std_size=24):
|
|
w, h = int(std_size / math.sqrt(ratio)), int(std_size * math.sqrt(ratio))
|
|
return (h * 32, w * 32)
|
|
|
|
def forward_pass(self, data: WarpCore.Data, extras: Extras, models: Models):
|
|
|
|
batch = data
|
|
ratio = batch["images"].shape[-2] / batch["images"].shape[-1]
|
|
shape_lr = self.get_target_lr_size(ratio)
|
|
with torch.no_grad():
|
|
conditions = self.get_conditions(batch, models, extras)
|
|
|
|
latents = self.encode_latents(batch, models, extras)
|
|
latents_lr = self.encode_latents(
|
|
batch, models, extras, target_size=shape_lr
|
|
)
|
|
|
|
flag_lr = random.random() < 0.5 or self.info.iter < 5000
|
|
|
|
if flag_lr:
|
|
noised, noise, target, logSNR, noise_cond, loss_weight = (
|
|
extras.gdf.diffuse(latents_lr, shift=1, loss_shift=1)
|
|
)
|
|
else:
|
|
noised, noise, target, logSNR, noise_cond, loss_weight = (
|
|
extras.gdf.diffuse(latents, shift=1, loss_shift=1)
|
|
)
|
|
if not flag_lr:
|
|
(
|
|
noised_lr,
|
|
noise_lr,
|
|
target_lr,
|
|
logSNR_lr,
|
|
noise_cond_lr,
|
|
loss_weight_lr,
|
|
) = extras.gdf.diffuse(
|
|
latents_lr,
|
|
shift=1,
|
|
loss_shift=1,
|
|
t=torch.ones(latents.shape[0]).to(latents.device) * 0.05,
|
|
)
|
|
|
|
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
|
|
|
|
if not flag_lr:
|
|
with torch.no_grad():
|
|
_, lr_enc_guide, lr_dec_guide = models.generator(
|
|
noised_lr, noise_cond_lr, reuire_f=True, **conditions
|
|
)
|
|
|
|
pred = models.generator(
|
|
noised,
|
|
noise_cond,
|
|
reuire_f=False,
|
|
lr_guide=(lr_enc_guide, lr_dec_guide) if not flag_lr else None,
|
|
**conditions,
|
|
)
|
|
loss = nn.functional.mse_loss(pred, target, reduction="none").mean(
|
|
dim=[1, 2, 3]
|
|
)
|
|
|
|
loss_adjusted = (loss * loss_weight).mean() / self.config.grad_accum_steps
|
|
|
|
if isinstance(extras.gdf.loss_weight, AdaptiveLossWeight):
|
|
extras.gdf.loss_weight.update_buckets(logSNR, loss)
|
|
return loss, loss_adjusted
|
|
|
|
def backward_pass(
|
|
self,
|
|
update,
|
|
loss_adjusted,
|
|
models: Models,
|
|
optimizers: TrainingCore.Optimizers,
|
|
schedulers: Schedulers,
|
|
):
|
|
|
|
if update:
|
|
|
|
torch.distributed.barrier()
|
|
loss_adjusted.backward()
|
|
|
|
grad_norm = nn.utils.clip_grad_norm_(
|
|
models.train_lora.module.parameters(), 1.0
|
|
)
|
|
optimizers_dict = optimizers.to_dict()
|
|
for k in optimizers_dict:
|
|
if k != "training":
|
|
optimizers_dict[k].step()
|
|
schedulers_dict = schedulers.to_dict()
|
|
for k in schedulers_dict:
|
|
if k != "training":
|
|
schedulers_dict[k].step()
|
|
for k in optimizers_dict:
|
|
if k != "training":
|
|
optimizers_dict[k].zero_grad(set_to_none=True)
|
|
self.info.total_steps += 1
|
|
else:
|
|
|
|
loss_adjusted.backward()
|
|
grad_norm = torch.tensor(0.0).to(self.device)
|
|
|
|
return grad_norm
|
|
|
|
def models_to_save(self):
|
|
return ["generator", "generator_ema", "trans_inr", "trans_inr_ema"]
|
|
|
|
def encode_latents(
|
|
self, batch: dict, models: Models, extras: Extras, target_size=None
|
|
) -> torch.Tensor:
|
|
|
|
images = batch["images"].to(self.device)
|
|
if target_size is not None:
|
|
images = F.interpolate(images, target_size)
|
|
|
|
return models.effnet(extras.effnet_preprocess(images))
|
|
|
|
def decode_latents(
|
|
self, latents: torch.Tensor, batch: dict, models: Models, extras: Extras
|
|
) -> torch.Tensor:
|
|
return models.previewer(latents)
|
|
|
|
def __init__(
|
|
self,
|
|
rank=0,
|
|
config_file_path=None,
|
|
config_dict=None,
|
|
device="cpu",
|
|
training=True,
|
|
world_size=1,
|
|
):
|
|
|
|
self.is_main_node = rank == 0
|
|
self.config: self.Config = self.setup_config(
|
|
config_file_path, config_dict, training
|
|
)
|
|
self.setup_ddp(self.config.experiment_id, single_gpu=world_size <= 1, rank=rank)
|
|
self.info: self.Info = self.setup_info()
|
|
print("in line 292", self.config.experiment_id, rank, world_size <= 1)
|
|
p = [i for i in range(2 * 768 // 32)]
|
|
p = [num / sum(p) for num in p]
|
|
self.rand_pro = p
|
|
self.res_list = [o for o in range(800, 2336, 32)]
|
|
|
|
def __call__(self, single_gpu=False):
|
|
|
|
if self.config.allow_tf32:
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
|
|
if self.is_main_node:
|
|
print()
|
|
print("**STARTIG JOB WITH CONFIG:**")
|
|
print(yaml.dump(self.config.to_dict(), default_flow_style=False))
|
|
print("------------------------------------")
|
|
print()
|
|
print("**INFO:**")
|
|
print(yaml.dump(vars(self.info), default_flow_style=False))
|
|
print("------------------------------------")
|
|
print()
|
|
print(
|
|
"in line 308",
|
|
self.is_main_node,
|
|
self.is_main_node,
|
|
self.process_id,
|
|
self.device,
|
|
)
|
|
# SETUP STUFF
|
|
extras = self.setup_extras_pre()
|
|
assert extras is not None, "setup_extras_pre() must return a DTO"
|
|
|
|
data = self.setup_data(extras)
|
|
assert data is not None, "setup_data() must return a DTO"
|
|
if self.is_main_node:
|
|
print("**DATA:**")
|
|
print(
|
|
yaml.dump(
|
|
{k: type(v).__name__ for k, v in data.to_dict().items()},
|
|
default_flow_style=False,
|
|
)
|
|
)
|
|
print("------------------------------------")
|
|
print()
|
|
|
|
models = self.setup_models(extras)
|
|
assert models is not None, "setup_models() must return a DTO"
|
|
if self.is_main_node:
|
|
print("**MODELS:**")
|
|
print(
|
|
yaml.dump(
|
|
{
|
|
k: f"{type(v).__name__} - {f'trainable params {sum(p.numel() for p in v.parameters() if p.requires_grad)}' if isinstance(v, nn.Module) else 'Not a nn.Module'}"
|
|
for k, v in models.to_dict().items()
|
|
},
|
|
default_flow_style=False,
|
|
)
|
|
)
|
|
print("------------------------------------")
|
|
print()
|
|
|
|
optimizers = self.setup_optimizers(extras, models)
|
|
assert optimizers is not None, "setup_optimizers() must return a DTO"
|
|
if self.is_main_node:
|
|
print("**OPTIMIZERS:**")
|
|
print(
|
|
yaml.dump(
|
|
{k: type(v).__name__ for k, v in optimizers.to_dict().items()},
|
|
default_flow_style=False,
|
|
)
|
|
)
|
|
print("------------------------------------")
|
|
print()
|
|
|
|
schedulers = self.setup_schedulers(extras, models, optimizers)
|
|
assert schedulers is not None, "setup_schedulers() must return a DTO"
|
|
if self.is_main_node:
|
|
print("**SCHEDULERS:**")
|
|
print(
|
|
yaml.dump(
|
|
{k: type(v).__name__ for k, v in schedulers.to_dict().items()},
|
|
default_flow_style=False,
|
|
)
|
|
)
|
|
print("------------------------------------")
|
|
print()
|
|
|
|
post_extras = self.setup_extras_post(extras, models, optimizers, schedulers)
|
|
assert post_extras is not None, "setup_extras_post() must return a DTO"
|
|
extras = self.Extras.from_dict({**extras.to_dict(), **post_extras.to_dict()})
|
|
if self.is_main_node:
|
|
print("**EXTRAS:**")
|
|
print(
|
|
yaml.dump(
|
|
{k: f"{v}" for k, v in extras.to_dict().items()},
|
|
default_flow_style=False,
|
|
)
|
|
)
|
|
print("------------------------------------")
|
|
print()
|
|
# -------
|
|
|
|
# TRAIN
|
|
if self.is_main_node:
|
|
print("**TRAINING STARTING...**")
|
|
self.train(data, extras, models, optimizers, schedulers)
|
|
|
|
if single_gpu is False:
|
|
barrier()
|
|
destroy_process_group()
|
|
if self.is_main_node:
|
|
print()
|
|
print("------------------------------------")
|
|
print()
|
|
print("**TRAINING COMPLETE**")
|
|
if self.config.wandb_project is not None:
|
|
wandb.alert(
|
|
title=f"Training {self.info.wandb_run_id} finished",
|
|
text=f"Training {self.info.wandb_run_id} finished",
|
|
)
|
|
|
|
def train(
|
|
self,
|
|
data: WarpCore.Data,
|
|
extras: WarpCore.Extras,
|
|
models: Models,
|
|
optimizers: TrainingCore.Optimizers,
|
|
schedulers: WarpCore.Schedulers,
|
|
):
|
|
start_iter = self.info.iter + 1
|
|
max_iters = self.config.updates * self.config.grad_accum_steps
|
|
if self.is_main_node:
|
|
print(f"STARTING AT STEP: {start_iter}/{max_iters}")
|
|
|
|
if self.is_main_node:
|
|
create_folder_if_necessary(
|
|
f"{self.config.output_path}/{self.config.experiment_id}/"
|
|
)
|
|
if "generator" in self.models_to_save():
|
|
models.generator.train()
|
|
|
|
iter_cnt = 0
|
|
epoch_cnt = 0
|
|
models.train_norm.train()
|
|
while True:
|
|
epoch_cnt += 1
|
|
if self.world_size > 1:
|
|
|
|
data.sampler.set_epoch(epoch_cnt)
|
|
for ggg in range(len(data.dataloader)):
|
|
iter_cnt += 1
|
|
# FORWARD PASS
|
|
|
|
loss, loss_adjusted = self.forward_pass(
|
|
next(data.iterator), extras, models
|
|
)
|
|
|
|
# # BACKWARD PASS
|
|
|
|
grad_norm = self.backward_pass(
|
|
iter_cnt % self.config.grad_accum_steps == 0
|
|
or iter_cnt == max_iters,
|
|
loss_adjusted,
|
|
models,
|
|
optimizers,
|
|
schedulers,
|
|
)
|
|
|
|
self.info.iter = iter_cnt
|
|
|
|
self.info.ema_loss = (
|
|
loss.mean().item()
|
|
if self.info.ema_loss is None
|
|
else self.info.ema_loss * 0.99 + loss.mean().item() * 0.01
|
|
)
|
|
|
|
if (
|
|
self.is_main_node
|
|
and np.isnan(loss.mean().item())
|
|
or np.isnan(grad_norm.item())
|
|
):
|
|
print(
|
|
f"gggg NaN value encountered in training run {self.info.wandb_run_id}",
|
|
f"Loss {loss.mean().item()} - Grad Norm {grad_norm.item()}. Run {self.info.wandb_run_id}",
|
|
)
|
|
|
|
if self.is_main_node:
|
|
logs = {
|
|
"loss": self.info.ema_loss,
|
|
"backward_loss": loss_adjusted.mean().item(),
|
|
"ema_loss": self.info.ema_loss,
|
|
"raw_ori_loss": loss.mean().item(),
|
|
"grad_norm": grad_norm.item(),
|
|
"lr": (
|
|
optimizers.generator.param_groups[0]["lr"]
|
|
if optimizers.generator is not None
|
|
else 0
|
|
),
|
|
"total_steps": self.info.total_steps,
|
|
}
|
|
|
|
print(
|
|
iter_cnt,
|
|
max_iters,
|
|
logs,
|
|
epoch_cnt,
|
|
)
|
|
|
|
if (
|
|
iter_cnt == 1
|
|
or iter_cnt % (self.config.save_every) == 0
|
|
or iter_cnt == max_iters
|
|
):
|
|
|
|
if np.isnan(loss.mean().item()):
|
|
if self.is_main_node and self.config.wandb_project is not None:
|
|
print(
|
|
f"NaN value encountered in training run {self.info.wandb_run_id}",
|
|
f"Loss {loss.mean().item()} - Grad Norm {grad_norm.item()}. Run {self.info.wandb_run_id}",
|
|
)
|
|
|
|
else:
|
|
if isinstance(extras.gdf.loss_weight, AdaptiveLossWeight):
|
|
self.info.adaptive_loss = {
|
|
"bucket_ranges": extras.gdf.loss_weight.bucket_ranges.tolist(),
|
|
"bucket_losses": extras.gdf.loss_weight.bucket_losses.tolist(),
|
|
}
|
|
|
|
if (
|
|
self.is_main_node
|
|
and iter_cnt
|
|
% (self.config.save_every * self.config.grad_accum_steps)
|
|
== 0
|
|
):
|
|
print(
|
|
"save model",
|
|
iter_cnt,
|
|
iter_cnt
|
|
% (
|
|
self.config.save_every
|
|
* self.config.grad_accum_steps
|
|
),
|
|
self.config.save_every,
|
|
self.config.grad_accum_steps,
|
|
)
|
|
torch.save(
|
|
models.train_lora.state_dict(),
|
|
f"{self.config.output_path}/{self.config.experiment_id}/train_lora.safetensors",
|
|
)
|
|
|
|
torch.save(
|
|
models.train_lora.state_dict(),
|
|
f"{self.config.output_path}/{self.config.experiment_id}/train_lora_{iter_cnt}.safetensors",
|
|
)
|
|
|
|
if (
|
|
iter_cnt == 1
|
|
or iter_cnt
|
|
% (self.config.save_every * self.config.grad_accum_steps)
|
|
== 0
|
|
or iter_cnt == max_iters
|
|
):
|
|
|
|
if self.is_main_node:
|
|
|
|
self.sample(models, data, extras)
|
|
if False:
|
|
param_changes = {
|
|
name: (param - initial_params[name]).norm().item()
|
|
for name, param in models.train_norm.named_parameters()
|
|
}
|
|
threshold = sorted(param_changes.values(), reverse=True)[
|
|
int(len(param_changes) * 0.1)
|
|
] # top 10%
|
|
important_params = [
|
|
name
|
|
for name, change in param_changes.items()
|
|
if change > threshold
|
|
]
|
|
print(
|
|
important_params, threshold, len(param_changes), self.process_id
|
|
)
|
|
json.dump(
|
|
important_params,
|
|
open(
|
|
f"{self.config.output_path}/{self.config.experiment_id}/param.json",
|
|
"w",
|
|
),
|
|
indent=4,
|
|
)
|
|
|
|
if self.info.iter >= max_iters:
|
|
break
|
|
|
|
def sample(self, models: Models, data: WarpCore.Data, extras: Extras):
|
|
|
|
models.generator.eval()
|
|
models.train_norm.eval()
|
|
with torch.no_grad():
|
|
batch = next(data.iterator)
|
|
ratio = batch["images"].shape[-2] / batch["images"].shape[-1]
|
|
|
|
shape_lr = self.get_target_lr_size(ratio)
|
|
conditions = self.get_conditions(
|
|
batch,
|
|
models,
|
|
extras,
|
|
is_eval=True,
|
|
is_unconditional=False,
|
|
eval_image_embeds=False,
|
|
)
|
|
unconditions = self.get_conditions(
|
|
batch,
|
|
models,
|
|
extras,
|
|
is_eval=True,
|
|
is_unconditional=True,
|
|
eval_image_embeds=False,
|
|
)
|
|
|
|
latents = self.encode_latents(batch, models, extras)
|
|
latents_lr = self.encode_latents(
|
|
batch, models, extras, target_size=shape_lr
|
|
)
|
|
|
|
if self.is_main_node:
|
|
|
|
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
|
|
|
|
*_, (sampled, _, _, sampled_lr) = extras.gdf.sample(
|
|
models.generator,
|
|
conditions,
|
|
latents.shape,
|
|
latents_lr.shape,
|
|
unconditions,
|
|
device=self.device,
|
|
**extras.sampling_configs,
|
|
)
|
|
|
|
sampled_ema = sampled
|
|
sampled_ema_lr = sampled_lr
|
|
|
|
if self.is_main_node:
|
|
print(
|
|
"sampling results hr latent shape ",
|
|
latents.shape,
|
|
"lr latent shape",
|
|
latents_lr.shape,
|
|
)
|
|
noised_images = torch.cat(
|
|
[
|
|
self.decode_latents(
|
|
latents[i : i + 1].float(), batch, models, extras
|
|
)
|
|
for i in range(len(latents))
|
|
],
|
|
dim=0,
|
|
)
|
|
|
|
sampled_images = torch.cat(
|
|
[
|
|
self.decode_latents(
|
|
sampled[i : i + 1].float(), batch, models, extras
|
|
)
|
|
for i in range(len(sampled))
|
|
],
|
|
dim=0,
|
|
)
|
|
sampled_images_ema = torch.cat(
|
|
[
|
|
self.decode_latents(
|
|
sampled_ema[i : i + 1].float(), batch, models, extras
|
|
)
|
|
for i in range(len(sampled_ema))
|
|
],
|
|
dim=0,
|
|
)
|
|
|
|
noised_images_lr = torch.cat(
|
|
[
|
|
self.decode_latents(
|
|
latents_lr[i : i + 1].float(), batch, models, extras
|
|
)
|
|
for i in range(len(latents_lr))
|
|
],
|
|
dim=0,
|
|
)
|
|
|
|
sampled_images_lr = torch.cat(
|
|
[
|
|
self.decode_latents(
|
|
sampled_lr[i : i + 1].float(), batch, models, extras
|
|
)
|
|
for i in range(len(sampled_lr))
|
|
],
|
|
dim=0,
|
|
)
|
|
sampled_images_ema_lr = torch.cat(
|
|
[
|
|
self.decode_latents(
|
|
sampled_ema_lr[i : i + 1].float(), batch, models, extras
|
|
)
|
|
for i in range(len(sampled_ema_lr))
|
|
],
|
|
dim=0,
|
|
)
|
|
|
|
images = batch["images"]
|
|
if images.size(-1) != noised_images.size(-1) or images.size(
|
|
-2
|
|
) != noised_images.size(-2):
|
|
images = nn.functional.interpolate(
|
|
images, size=noised_images.shape[-2:], mode="bicubic"
|
|
)
|
|
images_lr = nn.functional.interpolate(
|
|
images, size=noised_images_lr.shape[-2:], mode="bicubic"
|
|
)
|
|
|
|
collage_img = torch.cat(
|
|
[
|
|
torch.cat([i for i in images.cpu()], dim=-1),
|
|
torch.cat([i for i in noised_images.cpu()], dim=-1),
|
|
torch.cat([i for i in sampled_images.cpu()], dim=-1),
|
|
torch.cat([i for i in sampled_images_ema.cpu()], dim=-1),
|
|
],
|
|
dim=-2,
|
|
)
|
|
|
|
collage_img_lr = torch.cat(
|
|
[
|
|
torch.cat([i for i in images_lr.cpu()], dim=-1),
|
|
torch.cat([i for i in noised_images_lr.cpu()], dim=-1),
|
|
torch.cat([i for i in sampled_images_lr.cpu()], dim=-1),
|
|
torch.cat([i for i in sampled_images_ema_lr.cpu()], dim=-1),
|
|
],
|
|
dim=-2,
|
|
)
|
|
|
|
torchvision.utils.save_image(
|
|
collage_img,
|
|
f"{self.config.output_path}/{self.config.experiment_id}/{self.info.total_steps:06d}.jpg",
|
|
)
|
|
torchvision.utils.save_image(
|
|
collage_img_lr,
|
|
f"{self.config.output_path}/{self.config.experiment_id}/{self.info.total_steps:06d}_lr.jpg",
|
|
)
|
|
|
|
captions = batch["captions"]
|
|
if self.config.wandb_project is not None:
|
|
log_data = [
|
|
[captions[i]]
|
|
+ [wandb.Image(sampled_images[i])]
|
|
+ [wandb.Image(sampled_images_ema[i])]
|
|
+ [wandb.Image(images[i])]
|
|
for i in range(len(images))
|
|
]
|
|
log_table = wandb.Table(
|
|
data=log_data,
|
|
columns=["Captions", "Sampled", "Sampled EMA", "Orig"],
|
|
)
|
|
wandb.log({"Log": log_table})
|
|
|
|
if isinstance(extras.gdf.loss_weight, AdaptiveLossWeight):
|
|
plt.plot(
|
|
extras.gdf.loss_weight.bucket_ranges,
|
|
extras.gdf.loss_weight.bucket_losses[:-1],
|
|
)
|
|
plt.ylabel("Raw Loss")
|
|
plt.ylabel("LogSNR")
|
|
wandb.log({"Loss/LogSRN": plt})
|
|
|
|
models.generator.train()
|
|
models.train_norm.train()
|
|
print("finish sampling")
|
|
|
|
def sample_fortest(
|
|
self,
|
|
models: Models,
|
|
extras: Extras,
|
|
hr_shape,
|
|
lr_shape,
|
|
batch,
|
|
eval_image_embeds=False,
|
|
):
|
|
|
|
models.generator.eval()
|
|
models.trans_inr.eval()
|
|
with torch.no_grad():
|
|
|
|
if self.is_main_node:
|
|
conditions = self.get_conditions(
|
|
batch,
|
|
models,
|
|
extras,
|
|
is_eval=True,
|
|
is_unconditional=False,
|
|
eval_image_embeds=eval_image_embeds,
|
|
)
|
|
unconditions = self.get_conditions(
|
|
batch,
|
|
models,
|
|
extras,
|
|
is_eval=True,
|
|
is_unconditional=True,
|
|
eval_image_embeds=False,
|
|
)
|
|
|
|
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
|
|
|
|
*_, (sampled, _, _, sampled_lr) = extras.gdf.sample(
|
|
models.generator,
|
|
conditions,
|
|
hr_shape,
|
|
lr_shape,
|
|
unconditions,
|
|
device=self.device,
|
|
**extras.sampling_configs,
|
|
)
|
|
|
|
if models.generator_ema is not None:
|
|
|
|
*_, (sampled_ema, _, _, sampled_ema_lr) = extras.gdf.sample(
|
|
models.generator_ema,
|
|
conditions,
|
|
latents.shape,
|
|
latents_lr.shape,
|
|
unconditions,
|
|
device=self.device,
|
|
**extras.sampling_configs,
|
|
)
|
|
|
|
else:
|
|
sampled_ema = sampled
|
|
sampled_ema_lr = sampled_lr
|
|
|
|
return sampled, sampled_lr
|
|
|
|
|
|
def main_worker(rank, cfg):
|
|
print("Launching Script in main worker")
|
|
warpcore = WurstCore(config_file_path=cfg, rank=rank, world_size=get_world_size())
|
|
# core.fsdp_defaults['sharding_strategy'] = ShardingStrategy.NO_SHARD
|
|
|
|
# RUN TRAINING
|
|
warpcore(get_world_size() == 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
if get_master_ip() == "127.0.0.1":
|
|
|
|
mp.spawn(
|
|
main_worker,
|
|
nprocs=get_world_size(),
|
|
args=(sys.argv[1] if len(sys.argv) > 1 else None,),
|
|
)
|
|
else:
|
|
main_worker(
|
|
0,
|
|
sys.argv[1] if len(sys.argv) > 1 else None,
|
|
)
|