Add files using upload-large-folder tool

.gitattributes

@@ -2985,3 +2985,14 @@ results/versatile_diffusion/subj01/roi/12.png filter=lfs diff=lfs merge=lfs -tex
 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
+vdvae/header-image.png filter=lfs diff=lfs merge=lfs -text
+data/predicted_features/subj01/nsd_vdvae_nsdgeneral_pred_sub1_31l_alpha50k.npy filter=lfs diff=lfs merge=lfs -text
+versatile_diffusion/pretrained/kl-f8.pth filter=lfs diff=lfs merge=lfs -text
+data/processed_data/subj01/nsd_train_fmriavg_nsdgeneral_sub1.npy filter=lfs diff=lfs merge=lfs -text
+vdvae/model/imagenet64-iter-1600000-model-ema.th filter=lfs diff=lfs merge=lfs -text
+vdvae/model/imagenet64-iter-1600000-model.th filter=lfs diff=lfs merge=lfs -text
+data/predicted_features/subj01/nsd_clipvision_predtest_nsdgeneral.npy filter=lfs diff=lfs merge=lfs -text
+versatile_diffusion/pretrained/optimus-vae.pth filter=lfs diff=lfs merge=lfs -text
+vdvae/model/imagenet64-iter-1600000-opt.th filter=lfs diff=lfs merge=lfs -text
+data/extracted_features/subj01/nsd_vdvae_features_31l.npz filter=lfs diff=lfs merge=lfs -text
+data/extracted_features/subj01/nsd_cliptext_train.npy filter=lfs diff=lfs merge=lfs -text

data/extracted_features/subj01/nsd_cliptext_train.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe33228e2c7334c50c6b09d790c17c188787c46ed89addd37efe2c8f0dfc9e44
+size 4191086720

data/extracted_features/subj01/nsd_vdvae_features_31l.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a919a9bd590e71398e45fff63895051f50230151547f398b584e83aa434fc38f
+size 3588737796

data/predicted_features/subj01/nsd_clipvision_predtest_nsdgeneral.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba7dc9d374d06dc67d8b3b0fe88e2961bce2e44b549081c702023d94183d53a0
+size 1550585984

data/predicted_features/subj01/nsd_vdvae_nsdgeneral_pred_sub1_31l_alpha50k.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5db240ec4b4ad976b2f5ef847ae6277342565de10dadb581ae7748ba0b4a91a6
+size 716215936

data/processed_data/subj01/nsd_train_fmriavg_nsdgeneral_sub1.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26b874870b202e820f290617dcb15fe88af5908cdbf23b25c5a743bcf404aa12
+size 1114391456

vdvae/model/imagenet64-iter-1600000-model-ema.th

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d3b9e07177f5c5082e31685cc47f4e53f8557e4a5eefb72bca514445d3bc48e
+size 500977841

vdvae/model/imagenet64-iter-1600000-model.th

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1701a7e831b370ceaa5f3b57a9d8e7b610b9d10cffd3083dae11986eea95bbb
+size 501006513

vdvae/model/imagenet64-iter-1600000-opt.th

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:855bc24ebac829a7ae3cdb6fd8c40a91bde7faab4f5ed857a94a1480d03de6f8
+size 1001616051

versatile_diffusion/lib/model_zoo/optimus_models/vocab/bert_vocab_download_info.json

@@ -0,0 +1,15 @@
+{
+    "bert-base-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-vocab.txt",
+    "bert-large-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-vocab.txt",
+    "bert-base-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-vocab.txt",
+    "bert-large-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-vocab.txt",
+    "bert-base-multilingual-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-vocab.txt",
+    "bert-base-multilingual-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-vocab.txt",
+    "bert-base-chinese": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-vocab.txt",
+    "bert-base-german-cased": "https://int-deepset-models-bert.s3.eu-central-1.amazonaws.com/pytorch/bert-base-german-cased-vocab.txt",
+    "bert-large-uncased-whole-word-masking": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-vocab.txt",
+    "bert-large-cased-whole-word-masking": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-vocab.txt",
+    "bert-large-uncased-whole-word-masking-finetuned-squad": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-vocab.txt",
+    "bert-large-cased-whole-word-masking-finetuned-squad": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-vocab.txt",
+    "bert-base-cased-finetuned-mrpc": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-vocab.txt"
+}

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/configs/model/sd.yaml

@@ -0,0 +1,68 @@
+sd_base:
+  symbol: sd
+  find_unused_parameters: true
+
+sd_autoencoder:
+  type: autoencoderkl
+  args:
+    embed_dim: 4
+    monitor: val/rec_loss
+    ddconfig:
+      double_z: true
+      z_channels: 4
+      resolution: 256
+      in_channels: 3
+      out_ch: 3
+      ch: 128
+      ch_mult: [1, 2, 4, 4]
+      num_res_blocks: 2
+      attn_resolutions: []
+      dropout: 0.0
+    lossconfig:
+      target: torch.nn.Identity
+  pth: pretrained/kl-f8.pth
+
+sd_t2i:
+  super_cfg: sd_base
+  type: sd_t2i
+  args:
+    first_stage_config: MODEL(sd_autoencoder)
+    cond_stage_config: MODEL(clip_text_frozen)
+    unet_config: MODEL(openai_unet_sd)
+    beta_linear_start: 0.00085
+    beta_linear_end: 0.012
+    num_timesteps_cond: 1
+    timesteps: 1000
+    scale_factor: 0.18215
+    use_ema: true
+
+sd_t2i_noema:
+  super_cfg: sd
+  args:
+    use_ema: false
+
+#####################
+# sd with full clip #
+#####################
+
+sd_t2i_fullclip_backward_compatible:
+  super_cfg: sd_t2i
+  args:
+    cond_stage_config: MODEL(clip_frozen_encode_text_noproj)
+
+sd_t2i_fullclip_backward_compatible_noema:
+  super_cfg: sd_t2i_noema
+  args:
+    cond_stage_config: MODEL(clip_frozen_encode_text_noproj)
+
+sd_t2i_fullclip:
+  super_cfg: sd_t2i
+  args:
+    cond_stage_config: MODEL(clip_frozen_encode_text)
+
+sd_variation:
+  super_cfg: sd_t2i
+  type: sd_variation
+  args:
+    cond_stage_config: MODEL(clip_vision_frozen_justin)
+ 

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/configs/model/vd.yaml

@@ -0,0 +1,61 @@
+# vd_base:
+#   symbol: vd
+#   find_unused_parameters: true
+
+############
+# vd basic #
+############
+
+vd_basic:
+  super_cfg: sd_t2i
+  type: vd_basic
+  symbol: vd
+  find_unused_parameters: true
+  args:
+    cond_stage_config: MODEL(clip_frozen_encode_vision)
+
+vd_basic_noema:
+  super_cfg: vd_basic
+  args:
+    use_ema: false
+
+###################
+# vd dual-context #
+###################
+
+vd_dc:
+  super_cfg: sd_t2i_fullclip
+  type: vd_dc
+  symbol: vd
+  find_unused_parameters: true
+  args:
+    unet_config: MODEL(openai_unet_dual_context)
+
+vd_dc_noema:
+  super_cfg: vd_dc
+  args:
+    use_ema: false
+
+######
+# vd #
+######
+
+vd:
+  type: vd
+  symbol: vd
+  find_unused_parameters: true
+  args:
+    autokl_cfg: MODEL(sd_autoencoder)
+    optimus_cfg: MODEL(optimus_vae)
+    clip_cfg: MODEL(clip_frozen)
+    unet_config: MODEL(openai_unet_vd)
+    beta_linear_start: 0.00085
+    beta_linear_end: 0.012
+    timesteps: 1000
+    scale_factor: 0.18215
+    use_ema: true
+
+vd_noema:
+  super_cfg: vd
+  args:
+    use_ema: false

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/cfg_helper.py

@@ -0,0 +1,664 @@
+import os
+import os.path as osp
+import shutil
+import copy
+import time
+import pprint
+import numpy as np
+import torch
+import matplotlib
+import argparse
+import json
+import yaml
+from easydict import EasyDict as edict
+
+from .model_zoo import get_model
+
+############
+# cfg_bank #
+############
+
+def cfg_solvef(cmd, root):
+    if not isinstance(cmd, str):
+        return cmd
+    
+    if cmd.find('SAME')==0:
+        zoom = root
+        p = cmd[len('SAME'):].strip('()').split('.')
+        p = [pi.strip() for pi in p]
+        for pi in p:
+            try:
+                pi = int(pi)
+            except:
+                pass
+
+            try:
+                zoom = zoom[pi]
+            except:
+                return cmd
+        return cfg_solvef(zoom, root)
+
+    if cmd.find('SEARCH')==0:
+        zoom = root
+        p = cmd[len('SEARCH'):].strip('()').split('.')
+        p = [pi.strip() for pi in p]
+        find = True
+        # Depth first search
+        for pi in p:
+            try:
+                pi = int(pi)
+            except:
+                pass
+            
+            try:
+                zoom = zoom[pi]
+            except:
+                find = False
+                break
+
+        if find:
+            return cfg_solvef(zoom, root)
+        else:
+            if isinstance(root, dict):
+                for ri in root:
+                    rv = cfg_solvef(cmd, root[ri])
+                    if rv != cmd:
+                        return rv
+            if isinstance(root, list):
+                for ri in root:
+                    rv = cfg_solvef(cmd, ri)
+                    if rv != cmd:
+                        return rv
+            return cmd
+
+    if cmd.find('MODEL')==0:
+        goto = cmd[len('MODEL'):].strip('()')
+        return model_cfg_bank()(goto)
+
+    if cmd.find('DATASET')==0:
+        goto = cmd[len('DATASET'):].strip('()')
+        return dataset_cfg_bank()(goto)
+
+    return cmd
+
+def cfg_solve(cfg, cfg_root):
+    # The function solve cfg element such that 
+    #   all sorrogate input are settled.
+    #   (i.e. SAME(***) ) 
+    if isinstance(cfg, list):
+        for i in range(len(cfg)):
+            if isinstance(cfg[i], (list, dict)):
+                cfg[i] = cfg_solve(cfg[i], cfg_root)
+            else:
+                cfg[i] = cfg_solvef(cfg[i], cfg_root)
+    if isinstance(cfg, dict):
+        for k in cfg:
+            if isinstance(cfg[k], (list, dict)):
+                cfg[k] = cfg_solve(cfg[k], cfg_root)
+            else:
+                cfg[k] = cfg_solvef(cfg[k], cfg_root)        
+    return cfg
+
+class model_cfg_bank(object):
+    def __init__(self):
+        self.cfg_dir = osp.join('configs', 'model')
+        self.cfg_bank = edict()
+    
+    def __call__(self, name):
+        if name not in self.cfg_bank:
+            cfg_path = self.get_yaml_path(name)
+            with open(cfg_path, 'r') as f:
+                cfg_new = yaml.load(
+                    f, Loader=yaml.FullLoader)
+            cfg_new = edict(cfg_new)
+            self.cfg_bank.update(cfg_new)
+
+        cfg = self.cfg_bank[name]
+        cfg.name = name
+        if 'super_cfg' not in cfg:
+            cfg = cfg_solve(cfg, cfg)
+            self.cfg_bank[name] = cfg
+            return copy.deepcopy(cfg)
+
+        super_cfg = self.__call__(cfg.super_cfg)
+        # unlike other field,
+        # args will not be replaced but update.
+        if 'args' in cfg:
+            if 'args' in  super_cfg:
+                super_cfg.args.update(cfg.args)
+            else:
+                super_cfg.args = cfg.args
+            cfg.pop('args')
+
+        super_cfg.update(cfg)
+        super_cfg.pop('super_cfg')
+        cfg = super_cfg
+        try:
+            delete_args = cfg.pop('delete_args')
+        except:
+            delete_args = []
+
+        for dargs in delete_args:
+            cfg.args.pop(dargs)
+
+        cfg = cfg_solve(cfg, cfg)
+        self.cfg_bank[name] = cfg
+        return copy.deepcopy(cfg)
+
+    def get_yaml_path(self, name):
+        if name.find('ldm')==0:
+            return osp.join(
+                self.cfg_dir, 'ldm.yaml')
+        elif name.find('comodgan')==0:
+            return osp.join(
+                self.cfg_dir, 'comodgan.yaml')
+        elif name.find('stylegan')==0:
+            return osp.join(
+                self.cfg_dir, 'stylegan.yaml')
+        elif name.find('absgan')==0:
+            return osp.join(
+                self.cfg_dir, 'absgan.yaml')
+        elif name.find('ashgan')==0:
+            return osp.join(
+                self.cfg_dir, 'ashgan.yaml')
+        elif name.find('sr3')==0:
+            return osp.join(
+                self.cfg_dir, 'sr3.yaml')
+        elif name.find('specdiffsr')==0:
+            return osp.join(
+                self.cfg_dir, 'specdiffsr.yaml')
+        elif name.find('openai_unet')==0:
+            return osp.join(
+                self.cfg_dir, 'openai_unet.yaml')
+        elif name.find('clip')==0:
+            return osp.join(
+                self.cfg_dir, 'clip.yaml')
+        elif name.find('sd')==0:
+            return osp.join(
+                self.cfg_dir, 'sd.yaml')
+        elif name.find('vd')==0:
+            return osp.join(
+                self.cfg_dir, 'vd.yaml')
+        elif name.find('optimus')==0:
+            return osp.join(
+                self.cfg_dir, 'optimus.yaml')
+        else:
+            raise ValueError
+
+class dataset_cfg_bank(object):
+    def __init__(self):
+        self.cfg_dir = osp.join('configs', 'dataset')
+        self.cfg_bank = edict()
+
+    def __call__(self, name):
+        if name not in self.cfg_bank:
+            cfg_path = self.get_yaml_path(name)
+            with open(cfg_path, 'r') as f:
+                cfg_new = yaml.load(
+                    f, Loader=yaml.FullLoader)
+            cfg_new = edict(cfg_new)
+            self.cfg_bank.update(cfg_new)
+
+        cfg = self.cfg_bank[name]
+        cfg.name = name
+        if cfg.get('super_cfg', None) is None:
+            cfg = cfg_solve(cfg, cfg)
+            self.cfg_bank[name] = cfg
+            return copy.deepcopy(cfg)
+
+        super_cfg = self.__call__(cfg.super_cfg)
+        super_cfg.update(cfg)
+        cfg = super_cfg
+        cfg.super_cfg = None
+        try:
+            delete = cfg.pop('delete')
+        except:
+            delete = []
+
+        for dargs in delete:
+            cfg.pop(dargs)
+
+        cfg = cfg_solve(cfg, cfg)
+        self.cfg_bank[name] = cfg
+        return copy.deepcopy(cfg)
+
+    def get_yaml_path(self, name):
+        if name.find('cityscapes')==0:
+            return osp.join(
+                self.cfg_dir, 'cityscapes.yaml')
+        elif name.find('div2k')==0:
+            return osp.join(
+                self.cfg_dir, 'div2k.yaml')
+        elif name.find('gandiv2k')==0:
+            return osp.join(
+                self.cfg_dir, 'gandiv2k.yaml')
+        elif name.find('srbenchmark')==0:
+            return osp.join(
+                self.cfg_dir, 'srbenchmark.yaml')
+        elif name.find('imagedir')==0:
+            return osp.join(
+                self.cfg_dir, 'imagedir.yaml')
+        elif name.find('places2')==0:
+            return osp.join(
+                self.cfg_dir, 'places2.yaml')
+        elif name.find('ffhq')==0:
+            return osp.join(
+                self.cfg_dir, 'ffhq.yaml')
+        elif name.find('imcpt')==0:
+            return osp.join(
+                self.cfg_dir, 'imcpt.yaml')
+        elif name.find('texture')==0:
+            return osp.join(
+                self.cfg_dir, 'texture.yaml')
+        elif name.find('openimages')==0:
+            return osp.join(
+                self.cfg_dir, 'openimages.yaml')
+        elif name.find('laion2b')==0:
+            return osp.join(
+                self.cfg_dir, 'laion2b.yaml')
+        elif name.find('laionart')==0:
+            return osp.join(
+                self.cfg_dir, 'laionart.yaml')
+        elif name.find('celeba')==0:
+            return osp.join(
+                self.cfg_dir, 'celeba.yaml')
+        elif name.find('coyo')==0:
+            return osp.join(
+                self.cfg_dir, 'coyo.yaml')
+        elif name.find('pafc')==0:
+            return osp.join(
+                self.cfg_dir, 'pafc.yaml')
+        elif name.find('coco')==0:
+            return osp.join(
+                self.cfg_dir, 'coco.yaml')
+        else:
+            raise ValueError
+
+class experiment_cfg_bank(object):
+    def __init__(self):
+        self.cfg_dir = osp.join('configs', 'experiment')
+        self.cfg_bank = edict()
+
+    def __call__(self, name):
+        if name not in self.cfg_bank:
+            cfg_path = self.get_yaml_path(name)
+            with open(cfg_path, 'r') as f:
+                cfg = yaml.load(
+                    f, Loader=yaml.FullLoader)
+            cfg = edict(cfg)
+
+        cfg = cfg_solve(cfg, cfg)
+        cfg = cfg_solve(cfg, cfg) 
+        # twice for SEARCH
+        self.cfg_bank[name] = cfg
+        return copy.deepcopy(cfg)
+
+    def get_yaml_path(self, name):
+        return osp.join(
+            self.cfg_dir, name+'.yaml')
+
+def load_cfg_yaml(path):
+    if osp.isfile(path):
+        cfg_path = path
+    elif osp.isfile(osp.join('configs', 'experiment', path)):
+        cfg_path = osp.join('configs', 'experiment', path)
+    elif osp.isfile(osp.join('configs', 'experiment', path+'.yaml')):
+        cfg_path = osp.join('configs', 'experiment', path+'.yaml')
+    else:
+        assert False, 'No such config!'
+
+    with open(cfg_path, 'r') as f:
+        cfg = yaml.load(f, Loader=yaml.FullLoader)
+        cfg = edict(cfg)
+    cfg = cfg_solve(cfg, cfg)
+    cfg = cfg_solve(cfg, cfg)
+    return cfg
+
+##############
+# cfg_helper #
+##############
+
+def get_experiment_id(ref=None):
+    if ref is None:
+        time.sleep(0.5)
+        return int(time.time()*100)
+    else:
+        try:
+            return int(ref)
+        except:
+            pass
+        
+        _, ref = osp.split(ref)
+        ref = ref.split('_')[0]
+        try:
+            return int(ref)
+        except:
+            assert False, 'Invalid experiment ID!'
+
+def record_resume_cfg(path):
+    cnt = 0
+    while True:
+        if osp.exists(path+'.{:04d}'.format(cnt)):
+            cnt += 1
+            continue
+        shutil.copyfile(path, path+'.{:04d}'.format(cnt))
+        break
+
+def get_command_line_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--debug', action='store_true', default=False)
+    parser.add_argument('--config', type=str)
+    parser.add_argument('--gpu', nargs='+', type=int)
+
+    parser.add_argument('--node_rank', type=int, default=0)
+    parser.add_argument('--nodes', type=int, default=1)
+    parser.add_argument('--addr', type=str, default='127.0.0.1')
+    parser.add_argument('--port', type=int, default=11233)
+ 
+    parser.add_argument('--signature', nargs='+', type=str)
+    parser.add_argument('--seed', type=int)
+
+    parser.add_argument('--eval', type=str)
+    parser.add_argument('--eval_subdir', type=str)
+    parser.add_argument('--pretrained', type=str)
+
+    parser.add_argument('--resume_dir', type=str)
+    parser.add_argument('--resume_step', type=int)
+    parser.add_argument('--resume_weight', type=str)
+
+    args = parser.parse_args()
+
+    # Special handling the resume
+    if args.resume_dir is not None:
+        cfg = edict()
+        cfg.env = edict()
+        cfg.env.debug = args.debug
+        cfg.env.resume = edict()
+        cfg.env.resume.dir = args.resume_dir
+        cfg.env.resume.step = args.resume_step
+        cfg.env.resume.weight = args.resume_weight
+        return cfg
+
+    cfg = load_cfg_yaml(args.config)
+    cfg.env.debug = args.debug
+    cfg.env.gpu_device = [0] if args.gpu is None else list(args.gpu)
+    cfg.env.master_addr = args.addr
+    cfg.env.master_port = args.port
+    cfg.env.dist_url = 'tcp://{}:{}'.format(args.addr, args.port)
+    cfg.env.node_rank = args.node_rank
+    cfg.env.nodes = args.nodes
+
+    istrain = False if args.eval is not None else True
+    isdebug = cfg.env.debug
+
+    if istrain:
+        if isdebug:
+            cfg.env.experiment_id = 999999999999
+            cfg.train.signature = ['debug']
+        else:
+            cfg.env.experiment_id = get_experiment_id()
+            if args.signature is not None:
+                cfg.train.signature = args.signature
+    else:
+        if 'train' in cfg:
+            cfg.pop('train')
+        cfg.env.experiment_id = get_experiment_id(args.eval)
+        if args.signature is not None:
+            cfg.eval.signature = args.signature
+
+        if isdebug and (args.eval is None):
+            cfg.env.experiment_id = 999999999999
+            cfg.eval.signature = ['debug']
+
+        if args.eval_subdir is not None:
+            if isdebug:
+                cfg.eval.eval_subdir = 'debug'
+            else:
+                cfg.eval.eval_subdir = args.eval_subdir
+        if args.pretrained is not None:
+            cfg.eval.pretrained = args.pretrained 
+          # The override pretrained over the setting in cfg.model
+
+    if args.seed is not None:
+        cfg.env.rnd_seed = args.seed
+
+    return cfg
+
+def cfg_initiates(cfg):
+    cfge = cfg.env
+    isdebug = cfge.debug
+    isresume = 'resume' in cfge
+    istrain = 'train' in cfg
+    haseval = 'eval' in cfg
+    cfgt = cfg.train if istrain else None
+    cfgv = cfg.eval if haseval else None
+
+    ###############################
+    # get some environment params #
+    ###############################
+    
+    cfge.computer = os.uname()
+    cfge.torch_version = str(torch.__version__)
+
+    ##########
+    # resume #
+    ##########
+
+    if isresume:
+        resume_cfg_path = osp.join(cfge.resume.dir, 'config.yaml')
+        record_resume_cfg(resume_cfg_path)
+        with open(resume_cfg_path, 'r') as f:
+            cfg_resume = yaml.load(f, Loader=yaml.FullLoader)
+        cfg_resume = edict(cfg_resume)
+        cfg_resume.env.update(cfge)
+        cfg = cfg_resume
+        cfge = cfg.env
+        log_file = cfg.train.log_file
+
+        print('')
+        print('##########')
+        print('# resume #')
+        print('##########')
+        print('')
+        with open(log_file, 'a') as f:
+            print('', file=f)
+            print('##########', file=f)
+            print('# resume #', file=f)
+            print('##########', file=f)
+            print('', file=f)
+
+        pprint.pprint(cfg)
+        with open(log_file, 'a') as f:
+            pprint.pprint(cfg, f)
+
+    ####################
+    # node distributed #
+    ####################
+
+    if cfg.env.master_addr!='127.0.0.1':
+        os.environ['MASTER_ADDR'] = cfge.master_addr
+        os.environ['MASTER_PORT'] = '{}'.format(cfge.master_port)
+        if cfg.env.dist_backend=='nccl':
+            os.environ['NCCL_SOCKET_FAMILY'] = 'AF_INET'
+        if cfg.env.dist_backend=='gloo':
+            os.environ['GLOO_SOCKET_FAMILY'] = 'AF_INET'
+    
+    #######################
+    # cuda visible device #
+    #######################
+
+    os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(
+        [str(gid) for gid in cfge.gpu_device]) 
+
+    #####################
+    # return resume cfg #
+    #####################
+
+    if isresume:
+        return cfg
+
+    #############################################
+    # some misc setting that not need in resume #
+    #############################################
+
+    cfgm = cfg.model
+    cfge.gpu_count = len(cfge.gpu_device)
+
+    ##########################################
+    # align batch size and num worker config #
+    ##########################################
+ 
+    gpu_n = cfge.gpu_count * cfge.nodes
+    def align_batch_size(bs, bs_per_gpu): 
+        assert (bs is not None) or (bs_per_gpu is not None)
+        bs = bs_per_gpu * gpu_n if bs is None else bs
+        bs_per_gpu = bs // gpu_n if bs_per_gpu is None else bs_per_gpu
+        assert (bs == bs_per_gpu * gpu_n)
+        return bs, bs_per_gpu
+
+    if istrain:
+        cfgt.batch_size, cfgt.batch_size_per_gpu = \
+            align_batch_size(cfgt.batch_size, cfgt.batch_size_per_gpu)
+        cfgt.dataset_num_workers, cfgt.dataset_num_workers_per_gpu = \
+            align_batch_size(cfgt.dataset_num_workers, cfgt.dataset_num_workers_per_gpu)
+    if haseval:
+        cfgv.batch_size, cfgv.batch_size_per_gpu = \
+            align_batch_size(cfgv.batch_size, cfgv.batch_size_per_gpu)
+        cfgv.dataset_num_workers, cfgv.dataset_num_workers_per_gpu = \
+            align_batch_size(cfgv.dataset_num_workers, cfgv.dataset_num_workers_per_gpu)
+
+    ##################
+    # create log dir #
+    ##################
+
+    if istrain:
+        if not isdebug:
+            sig = cfgt.get('signature', [])
+            version = get_model().get_version(cfgm.type)
+            sig = sig + ['v{}'.format(version), 's{}'.format(cfge.rnd_seed)]
+        else:
+            sig = ['debug']
+
+        log_dir = [
+            cfge.log_root_dir, 
+            '{}_{}'.format(cfgm.symbol, cfgt.dataset.symbol),
+            '_'.join([str(cfge.experiment_id)] + sig)
+        ]
+        log_dir = osp.join(*log_dir)
+        log_file = osp.join(log_dir, 'train.log')
+        if not osp.exists(log_file):
+            os.makedirs(osp.dirname(log_file))
+        cfgt.log_dir = log_dir
+        cfgt.log_file = log_file
+
+        if haseval:
+            cfgv.log_dir = log_dir
+            cfgv.log_file = log_file
+    else:
+        model_symbol = cfgm.symbol
+        if cfgv.get('dataset', None) is None:
+            dataset_symbol = 'nodataset'
+        else:
+            dataset_symbol = cfgv.dataset.symbol
+
+        log_dir = osp.join(cfge.log_root_dir, '{}_{}'.format(model_symbol, dataset_symbol))
+        exp_dir = search_experiment_folder(log_dir, cfge.experiment_id)
+        if exp_dir is None:
+            if not isdebug:
+                sig = cfgv.get('signature', []) + ['evalonly']
+            else:
+                sig = ['debug']
+            exp_dir = '_'.join([str(cfge.experiment_id)] + sig)
+
+        eval_subdir = cfgv.get('eval_subdir', None)
+        # override subdir in debug mode (if eval_subdir is set)
+        eval_subdir = 'debug' if (eval_subdir is not None) and isdebug else eval_subdir 
+
+        if eval_subdir is not None:
+            log_dir = osp.join(log_dir, exp_dir, eval_subdir)
+        else:
+            log_dir = osp.join(log_dir, exp_dir)
+
+        disable_log_override = cfgv.get('disable_log_override', False)
+        if osp.isdir(log_dir):
+            if disable_log_override:
+                assert False, 'Override an exsited log_dir is disabled at [{}]'.format(log_dir)
+        else:
+            os.makedirs(log_dir)
+
+        log_file = osp.join(log_dir, 'eval.log')
+        cfgv.log_dir = log_dir
+        cfgv.log_file = log_file
+
+    ######################
+    # print and save cfg #
+    ######################
+
+    pprint.pprint(cfg)
+    with open(log_file, 'w') as f:
+        pprint.pprint(cfg, f)
+    with open(osp.join(log_dir, 'config.yaml'), 'w') as f:
+        yaml.dump(edict_2_dict(cfg), f)
+
+    #############
+    # save code #
+    #############
+
+    save_code = False
+    if istrain:
+        save_code = cfgt.get('save_code', False)
+    elif haseval:
+        save_code = cfgv.get('save_code', False)
+
+    if save_code:
+        codedir = osp.join(log_dir, 'code')
+        if osp.exists(codedir):
+            shutil.rmtree(codedir)
+        for d in ['configs', 'lib']:
+            fromcodedir = d
+            tocodedir = osp.join(codedir, d)
+            shutil.copytree(
+                fromcodedir, tocodedir, 
+                ignore=shutil.ignore_patterns(
+                    '*__pycache__*', '*build*'))
+        for codei in os.listdir('.'):
+            if osp.splitext(codei)[1] == 'py':
+                shutil.copy(codei, codedir)
+
+    #######################
+    # set matplotlib mode #
+    #######################
+
+    if 'matplotlib_mode' in cfge:
+        try:
+            matplotlib.use(cfge.matplotlib_mode)
+        except:
+            print('Warning: matplotlib mode [{}] failed to be set!'.format(cfge.matplotlib_mode))
+
+    return cfg
+
+def edict_2_dict(x):
+    if isinstance(x, dict):
+        xnew = {}
+        for k in x:
+            xnew[k] = edict_2_dict(x[k])
+        return xnew
+    elif isinstance(x, list):
+        xnew = []
+        for i in range(len(x)):
+            xnew.append( edict_2_dict(x[i]) )
+        return xnew
+    else:
+        return x
+
+def search_experiment_folder(root, exid):
+    target = None
+    for fi in os.listdir(root):
+        if not osp.isdir(osp.join(root, fi)):
+            continue
+        if int(fi.split('_')[0]) == exid:
+            if target is not None:
+                return None # duplicated
+            elif target is None:
+                target = fi
+    return target

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/cfg_holder.py

@@ -0,0 +1,28 @@
+import copy
+
+def singleton(class_):
+    instances = {}
+    def getinstance(*args, **kwargs):
+        if class_ not in instances:
+            instances[class_] = class_(*args, **kwargs)
+        return instances[class_]
+    return getinstance
+
+##############
+# cfg_holder #
+##############
+
+@singleton
+class cfg_unique_holder(object):
+    def __init__(self):
+        self.cfg = None
+        # this is use to track the main codes.
+        self.code = set()
+    def save_cfg(self, cfg):
+        self.cfg = copy.deepcopy(cfg)
+    def add_code(self, code):
+        """
+        A new main code is reached and 
+            its name is added.
+        """
+        self.code.add(code)

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/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/log/sd_nodataset/99999_evalonly/sd_variation/code/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/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/evaluator/__init__.py

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

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/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/log/sd_nodataset/99999_evalonly/sd_variation/code/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/log/sd_nodataset/99999_evalonly/sd_variation/code/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/log/sd_nodataset/99999_evalonly/sd_variation/code/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/log/sd_nodataset/99999_evalonly/sd_variation/code/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

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/attention.py

@@ -0,0 +1,435 @@
+from inspect import isfunction
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+from einops import rearrange, repeat
+
+from .diffusion_utils import checkpoint
+
+
+def exists(val):
+    return val is not None
+
+
+def uniq(arr):
+    return{el: True for el in arr}.keys()
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+
+def init_(tensor):
+    dim = tensor.shape[-1]
+    std = 1 / math.sqrt(dim)
+    tensor.uniform_(-std, std)
+    return tensor
+
+
+# feedforward
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        project_in = nn.Sequential(
+            nn.Linear(dim, inner_dim),
+            nn.GELU()
+        ) if not glu else GEGLU(dim, inner_dim)
+
+        self.net = nn.Sequential(
+            project_in,
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class LinearAttention(nn.Module):
+    def __init__(self, dim, heads=4, dim_head=32):
+        super().__init__()
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
+        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x)
+        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
+        k = k.softmax(dim=-1)  
+        context = torch.einsum('bhdn,bhen->bhde', k, v)
+        out = torch.einsum('bhde,bhdn->bhen', context, q)
+        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
+        return self.to_out(out)
+
+
+class SpatialSelfAttention(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.k = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.v = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=1,
+                                        stride=1,
+                                        padding=0)
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b,c,h,w = q.shape
+        q = rearrange(q, 'b c h w -> b (h w) c')
+        k = rearrange(k, 'b c h w -> b c (h w)')
+        w_ = torch.einsum('bij,bjk->bik', q, k)
+
+        w_ = w_ * (int(c)**(-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = rearrange(v, 'b c h w -> b c (h w)')
+        w_ = rearrange(w_, 'b i j -> b j i')
+        h_ = torch.einsum('bij,bjk->bik', v, w_)
+        h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
+        h_ = self.proj_out(h_)
+
+        return x+h_
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x, context=None, mask=None):
+        h = self.heads
+
+        q = self.to_q(x)
+        context = default(context, x)
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+
+        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+
+        if exists(mask):
+            mask = rearrange(mask, 'b ... -> b (...)')
+            max_neg_value = -torch.finfo(sim.dtype).max
+            mask = repeat(mask, 'b j -> (b h) () j', h=h)
+            sim.masked_fill_(~mask, max_neg_value)
+
+        # attention, what we cannot get enough of
+        attn = sim.softmax(dim=-1)
+
+        out = einsum('b i j, b j d -> b i d', attn, v)
+        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+        return self.to_out(out)
+
+
+class BasicTransformerBlock(nn.Module):
+    def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True,
+                 disable_self_attn=False):
+        super().__init__()
+        self.disable_self_attn = disable_self_attn
+        self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout,
+                                    context_dim=context_dim if self.disable_self_attn else None)  # is a self-attention if not self.disable_self_attn
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
+                                    heads=n_heads, dim_head=d_head, dropout=dropout)  # is self-attn if context is none
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+
+    def forward(self, x, context=None):
+        return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
+
+    def _forward(self, x, context=None):
+        x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
+        x = self.attn2(self.norm2(x), context=context) + x
+        x = self.ff(self.norm3(x)) + x
+        return x
+
+
+class SpatialTransformer(nn.Module):
+    """
+    Transformer block for image-like data.
+    First, project the input (aka embedding)
+    and reshape to b, t, d.
+    Then apply standard transformer action.
+    Finally, reshape to image
+    """
+    def __init__(self, in_channels, n_heads, d_head,
+                 depth=1, dropout=0., context_dim=None,
+                 disable_self_attn=False):
+        super().__init__()
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = Normalize(in_channels)
+
+        self.proj_in = nn.Conv2d(in_channels,
+                                 inner_dim,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+
+        self.transformer_blocks = nn.ModuleList(
+            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim,
+                                   disable_self_attn=disable_self_attn)
+                for d in range(depth)]
+        )
+
+        self.proj_out = zero_module(nn.Conv2d(inner_dim,
+                                              in_channels,
+                                              kernel_size=1,
+                                              stride=1,
+                                              padding=0))
+
+    def forward(self, x, context=None):
+        # note: if no context is given, cross-attention defaults to self-attention
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        x = self.proj_in(x)
+        x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
+        for block in self.transformer_blocks:
+            x = block(x, context=context)
+        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
+        x = self.proj_out(x)
+        return x + x_in
+
+
+##########################
+# transformer no context #
+##########################
+
+class BasicTransformerBlockNoContext(nn.Module):
+    def __init__(self, dim, n_heads, d_head, dropout=0., gated_ff=True, checkpoint=True):
+        super().__init__()
+        self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, 
+                                    dropout=dropout, context_dim=None)
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, 
+                                    dropout=dropout, context_dim=None)
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+
+    def forward(self, x):
+        return checkpoint(self._forward, (x,), self.parameters(), self.checkpoint)
+
+    def _forward(self, x):
+        x = self.attn1(self.norm1(x)) + x
+        x = self.attn2(self.norm2(x)) + x
+        x = self.ff(self.norm3(x)) + x
+        return x
+
+class SpatialTransformerNoContext(nn.Module):
+    """
+    Transformer block for image-like data.
+    First, project the input (aka embedding)
+    and reshape to b, t, d.
+    Then apply standard transformer action.
+    Finally, reshape to image
+    """
+    def __init__(self, in_channels, n_heads, d_head,
+                 depth=1, dropout=0.,):
+        super().__init__()
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = Normalize(in_channels)
+
+        self.proj_in = nn.Conv2d(in_channels,
+                                 inner_dim,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+
+        self.transformer_blocks = nn.ModuleList(
+            [BasicTransformerBlockNoContext(inner_dim, n_heads, d_head, dropout=dropout)
+                for d in range(depth)]
+        )
+
+        self.proj_out = zero_module(nn.Conv2d(inner_dim,
+                                              in_channels,
+                                              kernel_size=1,
+                                              stride=1,
+                                              padding=0))
+
+    def forward(self, x):
+        # note: if no context is given, cross-attention defaults to self-attention
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        x = self.proj_in(x)
+        x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
+        for block in self.transformer_blocks:
+            x = block(x)
+        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
+        x = self.proj_out(x)
+        return x + x_in
+
+
+#######################################
+# Spatial Transformer with Two Branch #
+#######################################
+
+class DualSpatialTransformer(nn.Module):
+    def __init__(self, in_channels, n_heads, d_head,
+                 depth=1, dropout=0., context_dim=None,
+                 disable_self_attn=False):
+        super().__init__()
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+
+        # First crossattn
+        self.norm_0 = Normalize(in_channels)
+        self.proj_in_0 = nn.Conv2d(
+            in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+        self.transformer_blocks_0 = nn.ModuleList(
+            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim,
+                                   disable_self_attn=disable_self_attn)
+                for d in range(depth)]
+        )
+        self.proj_out_0 = zero_module(nn.Conv2d(
+            inner_dim, in_channels, kernel_size=1, stride=1, padding=0))
+
+        # Second crossattn
+        self.norm_1 = Normalize(in_channels)
+        self.proj_in_1 = nn.Conv2d(
+            in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+        self.transformer_blocks_1 = nn.ModuleList(
+            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim,
+                                   disable_self_attn=disable_self_attn)
+                for d in range(depth)]
+        )
+        self.proj_out_1 = zero_module(nn.Conv2d(
+            inner_dim, in_channels, kernel_size=1, stride=1, padding=0))
+
+    def forward(self, x, context=None, which=None):
+        # note: if no context is given, cross-attention defaults to self-attention
+        b, c, h, w = x.shape
+        x_in = x
+        if which==0:
+            norm, proj_in, blocks, proj_out = \
+                self.norm_0, self.proj_in_0, self.transformer_blocks_0, self.proj_out_0
+        elif which==1:
+            norm, proj_in, blocks, proj_out = \
+                self.norm_1, self.proj_in_1, self.transformer_blocks_1, self.proj_out_1
+        else:
+            # assert False, 'DualSpatialTransformer forward with a invalid which branch!'
+            # import numpy.random as npr
+            # rwhich = 0 if npr.rand() < which else 1
+            # context = context[rwhich]
+            # if rwhich==0:
+            #     norm, proj_in, blocks, proj_out = \
+            #         self.norm_0, self.proj_in_0, self.transformer_blocks_0, self.proj_out_0
+            # elif rwhich==1:
+            #     norm, proj_in, blocks, proj_out = \
+            #         self.norm_1, self.proj_in_1, self.transformer_blocks_1, self.proj_out_1
+
+            # import numpy.random as npr
+            # rwhich = 0 if npr.rand() < 0.33 else 1
+            # if rwhich==0:
+            #     context = context[rwhich]
+            #     norm, proj_in, blocks, proj_out = \
+            #         self.norm_0, self.proj_in_0, self.transformer_blocks_0, self.proj_out_0
+            # else:
+
+            norm, proj_in, blocks, proj_out = \
+                self.norm_0, self.proj_in_0, self.transformer_blocks_0, self.proj_out_0
+            x0 = norm(x)
+            x0 = proj_in(x0)
+            x0 = rearrange(x0, 'b c h w -> b (h w) c').contiguous()
+            for block in blocks:
+                x0 = block(x0, context=context[0])
+            x0 = rearrange(x0, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
+            x0 = proj_out(x0)
+
+            norm, proj_in, blocks, proj_out = \
+                self.norm_1, self.proj_in_1, self.transformer_blocks_1, self.proj_out_1
+            x1 = norm(x)
+            x1 = proj_in(x1)
+            x1 = rearrange(x1, 'b c h w -> b (h w) c').contiguous()
+            for block in blocks:
+                x1 = block(x1, context=context[1])
+            x1 = rearrange(x1, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
+            x1 = proj_out(x1)
+            return x0*which + x1*(1-which) + x_in
+
+        x = norm(x)
+        x = proj_in(x)
+        x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
+        for block in blocks:
+            x = block(x, context=context)
+        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
+        x = proj_out(x)
+        return x + x_in

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/autoencoder.py

@@ -0,0 +1,428 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from contextlib import contextmanager
+from lib.model_zoo.common.get_model import get_model, register
+
+from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
+
+from .diffusion_modules import Encoder, Decoder
+from .distributions import DiagonalGaussianDistribution
+
+
+class VQModel(nn.Module):
+    def __init__(self,
+                 ddconfig,
+                 lossconfig,
+                 n_embed,
+                 embed_dim,
+                 ckpt_path=None,
+                 ignore_keys=[],
+                 image_key="image",
+                 colorize_nlabels=None,
+                 monitor=None,
+                 batch_resize_range=None,
+                 scheduler_config=None,
+                 lr_g_factor=1.0,
+                 remap=None,
+                 sane_index_shape=False, # tell vector quantizer to return indices as bhw
+                 use_ema=False
+                 ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.n_embed = n_embed
+        self.image_key = image_key
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        self.loss = instantiate_from_config(lossconfig)
+        self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25,
+                                        remap=remap,
+                                        sane_index_shape=sane_index_shape)
+        self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+        if colorize_nlabels is not None:
+            assert type(colorize_nlabels)==int
+            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
+        if monitor is not None:
+            self.monitor = monitor
+        self.batch_resize_range = batch_resize_range
+        if self.batch_resize_range is not None:
+            print(f"{self.__class__.__name__}: Using per-batch resizing in range {batch_resize_range}.")
+
+        self.use_ema = use_ema
+        if self.use_ema:
+            self.model_ema = LitEma(self)
+            print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+        self.scheduler_config = scheduler_config
+        self.lr_g_factor = lr_g_factor
+
+    @contextmanager
+    def ema_scope(self, context=None):
+        if self.use_ema:
+            self.model_ema.store(self.parameters())
+            self.model_ema.copy_to(self)
+            if context is not None:
+                print(f"{context}: Switched to EMA weights")
+        try:
+            yield None
+        finally:
+            if self.use_ema:
+                self.model_ema.restore(self.parameters())
+                if context is not None:
+                    print(f"{context}: Restored training weights")
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        sd = torch.load(path, map_location="cpu")["state_dict"]
+        keys = list(sd.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print("Deleting key {} from state_dict.".format(k))
+                    del sd[k]
+        missing, unexpected = self.load_state_dict(sd, strict=False)
+        print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
+        if len(missing) > 0:
+            print(f"Missing Keys: {missing}")
+            print(f"Unexpected Keys: {unexpected}")
+
+    def on_train_batch_end(self, *args, **kwargs):
+        if self.use_ema:
+            self.model_ema(self)
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        quant, emb_loss, info = self.quantize(h)
+        return quant, emb_loss, info
+
+    def encode_to_prequant(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        return h
+
+    def decode(self, quant):
+        quant = self.post_quant_conv(quant)
+        dec = self.decoder(quant)
+        return dec
+
+    def decode_code(self, code_b):
+        quant_b = self.quantize.embed_code(code_b)
+        dec = self.decode(quant_b)
+        return dec
+
+    def forward(self, input, return_pred_indices=False):
+        quant, diff, (_,_,ind) = self.encode(input)
+        dec = self.decode(quant)
+        if return_pred_indices:
+            return dec, diff, ind
+        return dec, diff
+
+    def get_input(self, batch, k):
+        x = batch[k]
+        if len(x.shape) == 3:
+            x = x[..., None]
+        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
+        if self.batch_resize_range is not None:
+            lower_size = self.batch_resize_range[0]
+            upper_size = self.batch_resize_range[1]
+            if self.global_step <= 4:
+                # do the first few batches with max size to avoid later oom
+                new_resize = upper_size
+            else:
+                new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16))
+            if new_resize != x.shape[2]:
+                x = F.interpolate(x, size=new_resize, mode="bicubic")
+            x = x.detach()
+        return x
+
+    def training_step(self, batch, batch_idx, optimizer_idx):
+        # https://github.com/pytorch/pytorch/issues/37142
+        # try not to fool the heuristics
+        x = self.get_input(batch, self.image_key)
+        xrec, qloss, ind = self(x, return_pred_indices=True)
+
+        if optimizer_idx == 0:
+            # autoencode
+            aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
+                                            last_layer=self.get_last_layer(), split="train",
+                                            predicted_indices=ind)
+
+            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
+            return aeloss
+
+        if optimizer_idx == 1:
+            # discriminator
+            discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
+                                            last_layer=self.get_last_layer(), split="train")
+            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
+            return discloss
+
+    def validation_step(self, batch, batch_idx):
+        log_dict = self._validation_step(batch, batch_idx)
+        with self.ema_scope():
+            log_dict_ema = self._validation_step(batch, batch_idx, suffix="_ema")
+        return log_dict
+
+    def _validation_step(self, batch, batch_idx, suffix=""):
+        x = self.get_input(batch, self.image_key)
+        xrec, qloss, ind = self(x, return_pred_indices=True)
+        aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0,
+                                        self.global_step,
+                                        last_layer=self.get_last_layer(),
+                                        split="val"+suffix,
+                                        predicted_indices=ind
+                                        )
+
+        discloss, log_dict_disc = self.loss(qloss, x, xrec, 1,
+                                            self.global_step,
+                                            last_layer=self.get_last_layer(),
+                                            split="val"+suffix,
+                                            predicted_indices=ind
+                                            )
+        rec_loss = log_dict_ae[f"val{suffix}/rec_loss"]
+        self.log(f"val{suffix}/rec_loss", rec_loss,
+                   prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
+        self.log(f"val{suffix}/aeloss", aeloss,
+                   prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
+        if version.parse(pl.__version__) >= version.parse('1.4.0'):
+            del log_dict_ae[f"val{suffix}/rec_loss"]
+        self.log_dict(log_dict_ae)
+        self.log_dict(log_dict_disc)
+        return self.log_dict
+
+    def configure_optimizers(self):
+        lr_d = self.learning_rate
+        lr_g = self.lr_g_factor*self.learning_rate
+        print("lr_d", lr_d)
+        print("lr_g", lr_g)
+        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
+                                  list(self.decoder.parameters())+
+                                  list(self.quantize.parameters())+
+                                  list(self.quant_conv.parameters())+
+                                  list(self.post_quant_conv.parameters()),
+                                  lr=lr_g, betas=(0.5, 0.9))
+        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
+                                    lr=lr_d, betas=(0.5, 0.9))
+
+        if self.scheduler_config is not None:
+            scheduler = instantiate_from_config(self.scheduler_config)
+
+            print("Setting up LambdaLR scheduler...")
+            scheduler = [
+                {
+                    'scheduler': LambdaLR(opt_ae, lr_lambda=scheduler.schedule),
+                    'interval': 'step',
+                    'frequency': 1
+                },
+                {
+                    'scheduler': LambdaLR(opt_disc, lr_lambda=scheduler.schedule),
+                    'interval': 'step',
+                    'frequency': 1
+                },
+            ]
+            return [opt_ae, opt_disc], scheduler
+        return [opt_ae, opt_disc], []
+
+    def get_last_layer(self):
+        return self.decoder.conv_out.weight
+
+    def log_images(self, batch, only_inputs=False, plot_ema=False, **kwargs):
+        log = dict()
+        x = self.get_input(batch, self.image_key)
+        x = x.to(self.device)
+        if only_inputs:
+            log["inputs"] = x
+            return log
+        xrec, _ = self(x)
+        if x.shape[1] > 3:
+            # colorize with random projection
+            assert xrec.shape[1] > 3
+            x = self.to_rgb(x)
+            xrec = self.to_rgb(xrec)
+        log["inputs"] = x
+        log["reconstructions"] = xrec
+        if plot_ema:
+            with self.ema_scope():
+                xrec_ema, _ = self(x)
+                if x.shape[1] > 3: xrec_ema = self.to_rgb(xrec_ema)
+                log["reconstructions_ema"] = xrec_ema
+        return log
+
+    def to_rgb(self, x):
+        assert self.image_key == "segmentation"
+        if not hasattr(self, "colorize"):
+            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
+        x = F.conv2d(x, weight=self.colorize)
+        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
+        return x
+
+
+class VQModelInterface(VQModel):
+    def __init__(self, embed_dim, *args, **kwargs):
+        super().__init__(embed_dim=embed_dim, *args, **kwargs)
+        self.embed_dim = embed_dim
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        return h
+
+    def decode(self, h, force_not_quantize=False):
+        # also go through quantization layer
+        if not force_not_quantize:
+            quant, emb_loss, info = self.quantize(h)
+        else:
+            quant = h
+        quant = self.post_quant_conv(quant)
+        dec = self.decoder(quant)
+        return dec
+
+
+@register('autoencoderkl')
+class AutoencoderKL(nn.Module):
+    def __init__(self,
+                 ddconfig,
+                 lossconfig,
+                 embed_dim,
+                 ckpt_path=None,
+                 ignore_keys=[],
+                 image_key="image",
+                 colorize_nlabels=None,
+                 monitor=None,):
+        super().__init__()
+        self.image_key = image_key
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        assert ddconfig["double_z"]
+        self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+        self.embed_dim = embed_dim
+        if colorize_nlabels is not None:
+            assert type(colorize_nlabels)==int
+            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
+        if monitor is not None:
+            self.monitor = monitor
+
+    def encode(self, x):
+        h = self.encoder(x)
+        moments = self.quant_conv(h)
+        posterior = DiagonalGaussianDistribution(moments)
+        return posterior
+
+    def decode(self, z):
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+        return dec
+
+    def forward(self, input, sample_posterior=True):
+        posterior = self.encode(input)
+        if sample_posterior:
+            z = posterior.sample()
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        return dec, posterior
+
+    def get_input(self, batch, k):
+        x = batch[k]
+        if len(x.shape) == 3:
+            x = x[..., None]
+        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
+        return x
+
+    def training_step(self, batch, batch_idx, optimizer_idx):
+        inputs = self.get_input(batch, self.image_key)
+        reconstructions, posterior = self(inputs)
+
+        if optimizer_idx == 0:
+            # train encoder+decoder+logvar
+            aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
+                                            last_layer=self.get_last_layer(), split="train")
+            self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
+            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False)
+            return aeloss
+
+        if optimizer_idx == 1:
+            # train the discriminator
+            discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
+                                                last_layer=self.get_last_layer(), split="train")
+
+            self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
+            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False)
+            return discloss
+
+    def validation_step(self, batch, batch_idx):
+        inputs = self.get_input(batch, self.image_key)
+        reconstructions, posterior = self(inputs)
+        aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
+                                        last_layer=self.get_last_layer(), split="val")
+
+        discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
+                                            last_layer=self.get_last_layer(), split="val")
+
+        self.log("val/rec_loss", log_dict_ae["val/rec_loss"])
+        self.log_dict(log_dict_ae)
+        self.log_dict(log_dict_disc)
+        return self.log_dict
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
+                                  list(self.decoder.parameters())+
+                                  list(self.quant_conv.parameters())+
+                                  list(self.post_quant_conv.parameters()),
+                                  lr=lr, betas=(0.5, 0.9))
+        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
+                                    lr=lr, betas=(0.5, 0.9))
+        return [opt_ae, opt_disc], []
+
+    def get_last_layer(self):
+        return self.decoder.conv_out.weight
+
+    @torch.no_grad()
+    def log_images(self, batch, only_inputs=False, **kwargs):
+        log = dict()
+        x = self.get_input(batch, self.image_key)
+        x = x.to(self.device)
+        if not only_inputs:
+            xrec, posterior = self(x)
+            if x.shape[1] > 3:
+                # colorize with random projection
+                assert xrec.shape[1] > 3
+                x = self.to_rgb(x)
+                xrec = self.to_rgb(xrec)
+            log["samples"] = self.decode(torch.randn_like(posterior.sample()))
+            log["reconstructions"] = xrec
+        log["inputs"] = x
+        return log
+
+    def to_rgb(self, x):
+        assert self.image_key == "segmentation"
+        if not hasattr(self, "colorize"):
+            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
+        x = F.conv2d(x, weight=self.colorize)
+        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
+        return x
+
+
+class IdentityFirstStage(nn.Module):
+    def __init__(self, *args, vq_interface=False, **kwargs):
+        self.vq_interface = vq_interface  # TODO: Should be true by default but check to not break older stuff
+        super().__init__()
+
+    def encode(self, x, *args, **kwargs):
+        return x
+
+    def decode(self, x, *args, **kwargs):
+        return x
+
+    def quantize(self, x, *args, **kwargs):
+        if self.vq_interface:
+            return x, None, [None, None, None]
+        return x
+
+    def forward(self, x, *args, **kwargs):
+        return x

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/bert.py

@@ -0,0 +1,142 @@
+import torch
+import torch.nn as nn
+from functools import partial
+
+# from ldm.modules.x_transformer import Encoder, TransformerWrapper  # TODO: can we directly rely on lucidrains code and simply add this as a reuirement? --> test
+
+
+class AbstractEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+
+class ClassEmbedder(nn.Module):
+    def __init__(self, embed_dim, n_classes=1000, key='class'):
+        super().__init__()
+        self.key = key
+        self.embedding = nn.Embedding(n_classes, embed_dim)
+
+    def forward(self, batch, key=None):
+        if key is None:
+            key = self.key
+        # this is for use in crossattn
+        c = batch[key][:, None]
+        c = self.embedding(c)
+        return c
+
+
+class TransformerEmbedder(AbstractEncoder):
+    """Some transformer encoder layers"""
+    def __init__(self, n_embed, n_layer, vocab_size, max_seq_len=77):
+        super().__init__()
+        self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len,
+                                              attn_layers=Encoder(dim=n_embed, depth=n_layer))
+
+    def forward(self, tokens):
+        z = self.transformer(tokens, return_embeddings=True)
+        return z
+
+    def encode(self, x):
+        return self(x)
+
+
+class BERTTokenizer(AbstractEncoder):
+    """ Uses a pretrained BERT tokenizer by huggingface. Vocab size: 30522 (?)"""
+    def __init__(self, device="cuda", vq_interface=True, max_length=77):
+        super().__init__()
+        from transformers import BertTokenizerFast  # TODO: add to reuquirements
+        self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
+        self.vq_interface = vq_interface
+        self.max_length = max_length
+
+    def forward(self, text):
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"]
+        return tokens
+
+    @torch.no_grad()
+    def encode(self, text):
+        tokens = self(text)
+        if not self.vq_interface:
+            return tokens
+        return None, None, [None, None, tokens]
+
+    def decode(self, text):
+        return text
+
+
+class BERTEmbedder(AbstractEncoder):
+    """Uses the BERT tokenizr model and add some transformer encoder layers"""
+    def __init__(self, n_embed, n_layer, vocab_size=30522, max_seq_len=77,
+                 ckpt_path=None, ignore_keys=[], device="cuda", use_tokenizer=True, embedding_dropout=0.0):
+        super().__init__()
+        self.use_tknz_fn = use_tokenizer
+        if self.use_tknz_fn:
+            self.tknz_fn = BERTTokenizer(vq_interface=False, max_length=max_seq_len)
+        self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len,
+                                              attn_layers=Encoder(dim=n_embed, depth=n_layer),
+                                              emb_dropout=embedding_dropout)
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)                                              
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        sd = torch.load(path, map_location="cpu")
+        keys = list(sd.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print("Deleting key {} from state_dict.".format(k))
+                    del sd[k]
+        missing, unexpected = self.load_state_dict(sd, strict=False)
+        print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")  
+
+    def forward(self, text):
+        if self.use_tknz_fn:
+            tokens = self.tknz_fn(text)
+        else:
+            tokens = text
+        device = self.transformer.token_emb.weight.device # a trick to get device
+        tokens = tokens.to(device)
+        z = self.transformer(tokens, return_embeddings=True)
+        return z
+
+    def encode(self, text):
+        # output of length 77
+        return self(text)
+
+
+class SpatialRescaler(nn.Module):
+    def __init__(self,
+                 n_stages=1,
+                 method='bilinear',
+                 multiplier=0.5,
+                 in_channels=3,
+                 out_channels=None,
+                 bias=False):
+        super().__init__()
+        self.n_stages = n_stages
+        assert self.n_stages >= 0
+        assert method in ['nearest','linear','bilinear','trilinear','bicubic','area']
+        self.multiplier = multiplier
+        self.interpolator = partial(torch.nn.functional.interpolate, mode=method)
+        self.remap_output = out_channels is not None
+        if self.remap_output:
+            print(f'Spatial Rescaler mapping from {in_channels} to {out_channels} channels after resizing.')
+            self.channel_mapper = nn.Conv2d(in_channels,out_channels,1,bias=bias)
+
+    def forward(self,x):
+        for stage in range(self.n_stages):
+            x = self.interpolator(x, scale_factor=self.multiplier)
+
+
+        if self.remap_output:
+            x = self.channel_mapper(x)
+        return x
+
+    def encode(self, x):
+        return self(x)

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/clip.py

@@ -0,0 +1,226 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from functools import partial
+from lib.model_zoo.common.get_model import register
+
+version = '0'
+symbol = 'clip'
+
+class AbstractEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError
+
+from transformers import CLIPTokenizer, CLIPTextModel
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+@register('clip_text_frozen', version)
+class FrozenCLIPTextEmbedder(AbstractEncoder):
+    """Uses the CLIP transformer encoder for text (from huggingface)"""
+    def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77):  # clip-vit-base-patch32
+        super().__init__()
+        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.transformer = CLIPTextModel.from_pretrained(version)
+        self.device = device
+        self.max_length = max_length   # TODO: typical value?
+        self.freeze()
+
+    def freeze(self):
+        self.transformer = self.transformer.eval()
+        #self.train = disabled_train
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"].to(self.device)
+        outputs = self.transformer(input_ids=tokens)
+        z = outputs.last_hidden_state
+        return z
+
+    def encode(self, text):
+        return self(text)
+
+from transformers import CLIPProcessor, CLIPVisionModel
+
+@register('clip_vision_frozen', version)
+class FrozenCLIPVisionEmbedder(AbstractEncoder):
+    def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77):  # clip-vit-base-patch32
+        super().__init__()
+        self.processor = CLIPProcessor.from_pretrained(version)
+        self.transformer = CLIPVisionModel.from_pretrained(version)
+        self.device = device
+        self.max_length = max_length   # TODO: typical value?
+        self.freeze()
+
+    def freeze(self):
+        self.transformer = self.transformer.eval()
+        #self.train = disabled_train
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, images):
+        inputs = self.processor(images=images, return_tensors="pt")
+        pixels = inputs['pixel_values'].to(self.device)
+        outputs = self.transformer(pixel_values=pixels)
+        z = outputs.last_hidden_state
+        return z
+
+    def encode(self, image):
+        return self(image)
+
+from transformers import CLIPModel
+
+@register('clip_frozen', version)
+class FrozenCLIP(AbstractEncoder):
+    def __init__(self, 
+                 version="openai/clip-vit-large-patch14", 
+                 max_length=77, 
+                 encode_type='encode_text',):  # clip-vit-base-patch32
+        super().__init__()
+        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.processor = CLIPProcessor.from_pretrained(version)
+        self.model = CLIPModel.from_pretrained(version)
+        self.max_length = max_length  # TODO: typical value?
+        self.encode_type = encode_type
+        self.pinv_text_projection = None
+        self.freeze()
+
+    def get_device(self):
+        # A trick to get device
+        return self.model.text_projection.weight.device
+
+    def freeze(self):
+        self.model = self.model.eval()
+        #self.train = disabled_train
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def encode_text_pooled(self, text):
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"].to(self.get_device())
+        return self.model.get_text_features(input_ids=tokens)
+
+    def encode_vision_pooled(self, images):
+        inputs = self.processor(images=images, return_tensors="pt")
+        pixels = inputs['pixel_values'].to(self.get_device())
+        return self.model.get_image_features(pixel_values=pixels)
+
+    def encode_text_noproj(self, text):
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"].to(self.get_device())
+        outputs = self.model.text_model(input_ids=tokens)
+        return outputs.last_hidden_state
+        
+    def encode_vision_noproj(self, images):
+        inputs = self.processor(images=images, return_tensors="pt")
+        pixels = inputs['pixel_values'].to(self.get_device())
+        outputs = self.model.vision_model(pixel_values=pixels)
+        return outputs.last_hidden_state
+
+    def encode_text_bug(self, text):
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"].to(self.get_device())
+        outputs = self.model.text_model(input_ids=tokens)
+        z = outputs.last_hidden_state
+        z_pooled = outputs.pooler_output
+        z = z / torch.norm(z_pooled.unsqueeze(1), dim=-1, keepdim=True)
+        return self.model.text_projection(z)
+
+    def encode_text(self, text):
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"].to(self.get_device())
+        outputs = self.model.text_model(input_ids=tokens)
+        z = self.model.text_projection(outputs.last_hidden_state)
+        z_pooled = self.model.text_projection(outputs.pooler_output)
+        z = z / torch.norm(z_pooled.unsqueeze(1), dim=-1, keepdim=True)
+        return z
+
+    def encode_vision(self, images):
+        z = self.encode_vision_noproj(images)
+        z = self.model.vision_model.post_layernorm(z)
+        z = self.model.visual_projection(z)
+        z_pooled = z[:, 0:1]
+        # z_pooled_normed = z_pooled / z_pooled.norm(dim=-1, keepdim=True)
+        z = z / torch.norm(z_pooled, dim=-1, keepdim=True)
+        return z
+
+    def encode_vision_pinvtext(self, images):
+        blank_text_encode_norm_avg = 28.9096
+        z = self.encode_vision(images)
+        if self.pinv_text_projection is None:
+            self.pinv_text_projection = torch.linalg.pinv(self.model.text_projection.weight).T
+        z = torch.matmul(z, self.pinv_text_projection)
+        # z = z / torch.norm(z[:, 0:1], dim=-1, keepdim=True)
+        z = z / torch.norm(z, dim=-1, keepdim=True)
+        z = z*blank_text_encode_norm_avg
+        # return z[:, 1:2].repeat(1, 77, 1)
+        z2 = self.encode_text_noproj('')
+        # z2[:, 1:77] = z[:, 0:76]
+        return torch.flip(z, dims=(1,))[:, 0:77]
+
+    def encode(self, *args, **kwargs):
+        return getattr(self, self.encode_type)(*args, **kwargs)
+
+#############################
+# copyed from justin's code #
+#############################
+
+@register('clip_vision_frozen_justin', version)
+class FrozenCLIPVisionEmbedder_Justin(AbstractEncoder):
+    """
+        Uses the CLIP image encoder.
+        """
+    def __init__(
+            self,
+            model='ViT-L/14',
+            jit=False,
+            device='cuda' if torch.cuda.is_available() else 'cpu',
+            antialias=False,
+        ):
+        super().__init__()
+        from . import clip_justin
+        self.model, _ = clip_justin.load(name=model, device=device, jit=jit)
+        self.device = device
+        self.antialias = antialias
+
+        self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
+        self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
+
+        # I didn't call this originally, but seems like it was frozen anyway
+        self.freeze()
+
+    def freeze(self):
+        self.transformer = self.model.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def preprocess(self, x):
+        import kornia
+        # Expects inputs in the range -1, 1
+        x = kornia.geometry.resize(x, (224, 224),
+                                   interpolation='bicubic',align_corners=True,
+                                   antialias=self.antialias)
+        x = (x + 1.) / 2.
+        # renormalize according to clip
+        x = kornia.enhance.normalize(x, self.mean, self.std)
+        return x
+
+    def forward(self, x):
+        # x is assumed to be in range [-1,1]
+        return self.model.encode_image(self.preprocess(x)).float()
+
+    def encode(self, im):
+        return self(im).unsqueeze(1)

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/clip_justin/__init__.py

@@ -0,0 +1 @@
+from .clip import load

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/clip_justin/clip.py

@@ -0,0 +1,237 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Any, Union, List
+from pkg_resources import packaging
+
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+from tqdm import tqdm
+
+from .model import build_model
+# from .simple_tokenizer import SimpleTokenizer as _Tokenizer
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
+    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__all__ = ["available_models", "load", "tokenize"]
+# _tokenizer = _Tokenizer()
+
+_MODELS = {
+    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
+    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
+    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
+    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
+    "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",
+    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
+    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
+    "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
+    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
+}
+
+
+def _download(url: str, root: str):
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+
+    expected_sha256 = url.split("/")[-2]
+    download_target = os.path.join(root, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
+            return download_target
+        else:
+            warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
+        raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def _transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+def available_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list(_MODELS.keys())
+
+
+def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None):
+    """Load a CLIP model
+
+    Parameters
+    ----------
+    name : str
+        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+
+    device : Union[str, torch.device]
+        The device to put the loaded model
+
+    jit : bool
+        Whether to load the optimized JIT model or more hackable non-JIT model (default).
+
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/clip"
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The CLIP model
+
+    preprocess : Callable[[PIL.Image], torch.Tensor]
+        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+    """
+    if name in _MODELS:
+        model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+    with open(model_path, 'rb') as opened_file:
+        try:
+            # loading JIT archive
+            model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
+            state_dict = None
+        except RuntimeError:
+            # loading saved state dict
+            if jit:
+                warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+                jit = False
+            state_dict = torch.load(opened_file, map_location="cpu")
+
+    if not jit:
+        model = build_model(state_dict or model.state_dict()).to(device)
+        if str(device) == "cpu":
+            model.float()
+        return model, _transform(model.visual.input_resolution)
+
+    # patch the device names
+    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+    def patch_device(module):
+        try:
+            graphs = [module.graph] if hasattr(module, "graph") else []
+        except RuntimeError:
+            graphs = []
+
+        if hasattr(module, "forward1"):
+            graphs.append(module.forward1.graph)
+
+        for graph in graphs:
+            for node in graph.findAllNodes("prim::Constant"):
+                if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
+                    node.copyAttributes(device_node)
+
+    model.apply(patch_device)
+    patch_device(model.encode_image)
+    patch_device(model.encode_text)
+
+    # patch dtype to float32 on CPU
+    if str(device) == "cpu":
+        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+        float_node = float_input.node()
+
+        def patch_float(module):
+            try:
+                graphs = [module.graph] if hasattr(module, "graph") else []
+            except RuntimeError:
+                graphs = []
+
+            if hasattr(module, "forward1"):
+                graphs.append(module.forward1.graph)
+
+            for graph in graphs:
+                for node in graph.findAllNodes("aten::to"):
+                    inputs = list(node.inputs())
+                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+                        if inputs[i].node()["value"] == 5:
+                            inputs[i].node().copyAttributes(float_node)
+
+        model.apply(patch_float)
+        patch_float(model.encode_image)
+        patch_float(model.encode_text)
+
+        model.float()
+
+    return model, _transform(model.input_resolution.item())
+
+
+# def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:
+#     """
+#     Returns the tokenized representation of given input string(s)
+
+#     Parameters
+#     ----------
+#     texts : Union[str, List[str]]
+#         An input string or a list of input strings to tokenize
+
+#     context_length : int
+#         The context length to use; all CLIP models use 77 as the context length
+
+#     truncate: bool
+#         Whether to truncate the text in case its encoding is longer than the context length
+
+#     Returns
+#     -------
+#     A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
+#     We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
+#     """
+#     if isinstance(texts, str):
+#         texts = [texts]
+
+#     sot_token = _tokenizer.encoder["<|startoftext|>"]
+#     eot_token = _tokenizer.encoder["<|endoftext|>"]
+#     all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+#     if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
+#         result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+#     else:
+#         result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
+
+#     for i, tokens in enumerate(all_tokens):
+#         if len(tokens) > context_length:
+#             if truncate:
+#                 tokens = tokens[:context_length]
+#                 tokens[-1] = eot_token
+#             else:
+#                 raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
+#         result[i, :len(tokens)] = torch.tensor(tokens)
+
+#     return result

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/clip_justin/model.py

@@ -0,0 +1,436 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x[:1], key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+        return x.squeeze(0)
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.relu3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            x = self.relu1(self.bn1(self.conv1(x)))
+            x = self.relu2(self.bn2(self.conv2(x)))
+            x = self.relu3(self.bn3(self.conv3(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.ln_post(x[:, 0, :])
+
+        if self.proj is not None:
+            x = x @ self.proj
+
+        return x
+
+
+class CLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int
+                 ):
+        super().__init__()
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisionTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim
+            )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features ** -0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")))
+
+    model = CLIP(
+        embed_dim,
+        image_resolution, vision_layers, vision_width, vision_patch_size,
+        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+    convert_weights(model)
+    model.load_state_dict(state_dict)
+    return model.eval()

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/clip_justin/simple_tokenizer.py

@@ -0,0 +1,132 @@
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str = default_bpe()):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/common/get_model.py

@@ -0,0 +1,120 @@
+from email.policy import strict
+import torch
+import torchvision.models
+import os.path as osp
+import copy
+from ...log_service import print_log 
+from .utils import \
+    get_total_param, get_total_param_sum, \
+    get_unit
+
+# def load_state_dict(net, model_path):
+#     if isinstance(net, dict):
+#         for ni, neti in net.items():
+#             paras = torch.load(model_path[ni], map_location=torch.device('cpu'))
+#             new_paras = neti.state_dict()
+#             new_paras.update(paras)
+#             neti.load_state_dict(new_paras)
+#     else:
+#         paras = torch.load(model_path, map_location=torch.device('cpu'))
+#         new_paras = net.state_dict()
+#         new_paras.update(paras)
+#         net.load_state_dict(new_paras)
+#     return
+
+# def save_state_dict(net, path):
+#     if isinstance(net, (torch.nn.DataParallel,
+#                         torch.nn.parallel.DistributedDataParallel)):
+#         torch.save(net.module.state_dict(), path)
+#     else:
+#         torch.save(net.state_dict(), path)
+
+def singleton(class_):
+    instances = {}
+    def getinstance(*args, **kwargs):
+        if class_ not in instances:
+            instances[class_] = class_(*args, **kwargs)
+        return instances[class_]
+    return getinstance
+
+def preprocess_model_args(args):
+    # If args has layer_units, get the corresponding
+    #     units.
+    # If args get backbone, get the backbone model.
+    args = copy.deepcopy(args)
+    if 'layer_units' in args:
+        layer_units = [
+            get_unit()(i) for i in args.layer_units
+        ]
+        args.layer_units = layer_units
+    if 'backbone' in args:
+        args.backbone = get_model()(args.backbone)
+    return args
+
+@singleton
+class get_model(object):
+    def __init__(self):
+        self.model = {}
+        self.version = {}
+
+    def register(self, model, name, version='x'):
+        self.model[name] = model
+        self.version[name] = version
+
+    def __call__(self, cfg, verbose=True):
+        """
+        Construct model based on the config. 
+        """
+        t = cfg.type
+
+        # the register is in each file
+        if t.find('ldm')==0:
+            from .. import ldm
+        elif t=='autoencoderkl':
+            from .. import autoencoder
+        elif t.find('clip')==0:
+            from .. import clip
+        elif t.find('sd')==0:
+            from .. import sd
+        elif t.find('vd')==0:
+            from .. import vd
+        elif t.find('openai_unet')==0:
+            from .. import openaimodel
+        elif t.find('optimus')==0:
+            from .. import optimus
+
+        args = preprocess_model_args(cfg.args)
+        net = self.model[t](**args)
+
+        if 'ckpt' in cfg:
+            checkpoint = torch.load(cfg.ckpt, map_location='cpu')
+            strict_sd = cfg.get('strict_sd', True)
+            net.load_state_dict(checkpoint['state_dict'], strict=strict_sd)
+            if verbose:
+                print_log('Load ckpt from {}'.format(cfg.ckpt))
+        elif 'pth' in cfg:
+            sd = torch.load(cfg.pth, map_location='cpu')
+            strict_sd = cfg.get('strict_sd', True)
+            net.load_state_dict(sd, strict=strict_sd)
+            if verbose:
+                print_log('Load pth from {}'.format(cfg.pth))
+
+        # display param_num & param_sum
+        if verbose:
+            print_log(
+                'Load {} with total {} parameters,' 
+                '{:.3f} parameter sum.'.format(
+                    t, 
+                    get_total_param(net), 
+                    get_total_param_sum(net) ))
+
+        return net
+
+    def get_version(self, name):
+        return self.version[name]
+
+def register(name, version='x'):
+    def wrapper(class_):
+        get_model().register(class_, name, version)
+        return class_
+    return wrapper

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/common/get_optimizer.py

@@ -0,0 +1,47 @@
+import torch
+import torch.optim as optim
+import numpy as np
+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
+
+class get_optimizer(object):
+    def __init__(self):
+        self.optimizer = {}
+        self.register(optim.SGD, 'sgd')
+        self.register(optim.Adam, 'adam')
+        self.register(optim.AdamW, 'adamw')
+
+    def register(self, optim, name):
+        self.optimizer[name] = optim
+
+    def __call__(self, net, cfg):
+        if cfg is None:
+            return None
+        t = cfg.type
+        if isinstance(net, (torch.nn.DataParallel,
+                            torch.nn.parallel.DistributedDataParallel)):
+            netm = net.module
+        else:
+            netm = net
+        pg = getattr(netm, 'parameter_group', None)
+
+        if pg is not None:
+            params = []
+            for group_name, module_or_para in pg.items():
+                if not isinstance(module_or_para, list):
+                    module_or_para = [module_or_para]
+
+                grouped_params = [mi.parameters() if isinstance(mi, torch.nn.Module) else [mi] for mi in module_or_para]
+                grouped_params = itertools.chain(*grouped_params)
+                pg_dict = {'params':grouped_params, 'name':group_name}
+                params.append(pg_dict)
+        else:
+            params = net.parameters()
+        return self.optimizer[t](params, lr=0, **cfg.args)

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/common/get_scheduler.py

@@ -0,0 +1,262 @@
+import torch
+import torch.optim as optim
+import numpy as np
+import copy
+from ... import sync
+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_scheduler(object):
+    def __init__(self):
+        self.lr_scheduler = {}
+
+    def register(self, lrsf, name):
+        self.lr_scheduler[name] = lrsf
+
+    def __call__(self, cfg):
+        if cfg is None:
+            return None
+        if isinstance(cfg, list):
+            schedulers = []
+            for ci in cfg:
+                t = ci.type
+                schedulers.append(
+                    self.lr_scheduler[t](**ci.args))
+            if len(schedulers) == 0:
+                raise ValueError
+            else:
+                return compose_scheduler(schedulers)
+        t = cfg.type
+        return self.lr_scheduler[t](**cfg.args)
+        
+
+def register(name):
+    def wrapper(class_):
+        get_scheduler().register(class_, name)
+        return class_
+    return wrapper
+
+class template_scheduler(object):
+    def __init__(self, step):
+        self.step = step
+
+    def __getitem__(self, idx):
+        raise ValueError
+
+    def set_lr(self, optim, new_lr, pg_lrscale=None):
+        """
+        Set Each parameter_groups in optim with new_lr
+        New_lr can be find according to the idx.
+        pg_lrscale tells how to scale each pg.
+        """
+        # new_lr = self.__getitem__(idx)
+        pg_lrscale = copy.deepcopy(pg_lrscale)
+        for pg in optim.param_groups:
+            if pg_lrscale is None:
+                pg['lr'] = new_lr
+            else:
+                pg['lr'] = new_lr * pg_lrscale.pop(pg['name'])
+        assert (pg_lrscale is None) or (len(pg_lrscale)==0), \
+            "pg_lrscale doesn't match pg"
+
+@register('constant')
+class constant_scheduler(template_scheduler):
+    def __init__(self, lr, step):
+        super().__init__(step)
+        self.lr = lr
+
+    def __getitem__(self, idx):
+        if idx >= self.step:
+            raise ValueError
+        return self.lr
+
+@register('poly')
+class poly_scheduler(template_scheduler):
+    def __init__(self, start_lr, end_lr, power, step):
+        super().__init__(step)
+        self.start_lr = start_lr
+        self.end_lr = end_lr
+        self.power = power
+
+    def __getitem__(self, idx):
+        if idx >= self.step:
+            raise ValueError
+        a, b = self.start_lr, self.end_lr
+        p, n = self.power, self.step
+        return b + (a-b)*((1-idx/n)**p)
+
+@register('linear')
+class linear_scheduler(template_scheduler):
+    def __init__(self, start_lr, end_lr, step):
+        super().__init__(step)
+        self.start_lr = start_lr
+        self.end_lr = end_lr
+
+    def __getitem__(self, idx):
+        if idx >= self.step:
+            raise ValueError
+        a, b, n = self.start_lr, self.end_lr, self.step
+        return b + (a-b)*(1-idx/n)
+
+@register('multistage')
+class constant_scheduler(template_scheduler):
+    def __init__(self, start_lr, milestones, gamma, step):
+        super().__init__(step)
+        self.start_lr = start_lr
+        m = [0] + milestones + [step]
+        lr_iter = start_lr
+        self.lr = []
+        for ms, me in zip(m[0:-1], m[1:]):
+            for _ in range(ms, me):
+                self.lr.append(lr_iter)
+            lr_iter *= gamma
+
+    def __getitem__(self, idx):
+        if idx >= self.step:
+            raise ValueError
+        return self.lr[idx]
+
+class compose_scheduler(template_scheduler):
+    def __init__(self, schedulers):
+        self.schedulers = schedulers
+        self.step = [si.step for si in schedulers]
+        self.step_milestone = []
+        acc = 0
+        for i in self.step:
+            acc += i
+            self.step_milestone.append(acc)
+        self.step = sum(self.step)
+
+    def __getitem__(self, idx):
+        if idx >= self.step:
+            raise ValueError
+        ms = self.step_milestone
+        for idx, (mi, mj) in enumerate(zip(ms[:-1], ms[1:])):
+            if mi <= idx < mj:
+                return self.schedulers[idx-mi]
+        raise ValueError
+
+####################
+# lambda schedular #
+####################
+
+class LambdaWarmUpCosineScheduler(template_scheduler):
+    """
+    note: use with a base_lr of 1.0
+    """
+    def __init__(self, 
+                 base_lr,
+                 warm_up_steps, 
+                 lr_min, lr_max, lr_start, max_decay_steps, verbosity_interval=0):
+        cfgt = cfguh().cfg.train
+        bs = cfgt.batch_size
+        if 'gradacc_every' not in cfgt:
+            print('Warning, gradacc_every is not found in xml, use 1 as default.')
+        acc = cfgt.get('gradacc_every', 1)
+        self.lr_multi = base_lr * bs * acc
+        self.lr_warm_up_steps = warm_up_steps
+        self.lr_start = lr_start
+        self.lr_min = lr_min
+        self.lr_max = lr_max
+        self.lr_max_decay_steps = max_decay_steps
+        self.last_lr = 0.
+        self.verbosity_interval = verbosity_interval
+
+    def schedule(self, n):
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: 
+                print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
+        if n < self.lr_warm_up_steps:
+            lr = (self.lr_max - self.lr_start) / self.lr_warm_up_steps * n + self.lr_start
+            self.last_lr = lr
+            return lr
+        else:
+            t = (n - self.lr_warm_up_steps) / (self.lr_max_decay_steps - self.lr_warm_up_steps)
+            t = min(t, 1.0)
+            lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
+                    1 + np.cos(t * np.pi))
+            self.last_lr = lr
+            return lr
+
+    def __getitem__(self, idx):
+        return self.schedule(idx) * self.lr_multi
+
+class LambdaWarmUpCosineScheduler2(template_scheduler):
+    """
+    supports repeated iterations, configurable via lists
+    note: use with a base_lr of 1.0.
+    """
+    def __init__(self, 
+                 base_lr,
+                 warm_up_steps, 
+                 f_min, f_max, f_start, cycle_lengths, verbosity_interval=0):
+        cfgt = cfguh().cfg.train
+        # bs = cfgt.batch_size
+        # if 'gradacc_every' not in cfgt:
+        #     print('Warning, gradacc_every is not found in xml, use 1 as default.')
+        # acc = cfgt.get('gradacc_every', 1)
+        # self.lr_multi = base_lr * bs * acc
+        self.lr_multi = base_lr
+        assert len(warm_up_steps) == len(f_min) == len(f_max) == len(f_start) == len(cycle_lengths)
+        self.lr_warm_up_steps = warm_up_steps
+        self.f_start = f_start
+        self.f_min = f_min
+        self.f_max = f_max
+        self.cycle_lengths = cycle_lengths
+        self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
+        self.last_f = 0.
+        self.verbosity_interval = verbosity_interval
+
+    def find_in_interval(self, n):
+        interval = 0
+        for cl in self.cum_cycles[1:]:
+            if n <= cl:
+                return interval
+            interval += 1
+
+    def schedule(self, n):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                                                       f"current cycle {cycle}")
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            t = (n - self.lr_warm_up_steps[cycle]) / (self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle])
+            t = min(t, 1.0)
+            f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (
+                    1 + np.cos(t * np.pi))
+            self.last_f = f
+            return f
+
+    def __getitem__(self, idx):
+        return self.schedule(idx) * self.lr_multi
+
+@register('stable_diffusion_linear')
+class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
+    def schedule(self, n):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: 
+                print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                      f"current cycle {cycle}")
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (self.cycle_lengths[cycle] - n) / (self.cycle_lengths[cycle])
+            self.last_f = f
+            return f

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/common/utils.py

@@ -0,0 +1,292 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import copy
+import functools
+import itertools
+
+import matplotlib.pyplot as plt
+
+########
+# unit #
+########
+
+def singleton(class_):
+    instances = {}
+    def getinstance(*args, **kwargs):
+        if class_ not in instances:
+            instances[class_] = class_(*args, **kwargs)
+        return instances[class_]
+    return getinstance
+
+def str2value(v):
+    v = v.strip()
+    try:
+        return int(v)
+    except:
+        pass
+    try:
+        return float(v)
+    except:
+        pass
+    if v in ('True', 'true'):
+        return True
+    elif v in ('False', 'false'):
+        return False
+    else:
+        return v
+
+@singleton
+class get_unit(object):
+    def __init__(self):
+        self.unit = {}
+        self.register('none', None)
+
+        # general convolution
+        self.register('conv'  , nn.Conv2d)
+        self.register('bn'    , nn.BatchNorm2d)
+        self.register('relu'  , nn.ReLU)
+        self.register('relu6' , nn.ReLU6)
+        self.register('lrelu' , nn.LeakyReLU)
+        self.register('dropout'  , nn.Dropout)
+        self.register('dropout2d', nn.Dropout2d)
+        self.register('sine',     Sine)
+        self.register('relusine', ReLUSine)
+
+    def register(self, 
+                 name, 
+                 unitf,):
+
+        self.unit[name] = unitf
+
+    def __call__(self, name):
+        if name is None:
+            return None
+        i = name.find('(')
+        i = len(name) if i==-1 else i
+        t = name[:i]
+        f = self.unit[t]
+        args = name[i:].strip('()')
+        if len(args) == 0:
+            args = {}
+            return f
+        else:
+            args = args.split('=')
+            args = [[','.join(i.split(',')[:-1]), i.split(',')[-1]] for i in args]
+            args = list(itertools.chain.from_iterable(args))
+            args = [i.strip() for i in args if len(i)>0]
+            kwargs = {}
+            for k, v in zip(args[::2], args[1::2]):
+                if v[0]=='(' and v[-1]==')':
+                    kwargs[k] = tuple([str2value(i) for i in v.strip('()').split(',')])
+                elif v[0]=='[' and v[-1]==']':
+                    kwargs[k] = [str2value(i) for i in v.strip('[]').split(',')]
+                else:
+                    kwargs[k] = str2value(v)
+            return functools.partial(f, **kwargs)
+
+def register(name):
+    def wrapper(class_):
+        get_unit().register(name, class_)
+        return class_
+    return wrapper
+
+class Sine(object):
+    def __init__(self, freq, gain=1):
+        self.freq = freq
+        self.gain = gain
+        self.repr = 'sine(freq={}, gain={})'.format(freq, gain)
+
+    def __call__(self, x, gain=1):
+        act_gain = self.gain * gain
+        return torch.sin(self.freq * x) * act_gain
+
+    def __repr__(self,):
+        return self.repr
+
+class ReLUSine(nn.Module):
+    def __init(self):
+        super().__init__()
+
+    def forward(self, input):
+        a = torch.sin(30 * input)
+        b = nn.ReLU(inplace=False)(input)
+        return a+b
+
+@register('lrelu_agc')
+# class lrelu_agc(nn.Module):
+class lrelu_agc(object):
+    """
+    The lrelu layer with alpha, gain and clamp
+    """
+    def __init__(self, alpha=0.1, gain=1, clamp=None):
+        # super().__init__()
+        self.alpha = alpha
+        if gain == 'sqrt_2':
+            self.gain = np.sqrt(2)
+        else:
+            self.gain = gain
+        self.clamp = clamp
+        self.repr = 'lrelu_agc(alpha={}, gain={}, clamp={})'.format(
+            alpha, gain, clamp)
+
+    # def forward(self, x, gain=1):
+    def __call__(self, x, gain=1):
+        x = F.leaky_relu(x, negative_slope=self.alpha, inplace=True)
+        act_gain = self.gain * gain
+        act_clamp = self.clamp * gain if self.clamp is not None else None
+        if act_gain != 1:
+            x = x * act_gain
+        if act_clamp is not None:
+            x = x.clamp(-act_clamp, act_clamp)
+        return x
+
+    def __repr__(self,):
+        return self.repr
+
+####################
+# spatial encoding #
+####################
+
+@register('se')
+class SpatialEncoding(nn.Module):
+    def __init__(self, 
+                 in_dim, 
+                 out_dim, 
+                 sigma = 6,
+                 cat_input=True,
+                 require_grad=False,):
+
+        super().__init__()
+        assert out_dim % (2*in_dim) == 0, "dimension must be dividable"
+
+        n = out_dim // 2 // in_dim
+        m = 2**np.linspace(0, sigma, n)
+        m = np.stack([m] + [np.zeros_like(m)]*(in_dim-1), axis=-1)
+        m = np.concatenate([np.roll(m, i, axis=-1) for i in range(in_dim)], axis=0)
+        self.emb = torch.FloatTensor(m)
+        if require_grad:
+            self.emb = nn.Parameter(self.emb, requires_grad=True)    
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.sigma = sigma
+        self.cat_input = cat_input
+        self.require_grad = require_grad
+
+    def forward(self, x, format='[n x c]'):
+        """
+        Args:
+            x: [n x m1],
+                m1 usually is 2
+        Outputs:
+            y: [n x m2]         
+                m2 dimention number
+        """
+        if format == '[bs x c x 2D]':
+            xshape = x.shape
+            x = x.permute(0, 2, 3, 1).contiguous()
+            x = x.view(-1, x.size(-1))
+        elif format == '[n x c]':
+            pass
+        else:
+            raise ValueError
+
+        if not self.require_grad:
+            self.emb = self.emb.to(x.device)
+        y = torch.mm(x, self.emb.T)
+        if self.cat_input:
+            z = torch.cat([x, torch.sin(y), torch.cos(y)], dim=-1)
+        else:
+            z = torch.cat([torch.sin(y), torch.cos(y)], dim=-1)
+
+        if format == '[bs x c x 2D]':
+            z = z.view(xshape[0], xshape[2], xshape[3], -1)
+            z = z.permute(0, 3, 1, 2).contiguous()
+        return z
+
+    def extra_repr(self):
+        outstr = 'SpatialEncoding (in={}, out={}, sigma={}, cat_input={}, require_grad={})'.format(
+            self.in_dim, self.out_dim, self.sigma, self.cat_input, self.require_grad)
+        return outstr
+
+@register('rffe')
+class RFFEncoding(SpatialEncoding):
+    """
+    Random Fourier Features
+    """
+    def __init__(self, 
+                 in_dim, 
+                 out_dim, 
+                 sigma = 6,
+                 cat_input=True,
+                 require_grad=False,):
+
+        super().__init__(in_dim, out_dim, sigma, cat_input, require_grad)
+        n = out_dim // 2
+        m = np.random.normal(0, sigma, size=(n, in_dim))
+        self.emb = torch.FloatTensor(m)
+        if require_grad:
+            self.emb = nn.Parameter(self.emb, requires_grad=True)    
+
+    def extra_repr(self):
+        outstr = 'RFFEncoding (in={}, out={}, sigma={}, cat_input={}, require_grad={})'.format(
+            self.in_dim, self.out_dim, self.sigma, self.cat_input, self.require_grad)
+        return outstr
+
+##########
+# helper #
+##########
+
+def freeze(net):
+    for m in net.modules():
+        if isinstance(m, (
+                nn.BatchNorm2d, 
+                nn.SyncBatchNorm,)):
+            # inplace_abn not supported
+            m.eval()
+    for pi in net.parameters():
+        pi.requires_grad = False
+    return net
+
+def common_init(m):
+    if isinstance(m, (
+            nn.Conv2d, 
+            nn.ConvTranspose2d,)):
+        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, (
+            nn.BatchNorm2d, 
+            nn.SyncBatchNorm,)):
+        nn.init.constant_(m.weight, 1)
+        nn.init.constant_(m.bias, 0)
+    else:
+        pass
+
+def init_module(module):
+    """
+    Args:
+        module: [nn.module] list or nn.module
+            a list of module to be initialized.
+    """
+    if isinstance(module, (list, tuple)):
+        module = list(module)
+    else:
+        module = [module]
+
+    for mi in module:
+        for mii in mi.modules():
+            common_init(mii)
+
+def get_total_param(net):
+    if getattr(net, 'parameters', None) is None:
+        return 0
+    return sum(p.numel() for p in net.parameters())
+
+def get_total_param_sum(net):
+    if getattr(net, 'parameters', None) is None:
+        return 0
+    with torch.no_grad():
+        s = sum(p.cpu().detach().numpy().sum().item() for p in net.parameters())
+    return s 

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/ddim.py

@@ -0,0 +1,341 @@
+"""SAMPLING ONLY."""
+
+import torch
+import numpy as np
+from tqdm import tqdm
+from functools import partial
+
+from .diffusion_utils import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
+
+
+class DDIMSampler(object):
+    def __init__(self, model, schedule="linear", **kwargs):
+        super().__init__()
+        self.model = model
+        self.ddpm_num_timesteps = model.num_timesteps
+        self.schedule = schedule
+
+    def register_buffer(self, name, attr):
+        if type(attr) == torch.Tensor:
+            if attr.device != torch.device("cuda"):
+                attr = attr.to(torch.device("cuda"))
+        setattr(self, name, attr)
+
+    def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
+        self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, 
+                                                  num_ddim_timesteps=ddim_num_steps,
+                                                  num_ddpm_timesteps=self.ddpm_num_timesteps,
+                                                  verbose=verbose)
+        alphas_cumprod = self.model.alphas_cumprod
+        assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
+        to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
+
+        self.register_buffer('betas', to_torch(self.model.betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
+        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
+
+        # ddim sampling parameters
+        ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(
+            alphacums=alphas_cumprod.cpu(),
+            ddim_timesteps=self.ddim_timesteps,
+            eta=ddim_eta,verbose=verbose)
+
+        self.register_buffer('ddim_sigmas', ddim_sigmas)
+        self.register_buffer('ddim_alphas', ddim_alphas)
+        self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
+        self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
+        sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
+            (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
+                        1 - self.alphas_cumprod / self.alphas_cumprod_prev))
+        self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
+
+    @torch.no_grad()
+    def sample(self,
+               S,
+               batch_size,
+               shape,
+               conditioning=None,
+               callback=None,
+               normals_sequence=None,
+               img_callback=None,
+               quantize_x0=False,
+               eta=0.,
+               mask=None,
+               x0=None,
+               temperature=1.,
+               noise_dropout=0.,
+               score_corrector=None,
+               corrector_kwargs=None,
+               verbose=True,
+               x_T=None,
+               log_every_t=100,
+               unconditional_guidance_scale=1.,
+               unconditional_conditioning=None,
+               # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
+               **kwargs
+               ):
+        if conditioning is not None:
+            if isinstance(conditioning, dict):
+                cbs = conditioning[list(conditioning.keys())[0]].shape[0]
+                if cbs != batch_size:
+                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+            else:
+                if conditioning.shape[0] != batch_size:
+                    print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
+
+        self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
+        # sampling
+        C, H, W = shape
+        size = (batch_size, C, H, W)
+        print(f'Data shape for DDIM sampling is {size}, eta {eta}')
+
+        samples, intermediates = self.ddim_sampling(conditioning, size,
+                                                    callback=callback,
+                                                    img_callback=img_callback,
+                                                    quantize_denoised=quantize_x0,
+                                                    mask=mask, x0=x0,
+                                                    ddim_use_original_steps=False,
+                                                    noise_dropout=noise_dropout,
+                                                    temperature=temperature,
+                                                    score_corrector=score_corrector,
+                                                    corrector_kwargs=corrector_kwargs,
+                                                    x_T=x_T,
+                                                    log_every_t=log_every_t,
+                                                    unconditional_guidance_scale=unconditional_guidance_scale,
+                                                    unconditional_conditioning=unconditional_conditioning,
+                                                    )
+        return samples, intermediates
+
+    @torch.no_grad()
+    def ddim_sampling(self, 
+                      cond, shape,
+                      x_T=None, 
+                      ddim_use_original_steps=False,
+                      callback=None, 
+                      timesteps=None, 
+                      quantize_denoised=False,
+                      mask=None, x0=None, 
+                      img_callback=None, log_every_t=100,
+                      temperature=1., 
+                      noise_dropout=0., 
+                      score_corrector=None, 
+                      corrector_kwargs=None,
+                      unconditional_guidance_scale=1., 
+                      unconditional_conditioning=None,):
+        device = torch.device('cuda:1')
+        b = shape[0]
+        if x_T is None:
+            img = torch.randn(shape, device=device)
+        else:
+            img = x_T.cuda(1)
+
+        if timesteps is None:
+            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+        elif timesteps is not None and not ddim_use_original_steps:
+            subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
+            timesteps = self.ddim_timesteps[:subset_end]
+
+        intermediates = {'x_inter': [img], 'pred_x0': [img]}
+        time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
+        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+        print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
+
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((b,), step, device=device, dtype=torch.long)
+
+            if mask is not None:
+                assert x0 is not None
+                img_orig = self.model.q_sample(x0, ts)  # TODO: deterministic forward pass?
+                img = img_orig * mask + (1. - mask) * img
+
+            outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
+                                      quantize_denoised=quantize_denoised, temperature=temperature,
+                                      noise_dropout=noise_dropout, score_corrector=score_corrector,
+                                      corrector_kwargs=corrector_kwargs,
+                                      unconditional_guidance_scale=unconditional_guidance_scale,
+                                      unconditional_conditioning=unconditional_conditioning)
+            img, pred_x0 = outs
+            if callback: callback(i)
+            if img_callback: img_callback(pred_x0, i)
+
+            if index % log_every_t == 0 or index == total_steps - 1:
+                intermediates['x_inter'].append(img)
+                intermediates['pred_x0'].append(pred_x0)
+
+        return img, intermediates
+
+    @torch.no_grad()
+    def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
+                      temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+                      unconditional_guidance_scale=1., unconditional_conditioning=None):
+        b, *_, device = *x.shape, 'cuda:1'
+        device = torch.device('cuda:1')
+
+        if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+            e_t = self.model.apply_model(x, t, c)
+        else:
+            x_in = torch.cat([x] * 2)
+            t_in = torch.cat([t] * 2)
+            c_in = torch.cat([unconditional_conditioning, c])
+            e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
+            e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
+
+        if score_corrector is not None:
+            assert self.model.parameterization == "eps"
+            e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
+
+        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+        alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+        sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+        # select parameters corresponding to the currently considered timestep
+        a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
+        a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
+        sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
+        sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
+        print(sqrt_one_minus_at.device)
+        # current prediction for x_0
+        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+        if quantize_denoised:
+            pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+        # direction pointing to x_t
+        dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+        noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+        if noise_dropout > 0.:
+            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+        x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+        return x_prev, pred_x0
+
+    # XX-added for forward then backward
+
+    @torch.no_grad()
+    def sample_fwdbwd(self,
+                      S,
+                      x0,
+                      S_ref,
+                      batch_size,
+                      shape,
+                      conditioning=None,
+                      callback=None,
+                      img_callback=None,
+                      quantize_x0=False,
+                      eta=0.,
+                      temperature=1.,
+                      noise_dropout=0.,
+                      score_corrector=None,
+                      corrector_kwargs=None,
+                      verbose=True,
+                      x_T=None,
+                      log_every_t=100,
+                      unconditional_guidance_scale=1.,
+                      unconditional_conditioning=None,
+                      # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
+                      **kwargs):
+        if conditioning is not None:
+            if isinstance(conditioning, dict):
+                cbs = conditioning[list(conditioning.keys())[0]].shape[0]
+                if cbs != batch_size:
+                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+            else:
+                if conditioning.shape[0] != batch_size:
+                    print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
+
+        self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
+        # sampling
+        C, H, W = shape
+        size = (batch_size, C, H, W)
+        print(f'Data shape for DDIM sampling is {size}, eta {eta}')
+
+        samples, intermediates = self.ddim_sampling_fwdbwd(
+            conditioning, size, x_0=x0, refstep=S_ref,
+            callback=callback,
+            img_callback=img_callback,
+            quantize_denoised=quantize_x0,
+            ddim_use_original_steps=False,
+            noise_dropout=noise_dropout,
+            temperature=temperature,
+            score_corrector=score_corrector,
+            corrector_kwargs=corrector_kwargs,
+            log_every_t=log_every_t,
+            unconditional_guidance_scale=unconditional_guidance_scale,
+            unconditional_conditioning=unconditional_conditioning,)
+
+        return samples, intermediates
+
+
+    @torch.no_grad()
+    def ddim_sampling_fwdbwd(self, 
+                             cond, 
+                             shape,
+                             x_0,
+                             refstep,
+                             ddim_use_original_steps=False,
+                             callback=None, 
+                             timesteps=None, 
+                             quantize_denoised=False,
+                             img_callback=None, 
+                             log_every_t=100,
+                             temperature=1., 
+                             noise_dropout=0., 
+                             score_corrector=None, 
+                             corrector_kwargs=None,
+                             unconditional_guidance_scale=1., 
+                             unconditional_conditioning=None, ):
+        '''
+        A function that forward diffuse x_0 to a reference step 
+            (reference step number is correlated with the timesteps
+             so the real time steps are timesteps[0:refstep] )
+        \i.e. x_0 -> x_k -> x_0'; k \in {0 .. len(timesteps)}
+        '''
+
+        device = self.model.betas.device
+        b = shape[0]
+
+        if timesteps is None:
+            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+        elif timesteps is not None and not ddim_use_original_steps:
+            subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
+            timesteps = self.ddim_timesteps[:subset_end]
+
+        timesteps = timesteps[0:refstep]
+
+        # forward diffusion
+        t = torch.full((b,), timesteps[-1], device=device, dtype=torch.long)
+        x_noisy = self.model.q_sample(x_start=x_0, t=t)
+        img = x_noisy
+
+        intermediates = {'x_inter': [x_noisy], 'pred_x0': [x_0]}
+        time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
+        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+        print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
+
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((b,), step, device=device, dtype=torch.long)
+            outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
+                                      quantize_denoised=quantize_denoised, temperature=temperature,
+                                      noise_dropout=noise_dropout, score_corrector=score_corrector,
+                                      corrector_kwargs=corrector_kwargs,
+                                      unconditional_guidance_scale=unconditional_guidance_scale,
+                                      unconditional_conditioning=unconditional_conditioning)
+            img, pred_x0 = outs
+            if callback: callback(i)
+            if img_callback: img_callback(pred_x0, i)
+
+            if index % log_every_t == 0 or index == total_steps - 1:
+                intermediates['x_inter'].append(img)
+                intermediates['pred_x0'].append(pred_x0)
+
+        return img, intermediates

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/ddim_dualcontext.py

@@ -0,0 +1,144 @@
+import torch
+import numpy as np
+from tqdm import tqdm
+from functools import partial
+
+from .diffusion_utils import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
+
+from .ddim import DDIMSampler
+
+class DDIMSampler_DualContext(DDIMSampler):
+    @torch.no_grad()
+    def sample_text(self, *args, **kwargs):
+        self.cond_type = 'prompt'
+        return self.sample(*args, **kwargs)
+
+    @torch.no_grad()
+    def sample_vision(self, *args, **kwargs):
+        self.cond_type = 'vision'
+        return self.sample(*args, **kwargs)
+
+    @torch.no_grad()
+    def sample_mixed(self, *args, **kwargs):
+        self.cond_type = kwargs.pop('cond_mixed_p')
+        return self.sample(*args, **kwargs)
+
+    @torch.no_grad()
+    def sample(self,
+               steps,
+               shape,
+               xt=None,
+               conditioning=None,
+               eta=0.,
+               temperature=1.,
+               noise_dropout=0.,
+               verbose=True,
+               log_every_t=100,
+               unconditional_guidance_scale=1.,
+               unconditional_conditioning=None,):
+
+        self.make_schedule(ddim_num_steps=steps, ddim_eta=eta, verbose=verbose)
+        # sampling
+        print(f'Data shape for DDIM sampling is {shape}, eta {eta}')
+
+        samples, intermediates = self.ddim_sampling(
+            conditioning, 
+            shape,
+            xt=xt,
+            ddim_use_original_steps=False,
+            noise_dropout=noise_dropout,
+            temperature=temperature,
+            log_every_t=log_every_t,
+            unconditional_guidance_scale=unconditional_guidance_scale,
+            unconditional_conditioning=unconditional_conditioning,)
+        return samples, intermediates
+
+    @torch.no_grad()
+    def ddim_sampling(self, 
+                      conditioning,
+                      shape,
+                      xt=None, 
+                      ddim_use_original_steps=False,
+                      timesteps=None, 
+                      log_every_t=100,
+                      temperature=1., 
+                      noise_dropout=0., 
+                      unconditional_guidance_scale=1., 
+                      unconditional_conditioning=None,):
+        device = self.model.betas.device
+        bs = shape[0]
+        if xt is None:
+            img = torch.randn(shape, device=device)
+        else:
+            img = xt
+
+        if timesteps is None:
+            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+        elif timesteps is not None and not ddim_use_original_steps:
+            subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
+            timesteps = self.ddim_timesteps[:subset_end]
+
+        intermediates = {'x_inter': [img], 'pred_x0': [img]}
+        time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
+        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+        print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
+
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((bs,), step, device=device, dtype=torch.long)
+
+            outs = self.p_sample_ddim(img, conditioning, ts, index=index, use_original_steps=ddim_use_original_steps,
+                                      temperature=temperature,
+                                      noise_dropout=noise_dropout,
+                                      unconditional_guidance_scale=unconditional_guidance_scale,
+                                      unconditional_conditioning=unconditional_conditioning)
+            img, pred_x0 = outs
+
+            if index % log_every_t == 0 or index == total_steps - 1:
+                intermediates['x_inter'].append(img)
+                intermediates['pred_x0'].append(pred_x0)
+
+        return img, intermediates
+
+    @torch.no_grad()
+    def p_sample_ddim(self, x, conditioning, t, index, repeat_noise=False, use_original_steps=False, 
+                      temperature=1., noise_dropout=0.,
+                      unconditional_guidance_scale=1., unconditional_conditioning=None):
+        b, *_, device = *x.shape, x.device
+
+        if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+            e_t = self.model.apply_model(x, t, conditioning, cond_type=self.cond_type)
+        else:
+            x_in = torch.cat([x] * 2)
+            t_in = torch.cat([t] * 2)
+            # c_in = torch.cat([unconditional_conditioning, conditioning])
+
+            # Added for vd-dc dual guidance
+            if isinstance(unconditional_conditioning, list):
+                c_in = [torch.cat([ui, ci]) for ui, ci in zip(unconditional_conditioning, conditioning)]
+            else:
+                c_in = torch.cat([unconditional_conditioning, conditioning])
+
+            e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in, cond_type=self.cond_type).chunk(2)
+            e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
+
+        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+        alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+        sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+        # select parameters corresponding to the currently considered timestep
+        a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
+        a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
+        sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
+        sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
+
+        # current prediction for x_0
+        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+        dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+        noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+        if noise_dropout > 0.:
+            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+        x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+        return x_prev, pred_x0

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/ddim_dualmodel.py

@@ -0,0 +1,244 @@
+import torch
+import numpy as np
+from tqdm import tqdm
+from functools import partial
+
+from .diffusion_utils import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
+
+from .ddim import DDIMSampler
+
+class DDIMSampler_DualModel(DDIMSampler):
+    def __init__(self, model_t2i, model_v2i, schedule="linear", **kwargs):
+        self.model = model_t2i
+        self.model_t2i = model_t2i
+        self.model_v2i = model_v2i
+        self.device = self.model_t2i.device
+        self.ddpm_num_timesteps = model_t2i.num_timesteps
+        self.schedule = schedule
+
+    @torch.no_grad()
+    def sample_text(self, *args, **kwargs):
+        self.cond_type = 'prompt'
+        self.p_sample_model_type = 't2i'
+        return self.sample(*args, **kwargs)
+
+    @torch.no_grad()
+    def sample_vision(self, *args, **kwargs):
+        self.cond_type = 'vision'
+        self.p_sample_model_type = 'v2i'
+        return self.sample(*args, **kwargs)
+
+    @torch.no_grad()
+    def sample(self,
+               steps,
+               shape,
+               xt=None,
+               conditioning=None,
+               eta=0.,
+               temperature=1.,
+               noise_dropout=0.,
+               verbose=True,
+               log_every_t=100,
+               unconditional_guidance_scale=1.,
+               unconditional_conditioning=None,):
+
+        self.make_schedule(ddim_num_steps=steps, ddim_eta=eta, verbose=verbose)
+        # sampling
+        print(f'Data shape for DDIM sampling is {shape}, eta {eta}')
+
+        samples, intermediates = self.ddim_sampling(
+            conditioning, 
+            shape,
+            xt=xt,
+            ddim_use_original_steps=False,
+            noise_dropout=noise_dropout,
+            temperature=temperature,
+            log_every_t=log_every_t,
+            unconditional_guidance_scale=unconditional_guidance_scale,
+            unconditional_conditioning=unconditional_conditioning,)
+        return samples, intermediates
+
+    @torch.no_grad()
+    def ddim_sampling(self, 
+                      conditioning,
+                      shape,
+                      xt=None, 
+                      ddim_use_original_steps=False,
+                      timesteps=None, 
+                      log_every_t=100,
+                      temperature=1., 
+                      noise_dropout=0., 
+                      unconditional_guidance_scale=1., 
+                      unconditional_conditioning=None,):
+        device = self.model.betas.device
+        bs = shape[0]
+        if xt is None:
+            img = torch.randn(shape, device=device)
+        else:
+            img = xt
+
+        if timesteps is None:
+            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+        elif timesteps is not None and not ddim_use_original_steps:
+            subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
+            timesteps = self.ddim_timesteps[:subset_end]
+
+        intermediates = {'x_inter': [img], 'pred_x0': [img]}
+        time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
+        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+        print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
+
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((bs,), step, device=device, dtype=torch.long)
+
+            outs = self.p_sample_ddim(img, conditioning, ts, index=index, use_original_steps=ddim_use_original_steps,
+                                      temperature=temperature,
+                                      noise_dropout=noise_dropout,
+                                      unconditional_guidance_scale=unconditional_guidance_scale,
+                                      unconditional_conditioning=unconditional_conditioning)
+            img, pred_x0 = outs
+
+            if index % log_every_t == 0 or index == total_steps - 1:
+                intermediates['x_inter'].append(img)
+                intermediates['pred_x0'].append(pred_x0)
+
+        return img, intermediates
+
+    @torch.no_grad()
+    def p_sample_ddim(self, x, conditioning, t, index, repeat_noise=False, use_original_steps=False, 
+                      temperature=1., noise_dropout=0.,
+                      unconditional_guidance_scale=1., unconditional_conditioning=None):
+        b, *_, device = *x.shape, x.device
+
+        if self.p_sample_model_type == 't2i':
+            apply_model = self.model_t2i.apply_model
+        elif self.p_sample_model_type == 'v2i':
+            apply_model = self.model_v2i.apply_model
+
+        if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+            e_t = apply_model(x, t, conditioning)
+        else:
+            x_in = torch.cat([x] * 2)
+            t_in = torch.cat([t] * 2)
+            c_in = torch.cat([unconditional_conditioning, conditioning])
+            e_t_uncond, e_t = apply_model(x_in, t_in, c_in).chunk(2)
+            e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
+
+        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+        alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+        sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+        # select parameters corresponding to the currently considered timestep
+        a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
+        a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
+        sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
+        sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
+
+        # current prediction for x_0
+        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+        dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+        noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+        if noise_dropout > 0.:
+            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+        x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+        return x_prev, pred_x0
+
+    @torch.no_grad()
+    def sample_mixed(self,
+                     steps,
+                     steps_t2i,
+                     steps_v2i,
+                     shape,
+                     xt=None,
+                     c_prompt=None,
+                     c_vision=None,
+                     eta=0.,
+                     temperature=1.,
+                     noise_dropout=0.,
+                     verbose=True,
+                     log_every_t=100,
+                     uc_scale=1.,
+                     uc_prompt=None,
+                     uc_vision=None,):
+
+        print(f'DDIM mixed sampling with shape {shape}, eta {eta}')
+        print(f'steps_t2i {steps_t2i}')
+        print(f'steps_v2i {steps_v2i}')
+
+        self.make_schedule(ddim_num_steps=steps, ddim_eta=eta, verbose=verbose)
+        self.ddim_timesteps_t2i = self.ddim_timesteps[steps_t2i]
+        self.ddim_timesteps_v2i = self.ddim_timesteps[steps_v2i]
+
+        samples, intermediates = self.ddim_sampling_mixed(
+            c_prompt, 
+            c_vision,
+            shape,
+            xt=xt,
+            noise_dropout=noise_dropout,
+            temperature=temperature,
+            log_every_t=log_every_t,
+            uc_scale=uc_scale,
+            uc_prompt=uc_prompt,
+            uc_vision=uc_vision, )
+        return samples, intermediates
+
+    @torch.no_grad()
+    def ddim_sampling_mixed(self, 
+                            c_prompt,
+                            c_vision,
+                            shape,
+                            xt=None, 
+                            log_every_t=100,
+                            temperature=1., 
+                            noise_dropout=0., 
+                            uc_scale=1., 
+                            uc_prompt=None,
+                            uc_vision=None, ):
+        device = self.device
+        bs = shape[0]
+        if xt is None:
+            img = torch.randn(shape, device=device)
+        else:
+            img = xt
+
+        timesteps = self.ddim_timesteps
+        intermediates = {'x_inter': [], 'pred_x0': []}
+        time_range = np.flip(timesteps)
+        total_steps = timesteps.shape[0]
+        print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
+
+        for i, step in enumerate(iterator):            
+            if step in self.ddim_timesteps_t2i:
+                self.p_sample_model_type = 't2i'
+                conditioning = c_prompt
+                unconditional_conditioning = uc_prompt
+            elif step in self.ddim_timesteps_v2i:
+                self.p_sample_model_type = 'v2i'
+                conditioning = c_vision
+                unconditional_conditioning = uc_vision
+            else:
+                raise ValueError # shouldn't reached
+
+            index = total_steps - i - 1
+            ts = torch.full((bs,), step, device=device, dtype=torch.long)
+            outs = self.p_sample_ddim(
+                img, conditioning, ts, 
+                index=index,
+                temperature=temperature,
+                noise_dropout=noise_dropout,
+                unconditional_guidance_scale=uc_scale,
+                unconditional_conditioning=unconditional_conditioning)
+            img, pred_x0 = outs
+
+            if index % log_every_t == 0 or index == total_steps - 1:
+                intermediates['x_inter'].append(img)
+                intermediates['pred_x0'].append(pred_x0)
+
+        return img, intermediates
+
+

versatile_diffusion/log/sd_nodataset/99999_evalonly/sd_variation/code/lib/model_zoo/ddim_vd.py

@@ -0,0 +1,293 @@
+import torch
+import numpy as np
+from tqdm import tqdm
+from functools import partial
+
+from .diffusion_utils import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
+
+from .ddim import DDIMSampler
+
+class DDIMSampler_VD(DDIMSampler):
+    @torch.no_grad()
+    def sample(self,
+               steps,
+               shape,
+               xt=None,
+               conditioning=None,
+               unconditional_guidance_scale=1.,
+               unconditional_conditioning=None,
+               xtype='image',
+               ctype='prompt',
+               eta=0.,
+               temperature=1.,
+               noise_dropout=0.,
+               verbose=True,
+               log_every_t=100,):
+
+        self.make_schedule(ddim_num_steps=steps, ddim_eta=eta, verbose=verbose)
+        print(f'Data shape for DDIM sampling is {shape}, eta {eta}')
+        samples, intermediates = self.ddim_sampling(
+            shape,
+            xt=xt,
+            conditioning=conditioning, 
+            unconditional_guidance_scale=unconditional_guidance_scale,
+            unconditional_conditioning=unconditional_conditioning,
+            xtype=xtype,
+            ctype=ctype,
+            ddim_use_original_steps=False,
+            noise_dropout=noise_dropout,
+            temperature=temperature,
+            log_every_t=log_every_t,)
+        return samples, intermediates
+
+    @torch.no_grad()
+    def ddim_sampling(self, 
+                      shape,
+                      xt=None,
+                      conditioning=None,
+                      unconditional_guidance_scale=1., 
+                      unconditional_conditioning=None,
+                      xtype='image',
+                      ctype='prompt',
+                      ddim_use_original_steps=False,
+                      timesteps=None, 
+                      noise_dropout=0., 
+                      temperature=1., 
+                      log_every_t=100,):
+
+        device = 1
+        bs = shape[0]
+        if xt is None:
+            xt = torch.randn(shape, device=device)
+
+        if timesteps is None:
+            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+        elif timesteps is not None and not ddim_use_original_steps:
+            subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
+            timesteps = self.ddim_timesteps[:subset_end]
+
+        intermediates = {'pred_xt': [], 'pred_x0': []}
+        time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
+        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+        # print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+        pred_xt = xt
+        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((bs,), step, device=device, dtype=torch.long)
+
+            outs = self.p_sample_ddim(
+                pred_xt, conditioning, ts, index, 
+                unconditional_guidance_scale=unconditional_guidance_scale,
+                unconditional_conditioning=unconditional_conditioning, 
+                xtype=xtype,
+                ctype=ctype,
+                use_original_steps=ddim_use_original_steps,
+                noise_dropout=noise_dropout,
+                temperature=temperature,)
+            pred_xt, pred_x0 = outs
+
+            if index % log_every_t == 0 or index == total_steps - 1:
+                intermediates['pred_xt'].append(pred_xt)
+                intermediates['pred_x0'].append(pred_x0)
+
+        return pred_xt, intermediates
+
+    @torch.no_grad()
+    def p_sample_ddim(self, x, conditioning, t, index, 
+                      unconditional_guidance_scale=1., 
+                      unconditional_conditioning=None, 
+                      xtype='image',
+                      ctype='prompt',
+                      repeat_noise=False, 
+                      use_original_steps=False, 
+                      noise_dropout=0.,
+                      temperature=1.,):
+
+        b, *_, device = *x.shape, x.device
+
+        if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+            e_t = self.model.apply_model(x, t, conditioning, xtype=xtype, ctype=ctype)
+        else:
+            x_in = torch.cat([x] * 2)
+            t_in = torch.cat([t] * 2)
+            c_in = torch.cat([unconditional_conditioning, conditioning])
+            e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in, xtype=xtype, ctype=ctype).chunk(2)
+            e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
+
+        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+        alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+        sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+        # select parameters corresponding to the currently considered timestep
+
+        if xtype == 'image':
+            extended_shape = (b, 1, 1, 1)
+        elif xtype == 'text':
+            extended_shape = (b, 1)
+
+        a_t = torch.full(extended_shape, alphas[index], device=device)
+        a_prev = torch.full(extended_shape, alphas_prev[index], device=device)
+        sigma_t = torch.full(extended_shape, sigmas[index], device=device)
+        sqrt_one_minus_at = torch.full(extended_shape, sqrt_one_minus_alphas[index],device=device)
+
+        # current prediction for x_0
+        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+        dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+        noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+        if noise_dropout > 0.:
+            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+        x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+        return x_prev, pred_x0
+
+class DDIMSampler_VD_DualContext(DDIMSampler_VD):
+    @torch.no_grad()
+    def sample_dc(self,
+               steps,
+               shape,
+               xt=None,
+               first_conditioning=None,
+               second_conditioning=None,
+               unconditional_guidance_scale=1.,
+               xtype='image',
+               first_ctype='prompt',
+               second_ctype='prompt',
+               eta=0.,
+               temperature=1.,
+               mixed_ratio=0.5,
+               noise_dropout=0.,
+               verbose=True,
+               log_every_t=100,):
+
+        self.make_schedule(ddim_num_steps=steps, ddim_eta=eta, verbose=verbose)
+        print(f'Data shape for DDIM sampling is {shape}, eta {eta}')
+        samples, intermediates = self.ddim_sampling_dc(
+            shape,
+            xt=xt,
+            first_conditioning=first_conditioning,
+            second_conditioning=second_conditioning,
+            unconditional_guidance_scale=unconditional_guidance_scale,
+            xtype=xtype,
+            first_ctype=first_ctype,
+            second_ctype=second_ctype,
+            ddim_use_original_steps=False,
+            noise_dropout=noise_dropout,
+            temperature=temperature,
+            log_every_t=log_every_t,
+            mixed_ratio=mixed_ratio, )
+        return samples, intermediates
+
+    @torch.no_grad()
+    def ddim_sampling_dc(self, 
+                      shape,
+                      xt=None,
+                      first_conditioning=None,
+                      second_conditioning=None,
+                      unconditional_guidance_scale=1., 
+                      xtype='image',
+                      first_ctype='prompt',
+                      second_ctype='prompt',
+                      ddim_use_original_steps=False,
+                      timesteps=None, 
+                      noise_dropout=0., 
+                      temperature=1.,
+                      mixed_ratio=0.5,
+                      log_every_t=100,):
+
+        device = self.model.device
+        bs = shape[0]
+        if xt is None:
+            xt = torch.randn(shape, device=device)
+
+        if timesteps is None:
+            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+        elif timesteps is not None and not ddim_use_original_steps:
+            subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
+            timesteps = self.ddim_timesteps[:subset_end]
+
+        intermediates = {'pred_xt': [], 'pred_x0': []}
+        time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
+        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+        # print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+        pred_xt = xt
+        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((bs,), step, device=device, dtype=torch.long)
+
+            outs = self.p_sample_ddim_dc(
+                pred_xt, 
+                first_conditioning, 
+                second_conditioning, 
+                ts, index, 
+                unconditional_guidance_scale=unconditional_guidance_scale,
+                xtype=xtype,
+                first_ctype=first_ctype,
+                second_ctype=second_ctype,
+                use_original_steps=ddim_use_original_steps,
+                noise_dropout=noise_dropout,
+                temperature=temperature,
+                mixed_ratio=mixed_ratio,)
+            pred_xt, pred_x0 = outs
+
+            if index % log_every_t == 0 or index == total_steps - 1:
+                intermediates['pred_xt'].append(pred_xt)
+                intermediates['pred_x0'].append(pred_x0)
+
+        return pred_xt, intermediates
+
+    @torch.no_grad()
+    def p_sample_ddim_dc(self, x, 
+                      first_conditioning,
+                      second_conditioning,
+                      t, index, 
+                      unconditional_guidance_scale=1., 
+                      xtype='image',
+                      first_ctype='prompt',
+                      second_ctype='prompt',
+                      repeat_noise=False, 
+                      use_original_steps=False, 
+                      noise_dropout=0.,
+                      temperature=1.,
+                      mixed_ratio=0.5,):
+
+        b, *_, device = *x.shape, x.device
+
+        x_in = torch.cat([x] * 2)
+        t_in = torch.cat([t] * 2)
+        first_c = torch.cat(first_conditioning)
+        second_c = torch.cat(second_conditioning)
+
+        e_t_uncond, e_t = self.model.apply_model_dc(
+            x_in, t_in, first_c, second_c, xtype=xtype, first_ctype=first_ctype, second_ctype=second_ctype, mixed_ratio=mixed_ratio).chunk(2)
+
+        # e_t_uncond, e_t = self.model.apply_model(x_in, t_in, first_c, xtype='image', ctype='vision').chunk(2)
+
+        e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
+
+        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+        alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+        sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+        # select parameters corresponding to the currently considered timestep
+
+        if xtype == 'image':
+            extended_shape = (b, 1, 1, 1)
+        elif xtype == 'text':
+            extended_shape = (b, 1)
+
+        a_t = torch.full(extended_shape, alphas[index], device=device)
+        a_prev = torch.full(extended_shape, alphas_prev[index], device=device)
+        sigma_t = torch.full(extended_shape, sigmas[index], device=device)
+        sqrt_one_minus_at = torch.full(extended_shape, sqrt_one_minus_alphas[index],device=device)
+
+        # current prediction for x_0
+        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+        dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+        noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+        if noise_dropout > 0.:
+            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+        x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+        return x_prev, pred_x0