New app first commit

.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,9 @@
+__pycache__
+.vscode/
+src/
+data/
+data
+log/
+log
+pretrained/
+pretrained

README.md

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

app.py

@@ -0,0 +1,1083 @@
+################################################################################
+# 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 self.which == 'v1.0':
+            sd = torch.load('pretrained/vd-four-flow-v1-0.pth', map_location='cpu')
+        net.load_state_dict(sd, strict=False)
+
+        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
+
+        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>'+
+            '<img src="file/assets/demo/misc/mask_inst1.gif" style="float:left;max-width:450px;">'+
+            '<img src="file/assets/demo/misc/mask_inst2.gif" style="float:left;max-width:450px;">'+
+            '<img src="file/assets/demo/misc/mask_inst3.gif" style="float:left;max-width:450px;">',)
+
+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>'+
+            '<img src="file/assets/demo/misc/mask_inst1.gif" style="float:left;max-width:450px;">'+
+            '<img src="file/assets/demo/misc/mask_inst2.gif" style="float:left;max-width:450px;">'+
+            '<img src="file/assets/demo/misc/mask_inst3.gif" style="float:left;max-width:450px;">',)
+
+###########
+# 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.84, 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;
+    }
+"""
+
+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>
+            """.format(' '+vd_inference.which))
+            # .format('')) #
+
+        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: center; max-width: 1200px; margin: 20px auto;">
+            <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>
+            """)
+
+    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,102 @@
+
+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
+  pth: pretrained/optimus-vae.pth
+  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

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

gradio_cached_examples/12/Image Result/23645e03-6435-4819-a746-2840be976419/captions.json

@@ -0,0 +1 @@
+{"/home/james/Project/vd-demo/gradio_cached_examples/12/Image Result/23645e03-6435-4819-a746-2840be976419/tmp9xugbhobbnp5ds0r.png": null, "/home/james/Project/vd-demo/gradio_cached_examples/12/Image Result/23645e03-6435-4819-a746-2840be976419/tmp0m_lns_xtd2zm06b.png": null}

gradio_cached_examples/12/log.csv

@@ -0,0 +1,2 @@
+Image Result,flag,username,timestamp
+/home/james/Project/vd-demo/gradio_cached_examples/12/Image Result/23645e03-6435-4819-a746-2840be976419,,,2023-02-07 08:44:12.243513

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