Spaces:
Running
on
Zero
Running
on
Zero
| #!/usr/bin/env python | |
| # -*- coding:utf-8 -*- | |
| # Power by Zongsheng Yue 2022-05-18 13:04:06 | |
| import os, sys, math, time, random, datetime | |
| import numpy as np | |
| from box import Box | |
| from pathlib import Path | |
| from loguru import logger | |
| from copy import deepcopy | |
| from omegaconf import OmegaConf | |
| from einops import rearrange | |
| from typing import Any, Dict, List, Optional, Tuple, Union | |
| from datapipe.datasets import create_dataset | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torch.utils.data as udata | |
| import torch.distributed as dist | |
| import torch.multiprocessing as mp | |
| import torchvision.utils as vutils | |
| from torch.nn.parallel import DistributedDataParallel as DDP | |
| from utils import util_net | |
| from utils import util_common | |
| from utils import util_image | |
| from utils.util_ops import append_dims | |
| import pyiqa | |
| from basicsr.utils import DiffJPEG, USMSharp | |
| from basicsr.utils.img_process_util import filter2D | |
| from basicsr.data.transforms import paired_random_crop | |
| from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt | |
| from diffusers import EulerDiscreteScheduler | |
| from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution | |
| from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img import retrieve_timesteps | |
| _base_seed = 10**6 | |
| _INTERPOLATION_MODE = 'bicubic' | |
| _Latent_bound = {'min':-10.0, 'max':10.0} | |
| _positive= 'Cinematic, high-contrast, photo-realistic, 8k, ultra HD, ' +\ | |
| 'meticulous detailing, hyper sharpness, perfect without deformations' | |
| _negative= 'Low quality, blurring, jpeg artifacts, deformed, over-smooth, cartoon, noisy,' +\ | |
| 'painting, drawing, sketch, oil painting' | |
| class TrainerBase: | |
| def __init__(self, configs): | |
| self.configs = configs | |
| # setup distributed training: self.num_gpus, self.rank | |
| self.setup_dist() | |
| # setup seed | |
| self.setup_seed() | |
| def setup_dist(self): | |
| num_gpus = torch.cuda.device_count() | |
| if num_gpus > 1: | |
| if mp.get_start_method(allow_none=True) is None: | |
| mp.set_start_method('spawn') | |
| rank = int(os.environ['LOCAL_RANK']) | |
| torch.cuda.set_device(rank % num_gpus) | |
| dist.init_process_group( | |
| timeout=datetime.timedelta(seconds=3600), | |
| backend='nccl', | |
| init_method='env://', | |
| ) | |
| self.num_gpus = num_gpus | |
| self.rank = int(os.environ['LOCAL_RANK']) if num_gpus > 1 else 0 | |
| def setup_seed(self, seed=None, global_seeding=None): | |
| if seed is None: | |
| seed = self.configs.train.get('seed', 12345) | |
| if global_seeding is None: | |
| global_seeding = self.configs.train.get('global_seeding', False) | |
| if not global_seeding: | |
| seed += self.rank | |
| torch.cuda.manual_seed(seed) | |
| else: | |
| torch.cuda.manual_seed_all(seed) | |
| random.seed(seed) | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| def init_logger(self): | |
| if self.configs.resume: | |
| assert self.configs.resume.endswith(".pth") | |
| save_dir = Path(self.configs.resume).parents[1] | |
| project_id = save_dir.name | |
| else: | |
| project_id = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M") | |
| save_dir = Path(self.configs.save_dir) / project_id | |
| if not save_dir.exists() and self.rank == 0: | |
| save_dir.mkdir(parents=True) | |
| # setting log counter | |
| if self.rank == 0: | |
| self.log_step = {phase: 1 for phase in ['train', 'val']} | |
| self.log_step_img = {phase: 1 for phase in ['train', 'val']} | |
| # text logging | |
| logtxet_path = save_dir / 'training.log' | |
| if self.rank == 0: | |
| if logtxet_path.exists(): | |
| assert self.configs.resume | |
| self.logger = logger | |
| self.logger.remove() | |
| self.logger.add(logtxet_path, format="{message}", mode='a', level='INFO') | |
| self.logger.add(sys.stdout, format="{message}") | |
| # tensorboard logging | |
| log_dir = save_dir / 'tf_logs' | |
| self.tf_logging = self.configs.train.tf_logging | |
| if self.rank == 0 and self.tf_logging: | |
| if not log_dir.exists(): | |
| log_dir.mkdir() | |
| self.writer = SummaryWriter(str(log_dir)) | |
| # checkpoint saving | |
| ckpt_dir = save_dir / 'ckpts' | |
| self.ckpt_dir = ckpt_dir | |
| if self.rank == 0 and (not ckpt_dir.exists()): | |
| ckpt_dir.mkdir() | |
| if 'ema_rate' in self.configs.train: | |
| self.ema_rate = self.configs.train.ema_rate | |
| assert isinstance(self.ema_rate, float), "Ema rate must be a float number" | |
| ema_ckpt_dir = save_dir / 'ema_ckpts' | |
| self.ema_ckpt_dir = ema_ckpt_dir | |
| if self.rank == 0 and (not ema_ckpt_dir.exists()): | |
| ema_ckpt_dir.mkdir() | |
| # save images into local disk | |
| self.local_logging = self.configs.train.local_logging | |
| if self.rank == 0 and self.local_logging: | |
| image_dir = save_dir / 'images' | |
| if not image_dir.exists(): | |
| (image_dir / 'train').mkdir(parents=True) | |
| (image_dir / 'val').mkdir(parents=True) | |
| self.image_dir = image_dir | |
| # logging the configurations | |
| if self.rank == 0: | |
| self.logger.info(OmegaConf.to_yaml(self.configs)) | |
| def close_logger(self): | |
| if self.rank == 0 and self.tf_logging: | |
| self.writer.close() | |
| def resume_from_ckpt(self): | |
| if self.configs.resume: | |
| assert self.configs.resume.endswith(".pth") and os.path.isfile(self.configs.resume) | |
| if self.rank == 0: | |
| self.logger.info(f"=> Loading checkpoint from {self.configs.resume}") | |
| ckpt = torch.load(self.configs.resume, map_location=f"cuda:{self.rank}") | |
| util_net.reload_model(self.model, ckpt['state_dict']) | |
| if self.configs.train.loss_coef.get('ldis', 0) > 0: | |
| util_net.reload_model(self.discriminator, ckpt['state_dict_dis']) | |
| torch.cuda.empty_cache() | |
| # learning rate scheduler | |
| self.iters_start = ckpt['iters_start'] | |
| for ii in range(1, self.iters_start+1): | |
| self.adjust_lr(ii) | |
| # logging | |
| if self.rank == 0: | |
| self.log_step = ckpt['log_step'] | |
| self.log_step_img = ckpt['log_step_img'] | |
| # EMA model | |
| if self.rank == 0 and hasattr(self.configs.train, 'ema_rate'): | |
| ema_ckpt_path = self.ema_ckpt_dir / ("ema_"+Path(self.configs.resume).name) | |
| self.logger.info(f"=> Loading EMA checkpoint from {str(ema_ckpt_path)}") | |
| ema_ckpt = torch.load(ema_ckpt_path, map_location=f"cuda:{self.rank}") | |
| util_net.reload_model(self.ema_model, ema_ckpt) | |
| torch.cuda.empty_cache() | |
| # AMP scaler | |
| if self.amp_scaler is not None: | |
| if "amp_scaler" in ckpt: | |
| self.amp_scaler.load_state_dict(ckpt["amp_scaler"]) | |
| if self.rank == 0: | |
| self.logger.info("Loading scaler from resumed state...") | |
| if self.configs.get('discriminator', None) is not None: | |
| if "amp_scaler_dis" in ckpt: | |
| self.amp_scaler_dis.load_state_dict(ckpt["amp_scaler_dis"]) | |
| if self.rank == 0: | |
| self.logger.info("Loading scaler (discriminator) from resumed state...") | |
| # reset the seed | |
| self.setup_seed(seed=self.iters_start) | |
| else: | |
| self.iters_start = 0 | |
| def setup_optimizaton(self): | |
| self.optimizer = torch.optim.AdamW(self.model.parameters(), | |
| lr=self.configs.train.lr, | |
| weight_decay=self.configs.train.weight_decay) | |
| # amp settings | |
| self.amp_scaler = torch.amp.GradScaler('cuda') if self.configs.train.use_amp else None | |
| if self.configs.train.lr_schedule == 'cosin': | |
| self.lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( | |
| optimizer=self.optimizer, | |
| T_max=self.configs.train.iterations - self.configs.train.warmup_iterations, | |
| eta_min=self.configs.train.lr_min, | |
| ) | |
| if self.configs.train.loss_coef.get('ldis', 0) > 0: | |
| self.optimizer_dis = torch.optim.Adam( | |
| self.discriminator.parameters(), | |
| lr=self.configs.train.lr_dis, | |
| weight_decay=self.configs.train.weight_decay_dis, | |
| ) | |
| self.amp_scaler_dis = torch.amp.GradScaler('cuda') if self.configs.train.use_amp else None | |
| def prepare_compiling(self): | |
| # https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_3#stable-diffusion-3 | |
| if not hasattr(self, "prepare_compiling_well") or (not self.prepare_compiling_well): | |
| torch.set_float32_matmul_precision("high") | |
| torch._inductor.config.conv_1x1_as_mm = True | |
| torch._inductor.config.coordinate_descent_tuning = True | |
| torch._inductor.config.epilogue_fusion = False | |
| torch._inductor.config.coordinate_descent_check_all_directions = True | |
| self.prepare_compiling_well = True | |
| def build_model(self): | |
| if self.configs.train.get("compile", True): | |
| self.prepare_compiling() | |
| params = self.configs.model.get('params', dict) | |
| model = util_common.get_obj_from_str(self.configs.model.target)(**params) | |
| model.cuda() | |
| if not self.configs.train.start_mode: # Loading the starting model for evaluation | |
| self.start_model = deepcopy(model) | |
| assert self.configs.model.ckpt_start_path is not None | |
| ckpt_start_path = self.configs.model.ckpt_start_path | |
| if self.rank == 0: | |
| self.logger.info(f"Loading the starting model from {ckpt_start_path}") | |
| ckpt = torch.load(ckpt_start_path, map_location=f"cuda:{self.rank}") | |
| if 'state_dict' in ckpt: | |
| ckpt = ckpt['state_dict'] | |
| util_net.reload_model(self.start_model, ckpt) | |
| self.freeze_model(self.start_model) | |
| self.start_model.eval() | |
| # delete the started timestep | |
| start_timestep = max(self.configs.train.timesteps) | |
| self.configs.train.timesteps.remove(start_timestep) | |
| # end_timestep = min(self.configs.train.timesteps) | |
| # self.configs.train.timesteps.remove(end_timestep) | |
| # setting the training model | |
| if self.configs.model.get('ckpt_path', None): # initialize if necessary | |
| ckpt_path = self.configs.model.ckpt_path | |
| if self.rank == 0: | |
| self.logger.info(f"Initializing model from {ckpt_path}") | |
| ckpt = torch.load(ckpt_path, map_location=f"cuda:{self.rank}") | |
| if 'state_dict' in ckpt: | |
| ckpt = ckpt['state_dict'] | |
| util_net.reload_model(model, ckpt) | |
| if self.configs.model.get("compile", False): | |
| if self.rank == 0: | |
| self.logger.info("Compile the model...") | |
| model.to(memory_format=torch.channels_last) | |
| model = torch.compile(model, mode="max-autotune", fullgraph=False) | |
| if self.num_gpus > 1: | |
| model = DDP(model, device_ids=[self.rank,]) # wrap the network | |
| if self.rank == 0 and hasattr(self.configs.train, 'ema_rate'): | |
| self.ema_model = deepcopy(model) | |
| self.freeze_model(self.ema_model) | |
| self.model = model | |
| # discriminator if necessary | |
| if self.configs.train.loss_coef.get('ldis', 0) > 0: | |
| assert hasattr(self.configs, 'discriminator') | |
| params = self.configs.discriminator.get('params', dict) | |
| discriminator = util_common.get_obj_from_str(self.configs.discriminator.target)(**params) | |
| discriminator.cuda() | |
| if self.configs.discriminator.get("compile", False): | |
| if self.rank == 0: | |
| self.logger.info("Compile the discriminator...") | |
| discriminator.to(memory_format=torch.channels_last) | |
| discriminator = torch.compile(discriminator, mode="max-autotune", fullgraph=False) | |
| if self.num_gpus > 1: | |
| discriminator = DDP(discriminator, device_ids=[self.rank,]) # wrap the network | |
| if self.configs.train.loss_coef.get('ldis', 0) > 0: | |
| if self.configs.discriminator.enable_grad_checkpoint: | |
| if self.rank == 0: | |
| self.logger.info("Activating gradient checkpointing for discriminator...") | |
| self.set_grad_checkpointing(discriminator) | |
| self.discriminator = discriminator | |
| # build the stable diffusion | |
| params = dict(self.configs.sd_pipe.params) | |
| torch_dtype = params.pop('torch_dtype') | |
| params['torch_dtype'] = get_torch_dtype(torch_dtype) | |
| # loading the fp16 robust vae for sdxl: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix | |
| if self.configs.get('vae_fp16', None) is not None: | |
| params_vae = dict(self.configs.vae_fp16.params) | |
| params_vae['torch_dtype'] = torch.float16 | |
| pipe_id = self.configs.vae_fp16.params.pretrained_model_name_or_path | |
| if self.rank == 0: | |
| self.logger.info(f'Loading improved vae from {pipe_id}...') | |
| vae_pipe = util_common.get_obj_from_str(self.configs.vae_fp16.target).from_pretrained(**params_vae) | |
| if self.rank == 0: | |
| self.logger.info('Loaded Done') | |
| params['vae'] = vae_pipe | |
| if ("StableDiffusion3" in self.configs.sd_pipe.target.split('.')[-1] | |
| and self.configs.sd_pipe.get("model_quantization", False)): | |
| if self.rank == 0: | |
| self.logger.info(f'Loading the quantized transformer for SD3...') | |
| nf4_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_compute_dtype=torch.bfloat16 | |
| ) | |
| params_model = dict(self.configs.model_nf4.params) | |
| torch_dtype = params_model.pop('torch_dtype') | |
| params_model['torch_dtype'] = get_torch_dtype(torch_dtype) | |
| params_model['quantization_config'] = nf4_config | |
| model_nf4 = util_common.get_obj_from_str(self.configs.model_nf4.target).from_pretrained( | |
| **params_model | |
| ) | |
| params['transformer'] = model_nf4 | |
| sd_pipe = util_common.get_obj_from_str(self.configs.sd_pipe.target).from_pretrained(**params) | |
| if self.configs.get('scheduler', None) is not None: | |
| pipe_id = self.configs.scheduler.target.split('.')[-1] | |
| if self.rank == 0: | |
| self.logger.info(f'Loading scheduler of {pipe_id}...') | |
| sd_pipe.scheduler = util_common.get_obj_from_str(self.configs.scheduler.target).from_config( | |
| sd_pipe.scheduler.config | |
| ) | |
| if self.rank == 0: | |
| self.logger.info('Loaded Done') | |
| if ("StableDiffusion3" in self.configs.sd_pipe.target.split('.')[-1] | |
| and self.configs.sd_pipe.get("model_quantization", False)): | |
| sd_pipe.enable_model_cpu_offload(gpu_id=self.rank,device='cuda') | |
| else: | |
| sd_pipe.to(f"cuda:{self.rank}") | |
| # freezing model parameters | |
| if hasattr(sd_pipe, 'unet'): | |
| self.freeze_model(sd_pipe.unet) | |
| if hasattr(sd_pipe, 'transformer'): | |
| self.freeze_model(sd_pipe.transformer) | |
| self.freeze_model(sd_pipe.vae) | |
| # compiling | |
| if self.configs.sd_pipe.get('compile', True): | |
| if self.rank == 0: | |
| self.logger.info('Compile the SD model...') | |
| sd_pipe.set_progress_bar_config(disable=True) | |
| if hasattr(sd_pipe, 'unet'): | |
| sd_pipe.unet.to(memory_format=torch.channels_last) | |
| sd_pipe.unet = torch.compile(sd_pipe.unet, mode="max-autotune", fullgraph=False) | |
| if hasattr(sd_pipe, 'transformer'): | |
| sd_pipe.transformer.to(memory_format=torch.channels_last) | |
| sd_pipe.transformer = torch.compile(sd_pipe.transformer, mode="max-autotune", fullgraph=False) | |
| sd_pipe.vae.to(memory_format=torch.channels_last) | |
| sd_pipe.vae = torch.compile(sd_pipe.vae, mode="max-autotune", fullgraph=True) | |
| # setting gradient checkpoint for vae | |
| if self.configs.sd_pipe.get("enable_grad_checkpoint_vae", True): | |
| if self.rank == 0: | |
| self.logger.info("Activating gradient checkpointing for VAE...") | |
| sd_pipe.vae._set_gradient_checkpointing(sd_pipe.vae.encoder) | |
| sd_pipe.vae._set_gradient_checkpointing(sd_pipe.vae.decoder) | |
| # setting gradient checkpoint for diffusion model | |
| if self.configs.sd_pipe.enable_grad_checkpoint: | |
| if self.rank == 0: | |
| self.logger.info("Activating gradient checkpointing for SD...") | |
| if hasattr(sd_pipe, 'unet'): | |
| self.set_grad_checkpointing(sd_pipe.unet) | |
| if hasattr(sd_pipe, 'transformer'): | |
| self.set_grad_checkpointing(sd_pipe.transformer) | |
| self.sd_pipe = sd_pipe | |
| # latent LPIPS loss | |
| if self.configs.train.loss_coef.get('llpips', 0) > 0: | |
| params = self.configs.llpips.get('params', dict) | |
| llpips_loss = util_common.get_obj_from_str(self.configs.llpips.target)(**params) | |
| llpips_loss.cuda() | |
| self.freeze_model(llpips_loss) | |
| # loading the pre-trained model | |
| ckpt_path = self.configs.llpips.ckpt_path | |
| self.load_model(llpips_loss, ckpt_path, tag='latent lpips') | |
| if self.configs.llpips.get("compile", True): | |
| if self.rank == 0: | |
| self.logger.info('Compile the llpips loss...') | |
| llpips_loss.to(memory_format=torch.channels_last) | |
| llpips_loss = torch.compile(llpips_loss, mode="max-autotune", fullgraph=True) | |
| self.llpips_loss = llpips_loss | |
| # model information | |
| self.print_model_info() | |
| torch.cuda.empty_cache() | |
| def set_grad_checkpointing(self, model): | |
| if hasattr(model, 'down_blocks'): | |
| for module in model.down_blocks: | |
| module.gradient_checkpointing = True | |
| module.training = True | |
| if hasattr(model, 'up_blocks'): | |
| for module in model.up_blocks: | |
| module.gradient_checkpointing = True | |
| module.training = True | |
| if hasattr(model, 'mid_blocks'): | |
| model.mid_block.gradient_checkpointing = True | |
| model.mid_block.training = True | |
| def build_dataloader(self): | |
| def _wrap_loader(loader): | |
| while True: yield from loader | |
| # make datasets | |
| datasets = {'train': create_dataset(self.configs.data.get('train', dict)), } | |
| if hasattr(self.configs.data, 'val') and self.rank == 0: | |
| datasets['val'] = create_dataset(self.configs.data.get('val', dict)) | |
| if self.rank == 0: | |
| for phase in datasets.keys(): | |
| length = len(datasets[phase]) | |
| self.logger.info('Number of images in {:s} data set: {:d}'.format(phase, length)) | |
| # make dataloaders | |
| if self.num_gpus > 1: | |
| sampler = udata.distributed.DistributedSampler( | |
| datasets['train'], | |
| num_replicas=self.num_gpus, | |
| rank=self.rank, | |
| ) | |
| else: | |
| sampler = None | |
| dataloaders = {'train': _wrap_loader(udata.DataLoader( | |
| datasets['train'], | |
| batch_size=self.configs.train.batch // self.num_gpus, | |
| shuffle=False if self.num_gpus > 1 else True, | |
| drop_last=True, | |
| num_workers=min(self.configs.train.num_workers, 4), | |
| pin_memory=True, | |
| prefetch_factor=self.configs.train.get('prefetch_factor', 2), | |
| worker_init_fn=my_worker_init_fn, | |
| sampler=sampler, | |
| ))} | |
| if hasattr(self.configs.data, 'val') and self.rank == 0: | |
| dataloaders['val'] = udata.DataLoader(datasets['val'], | |
| batch_size=self.configs.validate.batch, | |
| shuffle=False, | |
| drop_last=False, | |
| num_workers=0, | |
| pin_memory=True, | |
| ) | |
| self.datasets = datasets | |
| self.dataloaders = dataloaders | |
| self.sampler = sampler | |
| def print_model_info(self): | |
| if self.rank == 0: | |
| num_params = util_net.calculate_parameters(self.model) / 1000**2 | |
| # self.logger.info("Detailed network architecture:") | |
| # self.logger.info(self.model.__repr__()) | |
| if self.configs.train.get('use_fsdp', False): | |
| num_params *= self.num_gpus | |
| self.logger.info(f"Number of parameters: {num_params:.2f}M") | |
| if hasattr(self, 'discriminator'): | |
| num_params = util_net.calculate_parameters(self.discriminator) / 1000**2 | |
| self.logger.info(f"Number of parameters in discriminator: {num_params:.2f}M") | |
| def prepare_data(self, data, dtype=torch.float32, phase='train'): | |
| data = {key:value.cuda().to(dtype=dtype) for key, value in data.items()} | |
| return data | |
| def validation(self): | |
| pass | |
| def train(self): | |
| self.init_logger() # setup logger: self.logger | |
| self.build_dataloader() # prepare data: self.dataloaders, self.datasets, self.sampler | |
| self.build_model() # build model: self.model, self.loss | |
| self.setup_optimizaton() # setup optimization: self.optimzer, self.sheduler | |
| self.resume_from_ckpt() # resume if necessary | |
| self.model.train() | |
| num_iters_epoch = math.ceil(len(self.datasets['train']) / self.configs.train.batch) | |
| for ii in range(self.iters_start, self.configs.train.iterations): | |
| self.current_iters = ii + 1 | |
| # prepare data | |
| data = self.prepare_data(next(self.dataloaders['train']), phase='train') | |
| # training phase | |
| self.training_step(data) | |
| # update ema model | |
| if hasattr(self.configs.train, 'ema_rate') and self.rank == 0: | |
| self.update_ema_model() | |
| # validation phase | |
| if ((ii+1) % self.configs.train.save_freq == 0 and | |
| 'val' in self.dataloaders and | |
| self.rank == 0 | |
| ): | |
| self.validation() | |
| #update learning rate | |
| self.adjust_lr() | |
| # save checkpoint | |
| if (ii+1) % self.configs.train.save_freq == 0 and self.rank == 0: | |
| self.save_ckpt() | |
| if (ii+1) % num_iters_epoch == 0 and self.sampler is not None: | |
| self.sampler.set_epoch(ii+1) | |
| # close the tensorboard | |
| self.close_logger() | |
| def adjust_lr(self, current_iters=None): | |
| base_lr = self.configs.train.lr | |
| warmup_steps = self.configs.train.get("warmup_iterations", 0) | |
| current_iters = self.current_iters if current_iters is None else current_iters | |
| if current_iters <= warmup_steps: | |
| for params_group in self.optimizer.param_groups: | |
| params_group['lr'] = (current_iters / warmup_steps) * base_lr | |
| else: | |
| if hasattr(self, 'lr_scheduler'): | |
| self.lr_scheduler.step() | |
| def save_ckpt(self): | |
| ckpt_path = self.ckpt_dir / 'model_{:d}.pth'.format(self.current_iters) | |
| ckpt = { | |
| 'iters_start': self.current_iters, | |
| 'log_step': {phase:self.log_step[phase] for phase in ['train', 'val']}, | |
| 'log_step_img': {phase:self.log_step_img[phase] for phase in ['train', 'val']}, | |
| 'state_dict': self.model.state_dict(), | |
| } | |
| if self.amp_scaler is not None: | |
| ckpt['amp_scaler'] = self.amp_scaler.state_dict() | |
| if self.configs.train.loss_coef.get('ldis', 0) > 0: | |
| ckpt['state_dict_dis'] = self.discriminator.state_dict() | |
| if self.amp_scaler_dis is not None: | |
| ckpt['amp_scaler_dis'] = self.amp_scaler_dis.state_dict() | |
| torch.save(ckpt, ckpt_path) | |
| if hasattr(self.configs.train, 'ema_rate'): | |
| ema_ckpt_path = self.ema_ckpt_dir / 'ema_model_{:d}.pth'.format(self.current_iters) | |
| torch.save(self.ema_model.state_dict(), ema_ckpt_path) | |
| def logging_image(self, im_tensor, tag, phase, add_global_step=False, nrow=8): | |
| """ | |
| Args: | |
| im_tensor: b x c x h x w tensor | |
| im_tag: str | |
| phase: 'train' or 'val' | |
| nrow: number of displays in each row | |
| """ | |
| assert self.tf_logging or self.local_logging | |
| im_tensor = vutils.make_grid(im_tensor, nrow=nrow, normalize=True, scale_each=True) # c x H x W | |
| if self.local_logging: | |
| im_path = str(self.image_dir / phase / f"{tag}-{self.log_step_img[phase]}.png") | |
| im_np = im_tensor.cpu().permute(1,2,0).numpy() | |
| util_image.imwrite(im_np, im_path) | |
| if self.tf_logging: | |
| self.writer.add_image( | |
| f"{phase}-{tag}-{self.log_step_img[phase]}", | |
| im_tensor, | |
| self.log_step_img[phase], | |
| ) | |
| if add_global_step: | |
| self.log_step_img[phase] += 1 | |
| def logging_text(self, text_list, phase): | |
| """ | |
| Args: | |
| text_list: (b,) list | |
| phase: 'train' or 'val' | |
| """ | |
| assert self.local_logging | |
| if self.local_logging: | |
| text_path = str(self.image_dir / phase / f"text-{self.log_step_img[phase]}.txt") | |
| with open(text_path, 'w') as ff: | |
| for text in text_list: | |
| ff.write(text + '\n') | |
| def logging_metric(self, metrics, tag, phase, add_global_step=False): | |
| """ | |
| Args: | |
| metrics: dict | |
| tag: str | |
| phase: 'train' or 'val' | |
| """ | |
| if self.tf_logging: | |
| tag = f"{phase}-{tag}" | |
| if isinstance(metrics, dict): | |
| self.writer.add_scalars(tag, metrics, self.log_step[phase]) | |
| else: | |
| self.writer.add_scalar(tag, metrics, self.log_step[phase]) | |
| if add_global_step: | |
| self.log_step[phase] += 1 | |
| else: | |
| pass | |
| def load_model(self, model, ckpt_path=None, tag='model'): | |
| if self.rank == 0: | |
| self.logger.info(f'Loading {tag} from {ckpt_path}...') | |
| ckpt = torch.load(ckpt_path, map_location=f"cuda:{self.rank}") | |
| if 'state_dict' in ckpt: | |
| ckpt = ckpt['state_dict'] | |
| util_net.reload_model(model, ckpt) | |
| if self.rank == 0: | |
| self.logger.info('Loaded Done') | |
| def freeze_model(self, net): | |
| for params in net.parameters(): | |
| params.requires_grad = False | |
| def unfreeze_model(self, net): | |
| for params in net.parameters(): | |
| params.requires_grad = True | |
| def update_ema_model(self): | |
| decay = min(self.configs.train.ema_rate, (1 + self.current_iters) / (10 + self.current_iters)) | |
| target_params = dict(self.model.named_parameters()) | |
| # if hasattr(self.configs.train, 'ema_rate'): | |
| # with FSDP.summon_full_params(self.model, writeback=True): | |
| # target_params = dict(self.model.named_parameters()) | |
| # else: | |
| # target_params = dict(self.model.named_parameters()) | |
| one_minus_decay = 1.0 - decay | |
| for key, source_value in self.ema_model.named_parameters(): | |
| target_value = target_params[key] | |
| if target_value.requires_grad: | |
| source_value.sub_(one_minus_decay * (source_value - target_value.data)) | |
| class TrainerBaseSR(TrainerBase): | |
| def _dequeue_and_enqueue(self): | |
| """It is the training pair pool for increasing the diversity in a batch. | |
| Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a | |
| batch could not have different resize scaling factors. Therefore, we employ this training pair pool | |
| to increase the degradation diversity in a batch. | |
| """ | |
| # initialize | |
| b, c, h, w = self.lq.size() | |
| if not hasattr(self, 'queue_size'): | |
| self.queue_size = self.configs.degradation.get('queue_size', b*10) | |
| if not hasattr(self, 'queue_lr'): | |
| assert self.queue_size % b == 0, f'queue size {self.queue_size} should be divisible by batch size {b}' | |
| self.queue_lr = torch.zeros(self.queue_size, c, h, w).cuda() | |
| _, c, h, w = self.gt.size() | |
| self.queue_gt = torch.zeros(self.queue_size, c, h, w).cuda() | |
| _, c, h, w = self.gt_latent.size() | |
| self.queue_gt_latent = torch.zeros(self.queue_size, c, h, w).cuda() | |
| self.queue_txt = ["", ] * self.queue_size | |
| self.queue_ptr = 0 | |
| if self.queue_ptr == self.queue_size: # the pool is full | |
| # do dequeue and enqueue | |
| # shuffle | |
| idx = torch.randperm(self.queue_size) | |
| self.queue_lr = self.queue_lr[idx] | |
| self.queue_gt = self.queue_gt[idx] | |
| self.queue_gt_latent = self.queue_gt_latent[idx] | |
| self.queue_txt = [self.queue_txt[ii] for ii in idx] | |
| # get first b samples | |
| lq_dequeue = self.queue_lr[0:b, :, :, :].clone() | |
| gt_dequeue = self.queue_gt[0:b, :, :, :].clone() | |
| gt_latent_dequeue = self.queue_gt_latent[0:b, :, :, :].clone() | |
| txt_dequeue = deepcopy(self.queue_txt[0:b]) | |
| # update the queue | |
| self.queue_lr[0:b, :, :, :] = self.lq.clone() | |
| self.queue_gt[0:b, :, :, :] = self.gt.clone() | |
| self.queue_gt_latent[0:b, :, :, :] = self.gt_latent.clone() | |
| self.queue_txt[0:b] = deepcopy(self.txt) | |
| self.lq = lq_dequeue | |
| self.gt = gt_dequeue | |
| self.gt_latent = gt_latent_dequeue | |
| self.txt = txt_dequeue | |
| else: | |
| # only do enqueue | |
| self.queue_lr[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.lq.clone() | |
| self.queue_gt[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.gt.clone() | |
| self.queue_gt_latent[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.gt_latent.clone() | |
| self.queue_txt[self.queue_ptr:self.queue_ptr + b] = deepcopy(self.txt) | |
| self.queue_ptr = self.queue_ptr + b | |
| def prepare_data(self, data, phase='train'): | |
| if phase == 'train' and self.configs.data.get(phase).get('type') == 'realesrgan': | |
| if not hasattr(self, 'jpeger'): | |
| self.jpeger = DiffJPEG(differentiable=False).cuda() # simulate JPEG compression artifacts | |
| if (not hasattr(self, 'sharpener')) and self.configs.degradation.get('use_sharp', False): | |
| self.sharpener = USMSharp().cuda() | |
| im_gt = data['gt'].cuda() | |
| kernel1 = data['kernel1'].cuda() | |
| kernel2 = data['kernel2'].cuda() | |
| sinc_kernel = data['sinc_kernel'].cuda() | |
| ori_h, ori_w = im_gt.size()[2:4] | |
| if isinstance(self.configs.degradation.sf, int): | |
| sf = self.configs.degradation.sf | |
| else: | |
| assert len(self.configs.degradation.sf) == 2 | |
| sf = random.uniform(*self.configs.degradation.sf) | |
| if self.configs.degradation.use_sharp: | |
| im_gt = self.sharpener(im_gt) | |
| # ----------------------- The first degradation process ----------------------- # | |
| # blur | |
| out = filter2D(im_gt, kernel1) | |
| # random resize | |
| updown_type = random.choices( | |
| ['up', 'down', 'keep'], | |
| self.configs.degradation['resize_prob'], | |
| )[0] | |
| if updown_type == 'up': | |
| scale = random.uniform(1, self.configs.degradation['resize_range'][1]) | |
| elif updown_type == 'down': | |
| scale = random.uniform(self.configs.degradation['resize_range'][0], 1) | |
| else: | |
| scale = 1 | |
| mode = random.choice(['area', 'bilinear', 'bicubic']) | |
| out = F.interpolate(out, scale_factor=scale, mode=mode) | |
| # add noise | |
| gray_noise_prob = self.configs.degradation['gray_noise_prob'] | |
| if random.random() < self.configs.degradation['gaussian_noise_prob']: | |
| out = random_add_gaussian_noise_pt( | |
| out, | |
| sigma_range=self.configs.degradation['noise_range'], | |
| clip=True, | |
| rounds=False, | |
| gray_prob=gray_noise_prob, | |
| ) | |
| else: | |
| out = random_add_poisson_noise_pt( | |
| out, | |
| scale_range=self.configs.degradation['poisson_scale_range'], | |
| gray_prob=gray_noise_prob, | |
| clip=True, | |
| rounds=False) | |
| # JPEG compression | |
| jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.configs.degradation['jpeg_range']) | |
| out = torch.clamp(out, 0, 1) # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts | |
| out = self.jpeger(out, quality=jpeg_p) | |
| # ----------------------- The second degradation process ----------------------- # | |
| if random.random() < self.configs.degradation['second_order_prob']: | |
| # blur | |
| if random.random() < self.configs.degradation['second_blur_prob']: | |
| out = filter2D(out, kernel2) | |
| # random resize | |
| updown_type = random.choices( | |
| ['up', 'down', 'keep'], | |
| self.configs.degradation['resize_prob2'], | |
| )[0] | |
| if updown_type == 'up': | |
| scale = random.uniform(1, self.configs.degradation['resize_range2'][1]) | |
| elif updown_type == 'down': | |
| scale = random.uniform(self.configs.degradation['resize_range2'][0], 1) | |
| else: | |
| scale = 1 | |
| mode = random.choice(['area', 'bilinear', 'bicubic']) | |
| out = F.interpolate( | |
| out, | |
| size=(int(ori_h / sf * scale), int(ori_w / sf * scale)), | |
| mode=mode, | |
| ) | |
| # add noise | |
| gray_noise_prob = self.configs.degradation['gray_noise_prob2'] | |
| if random.random() < self.configs.degradation['gaussian_noise_prob2']: | |
| out = random_add_gaussian_noise_pt( | |
| out, | |
| sigma_range=self.configs.degradation['noise_range2'], | |
| clip=True, | |
| rounds=False, | |
| gray_prob=gray_noise_prob, | |
| ) | |
| else: | |
| out = random_add_poisson_noise_pt( | |
| out, | |
| scale_range=self.configs.degradation['poisson_scale_range2'], | |
| gray_prob=gray_noise_prob, | |
| clip=True, | |
| rounds=False, | |
| ) | |
| # JPEG compression + the final sinc filter | |
| # We also need to resize images to desired sizes. We group [resize back + sinc filter] together | |
| # as one operation. | |
| # We consider two orders: | |
| # 1. [resize back + sinc filter] + JPEG compression | |
| # 2. JPEG compression + [resize back + sinc filter] | |
| # Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines. | |
| if random.random() < 0.5: | |
| # resize back + the final sinc filter | |
| mode = random.choice(['area', 'bilinear', 'bicubic']) | |
| out = F.interpolate( | |
| out, | |
| size=(ori_h // sf, ori_w // sf), | |
| mode=mode, | |
| ) | |
| out = filter2D(out, sinc_kernel) | |
| # JPEG compression | |
| jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.configs.degradation['jpeg_range2']) | |
| out = torch.clamp(out, 0, 1) | |
| out = self.jpeger(out, quality=jpeg_p) | |
| else: | |
| # JPEG compression | |
| jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.configs.degradation['jpeg_range2']) | |
| out = torch.clamp(out, 0, 1) | |
| out = self.jpeger(out, quality=jpeg_p) | |
| # resize back + the final sinc filter | |
| mode = random.choice(['area', 'bilinear', 'bicubic']) | |
| out = F.interpolate( | |
| out, | |
| size=(ori_h // sf, ori_w // sf), | |
| mode=mode, | |
| ) | |
| out = filter2D(out, sinc_kernel) | |
| # resize back | |
| if self.configs.degradation.resize_back: | |
| out = F.interpolate(out, size=(ori_h, ori_w), mode=_INTERPOLATION_MODE) | |
| # clamp and round | |
| im_lq = torch.clamp((out * 255.0).round(), 0, 255) / 255. | |
| self.lq, self.gt, self.txt = im_lq, im_gt, data['txt'] | |
| if "gt_moment" not in data: | |
| self.gt_latent = self.encode_first_stage( | |
| im_gt.cuda(), | |
| center_input_sample=True, | |
| deterministic=self.configs.train.loss_coef.get('rkl', 0) > 0, | |
| ) | |
| else: | |
| self.gt_latent = self.encode_from_moment( | |
| data['gt_moment'].cuda(), | |
| deterministic=self.configs.train.loss_coef.get('rkl', 0) > 0, | |
| ) | |
| if (not self.configs.train.use_text) or self.configs.data.train.params.random_crop: | |
| self.txt = [_positive,] * im_lq.shape[0] | |
| # training pair pool | |
| self._dequeue_and_enqueue() | |
| self.lq = self.lq.contiguous() # for the warning: grad and param do not obey the gradient layout contract | |
| batch = {'lq':self.lq, 'gt':self.gt, 'gt_latent':self.gt_latent, 'txt':self.txt} | |
| elif phase == 'val': | |
| resolution = self.configs.data.train.params.gt_size // self.configs.degradation.sf | |
| batch = {} | |
| batch['lq'] = data['lq'].cuda() | |
| if 'gt' in data: | |
| batch['gt'] = data['gt'].cuda() | |
| batch['txt'] = [_positive, ] * data['lq'].shape[0] | |
| else: | |
| batch = {key:value.cuda().to(dtype=torch.float32) for key, value in data.items()} | |
| return batch | |
| def encode_from_moment(self, z, deterministic=True): | |
| dist = DiagonalGaussianDistribution(z) | |
| init_latents = dist.mode() if deterministic else dist.sample() | |
| latents_mean = latents_std = None | |
| if hasattr(self.sd_pipe.vae.config, "latents_mean") and self.sd_pipe.vae.config.latents_mean is not None: | |
| latents_mean = torch.tensor(self.sd_pipe.vae.config.latents_mean).view(1, 4, 1, 1) | |
| if hasattr(self.sd_pipe.vae.config, "latents_std") and self.sd_pipe.vae.config.latents_std is not None: | |
| latents_std = torch.tensor(self.sd_pipe.vae.config.latents_std).view(1, 4, 1, 1) | |
| scaling_factor = self.sd_pipe.vae.config.scaling_factor | |
| if latents_mean is not None and latents_std is not None: | |
| latents_mean = latents_mean.to(device=z.device, dtype=z.dtype) | |
| latents_std = latents_std.to(device=z.device, dtype=z.dtype) | |
| init_latents = (init_latents - latents_mean) * scaling_factor / latents_std | |
| else: | |
| init_latents = init_latents * scaling_factor | |
| return init_latents | |
| def encode_first_stage(self, x, deterministic=False, center_input_sample=True): | |
| if center_input_sample: | |
| x = x * 2.0 - 1.0 | |
| latents_mean = latents_std = None | |
| if hasattr(self.sd_pipe.vae.config, "latents_mean") and self.sd_pipe.vae.config.latents_mean is not None: | |
| latents_mean = torch.tensor(self.sd_pipe.vae.config.latents_mean).view(1, -1, 1, 1) | |
| if hasattr(self.sd_pipe.vae.config, "latents_std") and self.sd_pipe.vae.config.latents_std is not None: | |
| latents_std = torch.tensor(self.sd_pipe.vae.config.latents_std).view(1, -1, 1, 1) | |
| if deterministic: | |
| partial_encode = lambda xx: self.sd_pipe.vae.encode(xx).latent_dist.mode() | |
| else: | |
| partial_encode = lambda xx: self.sd_pipe.vae.encode(xx).latent_dist.sample() | |
| trunk_size = self.configs.sd_pipe.vae_split | |
| if trunk_size < x.shape[0]: | |
| init_latents = torch.cat([partial_encode(xx) for xx in x.split(trunk_size, 0)], dim=0) | |
| else: | |
| init_latents = partial_encode(x) | |
| scaling_factor = self.sd_pipe.vae.config.scaling_factor | |
| if latents_mean is not None and latents_std is not None: | |
| latents_mean = latents_mean.to(device=x.device, dtype=x.dtype) | |
| latents_std = latents_std.to(device=x.device, dtype=x.dtype) | |
| init_latents = (init_latents - latents_mean) * scaling_factor / latents_std | |
| else: | |
| init_latents = init_latents * scaling_factor | |
| return init_latents | |
| def decode_first_stage(self, z, clamp=True): | |
| z = z / self.sd_pipe.vae.config.scaling_factor | |
| trunk_size = 1 | |
| if trunk_size < z.shape[0]: | |
| out = torch.cat( | |
| [self.sd_pipe.vae.decode(xx).sample for xx in z.split(trunk_size, 0)], dim=0, | |
| ) | |
| else: | |
| out = self.sd_pipe.vae.decode(z).sample | |
| if clamp: | |
| out = out.clamp(-1.0, 1.0) | |
| return out | |
| def get_loss_from_discrimnator(self, logits_fake): | |
| if not (isinstance(logits_fake, list) or isinstance(logits_fake, tuple)): | |
| g_loss = -torch.mean(logits_fake, dim=list(range(1, logits_fake.ndim))) | |
| else: | |
| g_loss = -torch.mean(logits_fake[0], dim=list(range(1, logits_fake[0].ndim))) | |
| for current_logits in logits_fake[1:]: | |
| g_loss += -torch.mean(current_logits, dim=list(range(1, current_logits.ndim))) | |
| g_loss /= len(logits_fake) | |
| return g_loss | |
| def training_step(self, data): | |
| current_bs = data['gt'].shape[0] | |
| micro_bs = self.configs.train.microbatch | |
| num_grad_accumulate = math.ceil(current_bs / micro_bs) | |
| # grad zero | |
| self.model.zero_grad() | |
| # update generator | |
| if self.configs.train.loss_coef.get('ldis', 0) > 0: | |
| self.freeze_model(self.discriminator) # freeze discriminator | |
| z0_pred_list = [] | |
| tt_list = [] | |
| prompt_embeds_list = [] | |
| for jj in range(0, current_bs, micro_bs): | |
| micro_data = {key:value[jj:jj+micro_bs] for key, value in data.items()} | |
| last_batch = (jj+micro_bs >= current_bs) | |
| if last_batch or self.num_gpus <= 1: | |
| losses, z0_pred, zt_noisy, tt = self.backward_step(micro_data, num_grad_accumulate) | |
| else: | |
| with self.model.no_sync(): | |
| losses, z0_pred, zt_noisy, tt = self.backward_step(micro_data, num_grad_accumulate) | |
| if self.configs.train.loss_coef.get('ldis', 0) > 0: | |
| z0_pred_list.append(z0_pred.detach()) | |
| tt_list.append(tt) | |
| prompt_embeds_list.append(self.prompt_embeds.detach()) | |
| if self.configs.train.use_amp: | |
| self.amp_scaler.step(self.optimizer) | |
| self.amp_scaler.update() | |
| else: | |
| self.optimizer.step() | |
| # update discriminator | |
| if (self.configs.train.loss_coef.get('ldis', 0) > 0 and | |
| (self.current_iters < self.configs.train.dis_init_iterations | |
| or self.current_iters % self.configs.train.dis_update_freq == 0) | |
| ): | |
| # grad zero | |
| self.unfreeze_model(self.discriminator) # update discriminator | |
| self.discriminator.zero_grad() | |
| for ii, jj in enumerate(range(0, current_bs, micro_bs)): | |
| micro_data = {key:value[jj:jj+micro_bs] for key, value in data.items()} | |
| last_batch = (jj+micro_bs >= current_bs) | |
| target = micro_data['gt_latent'] | |
| inputs = z0_pred_list[ii] | |
| if last_batch or self.num_gpus <= 1: | |
| logits = self.dis_backward_step(target, inputs, tt_list[ii], prompt_embeds_list[ii]) | |
| else: | |
| with self.discriminator.no_sync(): | |
| logits = self.dis_backward_step( | |
| target, inputs, tt_list[ii], prompt_embeds_list[ii] | |
| ) | |
| # make logging | |
| if self.current_iters % self.configs.train.dis_update_freq == 0 and self.rank == 0: | |
| ndim = logits[0].ndim | |
| losses['real'] = logits[0].detach().mean(dim=list(range(1, ndim))) | |
| losses['fake'] = logits[1].detach().mean(dim=list(range(1, ndim))) | |
| if self.configs.train.use_amp: | |
| self.amp_scaler_dis.step(self.optimizer_dis) | |
| self.amp_scaler_dis.update() | |
| else: | |
| self.optimizer_dis.step() | |
| # make logging | |
| if self.rank == 0: | |
| self.log_step_train( | |
| losses, tt, micro_data, z0_pred, zt_noisy, z0_gt=micro_data['gt_latent'], | |
| ) | |
| def log_step_train(self, losses, tt, micro_data, z0_pred, zt_noisy, z0_gt=None, phase='train'): | |
| ''' | |
| param losses: a dict recording the loss informations | |
| ''' | |
| ''' | |
| param loss: a dict recording the loss informations | |
| param micro_data: batch data | |
| param tt: 1-D tensor, time steps | |
| ''' | |
| if hasattr(self.configs.train, 'timesteps'): | |
| if len(self.configs.train.timesteps) < 3: | |
| record_steps = sorted(self.configs.train.timesteps) | |
| else: | |
| record_steps = [min(self.configs.train.timesteps), | |
| max(self.configs.train.timesteps)] | |
| else: | |
| max_inference_steps = self.configs.train.max_inference_steps | |
| record_steps = [1, max_inference_steps//2, max_inference_steps] | |
| if ((self.current_iters // self.configs.train.dis_update_freq) % | |
| (self.configs.train.log_freq[0] // self.configs.train.dis_update_freq) == 1): | |
| self.loss_mean = {key:torch.zeros(size=(len(record_steps),), dtype=torch.float64) | |
| for key in losses.keys() if key not in ['real', 'fake']} | |
| if self.configs.train.loss_coef.get('ldis', 0) > 0: | |
| self.logit_mean = {key:torch.zeros(size=(len(record_steps),), dtype=torch.float64) | |
| for key in ['real', 'fake']} | |
| self.loss_count = torch.zeros(size=(len(record_steps),), dtype=torch.float64) | |
| for jj in range(len(record_steps)): | |
| for key, value in losses.items(): | |
| index = record_steps[jj] - 1 | |
| mask = torch.where(tt == index, torch.ones_like(tt), torch.zeros_like(tt)) | |
| assert value.shape == mask.shape | |
| current_loss = torch.sum(value.detach() * mask) | |
| if key in ['real', 'fake']: | |
| self.logit_mean[key][jj] += current_loss.item() | |
| else: | |
| self.loss_mean[key][jj] += current_loss.item() | |
| self.loss_count[jj] += mask.sum().item() | |
| if ((self.current_iters // self.configs.train.dis_update_freq) % | |
| (self.configs.train.log_freq[0] // self.configs.train.dis_update_freq) == 0): | |
| if torch.any(self.loss_count == 0): | |
| self.loss_count += 1e-4 | |
| for key in losses.keys(): | |
| if key in ['real', 'fake']: | |
| self.logit_mean[key] /= self.loss_count | |
| else: | |
| self.loss_mean[key] /= self.loss_count | |
| log_str = f"Train: {self.current_iters:06d}/{self.configs.train.iterations:06d}, " | |
| valid_keys = sorted([key for key in losses.keys() if key not in ['loss', 'real', 'fake']]) | |
| for ii, key in enumerate(valid_keys): | |
| if ii == 0: | |
| log_str += f"{key}" | |
| else: | |
| log_str += f"/{key}" | |
| if self.configs.train.loss_coef.get('ldis', 0) > 0: | |
| log_str += "/real/fake:" | |
| else: | |
| log_str += ":" | |
| for jj, current_record in enumerate(record_steps): | |
| for ii, key in enumerate(valid_keys): | |
| if ii == 0: | |
| if key in ['dis', 'ldis']: | |
| log_str += 't({:d}):{:+6.4f}'.format( | |
| current_record, | |
| self.loss_mean[key][jj].item(), | |
| ) | |
| elif key in ['lpips', 'ldif']: | |
| log_str += 't({:d}):{:4.2f}'.format( | |
| current_record, | |
| self.loss_mean[key][jj].item(), | |
| ) | |
| elif key == 'llpips': | |
| log_str += 't({:d}):{:5.3f}'.format( | |
| current_record, | |
| self.loss_mean[key][jj].item(), | |
| ) | |
| else: | |
| log_str += 't({:d}):{:.1e}'.format( | |
| current_record, | |
| self.loss_mean[key][jj].item(), | |
| ) | |
| else: | |
| if key in ['dis', 'ldis']: | |
| log_str += f"/{self.loss_mean[key][jj].item():+6.4f}" | |
| elif key in ['lpips', 'ldif']: | |
| log_str += f"/{self.loss_mean[key][jj].item():4.2f}" | |
| elif key == 'llpips': | |
| log_str += f"/{self.loss_mean[key][jj].item():5.3f}" | |
| else: | |
| log_str += f"/{self.loss_mean[key][jj].item():.1e}" | |
| if self.configs.train.loss_coef.get('ldis', 0) > 0: | |
| log_str += f"/{self.logit_mean['real'][jj].item():+4.2f}" | |
| log_str += f"/{self.logit_mean['fake'][jj].item():+4.2f}, " | |
| else: | |
| log_str += f", " | |
| log_str += 'lr:{:.1e}'.format(self.optimizer.param_groups[0]['lr']) | |
| self.logger.info(log_str) | |
| self.logging_metric(self.loss_mean, tag='Loss', phase=phase, add_global_step=True) | |
| if ((self.current_iters // self.configs.train.dis_update_freq) % | |
| (self.configs.train.log_freq[1] // self.configs.train.dis_update_freq) == 0): | |
| if zt_noisy is not None: | |
| xt_pred = self.decode_first_stage(zt_noisy.detach()) | |
| self.logging_image(xt_pred, tag='xt-noisy', phase=phase, add_global_step=False) | |
| if z0_pred is not None: | |
| x0_pred = self.decode_first_stage(z0_pred.detach()) | |
| self.logging_image(x0_pred, tag='x0-pred', phase=phase, add_global_step=False) | |
| if z0_gt is not None: | |
| x0_recon = self.decode_first_stage(z0_gt.detach()) | |
| self.logging_image(x0_recon, tag='x0-recons', phase=phase, add_global_step=False) | |
| if 'txt' in micro_data: | |
| self.logging_text(micro_data['txt'], phase=phase) | |
| self.logging_image(micro_data['lq'], tag='LQ', phase=phase, add_global_step=False) | |
| self.logging_image(micro_data['gt'], tag='GT', phase=phase, add_global_step=True) | |
| if ((self.current_iters // self.configs.train.dis_update_freq) % | |
| (self.configs.train.save_freq // self.configs.train.dis_update_freq) == 1): | |
| self.tic = time.time() | |
| if ((self.current_iters // self.configs.train.dis_update_freq) % | |
| (self.configs.train.save_freq // self.configs.train.dis_update_freq) == 0): | |
| self.toc = time.time() | |
| elaplsed = (self.toc - self.tic) | |
| self.logger.info(f"Elapsed time: {elaplsed:.2f}s") | |
| self.logger.info("="*100) | |
| def validation(self, phase='val'): | |
| torch.cuda.empty_cache() | |
| if not (self.configs.validate.use_ema and hasattr(self.configs.train, 'ema_rate')): | |
| self.model.eval() | |
| if self.configs.train.start_mode: | |
| start_noise_predictor = self.ema_model if self.configs.validate.use_ema else self.model | |
| intermediate_noise_predictor = None | |
| else: | |
| start_noise_predictor = self.start_model | |
| intermediate_noise_predictor = self.ema_model if self.configs.validate.use_ema else self.model | |
| num_iters_epoch = math.ceil(len(self.datasets[phase]) / self.configs.validate.batch) | |
| mean_psnr = mean_lpips = 0 | |
| for jj, data in enumerate(self.dataloaders[phase]): | |
| data = self.prepare_data(data, phase='val') | |
| with torch.amp.autocast('cuda'): | |
| xt_progressive, x0_progressive = self.sample( | |
| image_lq=data['lq'], | |
| prompt=[_positive,]*data['lq'].shape[0], | |
| target_size=tuple(data['gt'].shape[-2:]), | |
| start_noise_predictor=start_noise_predictor, | |
| intermediate_noise_predictor=intermediate_noise_predictor, | |
| ) | |
| x0 = xt_progressive[-1] | |
| num_inference_steps = len(xt_progressive) | |
| if 'gt' in data: | |
| if not hasattr(self, 'psnr_metric'): | |
| self.psnr_metric = pyiqa.create_metric( | |
| 'psnr', | |
| test_y_channel=self.configs.train.get('val_y_channel', True), | |
| color_space='ycbcr', | |
| device=torch.device("cuda"), | |
| ) | |
| if not hasattr(self, 'lpips_metric'): | |
| self.lpips_metric = pyiqa.create_metric( | |
| 'lpips-vgg', | |
| device=torch.device("cuda"), | |
| as_loss=False, | |
| ) | |
| x0_normalize = util_image.normalize_th(x0, mean=0.5, std=0.5, reverse=True) | |
| mean_psnr += self.psnr_metric(x0_normalize, data['gt']).sum().item() | |
| with torch.amp.autocast('cuda'), torch.no_grad(): | |
| mean_lpips += self.lpips_metric(x0_normalize, data['gt']).sum().item() | |
| if (jj + 1) % self.configs.validate.log_freq == 0: | |
| self.logger.info(f'Validation: {jj+1:02d}/{num_iters_epoch:02d}...') | |
| self.logging_image(data['gt'], tag='GT', phase=phase, add_global_step=False) | |
| xt_progressive = rearrange(torch.cat(xt_progressive, dim=1), 'b (k c) h w -> (b k) c h w', c=3) | |
| self.logging_image( | |
| xt_progressive, | |
| tag='sample-progress', | |
| phase=phase, | |
| add_global_step=False, | |
| nrow=num_inference_steps, | |
| ) | |
| x0_progressive = rearrange(torch.cat(x0_progressive, dim=1), 'b (k c) h w -> (b k) c h w', c=3) | |
| self.logging_image( | |
| x0_progressive, | |
| tag='x0-progress', | |
| phase=phase, | |
| add_global_step=False, | |
| nrow=num_inference_steps, | |
| ) | |
| self.logging_image(data['lq'], tag='LQ', phase=phase, add_global_step=True) | |
| if 'gt' in data: | |
| mean_psnr /= len(self.datasets[phase]) | |
| mean_lpips /= len(self.datasets[phase]) | |
| self.logger.info(f'Validation Metric: PSNR={mean_psnr:5.2f}, LPIPS={mean_lpips:6.4f}...') | |
| self.logging_metric(mean_psnr, tag='PSNR', phase=phase, add_global_step=False) | |
| self.logging_metric(mean_lpips, tag='LPIPS', phase=phase, add_global_step=True) | |
| self.logger.info("="*100) | |
| if not (self.configs.validate.use_ema and hasattr(self.configs.train, 'ema_rate')): | |
| self.model.train() | |
| torch.cuda.empty_cache() | |
| def backward_step(self, micro_data, num_grad_accumulate): | |
| loss_coef = self.configs.train.loss_coef | |
| losses = {} | |
| z0_gt = micro_data['gt_latent'] | |
| tt = torch.tensor( | |
| random.choices(self.configs.train.timesteps, k=z0_gt.shape[0]), | |
| dtype=torch.int64, | |
| device=f"cuda:{self.rank}", | |
| ) - 1 | |
| with torch.autocast(device_type="cuda", enabled=self.configs.train.use_amp): | |
| model_pred = self.model( | |
| micro_data['lq'], tt, sample_posterior=False, center_input_sample=True, | |
| ) | |
| z0_pred, zt_noisy_pred, z0_lq = self.sd_forward_step( | |
| prompt=micro_data['txt'], | |
| latents_hq=micro_data['gt_latent'], | |
| image_lq=micro_data['lq'], | |
| image_hq=micro_data['gt'], | |
| model_pred=model_pred, | |
| timesteps=tt, | |
| ) | |
| # diffusion loss | |
| if loss_coef.get('ldif', 0) > 0: | |
| if self.configs.train.loss_type == 'L2': | |
| ldif_loss = F.mse_loss(z0_pred, z0_gt, reduction='none') | |
| elif self.configs.train.loss_type == 'L1': | |
| ldif_loss = F.l1_loss(z0_pred, z0_gt, reduction='none') | |
| else: | |
| raise TypeError(f"Unsupported Loss type for Diffusion: {self.configs.train.loss_type}") | |
| ldif_loss = torch.mean(ldif_loss, dim=list(range(1, z0_gt.ndim))) | |
| losses['ldif'] = ldif_loss * loss_coef['ldif'] | |
| # Gaussian constraints | |
| if loss_coef.get('kl', 0) > 0: | |
| losses['kl'] = model_pred.kl() * loss_coef['kl'] | |
| if loss_coef.get('pkl', 0) > 0: | |
| losses['pkl'] = model_pred.partial_kl() * loss_coef['pkl'] | |
| if loss_coef.get('rkl', 0) > 0: | |
| other = Box( | |
| {'mean': z0_gt-z0_lq, | |
| 'var':torch.ones_like(z0_gt), | |
| 'logvar':torch.zeros_like(z0_gt)} | |
| ) | |
| losses['rkl'] = model_pred.kl(other) * loss_coef['rkl'] | |
| # discriminator loss | |
| if loss_coef.get('ldis', 0) > 0: | |
| if self.current_iters > self.configs.train.dis_init_iterations: | |
| logits_fake = self.discriminator( | |
| torch.clamp(z0_pred, min=_Latent_bound['min'], max=_Latent_bound['max']), | |
| timestep=tt, | |
| encoder_hidden_states=self.prompt_embeds, | |
| ) | |
| losses['ldis'] = self.get_loss_from_discrimnator(logits_fake) * loss_coef['ldis'] | |
| else: | |
| losses['ldis'] = torch.zeros((z0_gt.shape[0], ), dtype=torch.float32).cuda() | |
| # perceptual loss | |
| if loss_coef.get('llpips', 0) > 0: | |
| losses['llpips'] = self.llpips_loss(z0_pred, z0_gt).view(-1) * loss_coef['llpips'] | |
| for key in ['ldif', 'kl', 'rkl', 'pkl', 'ldis', 'llpips']: | |
| if loss_coef.get(key, 0) > 0: | |
| if not 'loss' in losses: | |
| losses['loss'] = losses[key] | |
| else: | |
| losses['loss'] = losses['loss'] + losses[key] | |
| loss = losses['loss'].mean() / num_grad_accumulate | |
| if self.amp_scaler is None: | |
| loss.backward() | |
| else: | |
| self.amp_scaler.scale(loss).backward() | |
| return losses, z0_pred, zt_noisy_pred, tt | |
| def dis_backward_step(self, target, inputs, tt, prompt_embeds): | |
| with torch.autocast(device_type="cuda", enabled=self.configs.train.use_amp): | |
| logits_real = self.discriminator(target, tt, prompt_embeds) | |
| inputs = inputs.clamp(min=_Latent_bound['min'], max=_Latent_bound['max']) | |
| logits_fake = self.discriminator(inputs, tt, prompt_embeds) | |
| loss = hinge_d_loss(logits_real, logits_fake).mean() | |
| if self.amp_scaler_dis is None: | |
| loss.backward() | |
| else: | |
| self.amp_scaler_dis.scale(loss).backward() | |
| return logits_real[-1], logits_fake[-1] | |
| def scale_sd_input( | |
| self, | |
| x:torch.Tensor, | |
| sigmas: torch.Tensor = None, | |
| timestep: torch.Tensor = None, | |
| ) : | |
| if sigmas is None: | |
| if not self.sd_pipe.scheduler.sigmas.numel() == (self.configs.sd_pipe.num_train_steps + 1): | |
| self.sd_pipe.scheduler = EulerDiscreteScheduler.from_pipe( | |
| self.configs.sd_pipe.params.pretrained_model_name_or_path, | |
| cache_dir=self.configs.sd_pipe.params.cache_dir, | |
| subfolder='scheduler', | |
| ) | |
| assert self.sd_pipe.scheduler.sigmas.numel() == (self.configs.sd_pipe.num_train_steps + 1) | |
| sigmas = self.sd_pipe.scheduler.sigmas.flip(0).to(x.device)[timestep] # (b,) | |
| sigmas = append_dims(sigmas, x.ndim) | |
| if sigmas.ndim < x.ndim: | |
| sigmas = append_dims(sigmas, x.ndim) | |
| out = x / ((sigmas**2 + 1) ** 0.5) | |
| return out | |
| def prepare_lq_latents( | |
| self, | |
| image_lq: torch.Tensor, | |
| timestep: torch.Tensor, | |
| height: int = 512, | |
| width: int = 512, | |
| start_noise_predictor: torch.nn.Module = None, | |
| ): | |
| """ | |
| Input: | |
| image_lq: low-quality image, torch.Tensor, range in [0, 1] | |
| hight, width: resolution for high-quality image | |
| """ | |
| image_lq_up = F.interpolate(image_lq, size=(height, width), mode='bicubic') | |
| init_latents = self.encode_first_stage( | |
| image_lq_up, deterministic=False, center_input_sample=True, | |
| ) | |
| if start_noise_predictor is None: | |
| model_pred = None | |
| else: | |
| model_pred = start_noise_predictor( | |
| image_lq, timestep, sample_posterior=False, center_input_sample=True, | |
| ) | |
| # get latents | |
| sigmas = self.sigmas_cache[timestep] | |
| sigmas = append_dims(sigmas, init_latents.ndim) | |
| latents = self.add_noise(init_latents, sigmas, model_pred) | |
| return latents | |
| def add_noise(self, latents, sigmas, model_pred=None): | |
| if sigmas.ndim < latents.ndim: | |
| sigmas = append_dims(sigmas, latents.ndim) | |
| if model_pred is None: | |
| noise = torch.randn_like(latents) | |
| zt_noisy = latents + sigmas * noise | |
| else: | |
| if self.configs.train.loss_coef.get('rkl', 0) > 0: | |
| mean, std = model_pred.mean, model_pred.std | |
| zt_noisy = latents + mean + sigmas * std * torch.randn_like(latents) | |
| else: | |
| zt_noisy = latents + sigmas * model_pred.sample() | |
| return zt_noisy | |
| def retrieve_timesteps(self): | |
| device=torch.device(f"cuda:{self.rank}") | |
| num_inference_steps = self.configs.train.get('num_inference_steps', 5) | |
| timesteps = np.linspace( | |
| max(self.configs.train.timesteps), 0, num_inference_steps, | |
| endpoint=False, dtype=np.int64, | |
| ) - 1 | |
| timesteps = torch.from_numpy(timesteps).to(device) | |
| self.sd_pipe.scheduler.timesteps = timesteps | |
| sigmas = self.sigmas_cache[timesteps.long()] | |
| sigma_last = torch.tensor([0,], dtype=torch.float32).to(device=sigmas.device) | |
| sigmas = torch.cat([sigmas, sigma_last]).type(torch.float32) | |
| self.sd_pipe.scheduler.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication | |
| self.sd_pipe.scheduler._step_index = None | |
| self.sd_pipe.scheduler._begin_index = None | |
| return self.sd_pipe.scheduler.timesteps, num_inference_steps | |
| class TrainerSDTurboSR(TrainerBaseSR): | |
| def sd_forward_step( | |
| self, | |
| prompt: Union[str, List[str]] = None, | |
| latents_hq: Optional[torch.Tensor] = None, | |
| image_lq: torch.Tensor = None, | |
| image_hq: torch.Tensor = None, | |
| model_pred: DiagonalGaussianDistribution = None, | |
| timesteps: List[int] = None, | |
| **kwargs, | |
| ): | |
| r""" | |
| Function invoked when calling the pipeline for generation. | |
| Args: | |
| prompt (`str` or `List[str]`, *optional*): | |
| The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. | |
| instead. | |
| image_lq (`torch.Tensor`): The low-quality image(s) for enhancement, range in [0, 1]. | |
| image_hq (`torch.Tensor`): The high-quality image(s) for enhancement, range in [0, 1]. | |
| noise_pred (`torch.Tensor`): Predicted noise by the noise prediction model | |
| latents_hq (`torch.Tensor`, *optional*): | |
| Pre-generated high-quality latents, sampled from a Gaussian distribution, to be used as inputs for image | |
| generation. If not provided, a latents tensor will be generated by sampling using vae . | |
| timesteps (`List[int]`, *optional*): | |
| Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument | |
| in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is | |
| passed will be used. Must be in descending order. | |
| aesthetic_score (`float`, *optional*, defaults to 6.0): | |
| Used to simulate an aesthetic score of the generated image by influencing the positive text condition. | |
| Part of SDXL's micro-conditioning as explained in section 2.2 of | |
| [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). | |
| negative_aesthetic_score (`float`, *optional*, defaults to 2.5): | |
| Part of SDXL's micro-conditioning as explained in section 2.2 of | |
| [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). Can be used to | |
| simulate an aesthetic score of the generated image by influencing the negative text condition. | |
| """ | |
| device=torch.device(f"cuda:{self.rank}") | |
| # Encode input prompt | |
| prompt_embeds, negative_prompt_embeds = self.sd_pipe.encode_prompt( | |
| prompt=prompt, | |
| device=device, | |
| num_images_per_prompt=1, | |
| do_classifier_free_guidance=False, | |
| ) | |
| self.prompt_embeds = prompt_embeds | |
| # select the noise level, self.scheduler.sigmas, [1001,], descending | |
| if not hasattr(self, 'sigmas_cache'): | |
| assert self.sd_pipe.scheduler.sigmas.numel() == (self.configs.sd_pipe.num_train_steps + 1) | |
| self.sigmas_cache = self.sd_pipe.scheduler.sigmas.flip(0)[1:].to(device) #ascending,1000 | |
| sigmas = self.sigmas_cache[timesteps] # (b,) | |
| # Prepare input for SD | |
| height, width = image_hq.shape[-2:] | |
| if self.configs.train.start_mode: | |
| image_lq_up = F.interpolate(image_lq, size=(height, width), mode='bicubic') | |
| zt_clean = self.encode_first_stage( | |
| image_lq_up, center_input_sample=True, | |
| deterministic=self.configs.train.loss_coef.get('rkl', 0) > 0, | |
| ) | |
| else: | |
| if latents_hq is None: | |
| zt_clean = self.encode_first_stage( | |
| image_hq, center_input_sample=True, deterministic=False, | |
| ) | |
| else: | |
| zt_clean = latents_hq | |
| sigmas = append_dims(sigmas, zt_clean.ndim) | |
| zt_noisy = self.add_noise(zt_clean, sigmas, model_pred) | |
| prompt_embeds = prompt_embeds.to(device) | |
| zt_noisy_scale = self.scale_sd_input(zt_noisy, sigmas) | |
| eps_pred = self.sd_pipe.unet( | |
| zt_noisy_scale, | |
| timesteps, | |
| encoder_hidden_states=prompt_embeds, | |
| timestep_cond=None, | |
| cross_attention_kwargs=None, | |
| added_cond_kwargs=None, | |
| return_dict=False, | |
| )[0] # eps-mode for sdxl and sdxl-refiner | |
| if self.configs.train.noise_detach: | |
| z0_pred = zt_noisy.detach() - sigmas * eps_pred | |
| else: | |
| z0_pred = zt_noisy - sigmas * eps_pred | |
| return z0_pred, zt_noisy, zt_clean | |
| def sample( | |
| self, | |
| image_lq: torch.Tensor, | |
| prompt: Union[str, List[str]] = None, | |
| target_size: Tuple[int, int] = (1024, 1024), | |
| start_noise_predictor: torch.nn.Module = None, | |
| intermediate_noise_predictor: torch.nn.Module = None, | |
| **kwargs, | |
| ): | |
| r""" | |
| Function invoked when calling the pipeline for generation. | |
| Args: | |
| prompt (`str` or `List[str]`, *optional*): | |
| The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. | |
| instead. | |
| image_lq (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`): | |
| The image(s) to modify with the pipeline. | |
| target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): | |
| The required height and width of the super-resolved image. | |
| strength (`float`, *optional*, defaults to 0.3): | |
| Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image` | |
| will be used as a starting point, adding more noise to it the larger the `strength`. The number of | |
| denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will | |
| be maximum and the denoising process will run for the full number of iterations specified in | |
| `num_inference_steps`. A value of 1, therefore, essentially ignores `image`. Note that in the case of | |
| `denoising_start` being declared as an integer, the value of `strength` will be ignored. | |
| num_inference_steps (`int`, *optional*, defaults to 50): | |
| The number of denoising steps. More denoising steps usually lead to a higher quality image at the | |
| expense of slower inference. | |
| timesteps (`List[int]`, *optional*): | |
| Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument | |
| in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is | |
| passed will be used. Must be in descending order. | |
| negative_prompt (`str` or `List[str]`, *optional*): | |
| The prompt or prompts not to guide the image generation. If not defined, one has to pass | |
| `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is | |
| less than `1`). | |
| """ | |
| device=torch.device(f"cuda:{self.rank}") | |
| batch_size = image_lq.shape[0] | |
| # Encode input prompt | |
| prompt_embeds, negative_prompt_embeds = self.sd_pipe.encode_prompt( | |
| prompt=prompt, | |
| device=device, | |
| num_images_per_prompt=1, | |
| do_classifier_free_guidance=False, | |
| ) | |
| timesteps, num_inference_steps = self.retrieve_timesteps() | |
| latent_timestep = timesteps[:1].repeat(batch_size) | |
| # Prepare latent variables | |
| height, width = target_size | |
| latents = self.prepare_lq_latents(image_lq, latent_timestep.long(), height, width, start_noise_predictor) | |
| # Prepare extra step kwargs. | |
| extra_step_kwargs = self.sd_pipe.prepare_extra_step_kwargs(None, 0.0) | |
| prompt_embeds = prompt_embeds.to(device) | |
| x0_progressive = [] | |
| images_progressive = [] | |
| for i, t in enumerate(timesteps): | |
| latents_scaled = self.sd_pipe.scheduler.scale_model_input(latents, t) | |
| # predict the noise residual | |
| eps_pred = self.sd_pipe.unet( | |
| latents_scaled, | |
| t, | |
| encoder_hidden_states=prompt_embeds, | |
| timestep_cond=None, | |
| added_cond_kwargs=None, | |
| return_dict=False, | |
| )[0] | |
| z0_pred = latents - self.sigmas_cache[t.long()] * eps_pred | |
| # compute the previous noisy sample x_t -> x_t-1 | |
| if intermediate_noise_predictor is not None and i + 1 < len(timesteps): | |
| t_next = timesteps[i+1] | |
| noise = intermediate_noise_predictor(image_lq, t_next, center_input_sample=True) | |
| else: | |
| noise = None | |
| extra_step_kwargs['noise'] = noise | |
| latents = self.sd_pipe.scheduler.step(eps_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] | |
| image = self.decode_first_stage(latents) | |
| images_progressive.append(image) | |
| x0_pred = self.decode_first_stage(z0_pred) | |
| x0_progressive.append(x0_pred) | |
| return images_progressive, x0_progressive | |
| def my_worker_init_fn(worker_id): | |
| np.random.seed(np.random.get_state()[1][0] + worker_id) | |
| def hinge_d_loss( | |
| logits_real: Union[torch.Tensor, List[torch.Tensor,]], | |
| logits_fake: Union[torch.Tensor, List[torch.Tensor,]], | |
| ): | |
| def _hinge_d_loss(logits_real, logits_fake): | |
| loss_real = F.relu(1.0 - logits_real) | |
| loss_fake = F.relu(1.0 + logits_fake) | |
| d_loss = 0.5 * (loss_real + loss_fake) | |
| loss = d_loss.mean(dim=list(range(1, logits_real.ndim))) | |
| return loss | |
| if not (isinstance(logits_real, list) or isinstance(logits_real, tuple)): | |
| loss = _hinge_d_loss(logits_real, logits_fake) | |
| else: | |
| loss = _hinge_d_loss(logits_real[0], logits_fake[0]) | |
| for xx, yy in zip(logits_real[1:], logits_fake[1:]): | |
| loss += _hinge_d_loss(xx, yy) | |
| loss /= len(logits_real) | |
| return loss | |
| def get_torch_dtype(torch_dtype: str): | |
| if torch_dtype == 'torch.float16': | |
| return torch.float16 | |
| elif torch_dtype == 'torch.bfloat16': | |
| return torch.bfloat16 | |
| elif torch_dtype == 'torch.float32': | |
| return torch.float32 | |
| else: | |
| raise ValueError(f'Unexpected torch dtype:{torch_dtype}') | |