""" https://github.com/CompVis/stable-diffusion/blob/21f890f9da3cfbeaba8e2ac3c425ee9e998d5229/ldm/models/diffusion/ddpm.py#L30 """ import copy import functools import json import os from pathlib import Path from pdb import set_trace as st from typing import Any from click import prompt import einops import blobfile as bf import imageio import numpy as np import torch as th import torch.distributed as dist import torchvision from PIL import Image from torch.nn.parallel.distributed import DistributedDataParallel as DDP from torch.optim import AdamW from torch.utils.tensorboard.writer import SummaryWriter from tqdm import tqdm from guided_diffusion import dist_util, logger from guided_diffusion.fp16_util import MixedPrecisionTrainer from guided_diffusion.nn import update_ema from guided_diffusion.resample import LossAwareSampler, UniformSampler # from .train_util import TrainLoop3DRec from guided_diffusion.train_util import (TrainLoop, calc_average_loss, find_ema_checkpoint, find_resume_checkpoint, get_blob_logdir, log_loss_dict, log_rec3d_loss_dict, parse_resume_step_from_filename) from guided_diffusion.gaussian_diffusion import ModelMeanType from ldm.modules.encoders.modules import FrozenClipImageEmbedder, TextEmbedder, FrozenCLIPTextEmbedder, FrozenOpenCLIPImagePredictionEmbedder, FrozenOpenCLIPImageEmbedder import dnnlib from dnnlib.util import requires_grad from dnnlib.util import calculate_adaptive_weight from ..train_util_diffusion import TrainLoop3DDiffusion from ..cvD.nvsD_canoD import TrainLoop3DcvD_nvsD_canoD from guided_diffusion.continuous_diffusion_utils import get_mixed_prediction, different_p_q_objectives, kl_per_group_vada, kl_balancer # from .train_util_diffusion_lsgm_noD_joint import TrainLoop3DDiffusionLSGMJointnoD # joint diffusion and rec class # from .controlLDM import TrainLoop3DDiffusionLSGM_Control # joint diffusion and rec class from .train_util_diffusion_lsgm_noD_joint import TrainLoop3DDiffusionLSGMJointnoD # joint diffusion and rec class # ! add new schedulers from https://github.com/Stability-AI/generative-models from .crossattn_cldm import TrainLoop3DDiffusionLSGM_crossattn # import SD stuffs from typing import Any, Dict, List, Optional, Tuple, Union from contextlib import contextmanager from omegaconf import ListConfig, OmegaConf from sgm.modules import UNCONDITIONAL_CONFIG from sgm.util import (default, disabled_train, get_obj_from_str, instantiate_from_config, log_txt_as_img) from transport import create_transport, Sampler # from sgm.sampling_utils.demo.streamlit_helpers import init_sampling class FlowMatchingEngine(TrainLoop3DDiffusionLSGM_crossattn): def __init__( self, *, rec_model, denoise_model, diffusion, sde_diffusion, control_model, control_key, only_mid_control, loss_class, data, eval_data, batch_size, microbatch, lr, ema_rate, log_interval, eval_interval, save_interval, resume_checkpoint, resume_cldm_checkpoint=None, use_fp16=False, fp16_scale_growth=0.001, schedule_sampler=None, weight_decay=0, lr_anneal_steps=0, iterations=10001, ignore_resume_opt=False, freeze_ae=False, denoised_ae=True, triplane_scaling_divider=10, use_amp=False, diffusion_input_size=224, normalize_clip_encoding=False, scale_clip_encoding=1, cfg_dropout_prob=0, cond_key='img_sr', use_eos_feature=False, compile=False, snr_type='lognorm', # denoiser_config, # conditioner_config: Union[None, Dict, ListConfig, # OmegaConf] = None, # sampler_config: Union[None, Dict, ListConfig, OmegaConf] = None, # loss_fn_config: Union[None, Dict, ListConfig, OmegaConf] = None, **kwargs): super().__init__(rec_model=rec_model, denoise_model=denoise_model, diffusion=diffusion, sde_diffusion=sde_diffusion, control_model=control_model, control_key=control_key, only_mid_control=only_mid_control, loss_class=loss_class, data=data, eval_data=eval_data, batch_size=batch_size, microbatch=microbatch, lr=lr, ema_rate=ema_rate, log_interval=log_interval, eval_interval=eval_interval, save_interval=save_interval, resume_checkpoint=resume_checkpoint, resume_cldm_checkpoint=resume_cldm_checkpoint, use_fp16=use_fp16, fp16_scale_growth=fp16_scale_growth, schedule_sampler=schedule_sampler, weight_decay=weight_decay, lr_anneal_steps=lr_anneal_steps, iterations=iterations, ignore_resume_opt=ignore_resume_opt, freeze_ae=freeze_ae, denoised_ae=denoised_ae, triplane_scaling_divider=triplane_scaling_divider, use_amp=use_amp, diffusion_input_size=diffusion_input_size, normalize_clip_encoding=normalize_clip_encoding, scale_clip_encoding=scale_clip_encoding, cfg_dropout_prob=cfg_dropout_prob, cond_key=cond_key, use_eos_feature=use_eos_feature, compile=compile, **kwargs) # ! sgm diffusion pipeline # ! reuse the conditioner if self.cond_key == 'caption': ldm_configs = OmegaConf.load( 'sgm/configs/t23d-clipl-compat-fm.yaml')['ldm_configs'] else: assert 'lognorm' in snr_type if snr_type == 'lognorm': # by default ldm_configs = OmegaConf.load( 'sgm/configs/img23d-clipl-compat-fm-lognorm.yaml')['ldm_configs'] # elif snr_type == 'lognorm-mv': # ldm_configs = OmegaConf.load( # 'sgm/configs/mv23d-clipl-compat-fm-lognorm-noclip.yaml')['ldm_configs'] elif snr_type == 'lognorm-mv-plucker': ldm_configs = OmegaConf.load( # 'sgm/configs/mv23d-plucker-clipl-compat-fm-lognorm.yaml')['ldm_configs'] 'sgm/configs/mv23d-plucker-clipl-compat-fm-lognorm-noclip.yaml')['ldm_configs'] else: ldm_configs = OmegaConf.load( 'sgm/configs/img23d-clipl-compat-fm.yaml')['ldm_configs'] self.loss_fn = ( instantiate_from_config(ldm_configs.loss_fn_config) # if loss_fn_config is not None # else None ) # self.denoiser = instantiate_from_config( # ldm_configs.denoiser_config).to(dist_util.dev()) self.transport_sampler = Sampler(self.loss_fn.transport) self.conditioner = instantiate_from_config( default(ldm_configs.conditioner_config, UNCONDITIONAL_CONFIG)).to(dist_util.dev()) # ! setup optimizer (with cond embedder params here) self._setup_opt2() self._load_model2() def _setup_opt(self): pass # see below def _setup_opt2(self): # ! add trainable conditioner parameters # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/models/diffusion.py#L219 # params = list(self.ddpm_model.parameters()) self.opt = AdamW([{ 'name': 'ddpm', 'params': self.ddpm_model.parameters(), }, ], lr=self.lr, weight_decay=self.weight_decay) embedder_params = [] for embedder in self.conditioner.embedders: if embedder.is_trainable: embedder_params = embedder_params + list(embedder.parameters()) if len(embedder_params) != 0: self.opt.add_param_group( { 'name': 'embedder', 'params': embedder_params, 'lr': self.lr*0.1, # smaller lr to finetune dino/clip } ) # if self.train_vae: # for rec_param_group in self._init_optim_groups(self.rec_model): # self.opt.add_param_group(rec_param_group) print(self.opt) def save(self, mp_trainer=None, model_name='ddpm'): # save embedder params also super().save(mp_trainer, model_name) # save embedder ckpt if dist_util.get_rank() == 0: for embedder in self.conditioner.embedders: if embedder.is_trainable: # embedder_params = embedder_params + list(embedder.parameters()) model_name = embedder.__class__.__name__ filename = f"embedder_{model_name}{(self.step+self.resume_step):07d}.pt" with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: th.save(embedder.state_dict(), f) dist_util.synchronize() def _load_model2(self): # ! load embedder for embedder in self.conditioner.embedders: if embedder.is_trainable: # embedder_params = embedder_params + list(embedder.parameters()) model_name = embedder.__class__.__name__ filename = f"embedder_{model_name}{(self.step+self.resume_step):07d}.pt" # embedder_FrozenDinov2ImageEmbedderMV2115000.pt # with bf.BlobFile(bf.join(get_blob_logdir(), filename), # "wb") as f: # th.save(embedder.state_dict(), f) split = self.resume_checkpoint.split("model") resume_checkpoint = str( Path(split[0]) / filename) if os.path.exists(resume_checkpoint): if dist.get_rank() == 0: logger.log( f"loading cond embedder from checkpoint: {resume_checkpoint}...") # if model is None: # model = self.model embedder.load_state_dict( dist_util.load_state_dict( resume_checkpoint, map_location=dist_util.dev(), )) else: logger.log(f'{resume_checkpoint} not found.') if dist_util.get_world_size() > 1: dist_util.sync_params(embedder.parameters()) def instantiate_cond_stage(self, normalize_clip_encoding, scale_clip_encoding, cfg_dropout_prob, use_eos_feature=False): pass # placeholder function. initialized in the self.__init__() using SD api # ! already merged def prepare_ddpm(self, eps, mode='p'): raise NotImplementedError('already implemented in self.denoiser') # merged from noD.py # use sota denoiser, loss_fn etc. def ldm_train_step(self, batch, behaviour='cano', *args, **kwargs): """ add sds grad to all ae predicted x_0 """ # ! enable the gradient of both models requires_grad(self.ddpm_model, True) self.mp_trainer.zero_grad() # !!!! if 'img' in batch: batch_size = batch['img'].shape[0] else: batch_size = len(batch['caption']) for i in range(0, batch_size, self.microbatch): micro = { k: v[i:i + self.microbatch].to(dist_util.dev()) if isinstance( v, th.Tensor) else v for k, v in batch.items() } # move condition to self.dtype # =================================== ae part =================================== # with th.cuda.amp.autocast(dtype=th.bfloat16, with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): loss = th.tensor(0.).to(dist_util.dev()) assert 'latent' in micro vae_out = {self.latent_name: micro['latent']} # else: # vae_out = self.ddp_rec_model( # img=micro['img_to_encoder'], # c=micro['c'], # behaviour='encoder_vae', # ) # pred: (B, 3, 64, 64) eps = vae_out[self.latent_name] / self.triplane_scaling_divider # eps = vae_out.pop(self.latent_name) # if 'bg_plane' in vae_out: # eps = th.cat((eps, vae_out['bg_plane']), # dim=1) # include background, B 12+4 32 32 # ! SD loss # cond = self.get_c_input(micro, bs=eps.shape[0]) micro['img-c'] = { 'img': micro['img'].to(self.dtype), 'c': micro['c'].to(self.dtype), } loss, loss_other_info = self.loss_fn(self.ddp_ddpm_model, # self.denoiser, self.conditioner, eps.to(self.dtype), micro) # type: ignore loss = loss.mean() log_rec3d_loss_dict({}) log_rec3d_loss_dict({ # 'eps_mean': # eps.mean(), # 'eps_std': # eps.std([1, 2, 3]).mean(0), # 'pred_x0_std': # loss_other_info['model_output'].std([1, 2, 3]).mean(0), "p_loss": loss, }) self.mp_trainer.backward(loss) # joint gradient descent # update ddpm accordingly self.mp_trainer.optimize(self.opt) # ! directly eval_cldm() for sampling. # if dist_util.get_rank() == 0 and self.step % 500 == 0: # self.log_control_images(vae_out, micro, loss_other_info) @th.inference_mode() def log_control_images(self, vae_out, micro, ddpm_ret): if 'posterior' in vae_out: vae_out.pop('posterior') # for calculating kl loss vae_out_for_pred = {self.latent_name: vae_out[self.latent_name][0:1].to(self.dtype)} with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): pred = self.ddp_rec_model(latent=vae_out_for_pred, c=micro['c'][0:1], behaviour=self.render_latent_behaviour) assert isinstance(pred, dict) pred_img = pred['image_raw'] if 'img' in micro: gt_img = micro['img'] else: gt_img = th.zeros_like(pred['image_raw']) if 'depth' in micro: gt_depth = micro['depth'] if gt_depth.ndim == 3: gt_depth = gt_depth.unsqueeze(1) gt_depth = (gt_depth - gt_depth.min()) / (gt_depth.max() - gt_depth.min()) else: gt_depth = th.zeros_like(gt_img[:, 0:1, ...]) if 'image_depth' in pred: pred_depth = pred['image_depth'] pred_depth = (pred_depth - pred_depth.min()) / (pred_depth.max() - pred_depth.min()) else: pred_depth = th.zeros_like(gt_depth) gt_img = self.pool_128(gt_img) gt_depth = self.pool_128(gt_depth) # cond = self.get_c_input(micro) # hint = th.cat(cond['c_concat'], 1) gt_vis = th.cat( [ gt_img, gt_img, gt_img, # self.pool_128(hint), # gt_img, gt_depth.repeat_interleave(3, dim=1) ], dim=-1)[0:1] # TODO, fail to load depth. range [0, 1] # eps_t_p_3D = eps_t_p.reshape(batch_size, eps_t_p.shape[1]//3, 3, -1) # B C 3 L # self.sampler noised_latent, sigmas, x_start = [ ddpm_ret[k] for k in ['noised_input', 'sigmas', 'model_output'] ] noised_latent = { 'latent_normalized_2Ddiffusion': noised_latent[0:1].to(self.dtype) * self.triplane_scaling_divider, } denoised_latent = { 'latent_normalized_2Ddiffusion': x_start[0:1].to(self.dtype) * self.triplane_scaling_divider, } with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): noised_ae_pred = self.ddp_rec_model( img=None, c=micro['c'][0:1], latent=noised_latent, behaviour=self.render_latent_behaviour) # pred_x0 = self.sde_diffusion._predict_x0_from_eps( # eps_t_p, pred_eps_p, logsnr_p) # for VAE loss, denosied latent # pred_xstart_3D denoised_ae_pred = self.ddp_rec_model( img=None, c=micro['c'][0:1], latent=denoised_latent, # latent=pred_x0[0:1] * self. # triplane_scaling_divider, # TODO, how to define the scale automatically? behaviour=self.render_latent_behaviour) pred_vis = th.cat( [ self.pool_128(img) for img in ( pred_img[0:1], noised_ae_pred['image_raw'][0:1], denoised_ae_pred['image_raw'][0:1], # controlnet result pred_depth[0:1].repeat_interleave(3, dim=1)) ], dim=-1) # B, 3, H, W if 'img' in micro: vis = th.cat([gt_vis, pred_vis], dim=-2)[0].permute(1, 2, 0).cpu() # ! pred in range[-1, 1] else: vis = pred_vis[0].permute(1, 2, 0).cpu() # vis_grid = torchvision.utils.make_grid(vis) # HWC vis = vis.numpy() * 127.5 + 127.5 vis = vis.clip(0, 255).astype(np.uint8) img_save_path = f'{logger.get_dir()}/{self.step+self.resume_step}denoised_{sigmas[0].item():3}.jpg' Image.fromarray(vis).save(img_save_path) # if self.cond_key == 'caption': # with open(f'{logger.get_dir()}/{self.step+self.resume_step}caption_{t_p[0].item():3}.txt', 'w') as f: # f.write(micro['caption'][0]) print('log denoised vis to: ', img_save_path) th.cuda.empty_cache() @th.no_grad() def sample( self, cond: Dict, uc: Union[Dict, None] = None, batch_size: int = 16, shape: Union[None, Tuple, List] = None, use_cfg=True, # cfg_scale=4, # default value in SiT # cfg_scale=1.5, # default value in SiT cfg_scale=4.0, # default value in SiT **kwargs, ): # self.sampler sample_fn = self.transport_sampler.sample_ode(num_steps=250, cfg=True) # default ode sampling setting. # th.manual_seed(42) # reproducible zs = th.randn(batch_size, *shape).to(dist_util.dev()).to(self.dtype) assert use_cfg # sample_model_kwargs = {'uc': uc, 'cond': cond} model_fn = self.ddpm_model.forward_with_cfg # default # ! prepare_inputs in VanillaCFG, for compat issue c_out = {} for k in cond: if k in ["vector", "crossattn", "concat"]: c_out[k] = th.cat((cond[k], uc[k]), 0) else: assert cond[k] == uc[k] c_out[k] = cond[k] sample_model_kwargs = {'context': c_out, 'cfg_scale': cfg_scale} zs = th.cat([zs, zs], 0) with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): samples = sample_fn(zs, model_fn, **sample_model_kwargs)[-1] samples, _ = samples.chunk(2, dim=0) # Remove null class samples return samples @th.inference_mode() def eval_cldm( self, prompt="", save_img=False, use_train_trajectory=False, camera=None, num_samples=1, num_instances=1, unconditional_guidance_scale=4.0, # default value in neural ode export_mesh=False, **kwargs, ): # ! slightly modified for new API. combined with # /cpfs01/shared/V2V/V2V_hdd/yslan/Repo/generative-models/sgm/models/diffusion.py:249 log_images() # TODO, support batch_size > 1 self.ddpm_model.eval() # assert unconditional_guidance_scale == 4.0 args = dnnlib.EasyDict( dict( batch_size=1, image_size=self.diffusion_input_size, denoise_in_channels=self.rec_model.decoder.triplane_decoder. out_chans, # type: ignore clip_denoised=False, class_cond=False)) model_kwargs = {} uc = None log = dict() ucg_keys = [self.cond_key] # i23d sampling_kwargs = {'cfg_scale': unconditional_guidance_scale} N = num_samples # hard coded, to update z_shape = ( N, self.ddpm_model.in_channels if not self.ddpm_model.roll_out else 3 * self.ddpm_model.in_channels, # type: ignore self.diffusion_input_size, self.diffusion_input_size) data = iter(self.data) def sample_and_save(batch_c,idx=0): with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): c, uc = self.conditioner.get_unconditional_conditioning( batch_c, force_uc_zero_embeddings=ucg_keys if len(self.conditioner.embedders) > 0 else [], ) for k in c: if isinstance(c[k], th.Tensor): # c[k], uc[k] = map(lambda y: y[k][:N].to(dist_util.dev()), # (c, uc)) assert c[k].shape[0] == 1 c[k], uc[k] = map(lambda y: y[k].repeat_interleave(N, 0).to(dist_util.dev()), (c, uc)) # support bs>1 sampling given a condition samples = self.sample(c, shape=z_shape[1:], uc=uc, batch_size=N, **sampling_kwargs) # st() # do rendering first # ! get c (Path(logger.get_dir())/f'{self.step+self.resume_step}').mkdir(exist_ok=True, parents=True) if 'img' in self.cond_key: img_save_path = f'{logger.get_dir()}/{self.step+self.resume_step}/imgcond-{idx}.jpg' if 'c' in self.cond_key: torchvision.utils.save_image(batch_c['img'][0], img_save_path, value_range=(-1,1), normalize=True, padding=0) # torch.Size([24, 6, 3, 256, 256]) else: torchvision.utils.save_image(batch_c['img'], img_save_path, value_range=(-1,1), normalize=True, padding=0) assert camera is not None batch = {'c': camera.clone()[:24]} # rendering for i in range(samples.shape[0]): th.cuda.empty_cache() # ! render sampled latent name_prefix = f'idx-{idx}-cfg={unconditional_guidance_scale}_sample-{i}' if self.cond_key == 'caption': name_prefix = f'{name_prefix}_{prompt}' with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): self.render_video_given_triplane( samples[i:i+1].to(self.dtype), self.rec_model, # compatible with join_model name_prefix=name_prefix, save_img=save_img, render_reference=batch, export_mesh=export_mesh, render_all=True) if self.cond_key == 'caption': batch_c = {self.cond_key: prompt} sample_and_save(batch_c) else: for idx, batch in enumerate(data): # batch = next(data) # using same cond here if self.cond_key == 'img-c': batch_c = { self.cond_key: { 'img': batch['img'].to(self.dtype).to(dist_util.dev()), 'c': batch['c'].to(self.dtype).to(dist_util.dev()), }, 'img': batch['img'].to(self.dtype).to(dist_util.dev()) # required by clip } else: batch_c = {self.cond_key: batch[self.cond_key].to(dist_util.dev()).to(self.dtype)} sample_and_save(batch_c, idx) self.ddpm_model.train() @th.inference_mode() def eval_i23d_and_export( self, inp_img, num_steps=250, seed=42, mesh_size=192, mesh_thres=10, unconditional_guidance_scale=4.0, # default value in neural ode # camera, prompt="", save_img=False, use_train_trajectory=False, num_samples=1, num_instances=1, export_mesh=True, **kwargs, ): # output_model, output_video = './logs/LSGM/inference/Objaverse/i23d/dit-L2/gradio_app/mesh/cfg=4.0_sample-0.obj', './logs/LSGM/inference/Objaverse/i23d/dit-L2/gradio_app/triplane_cfg=4.0_sample-0.mp4' # return output_model, output_video logger.log( num_steps, unconditional_guidance_scale, seed, mesh_size, mesh_thres, ) camera = th.load('assets/objv_eval_pose.pt', map_location=dist_util.dev())[:] inp_img = th.from_numpy(inp_img).permute(2,0,1).unsqueeze(0) / 127.5 - 1 # to [-1,1] # for gradio demo self.ddpm_model.eval() # assert unconditional_guidance_scale == 4.0 args = dnnlib.EasyDict( dict( batch_size=1, image_size=self.diffusion_input_size, denoise_in_channels=self.rec_model.decoder.triplane_decoder. out_chans, # type: ignore clip_denoised=False, class_cond=False)) model_kwargs = {} uc = None log = dict() ucg_keys = [self.cond_key] # i23d sampling_kwargs = {'cfg_scale': unconditional_guidance_scale, 'num_steps': num_steps, 'seed': seed} N = num_samples # hard coded, to update z_shape = ( N, self.ddpm_model.in_channels if not self.ddpm_model.roll_out else 3 * self.ddpm_model.in_channels, # type: ignore self.diffusion_input_size, self.diffusion_input_size) # data = iter(self.data) assert camera is not None batch = {'c': camera.clone()[:24]} def sample_and_save(batch_c): with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): c, uc = self.conditioner.get_unconditional_conditioning( batch_c, force_uc_zero_embeddings=ucg_keys if len(self.conditioner.embedders) > 0 else [], ) for k in c: if isinstance(c[k], th.Tensor): # c[k], uc[k] = map(lambda y: y[k][:N].to(dist_util.dev()), # (c, uc)) assert c[k].shape[0] == 1 c[k], uc[k] = map(lambda y: y[k].repeat_interleave(N, 0).to(dist_util.dev()), (c, uc)) # support bs>1 sampling given a condition samples = self.sample(c, shape=z_shape[1:], uc=uc, batch_size=N, **sampling_kwargs) # rendering all_vid_dump_path = [] all_mesh_dump_path = [] for i in range(samples.shape[0]): th.cuda.empty_cache() # ! render sampled latent name_prefix = f'cfg_{unconditional_guidance_scale}_sample-{i}' if self.cond_key == 'caption': name_prefix = f'{name_prefix}_{prompt}' with th.cuda.amp.autocast(dtype=self.dtype, enabled=self.mp_trainer.use_amp): vid_dump_path, mesh_dump_path = self.render_video_given_triplane( samples[i:i+1].to(self.dtype), self.rec_model, # compatible with join_model name_prefix=name_prefix, save_img=save_img, render_reference=batch, export_mesh=export_mesh, render_all=True, mesh_size=mesh_size, mesh_thres=mesh_thres) all_vid_dump_path.append(vid_dump_path) all_mesh_dump_path.append(mesh_dump_path) # return all_vid_dump_path, all_mesh_dump_path return all_vid_dump_path[0], all_mesh_dump_path[0] # for compat issue if self.cond_key == 'caption': batch_c = {self.cond_key: prompt} return sample_and_save(batch_c) else: # for idx, batch in enumerate(data): # batch = next(data) # using same cond here # if self.cond_key == 'img-c': # batch_c = { # self.cond_key: { # 'img': batch['img'].to(self.dtype).to(dist_util.dev()), # 'c': batch['c'].to(self.dtype).to(dist_util.dev()), # }, # 'img': batch['img'].to(self.dtype).to(dist_util.dev()) # required by clip # } # else: batch_c = {self.cond_key: inp_img.to(dist_util.dev()).to(self.dtype)} return sample_and_save(batch_c)