Duplicate from shi-labs/Versatile-Diffusion

Co-authored-by: Xingqian Xu <[email protected]>

.gitattributes

@@ -0,0 +1,36 @@
+*.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
+*.gif filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,11 @@
+__pycache__
+.vscode/
+src/
+data/
+data
+log/
+log
+pretrained/
+pretrained
+gradio_cached_examples/
+gradio_cached_examples

README.md

@@ -0,0 +1,15 @@
+---
+title: Versatile Diffusion
+emoji: null
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 3.17.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,1114 @@
+################################################################################
+# Copyright (C) 2023 Xingqian Xu - All Rights Reserved                         #
+#                                                                              #
+# Please visit Versatile Diffusion's arXiv paper for more details, link at     #
+# arxiv.org/abs/2211.08332                                                     #
+#                                                                              #
+# Besides, this work is also inspired by many established techniques including:#
+# Denoising Diffusion Probablistic Model; Denoising Diffusion Implicit Model;  #
+# Latent Diffusion Model; Stable Diffusion; Stable Diffusion - Img2Img; Stable #
+# Diffusion - Variation; ImageMixer; DreamBooth; Stable Diffusion - Lora; More #
+# Control for Free; Prompt-to-Prompt;                                          #
+#                                                                              #
+################################################################################
+
+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 types
+
+import torch
+import torchvision.transforms as tvtrans
+from lib.cfg_helper import model_cfg_bank
+from lib.model_zoo import get_model
+from cusomized_gradio_blocks import create_myexamples, customized_as_example, customized_postprocess
+
+n_sample_image = 2
+n_sample_text = 4
+cache_examples = True
+
+from lib.model_zoo.ddim import DDIMSampler
+
+##########
+# helper #
+##########
+
+def highlight_print(info):
+    print('')
+    print(''.join(['#']*(len(info)+4)))
+    print('# '+info+' #')
+    print(''.join(['#']*(len(info)+4)))
+    print('')
+
+def decompose(x, q=20, niter=100):
+    x_mean = x.mean(-1, keepdim=True)
+    x_input = x - x_mean
+    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]))
+    x_remain = x_input - x_lowrank
+    return u, s, v, x_mean, x_remain
+
+class adjust_rank(object):
+    def __init__(self, max_drop_rank=[1, 5], q=20):
+        self.max_semantic_drop_rank = max_drop_rank[0]
+        self.max_style_drop_rank = max_drop_rank[1]
+        self.q = q
+
+        def t2y0_semf_wrapper(t0, y00, t1, y01):
+            return lambda t: (np.exp((t-0.5)*2)-t0)/(t1-t0)*(y01-y00)+y00
+        t0, y00 = np.exp((0  -0.5)*2), -self.max_semantic_drop_rank
+        t1, y01 = np.exp((0.5-0.5)*2), 1
+        self.t2y0_semf = t2y0_semf_wrapper(t0, y00, t1, y01)
+
+        def x2y_semf_wrapper(x0, x1, y1):
+            return lambda x, y0: (x-x0)/(x1-x0)*(y1-y0)+y0
+        x0 = 0
+        x1, y1 = self.max_semantic_drop_rank+1, 1
+        self.x2y_semf = x2y_semf_wrapper(x0, x1, y1)
+        
+        def t2y0_styf_wrapper(t0, y00, t1, y01):
+            return lambda t: (np.exp((t-0.5)*2)-t0)/(t1-t0)*(y01-y00)+y00
+        t0, y00 = np.exp((1  -0.5)*2), -(q-self.max_style_drop_rank)
+        t1, y01 = np.exp((0.5-0.5)*2), 1
+        self.t2y0_styf = t2y0_styf_wrapper(t0, y00, t1, y01)
+
+        def x2y_styf_wrapper(x0, x1, y1):
+            return lambda x, y0: (x-x0)/(x1-x0)*(y1-y0)+y0
+        x0 = q-1
+        x1, y1 = self.max_style_drop_rank-1, 1
+        self.x2y_styf = x2y_styf_wrapper(x0, x1, y1)
+
+    def __call__(self, x, lvl):
+        if lvl == 0.5:
+            return x
+
+        if x.dtype == torch.float16:
+            fp16 = True
+            x = x.float()
+        else:
+            fp16 = False
+        std_save = x.std(axis=[-2, -1])
+
+        u, s, v, x_mean, x_remain = decompose(x, q=self.q)
+
+        if lvl < 0.5:
+            assert lvl>=0
+            for xi in range(0, self.max_semantic_drop_rank+1):
+                y0 = self.t2y0_semf(lvl)
+                yi = self.x2y_semf(xi, y0)
+                yi = 0 if yi<0 else yi
+                s[:, xi] *= yi
+
+        elif lvl > 0.5:
+            assert lvl <= 1
+            for xi in range(self.max_style_drop_rank, self.q):
+                y0 = self.t2y0_styf(lvl)
+                yi = self.x2y_styf(xi, y0)
+                yi = 0 if yi<0 else yi
+                s[:, xi] *= yi
+            x_remain = 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]))
+        x_new = x_lowrank + x_mean + x_remain
+
+        std_new = x_new.std(axis=[-2, -1])
+        x_new = x_new / std_new * std_save
+
+        if fp16:
+            x_new = x_new.half()
+
+        return x_new
+
+def remove_duplicate_word(tx):
+    def combine_words(input, length):
+        combined_inputs = []
+        if len(splitted_input)>1:
+            for i in range(len(input)-1):
+                combined_inputs.append(input[i]+" "+last_word_of(splitted_input[i+1],length)) #add the last word of the right-neighbour (overlapping) sequence (before it has expanded), which is the next word in the original sentence
+        return combined_inputs, length+1
+
+    def remove_duplicates(input, length):
+        bool_broke=False #this means we didn't find any duplicates here
+        for i in range(len(input) - length):
+            if input[i]==input[i + length]: #found a duplicate piece of sentence!
+                for j in range(0, length): #remove the overlapping sequences in reverse order
+                    del input[i + length - j]
+                bool_broke = True
+                break #break the for loop as the loop length does not matches the length of splitted_input anymore as we removed elements
+        if bool_broke:
+            return remove_duplicates(input, length) #if we found a duplicate, look for another duplicate of the same length
+        return input
+
+    def last_word_of(input, length):
+        splitted = input.split(" ")
+        if len(splitted)==0:
+            return input
+        else:
+            return splitted[length-1]
+
+    def split_and_puncsplit(text):
+        tx = text.split(" ")
+        txnew = []
+        for txi in tx:
+            txqueue=[]
+            while True:
+                if txi[0] in '([{':
+                    txqueue.extend([txi[:1], '<puncnext>'])
+                    txi = txi[1:]
+                    if len(txi) == 0:
+                        break
+                else:
+                    break
+            txnew += txqueue
+            txstack=[]
+            if len(txi) == 0:
+                continue
+            while True:
+                if txi[-1] in '?!.,:;}])':
+                    txstack = ['<puncnext>', txi[-1:]] + txstack
+                    txi = txi[:-1]
+                    if len(txi) == 0:
+                        break
+                else:
+                    break
+            if len(txi) != 0:
+                txnew += [txi]
+            txnew += txstack
+        return txnew
+
+    if tx == '':
+        return tx
+
+    splitted_input = split_and_puncsplit(tx)
+    word_length = 1
+    intermediate_output = False
+    while len(splitted_input)>1:
+        splitted_input = remove_duplicates(splitted_input, word_length)
+        if len(splitted_input)>1:
+            splitted_input, word_length = combine_words(splitted_input, word_length)
+        if intermediate_output:
+            print(splitted_input)
+            print(word_length)
+    output = splitted_input[0]
+    output = output.replace(' <puncnext> ', '')
+    return output
+
+def get_instruction(mode):
+    t2i_instruction = ["Generate image from text prompt."]
+    i2i_instruction = ["Generate image conditioned on reference image.",]
+    i2t_instruction = ["Generate text from reference image. "]
+    t2t_instruction = ["Generate text from reference text prompt. "]
+    dcg_instruction = ["Generate image conditioned on both text and image."]
+    tcg_instruction = ["Generate image conditioned on text and up to two images."]
+    mcg_instruction = ["Generate image from multiple contexts."]
+
+    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 == "Dual-Context":
+        return '\n'.join(dcg_instruction)
+    elif mode == "Triple-Context":
+        return '\n'.join(tcg_instruction)
+    elif mode == "Multi-Context":
+        return '\n'.join(mcg_instruction)
+    else:
+        assert False
+
+########
+# main #
+########
+class vd_dummy(object):
+    def __init__(self, *args, **kwarg):
+        self.which = 'Vdummy'
+    def inference_t2i(self, *args, **kwarg): pass
+    def inference_i2i(self, *args, **kwarg): pass
+    def inference_i2t(self, *args, **kwarg): pass
+    def inference_t2t(self, *args, **kwarg): pass
+    def inference_dcg(self, *args, **kwarg): pass
+    def inference_tcg(self, *args, **kwarg): pass
+    def inference_mcg(self, *args, **kwarg): 
+        return None, None
+
+class vd_inference(object):
+    def __init__(self, fp16=False, which='v2.0'):
+        highlight_print(which)
+        self.which = which
+
+        if self.which == 'v1.0':
+            cfgm = model_cfg_bank()('vd_four_flow_v1-0')
+        else:
+            assert False, 'Model type not supported'
+        net = get_model()(cfgm)
+
+        if fp16:
+            highlight_print('Running in FP16')
+            if self.which == 'v1.0':
+                net.ctx['text'].fp16 = True
+                net.ctx['image'].fp16 = True
+            net = net.half()
+            self.dtype = torch.float16
+        else:
+            self.dtype = torch.float32
+
+        if self.which == 'v1.0':
+            # if fp16:
+            #     sd = torch.load('pretrained/vd-four-flow-v1-0-fp16.pth', map_location='cpu')
+            # else:
+            #     sd = torch.load('pretrained/vd-four-flow-v1-0.pth', map_location='cpu')
+            from huggingface_hub import hf_hub_download
+            if fp16:
+                temppath = hf_hub_download('shi-labs/versatile-diffusion-model', 'pretrained_pth/vd-four-flow-v1-0-fp16.pth')
+            else:
+                temppath = hf_hub_download('shi-labs/versatile-diffusion-model', 'pretrained_pth/vd-four-flow-v1-0.pth')
+            sd = torch.load(temppath, map_location='cpu')
+
+        net.load_state_dict(sd, strict=False)
+
+        self.use_cuda = torch.cuda.is_available()
+        if self.use_cuda:
+            net.to('cuda')
+        self.net = net
+        self.sampler = DDIMSampler(net)
+
+        self.output_dim = [512, 512]
+        self.n_sample_image = n_sample_image
+        self.n_sample_text = n_sample_text
+        self.ddim_steps = 50
+        self.ddim_eta = 0.0
+        self.scale_textto = 7.5
+        self.image_latent_dim = 4
+        self.text_latent_dim = 768
+        self.text_temperature = 1
+
+        if which == 'v1.0':
+            self.adjust_rank_f = adjust_rank(max_drop_rank=[1, 5], q=20)
+            self.scale_imgto = 7.5
+            self.disentanglement_noglobal = True
+
+    def inference_t2i(self, text, seed):
+        n_samples = self.n_sample_image
+        scale = self.scale_textto
+        sampler = self.sampler
+        h, w = self.output_dim
+        u = self.net.ctx_encode([""], which='text').repeat(n_samples, 1, 1)
+        c = self.net.ctx_encode([text], which='text').repeat(n_samples, 1, 1)
+        shape = [n_samples, self.image_latent_dim, h//8, w//8]
+        np.random.seed(seed)
+        torch.manual_seed(seed + 100)
+        x, _ = sampler.sample(
+            steps=self.ddim_steps,
+            x_info={'type':'image'},
+            c_info={'type':'text', 'conditioning':c, 'unconditional_conditioning':u, 
+                    'unconditional_guidance_scale':scale},
+            shape=shape,
+            verbose=False,
+            eta=self.ddim_eta)
+        im = self.net.vae_decode(x, which='image')
+        im = [tvtrans.ToPILImage()(i) for i in im]
+        return im
+
+    def inference_i2i(self, im, fid_lvl, fcs_lvl, clr_adj, seed):
+        n_samples = self.n_sample_image
+        scale = self.scale_imgto
+        sampler = self.sampler
+        h, w = self.output_dim
+        device = self.net.device
+
+        BICUBIC = PIL.Image.Resampling.BICUBIC
+        im = im.resize([w, h], resample=BICUBIC)
+
+        if fid_lvl == 1:
+            return [im]*n_samples
+
+        cx = tvtrans.ToTensor()(im)[None].to(device).to(self.dtype)
+
+        c = self.net.ctx_encode(cx, which='image')
+        if self.disentanglement_noglobal:
+            c_glb = c[:, 0:1]
+            c_loc = c[:, 1: ]
+            c_loc = self.adjust_rank_f(c_loc, fcs_lvl)
+            c = torch.cat([c_glb, c_loc], dim=1).repeat(n_samples, 1, 1)
+        else:
+            c = self.adjust_rank_f(c, fcs_lvl).repeat(n_samples, 1, 1)
+        u = torch.zeros_like(c)
+
+        shape = [n_samples, self.image_latent_dim, h//8, w//8]
+        np.random.seed(seed)
+        torch.manual_seed(seed + 100)
+        if fid_lvl!=0:
+            x0 = self.net.vae_encode(cx, which='image').repeat(n_samples, 1, 1, 1)
+            step = int(self.ddim_steps * (1-fid_lvl))
+            x, _ = sampler.sample(
+                steps=self.ddim_steps,
+                x_info={'type':'image', 'x0':x0, 'x0_forward_timesteps':step},
+                c_info={'type':'image', 'conditioning':c, 'unconditional_conditioning':u, 
+                        'unconditional_guidance_scale':scale},
+                shape=shape,
+                verbose=False,
+                eta=self.ddim_eta)
+        else:
+            x, _ = sampler.sample(
+                steps=self.ddim_steps,
+                x_info={'type':'image',},
+                c_info={'type':'image', 'conditioning':c, 'unconditional_conditioning':u, 
+                        'unconditional_guidance_scale':scale},
+                shape=shape,
+                verbose=False,
+                eta=self.ddim_eta)
+
+        imout = self.net.vae_decode(x, which='image')
+
+        if clr_adj == 'Simple':
+            cx_mean = cx.view(3, -1).mean(-1)[:, None, None]
+            cx_std  = cx.view(3, -1).std(-1)[:, None, None]
+            imout_mean = [imouti.view(3, -1).mean(-1)[:, None, None] for imouti in imout]
+            imout_std  = [imouti.view(3, -1).std(-1)[:, None, None] for imouti in imout]
+            imout = [(ii-mi)/si*cx_std+cx_mean for ii, mi, si in zip(imout, imout_mean, imout_std)]
+            imout = [torch.clamp(ii, 0, 1) for ii in imout]
+
+        imout = [tvtrans.ToPILImage()(i) for i in imout]
+        return imout
+
+    def inference_i2t(self, im, seed):
+        n_samples = self.n_sample_text
+        scale = self.scale_imgto
+        sampler = self.sampler
+        h, w = self.output_dim
+        device = self.net.device
+
+        BICUBIC = PIL.Image.Resampling.BICUBIC
+        im = im.resize([w, h], resample=BICUBIC)
+
+        cx = tvtrans.ToTensor()(im)[None].to(device)
+        c = self.net.ctx_encode(cx, which='image').repeat(n_samples, 1, 1)
+        u = self.net.ctx_encode(torch.zeros_like(cx), which='image').repeat(n_samples, 1, 1)
+
+        shape = [n_samples, self.text_latent_dim]
+        np.random.seed(seed)
+        torch.manual_seed(seed + 100)
+        x, _ = sampler.sample(
+            steps=self.ddim_steps,
+            x_info={'type':'text',},
+            c_info={'type':'image', 'conditioning':c, 'unconditional_conditioning':u, 
+                    'unconditional_guidance_scale':scale},
+            shape=shape,
+            verbose=False,
+            eta=self.ddim_eta)
+        tx = self.net.vae_decode(x, which='text', temperature=self.text_temperature)
+        tx = [remove_duplicate_word(txi) for txi in tx]
+        tx_combined = '\n'.join(tx)
+        return tx_combined
+
+    def inference_t2t(self, text, seed):
+        n_samples = self.n_sample_text
+        scale = self.scale_textto
+        sampler = self.sampler
+        u = self.net.ctx_encode([""], which='text').repeat(n_samples, 1, 1)
+        c = self.net.ctx_encode([text], which='text').repeat(n_samples, 1, 1)
+        shape = [n_samples, self.text_latent_dim]
+        np.random.seed(seed)
+        torch.manual_seed(seed + 100)
+        x, _ = sampler.sample(
+            steps=self.ddim_steps,
+            x_info={'type':'text',},
+            c_info={'type':'text', 'conditioning':c, 'unconditional_conditioning':u, 
+                    'unconditional_guidance_scale':scale},
+            shape=shape,
+            verbose=False,
+            eta=self.ddim_eta)
+        tx = self.net.vae_decode(x, which='text', temperature=self.text_temperature)
+        tx = [remove_duplicate_word(txi) for txi in tx]
+        tx_combined = '\n'.join(tx)
+        return tx_combined
+
+    def inference_dcg(self, imctx, fcs_lvl, textctx, textstrength, seed):
+        n_samples = self.n_sample_image
+        sampler = self.sampler
+        h, w = self.output_dim
+        device = self.net.device
+
+        c_info_list = []
+
+        if (textctx is not None) and (textctx != "") and (textstrength != 0):
+            ut = self.net.ctx_encode([""], which='text').repeat(n_samples, 1, 1)
+            ct = self.net.ctx_encode([textctx], which='text').repeat(n_samples, 1, 1)
+            scale = self.scale_imgto*(1-textstrength) + self.scale_textto*textstrength
+
+            c_info_list.append({
+                'type':'text', 
+                'conditioning':ct, 
+                'unconditional_conditioning':ut,
+                'unconditional_guidance_scale':scale,
+                'ratio': textstrength, })
+        else:
+            scale = self.scale_imgto
+            textstrength = 0
+
+        BICUBIC = PIL.Image.Resampling.BICUBIC
+        cx = imctx.resize([w, h], resample=BICUBIC)
+        cx = tvtrans.ToTensor()(cx)[None].to(device).to(self.dtype)
+        ci = self.net.ctx_encode(cx, which='image')
+
+        if self.disentanglement_noglobal:
+            ci_glb = ci[:, 0:1]
+            ci_loc = ci[:, 1: ]
+            ci_loc = self.adjust_rank_f(ci_loc, fcs_lvl)
+            ci = torch.cat([ci_glb, ci_loc], dim=1).repeat(n_samples, 1, 1)
+        else:
+            ci = self.adjust_rank_f(ci, fcs_lvl).repeat(n_samples, 1, 1)
+
+        c_info_list.append({
+            'type':'image', 
+            'conditioning':ci, 
+            'unconditional_conditioning':torch.zeros_like(ci),
+            'unconditional_guidance_scale':scale,
+            'ratio': (1-textstrength), })
+
+        shape = [n_samples, self.image_latent_dim, h//8, w//8]
+        np.random.seed(seed)
+        torch.manual_seed(seed + 100)
+        x, _ = sampler.sample_multicontext(
+            steps=self.ddim_steps,
+            x_info={'type':'image',},
+            c_info_list=c_info_list,
+            shape=shape,
+            verbose=False,
+            eta=self.ddim_eta)
+
+        imout = self.net.vae_decode(x, which='image')
+        imout = [tvtrans.ToPILImage()(i) for i in imout]
+        return imout
+
+    def inference_tcg(self, *args):
+        args_imag = list(args[0:10]) + [None, None, None, None, None]*2
+        args_rest = args[10:]
+        imin, imout = self.inference_mcg(*args_imag, *args_rest)
+        return imin, imout
+
+    def inference_mcg(self, *args):
+        imctx = [args[0:5], args[5:10], args[10:15], args[15:20]]
+        textctx, textstrength, seed = args[20:]
+
+        n_samples = self.n_sample_image
+        sampler = self.sampler
+        h, w = self.output_dim
+        device = self.net.device
+
+        c_info_list = []
+
+        if (textctx is not None) and (textctx != "") and (textstrength != 0):
+            ut = self.net.ctx_encode([""], which='text').repeat(n_samples, 1, 1)
+            ct = self.net.ctx_encode([textctx], which='text').repeat(n_samples, 1, 1)
+            scale = self.scale_imgto*(1-textstrength) + self.scale_textto*textstrength
+
+            c_info_list.append({
+                'type':'text', 
+                'conditioning':ct, 
+                'unconditional_conditioning':ut,
+                'unconditional_guidance_scale':scale,
+                'ratio': textstrength, })
+        else:
+            scale = self.scale_imgto
+            textstrength = 0
+
+        input_save = []
+        imc = []
+        for im, imm, strength, fcs_lvl, use_mask in imctx:
+            if (im is None) and (imm is None):
+                continue
+            BILINEAR = PIL.Image.Resampling.BILINEAR
+            BICUBIC = PIL.Image.Resampling.BICUBIC
+            if use_mask:
+                cx = imm['image'].resize([w, h], resample=BICUBIC)
+                cx = tvtrans.ToTensor()(cx)[None].to(self.dtype).to(device)
+                m = imm['mask'].resize([w, h], resample=BILINEAR)
+                m = tvtrans.ToTensor()(m)[None, 0:1].to(self.dtype).to(device)
+                m = (1-m)
+                cx_show = cx*m
+                ci = self.net.ctx_encode(cx, which='image', masks=m)
+            else:
+                cx = im.resize([w, h], resample=BICUBIC)
+                cx = tvtrans.ToTensor()(cx)[None].to(self.dtype).to(device)
+                ci = self.net.ctx_encode(cx, which='image')
+                cx_show = cx
+
+            input_save.append(tvtrans.ToPILImage()(cx_show[0]))
+
+            if self.disentanglement_noglobal:
+                ci_glb = ci[:, 0:1]
+                ci_loc = ci[:, 1: ]
+                ci_loc = self.adjust_rank_f(ci_loc, fcs_lvl)
+                ci = torch.cat([ci_glb, ci_loc], dim=1).repeat(n_samples, 1, 1)
+            else:
+                ci = self.adjust_rank_f(ci, fcs_lvl).repeat(n_samples, 1, 1)
+            imc.append(ci * strength)
+
+        cis = torch.cat(imc, dim=1)
+        c_info_list.append({
+            'type':'image', 
+            'conditioning':cis, 
+            'unconditional_conditioning':torch.zeros_like(cis),
+            'unconditional_guidance_scale':scale,
+            'ratio': (1-textstrength), })
+
+        shape = [n_samples, self.image_latent_dim, h//8, w//8]
+        np.random.seed(seed)
+        torch.manual_seed(seed + 100)
+        x, _ = sampler.sample_multicontext(
+            steps=self.ddim_steps,
+            x_info={'type':'image',},
+            c_info_list=c_info_list,
+            shape=shape,
+            verbose=False,
+            eta=self.ddim_eta)
+
+        imout = self.net.vae_decode(x, which='image')
+        imout = [tvtrans.ToPILImage()(i) for i in imout]
+        return input_save, imout
+
+# vd_inference = vd_dummy()
+vd_inference = vd_inference(which='v1.0', fp16=True)
+
+#################
+# sub interface #
+#################
+
+def t2i_interface(with_example=False):
+    gr.HTML('<p id=myinst>&nbsp Description: ' + get_instruction("Text-to-Image") + '</p>')
+    with gr.Row():
+        with gr.Column():
+            text = gr.Textbox(lines=4, placeholder="Input prompt...", label='Text Input')
+            seed = gr.Number(20, label="Seed", precision=0)
+            button = gr.Button("Run")
+        with gr.Column():
+            img_output = gr.Gallery(label="Image Result", elem_id='customized_imbox').style(grid=n_sample_image)
+
+    button.click(
+        vd_inference.inference_t2i,
+        inputs=[text, seed],
+        outputs=[img_output])
+
+    if with_example:
+        gr.Examples(
+            label='Examples',
+            examples=get_example('Text-to-Image'),
+            fn=vd_inference.inference_t2i,
+            inputs=[text, seed],
+            outputs=[img_output],
+            cache_examples=cache_examples),
+
+def i2i_interface(with_example=False):
+    gr.HTML('<p id=myinst>&nbsp Description: ' + get_instruction("Image-Variation") + '</p>')
+    with gr.Row():
+        with gr.Column():
+            img_input = gr.Image(label='Image Input', type='pil', elem_id='customized_imbox')
+            sim_flag  = gr.Checkbox(label='Show Detail Controls')
+            with gr.Row():
+                fid_lvl = gr.Slider(label="Fidelity (Dislike -- Same)", minimum=0, maximum=1, value=0, step=0.02, visible=False)
+                fcs_lvl = gr.Slider(label="Focus (Semantic -- Style)", minimum=0, maximum=1, value=0.5, step=0.02, visible=False)
+            clr_adj = gr.Radio(label="Color Adjustment", choices=["None", "Simple"], value='Simple', visible=False)
+            explain = gr.HTML('<p id=myinst>&nbsp Fidelity: How likely the output image looks like the referece image (0-dislike (default), 1-same).</p>'+
+                              '<p id=myinst>&nbsp Focus: What the output image should focused on (0-semantic, 0.5-balanced (default), 1-style).</p>', 
+                              visible=False)
+            seed = gr.Number(20, label="Seed", precision=0)
+            button = gr.Button("Run")
+        with gr.Column():
+            img_output = gr.Gallery(label="Image Result", elem_id='customized_imbox').style(grid=n_sample_image)
+
+        sim_flag.change(
+            fn=lambda x: {
+                explain : gr.update(visible=x), 
+                fid_lvl : gr.update(visible=x), 
+                fcs_lvl : gr.update(visible=x), 
+                clr_adj : gr.update(visible=x), },
+            inputs=sim_flag,
+            outputs=[explain, fid_lvl, fcs_lvl, clr_adj, seed],)
+
+    button.click(
+        vd_inference.inference_i2i,
+        inputs=[img_input, fid_lvl, fcs_lvl, clr_adj, seed],
+        outputs=[img_output])
+
+    if with_example:
+        gr.Examples(
+            label='Examples',
+            examples=get_example('Image-Variation'),
+            fn=vd_inference.inference_i2i,
+            inputs=[img_input, fid_lvl, fcs_lvl, clr_adj, seed],
+            outputs=[img_output],
+            cache_examples=cache_examples),
+
+def i2t_interface(with_example=False):
+    gr.HTML('<p id=myinst>&nbsp Description: ' + get_instruction("Image-to-Text") + '</p>')
+    with gr.Row():
+        with gr.Column():
+            img_input = gr.Image(label='Image Input', type='pil', elem_id='customized_imbox')
+            seed = gr.Number(20, label="Seed", precision=0)
+            button = gr.Button("Run")
+        with gr.Column():
+            txt_output = gr.Textbox(lines=4, label='Text Result')
+
+    button.click(
+        vd_inference.inference_i2t,
+        inputs=[img_input, seed],
+        outputs=[txt_output])
+
+    if with_example:
+        gr.Examples(
+            label='Examples',
+            examples=get_example('Image-to-Text'),
+            fn=vd_inference.inference_i2t,
+            inputs=[img_input, seed],
+            outputs=[txt_output],
+            cache_examples=cache_examples),
+
+def t2t_interface(with_example=False):
+    gr.HTML('<p id=myinst>&nbsp Description: ' + get_instruction("Text-Variation") + '</p>')
+    with gr.Row():
+        with gr.Column():
+            text = gr.Textbox(lines=4, placeholder="Input prompt...", label='Text Input')
+            seed = gr.Number(20, label="Seed", precision=0)
+            button = gr.Button("Run")
+        with gr.Column():
+            txt_output = gr.Textbox(lines=4, label='Text Result')
+
+    button.click(
+        vd_inference.inference_t2t,
+        inputs=[text, seed],
+        outputs=[txt_output])
+
+    if with_example:
+        gr.Examples(
+            label='Examples',
+            examples=get_example('Text-Variation'),
+            fn=vd_inference.inference_t2t,
+            inputs=[text, seed],
+            outputs=[txt_output],
+            cache_examples=cache_examples, )
+
+class image_mimage_swap(object):
+    def __init__(self, block0, block1):
+        self.block0 = block0
+        self.block1 = block1
+        self.which_update = 'both'
+
+    def __call__(self, x0, x1, flag):
+        if self.which_update == 'both':
+            return self.update_both(x0, x1, flag)
+        elif self.which_update == 'visible':
+            return self.update_visible(x0, x1, flag)
+        elif self.which_update == 'visible_oneoff':
+            return self.update_visible_oneoff(x0, x1, flag)
+        else:
+            assert False
+
+    def update_both(self, x0, x1, flag):
+        if flag:
+            ug0 = gr.update(visible=False)
+            if x0 is None:
+                ug1 = gr.update(value=None, visible=True)
+            else:
+                if (x1 is not None) and ('mask' in x1):
+                    value1 = {'image':x0, 'mask':x1['mask']}
+                else:
+                    value1 = {'image':x0, 'mask':None}
+                ug1 = gr.update(value=value1, visible=True)
+        else:
+            if (x1 is not None) and ('image' in x1):
+                value0 = x1['image']
+            else:
+                value0 = None
+            ug0 = gr.update(value=value0, visible=True)
+            ug1 = gr.update(visible=False)
+        return {
+            self.block0 : ug0,
+            self.block1 : ug1,}
+
+    def update_visible(self, x0, x1, flag):
+        return {
+            self.block0 : gr.update(visible=not flag),
+            self.block1 : gr.update(visible=flag), }
+
+    def update_visible_oneoff(self, x0, x1, flag):
+        self.which_update = 'both'
+        return {
+            self.block0 : gr.update(visible=not flag),
+            self.block1 : gr.update(visible=flag), }
+
+class example_visible_only_hack(object):
+    def __init__(self, checkbox_list, functor_list):
+        self.checkbox_list = checkbox_list
+        self.functor_list = functor_list
+
+    def __call__(self, *args):
+        for bi, fi, vi in zip(self.checkbox_list, self.functor_list, args):
+            if bi.value != vi:
+                fi.which_update = 'visible_oneoff'
+
+def dcg_interface(with_example=False):
+    gr.HTML('<p id=myinst>&nbsp Description: ' + get_instruction("Dual-Context") + '</p>')
+    with gr.Row():
+        input_session = []
+        with gr.Column():
+            img = gr.Image(label='Image Input', type='pil', elem_id='customized_imbox')
+            fcs = gr.Slider(label="Focus (Semantic -- Style)", minimum=0, maximum=1, value=0.5, step=0.02)
+            gr.HTML('<p id=myinst>&nbsp Focus: Focus on what aspect of the image? (0-semantic, 0.5-balanced (default), 1-style).</p>')
+
+            text = gr.Textbox(lines=2, placeholder="Input prompt...", label='Text Input')
+            tstrength = gr.Slider(label="Text Domination (NoEffect -- TextOnly)", minimum=0, maximum=1, value=0, step=0.02)
+
+            seed = gr.Number(20, label="Seed", precision=0)
+            button = gr.Button("Run")
+
+        with gr.Column():
+            output_gallary = gr.Gallery(label="Image Result", elem_id='customized_imbox').style(grid=n_sample_image)
+
+    input_list = []
+    for i in input_session:
+        input_list += i
+    button.click(
+        vd_inference.inference_dcg, 
+        inputs=[img, fcs, text, tstrength, seed],
+        outputs=[output_gallary])
+
+    if with_example:
+        gr.Examples(
+            label='Examples',
+            examples=get_example('Dual-Context'),
+            fn=vd_inference.inference_dcg,
+            inputs=[img, fcs, text, tstrength, seed],
+            outputs=[output_gallary],
+            cache_examples=cache_examples)
+
+def tcg_interface(with_example=False):
+    gr.HTML('<p id=myinst>&nbsp Description: ' + get_instruction("Triple-Context") + '</p>')
+    with gr.Row():
+        input_session = []
+        with gr.Column(min_width=940):
+            with gr.Row():
+                with gr.Column():
+                    img0  = gr.Image(label='Image Input', type='pil', elem_id='customized_imbox')
+                    img0.as_example = types.MethodType(customized_as_example, img0)
+                    imgm0 = gr.Image(label='Image Input with Mask', type='pil', elem_id='customized_imbox', tool='sketch', source="upload", visible=False)
+                    imgm0.postprocess = types.MethodType(customized_postprocess, imgm0)
+                    imgm0.as_example = types.MethodType(customized_as_example, imgm0)
+                    istrength0 = gr.Slider(label="Weight", minimum=0, maximum=1, value=1, step=0.02)
+                    fcs0 = gr.Slider(label="Focus (Semantic -- Style)", minimum=0, maximum=1, value=0.5, step=0.02)
+                    msk0 = gr.Checkbox(label='Use mask?')
+                    swapf0 = image_mimage_swap(img0, imgm0)
+
+                    msk0.change(
+                        fn=swapf0,
+                        inputs=[img0, imgm0, msk0],
+                        outputs=[img0, imgm0],)
+                    input_session.append([img0, imgm0, istrength0, fcs0, msk0])
+
+                with gr.Column():
+                    img1  = gr.Image(label='Image Input', type='pil', elem_id='customized_imbox')
+                    img1.as_example = types.MethodType(customized_as_example, img1)
+                    imgm1 = gr.Image(label='Image Input with Mask', type='pil', elem_id='customized_imbox', tool='sketch', source="upload", visible=False)
+                    imgm1.postprocess = types.MethodType(customized_postprocess, imgm1)
+                    imgm1.as_example = types.MethodType(customized_as_example, imgm1)
+                    istrength1 = gr.Slider(label="Weight", minimum=0, maximum=1, value=1, step=0.02)
+                    fcs1 = gr.Slider(label="Focus (Semantic -- Style)", minimum=0, maximum=1, value=0.5, step=0.02)
+                    msk1 = gr.Checkbox(label='Use mask?')
+                    swapf1 = image_mimage_swap(img1, imgm1)
+
+                    msk1.change(
+                        fn=swapf1,
+                        inputs=[img1, imgm1, msk1],
+                        outputs=[img1, imgm1],)
+                    input_session.append([img1, imgm1, istrength1, fcs1, msk1])
+
+            gr.HTML('<p id=myinst>&nbsp Weight: The strength of the reference image. This weight is subject to <u>Text Domination</u>).</p>'+
+                    '<p id=myinst>&nbsp Focus: Focus on what aspect of the image? (0-semantic, 0.5-balanced (default), 1-style).</p>'+
+                    '<p id=myinst>&nbsp Mask: Remove regions on reference image so they will not influence the output.</p>',)
+
+            text = gr.Textbox(lines=2, placeholder="Input prompt...", label='Text Input')
+            tstrength = gr.Slider(label="Text Domination (NoEffect -- TextOnly)", minimum=0, maximum=1, value=0, step=0.02)
+
+            seed = gr.Number(20, label="Seed", precision=0)
+            button = gr.Button("Run")
+
+        with gr.Column(min_width=470):
+            input_gallary = gr.Gallery(label="Input Display", elem_id="customized_imbox").style(grid=2)
+            output_gallary = gr.Gallery(label="Image Result", elem_id="customized_imbox").style(grid=n_sample_image)
+
+    input_list = []
+    for i in input_session:
+        input_list += i
+    input_list += [text, tstrength, seed]
+    button.click(
+        vd_inference.inference_tcg, 
+        inputs=input_list,
+        outputs=[input_gallary, output_gallary])
+
+    if with_example:
+        create_myexamples(
+            label='Examples',
+            examples=get_example('Triple-Context'),
+            fn=vd_inference.inference_tcg,
+            inputs=input_list,
+            outputs=[input_gallary, output_gallary, ],
+            cache_examples=cache_examples, )
+
+    gr.HTML('<br><p id=myinst>&nbsp How to add mask: Please see the following instructions.</p><br>'+
+            '<div id="maskinst">'+
+                '<img src="file/assets/demo/misc/mask_inst1.gif">'+
+                '<img src="file/assets/demo/misc/mask_inst2.gif">'+
+                '<img src="file/assets/demo/misc/mask_inst3.gif">'+
+            '</div>')
+
+def mcg_interface(with_example=False):
+    num_img_input = 4
+    gr.HTML('<p id=myinst>&nbsp Description: ' + get_instruction("Multi-Context") + '</p>')
+    with gr.Row():
+        input_session = []
+        with gr.Column():
+            for idx in range(num_img_input):
+                with gr.Tab('Image{}'.format(idx+1)):
+                    img = gr.Image(label='Image Input', type='pil', elem_id='customized_imbox')
+                    img.as_example = types.MethodType(customized_as_example, img)
+                    imgm = gr.Image(label='Image Input with Mask', type='pil', elem_id='customized_imbox', tool='sketch', source="upload", visible=False)
+                    imgm.postprocess = types.MethodType(customized_postprocess, imgm)
+                    imgm.as_example = types.MethodType(customized_as_example, imgm)
+
+                    with gr.Row():
+                        istrength = gr.Slider(label="Weight", minimum=0, maximum=1, value=1, step=0.02)
+                        fcs = gr.Slider(label="Focus (Semantic -- Style)", minimum=0, maximum=1, value=0.5, step=0.02)
+                    msk = gr.Checkbox(label='Use mask?')
+                    gr.HTML('<p id=myinst>&nbsp Weight: The strength of the reference image. This weight is subject to <u>Text Domination</u>).</p>'+
+                            '<p id=myinst>&nbsp Focus: Focus on what aspect of the image? (0-semantic, 0.5-balanced (default), 1-style).</p>'+
+                            '<p id=myinst>&nbsp Mask: Remove regions on reference image so they will not influence the output.</p>',)
+
+                    msk.change(
+                        fn=image_mimage_swap(img, imgm),
+                        inputs=[img, imgm, msk],
+                        outputs=[img, imgm],)
+                    input_session.append([img, imgm, istrength, fcs, msk])
+
+            text = gr.Textbox(lines=2, placeholder="Input prompt...", label='Text Input')
+            tstrength = gr.Slider(label="Text Domination (NoEffect -- TextOnly)", minimum=0, maximum=1, value=0, step=0.02)
+
+            seed = gr.Number(20, label="Seed", precision=0)
+            button = gr.Button("Run")
+
+
+        with gr.Column():
+            input_gallary = gr.Gallery(label="Input Display", elem_id='customized_imbox').style(grid=4)
+            output_gallary = gr.Gallery(label="Image Result", elem_id='customized_imbox').style(grid=n_sample_image)
+
+    input_list = []
+    for i in input_session:
+        input_list += i
+    input_list += [text, tstrength, seed]
+    button.click(
+        vd_inference.inference_mcg, 
+        inputs=input_list,
+        outputs=[input_gallary, output_gallary], )
+
+    if with_example:
+        create_myexamples(
+            label='Examples',
+            examples=get_example('Multi-Context'),
+            fn=vd_inference.inference_mcg,
+            inputs=input_list,
+            outputs=[input_gallary, output_gallary],
+            cache_examples=cache_examples, )
+
+    gr.HTML('<br><p id=myinst>&nbsp How to add mask: Please see the following instructions.</p><br>'+
+            '<div id="maskinst">'+
+                '<img src="file/assets/demo/misc/mask_inst1.gif">'+
+                '<img src="file/assets/demo/misc/mask_inst2.gif">'+
+                '<img src="file/assets/demo/misc/mask_inst3.gif">'+
+            '</div>')
+
+###########
+# Example #
+###########
+
+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 landscape with mountains and rivers', 20],
+        ]
+    elif mode == "Image-Variation":
+        case = [
+            ['assets/demo/reg_example/ghibli.jpg', 0, 0.5, 'None', 20],
+            ['assets/demo/reg_example/ghibli.jpg', 0.5, 0.5, 'None', 20],
+            ['assets/demo/reg_example/matisse.jpg', 0, 0, 'None', 20],
+            ['assets/demo/reg_example/matisse.jpg', 0, 1, 'Simple', 20],
+            ['assets/demo/reg_example/vermeer.jpg', 0.2, 0.3, 'None', 30],
+        ]
+    elif mode == "Image-to-Text":
+        case = [
+            ['assets/demo/reg_example/house_by_lake.jpg', 20],
+        ]
+    elif mode == "Text-Variation":
+        case = [
+            ['heavy arms gundam penguin mech', 20],
+        ]
+    elif mode == "Dual-Context":
+        case = [
+            ['assets/demo/reg_example/benz.jpg', 0.5, 'cyberpunk 2077', 0.7, 22],
+            ['assets/demo/reg_example/ghibli.jpg', 1, 'Red maple on a hill in golden Autumn.', 0.5, 21],
+        ]
+    elif mode == "Triple-Context":
+        case = [
+            [
+                'assets/demo/reg_example/night_light.jpg', None, 1   , 0.5, False,
+                'assets/demo/reg_example/paris.jpg'      , None, 0.94, 0.5, False,
+                "snow on the street", 0.4, 28],
+            [
+                'assets/demo/tcg_example/e1i0.jpg', None, 1  , 0.5, False,
+                'assets/demo/tcg_example/e1i1.jpg', None, 0.94, 0.5, False,
+                "a painting of an elegant woman in front of the moon", 0.2, 217],
+            [
+                'assets/demo/tcg_example/e2i0.jpg',  None, 1, 0.5, False,
+                'assets/demo/reg_example/paris.jpg', None, 1, 0.5, False,
+                "", 0, 29],
+            [
+                'assets/demo/tcg_example/e0i0.jpg', None, 1  , 0.5, False,
+                'assets/demo/tcg_example/e0i1.jpg', None, 0.9, 0.5, False,
+                "rose blooms on the tree", 0.2, 20],
+            [
+                'assets/demo/reg_example/ghibli.jpg', None, 1   , 1  , False,
+                'assets/demo/reg_example/space.jpg' , None, 0.88, 0.5, False,
+                "", 0, 20],
+            [
+                'assets/demo/reg_example/train.jpg'  , None, 0.8, 0.5, False,
+                'assets/demo/reg_example/matisse.jpg', None, 1  , 1  , False,
+                "", 0, 20],
+        ]
+    elif mode == "Multi-Context":
+        case = [
+            [
+                'assets/demo/mcg_example/e0i0.jpg', None, 1, 0.5, False,
+                'assets/demo/mcg_example/e0i1.jpg', None, 1, 0.5, False,
+                'assets/demo/mcg_example/e0i2.jpg', None, 0.86, 0.5, False,
+                None, None, 1, 0.5, False,
+                "", 0, 20],
+        ]
+    else:
+        raise ValueError
+    return case
+
+#############
+# Interface #
+#############
+
+css = """
+    #customized_imbox {
+        min-height: 450px;
+    }
+    #customized_imbox>div[data-testid="image"] {
+        min-height: 450px;
+    }
+    #customized_imbox>div[data-testid="image"]>div {
+        min-height: 450px;
+    }
+    #customized_imbox>div[data-testid="image"]>iframe {
+        min-height: 450px;
+    }
+    #customized_imbox>div.unpadded_box {
+        min-height: 450px;
+    }
+    #myinst {
+        font-size: 0.8rem; 
+        margin: 0rem;
+        color: #6B7280;
+    }
+    #maskinst {
+        text-align: justify;
+        min-width: 1200px;
+    }
+    #maskinst>img {
+        min-width:399px;
+        max-width:450px;
+        vertical-align: top;
+        display: inline-block;
+    }
+    #maskinst:after {
+        content: "";
+        width: 100%;
+        display: inline-block;
+    }
+"""
+
+if True:
+    with gr.Blocks(css=css) as demo:
+        gr.HTML(
+            """
+            <div style="text-align: center; max-width: 1200px; margin: 20px auto;">
+            <h1 style="font-weight: 900; font-size: 3rem; margin: 0rem">
+                Versatile Diffusion
+            </h1>
+            <h2 style="font-weight: 450; font-size: 1rem; margin-top: 0.5rem; margin-bottom: 0.5rem">
+            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>
+            <h3 style="font-weight: 450; font-size: 1rem; margin: 0rem">
+            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>
+            """)
+
+        with gr.Tab('Text-to-Image'):
+            t2i_interface(with_example=True)
+        with gr.Tab('Image-Variation'):
+            i2i_interface(with_example=True)
+        with gr.Tab('Image-to-Text'):
+            i2t_interface(with_example=True)
+        with gr.Tab('Text-Variation'):
+            t2t_interface(with_example=True)
+        with gr.Tab('Dual-Context Image-Generation'):
+            dcg_interface(with_example=True)
+        with gr.Tab('Triple-Context Image-Blender'):
+            tcg_interface(with_example=True)
+        with gr.Tab('Multi-Context Image-Blender'):
+            mcg_interface(with_example=True)
+
+        gr.HTML(
+            """
+            <div style="text-align: justify; max-width: 1200px; margin: 20px auto;">
+            <h3 style="font-weight: 450; font-size: 0.8rem; margin: 0rem">
+            <b>Version</b>: {}
+            </h3>
+            <h3 style="font-weight: 450; font-size: 0.8rem; margin: 0rem">
+            <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>
+            <h3 style="font-weight: 450; font-size: 0.8rem; margin: 0rem">
+            <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>
+            """.format(' '+vd_inference.which))
+
+    # demo.launch(share=True)
+    demo.launch(debug=True)

configs/model/autokl.yaml

@@ -0,0 +1,23 @@
+autokl:
+  symbol: autokl
+  find_unused_parameters: false
+
+autokl_v1:
+  super_cfg: autokl
+  type: autoencoderkl
+  args:
+    embed_dim: 4
+    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: null
+  # pth: pretrained/kl-f8.pth
+  hfm: ['shi-labs/versatile-diffusion-model', 'pretrained_pth/kl-f8.pth']

configs/model/clip.yaml

@@ -0,0 +1,13 @@
+clip:
+  symbol: clip
+  args: {}
+
+clip_text_context_encoder:
+  super_cfg: clip
+  type: clip_text_context_encoder
+  args: {}
+
+clip_image_context_encoder:
+  super_cfg: clip
+  type: clip_image_context_encoder
+  args: {}

configs/model/openai_unet.yaml

@@ -0,0 +1,96 @@
+#########
+# v1 2d #
+#########
+
+openai_unet_2d_v1:
+  type: openai_unet_2d_next
+  args:
+    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 ]
+    num_heads: 8
+    context_dim: 768
+    use_checkpoint: True
+    parts: [global, data, context]
+
+openai_unet_2d_v1_g:
+  super_cfg: openai_unet_2d_v1
+  args:
+    parts: [global]
+
+openai_unet_2d_v1_d:
+  super_cfg: openai_unet_2d_v1
+  args:
+    parts: [data]
+
+openai_unet_2d_v1_c:
+  super_cfg: openai_unet_2d_v1
+  args:
+    parts: [context]
+
+openai_unet_2d_v1_gd:
+  super_cfg: openai_unet_2d_v1
+  args:
+    parts: [global, data]
+
+openai_unet_2d_v1_gc:
+  super_cfg: openai_unet_2d_v1
+  args:
+    parts: [global, context]
+
+openai_unet_2d_v1_dc:
+  super_cfg: openai_unet_2d_v1
+  args:
+    parts: [data, context]
+
+#########
+# v1 0d #
+#########
+
+openai_unet_0d_v1:
+  type: openai_unet_0d_next
+  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
+    parts: [global, data, context]
+
+openai_unet_0d_v1_g:
+  super_cfg: openai_unet_0d_v1
+  args:
+    parts: [global]
+
+openai_unet_0d_v1_d:
+  super_cfg: openai_unet_0d_v1
+  args:
+    parts: [data]
+
+openai_unet_0d_v1_c:
+  super_cfg: openai_unet_0d_v1
+  args:
+    parts: [context]
+
+openai_unet_0d_v1_gd:
+  super_cfg: openai_unet_0d_v1
+  args:
+    parts: [global, data]
+
+openai_unet_0d_v1_gc:
+  super_cfg: openai_unet_0d_v1
+  args:
+    parts: [global, context]
+
+openai_unet_0d_v1_dc:
+  super_cfg: openai_unet_0d_v1
+  args:
+    parts: [data, context]

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_v1:
+  super_cfg: optimus
+  type: optimus_vae_next
+  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
+  # pth: pretrained/optimus-vae.pth
+  hfm: ['shi-labs/versatile-diffusion-model', 'pretrained_pth/optimus-vae.pth']

configs/model/vd.yaml

@@ -0,0 +1,29 @@
+vd_base:
+  symbol: vd
+  find_unused_parameters: true
+  type: vd_v2_0
+  args:
+    beta_linear_start: 0.00085
+    beta_linear_end: 0.012
+    timesteps: 1000
+    use_ema: false
+
+###########
+# vd v1.0 #
+###########
+
+vd_four_flow_v1-0:
+  super_cfg: vd_base
+  args:
+    vae_cfg_list:
+      - [image, MODEL(autokl_v1)]
+      - [text, MODEL(optimus_v1)]
+    ctx_cfg_list:
+      - [image, MODEL(clip_image_context_encoder)]
+      - [text, MODEL(clip_text_context_encoder)]
+    diffuser_cfg_list:
+      - [image, MODEL(openai_unet_2d_v1)]
+      - [text, MODEL(openai_unet_0d_v1_dc)]
+    global_layer_ptr: image
+    latent_scale_factor: 
+      image: 0.18215

cusomized_gradio_blocks.py

@@ -0,0 +1,271 @@
+from __future__ import annotations
+
+import ast
+import csv
+import inspect
+import os
+import subprocess
+import tempfile
+import threading
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Tuple
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import PIL
+import PIL.Image
+
+import gradio
+from gradio import components, processing_utils, routes, utils
+from gradio.context import Context
+from gradio.documentation import document, set_documentation_group
+from gradio.flagging import CSVLogger
+
+if TYPE_CHECKING:  # Only import for type checking (to avoid circular imports).
+    from gradio.components import IOComponent
+
+CACHED_FOLDER = "gradio_cached_examples"
+LOG_FILE = "log.csv"
+
+def create_myexamples(
+    examples: List[Any] | List[List[Any]] | str,
+    inputs: IOComponent | List[IOComponent],
+    outputs: IOComponent | List[IOComponent] | None = None,
+    fn: Callable | None = None,
+    cache_examples: bool = False,
+    examples_per_page: int = 10,
+    _api_mode: bool = False,
+    label: str | None = None,
+    elem_id: str | None = None,
+    run_on_click: bool = False,
+    preprocess: bool = True,
+    postprocess: bool = True,
+    batch: bool = False,):
+    """Top-level synchronous function that creates Examples. Provided for backwards compatibility, i.e. so that gr.Examples(...) can be used to create the Examples component."""
+    examples_obj = MyExamples(
+        examples=examples,
+        inputs=inputs,
+        outputs=outputs,
+        fn=fn,
+        cache_examples=cache_examples,
+        examples_per_page=examples_per_page,
+        _api_mode=_api_mode,
+        label=label,
+        elem_id=elem_id,
+        run_on_click=run_on_click,
+        preprocess=preprocess,
+        postprocess=postprocess,
+        batch=batch,
+        _initiated_directly=False,
+    )
+    utils.synchronize_async(examples_obj.create)
+    return examples_obj
+
+class MyExamples(gradio.helpers.Examples):
+    def __init__(
+            self,
+            examples: List[Any] | List[List[Any]] | str,
+            inputs: IOComponent | List[IOComponent],
+            outputs: IOComponent | List[IOComponent] | None = None,
+            fn: Callable | None = None,
+            cache_examples: bool = False,
+            examples_per_page: int = 10,
+            _api_mode: bool = False,
+            label: str | None = "Examples",
+            elem_id: str | None = None,
+            run_on_click: bool = False,
+            preprocess: bool = True,
+            postprocess: bool = True,
+            batch: bool = False,
+            _initiated_directly: bool = True,):
+
+        if _initiated_directly:
+            warnings.warn(
+                "Please use gr.Examples(...) instead of gr.examples.Examples(...) to create the Examples.",
+            )
+
+        if cache_examples and (fn is None or outputs is None):
+            raise ValueError("If caching examples, `fn` and `outputs` must be provided")
+
+        if not isinstance(inputs, list):
+            inputs = [inputs]
+        if outputs and not isinstance(outputs, list):
+            outputs = [outputs]
+
+        working_directory = Path().absolute()
+
+        if examples is None:
+            raise ValueError("The parameter `examples` cannot be None")
+        elif isinstance(examples, list) and (
+            len(examples) == 0 or isinstance(examples[0], list)
+        ):
+            pass
+        elif (
+            isinstance(examples, list) and len(inputs) == 1
+        ):  # If there is only one input component, examples can be provided as a regular list instead of a list of lists
+            examples = [[e] for e in examples]
+        elif isinstance(examples, str):
+            if not Path(examples).exists():
+                raise FileNotFoundError(
+                    "Could not find examples directory: " + examples
+                )
+            working_directory = examples
+            if not (Path(examples) / LOG_FILE).exists():
+                if len(inputs) == 1:
+                    examples = [[e] for e in os.listdir(examples)]
+                else:
+                    raise FileNotFoundError(
+                        "Could not find log file (required for multiple inputs): "
+                        + LOG_FILE
+                    )
+            else:
+                with open(Path(examples) / LOG_FILE) as logs:
+                    examples = list(csv.reader(logs))
+                    examples = [
+                        examples[i][: len(inputs)] for i in range(1, len(examples))
+                    ]  # remove header and unnecessary columns
+
+        else:
+            raise ValueError(
+                "The parameter `examples` must either be a string directory or a list"
+                "(if there is only 1 input component) or (more generally), a nested "
+                "list, where each sublist represents a set of inputs."
+            )
+
+        input_has_examples = [False] * len(inputs)
+        for example in examples:
+            for idx, example_for_input in enumerate(example):
+                # if not (example_for_input is None):
+                if True:
+                    try:
+                        input_has_examples[idx] = True
+                    except IndexError:
+                        pass  # If there are more example components than inputs, ignore. This can sometimes be intentional (e.g. loading from a log file where outputs and timestamps are also logged)
+
+        inputs_with_examples = [
+            inp for (inp, keep) in zip(inputs, input_has_examples) if keep
+        ]
+        non_none_examples = [
+            [ex for (ex, keep) in zip(example, input_has_examples) if keep]
+            for example in examples
+        ]
+
+        self.examples = examples
+        self.non_none_examples = non_none_examples
+        self.inputs = inputs
+        self.inputs_with_examples = inputs_with_examples
+        self.outputs = outputs
+        self.fn = fn
+        self.cache_examples = cache_examples
+        self._api_mode = _api_mode
+        self.preprocess = preprocess
+        self.postprocess = postprocess
+        self.batch = batch
+
+        with utils.set_directory(working_directory):
+            self.processed_examples = [
+                [
+                    component.postprocess(sample)
+                    for component, sample in zip(inputs, example)
+                ]
+                for example in examples
+            ]
+        self.non_none_processed_examples = [
+            [ex for (ex, keep) in zip(example, input_has_examples) if keep]
+            for example in self.processed_examples
+        ]
+        if cache_examples:
+            for example in self.examples:
+                if len([ex for ex in example if ex is not None]) != len(self.inputs):
+                    warnings.warn(
+                        "Examples are being cached but not all input components have "
+                        "example values. This may result in an exception being thrown by "
+                        "your function. If you do get an error while caching examples, make "
+                        "sure all of your inputs have example values for all of your examples "
+                        "or you provide default values for those particular parameters in your function."
+                    )
+                    break
+
+        with utils.set_directory(working_directory):
+            self.dataset = components.Dataset(
+                components=inputs_with_examples,
+                samples=non_none_examples,
+                type="index",
+                label=label,
+                samples_per_page=examples_per_page,
+                elem_id=elem_id,
+            )
+
+        self.cached_folder = Path(CACHED_FOLDER) / str(self.dataset._id)
+        self.cached_file = Path(self.cached_folder) / "log.csv"
+        self.cache_examples = cache_examples
+        self.run_on_click = run_on_click
+
+from gradio import utils, processing_utils
+from PIL import Image as _Image
+from pathlib import Path
+import numpy as np
+
+def customized_postprocess(self, y):
+    if y is None:
+        return None
+
+    if isinstance(y, dict):
+        if self.tool == "sketch" and self.source in ["upload", "webcam"]:
+            y, mask = y["image"], y["mask"]
+            if y is None:
+                return None
+            elif isinstance(y, np.ndarray):
+                im = processing_utils.encode_array_to_base64(y)
+            elif isinstance(y, _Image.Image):
+                im = processing_utils.encode_pil_to_base64(y)
+            elif isinstance(y, (str, Path)):
+                im = processing_utils.encode_url_or_file_to_base64(y)
+            else:
+                raise ValueError("Cannot process this value as an Image")
+            im = self._format_image(im)
+
+            if mask is None:
+                return im
+            elif isinstance(y, np.ndarray):
+                mask_im = processing_utils.encode_array_to_base64(mask)
+            elif isinstance(y, _Image.Image):
+                mask_im = processing_utils.encode_pil_to_base64(mask)
+            elif isinstance(y, (str, Path)):
+                mask_im = processing_utils.encode_url_or_file_to_base64(mask)
+            else:
+                raise ValueError("Cannot process this value as an Image")
+
+            return {"image": im, "mask" : mask_im,}
+
+    elif isinstance(y, np.ndarray):
+        return processing_utils.encode_array_to_base64(y)
+    elif isinstance(y, _Image.Image):
+        return processing_utils.encode_pil_to_base64(y)
+    elif isinstance(y, (str, Path)):
+        return processing_utils.encode_url_or_file_to_base64(y)
+    else:
+        raise ValueError("Cannot process this value as an Image")
+
+# def customized_as_example(self, input_data=None):
+#     if input_data is None:
+#         return str('assets/demo/misc/noimage.jpg')
+#     elif isinstance(input_data, dict):
+#         im = np.array(PIL.Image.open(input_data["image"])).astype(float)
+#         mask = np.array(PIL.Image.open(input_data["mask"])).astype(float)/255
+#         imm = (im * (1-mask)).astype(np.uint8)
+#         import time
+#         ctime = int(time.time()*100)
+#         impath = 'assets/demo/temp/temp_{}.png'.format(ctime)
+#         PIL.Image.fromarray(imm).save(impath)
+#         return str(utils.abspath(impath))
+#     else:
+#         return str(utils.abspath(input_data))
+
+def customized_as_example(self, input_data=None):
+    if input_data is None:
+        return str('assets/demo/misc/noimage.jpg')
+    else:
+        return str(utils.abspath(input_data))

lib/cfg_helper.py

@@ -0,0 +1,612 @@
+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('openai_unet')==0:
+            return osp.join(
+                self.cfg_dir, 'openai_unet.yaml')
+        elif (name.find('clip')==0) or (name.find('openclip')==0):
+            return osp.join(
+                self.cfg_dir, 'clip.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')
+        elif name.find('autokl')==0:
+            return osp.join(
+                self.cfg_dir, 'autokl.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('laion2b')==0:
+            return osp.join(
+                self.cfg_dir, 'laion2b.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)
+    parser.add_argument('--node_list', nargs='+', type=str)
+    parser.add_argument('--nodes', type=int)
+    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)
+
+    if args.node_list is None:
+        cfg.env.node_rank = 0 if args.node_rank is None else args.node_rank
+        cfg.env.nodes = 1 if args.nodes is None else args.nodes
+    else:
+        import socket
+        hostname = socket.gethostname()
+        assert cfg.env.master_addr == args.node_list[0] 
+        cfg.env.node_rank = args.node_list.index(hostname)
+        cfg.env.nodes = len(args.node_list)
+        cfg.env.node_list = args.node_list
+
+    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', [])
+            sig = sig + ['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)
+    if cfge.node_rank==0:
+        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)
+    else:
+        with open(osp.join(log_dir, 'config.yaml.{}'.format(cfge.node_rank)), '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)
+    save_code = save_code and (cfge.node_rank==0)
+
+    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/autokl.py

@@ -0,0 +1,140 @@
+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 .autokl_modules import Encoder, Decoder
+from .distributions import DiagonalGaussianDistribution
+
+from .autokl_utils import LPIPSWithDiscriminator
+
+@register('autoencoderkl')
+class AutoencoderKL(nn.Module):
+    def __init__(self,
+                 ddconfig,
+                 lossconfig,
+                 embed_dim,):
+        super().__init__()
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        if lossconfig is not None:
+            self.loss = LPIPSWithDiscriminator(**lossconfig)
+        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
+
+    @torch.no_grad()
+    def encode(self, x, out_posterior=False):
+        return self.encode_trainable(x, out_posterior)
+
+    def encode_trainable(self, x, out_posterior=False):
+        x = x*2-1
+        h = self.encoder(x)
+        moments = self.quant_conv(h)
+        posterior = DiagonalGaussianDistribution(moments)
+        if out_posterior:
+            return posterior
+        else:
+            return posterior.sample()
+
+    @torch.no_grad()
+    def decode(self, z):
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+        dec = torch.clamp((dec+1)/2, 0, 1)
+        return dec
+
+    def decode_trainable(self, z):
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+        dec = (dec+1)/2
+        return dec
+
+    def apply_model(self, input, sample_posterior=True):
+        posterior = self.encode_trainable(input, out_posterior=True)
+        if sample_posterior:
+            z = posterior.sample()
+        else:
+            z = posterior.mode()
+        dec = self.decode_trainable(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 forward(self, x, optimizer_idx, global_step):
+        reconstructions, posterior = self.apply_model(x)
+
+        if optimizer_idx == 0:
+            # train encoder+decoder+logvar
+            aeloss, log_dict_ae = self.loss(x, reconstructions, posterior, optimizer_idx, global_step=global_step,
+                                            last_layer=self.get_last_layer(), split="train")
+            return aeloss, log_dict_ae
+
+        if optimizer_idx == 1:
+            # train the discriminator
+            discloss, log_dict_disc = self.loss(x, reconstructions, posterior, optimizer_idx, global_step=global_step,
+                                                last_layer=self.get_last_layer(), split="train")
+
+            return discloss, log_dict_disc
+
+    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

lib/model_zoo/autokl_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/autokl_utils.py

@@ -0,0 +1,400 @@
+import torch
+import torch.nn as nn
+import functools
+
+class ActNorm(nn.Module):
+    def __init__(self, num_features, logdet=False, affine=True,
+                 allow_reverse_init=False):
+        assert affine
+        super().__init__()
+        self.logdet = logdet
+        self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
+        self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
+        self.allow_reverse_init = allow_reverse_init
+
+        self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
+
+    def initialize(self, input):
+        with torch.no_grad():
+            flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
+            mean = (
+                flatten.mean(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+            std = (
+                flatten.std(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+
+            self.loc.data.copy_(-mean)
+            self.scale.data.copy_(1 / (std + 1e-6))
+
+    def forward(self, input, reverse=False):
+        if reverse:
+            return self.reverse(input)
+        if len(input.shape) == 2:
+            input = input[:,:,None,None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        _, _, height, width = input.shape
+
+        if self.training and self.initialized.item() == 0:
+            self.initialize(input)
+            self.initialized.fill_(1)
+
+        h = self.scale * (input + self.loc)
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+
+        if self.logdet:
+            log_abs = torch.log(torch.abs(self.scale))
+            logdet = height*width*torch.sum(log_abs)
+            logdet = logdet * torch.ones(input.shape[0]).to(input)
+            return h, logdet
+
+        return h
+
+    def reverse(self, output):
+        if self.training and self.initialized.item() == 0:
+            if not self.allow_reverse_init:
+                raise RuntimeError(
+                    "Initializing ActNorm in reverse direction is "
+                    "disabled by default. Use allow_reverse_init=True to enable."
+                )
+            else:
+                self.initialize(output)
+                self.initialized.fill_(1)
+
+        if len(output.shape) == 2:
+            output = output[:,:,None,None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        h = output / self.scale - self.loc
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+        return h
+
+#################
+# Discriminator #
+#################
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+class NLayerDiscriminator(nn.Module):
+    """Defines a PatchGAN discriminator as in Pix2Pix
+        --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+    """
+    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
+        """Construct a PatchGAN discriminator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            norm_layer      -- normalization layer
+        """
+        super(NLayerDiscriminator, self).__init__()
+        if not use_actnorm:
+            norm_layer = nn.BatchNorm2d
+        else:
+            norm_layer = ActNorm
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func != nn.BatchNorm2d
+        else:
+            use_bias = norm_layer != nn.BatchNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # gradually increase the number of filters
+            nf_mult_prev = nf_mult
+            nf_mult = min(2 ** n, 8)
+            sequence += [
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True)
+            ]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2 ** n_layers, 8)
+        sequence += [
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        sequence += [
+            nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+        self.main = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.main(input)
+
+#########
+# LPIPS #
+#########
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super(ScalingLayer, self).__init__()
+        self.register_buffer('shift', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
+        self.register_buffer('scale', torch.Tensor([.458, .448, .450])[None, :, None, None])
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+
+class NetLinLayer(nn.Module):
+    """ A single linear layer which does a 1x1 conv """
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super(NetLinLayer, self).__init__()
+        layers = [nn.Dropout(), ] if (use_dropout) else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False), ]
+        self.model = nn.Sequential(*layers)
+
+from collections import namedtuple
+from torchvision import models
+from torchvision.models import VGG16_Weights
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(vgg16, self).__init__()
+        if pretrained:
+            vgg_pretrained_features = models.vgg16(weights=VGG16_Weights.IMAGENET1K_V1).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+        return out
+
+def normalize_tensor(x,eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(x**2,dim=1,keepdim=True))
+    return x/(norm_factor+eps)
+
+def spatial_average(x, keepdim=True):
+    return x.mean([2,3],keepdim=keepdim)
+
+def get_ckpt_path(*args, **kwargs):
+    return 'pretrained/lpips.pth'
+
+class LPIPS(nn.Module):
+    # Learned perceptual metric
+    def __init__(self, use_dropout=True):
+        super().__init__()
+        self.scaling_layer = ScalingLayer()
+        self.chns = [64, 128, 256, 512, 512]  # vg16 features
+        self.net = vgg16(pretrained=True, requires_grad=False)
+        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+        self.load_from_pretrained()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def load_from_pretrained(self, name="vgg_lpips"):
+        ckpt = get_ckpt_path(name, "taming/modules/autoencoder/lpips")
+        self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+        print("loaded pretrained LPIPS loss from {}".format(ckpt))
+
+    @classmethod
+    def from_pretrained(cls, name="vgg_lpips"):
+        if name != "vgg_lpips":
+            raise NotImplementedError
+        model = cls()
+        ckpt = get_ckpt_path(name)
+        model.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+        return model
+
+    def forward(self, input, target):
+        in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
+        outs0, outs1 = self.net(in0_input), self.net(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+        for kk in range(len(self.chns)):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+
+        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))]
+        val = res[0]
+        for l in range(1, len(self.chns)):
+            val += res[l]
+        return val
+
+############
+# The loss #
+############
+
+def adopt_weight(weight, global_step, threshold=0, value=0.):
+    if global_step < threshold:
+        weight = value
+    return weight
+
+def hinge_d_loss(logits_real, logits_fake):
+    loss_real = torch.mean(F.relu(1. - logits_real))
+    loss_fake = torch.mean(F.relu(1. + logits_fake))
+    d_loss = 0.5 * (loss_real + loss_fake)
+    return d_loss
+
+def vanilla_d_loss(logits_real, logits_fake):
+    d_loss = 0.5 * (
+        torch.mean(torch.nn.functional.softplus(-logits_real)) +
+        torch.mean(torch.nn.functional.softplus(logits_fake)))
+    return d_loss
+
+class LPIPSWithDiscriminator(nn.Module):
+    def __init__(self, disc_start, logvar_init=0.0, kl_weight=1.0, pixelloss_weight=1.0,
+                 disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0,
+                 perceptual_weight=1.0, use_actnorm=False, disc_conditional=False,
+                 disc_loss="hinge"):
+
+        super().__init__()
+        assert disc_loss in ["hinge", "vanilla"]
+        self.kl_weight = kl_weight
+        self.pixel_weight = pixelloss_weight
+        self.perceptual_loss = LPIPS().eval()
+        self.perceptual_weight = perceptual_weight
+        # output log variance
+        self.logvar = nn.Parameter(torch.ones(size=()) * logvar_init)
+
+        self.discriminator = NLayerDiscriminator(input_nc=disc_in_channels,
+                                                 n_layers=disc_num_layers,
+                                                 use_actnorm=use_actnorm
+                                                 ).apply(weights_init)
+        self.discriminator_iter_start = disc_start
+        self.disc_loss = hinge_d_loss if disc_loss == "hinge" else vanilla_d_loss
+        self.disc_factor = disc_factor
+        self.discriminator_weight = disc_weight
+        self.disc_conditional = disc_conditional
+
+    def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
+        if last_layer is not None:
+            nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
+            g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
+        else:
+            nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0]
+            g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0]
+
+        d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+        d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
+        d_weight = d_weight * self.discriminator_weight
+        return d_weight
+
+    def forward(self, inputs, reconstructions, posteriors, optimizer_idx,
+                global_step, last_layer=None, cond=None, split="train",
+                weights=None):
+        rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous())
+        if self.perceptual_weight > 0:
+            p_loss = self.perceptual_loss(inputs.contiguous(), reconstructions.contiguous())
+            rec_loss = rec_loss + self.perceptual_weight * p_loss
+
+        nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar
+        weighted_nll_loss = nll_loss
+        if weights is not None:
+            weighted_nll_loss = weights*nll_loss
+        weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]
+        nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
+        kl_loss = posteriors.kl()
+        kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
+
+        # now the GAN part
+        if optimizer_idx == 0:
+            # generator update
+            if cond is None:
+                assert not self.disc_conditional
+                logits_fake = self.discriminator(reconstructions.contiguous())
+            else:
+                assert self.disc_conditional
+                logits_fake = self.discriminator(torch.cat((reconstructions.contiguous(), cond), dim=1))
+            g_loss = -torch.mean(logits_fake)
+
+            if self.disc_factor > 0.0:
+                try:
+                    d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer)
+                except RuntimeError:
+                    assert not self.training
+                    d_weight = torch.tensor(0.0)
+            else:
+                d_weight = torch.tensor(0.0)
+
+            disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start)
+            loss = weighted_nll_loss + self.kl_weight * kl_loss + d_weight * disc_factor * g_loss
+
+            log = {"Loss": loss.clone().detach().mean(), 
+                   "logvar": self.logvar.detach(),
+                   "loss_kl": kl_loss.detach().mean(), 
+                   "loss_nll": nll_loss.detach().mean(),
+                   "loss_rec": rec_loss.detach().mean(),
+                   "d_weight": d_weight.detach(),
+                   "disc_factor": torch.tensor(disc_factor),
+                   "loss_g": g_loss.detach().mean(),
+                   }
+            return loss, log
+
+        if optimizer_idx == 1:
+            # second pass for discriminator update
+            if cond is None:
+                logits_real = self.discriminator(inputs.contiguous().detach())
+                logits_fake = self.discriminator(reconstructions.contiguous().detach())
+            else:
+                logits_real = self.discriminator(torch.cat((inputs.contiguous().detach(), cond), dim=1))
+                logits_fake = self.discriminator(torch.cat((reconstructions.contiguous().detach(), cond), dim=1))
+
+            disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start)
+            d_loss = disc_factor * self.disc_loss(logits_real, logits_fake)
+
+            log = {"Loss": d_loss.clone().detach().mean(),
+                   "loss_disc": d_loss.clone().detach().mean(),
+                   "logits_real": logits_real.detach().mean(),
+                   "logits_fake": logits_fake.detach().mean()
+                   }
+            return d_loss, log