Duplicate from shi-labs/Versatile-Diffusion

Co-authored-by: Xingqian Xu <JamesXu@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,7 @@
+__pycache__
+.vscode/
+src/
+data/
+data
+log/
+log

README.md

@@ -0,0 +1,15 @@
+---
+title: Versatile Diffusion
+emoji: null
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 3.9.1
+app_file: app.py
+pinned: false
+license: mit
+python_version: 3.8.5
+duplicated_from: shi-labs/Versatile-Diffusion
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,845 @@
+import gradio as gr
+import os
+import PIL
+from PIL import Image
+from pathlib import Path
+import numpy as np
+import numpy.random as npr
+from contextlib import nullcontext
+
+import torch
+import torchvision.transforms as tvtrans
+from lib.cfg_helper import model_cfg_bank
+from lib.model_zoo import get_model
+
+n_sample_image_default = 2
+n_sample_text_default = 4
+cache_examples = True
+hfm_repo_id = 'shi-labs/versatile-diffusion-model'
+hfm_filename = 'pretrained_pth/vd-four-flow-v1-0-fp16.pth' 
+
+def highlight_print(info):
+    print('')
+    print(''.join(['#']*(len(info)+4)))
+    print('# '+info+' #')
+    print(''.join(['#']*(len(info)+4)))
+    print('')
+
+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
+
+class vd_inference(object):
+    def __init__(self, pth=None, hfm_repo=None, fp16=False, device=0):
+        cfgm_name = 'vd_noema'
+        cfgm = model_cfg_bank()('vd_noema')
+        net = get_model()(cfgm)
+        if fp16:
+            highlight_print('Running in FP16')
+            net.clip.fp16 = True
+            net = net.half()
+        if pth is not None:
+            sd = torch.load(pth, map_location='cpu')
+            print('Load pretrained weight from {}'.format(pth))
+        else:
+            from huggingface_hub import hf_hub_download
+            temppath = hf_hub_download(hfm_repo[0], hfm_repo[1])
+            sd = torch.load(temppath, map_location='cpu')
+            print('Load pretrained weight from {}/{}'.format(*hfm_repo))
+
+        net.load_state_dict(sd, strict=False)
+        net.to(device)
+
+        self.device = device
+        self.model_name = cfgm_name
+        self.net = net
+        self.fp16 = fp16
+        from lib.model_zoo.ddim_vd import DDIMSampler_VD
+        self.sampler = DDIMSampler_VD(net)
+
+    def regularize_image(self, x):
+        BICUBIC = PIL.Image.Resampling.BICUBIC
+        if isinstance(x, str):
+            x = Image.open(x).resize([512, 512], resample=BICUBIC)
+            x = tvtrans.ToTensor()(x)
+        elif isinstance(x, PIL.Image.Image):
+            x = x.resize([512, 512], resample=BICUBIC)
+            x = tvtrans.ToTensor()(x)
+        elif isinstance(x, np.ndarray):
+            x = PIL.Image.fromarray(x).resize([512, 512], resample=BICUBIC)
+            x = tvtrans.ToTensor()(x)
+        elif isinstance(x, torch.Tensor):
+            pass
+        else:
+            assert False, 'Unknown image type'
+
+        assert (x.shape[1]==512) & (x.shape[2]==512), \
+            'Wrong image size'
+        x = x.to(self.device)
+        if self.fp16:
+            x = x.half()
+        return x
+
+    def decode(self, z, xtype, ctype, color_adj='None', color_adj_to=None):
+        net = self.net
+        if xtype == 'image':
+            x = net.autokl_decode(z)
+
+            color_adj_flag = (color_adj!='none') and (color_adj!='None') and (color_adj is not None)
+            color_adj_simple = (color_adj=='Simple') or color_adj=='simple'
+            color_adj_keep_ratio = 0.5
+
+            if color_adj_flag and (ctype=='vision'):
+                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=color_adj_keep_ratio, simple=color_adj_simple)
+                    x_adj.append(xi_adj)
+                x = x_adj
+            else:
+                x = torch.clamp((x+1.0)/2.0, min=0.0, max=1.0)
+                x = [tvtrans.ToPILImage()(xi) for xi in x]
+            return x
+
+        elif xtype == 'text':
+            prompt_temperature = 1.0
+            prompt_merge_same_adj_word = True
+            x = net.optimus_decode(z, temperature=prompt_temperature)
+            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
+            return x
+
+    def inference(self, xtype, cin, ctype, scale=7.5, n_samples=None, color_adj=None,):
+        net = self.net
+        sampler = self.sampler
+        ddim_steps = 50
+        ddim_eta = 0.0
+
+        if xtype == 'image':
+            n_samples = n_sample_image_default if n_samples is None else n_samples
+        elif xtype == 'text':
+            n_samples = n_sample_text_default if n_samples is None else n_samples
+
+        if ctype in ['prompt', 'text']:
+            c = net.clip_encode_text(n_samples * [cin])
+            u = None
+            if scale != 1.0:
+                u = net.clip_encode_text(n_samples * [""])
+
+        elif ctype in ['vision', 'image']:
+            cin = self.regularize_image(cin)
+            ctemp = cin*2 - 1
+            ctemp = ctemp[None].repeat(n_samples, 1, 1, 1)
+            c = net.clip_encode_vision(ctemp)
+            u = None
+            if scale != 1.0:
+                dummy = torch.zeros_like(ctemp)
+                u = net.clip_encode_vision(dummy)
+
+        u, c = [u.half(), c.half()] if self.fp16 else [u, c]
+
+        if xtype == 'image':
+            h, w = [512, 512]
+            shape = [n_samples, 4, h//8, w//8]
+            z, _ = sampler.sample(
+                steps=ddim_steps,
+                shape=shape,
+                conditioning=c,
+                unconditional_guidance_scale=scale,
+                unconditional_conditioning=u,
+                xtype=xtype, ctype=ctype,
+                eta=ddim_eta,
+                verbose=False,)
+            x = self.decode(z, xtype, ctype, color_adj=color_adj, color_adj_to=cin)
+            return x
+
+        elif xtype == 'text':
+            n = 768
+            shape = [n_samples, n]
+            z, _ = sampler.sample(
+                steps=ddim_steps,
+                shape=shape,
+                conditioning=c,
+                unconditional_guidance_scale=scale,
+                unconditional_conditioning=u,
+                xtype=xtype, ctype=ctype,
+                eta=ddim_eta,
+                verbose=False,)
+            x = self.decode(z, xtype, ctype)
+            return x
+
+    def application_disensemble(self, cin, n_samples=None, level=0, color_adj=None,):
+        net = self.net
+        scale = 7.5
+        sampler = self.sampler
+        ddim_steps = 50
+        ddim_eta = 0.0
+        n_samples = n_sample_image_default if n_samples is None else n_samples
+
+        cin = self.regularize_image(cin)
+        ctemp = cin*2 - 1
+        ctemp = ctemp[None].repeat(n_samples, 1, 1, 1)
+        c = net.clip_encode_vision(ctemp)
+        u = None
+        if scale != 1.0:
+            dummy = torch.zeros_like(ctemp)
+            u = net.clip_encode_vision(dummy)
+        u, c = [u.half(), c.half()] if self.fp16 else [u, c]
+
+        if level == 0:
+            pass
+        else:
+            c_glb = c[:, 0:1]
+            c_loc = c[:, 1: ]
+            u_glb = u[:, 0:1]
+            u_loc = u[:, 1: ]
+
+            if level == -1:
+                c_loc = self.remove_low_rank(c_loc, demean=True, q=50, q_remove=1)
+                u_loc = self.remove_low_rank(u_loc, demean=True, q=50, q_remove=1)
+            if level == -2:
+                c_loc = self.remove_low_rank(c_loc, demean=True, q=50, q_remove=2)
+                u_loc = self.remove_low_rank(u_loc, demean=True, q=50, q_remove=2)
+            if level == 1:
+                c_loc = self.find_low_rank(c_loc, demean=True, q=10)
+                u_loc = self.find_low_rank(u_loc, demean=True, q=10)
+            if level == 2:
+                c_loc = self.find_low_rank(c_loc, demean=True, q=2)
+                u_loc = self.find_low_rank(u_loc, demean=True, q=2)
+
+            c = torch.cat([c_glb, c_loc], dim=1)
+            u = torch.cat([u_glb, u_loc], dim=1)
+
+        h, w = [512, 512]
+        shape = [n_samples, 4, h//8, w//8]
+        z, _ = sampler.sample(
+            steps=ddim_steps,
+            shape=shape,
+            conditioning=c,
+            unconditional_guidance_scale=scale,
+            unconditional_conditioning=u,
+            xtype='image', ctype='vision',
+            eta=ddim_eta,
+            verbose=False,)
+        x = self.decode(z, 'image', 'vision', color_adj=color_adj, color_adj_to=cin)
+        return x
+
+    def find_low_rank(self, x, demean=True, q=20, niter=10):
+        if demean:
+            x_mean = x.mean(-1, keepdim=True)
+            x_input = x - x_mean
+        else:
+            x_input = x
+
+        if x_input.dtype == torch.float16:
+            fp16 = True
+            x_input = x_input.float()
+        else:
+            fp16 = False
+
+        u, s, v = torch.pca_lowrank(x_input, q=q, center=False, niter=niter)
+        ss = torch.stack([torch.diag(si) for si in s])
+        x_lowrank = torch.bmm(torch.bmm(u, ss), torch.permute(v, [0, 2, 1]))        
+
+        if fp16:
+            x_lowrank = x_lowrank.half()
+
+        if demean:
+            x_lowrank += x_mean
+        return x_lowrank
+
+    def remove_low_rank(self, x, demean=True, q=20, niter=10, q_remove=10):
+        if demean:
+            x_mean = x.mean(-1, keepdim=True)
+            x_input = x - x_mean
+        else:
+            x_input = x
+
+        if x_input.dtype == torch.float16:
+            fp16 = True
+            x_input = x_input.float()
+        else:
+            fp16 = False
+
+        u, s, v = torch.pca_lowrank(x_input, q=q, center=False, niter=niter)
+        s[:, 0:q_remove] = 0
+        ss = torch.stack([torch.diag(si) for si in s])
+        x_lowrank = torch.bmm(torch.bmm(u, ss), torch.permute(v, [0, 2, 1]))        
+
+        if fp16:
+            x_lowrank = x_lowrank.half()
+
+        if demean:
+            x_lowrank += x_mean
+        return x_lowrank
+
+    def application_dualguided(self, cim, ctx, n_samples=None, mixing=0.5, color_adj=None, ):
+        net = self.net
+        scale = 7.5
+        sampler = self.sampler
+        ddim_steps = 50
+        ddim_eta = 0.0
+        n_samples = n_sample_image_default if n_samples is None else n_samples
+
+        ctemp0 = self.regularize_image(cim)
+        ctemp1 = ctemp0*2 - 1
+        ctemp1 = ctemp1[None].repeat(n_samples, 1, 1, 1)
+        cim = net.clip_encode_vision(ctemp1)
+        uim = None
+        if scale != 1.0:
+            dummy = torch.zeros_like(ctemp1)
+            uim = net.clip_encode_vision(dummy)
+
+        ctx = net.clip_encode_text(n_samples * [ctx])
+        utx = None
+        if scale != 1.0:
+            utx = net.clip_encode_text(n_samples * [""])
+
+        uim, cim = [uim.half(), cim.half()] if self.fp16 else [uim, cim]
+        utx, ctx = [utx.half(), ctx.half()] if self.fp16 else [utx, ctx]
+
+        h, w = [512, 512]
+        shape = [n_samples, 4, h//8, w//8]
+
+        z, _ = sampler.sample_dc(
+            steps=ddim_steps,
+            shape=shape,
+            first_conditioning=[uim, cim],
+            second_conditioning=[utx, ctx],
+            unconditional_guidance_scale=scale,
+            xtype='image', 
+            first_ctype='vision',
+            second_ctype='prompt',
+            eta=ddim_eta,
+            verbose=False,
+            mixed_ratio=(1-mixing), )
+        x = self.decode(z, 'image', 'vision', color_adj=color_adj, color_adj_to=ctemp0)
+        return x
+
+    def application_i2t2i(self, cim, ctx_n, ctx_p, n_samples=None, color_adj=None,):
+        net = self.net
+        scale = 7.5
+        sampler = self.sampler
+        ddim_steps = 50
+        ddim_eta = 0.0
+        prompt_temperature = 1.0
+        n_samples = n_sample_image_default if n_samples is None else n_samples
+
+        ctemp0 = self.regularize_image(cim)
+        ctemp1 = ctemp0*2 - 1
+        ctemp1 = ctemp1[None].repeat(n_samples, 1, 1, 1)
+        cim = net.clip_encode_vision(ctemp1)
+        uim = None
+        if scale != 1.0:
+            dummy = torch.zeros_like(ctemp1)
+            uim = net.clip_encode_vision(dummy)
+
+        uim, cim = [uim.half(), cim.half()] if self.fp16 else [uim, cim]
+
+        n = 768
+        shape = [n_samples, n]
+        zt, _ = sampler.sample(
+            steps=ddim_steps,
+            shape=shape,
+            conditioning=cim,
+            unconditional_guidance_scale=scale,
+            unconditional_conditioning=uim,
+            xtype='text', ctype='vision',
+            eta=ddim_eta,
+            verbose=False,)
+        ztn = net.optimus_encode([ctx_n])
+        ztp = net.optimus_encode([ctx_p])
+
+        ztn_norm = ztn / ztn.norm(dim=1)
+        zt_proj_mag = torch.matmul(zt, ztn_norm[0])
+        zt_perp = zt - zt_proj_mag[:, None] * ztn_norm
+        zt_newd = zt_perp + ztp
+        ctx_new = net.optimus_decode(zt_newd, temperature=prompt_temperature)
+
+        ctx_new = net.clip_encode_text(ctx_new)
+        ctx_p = net.clip_encode_text([ctx_p])
+        ctx_new = torch.cat([ctx_new, ctx_p.repeat(n_samples, 1, 1)], dim=1)
+        utx_new = net.clip_encode_text(n_samples * [""])
+        utx_new = torch.cat([utx_new, utx_new], dim=1)
+
+        cim_loc = cim[:, 1: ]
+        cim_loc_new = self.find_low_rank(cim_loc, demean=True, q=10)
+        cim_new = cim_loc_new
+        uim_new = uim[:, 1:]
+        
+        h, w = [512, 512]
+        shape = [n_samples, 4, h//8, w//8]
+        z, _ = sampler.sample_dc(
+            steps=ddim_steps,
+            shape=shape,
+            first_conditioning=[uim_new, cim_new],
+            second_conditioning=[utx_new, ctx_new],
+            unconditional_guidance_scale=scale,
+            xtype='image', 
+            first_ctype='vision',
+            second_ctype='prompt',
+            eta=ddim_eta,
+            verbose=False,
+            mixed_ratio=0.33, )
+
+        x = self.decode(z, 'image', 'vision', color_adj=color_adj, color_adj_to=ctemp0)
+        return x
+
+vd_inference = vd_inference(hfm_repo=[hfm_repo_id, hfm_filename], fp16=True, device='cuda')
+
+def main(mode,
+         image=None,
+         prompt=None,
+         nprompt=None,
+         pprompt=None,
+         color_adj=None,
+         disentanglement_level=None,
+         dual_guided_mixing=None,
+         seed=0,):
+
+    if seed<0:
+        seed = 0
+    np.random.seed(seed)
+    torch.manual_seed(seed+100)
+
+    if mode == 'Text-to-Image':
+        if (prompt is None) or (prompt == ""):
+            return None, None
+        with torch.no_grad():
+            rv = vd_inference.inference(
+                xtype = 'image',
+                cin = prompt,
+                ctype = 'prompt', )
+        return rv, None
+    elif mode == 'Image-Variation':
+        if image is None:
+            return None, None
+        with torch.no_grad():
+            rv = vd_inference.inference(
+                xtype = 'image',
+                cin = image,
+                ctype = 'vision',
+                color_adj = color_adj,)
+        return rv, None
+    elif mode == 'Image-to-Text':
+        if image is None:
+            return None, None
+        with torch.no_grad():
+            rv = vd_inference.inference(
+                xtype = 'text',
+                cin = image,
+                ctype = 'vision',)
+        return None, '\n'.join(rv)
+    elif mode == 'Text-Variation':
+        if prompt is None:
+            return None, None
+        with torch.no_grad():
+            rv = vd_inference.inference(
+                xtype = 'text',
+                cin = prompt,
+                ctype = 'prompt',)
+        return None, '\n'.join(rv)
+    elif mode == 'Disentanglement':
+        if image is None:
+            return None, None
+        with torch.no_grad():
+            rv = vd_inference.application_disensemble(
+                cin = image,
+                level = disentanglement_level,
+                color_adj = color_adj,)
+        return rv, None
+    elif mode == 'Dual-Guided':
+        if (image is None) or (prompt is None) or (prompt==""):
+            return None, None
+        with torch.no_grad():
+            rv = vd_inference.application_dualguided(
+                cim = image,
+                ctx = prompt,
+                mixing = dual_guided_mixing,
+                color_adj = color_adj,)
+        return rv, None
+    elif mode == 'Latent-I2T2I':
+        if (image is None) or (nprompt is None) or (nprompt=="") \
+                or (pprompt is None) or (pprompt==""):
+            return None, None
+        with torch.no_grad():
+            rv = vd_inference.application_i2t2i(
+                cim = image,
+                ctx_n = nprompt,
+                ctx_p = pprompt,
+                color_adj = color_adj,)
+        return rv, None
+    else:
+        assert False, "No such mode!"
+
+def get_instruction(mode):
+    t2i_instruction = ["Generate image from text prompt."]
+    i2i_instruction = [
+        "Generate image conditioned on reference image.", 
+        "Color Calibration provide an opinion to adjust image color according to reference image.", ]
+    i2t_instruction = ["Generate text from reference image."]
+    t2t_instruction = ["Generate text from reference text prompt. (Model insufficiently trained, thus results are still experimental)"]
+    dis_instruction = [
+        "Generate a variation of reference image that disentangled for semantic or style.",
+        "Color Calibration provide an opinion to adjust image color according to reference image.",
+        "Disentanglement level controls the level of focus towards semantic (-2, -1) or style (1 2). Level 0 serves as Image-Variation.", ]
+    dug_instruction = [
+        "Generate image from dual guidance of reference image and text prompt.",
+        "Color Calibration provide an opinion to adjust image color according to reference image.",
+        "Guidance Mixing provides linear balances between image and text context. (0 towards image, 1 towards text)", ]
+    iti_instruction = [
+        "Generate image variations via image-to-text, text-latent-editing, and then text-to-image. (Still under exploration)",
+        "Color Calibration provide an opinion to adjust image color according to reference image.",
+        "Input prompt that will be substract from text/text latent code.",
+        "Input prompt that will be added to text/text latent code.", ]
+
+    if mode == "Text-to-Image":
+        return '\n'.join(t2i_instruction)
+    elif mode == "Image-Variation":
+        return '\n'.join(i2i_instruction)
+    elif mode == "Image-to-Text":
+        return '\n'.join(i2t_instruction)
+    elif mode == "Text-Variation":
+        return '\n'.join(t2t_instruction)
+    elif mode == "Disentanglement":
+        return '\n'.join(dis_instruction)
+    elif mode == "Dual-Guided":
+        return '\n'.join(dug_instruction)
+    elif mode == "Latent-I2T2I":
+        return '\n'.join(iti_instruction)
+
+#############
+# Interface #
+#############
+
+if True:
+    img_output = gr.Gallery(label="Image Result").style(grid=n_sample_image_default)
+    txt_output = gr.Textbox(lines=4, label='Text Result', visible=False)
+
+    with gr.Blocks() as demo:
+        gr.HTML(
+            """
+            <div style="position: relative; float: left; text-align: center; width: 60%; min-width:600px; height: 160px; margin: 20px 0 20px 20%;">
+              <h1 style="font-weight: 900; font-size: 3rem;">
+                Versatile Diffusion
+              </h1>
+              <br>
+              <h2 style="font-weight: 450; font-size: 1rem;">
+                We built <b>Versatile Diffusion (VD), the first unified multi-flow multimodal diffusion framework</b>, as a step towards <b>Universal Generative AI</b>. 
+                VD can natively support image-to-text, image-variation, text-to-image, and text-variation, 
+                and can be further extended to other applications such as 
+                semantic-style disentanglement, image-text dual-guided generation, latent image-to-text-to-image editing, and more. 
+                Future versions will support more modalities such as speech, music, video and 3D. 
+              </h2>
+              <br>
+              <h3>Xingqian Xu, Atlas Wang, Eric Zhang, Kai Wang, 
+                and <a href="https://www.humphreyshi.com/home">Humphrey Shi</a> 
+                [<a href="https://arxiv.org/abs/2211.08332" style="color:blue;">arXiv</a>] 
+                [<a href="https://github.com/SHI-Labs/Versatile-Diffusion" style="color:blue;">GitHub</a>]
+              </h3>
+            </div>
+            <div style="position: relative; float: right; width: 19.9%; min-width:200px; margin: 20px auto;">
+              <img src="https://huggingface.co/spaces/shi-labs/Versatile-Diffusion/resolve/main/assets/figures/share_instruction.png">
+            </div>
+            """)
+        mode_input = gr.Radio([
+            "Text-to-Image", "Image-Variation", "Image-to-Text", "Text-Variation",
+            "Disentanglement", "Dual-Guided", "Latent-I2T2I"], value='Text-to-Image', label="VD Flows and Applications")
+
+        instruction = gr.Textbox(get_instruction("Text-to-Image"), label='Info')
+
+        with gr.Row():
+            with gr.Column():
+                img_input = gr.Image(label='Image Input', visible=False)
+                txt_input = gr.Textbox(lines=4, placeholder="Input prompt...", label='Text Input')
+                ntxt_input = gr.Textbox(label='Remove Prompt', visible=False)
+                ptxt_input = gr.Textbox(label='Add Prompt', visible=False)
+                coladj_input = gr.Radio(["None", "Simple"], value='Simple', label="Color Calibration", visible=False)
+                dislvl_input = gr.Slider(-2, 2, value=0, step=1, label="Disentanglement level", visible=False)
+                dguide_input = gr.Slider(0, 1, value=0.5, step=0.01, label="Guidance Mixing", visible=False)
+                seed_input = gr.Number(100, label="Seed", precision=0)
+
+                btn = gr.Button("Run")
+                btn.click(
+                    main, 
+                    inputs=[
+                        mode_input,
+                        img_input,
+                        txt_input,
+                        ntxt_input,
+                        ptxt_input,
+                        coladj_input,
+                        dislvl_input,
+                        dguide_input,
+                        seed_input, ],
+                    outputs=[img_output, txt_output])
+
+            with gr.Column():
+                img_output.render()
+                txt_output.render()
+
+        example_mode = [
+            "Text-to-Image", 
+            "Image-Variation", 
+            "Image-to-Text", 
+            "Text-Variation", 
+            "Disentanglement", 
+            "Dual-Guided", 
+            "Latent-I2T2I"]
+
+        def get_example(mode):
+            if mode == 'Text-to-Image':
+                case = [
+                    ['a dream of a village in china, by Caspar David Friedrich, matte painting trending on artstation HQ', 23],
+                    ['a beautiful grand nebula in the universe', 24],
+                    ['heavy arms gundam penguin mech', 25],
+                ]
+            elif mode == "Image-Variation":
+                case = [
+                    ['assets/space.jpg', 'None', 26],
+                    ['assets/train.jpg', 'Simple', 27],
+                ]
+            elif mode == "Image-to-Text":
+                case = [
+                    ['assets/boy_and_girl.jpg' , 28],
+                    ['assets/house_by_lake.jpg', 29],
+                ]
+            elif mode == "Text-Variation":
+                case = [
+                    ['a dream of a village in china, by Caspar David Friedrich, matte painting trending on artstation HQ' , 32],
+                    ['a beautiful grand nebula in the universe' , 33],
+                    ['heavy arms gundam penguin mech', 34],
+                ]
+            elif mode == "Disentanglement":
+                case = [
+                    ['assets/vermeer.jpg', 'Simple', -2, 30],
+                    ['assets/matisse.jpg', 'Simple',  2, 31],
+                ]
+            elif mode == "Dual-Guided":
+                case = [
+                    ['assets/benz.jpg',    'cyberpunk 2077', 'Simple', 0.75, 22],
+                    ['assets/vermeer.jpg', 'a girl with a diamond necklace',  'Simple', 0.66, 21],
+                ]
+            elif mode == "Latent-I2T2I":
+                case = [
+                    ['assets/ghibli.jpg',  'white house', 'tall castle', 'Simple', 20],
+                    ['assets/matisse.jpg', 'fruits and bottles on the table', 'flowers on the table', 'Simple', 21],
+                ]
+            else:
+                raise ValueError
+            case = [[mode] + casei for casei in case]
+            return case
+
+        def get_example_iof(mode):
+            if mode == 'Text-to-Image':
+                inps = [txt_input, seed_input]
+                oups = [img_output]
+                fn = lambda m, x, y: \
+                    main(mode=m, prompt=x, seed=y)[0]
+            elif mode == "Image-Variation":
+                inps = [img_input, coladj_input, seed_input]
+                oups = [img_output]
+                fn = lambda m, x, y, z: \
+                    main(mode=m, image=x, color_adj=y, seed=z)[0]
+            elif mode == "Image-to-Text":
+                inps = [img_input, seed_input]
+                oups = [txt_output]
+                fn = lambda m, x, y: \
+                    main(mode=m, image=x, seed=y)[1]
+            elif mode == "Text-Variation":
+                inps = [txt_input, seed_input]
+                oups = [txt_output]
+                fn = lambda m, x, y: \
+                    main(mode=m, prompt=x, seed=y)[1]
+            elif mode == "Disentanglement":
+                inps = [img_input, coladj_input, dislvl_input, seed_input]
+                oups = [img_output]
+                fn = lambda m, x, y, z, w: \
+                    main(mode=m, image=x, color_adj=y, disentanglement_level=z, seed=w)[0]
+            elif mode == "Dual-Guided":
+                inps = [img_input, txt_input, coladj_input, dguide_input, seed_input]
+                oups = [img_output]
+                fn = lambda m, x, y, z, w, u: \
+                    main(mode=m, image=x, prompt=y, color_adj=z, dual_guided_mixing=w, seed=u)[0]
+            elif mode == "Latent-I2T2I":
+                inps = [img_input, ntxt_input, ptxt_input, coladj_input, seed_input]
+                oups = [img_output]
+                fn = lambda m, x, y, z, w, u: \
+                    main(mode=m, image=x, nprompt=y, pprompt=z, color_adj=w, seed=u)[0]
+            else:
+                raise ValueError
+            return [mode_input]+inps, oups, fn
+
+        with gr.Row():
+            for emode in example_mode[0:4]:
+                with gr.Column():
+                    gr.Examples(
+                        label=emode+' Examples',
+                        examples=get_example(emode),
+                        inputs=get_example_iof(emode)[0],
+                        outputs=get_example_iof(emode)[1],
+                        fn = get_example_iof(emode)[2],
+                        cache_examples=cache_examples),
+        with gr.Row():
+            for emode in example_mode[4:7]:
+                with gr.Column():
+                    gr.Examples(
+                        label=emode+' Examples',
+                        examples=get_example(emode),
+                        inputs=get_example_iof(emode)[0],
+                        outputs=get_example_iof(emode)[1],
+                        fn = get_example_iof(emode)[2],
+                        cache_examples=cache_examples),
+
+        mode_input.change(
+            fn=lambda x: gr.update(value=get_instruction(x)),
+            inputs=mode_input,
+            outputs=instruction,)
+
+        mode_input.change(
+            fn=lambda x: gr.update(visible=(x not in ['Text-to-Image', 'Text-Variation'])),
+            inputs=mode_input,
+            outputs=img_input,)
+
+        mode_input.change(
+            fn=lambda x: gr.update(visible=(x in ['Text-to-Image', 'Text-Variation', 'Dual-Guided'])),
+            inputs=mode_input,
+            outputs=txt_input,)
+
+        mode_input.change(
+            fn=lambda x: gr.update(visible=(x in ['Latent-I2T2I'])),
+            inputs=mode_input,
+            outputs=ntxt_input,)
+        mode_input.change(
+            fn=lambda x: gr.update(visible=(x in ['Latent-I2T2I'])),
+            inputs=mode_input,
+            outputs=ptxt_input,)
+
+        mode_input.change(
+            fn=lambda x: gr.update(visible=(x not in ['Text-to-Image', 'Image-to-Text', 'Text-Variation'])),
+            inputs=mode_input,
+            outputs=coladj_input,)
+
+        mode_input.change(
+            fn=lambda x: gr.update(visible=(x=='Disentanglement')),
+            inputs=mode_input,
+            outputs=dislvl_input,)
+
+        mode_input.change(
+            fn=lambda x: gr.update(visible=(x=='Dual-Guided')),
+            inputs=mode_input,
+            outputs=dguide_input,)
+
+        mode_input.change(
+            fn=lambda x: gr.update(visible=(x not in ['Image-to-Text', 'Text-Variation'])),
+            inputs=mode_input,
+            outputs=img_output,)
+        mode_input.change(
+            fn=lambda x: gr.update(visible=(x in ['Image-to-Text', 'Text-Variation'])),
+            inputs=mode_input,
+            outputs=txt_output,)
+
+        gr.HTML(
+            """
+            <div style="text-align: center; max-width: 1200px; margin: 20px auto;">
+            <h3>
+            <b>Caution</b>: 
+            We would like the raise the awareness of users of this demo of its potential issues and concerns.
+            Like previous large foundation models, Versatile Diffusion could be problematic in some cases, partially due to the imperfect training data and pretrained network (VAEs / context encoders) with limited scope.
+            In its future research phase, VD may do better on tasks such as text-to-image, image-to-text, etc., with the help of more powerful VAEs, more sophisticated network designs, and more cleaned data.
+            So far, we keep all features available for research testing both to show the great potential of the VD framework and to collect important feedback to improve the model in the future.
+            We welcome researchers and users to report issues with the HuggingFace community discussion feature or email the authors.
+            </h3>
+            <br>
+            <h3>
+            <b>Biases and content acknowledgement</b>:
+            Beware that VD may output content that reinforces or exacerbates societal biases, as well as realistic faces, pornography, and violence. 
+            VD was trained on the LAION-2B dataset, which scraped non-curated online images and text, and may contained unintended exceptions as we removed illegal content. 
+            VD in this demo is meant only for research purposes.
+            </h3>
+            </div>
+            """)
+
+    # demo.launch(share=True)
+    demo.launch(debug=True)

configs/model/clip.yaml

@@ -0,0 +1,50 @@
+
+clip:
+  symbol: clip
+  args: {}
+
+clip_frozen:
+  super_cfg: clip
+  type: clip_frozen
+  args: {}
+
+clip_text_frozen:
+  super_cfg: clip
+  type: clip_text_frozen
+  args: {}
+
+clip_vision_frozen:
+  super_cfg: clip
+  type: clip_vision_frozen
+  args: {}
+
+############################
+# clip with focused encode #
+############################
+
+clip_frozen_encode_text:
+  super_cfg: clip
+  type: clip_frozen
+  args:
+    encode_type : encode_text
+
+clip_frozen_encode_vision:
+  super_cfg: clip
+  type: clip_frozen
+  args:
+    encode_type : encode_vision
+
+clip_frozen_encode_text_noproj:
+  super_cfg: clip
+  type: clip_frozen
+  args:
+    encode_type : encode_text_noproj
+
+#####################################
+# clip vision forzen justin version #
+#####################################
+
+clip_vision_frozen_justin:
+  super_cfg: clip
+  type: clip_vision_frozen_justin
+  args: {}

configs/model/openai_unet.yaml

@@ -0,0 +1,72 @@
+openai_unet_sd:
+  type: openai_unet
+  args:
+    image_size: null # no use
+    in_channels: 4
+    out_channels: 4
+    model_channels: 320
+    attention_resolutions: [ 4, 2, 1 ]
+    num_res_blocks: [ 2, 2, 2, 2 ]
+    channel_mult: [ 1, 2, 4, 4 ]
+    # disable_self_attentions: [ False, False, False, False ]  # converts the self-attention to a cross-attention layer if true
+    num_heads: 8
+    use_spatial_transformer: True
+    transformer_depth: 1
+    context_dim: 768
+    use_checkpoint: True
+    legacy: False
+
+openai_unet_dual_context:
+  super_cfg: openai_unet_sd
+  type: openai_unet_dual_context
+
+########################
+# Code cleaned version #
+########################
+
+openai_unet_2d:
+  type: openai_unet_2d
+  args:
+    input_channels: 4
+    model_channels: 320
+    output_channels: 4
+    num_noattn_blocks: [ 2, 2, 2, 2 ]
+    channel_mult: [ 1, 2, 4, 4 ]
+    with_attn: [true, true, true, false]
+    num_heads: 8
+    context_dim: 768
+    use_checkpoint: True
+
+openai_unet_0d:
+  type: openai_unet_0d
+  args:
+    input_channels: 768
+    model_channels: 320
+    output_channels: 768
+    num_noattn_blocks: [ 2, 2, 2, 2 ]
+    channel_mult: [ 1, 2, 4, 4 ]
+    with_attn: [true, true, true, false]
+    num_heads: 8
+    context_dim: 768
+    use_checkpoint: True
+
+openai_unet_0dmd:
+  type: openai_unet_0dmd
+  args:
+    input_channels: 768
+    model_channels: 320
+    output_channels: 768
+    num_noattn_blocks: [ 2, 2, 2, 2 ]
+    channel_mult: [ 1, 2, 4, 4 ]
+    second_dim: [ 4, 4, 4, 4 ]
+    with_attn: [true, true, true, false]
+    num_heads: 8
+    context_dim: 768
+    use_checkpoint: True
+
+openai_unet_vd:
+  type: openai_unet_vd
+  args:
+    unet_image_cfg: MODEL(openai_unet_2d)
+    unet_text_cfg: MODEL(openai_unet_0dmd)
+  

configs/model/optimus.yaml

@@ -0,0 +1,103 @@
+
+optimus:
+  symbol: optimus
+  find_unused_parameters: false
+  args: {}
+
+optimus_bert_encoder:
+  super_cfg: optimus
+  type: optimus_bert_connector
+  # pth: pretrained/optimus_bert_encoder.pth
+  args:
+    config:
+      architectures: 
+        - BertForMaskedLM
+      attention_probs_dropout_prob: 0.1
+      finetuning_task: null
+      hidden_act: gelu
+      hidden_dropout_prob: 0.1
+      hidden_size: 768
+      initializer_range: 0.02
+      intermediate_size: 3072
+      layer_norm_eps: 1.e-12
+      max_position_embeddings: 512
+      num_attention_heads: 12
+      num_hidden_layers: 12
+      num_labels: 2
+      output_attentions: false
+      output_hidden_states: false
+      pruned_heads: {}
+      torchscript: false
+      type_vocab_size: 2
+      vocab_size: 28996
+    latent_size: 768
+
+optimus_bert_tokenizer:
+  super_cfg: optimus
+  type: optimus_bert_tokenizer
+  args:
+    do_lower_case: false 
+    max_len: 512
+    vocab_file: lib/model_zoo/optimus_models/vocab/bert-base-cased-vocab.txt
+
+optimus_gpt2_decoder:
+  super_cfg: optimus
+  type: optimus_gpt2_connector
+  # pth: pretrained/optimus_gpt2_decoder.pth
+  args:
+    config:
+      architectures:
+        - GPT2LMHeadModel
+      attn_pdrop: 0.1
+      embd_pdrop: 0.1
+      finetuning_task: null
+      hidden_size: 768
+      initializer_range: 0.02
+      latent_size: 768
+      layer_norm_epsilon: 1.e-05
+      max_position_embeddings: 1024
+      n_ctx: 1024
+      n_embd: 768
+      n_head: 12
+      n_layer: 12
+      n_positions: 1024
+      num_attention_heads: 12
+      num_hidden_layers: 12
+      num_labels: 1
+      output_attentions: false
+      output_hidden_states: false
+      pretrained_config_archive_map:
+        gpt2        : https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-config.json
+        gpt2-medium : https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-medium-config.json
+        gpt2-large  : https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-large-config.json
+      pruned_heads: {}
+      resid_pdrop: 0.1
+      summary_activation: null
+      summary_first_dropout: 0.1
+      summary_proj_to_labels: true
+      summary_type: cls_index
+      summary_use_proj: true
+      torchscript: false
+      vocab_size: 50260
+
+optimus_gpt2_tokenizer:
+  super_cfg: optimus
+  type: optimus_gpt2_tokenizer
+  args:
+    do_lower_case: false 
+    max_len: 1024
+    vocab_file: lib/model_zoo/optimus_models/vocab/gpt2-vocab.json
+    merges_file: lib/model_zoo/optimus_models/vocab/gpt2-merges.txt
+
+optimus_vae:
+  super_cfg: optimus
+  type: optimus_vae
+  args:
+    encoder: MODEL(optimus_bert_encoder)
+    decoder: MODEL(optimus_gpt2_decoder)
+    tokenizer_encoder: MODEL(optimus_bert_tokenizer)
+    tokenizer_decoder: MODEL(optimus_gpt2_tokenizer)
+    args:
+      latent_size: 768
+  hfm: ['shi-labs/versatile-diffusion-model', 'pretrained_pth/optimus-vae.pth']
+  # pth: pretrained/optimus-vae.pth

configs/model/sd.yaml

@@ -0,0 +1,69 @@
+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
+  hfm: ['shi-labs/versatile-diffusion-model', 'pretrained_pth/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)
+ 

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

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

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)

lib/log_service.py

@@ -0,0 +1,166 @@
+import timeit
+import numpy as np
+import os
+import os.path as osp
+import shutil
+import copy
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+from .cfg_holder import cfg_unique_holder as cfguh
+from . import sync
+
+print_console_local_rank0_only = True
+
+def print_log(*console_info):
+    local_rank = sync.get_rank('local')
+    if print_console_local_rank0_only and (local_rank!=0):
+        return
+    console_info = [str(i) for i in console_info]
+    console_info = ' '.join(console_info)
+    print(console_info)
+
+    if local_rank!=0:
+        return
+
+    log_file = None
+    try:
+        log_file = cfguh().cfg.train.log_file
+    except:
+        try:
+            log_file = cfguh().cfg.eval.log_file
+        except:
+            return
+    if log_file is not None:
+        with open(log_file, 'a') as f:
+            f.write(console_info + '\n')
+
+class distributed_log_manager(object):
+    def __init__(self):
+        self.sum = {}
+        self.cnt = {}
+        self.time_check = timeit.default_timer()
+
+        cfgt = cfguh().cfg.train
+        use_tensorboard = getattr(cfgt, 'log_tensorboard', False)
+
+        self.ddp = sync.is_ddp()
+        self.rank = sync.get_rank('local')
+        self.world_size = sync.get_world_size('local')
+
+        self.tb = None
+        if use_tensorboard and (self.rank==0):
+            import tensorboardX
+            monitoring_dir = osp.join(cfguh().cfg.train.log_dir, 'tensorboard')
+            self.tb = tensorboardX.SummaryWriter(osp.join(monitoring_dir))
+
+    def accumulate(self, n, **data):
+        if n < 0:
+            raise ValueError
+
+        for itemn, di in data.items():
+            if itemn in self.sum:
+                self.sum[itemn] += di * n
+                self.cnt[itemn] += n
+            else:
+                self.sum[itemn] = di * n
+                self.cnt[itemn] = n
+
+    def get_mean_value_dict(self):
+        value_gather = [
+            self.sum[itemn]/self.cnt[itemn] \
+                for itemn in sorted(self.sum.keys()) ]
+
+        value_gather_tensor = torch.FloatTensor(value_gather).to(self.rank)
+        if self.ddp:
+            dist.all_reduce(value_gather_tensor, op=dist.ReduceOp.SUM)
+            value_gather_tensor /= self.world_size
+
+        mean = {}
+        for idx, itemn in enumerate(sorted(self.sum.keys())):
+            mean[itemn] = value_gather_tensor[idx].item()
+        return mean
+
+    def tensorboard_log(self, step, data, mode='train', **extra):
+        if self.tb is None:
+            return
+        if mode == 'train':
+            self.tb.add_scalar('other/epochn', extra['epochn'], step)
+            if 'lr' in extra:
+                self.tb.add_scalar('other/lr', extra['lr'], step)
+            for itemn, di in data.items():
+                if itemn.find('loss') == 0:
+                    self.tb.add_scalar('loss/'+itemn,  di, step)
+                elif itemn == 'Loss':
+                    self.tb.add_scalar('Loss',  di, step)
+                else:
+                    self.tb.add_scalar('other/'+itemn, di, step)
+        elif mode == 'eval':
+            if isinstance(data, dict):
+                for itemn, di in data.items():
+                    self.tb.add_scalar('eval/'+itemn, di, step)
+            else:
+                self.tb.add_scalar('eval', data, step)
+        return
+
+    def train_summary(self, itern, epochn, samplen, lr, tbstep=None):
+        console_info = [
+            'Iter:{}'.format(itern),
+            'Epoch:{}'.format(epochn),
+            'Sample:{}'.format(samplen),]
+
+        if lr is not None:
+            console_info += ['LR:{:.4E}'.format(lr)]
+
+        mean = self.get_mean_value_dict()
+
+        tbstep = itern if tbstep is None else tbstep
+        self.tensorboard_log(
+            tbstep, mean, mode='train',
+            itern=itern, epochn=epochn, lr=lr)
+
+        loss = mean.pop('Loss')
+        mean_info = ['Loss:{:.4f}'.format(loss)] + [
+            '{}:{:.4f}'.format(itemn, mean[itemn]) \
+                for itemn in sorted(mean.keys()) \
+                    if itemn.find('loss') == 0
+        ]
+        console_info += mean_info
+        console_info.append('Time:{:.2f}s'.format(
+            timeit.default_timer() - self.time_check))
+        return ' , '.join(console_info)
+
+    def clear(self):
+        self.sum = {}
+        self.cnt = {}
+        self.time_check = timeit.default_timer()
+
+    def tensorboard_close(self):
+        if self.tb is not None:
+            self.tb.close()
+
+# ----- also include some small utils -----
+
+def torch_to_numpy(*argv):
+    if len(argv) > 1:
+        data = list(argv)
+    else:
+        data = argv[0]
+
+    if isinstance(data, torch.Tensor):
+        return data.to('cpu').detach().numpy()
+
+    elif isinstance(data, (list, tuple)):
+        out = []
+        for di in data:
+            out.append(torch_to_numpy(di))
+        return out
+
+    elif isinstance(data, dict):
+        out = {}
+        for ni, di in data.items():
+            out[ni] = torch_to_numpy(di)
+        return out
+    
+    else:
+        return data

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

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

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

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)

lib/model_zoo/clip.py

@@ -0,0 +1,178 @@
+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, CLIPModel
+
+@register('clip_frozen', version)
+class FrozenCLIP(AbstractEncoder):
+    def __init__(self, 
+                 version="openai/clip-vit-large-patch14", 
+                 max_length=77, 
+                 encode_type='encode_text',
+                 fp16=False, ):
+        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.fp16 = fp16
+        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())
+        outputs = self.model.get_text_features(input_ids=tokens)
+        return outputs
+
+    def encode_vision_pooled(self, images):
+        inputs = self.processor(images=images, return_tensors="pt")
+        pixels = inputs['pixel_values'].half() if self.fp16 else inputs['pixel_values']
+        pixels = pixels.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'].half() if self.fp16 else inputs['pixel_values']
+        pixels = pixels.to(self.get_device())
+        outputs = self.model.vision_model(pixel_values=pixels)
+        return outputs.last_hidden_state
+
+    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(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)

lib/model_zoo/clip_justin/__init__.py

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

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

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()

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

lib/model_zoo/common/get_model.py

@@ -0,0 +1,128 @@
+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))
+        elif 'hfm' in cfg:
+            from huggingface_hub import hf_hub_download
+            temppath = hf_hub_download(cfg.hfm[0], cfg.hfm[1])
+            sd = torch.load(temppath, 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.hfm))
+
+        # 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

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)

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

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 

lib/model_zoo/ddim.py

@@ -0,0 +1,216 @@
+"""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 = self.model.betas.device
+        b = shape[0]
+        if x_T is None:
+            img = torch.randn(shape, device=device)
+        else:
+            img = x_T
+
+        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, x.device
+
+        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)
+
+        # 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, 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

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, 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

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
+
+

lib/model_zoo/ddim_vd.py

@@ -0,0 +1,290 @@
+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 = self.model.device
+        bs = shape[0]
+        if xt is None:
+            xt = torch.randn(shape, device=device, dtype=conditioning.dtype)
+
+        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, dtype=x.dtype)
+        a_prev = torch.full(extended_shape, alphas_prev[index], device=device, dtype=x.dtype)
+        sigma_t = torch.full(extended_shape, sigmas[index], device=device, dtype=x.dtype)
+        sqrt_one_minus_at = torch.full(extended_shape, sqrt_one_minus_alphas[index], device=device, dtype=x.dtype)
+
+        # 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, 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_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, dtype=first_conditioning[1].dtype)
+
+        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 = 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, dtype=x.dtype)
+        a_prev = torch.full(extended_shape, alphas_prev[index], device=device, dtype=x.dtype)
+        sigma_t = torch.full(extended_shape, sigmas[index], device=device, dtype=x.dtype)
+        sqrt_one_minus_at = torch.full(extended_shape, sqrt_one_minus_alphas[index], device=device, dtype=x.dtype)
+
+        # 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, 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

lib/model_zoo/diffusion_modules.py

@@ -0,0 +1,835 @@
+# pytorch_diffusion + derived encoder decoder
+import math
+import torch
+import torch.nn as nn
+import numpy as np
+from einops import rearrange
+
+# from .diffusion_utils import instantiate_from_config
+from .attention import LinearAttention
+
+
+def get_timestep_embedding(timesteps, embedding_dim):
+    """
+    This matches the implementation in Denoising Diffusion Probabilistic Models:
+    From Fairseq.
+    Build sinusoidal embeddings.
+    This matches the implementation in tensor2tensor, but differs slightly
+    from the description in Section 3.5 of "Attention Is All You Need".
+    """
+    assert len(timesteps.shape) == 1
+
+    half_dim = embedding_dim // 2
+    emb = math.log(10000) / (half_dim - 1)
+    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
+    emb = emb.to(device=timesteps.device)
+    emb = timesteps.float()[:, None] * emb[None, :]
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+    if embedding_dim % 2 == 1:  # zero pad
+        emb = torch.nn.functional.pad(emb, (0,1,0,0))
+    return emb
+
+
+def nonlinearity(x):
+    # swish
+    return x*torch.sigmoid(x)
+
+
+def Normalize(in_channels, num_groups=32):
+    return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=0)
+
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0,1,0,1)
+            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
+                 dropout, temb_channels=512):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv2d(in_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels,
+                                             out_channels)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(out_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(in_channels,
+                                                     out_channels,
+                                                     kernel_size=3,
+                                                     stride=1,
+                                                     padding=1)
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(in_channels,
+                                                    out_channels,
+                                                    kernel_size=1,
+                                                    stride=1,
+                                                    padding=0)
+
+    def forward(self, x, temb):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None]
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x+h
+
+
+class LinAttnBlock(LinearAttention):
+    """to match AttnBlock usage"""
+    def __init__(self, in_channels):
+        super().__init__(dim=in_channels, heads=1, dim_head=in_channels)
+
+
+class AttnBlock(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 = q.reshape(b,c,h*w)
+        q = q.permute(0,2,1)   # b,hw,c
+        k = k.reshape(b,c,h*w) # b,c,hw
+        w_ = torch.bmm(q,k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c)**(-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b,c,h*w)
+        w_ = w_.permute(0,2,1)   # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v,w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b,c,h,w)
+
+        h_ = self.proj_out(h_)
+
+        return x+h_
+
+
+def make_attn(in_channels, attn_type="vanilla"):
+    assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown'
+    print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
+    if attn_type == "vanilla":
+        return AttnBlock(in_channels)
+    elif attn_type == "none":
+        return nn.Identity(in_channels)
+    else:
+        return LinAttnBlock(in_channels)
+
+
+class Model(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, use_timestep=True, use_linear_attn=False, attn_type="vanilla"):
+        super().__init__()
+        if use_linear_attn: attn_type = "linear"
+        self.ch = ch
+        self.temb_ch = self.ch*4
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        self.use_timestep = use_timestep
+        if self.use_timestep:
+            # timestep embedding
+            self.temb = nn.Module()
+            self.temb.dense = nn.ModuleList([
+                torch.nn.Linear(self.ch,
+                                self.temb_ch),
+                torch.nn.Linear(self.temb_ch,
+                                self.temb_ch),
+            ])
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(in_channels,
+                                       self.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,)+tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(make_attn(block_in, attn_type=attn_type))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions-1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            skip_in = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks+1):
+                if i_block == self.num_res_blocks:
+                    skip_in = ch*in_ch_mult[i_level]
+                block.append(ResnetBlock(in_channels=block_in+skip_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(make_attn(block_in, attn_type=attn_type))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up) # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x, t=None, context=None):
+        #assert x.shape[2] == x.shape[3] == self.resolution
+        if context is not None:
+            # assume aligned context, cat along channel axis
+            x = torch.cat((x, context), dim=1)
+        if self.use_timestep:
+            # timestep embedding
+            assert t is not None
+            temb = get_timestep_embedding(t, self.ch)
+            temb = self.temb.dense[0](temb)
+            temb = nonlinearity(temb)
+            temb = self.temb.dense[1](temb)
+        else:
+            temb = None
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions-1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](
+                    torch.cat([h, hs.pop()], dim=1), temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+    def get_last_layer(self):
+        return self.conv_out.weight
+
+
+class Encoder(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla",
+                 **ignore_kwargs):
+        super().__init__()
+        if use_linear_attn: attn_type = "linear"
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(in_channels,
+                                       self.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,)+tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(make_attn(block_in, attn_type=attn_type))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions-1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        2*z_channels if double_z else z_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        # timestep embedding
+        temb = None
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions-1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
+                 attn_type="vanilla", **ignorekwargs):
+        super().__init__()
+        if use_linear_attn: attn_type = "linear"
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.give_pre_end = give_pre_end
+        self.tanh_out = tanh_out
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        in_ch_mult = (1,)+tuple(ch_mult)
+        block_in = ch*ch_mult[self.num_resolutions-1]
+        curr_res = resolution // 2**(self.num_resolutions-1)
+        self.z_shape = (1,z_channels,curr_res,curr_res)
+        print("Working with z of shape {} = {} dimensions.".format(
+            self.z_shape, np.prod(self.z_shape)))
+
+        # z to block_in
+        self.conv_in = torch.nn.Conv2d(z_channels,
+                                       block_in,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks+1):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(make_attn(block_in, attn_type=attn_type))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up) # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, z):
+        #assert z.shape[1:] == self.z_shape[1:]
+        self.last_z_shape = z.shape
+
+        # timestep embedding
+        temb = None
+
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](h, temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        if self.give_pre_end:
+            return h
+
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        if self.tanh_out:
+            h = torch.tanh(h)
+        return h
+
+
+class SimpleDecoder(nn.Module):
+    def __init__(self, in_channels, out_channels, *args, **kwargs):
+        super().__init__()
+        self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
+                                     ResnetBlock(in_channels=in_channels,
+                                                 out_channels=2 * in_channels,
+                                                 temb_channels=0, dropout=0.0),
+                                     ResnetBlock(in_channels=2 * in_channels,
+                                                out_channels=4 * in_channels,
+                                                temb_channels=0, dropout=0.0),
+                                     ResnetBlock(in_channels=4 * in_channels,
+                                                out_channels=2 * in_channels,
+                                                temb_channels=0, dropout=0.0),
+                                     nn.Conv2d(2*in_channels, in_channels, 1),
+                                     Upsample(in_channels, with_conv=True)])
+        # end
+        self.norm_out = Normalize(in_channels)
+        self.conv_out = torch.nn.Conv2d(in_channels,
+                                        out_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        for i, layer in enumerate(self.model):
+            if i in [1,2,3]:
+                x = layer(x, None)
+            else:
+                x = layer(x)
+
+        h = self.norm_out(x)
+        h = nonlinearity(h)
+        x = self.conv_out(h)
+        return x
+
+
+class UpsampleDecoder(nn.Module):
+    def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution,
+                 ch_mult=(2,2), dropout=0.0):
+        super().__init__()
+        # upsampling
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        block_in = in_channels
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.res_blocks = nn.ModuleList()
+        self.upsample_blocks = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            res_block = []
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                res_block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+            self.res_blocks.append(nn.ModuleList(res_block))
+            if i_level != self.num_resolutions - 1:
+                self.upsample_blocks.append(Upsample(block_in, True))
+                curr_res = curr_res * 2
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        # upsampling
+        h = x
+        for k, i_level in enumerate(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.res_blocks[i_level][i_block](h, None)
+            if i_level != self.num_resolutions - 1:
+                h = self.upsample_blocks[k](h)
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class LatentRescaler(nn.Module):
+    def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2):
+        super().__init__()
+        # residual block, interpolate, residual block
+        self.factor = factor
+        self.conv_in = nn.Conv2d(in_channels,
+                                 mid_channels,
+                                 kernel_size=3,
+                                 stride=1,
+                                 padding=1)
+        self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
+                                                     out_channels=mid_channels,
+                                                     temb_channels=0,
+                                                     dropout=0.0) for _ in range(depth)])
+        self.attn = AttnBlock(mid_channels)
+        self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
+                                                     out_channels=mid_channels,
+                                                     temb_channels=0,
+                                                     dropout=0.0) for _ in range(depth)])
+
+        self.conv_out = nn.Conv2d(mid_channels,
+                                  out_channels,
+                                  kernel_size=1,
+                                  )
+
+    def forward(self, x):
+        x = self.conv_in(x)
+        for block in self.res_block1:
+            x = block(x, None)
+        x = torch.nn.functional.interpolate(x, size=(int(round(x.shape[2]*self.factor)), int(round(x.shape[3]*self.factor))))
+        x = self.attn(x)
+        for block in self.res_block2:
+            x = block(x, None)
+        x = self.conv_out(x)
+        return x
+
+
+class MergedRescaleEncoder(nn.Module):
+    def __init__(self, in_channels, ch, resolution, out_ch, num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True,
+                 ch_mult=(1,2,4,8), rescale_factor=1.0, rescale_module_depth=1):
+        super().__init__()
+        intermediate_chn = ch * ch_mult[-1]
+        self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult,
+                               z_channels=intermediate_chn, double_z=False, resolution=resolution,
+                               attn_resolutions=attn_resolutions, dropout=dropout, resamp_with_conv=resamp_with_conv,
+                               out_ch=None)
+        self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn,
+                                       mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth)
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.rescaler(x)
+        return x
+
+
+class MergedRescaleDecoder(nn.Module):
+    def __init__(self, z_channels, out_ch, resolution, num_res_blocks, attn_resolutions, ch, ch_mult=(1,2,4,8),
+                 dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1):
+        super().__init__()
+        tmp_chn = z_channels*ch_mult[-1]
+        self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_resolutions=attn_resolutions, dropout=dropout,
+                               resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks,
+                               ch_mult=ch_mult, resolution=resolution, ch=ch)
+        self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn,
+                                       out_channels=tmp_chn, depth=rescale_module_depth)
+
+    def forward(self, x):
+        x = self.rescaler(x)
+        x = self.decoder(x)
+        return x
+
+
+class Upsampler(nn.Module):
+    def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2):
+        super().__init__()
+        assert out_size >= in_size
+        num_blocks = int(np.log2(out_size//in_size))+1
+        factor_up = 1.+ (out_size % in_size)
+        print(f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}")
+        self.rescaler = LatentRescaler(factor=factor_up, in_channels=in_channels, mid_channels=2*in_channels,
+                                       out_channels=in_channels)
+        self.decoder = Decoder(out_ch=out_channels, resolution=out_size, z_channels=in_channels, num_res_blocks=2,
+                               attn_resolutions=[], in_channels=None, ch=in_channels,
+                               ch_mult=[ch_mult for _ in range(num_blocks)])
+
+    def forward(self, x):
+        x = self.rescaler(x)
+        x = self.decoder(x)
+        return x
+
+
+class Resize(nn.Module):
+    def __init__(self, in_channels=None, learned=False, mode="bilinear"):
+        super().__init__()
+        self.with_conv = learned
+        self.mode = mode
+        if self.with_conv:
+            print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode")
+            raise NotImplementedError()
+            assert in_channels is not None
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=4,
+                                        stride=2,
+                                        padding=1)
+
+    def forward(self, x, scale_factor=1.0):
+        if scale_factor==1.0:
+            return x
+        else:
+            x = torch.nn.functional.interpolate(x, mode=self.mode, align_corners=False, scale_factor=scale_factor)
+        return x
+
+class FirstStagePostProcessor(nn.Module):
+
+    def __init__(self, ch_mult:list, in_channels,
+                 pretrained_model:nn.Module=None,
+                 reshape=False,
+                 n_channels=None,
+                 dropout=0.,
+                 pretrained_config=None):
+        super().__init__()
+        if pretrained_config is None:
+            assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
+            self.pretrained_model = pretrained_model
+        else:
+            assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
+            self.instantiate_pretrained(pretrained_config)
+
+        self.do_reshape = reshape
+
+        if n_channels is None:
+            n_channels = self.pretrained_model.encoder.ch
+
+        self.proj_norm = Normalize(in_channels,num_groups=in_channels//2)
+        self.proj = nn.Conv2d(in_channels,n_channels,kernel_size=3,
+                            stride=1,padding=1)
+
+        blocks = []
+        downs = []
+        ch_in = n_channels
+        for m in ch_mult:
+            blocks.append(ResnetBlock(in_channels=ch_in,out_channels=m*n_channels,dropout=dropout))
+            ch_in = m * n_channels
+            downs.append(Downsample(ch_in, with_conv=False))
+
+        self.model = nn.ModuleList(blocks)
+        self.downsampler = nn.ModuleList(downs)
+
+
+    def instantiate_pretrained(self, config):
+        model = instantiate_from_config(config)
+        self.pretrained_model = model.eval()
+        # self.pretrained_model.train = False
+        for param in self.pretrained_model.parameters():
+            param.requires_grad = False
+
+
+    @torch.no_grad()
+    def encode_with_pretrained(self,x):
+        c = self.pretrained_model.encode(x)
+        if isinstance(c, DiagonalGaussianDistribution):
+            c = c.mode()
+        return  c
+
+    def forward(self,x):
+        z_fs = self.encode_with_pretrained(x)
+        z = self.proj_norm(z_fs)
+        z = self.proj(z)
+        z = nonlinearity(z)
+
+        for submodel, downmodel in zip(self.model,self.downsampler):
+            z = submodel(z,temb=None)
+            z = downmodel(z)
+
+        if self.do_reshape:
+            z = rearrange(z,'b c h w -> b (h w) c')
+        return z
+

lib/model_zoo/diffusion_utils.py

@@ -0,0 +1,250 @@
+import os
+import math
+import torch
+import torch.nn as nn
+import numpy as np
+from einops import repeat
+
+def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+    if schedule == "linear":
+        betas = (
+                torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2
+        )
+
+    elif schedule == "cosine":
+        timesteps = (
+                torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s
+        )
+        alphas = timesteps / (1 + cosine_s) * np.pi / 2
+        alphas = torch.cos(alphas).pow(2)
+        alphas = alphas / alphas[0]
+        betas = 1 - alphas[1:] / alphas[:-1]
+        betas = np.clip(betas, a_min=0, a_max=0.999)
+
+    elif schedule == "sqrt_linear":
+        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
+    elif schedule == "sqrt":
+        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5
+    else:
+        raise ValueError(f"schedule '{schedule}' unknown.")
+    return betas.numpy()
+
+def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):
+    if ddim_discr_method == 'uniform':
+        c = num_ddpm_timesteps // num_ddim_timesteps
+        ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
+    elif ddim_discr_method == 'quad':
+        ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
+    else:
+        raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
+
+    # assert ddim_timesteps.shape[0] == num_ddim_timesteps
+    # add one to get the final alpha values right (the ones from first scale to data during sampling)
+    steps_out = ddim_timesteps + 1
+    if verbose:
+        print(f'Selected timesteps for ddim sampler: {steps_out}')
+    return steps_out
+
+def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
+    # select alphas for computing the variance schedule
+    alphas = alphacums[ddim_timesteps]
+    alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
+
+    # according the the formula provided in https://arxiv.org/abs/2010.02502
+    sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
+    if verbose:
+        print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')
+        print(f'For the chosen value of eta, which is {eta}, '
+              f'this results in the following sigma_t schedule for ddim sampler {sigmas}')
+    return sigmas, alphas, alphas_prev
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+def extract_into_tensor(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+def checkpoint(func, inputs, params, flag):
+    """
+    Evaluate a function without caching intermediate activations, allowing for
+    reduced memory at the expense of extra compute in the backward pass.
+    :param func: the function to evaluate.
+    :param inputs: the argument sequence to pass to `func`.
+    :param params: a sequence of parameters `func` depends on but does not
+                   explicitly take as arguments.
+    :param flag: if False, disable gradient checkpointing.
+    """
+    if flag:
+        args = tuple(inputs) + tuple(params)
+        return CheckpointFunction.apply(func, len(inputs), *args)
+    else:
+        return func(*inputs)
+
+class CheckpointFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, run_function, length, *args):
+        ctx.run_function = run_function
+        ctx.input_tensors = list(args[:length])
+        ctx.input_params = list(args[length:])
+
+        with torch.no_grad():
+            output_tensors = ctx.run_function(*ctx.input_tensors)
+        return output_tensors
+
+    @staticmethod
+    def backward(ctx, *output_grads):
+        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
+        with torch.enable_grad():
+            # Fixes a bug where the first op in run_function modifies the
+            # Tensor storage in place, which is not allowed for detach()'d
+            # Tensors.
+            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
+            output_tensors = ctx.run_function(*shallow_copies)
+        input_grads = torch.autograd.grad(
+            output_tensors,
+            ctx.input_tensors + ctx.input_params,
+            output_grads,
+            allow_unused=True,
+        )
+        del ctx.input_tensors
+        del ctx.input_params
+        del output_tensors
+        return (None, None) + input_grads
+
+def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+    """
+    Create sinusoidal timestep embeddings.
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    if not repeat_only:
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=timesteps.device)
+        args = timesteps[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    else:
+        embedding = repeat(timesteps, 'b -> b d', d=dim)
+    return embedding
+
+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 scale_module(module, scale):
+    """
+    Scale the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().mul_(scale)
+    return module
+
+def mean_flat(tensor):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+def normalization(channels):
+    """
+    Make a standard normalization layer.
+    :param channels: number of input channels.
+    :return: an nn.Module for normalization.
+    """
+    return GroupNorm32(32, channels)
+
+# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
+class SiLU(nn.Module):
+    def forward(self, x):
+        return x * torch.sigmoid(x)
+
+class GroupNorm32(nn.GroupNorm):
+    def forward(self, x):
+        # return super().forward(x.float()).type(x.dtype)
+        return super().forward(x)
+
+def conv_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D convolution module.
+    """
+    if dims == 1:
+        return nn.Conv1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.Conv3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+def linear(*args, **kwargs):
+    """
+    Create a linear module.
+    """
+    return nn.Linear(*args, **kwargs)
+
+def avg_pool_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D average pooling module.
+    """
+    if dims == 1:
+        return nn.AvgPool1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.AvgPool2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.AvgPool3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+class HybridConditioner(nn.Module):
+
+    def __init__(self, c_concat_config, c_crossattn_config):
+        super().__init__()
+        self.concat_conditioner = instantiate_from_config(c_concat_config)
+        self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)
+
+    def forward(self, c_concat, c_crossattn):
+        c_concat = self.concat_conditioner(c_concat)
+        c_crossattn = self.crossattn_conditioner(c_crossattn)
+        return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
+
+def noise_like(x, repeat=False):
+    noise = torch.randn_like(x)
+    if repeat:
+        bs = x.shape[0]
+        noise = noise[0:1].repeat(bs, *((1,) * (len(x.shape) - 1)))
+    return noise
+
+##########################
+# inherit from ldm.utils #
+##########################
+
+def count_params(model, verbose=False):
+    total_params = sum(p.numel() for p in model.parameters())
+    if verbose:
+        print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.")
+    return total_params

lib/model_zoo/distributions.py

@@ -0,0 +1,92 @@
+import torch
+import numpy as np
+
+
+class AbstractDistribution:
+    def sample(self):
+        raise NotImplementedError()
+
+    def mode(self):
+        raise NotImplementedError()
+
+
+class DiracDistribution(AbstractDistribution):
+    def __init__(self, value):
+        self.value = value
+
+    def sample(self):
+        return self.value
+
+    def mode(self):
+        return self.value
+
+
+class DiagonalGaussianDistribution(object):
+    def __init__(self, parameters, deterministic=False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
+
+    def sample(self):
+        x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
+        return x
+
+    def kl(self, other=None):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        else:
+            if other is None:
+                return 0.5 * torch.sum(torch.pow(self.mean, 2)
+                                       + self.var - 1.0 - self.logvar,
+                                       dim=[1, 2, 3])
+            else:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean - other.mean, 2) / other.var
+                    + self.var / other.var - 1.0 - self.logvar + other.logvar,
+                    dim=[1, 2, 3])
+
+    def nll(self, sample, dims=[1,2,3]):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims)
+
+    def mode(self):
+        return self.mean
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12
+    Compute the KL divergence between two gaussians.
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, torch.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for torch.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor)
+        for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+        -1.0
+        + logvar2
+        - logvar1
+        + torch.exp(logvar1 - logvar2)
+        + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
+    )