Add files using upload-large-folder tool

.gitattributes

@@ -2984,3 +2984,4 @@ results/versatile_diffusion/subj01/97.png filter=lfs diff=lfs merge=lfs -text
 results/versatile_diffusion/subj01/roi/12.png filter=lfs diff=lfs merge=lfs -text
 results/versatile_diffusion/subj01/roi/2.png filter=lfs diff=lfs merge=lfs -text
 results/versatile_diffusion/subj01/roi/3.png filter=lfs diff=lfs merge=lfs -text
+results/versatile_diffusion/subj01/roi/4.png filter=lfs diff=lfs merge=lfs -text

versatile_diffusion/lib/data_factory/common/__init__.py

@@ -0,0 +1,6 @@
+from .ds_base import ds_base, collate, register as regdataset
+from .ds_loader import pre_loader_checkings, register as regloader
+from .ds_transform import TBase, have, register as regtrans
+from .ds_estimator import register as regestmat
+from .ds_formatter import register as regformat
+from .ds_sampler import register as regsampler

versatile_diffusion/lib/data_factory/common/ds_base.py

@@ -0,0 +1,280 @@
+import os
+import os.path as osp
+import numpy as np
+import numpy.random as npr
+import torch
+import torch.distributed as dist
+import torchvision
+import copy
+import itertools
+
+from ... import sync
+from ...cfg_holder import cfg_unique_holder as cfguh
+from ...log_service import print_log
+
+import torch.distributed as dist
+from multiprocessing import shared_memory
+
+# import multiprocessing
+# if hasattr(multiprocessing, "shared_memory"):
+#     from multiprocessing import shared_memory
+# else:
+#     # workaround for single gpu inference on colab
+#     shared_memory = None
+    
+import pickle
+import hashlib
+import random
+
+class ds_base(torch.utils.data.Dataset):
+    def __init__(self, 
+                 cfg, 
+                 loader = None, 
+                 estimator = None, 
+                 transforms = None, 
+                 formatter = None):
+
+        self.cfg = cfg
+        self.load_info = None
+        self.init_load_info()
+        self.loader = loader
+        self.transforms = transforms
+        self.formatter = formatter
+
+        if self.load_info is not None:
+            load_info_order_by = getattr(self.cfg, 'load_info_order_by', 'default')
+            if load_info_order_by == 'default':
+                self.load_info = sorted(self.load_info, key=lambda x:x['unique_id'])
+            else: 
+                try:
+                    load_info_order_by, reverse = load_info_order_by.split('|')
+                    reverse = reverse == 'reverse'
+                except:
+                    reverse = False
+                self.load_info = sorted(
+                    self.load_info, key=lambda x:x[load_info_order_by], reverse=reverse)
+
+        load_info_add_idx = getattr(self.cfg, 'load_info_add_idx', True)
+        if (self.load_info is not None) and load_info_add_idx:
+            for idx, info in enumerate(self.load_info):
+                info['idx'] = idx
+
+        if estimator is not None:
+            self.load_info = estimator(self.load_info)
+
+        self.try_sample = getattr(self.cfg, 'try_sample', None)
+        if self.try_sample is not None:
+            try: 
+                start, end = self.try_sample
+            except:
+                start, end = 0, self.try_sample
+            self.load_info = self.load_info[start:end]
+
+        self.repeat = getattr(self.cfg, 'repeat', 1)
+
+        pick = getattr(self.cfg, 'pick', None)
+        if pick is not None:
+            self.load_info = [i for i in self.load_info if i['filename'] in pick]
+
+        #########
+        # cache #
+        #########
+
+        self.cache_sm = getattr(self.cfg, 'cache_sm', False)
+        self.cache_cnt = 0
+        if self.cache_sm:
+            self.cache_pct = getattr(self.cfg, 'cache_pct', 0)
+            cache_unique_id = sync.nodewise_sync().random_sync_id()
+            self.cache_unique_id = hashlib.sha256(pickle.dumps(cache_unique_id)).hexdigest()
+            self.__cache__(self.cache_pct)
+
+        #######
+        # log #
+        #######
+
+        if self.load_info is not None:
+            console_info = '{}: '.format(self.__class__.__name__)
+            console_info += 'total {} unique images, '.format(len(self.load_info))
+            console_info += 'total {} unique sample. Cached {}. Repeat {} times.'.format(
+                len(self.load_info), self.cache_cnt, self.repeat)
+        else:
+            console_info = '{}: load_info not ready.'.format(self.__class__.__name__)
+        print_log(console_info)
+
+    def init_load_info(self):
+        # implement by sub class
+        pass
+
+    def __len__(self):
+        return len(self.load_info)*self.repeat
+
+    def __cache__(self, pct):
+        if pct == 0:
+            self.cache_cnt = 0
+            return
+        self.cache_cnt = int(len(self.load_info)*pct)
+        if not self.cache_sm:
+            for i in range(self.cache_cnt):
+                self.load_info[i] = self.loader(self.load_info[i])
+            return
+
+        for i in range(self.cache_cnt):
+            shm_name = str(self.load_info[i]['unique_id']) + '_' + self.cache_unique_id
+            if i % self.local_world_size == self.local_rank:
+                data = pickle.dumps(self.loader(self.load_info[i]))
+                datan = len(data)
+                # self.print_smname_to_file(shm_name)
+                shm = shared_memory.SharedMemory(
+                    name=shm_name, create=True, size=datan)
+                shm.buf[0:datan] = data[0:datan]
+                shm.close()
+                self.load_info[i] = shm_name
+            else:
+                self.load_info[i] = shm_name
+        dist.barrier()
+
+    def __getitem__(self, idx):
+        idx = idx%len(self.load_info)
+        # element = copy.deepcopy(self.load_info[idx])
+
+        # 0730 try shared memory
+        element = copy.deepcopy(self.load_info[idx])
+        if isinstance(element, str):
+            shm = shared_memory.SharedMemory(name=element)
+            element = pickle.loads(shm.buf)
+            shm.close()
+        else:
+            element = copy.deepcopy(element)
+            element['load_info_ptr'] = self.load_info
+
+        if idx >= self.cache_cnt:
+            element = self.loader(element)
+        if self.transforms is not None:
+            element = self.transforms(element)
+        if self.formatter is not None:
+            return self.formatter(element)
+        else:
+            return element
+
+    # 0730 try shared memory
+    def __del__(self):
+        # Clean the shared memory
+        for infoi in self.load_info:
+            if isinstance(infoi, str) and (self.local_rank==0):
+                shm = shared_memory.SharedMemory(name=infoi)
+                shm.close()
+                shm.unlink()
+
+    def print_smname_to_file(self, smname):
+        try:
+            log_file = cfguh().cfg.train.log_file
+        except:
+            try:
+                log_file = cfguh().cfg.eval.log_file
+            except:
+                raise ValueError
+        # a trick to use the log_file path
+        sm_file = log_file.replace('.log', '.smname')
+        with open(sm_file, 'a') as f:
+            f.write(smname + '\n')
+
+def singleton(class_):
+    instances = {}
+    def getinstance(*args, **kwargs):
+        if class_ not in instances:
+            instances[class_] = class_(*args, **kwargs)
+        return instances[class_]
+    return getinstance
+
+from .ds_loader import get_loader
+from .ds_transform import get_transform
+from .ds_estimator import get_estimator
+from .ds_formatter import get_formatter
+
+@singleton
+class get_dataset(object):
+    def __init__(self):
+        self.dataset = {}
+
+    def register(self, ds):
+        self.dataset[ds.__name__] = ds
+
+    def __call__(self, cfg):
+        if cfg is None:
+            return None
+        t = cfg.type
+        if t is None:
+            return None
+        elif t in ['laion2b', 'laion2b_dummy', 
+                   'laion2b_webdataset', 
+                   'laion2b_webdataset_sdofficial', ]:
+            from .. import ds_laion2b
+        elif t in ['coyo', 'coyo_dummy', 
+                   'coyo_webdataset', ]:
+            from .. import ds_coyo_webdataset
+        elif t in ['laionart', 'laionart_dummy', 
+                   'laionart_webdataset', ]:
+            from .. import ds_laionart
+        elif t in ['celeba']:
+            from .. import ds_celeba
+        elif t in ['div2k']:
+            from .. import ds_div2k
+        elif t in ['pafc']:
+            from .. import ds_pafc
+        elif t in ['coco_caption']:
+            from .. import ds_coco
+        else:
+            raise ValueError
+
+        loader    = get_loader()   (cfg.get('loader'   , None))
+        transform = get_transform()(cfg.get('transform', None))
+        estimator = get_estimator()(cfg.get('estimator', None))
+        formatter = get_formatter()(cfg.get('formatter', None))
+
+        return self.dataset[t](
+            cfg, loader, estimator, 
+            transform, formatter)
+
+def register():
+    def wrapper(class_):
+        get_dataset().register(class_)
+        return class_
+    return wrapper
+
+# some other helpers
+
+class collate(object):
+    """
+        Modified from torch.utils.data._utils.collate
+        It handle list different from the default.
+            List collate just by append each other.
+    """
+    def __init__(self):
+        self.default_collate = \
+            torch.utils.data._utils.collate.default_collate
+
+    def __call__(self, batch):
+        """
+        Args:
+            batch: [data, data] -or- [(data1, data2, ...), (data1, data2, ...)]
+        This function will not be used as induction function
+        """
+        elem = batch[0]
+        if not (elem, (tuple, list)):
+            return self.default_collate(batch)
+        
+        rv = []
+        # transposed
+        for i in zip(*batch):
+            if isinstance(i[0], list):
+                if len(i[0]) != 1:
+                    raise ValueError
+                try:
+                    i = [[self.default_collate(ii).squeeze(0)] for ii in i]
+                except:
+                    pass
+                rvi = list(itertools.chain.from_iterable(i))
+                rv.append(rvi) # list concat
+            else:
+                rv.append(self.default_collate(i))
+        return rv

versatile_diffusion/lib/data_factory/common/ds_estimator.py

@@ -0,0 +1,85 @@
+import os.path as osp
+import numpy as np
+import numpy.random as npr
+import PIL
+import cv2
+
+import torch
+import torchvision
+import xml.etree.ElementTree as ET
+import json
+import copy
+import math
+
+def singleton(class_):
+    instances = {}
+    def getinstance(*args, **kwargs):
+        if class_ not in instances:
+            instances[class_] = class_(*args, **kwargs)
+        return instances[class_]
+    return getinstance
+
+@singleton
+class get_estimator(object):
+    def __init__(self):
+        self.estimator = {}
+
+    def register(self, estimf):
+        self.estimator[estimf.__name__] = estimf
+
+    def __call__(self, cfg):
+        if cfg is None:
+            return None
+        t = cfg.type
+        return self.estimator[t](**cfg.args)
+
+def register():
+    def wrapper(class_):
+        get_estimator().register(class_)
+        return class_
+    return wrapper
+
+@register()
+class PickFileEstimator(object):
+    """
+    This is an estimator that filter load_info
+        using the provided filelist
+    """
+    def __init__(self, 
+                 filelist = None,
+                 repeat_n = 1):
+        """
+        Args:
+            filelist: a list of string gives the name of images 
+                we would like to visualize, evaluate or train. 
+            repeat_n: int, times these images will be repeated
+        """
+        self.filelist = filelist
+        self.repeat_n = repeat_n
+
+    def __call__(self, load_info):
+        load_info_new = []
+        for info in load_info:
+            if os.path.basename(info['image_path']).split('.')[0] in self.filelist:
+                load_info_new.append(info)
+        return load_info_new * self.repeat_n
+
+@register()
+class PickIndexEstimator(object):
+    """
+    This is an estimator that filter load_info
+        using the provided indices
+    """
+    def __init__(self, 
+                 indexlist = None,
+                 **kwargs):
+        """
+        Args:
+            indexlist: [] of int.
+                the indices to be filtered out. 
+        """
+        self.indexlist = indexlist
+
+    def __call__(self, load_info):
+        load_info_new = [load_info[i] for i in self.indexlist]
+        return load_info_new

versatile_diffusion/lib/data_factory/common/ds_formatter.py

@@ -0,0 +1,39 @@
+import os
+import os.path as osp
+import numpy as np
+import numpy.random as npr
+import torch
+import cv2
+import scipy.ndimage
+from PIL import Image
+import copy
+import gc
+import itertools
+
+def singleton(class_):
+    instances = {}
+    def getinstance(*args, **kwargs):
+        if class_ not in instances:
+            instances[class_] = class_(*args, **kwargs)
+        return instances[class_]
+    return getinstance
+
+@singleton
+class get_formatter(object):
+    def __init__(self):
+        self.formatter = {}
+
+    def register(self, formatf):
+        self.formatter[formatf.__name__] = formatf
+
+    def __call__(self, cfg):
+        if cfg is None:
+            return None
+        t = cfg.type
+        return self.formatter[t](**cfg.args)
+
+def register():
+    def wrapper(class_):
+        get_formatter().register(class_)
+        return class_
+    return wrapper

versatile_diffusion/lib/data_factory/common/ds_loader.py

@@ -0,0 +1,97 @@
+import os.path as osp
+import numpy as np
+import numpy.random as npr
+import PIL
+import cv2
+
+import torch
+import torchvision
+import xml.etree.ElementTree as ET
+import json
+import copy
+
+from ...cfg_holder import cfg_unique_holder as cfguh
+
+def singleton(class_):
+    instances = {}
+    def getinstance(*args, **kwargs):
+        if class_ not in instances:
+            instances[class_] = class_(*args, **kwargs)
+        return instances[class_]
+    return getinstance
+
+@singleton
+class get_loader(object):
+    def __init__(self):
+        self.loader = {}
+
+    def register(self, loadf):
+        self.loader[loadf.__name__] = loadf
+
+    def __call__(self, cfg):
+        if cfg is None:
+            return None
+        if isinstance(cfg, list):
+            loader = []
+            for ci in cfg:
+                t = ci.type
+                loader.append(self.loader[t](**ci.args))
+            return compose(loader)
+        t = cfg.type
+        return self.loader[t](**cfg.args)
+
+class compose(object):
+    def __init__(self, loaders):
+        self.loaders = loaders
+
+    def __call__(self, element):
+        for l in self.loaders:
+            element = l(element)
+        return element
+    
+    def __getitem__(self, idx):
+        return self.loaders[idx]
+
+def register():
+    def wrapper(class_):
+        get_loader().register(class_)
+        return class_
+    return wrapper
+
+def pre_loader_checkings(ltype):
+    lpath = ltype+'_path'
+    # cache feature added on 20201021
+    lcache = ltype+'_cache'
+    def wrapper(func):
+        def inner(self, element):
+            if lcache in element:
+                # cache feature added on 20201021
+                data = element[lcache]
+            else:
+                if ltype in element:
+                    raise ValueError
+                if lpath not in element:
+                    raise ValueError
+
+                if element[lpath] is None:
+                    data = None
+                else:
+                    data = func(self, element[lpath], element)
+            element[ltype] = data
+
+            if ltype == 'image':
+                if isinstance(data, np.ndarray):
+                    imsize = data.shape[-2:]
+                elif isinstance(data, PIL.Image.Image):
+                    imsize = data.size[::-1]
+                elif isinstance(data, torch.Tensor):
+                    imsize = [data.size(-2), data.size(-1)]
+                elif data is None:
+                    imsize = None
+                else:
+                    raise ValueError
+                element['imsize'] = imsize
+                element['imsize_current'] = copy.deepcopy(imsize)
+            return element
+        return inner
+    return wrapper

versatile_diffusion/lib/data_factory/common/ds_sampler.py

@@ -0,0 +1,273 @@
+from tokenize import group
+import torch
+import numpy as np
+import numpy.random as npr
+import torch.distributed as dist
+import math
+
+from ...log_service import print_log
+from ... import sync
+
+def singleton(class_):
+    instances = {}
+    def getinstance(*args, **kwargs):
+        if class_ not in instances:
+            instances[class_] = class_(*args, **kwargs)
+        return instances[class_]
+    return getinstance
+
+@singleton
+class get_sampler(object):
+    def __init__(self):
+        self.sampler = {}
+
+    def register(self, sampler):
+        self.sampler[sampler.__name__] = sampler
+
+    def __call__(self, dataset, cfg):
+        if cfg == 'default_train':
+            return GlobalDistributedSampler(dataset, shuffle=True, extend=False)
+        elif cfg == 'default_eval':
+            return GlobalDistributedSampler(dataset, shuffle=False, extend=True)
+        else:
+            t = cfg.type
+            return self.sampler[t](dataset=dataset, **cfg.args)
+
+def register():
+    def wrapper(class_):
+        get_sampler().register(class_)
+        return class_
+    return wrapper
+
+######################
+# DistributedSampler #
+######################
+
+@register()
+class GlobalDistributedSampler(torch.utils.data.Sampler):
+    """
+    This is a distributed sampler that sync accross gpus and nodes.
+    """
+    def __init__(self, 
+                 dataset, 
+                 shuffle=True,
+                 extend=False,):
+        """
+        Arguments:
+            dataset: Dataset used for sampling.
+            shuffle: If true, sampler will shuffle the indices
+            extend: If true, sampler will extend the indices that can be even distributed by ranks 
+                otherwise sampler will truncate the indices to make it even.
+        """
+        self.ddp = sync.is_ddp()
+        self.rank = sync.get_rank('global')
+        self.world_size = sync.get_world_size('global')
+        self.dataset = dataset
+        self.shuffle = shuffle
+        self.extend = extend
+
+        num_samples = len(dataset) // self.world_size
+        if extend and (len(dataset)%self.world_size != 0):
+            num_samples+=1
+        self.num_samples = num_samples
+        self.total_size = num_samples * self.world_size
+
+    def __iter__(self):
+        indices = self.get_sync_order()
+        if self.extend:
+            # extend using the front indices
+            indices = indices+indices[0:self.total_size-len(indices)]
+        else:
+            # truncate
+            indices = indices[0:self.total_size]
+        # subsample
+        indices = indices[self.rank : len(indices) : self.world_size]
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+    def get_sync_order(self):
+        if self.shuffle:
+            indices = torch.randperm(len(self.dataset)).to(self.rank)
+            if self.ddp:
+                dist.broadcast(indices, src=0)
+            indices = indices.to('cpu').tolist()
+        else:
+            indices = list(range(len(self.dataset)))
+        print_log('Sampler : {}'.format(str(indices[0:5])) )
+        return indices
+
+@register()
+class LocalDistributedSampler(GlobalDistributedSampler):
+    """
+    This is a distributed sampler that sync across gpus within the nodes.
+        But not sync across nodes.
+    """
+    def __init__(self, 
+                 dataset, 
+                 shuffle=True,
+                 extend=False,):
+        super().__init__(dataset, shuffle, extend)
+        self.rank = sync.get_rank('local')
+        self.world_size = sync.get_world_size('local')
+
+    def get_sync_order(self):
+        if self.shuffle:
+            if self.rank == 0:
+                indices = list(npr.permutation(len(self.dataset)))
+                sync.nodewise_sync().broadcast_r0(indices)
+            else:
+                indices = sync.nodewise_sync().broadcast_r0(None)
+        else:
+            indices = list(range(len(self.dataset)))
+        print_log('Sampler : {}'.format(str(indices[0:5])) )
+        return indices
+
+############################
+# random sample with group #
+############################
+# Deprecated
+
+@register()
+class GroupSampler(torch.utils.data.Sampler):
+    """
+    This is a new DistributedSampler that sample all index according to group.
+    i.e. 
+    if group_size=3, num_replicas=2, train mode:
+        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10
+            ==> (group) [0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10]
+            ==> (distribute) process0: [3, 4, 5], (leftover [6, 7, 8, 9, 10])
+                             process1: [0, 1, 2]
+            ==> (group leftover) process0: [3, 4, 5], (leftover [6, 7], [8, 9], 10)
+                                 process1: [0, 1, 2]
+            ==> (distribute) process0: [3, 4, 5], [6, 7] (remove 10)
+                             process1: [0, 1, 2], [8, 9]
+
+        it will avoid_batchsize=1:
+        0, 1, 2, 3, 4, 5, 6, 7, 8,
+            ==> (group) [0, 1, 2], [3, 4, 5], [6, 7, 8]
+            ==> (distribute) process0: [3, 4, 5], (leftover [6, 7, 8])
+                             process1: [0, 1, 2]
+            ==> (group leftover) process0: [3, 4, 5], (leftover [6], [7], [8])
+                                 process1: [0, 1, 2]
+            ==> (distribute) process0: [3, 4, 5], (remove 6, 7, 8) (because distribute make batchsize 1)
+                             process1: [0, 1, 2]
+
+    if group_size=3, num_replicas=2, eval mode:
+        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10
+            ==> (extend) 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 10
+            ==> (group) [0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 10]
+            ==> (distribute) process0: [0, 1, 2], [6, 7, 8],
+                             process1: [3, 4, 5], [9, 10, 10]
+    """
+
+    def __init__(self, 
+                 dataset, 
+                 group_size,
+                 num_replicas=None, 
+                 rank=None, 
+                 mode='train',):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise ValueError
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise ValueError
+            rank = dist.get_rank()
+
+        self.dataset = dataset
+        self.len_dataset = len(dataset)
+        self.group_size = group_size
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.mode = mode
+        len_dataset = self.len_dataset
+
+        if (len_dataset % num_replicas != 0) and (mode == 'train'):
+            # drop the non_aligned
+            aligned_indices = np.arange(len_dataset)[:-(len_dataset % num_replicas)]
+            aligned_len_dataset = aligned_indices.shape[0]
+        elif (len_dataset % num_replicas != 0) and (mode == 'eval'):
+            extend = np.array([len_dataset-1 for _ in range(num_replicas - len_dataset % num_replicas)])
+            aligned_indices = np.concatenate([range(len_dataset), extend])
+            aligned_len_dataset = aligned_indices.shape[0]
+        else:
+            aligned_indices = np.arange(len_dataset)
+            aligned_len_dataset = len_dataset
+
+        num_even_distributed_groups = aligned_len_dataset // (group_size * num_replicas)
+        num_even = num_even_distributed_groups * group_size * num_replicas
+
+        self.regular_groups = aligned_indices[0:num_even].reshape(-1, group_size)
+        self.leftover_groups = aligned_indices[num_even:].reshape(num_replicas, -1)
+
+        if self.leftover_groups.size == 0:
+            self.leftover_groups = None
+        elif (self.leftover_groups.shape[-1]==1) and (mode == 'train'):
+            # avoid bs=1
+            self.leftover_groups = None
+
+        # a urly way to modify dataset.load_info according to the grouping
+        for groupi in self.regular_groups:
+            for idx in groupi:
+                idx_lowerbd = groupi[0]
+                idx_upperbd = groupi[-1]
+                idx_reference = (idx_lowerbd+idx_upperbd)//2
+                dataset.load_info[idx]['ref_size'] = dataset.load_info[idx_reference]['image_size']
+        if self.leftover_groups is not None:
+            for groupi in self.leftover_groups:
+                for idx in groupi:
+                    idx_lowerbd = groupi[0]
+                    idx_upperbd = groupi[-1]
+                    idx_reference = (idx_lowerbd+idx_upperbd)//2
+                    dataset.load_info[idx]['ref_size'] = dataset.load_info[idx_reference]['image_size']
+
+    def concat(self, nparrays, axis=0):
+        # a helper for save concaternation
+        nparrays = [i for i in nparrays if i.size > 0]
+        return np.concatenate(nparrays, axis=axis)
+
+    def __iter__(self):
+        indices = self.get_sync_order()
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+    def get_sync_order(self):
+        # g = torch.Generator()
+        # g.manual_seed(self.epoch)
+
+        mode         = self.mode
+        rank         = self.rank
+        num_replicas = self.num_replicas
+        group_size   = self.group_size
+        num_groups = len(self.regular_groups)
+
+        if mode == 'train':
+            g_indices = torch.randperm(num_groups).to(rank)
+            dist.broadcast(g_indices, src=0)
+            g_indices = g_indices.to('cpu').tolist()
+            num_groups_per_rank = num_groups // num_replicas
+            groups = self.regular_groups[g_indices][num_groups_per_rank*rank : num_groups_per_rank*(rank+1)]
+            indices = groups.flatten()
+
+            if self.leftover_groups is not None:
+                leftg_indices = torch.randperm(len(self.leftover_groups)).to(rank)
+                dist.broadcast(leftg_indices, src=0)
+                leftg_indices = leftg_indices.to('cpu').tolist()
+                last = self.leftover_groups[leftg_indices][rank]
+                indices = np.concatenate([indices, last], axis=0)
+        elif mode == 'eval':
+            groups = self.regular_groups.reshape(-1, num_replicas, group_size)[:, rank, :]
+            indices = groups.flatten()
+            if self.leftover_groups is not None:
+                last = self.leftover_groups[rank]
+                indices = np.concatenate([indices, last], axis=0)
+        else:
+            raise ValueError
+        
+        print_log('Sampler RANK {} : {}'.format(rank, str(indices[0:group_size+1])))
+        return indices

versatile_diffusion/lib/data_factory/common/ds_transform.py

@@ -0,0 +1,178 @@
+import os.path as osp
+import numpy as np
+import numpy.random as npr
+import PIL
+import cv2
+
+import torch
+import torchvision
+import xml.etree.ElementTree as ET
+import json
+import copy
+import math
+
+def singleton(class_):
+    instances = {}
+    def getinstance(*args, **kwargs):
+        if class_ not in instances:
+            instances[class_] = class_(*args, **kwargs)
+        return instances[class_]
+    return getinstance
+
+@singleton
+class get_transform(object):
+    def __init__(self):
+        self.transform = {}
+
+    def register(self, transf):
+        self.transform[transf.__name__] = transf
+
+    def __call__(self, cfg):
+        if cfg is None:
+            return None
+        if isinstance(cfg, list):
+            loader = []
+            for ci in cfg:
+                t = ci.type
+                loader.append(self.transform[t](**ci.args))
+            return compose(loader)
+        t = cfg.type
+        return self.transform[t](**cfg.args)
+
+def register():
+    def wrapper(class_):
+        get_transform().register(class_)
+        return class_
+    return wrapper
+
+def have(must=[], may=[]):
+    """
+    The nextgen decorator that have two list of
+        input tells what category the transform
+        will operate on. 
+    Args:
+        must: [] of str,
+            the names of the items that must be included
+            inside the element. 
+            If element[name] exist: do the transform 
+            If element[name] is None: raise Exception.
+            If element[name] not exist: raise Exception.
+        may: [] of str,
+            the names of the items that may be contained 
+            inside the element for transform. 
+            If element[name] exist: do the transform 
+            If element[name] is None: ignore it.
+            If element[name] not exist: ignore it.
+    """
+    def route(self, item, e, d):
+        """
+        Route the element to a proper function
+            for calculation.
+        Args:
+            self: object,
+                the transform functor.
+            item: str,
+                the item name of the data.
+            e: {},
+                the element
+            d: nparray, tensor or PIL.Image,
+                the data to transform.
+        """
+        if isinstance(d, np.ndarray):
+            dtype = 'nparray'
+        elif isinstance(d, torch.Tensor):
+            dtype = 'tensor'
+        elif isinstance(d, PIL.Image.Image):
+            dtype = 'pilimage'
+        else:
+            raise ValueError
+
+        # find function by order
+        f = None
+        for attrname in [
+                'exec_{}_{}'.format(item, dtype),
+                'exec_{}'.format(item),
+                'exec_{}'.format(dtype),
+                'exec']:
+            f = getattr(self, attrname, None)
+            if f is not None:
+                break
+        d, e = f(d, e)
+        e[item] = d
+        return e
+
+    def wrapper(func):
+        def inner(self, e): 
+            e['imsize_previous'] = e['imsize_current']
+            imsize_tag_cnt = 0
+            imsize_tag = 'imsize_before_' + self.__class__.__name__
+            while True:
+                if imsize_tag_cnt != 0:
+                    tag = imsize_tag + str(imsize_tag_cnt)
+                else:
+                    tag = imsize_tag
+                if not tag in e:
+                    e[tag] = e['imsize_current']
+                    break
+                imsize_tag_cnt += 1
+            
+            e = func(self, e)
+            # must transform list
+            for item in must:
+                try:
+                    d = e[item]
+                except:
+                    raise ValueError
+                if d is None:
+                    raise ValueError
+                e = route(self, item, e, d)
+            # may transform list
+            for item in may:
+                try:
+                    d = e[item]
+                except:
+                    d = None
+                if d is not None:
+                    e = route(self, item, e, d)
+            return e
+        return inner
+    return wrapper
+
+class compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, element):
+        for t in self.transforms:
+            element = t(element)
+        return element
+
+class TBase(object):
+    def __init__(self):
+        pass
+
+    def exec(self, data, element):
+        raise ValueError
+
+    def rand(self, 
+             uid,
+             tag, 
+             rand_f, 
+             *args,
+             **kwargs):
+        """
+        Args:
+            uid: string element['unique_id']
+            tag: string tells the tag uses when tracking the random number.
+                Or the tag to restore the tracked random number.
+            rand_f: the random function use to generate random number. 
+            **kwargs: the argument for the given random function.
+        """
+        # if rnduh().hdata is not None:
+        #     return rnduh().get_history(uid, self.__class__.__name__, tag)
+        # if rnduh().record_path is None:
+        #     return rand_f(*args, **kwargs)
+        # the special mode to create the random file.
+        d = rand_f(*args, **kwargs)
+        # rnduh().record(uid, self.__class__.__name__, tag, d)
+        return d

versatile_diffusion/lib/data_factory/ds_laion2b_webdataset.py

@@ -0,0 +1,221 @@
+import os
+import os.path as osp
+import numpy as np
+import numpy.random as npr
+import torch
+import torch.distributed as dist
+import torchvision.transforms as tvtrans
+import PIL.Image
+PIL.Image.MAX_IMAGE_PIXELS = None
+import math
+import json
+import copy
+import pickle
+from multiprocessing import shared_memory
+import time
+from .common import *
+from ..log_service import print_log
+
+from lib import visual_service as vis
+from .. import sync
+
+import webdataset as wds
+
+###################################################
+# this is a special ds that use webdataset mainly #
+###################################################
+
+@regdataset()
+class laion2b_dummy(ds_base):
+    def init_load_info(self):
+        self.load_info = []
+
+@regdataset()
+class laion2b_webdataset(ds_base):
+    def init_load_info(self):
+        self.load_info = []
+
+    def make_loader(self, batch_size, num_workers, train=True):
+        cfg = self.cfg
+        self.root_dir = cfg.root_dir
+
+        interpolation_mode = tvtrans.InterpolationMode.BICUBIC
+        if train:
+            trans = [
+                tvtrans.Resize(cfg.scale, interpolation=interpolation_mode),
+                tvtrans.RandomCrop(cfg.scale),
+                tvtrans.ToTensor(),]
+        else:
+            trans = [
+                tvtrans.Resize(cfg.scale, interpolation=interpolation_mode),
+                tvtrans.CenterCrop(cfg.scale),
+                tvtrans.ToTensor(),]
+
+        trans = tvtrans.Compose(trans)
+
+        trans_dict = {'jpg': trans}
+        postprocess = customized_postprocess
+
+        shuffle = cfg.get('shuffle', 10000)
+        shardshuffle = shuffle > 0
+        node_world_size = sync.get_world_size('node')
+        nodesplitter = wds.shardlists.split_by_node \
+            if node_world_size==1 else wds.shardlists.single_node_only
+
+        tars = [osp.join(self.root_dir, 'data', i) for i in os.listdir(osp.join(self.root_dir, 'data')) 
+            if osp.splitext(i)[1]=='.tar']
+        tars = sorted(tars)
+
+        dset = wds.WebDataset(
+            tars,
+            nodesplitter=nodesplitter,
+            shardshuffle=shardshuffle,
+            handler=wds.warn_and_continue).repeat().shuffle(shuffle)
+
+        print_log(f'Loading webdataset with {len(dset.pipeline[0].urls)} shards.')
+        self.min_size = cfg.get('min_size', None)
+        self.max_pwatermark = cfg.get('max_pwatermark', None)
+        dset = (dset
+            .select(self.filter_keys)
+            .decode('pil', handler=wds.warn_and_continue)
+            .select(self.filter_size)
+            .map_dict(**trans_dict, handler=wds.warn_and_continue))
+
+        if postprocess is not None:
+            dset = dset.map(postprocess)
+ 
+        dset.batched(batch_size, partial=False)
+ 
+        loader = wds.WebLoader(
+            dset, 
+            batch_size=None,
+            shuffle=False,
+            num_workers=num_workers, )
+        return loader
+
+    def filter_size(self, x):
+        try:
+            valid = True
+            if self.min_size is not None and self.min_size > 1:
+                try:
+                    valid = valid and x['json']['original_width'] >= self.min_size and \
+                        x['json']['original_height'] >= self.min_size
+                except Exception:
+                    valid = False
+            if self.max_pwatermark is not None and self.max_pwatermark < 1.0:
+                try:
+                    valid = valid and  x['json']['pwatermark'] <= self.max_pwatermark
+                except Exception:
+                    valid = False
+            return valid
+        except Exception:
+            return False
+
+    def filter_keys(self, x):
+        try:
+            return ("jpg" in x) and ("txt" in x)
+        except Exception:
+            return False
+
+    def train_dataloader(self):
+        return self.make_loader(self.train)
+
+    def val_dataloader(self):
+        return self.make_loader(self.validation, train=False)
+
+    def test_dataloader(self):
+        return self.make_loader(self.test, train=False)
+
+def customized_postprocess(element):
+    return element['jpg']*2-1, element['txt'], element['__key__']
+
+def dict_collation_fn(samples, combine_tensors=True, combine_scalars=True):
+    keys = set.intersection(*[set(sample.keys()) for sample in samples])
+    batched = {key: [] for key in keys}
+
+    for s in samples:
+        [batched[key].append(s[key]) for key in batched]
+
+    result = {}
+    for key in batched:
+        if isinstance(batched[key][0], (int, float)):
+            if combine_scalars:
+                result[key] = np.array(list(batched[key]))
+        elif isinstance(batched[key][0], torch.Tensor):
+            if combine_tensors:
+                result[key] = torch.stack(list(batched[key]))
+        elif isinstance(batched[key][0], np.ndarray):
+            if combine_tensors:
+                result[key] = np.array(list(batched[key]))
+        else:
+            result[key] = list(batched[key])
+    return result
+
+###################
+# for sd official #
+###################
+
+def customized_postprocess_sdofficial(element):
+    return {
+        'jpg': element['jpg']*2-1, 
+        'txt': element['txt'], }
+
+@regdataset()
+class laion2b_webdataset_sdofficial(laion2b_webdataset):
+    def make_loader(self, batch_size, num_workers, train=True):
+        cfg = self.cfg
+        self.root_dir = cfg.root_dir
+
+        interpolation_mode = tvtrans.InterpolationMode.BICUBIC
+        if train:
+            trans = [
+                tvtrans.Resize(cfg.scale, interpolation=interpolation_mode),
+                tvtrans.RandomCrop(cfg.scale),
+                tvtrans.ToTensor(),]
+        else:
+            trans = [
+                tvtrans.Resize(cfg.scale, interpolation=interpolation_mode),
+                tvtrans.CenterCrop(cfg.scale),
+                tvtrans.ToTensor(),]
+
+        trans = tvtrans.Compose(trans)
+
+        trans_dict = {'jpg': trans}
+        postprocess = customized_postprocess_sdofficial
+
+        shuffle = 10000
+        shardshuffle = shuffle > 0
+        node_world_size = 1
+        nodesplitter = wds.shardlists.split_by_node \
+            if node_world_size==1 else wds.shardlists.single_node_only
+
+        tars = [osp.join(self.root_dir, 'data', i) for i in os.listdir(osp.join(self.root_dir, 'data')) 
+            if osp.splitext(i)[1]=='.tar']
+        tars = sorted(tars)
+
+        dset = wds.WebDataset(
+            tars,
+            nodesplitter=nodesplitter,
+            shardshuffle=shardshuffle,
+            handler=wds.warn_and_continue).repeat().shuffle(shuffle)
+
+        print(f'Loading webdataset with {len(dset.pipeline[0].urls)} shards.')
+        self.min_size = cfg.get('min_size', None)
+        self.max_pwatermark = cfg.get('max_pwatermark', None)
+        dset = (dset
+            .select(self.filter_keys)
+            .decode('pil', handler=wds.warn_and_continue)
+            .select(self.filter_size)
+            .map_dict(**trans_dict, handler=wds.warn_and_continue))
+
+        if postprocess is not None:
+            dset = dset.map(postprocess)
+ 
+        dset.batched(batch_size, partial=False, collation_fn=dict_collation_fn)
+ 
+        loader = wds.WebLoader(
+            dset, 
+            batch_size=None,
+            shuffle=False,
+            num_workers=num_workers, )
+        return loader

versatile_diffusion/lib/evaluator/__init__.py

@@ -0,0 +1 @@
+from .eva_base import get_evaluator

versatile_diffusion/lib/evaluator/eva_base.py

@@ -0,0 +1,293 @@
+import torch
+import torch.distributed as dist
+
+import os
+import os.path as osp
+import numpy as np
+import cv2
+import copy
+import json
+
+from ..log_service import print_log
+
+def singleton(class_):
+    instances = {}
+    def getinstance(*args, **kwargs):
+        if class_ not in instances:
+            instances[class_] = class_(*args, **kwargs)
+        return instances[class_]
+    return getinstance
+
+@singleton
+class get_evaluator(object):
+    def __init__(self):
+        self.evaluator = {}
+
+    def register(self, evaf, name):
+        self.evaluator[name] = evaf
+
+    def __call__(self, pipeline_cfg=None):
+        if pipeline_cfg is None:
+            from . import eva_null
+            return self.evaluator['null']()
+
+        if not isinstance(pipeline_cfg, list):
+            t = pipeline_cfg.type
+            if t == 'miou':
+                from . import eva_miou
+            if t == 'psnr':
+                from . import eva_psnr
+            if t == 'ssim':
+                from . import eva_ssim
+            if t == 'lpips':
+                from . import eva_lpips
+            if t == 'fid':
+                from . import eva_fid
+            return self.evaluator[t](**pipeline_cfg.args)
+
+        evaluator = []
+        for ci in pipeline_cfg:
+            t = ci.type
+            if t == 'miou':
+                from . import eva_miou
+            if t == 'psnr':
+                from . import eva_psnr
+            if t == 'ssim':
+                from . import eva_ssim
+            if t == 'lpips':
+                from . import eva_lpips
+            if t == 'fid':
+                from . import eva_fid
+            evaluator.append(
+                self.evaluator[t](**ci.args))
+        if len(evaluator) == 0:
+            return None
+        else:
+            return compose(evaluator)
+
+def register(name):
+    def wrapper(class_):
+        get_evaluator().register(class_, name)
+        return class_
+    return wrapper
+
+class base_evaluator(object):
+    def __init__(self, 
+                 **args):
+        '''
+        Args:
+            sample_n, int,
+                the total number of sample. used in 
+                distributed sync
+        '''
+        if not dist.is_available():
+            raise ValueError
+        self.world_size = dist.get_world_size()
+        self.rank = dist.get_rank()
+        self.sample_n = None
+        self.final = {}
+
+    def sync(self, data):
+        """
+        Args:
+            data: any,
+                the data needs to be broadcasted
+        """
+        if data is None:
+            return None
+
+        if isinstance(data, tuple):
+            data = list(data)
+
+        if isinstance(data, list):
+            data_list = []
+            for datai in data:
+                data_list.append(self.sync(datai))
+            data = [[*i] for i in zip(*data_list)]
+            return data
+
+        data = [
+            self.sync_(data, ranki)
+                for ranki in range(self.world_size)
+        ]
+        return data
+
+    def sync_(self, data, rank):
+
+        t = type(data)
+        is_broadcast = rank == self.rank
+
+        if t is np.ndarray:
+            dtrans = data
+            dt = data.dtype
+            if dt in [
+                    int,
+                    np.bool,
+                    np.uint8, 
+                    np.int8,
+                    np.int16,
+                    np.int32,
+                    np.int64,]:
+                dtt = torch.int64
+            elif dt in [
+                    float,
+                    np.float16,
+                    np.float32,
+                    np.float64,]:
+                dtt = torch.float64
+
+        elif t is str:
+            dtrans = np.array(
+                [ord(c) for c in data],
+                dtype = np.int64
+            )
+            dt = np.int64
+            dtt = torch.int64
+        else:
+            raise ValueError
+
+        if is_broadcast:
+            n = len(dtrans.shape)
+            n = torch.tensor(n).long()
+
+            n = n.to(self.rank)
+            dist.broadcast(n, src=rank)
+
+            n = list(dtrans.shape)
+            n = torch.tensor(n).long()
+            n = n.to(self.rank)
+            dist.broadcast(n, src=rank)
+
+            n = torch.tensor(dtrans, dtype=dtt)
+            n = n.to(self.rank)
+            dist.broadcast(n, src=rank)
+            return data
+
+        n = torch.tensor(0).long()
+        n = n.to(self.rank)
+        dist.broadcast(n, src=rank)
+        n = n.item()
+
+        n = torch.zeros(n).long()
+        n = n.to(self.rank)
+        dist.broadcast(n, src=rank)
+        n = list(n.to('cpu').numpy())
+
+        n = torch.zeros(n, dtype=dtt)
+        n = n.to(self.rank)
+        dist.broadcast(n, src=rank)
+        n = n.to('cpu').numpy().astype(dt)
+
+        if t is np.ndarray:
+            return n
+        elif t is str:
+            n = ''.join([chr(c) for c in n])
+            return n
+
+    def zipzap_arrange(self, data):
+        '''
+        Order the data so it range like this:
+            input [[0, 2, 4, 6], [1, 3, 5, 7]] -> output [0, 1, 2, 3, 4, 5, ...]
+        '''
+        if isinstance(data[0], list):
+            data_new = []
+            maxlen = max([len(i) for i in data])
+            totlen = sum([len(i) for i in data])
+            cnt = 0
+            for idx in range(maxlen):
+                for datai in data:
+                    data_new += [datai[idx]]
+                    cnt += 1
+                    if cnt >= totlen:
+                        break
+            return data_new
+
+        elif isinstance(data[0], np.ndarray):
+            maxlen = max([i.shape[0] for i in data])
+            totlen = sum([i.shape[0] for i in data])
+            datai_shape = data[0].shape[1:]
+            data = [
+                np.concatenate(datai, np.zeros(maxlen-datai.shape[0], *datai_shape), axis=0)
+                if datai.shape[0] < maxlen else datai
+                for datai in data
+            ] # even the array
+            data = np.stack(data, axis=1).reshape(-1, *datai_shape)
+            data = data[:totlen]
+            return data
+
+        else:
+            raise NotImplementedError
+
+    def add_batch(self, **args):
+        raise NotImplementedError
+
+    def set_sample_n(self, sample_n):
+        self.sample_n = sample_n
+
+    def compute(self):
+        raise NotImplementedError
+
+    # Function needed in training to judge which 
+    #   evaluated number is better
+    def isbetter(self, old, new):
+        return new>old
+
+    def one_line_summary(self):
+        print_log('Evaluator display')
+
+    def save(self, path):
+        if not osp.exists(path):
+            os.makedirs(path)
+        ofile = osp.join(path, 'result.json')
+        with open(ofile, 'w') as f:
+            json.dump(self.final, f, indent=4)
+
+    def clear_data(self):
+        raise NotImplementedError
+
+class compose(object):
+    def __init__(self, pipeline):
+        self.pipeline = pipeline
+        self.sample_n = None
+        self.final = {}
+
+    def add_batch(self, *args, **kwargs):
+        for pi in self.pipeline:
+            pi.add_batch(*args, **kwargs)
+
+    def set_sample_n(self, sample_n):
+        self.sample_n = sample_n
+        for pi in self.pipeline:
+            pi.set_sample_n(sample_n)
+
+    def compute(self):
+        rv = {}
+        for pi in self.pipeline:
+            rv[pi.symbol] = pi.compute()
+            self.final[pi.symbol] = pi.final
+        return rv
+
+    def isbetter(self, old, new):
+        check = 0
+        for pi in self.pipeline:
+            if pi.isbetter(old, new):
+                check+=1
+        if check/len(self.pipeline)>0.5:
+            return True
+        else:
+            return False
+
+    def one_line_summary(self):
+        for pi in self.pipeline:
+            pi.one_line_summary()
+
+    def save(self, path):
+        if not osp.exists(path):
+            os.makedirs(path)
+        ofile = osp.join(path, 'result.json')
+        with open(ofile, 'w') as f:
+            json.dump(self.final, f, indent=4)
+
+    def clear_data(self):
+        for pi in self.pipeline:
+            pi.clear_data()

versatile_diffusion/lib/evaluator/eva_null.py

@@ -0,0 +1,26 @@
+import torch
+import numpy as np
+import lpips
+
+from .. import nputils
+from ..log_service import print_log
+
+from .eva_base import base_evaluator, register
+
+@register('null')
+class null_evaluator(base_evaluator):
+    def __init__(self, **dummy):
+        super().__init__()
+
+    def add_batch(self, 
+                  **dummy):
+        pass
+
+    def compute(self):
+        return None
+
+    def one_line_summary(self):
+        print_log('Evaluator null')
+
+    def clear_data(self):
+        pass

versatile_diffusion/lib/experiments/sd_default.py

@@ -0,0 +1,441 @@
+import torch
+import torch.distributed as dist
+from torchvision import transforms as tvtrans
+import os
+import os.path as osp
+import time
+import timeit
+import copy
+import json
+import pickle
+import PIL.Image
+import numpy as np
+from datetime import datetime
+from easydict import EasyDict as edict
+from collections import OrderedDict
+
+from lib.cfg_holder import cfg_unique_holder as cfguh
+from lib.data_factory import get_dataset, get_sampler, collate
+from lib.model_zoo import \
+    get_model, get_optimizer, get_scheduler
+from lib.log_service import print_log
+
+from ..utils import train as train_base
+from ..utils import eval as eval_base
+from ..utils import train_stage as tsbase
+from ..utils import eval_stage as esbase
+from .. import sync
+
+###############
+# some helper #
+###############
+
+def atomic_save(cfg, net, opt, step, path):
+    if isinstance(net, (torch.nn.DataParallel,
+                        torch.nn.parallel.DistributedDataParallel)):
+        netm = net.module
+    else:
+        netm = net
+    sd = netm.state_dict()
+    slimmed_sd = [(ki, vi) for ki, vi in sd.items()
+        if ki.find('first_stage_model')!=0 and ki.find('cond_stage_model')!=0]
+
+    checkpoint = {
+        "config" : cfg,
+        "state_dict" : OrderedDict(slimmed_sd),
+        "step" : step}
+    if opt is not None:
+        checkpoint['optimizer_states'] = opt.state_dict()
+    import io
+    import fsspec
+    bytesbuffer = io.BytesIO()
+    torch.save(checkpoint, bytesbuffer)
+    with fsspec.open(path, "wb") as f:
+        f.write(bytesbuffer.getvalue())
+
+def load_state_dict(net, cfg):
+    pretrained_pth_full  = cfg.get('pretrained_pth_full' , None)
+    pretrained_ckpt_full = cfg.get('pretrained_ckpt_full', None)
+    pretrained_pth       = cfg.get('pretrained_pth'      , None)
+    pretrained_ckpt      = cfg.get('pretrained_ckpt'     , None)
+    pretrained_pth_dm    = cfg.get('pretrained_pth_dm'   , None)
+    pretrained_pth_ema   = cfg.get('pretrained_pth_ema'  , None)
+    strict_sd = cfg.get('strict_sd', False)
+    errmsg = "Overlapped model state_dict! This is undesired behavior!"
+
+    if pretrained_pth_full is not None or pretrained_ckpt_full is not None:
+        assert (pretrained_pth is None) and \
+               (pretrained_ckpt is None) and \
+               (pretrained_pth_dm is None) and \
+               (pretrained_pth_ema is None), errmsg            
+        if pretrained_pth_full is not None:
+            target_file = pretrained_pth_full
+            sd = torch.load(target_file, map_location='cpu')
+            assert pretrained_ckpt is None, errmsg
+        else:
+            target_file = pretrained_ckpt_full
+            sd = torch.load(target_file, map_location='cpu')['state_dict']
+        print_log('Load full model from [{}] strict [{}].'.format(
+            target_file, strict_sd))
+        net.load_state_dict(sd, strict=strict_sd)
+
+    if pretrained_pth is not None or pretrained_ckpt is not None:
+        assert (pretrained_ckpt_full is None) and \
+               (pretrained_pth_full is None) and \
+               (pretrained_pth_dm is None) and \
+               (pretrained_pth_ema is None), errmsg
+        if pretrained_pth is not None:
+            target_file = pretrained_pth
+            sd = torch.load(target_file, map_location='cpu')
+            assert pretrained_ckpt is None, errmsg
+        else:
+            target_file = pretrained_ckpt
+            sd = torch.load(target_file, map_location='cpu')['state_dict']
+        print_log('Load model from [{}] strict [{}].'.format(
+            target_file, strict_sd))
+        sd_extra = [(ki, vi) for ki, vi in net.state_dict().items() \
+            if ki.find('first_stage_model')==0 or ki.find('cond_stage_model')==0]
+        sd.update(OrderedDict(sd_extra))
+        net.load_state_dict(sd, strict=strict_sd)
+
+    if pretrained_pth_dm is not None:
+        assert (pretrained_ckpt_full is None) and \
+               (pretrained_pth_full is None) and \
+               (pretrained_pth is None) and \
+               (pretrained_ckpt is None), errmsg
+        print_log('Load diffusion model from [{}] strict [{}].'.format(
+            pretrained_pth_dm, strict_sd))
+        sd = torch.load(pretrained_pth_dm, map_location='cpu')
+        net.model.diffusion_model.load_state_dict(sd, strict=strict_sd)
+
+    if pretrained_pth_ema is not None:
+        assert (pretrained_ckpt_full is None) and \
+               (pretrained_pth_full is None) and \
+               (pretrained_pth is None) and \
+               (pretrained_ckpt is None), errmsg
+        print_log('Load unet ema model from [{}] strict [{}].'.format(
+            pretrained_pth_ema, strict_sd))
+        sd = torch.load(pretrained_pth_ema, map_location='cpu')
+        net.model_ema.load_state_dict(sd, strict=strict_sd)
+
+def auto_merge_imlist(imlist, max=64):
+    imlist = imlist[0:max]
+    h, w = imlist[0].shape[0:2]
+    num_images = len(imlist)
+    num_row = int(np.sqrt(num_images))
+    num_col = num_images//num_row + 1 if num_images%num_row!=0 else num_images//num_row
+    canvas = np.zeros([num_row*h, num_col*w, 3], dtype=np.uint8)
+    for idx, im in enumerate(imlist):
+        hi = (idx // num_col) * h
+        wi = (idx % num_col) * w
+        canvas[hi:hi+h, wi:wi+w, :] = im
+    return canvas
+
+def latent2im(net, latent):
+    single_input = len(latent.shape) == 3
+    if single_input:
+        latent = latent[None]
+    im = net.decode_image(latent.to(net.device))
+    im = torch.clamp((im+1.0)/2.0, min=0.0, max=1.0)
+    im = [tvtrans.ToPILImage()(i) for i in im]
+    if single_input:
+        im = im[0]
+    return im
+
+def im2latent(net, im):
+    single_input = not isinstance(im, list)
+    if single_input:
+        im = [im]
+    im = torch.stack([tvtrans.ToTensor()(i) for i in im], dim=0)
+    im = (im*2-1).to(net.device)
+    z = net.encode_image(im)
+    if single_input:
+        z = z[0]
+    return z
+
+class color_adjust(object):
+    def __init__(self, ref_from, ref_to):
+        x0, m0, std0 = self.get_data_and_stat(ref_from)
+        x1, m1, std1 = self.get_data_and_stat(ref_to)
+        self.ref_from_stat = (m0, std0)
+        self.ref_to_stat   = (m1, std1)
+        self.ref_from = self.preprocess(x0).reshape(-1, 3)
+        self.ref_to = x1.reshape(-1, 3)
+
+    def get_data_and_stat(self, x):
+        if isinstance(x, str):
+            x = np.array(PIL.Image.open(x))
+        elif isinstance(x, PIL.Image.Image):
+            x = np.array(x)
+        elif isinstance(x, torch.Tensor):
+            x = torch.clamp(x, min=0.0, max=1.0)
+            x = np.array(tvtrans.ToPILImage()(x))
+        elif isinstance(x, np.ndarray):
+            pass
+        else:
+            raise ValueError
+        x = x.astype(float)
+        m = np.reshape(x, (-1, 3)).mean(0)
+        s = np.reshape(x, (-1, 3)).std(0)
+        return x, m, s
+
+    def preprocess(self, x):
+        m0, s0 = self.ref_from_stat
+        m1, s1 = self.ref_to_stat
+        y = ((x-m0)/s0)*s1 + m1
+        return y
+
+    def __call__(self, xin, keep=0, simple=False):
+        xin, _, _ = self.get_data_and_stat(xin)
+        x = self.preprocess(xin)
+        if simple: 
+            y = (x*(1-keep) + xin*keep)
+            y = np.clip(y, 0, 255).astype(np.uint8)
+            return y
+
+        h, w = x.shape[:2]
+        x = x.reshape(-1, 3)
+        y = []
+        for chi in range(3):
+            yi = self.pdf_transfer_1d(self.ref_from[:, chi], self.ref_to[:, chi], x[:, chi])
+            y.append(yi)
+
+        y = np.stack(y, axis=1)
+        y = y.reshape(h, w, 3)
+        y = (y.astype(float)*(1-keep) + xin.astype(float)*keep)
+        y = np.clip(y, 0, 255).astype(np.uint8)
+        return y
+
+    def pdf_transfer_1d(self, arr_fo, arr_to, arr_in, n=600):
+        arr = np.concatenate((arr_fo, arr_to))
+        min_v = arr.min() - 1e-6
+        max_v = arr.max() + 1e-6
+        min_vto = arr_to.min() - 1e-6
+        max_vto = arr_to.max() + 1e-6
+        xs = np.array(
+            [min_v + (max_v - min_v) * i / n for i in range(n + 1)])
+        hist_fo, _ = np.histogram(arr_fo, xs)
+        hist_to, _ = np.histogram(arr_to, xs)
+        xs = xs[:-1]
+        # compute probability distribution
+        cum_fo = np.cumsum(hist_fo)
+        cum_to = np.cumsum(hist_to)
+        d_fo = cum_fo / cum_fo[-1]
+        d_to = cum_to / cum_to[-1]
+        # transfer
+        t_d = np.interp(d_fo, d_to, xs)
+        t_d[d_fo <= d_to[ 0]] = min_vto
+        t_d[d_fo >= d_to[-1]] = max_vto
+        arr_out = np.interp(arr_in, xs, t_d)
+        return arr_out
+
+########
+# main #
+########
+
+class eval(eval_base):
+    def prepare_model(self):
+        cfg = cfguh().cfg
+        net = get_model()(cfg.model)
+        if cfg.env.cuda:
+            net.to(self.local_rank)
+            load_state_dict(net, cfg.eval) #<--- added
+            net = torch.nn.parallel.DistributedDataParallel(
+                net, device_ids=[self.local_rank], 
+                find_unused_parameters=True)
+        net.eval()
+        return {'net' : net,}
+
+class eval_stage(esbase):
+    """
+    This is eval stage that can check comprehensive results
+    """
+    def __init__(self):
+        from ..model_zoo.ddim import DDIMSampler
+        self.sampler = DDIMSampler
+
+    def get_net(self, paras):
+        return paras['net']
+
+    def get_image_path(self):
+        if 'train' in cfguh().cfg:
+            log_dir = cfguh().cfg.train.log_dir
+        else:
+            log_dir = cfguh().cfg.eval.log_dir
+        return os.path.join(log_dir, "udemo")
+
+    @torch.no_grad()
+    def sample(self, net, sampler, prompt, output_dim, scale, n_samples, ddim_steps, ddim_eta):
+        h, w = output_dim
+        uc = None
+        if scale != 1.0:
+            uc = net.get_learned_conditioning(n_samples * [""])
+        c = net.get_learned_conditioning(n_samples * [prompt])
+        shape = [4, h//8, w//8]
+        rv = sampler.sample(
+            S=ddim_steps,
+            conditioning=c,
+            batch_size=n_samples,
+            shape=shape,
+            verbose=False,
+            unconditional_guidance_scale=scale,
+            unconditional_conditioning=uc,
+            eta=ddim_eta)
+        return rv
+
+    def save_images(self, pil_list, name, path, suffix=''):
+        canvas = auto_merge_imlist([np.array(i) for i in pil_list])
+        image_name = '{}{}.png'.format(name, suffix)
+        PIL.Image.fromarray(canvas).save(osp.join(path, image_name))
+
+    def __call__(self, **paras):
+        cfg = cfguh().cfg
+        cfgv = cfg.eval
+
+        net = paras['net']
+        eval_cnt = paras.get('eval_cnt', None)
+        fix_seed = cfgv.get('fix_seed', False)
+
+        LRANK = sync.get_rank('local')
+        LWSIZE = sync.get_world_size('local')
+
+        image_path = self.get_image_path()
+        self.create_dir(image_path)
+        eval_cnt = paras.get('eval_cnt', None)
+        suffix='' if eval_cnt is None else '_itern'+str(eval_cnt)
+
+        if isinstance(net, (torch.nn.DataParallel,
+                            torch.nn.parallel.DistributedDataParallel)):
+            netm = net.module
+        else:
+            netm = net
+
+        with_ema = getattr(netm, 'model_ema', None) is not None
+        sampler = self.sampler(netm)
+        setattr(netm, 'device', LRANK) # Trick
+
+        replicate = cfgv.get('replicate', 1)
+        conditioning = cfgv.conditioning * replicate
+        conditioning_local = conditioning[LRANK : len(conditioning) : LWSIZE]
+        seed_increment = [i for i in range(len(conditioning))][LRANK : len(conditioning) : LWSIZE]
+
+        for prompti, seedi in zip(conditioning_local, seed_increment):
+            if prompti == 'SKIP':
+                continue
+            draw_filename = prompti.strip().replace(' ', '-')
+            if fix_seed:
+                np.random.seed(cfg.env.rnd_seed + seedi)
+                torch.manual_seed(cfg.env.rnd_seed + seedi + 100)
+                suffixi = suffix + "_seed{}".format(cfg.env.rnd_seed + seedi + 100)
+            else:
+                suffixi = suffix
+
+            if with_ema:
+                with netm.ema_scope():
+                    x, _ = self.sample(netm, sampler, prompti, **cfgv.sample)
+            else:
+                x, _ = self.sample(netm, sampler, prompti, **cfgv.sample)
+
+            demo_image = latent2im(netm, x)
+            self.save_images(demo_image, draw_filename, image_path, suffix=suffixi)
+
+        if eval_cnt is not None:
+            print_log('Demo printed for {}'.format(eval_cnt))
+        return {}
+
+##################
+# eval variation #
+##################
+
+class eval_stage_variation(eval_stage):
+    @torch.no_grad()
+    def sample(self, net, sampler, visual_hint, output_dim, scale, n_samples, ddim_steps, ddim_eta):
+        h, w = output_dim
+        vh = tvtrans.ToTensor()(PIL.Image.open(visual_hint))[None].to(net.device)
+        c = net.get_learned_conditioning(vh)
+        c = c.repeat(n_samples, 1, 1)
+        uc = None
+        if scale != 1.0:
+            dummy = torch.zeros_like(vh)
+            uc = net.get_learned_conditioning(dummy)
+            uc = uc.repeat(n_samples, 1, 1)
+
+        shape = [4, h//8, w//8]
+        rv = sampler.sample(
+            S=ddim_steps,
+            conditioning=c,
+            batch_size=n_samples,
+            shape=shape,
+            verbose=False,
+            unconditional_guidance_scale=scale,
+            unconditional_conditioning=uc,
+            eta=ddim_eta)
+        return rv
+
+    def __call__(self, **paras):
+        cfg = cfguh().cfg
+        cfgv = cfg.eval
+
+        net = paras['net']
+        eval_cnt = paras.get('eval_cnt', None)
+        fix_seed = cfgv.get('fix_seed', False)
+
+        LRANK = sync.get_rank('local')
+        LWSIZE = sync.get_world_size('local')
+
+        image_path = self.get_image_path()
+        self.create_dir(image_path)
+        eval_cnt = paras.get('eval_cnt', None)
+        suffix='' if eval_cnt is None else '_'+str(eval_cnt)
+
+        if isinstance(net, (torch.nn.DataParallel,
+                            torch.nn.parallel.DistributedDataParallel)):
+            netm = net.module
+        else:
+            netm = net
+
+        with_ema = getattr(netm, 'model_ema', None) is not None
+        sampler = self.sampler(netm)
+        setattr(netm, 'device', LRANK) # Trick
+
+        color_adj = cfguh().cfg.eval.get('color_adj', False)
+        color_adj_keep_ratio = cfguh().cfg.eval.get('color_adj_keep_ratio', 0.5)
+        color_adj_simple = cfguh().cfg.eval.get('color_adj_simple', True)
+
+        replicate = cfgv.get('replicate', 1)
+        conditioning = cfgv.conditioning * replicate
+        conditioning_local = conditioning[LRANK : len(conditioning) : LWSIZE]
+        seed_increment = [i for i in range(len(conditioning))][LRANK : len(conditioning) : LWSIZE]
+
+        for ci, seedi in zip(conditioning_local, seed_increment):
+            if ci == 'SKIP':
+                continue
+
+            draw_filename = osp.splitext(osp.basename(ci))[0]
+
+            if fix_seed:
+                np.random.seed(cfg.env.rnd_seed + seedi)
+                torch.manual_seed(cfg.env.rnd_seed + seedi + 100)
+                suffixi = suffix + "_seed{}".format(cfg.env.rnd_seed + seedi + 100)
+            else:
+                suffixi = suffix
+
+            if with_ema:
+                with netm.ema_scope():
+                    x, _ = self.sample(netm, sampler, ci, **cfgv.sample)
+            else:
+                x, _ = self.sample(netm, sampler, ci, **cfgv.sample)
+
+            demo_image = latent2im(netm, x)
+            if color_adj:
+                x_adj = []
+                for demoi in demo_image:
+                    color_adj_f = color_adjust(ref_from=demoi, ref_to=ci)
+                    xi_adj = color_adj_f(demoi, keep=color_adj_keep_ratio, simple=color_adj_simple)
+                    x_adj.append(xi_adj)
+                demo_image = x_adj
+            self.save_images(demo_image, draw_filename, image_path, suffix=suffixi)
+
+        if eval_cnt is not None:
+            print_log('Demo printed for {}'.format(eval_cnt))
+        return {}

versatile_diffusion/lib/experiments/vd_default.py

@@ -0,0 +1,549 @@
+import torch
+import torch.distributed as dist
+from torchvision import transforms as tvtrans
+import os
+import os.path as osp
+import time
+import timeit
+import copy
+import json
+import pickle
+import PIL.Image
+import numpy as np
+from datetime import datetime
+from easydict import EasyDict as edict
+from collections import OrderedDict
+
+from lib.cfg_holder import cfg_unique_holder as cfguh
+from lib.data_factory import get_dataset, get_sampler, collate
+from lib.model_zoo import \
+    get_model, get_optimizer, get_scheduler
+from lib.log_service import print_log
+
+from ..utils import train as train_base
+from ..utils import eval as eval_base
+from ..utils import train_stage as tsbase
+from ..utils import eval_stage as esbase
+from .. import sync
+
+from .sd_default import auto_merge_imlist, latent2im, color_adjust
+from .sd_default import eval as eval_base
+from .sd_default import eval_stage as eval_stage_base
+
+###############
+# some helper #
+###############
+
+def atomic_save(cfg, net, opt, step, path):
+    if isinstance(net, (torch.nn.DataParallel,
+                        torch.nn.parallel.DistributedDataParallel)):
+        netm = net.module
+    else:
+        netm = net
+    sd = netm.state_dict()
+    slimmed_sd = [(ki, vi) for ki, vi in sd.items()
+        if ki.find('autokl')!=0 and ki.find('optimus')!=0 and ki.find('clip')!=0]
+
+    checkpoint = {
+        "config" : cfg,
+        "state_dict" : OrderedDict(slimmed_sd),
+        "step" : step}
+    if opt is not None:
+        checkpoint['optimizer_states'] = opt.state_dict()
+    import io
+    import fsspec
+    bytesbuffer = io.BytesIO()
+    torch.save(checkpoint, bytesbuffer)
+    with fsspec.open(path, "wb") as f:
+        f.write(bytesbuffer.getvalue())
+
+def load_state_dict(net, cfg):
+    pretrained_pth_full  = cfg.get('pretrained_pth_full' , None)
+    pretrained_ckpt_full = cfg.get('pretrained_ckpt_full', None)
+    pretrained_pth       = cfg.get('pretrained_pth'      , None)
+    pretrained_ckpt      = cfg.get('pretrained_ckpt'     , None)
+    pretrained_pth_dm    = cfg.get('pretrained_pth_dm'   , None)
+    pretrained_pth_ema   = cfg.get('pretrained_pth_ema'  , None)
+    strict_sd = cfg.get('strict_sd', False)
+    errmsg = "Overlapped model state_dict! This is undesired behavior!"
+
+    if pretrained_pth_full is not None or pretrained_ckpt_full is not None:
+        assert (pretrained_pth is None) and \
+               (pretrained_ckpt is None) and \
+               (pretrained_pth_dm is None) and \
+               (pretrained_pth_ema is None), errmsg            
+        if pretrained_pth_full is not None:
+            target_file = pretrained_pth_full
+            sd = torch.load(target_file, map_location='cpu')
+            assert pretrained_ckpt is None, errmsg
+        else:
+            target_file = pretrained_ckpt_full
+            sd = torch.load(target_file, map_location='cpu')['state_dict']
+        print_log('Load full model from [{}] strict [{}].'.format(
+            target_file, strict_sd))
+        net.load_state_dict(sd, strict=strict_sd)
+
+    if pretrained_pth is not None or pretrained_ckpt is not None:
+        assert (pretrained_ckpt_full is None) and \
+               (pretrained_pth_full is None) and \
+               (pretrained_pth_dm is None) and \
+               (pretrained_pth_ema is None), errmsg
+        if pretrained_pth is not None:
+            target_file = pretrained_pth
+            sd = torch.load(target_file, map_location='cpu')
+            assert pretrained_ckpt is None, errmsg
+        else:
+            target_file = pretrained_ckpt
+            sd = torch.load(target_file, map_location='cpu')['state_dict']
+        print_log('Load model from [{}] strict [{}].'.format(
+            target_file, strict_sd))
+        sd_extra = [(ki, vi) for ki, vi in net.state_dict().items() \
+            if ki.find('autokl')==0 or ki.find('optimus')==0 or ki.find('clip')==0]
+        sd.update(OrderedDict(sd_extra))
+        net.load_state_dict(sd, strict=strict_sd)
+
+    if pretrained_pth_dm is not None:
+        assert (pretrained_ckpt_full is None) and \
+               (pretrained_pth_full is None) and \
+               (pretrained_pth is None) and \
+               (pretrained_ckpt is None), errmsg
+        print_log('Load diffusion model from [{}] strict [{}].'.format(
+            pretrained_pth_dm, strict_sd))
+        sd = torch.load(pretrained_pth_dm, map_location='cpu')
+        net.model.diffusion_model.load_state_dict(sd, strict=strict_sd)
+
+    if pretrained_pth_ema is not None:
+        assert (pretrained_ckpt_full is None) and \
+               (pretrained_pth_full is None) and \
+               (pretrained_pth is None) and \
+               (pretrained_ckpt is None), errmsg
+        print_log('Load unet ema model from [{}] strict [{}].'.format(
+            pretrained_pth_ema, strict_sd))
+        sd = torch.load(pretrained_pth_ema, map_location='cpu')
+        net.model_ema.load_state_dict(sd, strict=strict_sd)
+
+###################
+# official stages #
+###################
+
+class eval(eval_base):
+    pass
+
+class eval_stage(eval_stage_base):
+    """
+    Evaluation of both prompt and vision
+    """
+    def __init__(self):
+        from ..model_zoo.ddim_vd import DDIMSampler_VD
+        self.sampler = DDIMSampler_VD
+
+    @torch.no_grad()
+    def sample(
+            self, net, sampler, context, otype, ctype, image_output_dim, text_latent_dim,
+            scale, n_samples, ddim_steps, ddim_eta):
+        if ctype == 'prompt':
+            c = net.clip_encode_text(n_samples * [context])
+            uc = None
+            if scale != 1.0:
+                uc = net.clip_encode_text(n_samples * [""])
+        elif ctype == 'vision':
+            context = context[None].repeat(n_samples, 1, 1, 1)
+            c = net.clip_encode_vision(context)
+            uc = None
+            if scale != 1.0:
+                dummy = torch.zeros_like(context)
+                uc = net.clip_encode_vision(dummy)
+
+        if otype == 'image':
+            h, w = image_output_dim
+            shape = [n_samples, 4, h//8, w//8]
+            rv = sampler.sample(
+                steps=ddim_steps,
+                shape=shape,
+                conditioning=c,
+                unconditional_guidance_scale=scale,
+                unconditional_conditioning=uc,
+                xtype=otype, ctype=ctype,
+                eta=ddim_eta,
+                verbose=False,)
+        elif otype == 'text':
+            n = text_latent_dim
+            shape = [n_samples, n]
+            rv = sampler.sample(
+                steps=ddim_steps,
+                shape=shape,
+                conditioning=c,
+                unconditional_guidance_scale=scale,
+                unconditional_conditioning=uc,
+                xtype=otype, ctype=ctype,
+                eta=ddim_eta,
+                verbose=False,)
+
+        return rv
+
+    def decode_and_save(
+            self, netm, z, xtype, ctype, path, name, suffix,
+            color_adj=False, color_adj_to=None):
+        if xtype == 'image':
+            x = netm.autokl_decode(z)
+            name = 't2i_'+name if ctype == 'prompt' else 'v2i_'+name
+            if color_adj and (ctype=='vision'):
+                keep_ratio = cfguh().cfg.eval.get('color_adj_keep_ratio', 0.5)
+                simple = cfguh().cfg.eval.get('color_adj_simple', True)
+                x_adj = []
+                for xi in x:
+                    color_adj_f = color_adjust(ref_from=(xi+1)/2, ref_to=color_adj_to)
+                    xi_adj = color_adj_f((xi+1)/2, keep=keep_ratio, simple=simple)
+                    x_adj.append(xi_adj)
+                x = x_adj
+            self.save_images(x, name, path, suffix=suffix)
+        elif xtype == 'text':
+            prompt_temperature = cfguh().cfg.eval.get('prompt_temperature', 1.0)
+            x = netm.optimus_decode(z, temperature=prompt_temperature)
+            name = 't2t_'+name if ctype == 'prompt' else 'v2t_'+name
+            prompt_merge_same_adj_word = cfguh().cfg.eval.get('prompt_merge_same_adj_word', False)
+            if prompt_merge_same_adj_word:
+                xnew = []
+                for xi in x:
+                    xi_split = xi.split()
+                    xinew = []
+                    for idxi, wi in enumerate(xi_split):
+                        if idxi!=0 and wi==xi_split[idxi-1]:
+                            continue
+                        xinew.append(wi)
+                    xnew.append(' '.join(xinew))
+                x = xnew
+            self.save_text(x, name, path, suffix=suffix)
+
+    def save_images(self, x, name, path, suffix=''):
+        if isinstance(x, torch.Tensor):
+            single_input = len(x.shape) == 3
+            if single_input:
+                x = x[None]
+            x = torch.clamp((x+1.0)/2.0, min=0.0, max=1.0)
+            x = [tvtrans.ToPILImage()(xi) for xi in x]
+            xlist = [np.array(xi) for xi in x]
+        elif isinstance(x, list):
+            xlist = x
+        canvas = auto_merge_imlist(xlist)
+        image_name = '{}{}.png'.format(name, suffix)
+        PIL.Image.fromarray(canvas).save(osp.join(path, image_name))
+
+    def save_text(self, x, name, path, suffix=''):
+        file_name = '{}{}.txt'.format(name, suffix)
+        with open(osp.join(path, file_name) ,'w') as f:
+            for xi in x:
+                f.write(xi+'\n')
+
+    def __call__(self, **paras):
+        cfg = cfguh().cfg
+        cfgv = cfg.eval
+
+        net = self.get_net(paras)
+        eval_cnt = paras.get('eval_cnt', None)
+        fix_seed = cfgv.get('fix_seed', False)
+
+        LRANK = sync.get_rank('local')
+        LWSIZE = sync.get_world_size('local')
+
+        output_path = self.get_image_path()
+        self.create_dir(output_path)
+        eval_cnt = paras.get('eval_cnt', None)
+        suffix='' if eval_cnt is None else '_'+str(eval_cnt)
+
+        if isinstance(net, (torch.nn.DataParallel,
+                            torch.nn.parallel.DistributedDataParallel)):
+            netm = net.module
+        else:
+            netm = net
+
+        with_ema = getattr(netm, 'model_ema', None) is not None
+        sampler = self.sampler(netm)
+        setattr(netm, 'device', LRANK) # Trick
+
+        color_adj = cfguh().cfg.eval.get('color_adj', False)
+
+        replicate = cfgv.get('replicate', 1)
+        conditioning = cfgv.conditioning * replicate
+        conditioning_local = conditioning[LRANK : len(conditioning) : LWSIZE]
+        seed_increment = [i for i in range(len(conditioning))][LRANK : len(conditioning) : LWSIZE]
+
+        for conditioningi, seedi in zip(conditioning_local, seed_increment):
+            if conditioningi == 'SKIP':
+                continue
+
+            ci, otypei = conditioningi
+
+            if osp.isfile(ci):
+                # is vision
+                output_name = osp.splitext(osp.basename(ci))[0]
+                ci = tvtrans.ToTensor()(PIL.Image.open(ci))
+                ci = ci*2 - 1
+                ctypei = 'vision'
+            else:
+                # is prompt
+                output_name = ci.strip().replace(' ', '-')
+                ctypei = 'prompt'
+
+            if fix_seed:
+                np.random.seed(cfg.env.rnd_seed + seedi)
+                torch.manual_seed(cfg.env.rnd_seed + seedi + 100)
+                suffixi = suffix + "_seed{}".format(cfg.env.rnd_seed + seedi + 100)
+            else:
+                suffixi = suffix
+
+            if with_ema:
+                with netm.ema_scope():
+                    z, _ = self.sample(netm, sampler, ci, otypei, ctypei, **cfgv.sample)
+            else:
+                z, _ = self.sample(netm, sampler, ci, otypei, ctypei, **cfgv.sample)
+
+            self.decode_and_save(
+                netm, z, otypei, ctypei, output_path, output_name, suffixi,
+                color_adj=color_adj, color_adj_to=conditioningi[0],)
+
+        if eval_cnt is not None:
+            print_log('Demo printed for {}'.format(eval_cnt))
+        return {}
+
+################
+# basic stages #
+################
+
+class eval_stage_basic(eval_stage_base):
+    @torch.no_grad()
+    def sample(self, net, sampler, visual_hint, output_dim, scale, n_samples, ddim_steps, ddim_eta):
+        h, w = output_dim
+        vh = PIL.Image.open(visual_hint)
+        c = net.clip_encode_vision(n_samples * [vh])
+        uc = None
+        if scale != 1.0:
+            dummy = torch.zeros_like(tvtrans.ToTensor()(vh))
+            uc = net.clip_encode_vision(n_samples * [dummy])
+
+        shape = [4, h//8, w//8]
+        rv = sampler.sample(
+            S=ddim_steps,
+            conditioning=c,
+            batch_size=n_samples,
+            shape=shape,
+            verbose=False,
+            unconditional_guidance_scale=scale,
+            unconditional_conditioning=uc,
+            eta=ddim_eta)
+        return rv
+
+    def __call__(self, **paras):
+        cfg = cfguh().cfg
+        cfgv = cfg.eval
+
+        net = paras['net']
+        eval_cnt = paras.get('eval_cnt', None)
+        fix_seed = cfgv.get('fix_seed', False)
+
+        LRANK = sync.get_rank('local')
+        LWSIZE = sync.get_world_size('local')
+
+        image_path = self.get_image_path()
+        self.create_dir(image_path)
+        eval_cnt = paras.get('eval_cnt', None)
+        suffix='' if eval_cnt is None else '_'+str(eval_cnt)
+
+        if isinstance(net, (torch.nn.DataParallel,
+                            torch.nn.parallel.DistributedDataParallel)):
+            netm = net.module
+        else:
+            netm = net
+
+        with_ema = getattr(netm, 'model_ema', None) is not None
+        sampler = self.sampler(netm)
+        setattr(netm, 'device', LRANK) # Trick
+
+        color_adj = cfguh().cfg.eval.get('color_adj', False)
+        color_adj_keep_ratio = cfguh().cfg.eval.get('color_adj_keep_ratio', 0.5)
+        color_adj_simple = cfguh().cfg.eval.get('color_adj_simple', True)
+
+        replicate = cfgv.get('replicate', 1)
+        conditioning = cfgv.conditioning * replicate
+        conditioning_local = conditioning[LRANK : len(conditioning) : LWSIZE]
+        seed_increment = [i for i in range(len(conditioning))][LRANK : len(conditioning) : LWSIZE]
+
+        for ci, seedi in zip(conditioning_local, seed_increment):
+            if ci == 'SKIP':
+                continue
+            draw_filename = osp.splitext(osp.basename(ci))[0]
+            if fix_seed:
+                np.random.seed(cfg.env.rnd_seed + seedi)
+                torch.manual_seed(cfg.env.rnd_seed + seedi + 100)
+                suffixi = suffix + "_seed{}".format(cfg.env.rnd_seed + seedi + 100)
+            else:
+                suffixi = suffix
+
+            if with_ema:
+                with netm.ema_scope():
+                    x, _ = self.sample(netm, sampler, ci, **cfgv.sample)
+            else:
+                x, _ = self.sample(netm, sampler, ci, **cfgv.sample)
+
+            demo_image = latent2im(netm, x)
+            if color_adj:
+                x_adj = []
+                for demoi in demo_image:
+                    color_adj_f = color_adjust(ref_from=demoi, ref_to=ci)
+                    xi_adj = color_adj_f(demoi, keep=color_adj_keep_ratio, simple=color_adj_simple)
+                    x_adj.append(xi_adj)
+                demo_image = x_adj
+            self.save_images(demo_image, draw_filename, image_path, suffix=suffixi)
+
+        if eval_cnt is not None:
+            print_log('Demo printed for {}'.format(eval_cnt))
+        return {}
+
+#######################
+# dual context stages #
+#######################
+
+class eval_stage_dc(eval_stage_base):
+    def __init__(self):
+        from ..model_zoo.ddim_dualcontext import DDIMSampler_DualContext
+        self.sampler = DDIMSampler_DualContext
+
+    @torch.no_grad()
+    def sample(
+            self, net, sampler, conditioning, output_dim, 
+            scale, n_samples, ddim_steps, ddim_eta):
+        ctype, cvalue =conditioning
+        if ctype == 'prompt':
+            return self.sample_text(
+                net, sampler, cvalue, output_dim,
+                scale, n_samples, ddim_steps, ddim_eta)
+        elif ctype == 'vision':
+            return self.sample_vision(
+                net, sampler, cvalue, output_dim,
+                scale, n_samples, ddim_steps, ddim_eta)
+        else:
+            raise ValueError
+
+    @torch.no_grad()
+    def sample_text(
+            self, net, sampler, prompt, output_dim, 
+            scale, n_samples, ddim_steps, ddim_eta):
+        h, w = output_dim
+        uc = None
+        if scale != 1.0:
+            uc = net.clip_encode_text(n_samples * [""])
+        c = net.clip_encode_text(n_samples * [prompt])
+        shape = [n_samples, 4, h//8, w//8]
+        rv = sampler.sample_text(
+            steps=ddim_steps,
+            shape=shape,
+            conditioning=c,
+            unconditional_guidance_scale=scale,
+            unconditional_conditioning=uc,
+            eta=ddim_eta,
+            verbose=False,)
+        return rv
+
+    @torch.no_grad()
+    def sample_vision(
+            self, net, sampler, visual_hint, output_dim, 
+            scale, n_samples, ddim_steps, ddim_eta):
+        h, w = output_dim
+        if len(visual_hint.shape) == 3:
+            visual_hint=visual_hint[None].repeat(n_samples, 1, 1, 1)
+        else:
+            raise ValueError
+
+        c = net.clip_encode_vision(visual_hint)
+        uc = None
+        if scale != 1.0:
+            visual_hint_blank = torch.zeros_like(visual_hint)
+            uc = net.clip_encode_vision(visual_hint_blank)
+
+        shape = [n_samples, 4, h//8, w//8]
+        rv = sampler.sample_vision(
+            steps=ddim_steps,
+            shape=shape,
+            conditioning=c,
+            unconditional_guidance_scale=scale,
+            unconditional_conditioning=uc,
+            eta=ddim_eta,
+            verbose=False,)
+        return rv
+
+    def __call__(self, **paras):
+        cfg = cfguh().cfg
+        cfgv = cfg.eval
+
+        net = self.get_net(paras)
+        eval_cnt = paras.get('eval_cnt', None)
+        fix_seed = cfgv.get('fix_seed', False)
+
+        LRANK = sync.get_rank('local')
+        LWSIZE = sync.get_world_size('local')
+
+        image_path = self.get_image_path()
+        self.create_dir(image_path)
+        suffix='' if eval_cnt is None else '_'+str(eval_cnt)
+
+        if isinstance(net, (torch.nn.DataParallel,
+                            torch.nn.parallel.DistributedDataParallel)):
+            netm = net.module
+        else:
+            netm = net
+
+        with_ema = getattr(netm, 'model_ema', None) is not None
+        sampler = self.sampler(netm)
+        setattr(netm, 'device', LRANK) # Trick
+
+        color_adj = cfguh().cfg.eval.get('color_adj', False)
+        color_adj_keep_ratio = cfguh().cfg.eval.get('color_adj_keep_ratio', 0.5)
+        color_adj_simple = cfguh().cfg.eval.get('color_adj_simple', True)
+
+        replicate = cfgv.get('replicate', 1)
+        conditioning = cfgv.conditioning * replicate
+        conditioning_local = conditioning[LRANK : len(conditioning) : LWSIZE]
+        seed_increment = [i for i in range(len(conditioning))][LRANK : len(conditioning) : LWSIZE]
+
+        for ci, seedi in zip(conditioning_local, seed_increment):
+            if ci == 'SKIP':
+                continue
+
+            if osp.isfile(ci):
+                # is vision
+                draw_filename = 'v2i_' + osp.splitext(osp.basename(ci))[0]
+                ci = tvtrans.ToTensor()(PIL.Image.open(ci))
+                ci = ci*2 - 1
+                ci = ('vision', ci)
+            else:
+                # is prompt
+                draw_filename = 't2i_' + ci.strip().replace(' ', '-')
+                ci = ('prompt', ci)
+
+            if fix_seed:
+                np.random.seed(cfg.env.rnd_seed + seedi)
+                torch.manual_seed(cfg.env.rnd_seed + seedi + 100)
+                suffixi = suffix + "_seed{}".format(cfg.env.rnd_seed + seedi + 100)
+            else:
+                suffixi = suffix
+
+            if with_ema:
+                with netm.ema_scope():
+                    x, _ = self.sample(netm, sampler, ci, **cfgv.sample)
+            else:
+                x, _ = self.sample(netm, sampler, ci, **cfgv.sample)
+
+            demo_image = latent2im(netm, x)
+            if color_adj and ci[0] == 'vision':
+                x_adj = []
+                for demoi in demo_image:
+                    color_adj_f = color_adjust(ref_from=demoi, ref_to=ci[1])
+                    xi_adj = color_adj_f(demoi, keep=color_adj_keep_ratio, simple=color_adj_simple)
+                    x_adj.append(xi_adj)
+                demo_image = x_adj
+            self.save_images(demo_image, draw_filename, image_path, suffix=suffixi)
+
+        if eval_cnt is not None:
+            print_log('Demo printed for {}'.format(eval_cnt))
+        return {}
+

versatile_diffusion/lib/model_zoo/__init__.py

@@ -0,0 +1,4 @@
+from .common.get_model import get_model
+from .common.get_optimizer import get_optimizer
+from .common.get_scheduler import get_scheduler
+from .common.utils import get_unit