code pushed

README.md

@@ -1,7 +1,7 @@
 ---
 title: Prompt-Free Diffusion
 emoji: 👀
-colorFrom: orange
+colorFrom: red
 colorTo: blue
 sdk: gradio
 sdk_version: 3.32.0

app.py

@@ -0,0 +1,494 @@
+################################################################################
+# Copyright (C) 2023 Xingqian Xu - All Rights Reserved                         #
+#                                                                              #
+# Please visit Prompt-Free-Diffusion's arXiv paper for more details, link at   #
+# arxiv.org/abs/2305.16223                                                     #
+#                                                                              #
+################################################################################
+
+import gradio as gr
+import os.path as osp
+from PIL import Image
+import numpy as np
+import time
+
+import torch
+import torchvision.transforms as tvtrans
+from lib.cfg_helper import model_cfg_bank
+from lib.model_zoo import get_model
+
+from collections import OrderedDict
+from lib.model_zoo.ddim import DDIMSampler
+
+n_sample_image = 1
+
+controlnet_path = OrderedDict([
+    ['canny'             , ('canny'   , 'pretrained/controlnet/control_sd15_canny_slimmed.safetensors')],
+    ['canny_v11p'        , ('canny'   , 'pretrained/controlnet/control_v11p_sd15_canny_slimmed.safetensors')],
+    ['depth'             , ('depth'   , 'pretrained/controlnet/control_sd15_depth_slimmed.safetensors')],
+    ['hed'               , ('hed'     , 'pretrained/controlnet/control_sd15_hed_slimmed.safetensors')],
+    ['mlsd'              , ('mlsd'    , 'pretrained/controlnet/control_sd15_mlsd_slimmed.safetensors')],
+    ['mlsd_v11p'         , ('mlsd'    , 'pretrained/controlnet/control_v11p_sd15_mlsd_slimmed.safetensors')],
+    ['normal'            , ('normal'  , 'pretrained/controlnet/control_sd15_normal_slimmed.safetensors')],
+    ['openpose'          , ('openpose', 'pretrained/controlnet/control_sd15_openpose_slimmed.safetensors')],
+    ['openpose_v11p'     , ('openpose', 'pretrained/controlnet/control_v11p_sd15_openpose_slimmed.safetensors')],
+    ['scribble'          , ('scribble', 'pretrained/controlnet/control_sd15_scribble_slimmed.safetensors')],
+    ['softedge_v11p'     , ('scribble', 'pretrained/controlnet/control_v11p_sd15_softedge_slimmed.safetensors')],
+    ['seg'               , ('none'    , 'pretrained/controlnet/control_sd15_seg_slimmed.safetensors')],
+    ['lineart_v11p'      , ('none'    , 'pretrained/controlnet/control_v11p_sd15_lineart_slimmed.safetensors')],
+    ['lineart_anime_v11p', ('none'    , 'pretrained/controlnet/control_v11p_sd15s2_lineart_anime_slimmed.safetensors')],
+])
+
+preprocess_method = [
+    'canny'                ,
+    'depth'                ,
+    'hed'                  ,
+    'mlsd'                 ,
+    'normal'               ,
+    'openpose'             ,
+    'openpose_withface'    ,
+    'openpose_withfacehand',
+    'scribble'             ,
+    'none'                 ,
+]
+
+diffuser_path = OrderedDict([
+    ['SD-v1.5'             , 'pretrained/pfd/diffuser/SD-v1-5.safetensors'],
+    ['OpenJouney-v4'       , 'pretrained/pfd/diffuser/OpenJouney-v4.safetensors'],
+    ['Deliberate-v2.0'     , 'pretrained/pfd/diffuser/Deliberate-v2-0.safetensors'],
+    ['RealisticVision-v2.0', 'pretrained/pfd/diffuser/RealisticVision-v2-0.safetensors'],
+    ['Anything-v4'         , 'pretrained/pfd/diffuser/Anything-v4.safetensors'],
+    ['Oam-v3'              , 'pretrained/pfd/diffuser/AbyssOrangeMix-v3.safetensors'],
+    ['Oam-v2'              , 'pretrained/pfd/diffuser/AbyssOrangeMix-v2.safetensors'],
+])
+
+ctxencoder_path = OrderedDict([
+    ['SeeCoder'      , 'pretrained/pfd/seecoder/seecoder-v1-0.safetensors'],
+    ['SeeCoder-PA'   , 'pretrained/pfd/seecoder/seecoder-pa-v1-0.safetensors'],
+    ['SeeCoder-Anime', 'pretrained/pfd/seecoder/seecoder-anime-v1-0.safetensors'],
+])
+
+##########
+# helper #
+##########
+
+def highlight_print(info):
+    print('')
+    print(''.join(['#']*(len(info)+4)))
+    print('# '+info+' #')
+    print(''.join(['#']*(len(info)+4)))
+    print('')
+
+def load_sd_from_file(target):
+    if osp.splitext(target)[-1] == '.ckpt':
+        sd = torch.load(target, map_location='cpu')['state_dict']
+    elif osp.splitext(target)[-1] == '.pth':
+        sd = torch.load(target, map_location='cpu')
+    elif osp.splitext(target)[-1] == '.safetensors':
+        from safetensors.torch import load_file as stload
+        sd = OrderedDict(stload(target, device='cpu'))
+    else:
+        assert False, "File type must be .ckpt or .pth or .safetensors"
+    return sd
+
+########
+# main #
+########
+
+class prompt_free_diffusion(object):
+    def __init__(self, 
+                 fp16=False, 
+                 tag_ctx=None,
+                 tag_diffuser=None,
+                 tag_ctl=None,):
+
+        self.tag_ctx = tag_ctx
+        self.tag_diffuser = tag_diffuser
+        self.tag_ctl = tag_ctl
+        self.strict_sd = True
+
+        cfgm = model_cfg_bank()('pfd_seecoder_with_controlnet')
+        self.net = get_model()(cfgm)
+        
+        self.action_load_ctx(tag_ctx)
+        self.action_load_diffuser(tag_diffuser)
+        self.action_load_ctl(tag_ctl)
+ 
+        if fp16:
+            highlight_print('Running in FP16')
+            self.net.ctx['image'].fp16 = True
+            self.net = self.net.half()
+            self.dtype = torch.float16
+        else:
+            self.dtype = torch.float32
+
+        self.use_cuda = torch.cuda.is_available()
+        if self.use_cuda:
+            self.net.to('cuda')
+
+        self.net.eval()
+        self.sampler = DDIMSampler(self.net)
+
+        self.n_sample_image = n_sample_image
+        self.ddim_steps = 50
+        self.ddim_eta = 0.0
+        self.image_latent_dim = 4
+
+    def load_ctx(self, pretrained):
+        sd = load_sd_from_file(pretrained)
+        sd_extra = [(ki, vi) for ki, vi in self.net.state_dict().items() \
+            if ki.find('ctx.')!=0]
+        sd.update(OrderedDict(sd_extra))
+
+        self.net.load_state_dict(sd, strict=True)
+        print('Load context encoder from [{}] strict [{}].'.format(pretrained, True))
+
+    def load_diffuser(self, pretrained):
+        sd = load_sd_from_file(pretrained)
+        if len([ki for ki in sd.keys() if ki.find('diffuser.image.context_blocks.')==0]) == 0:
+            sd = [(
+                ki.replace('diffuser.text.context_blocks.', 'diffuser.image.context_blocks.'), vi) 
+                    for ki, vi in sd.items()]
+            sd = OrderedDict(sd)
+        sd_extra = [(ki, vi) for ki, vi in self.net.state_dict().items() \
+            if ki.find('diffuser.')!=0]
+        sd.update(OrderedDict(sd_extra))
+        self.net.load_state_dict(sd, strict=True)
+        print('Load diffuser from [{}] strict [{}].'.format(pretrained, True))
+
+    def load_ctl(self, pretrained):
+        sd = load_sd_from_file(pretrained)
+        self.net.ctl.load_state_dict(sd, strict=True)
+        print('Load controlnet from [{}] strict [{}].'.format(pretrained, True))
+
+    def action_load_ctx(self, tag):
+        pretrained = ctxencoder_path[tag]
+        if tag == 'SeeCoder-PA':
+            from lib.model_zoo.seecoder import PPE_MLP
+            pe_layer = \
+                PPE_MLP(freq_num=20, freq_max=None, out_channel=768, mlp_layer=3)
+            if self.dtype == torch.float16:
+                pe_layer = pe_layer.half()
+            if self.use_cuda:
+                pe_layer.to('cuda')
+            pe_layer.eval()
+            self.net.ctx['image'].qtransformer.pe_layer = pe_layer
+        else:
+            self.net.ctx['image'].qtransformer.pe_layer = None
+        if pretrained is not None:
+            self.load_ctx(pretrained)
+        self.tag_ctx = tag
+        return tag
+
+    def action_load_diffuser(self, tag):
+        pretrained = diffuser_path[tag]
+        if pretrained is not None:
+            self.load_diffuser(pretrained)
+        self.tag_diffuser = tag
+        return tag
+
+    def action_load_ctl(self, tag):
+        pretrained = controlnet_path[tag][1]
+        if pretrained is not None:
+            self.load_ctl(pretrained)
+        self.tag_ctl = tag
+        return tag
+
+    def action_autoset_hw(self, imctl):
+        if imctl is None:
+            return 512, 512
+        w, h = imctl.size
+        w = w//64 * 64
+        h = h//64 * 64
+        w = w if w >=512 else 512
+        w = w if w <=1536 else 1536
+        h = h if h >=512 else 512
+        h = h if h <=1536 else 1536
+        return h, w
+
+    def action_autoset_method(self, tag):
+        return controlnet_path[tag][0]
+
+    def action_inference(
+            self, im, imctl, ctl_method, do_preprocess, 
+            h, w, ugscale, seed, 
+            tag_ctx, tag_diffuser, tag_ctl,):
+
+        if tag_ctx != self.tag_ctx:
+            self.action_load_ctx(tag_ctx)
+        if tag_diffuser != self.tag_diffuser:
+            self.action_load_diffuser(tag_diffuser)
+        if tag_ctl != self.tag_ctl:
+            self.action_load_ctl(tag_ctl)
+
+        n_samples = self.n_sample_image
+
+        sampler = self.sampler
+        device = self.net.device
+
+        w = w//64 * 64
+        h = h//64 * 64
+        if imctl is not None:
+            imctl = imctl.resize([w, h], Image.Resampling.BICUBIC)
+
+        craw = tvtrans.ToTensor()(im)[None].to(device).to(self.dtype)
+        c = self.net.ctx_encode(craw, which='image').repeat(n_samples, 1, 1)
+        u = torch.zeros_like(c)
+
+        if tag_ctx in ["SeeCoder-Anime"]:
+            u = torch.load('assets/anime_ug.pth')[None].to(device).to(self.dtype)
+            pad = c.size(1) - u.size(1)
+            u = torch.cat([u, torch.zeros_like(u[:, 0:1].repeat(1, pad, 1))], axis=1)
+
+        if tag_ctl != 'none':
+            ccraw = tvtrans.ToTensor()(imctl)[None].to(device).to(self.dtype)
+            if do_preprocess:
+                cc = self.net.ctl.preprocess(ccraw, type=ctl_method, size=[h, w])
+                cc = cc.to(self.dtype)
+            else:
+                cc = ccraw
+        else:
+            cc = None
+
+        shape = [n_samples, self.image_latent_dim, h//8, w//8]
+
+        if seed < 0:
+            np.random.seed(int(time.time()))
+            torch.manual_seed(-seed + 100)
+        else:
+            np.random.seed(seed + 100)
+            torch.manual_seed(seed)
+
+        x, _ = sampler.sample(
+            steps=self.ddim_steps,
+            x_info={'type':'image',},
+            c_info={'type':'image', 'conditioning':c, 'unconditional_conditioning':u, 
+                    'unconditional_guidance_scale':ugscale,
+                    'control':cc,},
+            shape=shape,
+            verbose=False,
+            eta=self.ddim_eta)
+
+        ccout = [tvtrans.ToPILImage()(i) for i in cc] if cc is not None else []
+        imout = self.net.vae_decode(x, which='image')
+        imout = [tvtrans.ToPILImage()(i) for i in imout]
+        return imout + ccout
+
+pfd_inference = prompt_free_diffusion(
+    fp16=True, tag_ctx = 'SeeCoder', tag_diffuser = 'Deliberate-v2.0', tag_ctl = 'canny',)
+
+#################
+# sub interface #
+#################
+
+cache_examples = True
+
+def get_example():
+    case = [
+        [
+            'assets/examples/ghibli-input.jpg', 
+            'assets/examples/ghibli-canny.png', 
+            'canny', False, 
+            768, 1024, 1.8, 23, 
+            'SeeCoder', 'Deliberate-v2.0', 'canny', ],
+        [
+            'assets/examples/astronautridinghouse-input.jpg', 
+            'assets/examples/astronautridinghouse-canny.png', 
+            'canny', False, 
+            512, 768, 2.0, 21, 
+            'SeeCoder', 'Deliberate-v2.0', 'canny', ],
+        [
+            'assets/examples/grassland-input.jpg', 
+            'assets/examples/grassland-scribble.png', 
+            'scribble', False, 
+            768, 512, 2.0, 41, 
+            'SeeCoder', 'Deliberate-v2.0', 'scribble', ],
+        [
+            'assets/examples/jeep-input.jpg', 
+            'assets/examples/jeep-depth.png', 
+            'depth', False, 
+            512, 768, 2.0, 30, 
+            'SeeCoder', 'Deliberate-v2.0', 'depth', ],
+        [
+            'assets/examples/bedroom-input.jpg', 
+            'assets/examples/bedroom-mlsd.png', 
+            'mlsd', False, 
+            512, 512, 2.0, 31, 
+            'SeeCoder', 'Deliberate-v2.0', 'mlsd', ],
+        [
+            'assets/examples/nightstreet-input.jpg', 
+            'assets/examples/nightstreet-canny.png', 
+            'canny', False, 
+            768, 512, 2.3, 20, 
+            'SeeCoder', 'Deliberate-v2.0', 'canny', ],
+        [
+            'assets/examples/woodcar-input.jpg', 
+            'assets/examples/woodcar-depth.png', 
+            'depth', False, 
+            768, 512, 2.0, 20, 
+            'SeeCoder', 'Deliberate-v2.0', 'depth', ],
+        [
+            'assets/examples-anime/miku.jpg', 
+            'assets/examples-anime/miku-canny.png', 
+            'canny', False, 
+            768, 576, 1.5, 22, 
+            'SeeCoder-Anime', 'Anything-v4', 'canny', ],
+        [
+            'assets/examples-anime/random0.jpg', 
+            'assets/examples-anime/pose.png', 
+            'openpose', False, 
+            768, 1536, 2.0, 41, 
+            'SeeCoder-Anime', 'Oam-v2', 'openpose_v11p', ],
+        [
+            'assets/examples-anime/random1.jpg', 
+            'assets/examples-anime/pose.png', 
+            'openpose', False, 
+            768, 1536, 2.5, 28, 
+            'SeeCoder-Anime', 'Oam-v2', 'openpose_v11p', ], 
+        [
+            'assets/examples-anime/camping.jpg', 
+            'assets/examples-anime/pose.png', 
+            'openpose', False, 
+            768, 1536, 2.0, 35, 
+            'SeeCoder-Anime', 'Anything-v4', 'openpose_v11p', ],
+        [
+            'assets/examples-anime/hanfu_girl.jpg', 
+            'assets/examples-anime/pose.png', 
+            'openpose', False, 
+            768, 1536, 2.0, 20, 
+            'SeeCoder-Anime', 'Anything-v4', 'openpose_v11p', ],
+    ]
+    return case
+
+def interface():
+    with gr.Row():
+        with gr.Column():
+            img_input = gr.Image(label='Image Input', type='pil', elem_id='customized_imbox')
+            with gr.Row():
+                out_width  = gr.Slider(label="Width" , minimum=512, maximum=1536, value=512, step=64, visible=True)
+                out_height = gr.Slider(label="Height", minimum=512, maximum=1536, value=512, step=64, visible=True)
+            with gr.Row():
+                scl_lvl = gr.Slider(label="CFGScale", minimum=0, maximum=10, value=2, step=0.01, visible=True)
+                seed = gr.Number(20, label="Seed", precision=0)
+            with gr.Row():
+                tag_ctx = gr.Dropdown(label='Context Encoder', choices=[pi for pi in ctxencoder_path.keys()], value='SeeCoder')
+                tag_diffuser = gr.Dropdown(label='Diffuser', choices=[pi for pi in diffuser_path.keys()], value='Deliberate-v2.0')
+            button = gr.Button("Run")
+        with gr.Column():
+            ctl_input = gr.Image(label='Control Input', type='pil', elem_id='customized_imbox')
+            do_preprocess = gr.Checkbox(label='Preprocess', value=False)
+            with gr.Row():
+                ctl_method = gr.Dropdown(label='Preprocess Type', choices=preprocess_method, value='canny')
+                tag_ctl    = gr.Dropdown(label='ControlNet',      choices=[pi for pi in controlnet_path.keys()], value='canny')
+        with gr.Column():
+            img_output = gr.Gallery(label="Image Result", elem_id='customized_imbox').style(grid=n_sample_image+1)
+
+    tag_ctl.change(
+        pfd_inference.action_autoset_method,
+        inputs = [tag_ctl],
+        outputs = [ctl_method],)
+
+    ctl_input.change(
+        pfd_inference.action_autoset_hw,
+        inputs = [ctl_input],
+        outputs = [out_height, out_width],)
+
+    # tag_ctx.change(
+    #     pfd_inference.action_load_ctx,
+    #     inputs = [tag_ctx],
+    #     outputs = [tag_ctx],)
+
+    # tag_diffuser.change(
+    #     pfd_inference.action_load_diffuser,
+    #     inputs = [tag_diffuser],
+    #     outputs = [tag_diffuser],)
+
+    # tag_ctl.change(
+    #     pfd_inference.action_load_ctl,
+    #     inputs = [tag_ctl],
+    #     outputs = [tag_ctl],)
+
+    button.click(
+        pfd_inference.action_inference,
+        inputs=[img_input, ctl_input, ctl_method, do_preprocess, 
+                out_height, out_width, scl_lvl, seed, 
+                tag_ctx, tag_diffuser, tag_ctl, ],
+        outputs=[img_output])
+    
+    gr.Examples(
+        label='Examples', 
+        examples=get_example(), 
+        fn=pfd_inference.action_inference,
+        inputs=[img_input, ctl_input, ctl_method, do_preprocess,
+                out_height, out_width, scl_lvl, seed, 
+                tag_ctx, tag_diffuser, tag_ctl, ],
+        outputs=[img_output],
+        cache_examples=cache_examples,)
+
+#############
+# 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">
+                Prompt-Free Diffusion
+            </h1>
+            </div>
+            """)
+
+        interface()
+
+        # 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>
+        #     </div>
+        #     """.format(' '+str(pfd_inference.pretrained)))
+
+    # demo.launch(server_name="0.0.0.0", server_port=7992)
+    # demo.launch()
+    demo.launch(debug=True)

configs/model/autokl.yaml

@@ -0,0 +1,26 @@
+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
+
+autokl_v2:
+  super_cfg: autokl_v1
+  pth: pretrained/pfd/vae/sd-v2-0-base-autokl.pth

configs/model/clip.yaml

@@ -0,0 +1,12 @@
+################
+# clip for sd1 #
+################
+
+clip:
+  symbol: clip
+  args: {}
+
+clip_text_context_encoder_sdv1:
+  super_cfg: clip
+  type: clip_text_context_encoder_sdv1
+  args: {}

configs/model/controlnet.yaml

@@ -0,0 +1,18 @@
+controlnet:
+  symbol: controlnet
+  type: controlnet
+  find_unused_parameters: false
+  args:
+    image_size: 32 # unused
+    in_channels: 4
+    hint_channels: 3
+    model_channels: 320
+    attention_resolutions: [ 4, 2, 1 ]
+    num_res_blocks: 2
+    channel_mult: [ 1, 2, 4, 4 ]
+    num_heads: 8
+    use_spatial_transformer: True
+    transformer_depth: 1
+    context_dim: 768
+    use_checkpoint: True
+    legacy: False

configs/model/openai_unet.yaml

@@ -0,0 +1,35 @@
+openai_unet_sd:
+  type: openai_unet
+  args:
+    image_size: null # no use
+    in_channels: 4
+    out_channels: 4
+    model_channels: 320
+    attention_resolutions: [ 4, 2, 1 ]
+    num_res_blocks: [ 2, 2, 2, 2 ]
+    channel_mult: [ 1, 2, 4, 4 ]
+    # disable_self_attentions: [ False, False, False, False ]  # converts the self-attention to a cross-attention layer if true
+    num_heads: 8
+    use_spatial_transformer: True
+    transformer_depth: 1
+    context_dim: 768
+    use_checkpoint: True
+    legacy: False
+
+#########
+# 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: False
+    parts: [global, data, context]

configs/model/pfd.yaml

@@ -0,0 +1,33 @@
+pfd_base:
+  symbol: pfd
+  find_unused_parameters: true
+  type: pfd
+  args:
+    beta_linear_start: 0.00085
+    beta_linear_end: 0.012
+    timesteps: 1000
+    use_ema: false
+
+pfd_seecoder:
+  super_cfg: pfd_base
+  args:
+    vae_cfg_list:
+      - [image, MODEL(autokl_v2)]
+    ctx_cfg_list:
+      - [image, MODEL(seecoder)]
+    diffuser_cfg_list:
+      - [image, MODEL(openai_unet_2d_v1)]
+    latent_scale_factor: 
+      image: 0.18215
+
+pdf_seecoder_pa:
+  super_cfg: pfd_seecoder
+  args:
+    ctx_cfg_list:
+      - [image, MODEL(seecoder_pa)]
+
+pfd_seecoder_with_controlnet:
+  super_cfg: pfd_seecoder
+  type: pfd_with_control
+  args:
+    ctl_cfg: MODEL(controlnet)

configs/model/seecoder.yaml

@@ -0,0 +1,62 @@
+seecoder_base:
+  symbol: seecoder
+  args: {}
+
+seecoder:
+  super_cfg: seecoder_base
+  type: seecoder
+  args:
+    imencoder_cfg : MODEL(swin_large)
+    imdecoder_cfg : MODEL(seecoder_decoder)
+    qtransformer_cfg : MODEL(seecoder_query_transformer)
+
+seecoder_pa:
+  super_cfg: seet
+  type: seecoder
+  args:
+    imencoder_cfg : MODEL(swin_large)
+    imdecoder_cfg : MODEL(seecoder_decoder)
+    qtransformer_cfg : MODEL(seecoder_query_transformer_position_aware)
+
+###########
+# decoder #
+###########
+
+seecoder_decoder:
+  super_cfg: seecoder_base
+  type: seecoder_decoder
+  args:
+    inchannels:
+      res3:  384
+      res4:  768
+      res5: 1536
+    trans_input_tags: ['res3', 'res4', 'res5']
+    trans_dim: 768
+    trans_dropout: 0.1
+    trans_nheads: 8
+    trans_feedforward_dim: 1024
+    trans_num_layers: 6
+
+#####################
+# query_transformer #
+#####################
+
+seecoder_query_transformer:
+  super_cfg: seecoder_base
+  type: seecoder_query_transformer
+  args:
+    in_channels : 768
+    hidden_dim: 768
+    num_queries: [4, 144]
+    nheads: 8
+    num_layers: 9
+    feedforward_dim: 2048
+    pre_norm: False
+    num_feature_levels: 3
+    enforce_input_project: False
+    with_fea2d_pos: false
+
+seecoder_query_transformer_position_aware:
+  super_cfg: seecoder_query_transformer
+  args:
+    with_fea2d_pos: true

configs/model/swin.yaml

@@ -0,0 +1,32 @@
+swin:
+  symbol: swin
+  args: {}
+
+swin_base:
+  super_cfg: swin
+  type: swin
+  args:
+    embed_dim: 128
+    depths: [ 2, 2, 18, 2 ]
+    num_heads: [ 4, 8, 16, 32 ]
+    window_size: 7
+    ape: False
+    drop_path_rate: 0.3
+    patch_norm: True
+  pretrained: pretrained/swin/swin_base_patch4_window7_224_22k.pth
+  strict_sd: False
+
+swin_large:
+  super_cfg: swin
+  type: swin
+  args:
+    embed_dim: 192
+    depths: [ 2, 2, 18, 2 ]
+    num_heads: [ 6, 12, 24, 48 ]
+    window_size: 12
+    ape: False
+    drop_path_rate: 0.3
+    patch_norm: True
+  pretrained: pretrained/swin/swin_large_patch4_window12_384_22k.pth
+  strict_sd: False
+

lib/cfg_helper.py

@@ -0,0 +1,666 @@
+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:
+            return osp.join(
+                self.cfg_dir, 'clip.yaml')
+        elif name.find('autokl')==0:
+            return osp.join(
+                self.cfg_dir, 'autokl.yaml')
+        elif name.find('controlnet')==0:
+            return osp.join(
+                self.cfg_dir, 'controlnet.yaml')
+        elif name.find('swin')==0:
+            return osp.join(
+                self.cfg_dir, 'swin.yaml')
+        elif name.find('pfd')==0:
+            return osp.join(
+                self.cfg_dir, 'pfd.yaml')
+        elif name.find('seecoder')==0:
+            return osp.join(
+                self.cfg_dir, 'seecoder.yaml')
+        else:
+            raise ValueError
+
+class dataset_cfg_bank(object):
+    def __init__(self):
+        self.cfg_dir = osp.join('configs', 'dataset')
+        self.cfg_bank = edict()
+
+    def __call__(self, name):
+        if name not in self.cfg_bank:
+            cfg_path = self.get_yaml_path(name)
+            with open(cfg_path, 'r') as f:
+                cfg_new = yaml.load(
+                    f, Loader=yaml.FullLoader)
+            cfg_new = edict(cfg_new)
+            self.cfg_bank.update(cfg_new)
+
+        cfg = self.cfg_bank[name]
+        cfg.name = name
+        if cfg.get('super_cfg', None) is None:
+            cfg = cfg_solve(cfg, cfg)
+            self.cfg_bank[name] = cfg
+            return copy.deepcopy(cfg)
+
+        super_cfg = self.__call__(cfg.super_cfg)
+        super_cfg.update(cfg)
+        cfg = super_cfg
+        cfg.super_cfg = None
+        try:
+            delete = cfg.pop('delete')
+        except:
+            delete = []
+
+        for dargs in delete:
+            cfg.pop(dargs)
+
+        cfg = cfg_solve(cfg, cfg)
+        self.cfg_bank[name] = cfg
+        return copy.deepcopy(cfg)
+
+    def get_yaml_path(self, name):
+        if name.find('cityscapes')==0:
+            return osp.join(
+                self.cfg_dir, 'cityscapes.yaml')
+        elif name.find('div2k')==0:
+            return osp.join(
+                self.cfg_dir, 'div2k.yaml')
+        elif name.find('gandiv2k')==0:
+            return osp.join(
+                self.cfg_dir, 'gandiv2k.yaml')
+        elif name.find('srbenchmark')==0:
+            return osp.join(
+                self.cfg_dir, 'srbenchmark.yaml')
+        elif name.find('imagedir')==0:
+            return osp.join(
+                self.cfg_dir, 'imagedir.yaml')
+        elif name.find('places2')==0:
+            return osp.join(
+                self.cfg_dir, 'places2.yaml')
+        elif name.find('ffhq')==0:
+            return osp.join(
+                self.cfg_dir, 'ffhq.yaml')
+        elif name.find('imcpt')==0:
+            return osp.join(
+                self.cfg_dir, 'imcpt.yaml')
+        elif name.find('texture')==0:
+            return osp.join(
+                self.cfg_dir, 'texture.yaml')
+        elif name.find('openimages')==0:
+            return osp.join(
+                self.cfg_dir, 'openimages.yaml')
+        elif name.find('laion2b')==0:
+            return osp.join(
+                self.cfg_dir, 'laion2b.yaml')
+        elif name.find('laionart')==0:
+            return osp.join(
+                self.cfg_dir, 'laionart.yaml')
+        elif name.find('celeba')==0:
+            return osp.join(
+                self.cfg_dir, 'celeba.yaml')
+        elif name.find('coyo')==0:
+            return osp.join(
+                self.cfg_dir, 'coyo.yaml')
+        elif name.find('pafc')==0:
+            return osp.join(
+                self.cfg_dir, 'pafc.yaml')
+        elif name.find('coco')==0:
+            return osp.join(
+                self.cfg_dir, 'coco.yaml')
+        elif name.find('genai')==0:
+            return osp.join(
+                self.cfg_dir, 'genai.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,165 @@
+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
+
+def print_log(*console_info):
+    grank, lrank, _ = sync.get_rank('all')
+    if lrank!=0:
+        return
+
+    console_info = [str(i) for i in console_info]
+    console_info = ' '.join(console_info)
+    print(console_info)
+
+    if grank!=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
+        self.ddp = sync.is_ddp()
+        self.grank, self.lrank, _ = sync.get_rank('all')
+        self.gwsize = sync.get_world_size('global')
+
+        use_tensorboard = cfgt.get('log_tensorboard', False) and (self.grank==0)
+
+        self.tb = None
+        if use_tensorboard:
+            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.lrank)
+        if self.ddp:
+            dist.all_reduce(value_gather_tensor, op=dist.ReduceOp.SUM)
+            value_gather_tensor /= self.gwsize
+
+        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) and (extra['lr'] is not None):
+                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,540 @@
+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
+
+try:
+    import xformers
+    import xformers.ops
+    XFORMERS_IS_AVAILBLE = True
+except:
+    XFORMERS_IS_AVAILBLE = False
+
+
+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.inner_dim = inner_dim
+
+        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)
+
+    def forward_next(self, x, context=None, mask=None):
+        assert mask is None, 'not supported yet'
+        x0 = rearrange(x, 'b n c -> n b c')
+        if context is not None:
+            c0 = rearrange(context, 'b n c -> n b c')
+        else:
+            c0 = x0
+        r, _ = F.multi_head_attention_forward(
+            x0, c0, c0,
+            embed_dim_to_check = self.inner_dim,
+            num_heads = self.heads,
+            in_proj_weight = None, in_proj_bias = None,
+            bias_k = None, bias_v = None,
+            add_zero_attn = False, dropout_p = 0,
+            out_proj_weight = self.to_out[0].weight,
+            out_proj_bias = self.to_out[0].bias,
+            use_separate_proj_weight = True,
+            q_proj_weight = self.to_q.weight,
+            k_proj_weight = self.to_k.weight,
+            v_proj_weight = self.to_v.weight,)
+        r = rearrange(r, 'n b c -> b n c')
+        r = self.to_out[1](r) # dropout
+        return r
+
+
+class MemoryEfficientCrossAttention(nn.Module):
+    # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
+        super().__init__()
+        print(f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using "
+              f"{heads} heads.")
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+
+        self.heads = heads
+        self.dim_head = dim_head
+
+        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))
+        self.attention_op: Optional[Any] = None
+
+    def forward(self, x, context=None, mask=None):
+        q = self.to_q(x)
+        context = default(context, x)
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        b, _, _ = q.shape
+        q, k, v = map(
+            lambda t: t.unsqueeze(3)
+            .reshape(b, t.shape[1], self.heads, self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b * self.heads, t.shape[1], self.dim_head)
+            .contiguous(),
+            (q, k, v),
+        )
+
+        # actually compute the attention, what we cannot get enough of
+        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
+
+        if exists(mask):
+            raise NotImplementedError
+        out = (
+            out.unsqueeze(0)
+            .reshape(b, self.heads, out.shape[1], self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b, out.shape[1], self.heads * self.dim_head)
+        )
+        return self.to_out(out)
+
+
+class BasicTransformerBlock(nn.Module):
+    ATTENTION_MODES = {
+        "softmax": CrossAttention,  # vanilla attention
+        "softmax-xformers": MemoryEfficientCrossAttention
+    }
+    def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True,
+                 disable_self_attn=False):
+        super().__init__()
+        attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILBLE else "softmax"
+        assert attn_mode in self.ATTENTION_MODES
+        attn_cls = self.ATTENTION_MODES[attn_mode]
+        self.disable_self_attn = disable_self_attn
+        self.attn1 = attn_cls(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 = attn_cls(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
+    NEW: use_linear for more efficiency instead of the 1x1 convs
+    """
+    def __init__(self, in_channels, n_heads, d_head,
+                 depth=1, dropout=0., context_dim=None,
+                 disable_self_attn=False, use_linear=False,
+                 use_checkpoint=True):
+        super().__init__()
+        if exists(context_dim) and not isinstance(context_dim, list):
+            context_dim = [context_dim]
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = Normalize(in_channels)
+        if not use_linear:
+            self.proj_in = nn.Conv2d(in_channels,
+                                     inner_dim,
+                                     kernel_size=1,
+                                     stride=1,
+                                     padding=0)
+        else:
+            self.proj_in = nn.Linear(in_channels, inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim[d],
+                                   disable_self_attn=disable_self_attn, checkpoint=use_checkpoint)
+                for d in range(depth)]
+        )
+        if not use_linear:
+            self.proj_out = zero_module(nn.Conv2d(inner_dim,
+                                                  in_channels,
+                                                  kernel_size=1,
+                                                  stride=1,
+                                                  padding=0))
+        else:
+            self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
+        self.use_linear = use_linear
+
+    def forward(self, x, context=None):
+        # note: if no context is given, cross-attention defaults to self-attention
+        if not isinstance(context, list):
+            context = [context]
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        if not self.use_linear:
+            x = self.proj_in(x)
+        x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
+        if self.use_linear:
+            x = self.proj_in(x)
+        for i, block in enumerate(self.transformer_blocks):
+            x = block(x, context=context[i])
+        if self.use_linear:
+            x = self.proj_out(x)
+        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
+        if not self.use_linear:
+            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,166 @@
+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):
+        dec = self.decode_trainable(z)
+        dec = torch.clamp(dec, 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
+
+@register('autoencoderkl_customnorm')
+class AutoencoderKL_CustomNorm(AutoencoderKL):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.mean = torch.Tensor([0.48145466, 0.4578275, 0.40821073])
+        self.std  = torch.Tensor([0.26862954, 0.26130258, 0.27577711])
+
+    def encode_trainable(self, x, out_posterior=False):
+        m = self.mean[None, :, None, None].to(z.device).to(z.dtype)
+        s = self.std[None, :, None, None].to(z.device).to(z.dtype)
+        x = (x-m)/s
+        h = self.encoder(x)
+        moments = self.quant_conv(h)
+        posterior = DiagonalGaussianDistribution(moments)
+        if out_posterior:
+            return posterior
+        else:
+            return posterior.sample()
+
+    def decode_trainable(self, z):
+        m = self.mean[None, :, None, None].to(z.device).to(z.dtype)
+        s = self.std[None, :, None, None].to(z.device).to(z.dtype)
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+        dec = (dec+1)/2
+        return dec

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

lib/model_zoo/clip.py

@@ -0,0 +1,788 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from functools import partial
+from lib.model_zoo.common.get_model import register
+
+symbol = 'clip'
+
+class AbstractEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError
+
+from transformers import CLIPTokenizer, CLIPTextModel
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+@register('clip_text_context_encoder_sdv1')
+class CLIPTextContextEncoderSDv1(AbstractEncoder):
+    """Uses the CLIP transformer encoder for text (from huggingface)"""
+    def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77, freeze=True):  # clip-vit-base-patch32
+        super().__init__()
+        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.transformer = CLIPTextModel.from_pretrained(version)
+        self.device = device
+        self.max_length = max_length
+        if freeze:
+            self.freeze()
+
+    def freeze(self):
+        self.transformer = self.transformer.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        with torch.no_grad():
+            batch_encoding = self.tokenizer(
+                text, truncation=True, max_length=self.max_length, return_length=True,
+                return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+            tokens = batch_encoding["input_ids"].to(self.device)
+        max_token_n = self.transformer.text_model.embeddings.position_ids.shape[1]
+        positional_ids = torch.arange(max_token_n)[None].to(self.device)
+        outputs = self.transformer(
+            input_ids=tokens, 
+            position_ids=positional_ids, )
+        z = outputs.last_hidden_state
+        return z
+
+    def encode(self, text):
+        return self(text)
+
+#############################
+# copyed from justin's code #
+#############################
+
+@register('clip_image_context_encoder_justin')
+class CLIPImageContextEncoderJustin(AbstractEncoder):
+    """
+        Uses the CLIP image encoder.
+        """
+    def __init__(
+            self,
+            model='ViT-L/14',
+            jit=False,
+            device='cuda' if torch.cuda.is_available() else 'cpu',
+            antialias=False,
+        ):
+        super().__init__()
+        from . import clip_justin
+        self.model, _ = clip_justin.load(name=model, device=device, jit=jit)
+        self.device = device
+        self.antialias = antialias
+
+        self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
+        self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
+
+        # I didn't call this originally, but seems like it was frozen anyway
+        self.freeze()
+
+    def freeze(self):
+        self.transformer = self.model.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def preprocess(self, x):
+        import kornia
+        # Expects inputs in the range -1, 1
+        x = kornia.geometry.resize(x, (224, 224),
+                                   interpolation='bicubic',align_corners=True,
+                                   antialias=self.antialias)
+        x = (x + 1.) / 2.
+        # renormalize according to clip
+        x = kornia.enhance.normalize(x, self.mean, self.std)
+        return x
+
+    def forward(self, x):
+        # x is assumed to be in range [-1,1]
+        return self.model.encode_image(self.preprocess(x)).float()
+
+    def encode(self, im):
+        return self(im).unsqueeze(1)
+
+###############
+# for vd next #
+###############
+
+from transformers import CLIPModel
+
+@register('clip_text_context_encoder')
+class CLIPTextContextEncoder(AbstractEncoder):
+    def __init__(self, 
+                 version="openai/clip-vit-large-patch14", 
+                 max_length=77, 
+                 fp16=False, ):
+        super().__init__()
+        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.model = CLIPModel.from_pretrained(version)
+        self.max_length = max_length
+        self.fp16 = fp16
+        self.freeze()
+
+    def get_device(self):
+        # A trick to get device
+        return self.model.text_projection.weight.device
+
+    def freeze(self):
+        self.model = self.model.eval()
+        self.train = disabled_train
+        for param in self.parameters():
+            param.requires_grad = False
+        
+    def encode(self, text):
+        batch_encoding = self.tokenizer(
+            text, truncation=True, max_length=self.max_length, return_length=True,
+            return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"].to(self.get_device())
+        outputs = self.model.text_model(input_ids=tokens)
+        z = self.model.text_projection(outputs.last_hidden_state)
+        z_pooled = self.model.text_projection(outputs.pooler_output)
+        z = z / torch.norm(z_pooled.unsqueeze(1), dim=-1, keepdim=True)
+        return z
+
+from transformers import CLIPProcessor
+
+@register('clip_image_context_encoder')
+class CLIPImageContextEncoder(AbstractEncoder):
+    def __init__(self, 
+                 version="openai/clip-vit-large-patch14", 
+                 fp16=False, ):
+        super().__init__()
+        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.processor = CLIPProcessor.from_pretrained(version)
+        self.model = CLIPModel.from_pretrained(version)
+        self.fp16 = fp16
+        self.freeze()
+
+    def get_device(self):
+        # A trick to get device
+        return self.model.text_projection.weight.device
+
+    def freeze(self):
+        self.model = self.model.eval()
+        self.train = disabled_train
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def _encode(self, images):
+        if isinstance(images, torch.Tensor):
+            import torchvision.transforms as tvtrans
+            images = [tvtrans.ToPILImage()(i) for i in images]
+        inputs = self.processor(images=images, return_tensors="pt")
+        pixels = inputs['pixel_values'].half() if self.fp16 else inputs['pixel_values']
+        pixels = pixels.to(self.get_device())
+        outputs = self.model.vision_model(pixel_values=pixels)
+        z = outputs.last_hidden_state
+        z = self.model.vision_model.post_layernorm(z)
+        z = self.model.visual_projection(z)
+        z_pooled = z[:, 0:1]
+        z = z / torch.norm(z_pooled, dim=-1, keepdim=True)
+        return z
+
+    @torch.no_grad()
+    def _encode_wmask(self, images, masks):
+        assert isinstance(masks, torch.Tensor)
+        assert (len(masks.shape)==4) and (masks.shape[1]==1)
+        masks = torch.clamp(masks, 0, 1)
+        masked_images = images*masks
+        masks = masks.float()
+        masks = F.interpolate(masks, [224, 224], mode='bilinear')
+        if masks.sum() == masks.numel():
+            return self._encode(images)
+
+        device = images.device
+        dtype = images.dtype
+        gscale = masks.mean(axis=[1, 2, 3], keepdim=True).flatten(2)
+
+        vtoken_kernel_size = self.model.vision_model.embeddings.patch_embedding.kernel_size
+        vtoken_stride = self.model.vision_model.embeddings.patch_embedding.stride
+        mask_kernal = torch.ones([1, 1, *vtoken_kernel_size], device=device, requires_grad=False).float()
+        vtoken_mask = torch.nn.functional.conv2d(masks, mask_kernal, stride=vtoken_stride).flatten(2).transpose(1, 2)
+        vtoken_mask = vtoken_mask/np.prod(vtoken_kernel_size)
+        vtoken_mask = torch.concat([gscale, vtoken_mask], axis=1)
+
+        import types
+        def customized_embedding_forward(self, pixel_values):
+            batch_size = pixel_values.shape[0]
+            patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
+            patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+
+            class_embeds = self.class_embedding.expand(batch_size, 1, -1)
+            embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+            embeddings = embeddings + self.position_embedding(self.position_ids)
+            embeddings = embeddings*vtoken_mask.to(embeddings.dtype)
+            return embeddings
+
+        old_forward = self.model.vision_model.embeddings.forward
+        self.model.vision_model.embeddings.forward = types.MethodType(
+            customized_embedding_forward, self.model.vision_model.embeddings)
+
+        z = self._encode(images)
+        self.model.vision_model.embeddings.forward = old_forward
+        z = z * vtoken_mask.to(dtype)
+        return z
+
+    # def _encode_wmask(self, images, masks):
+    #     assert isinstance(masks, torch.Tensor)
+    #     assert (len(masks.shape)==4) and (masks.shape[1]==1)
+    #     masks = torch.clamp(masks, 0, 1)
+    #     masks = masks.float()
+    #     masks = F.interpolate(masks, [224, 224], mode='bilinear')
+    #     if masks.sum() == masks.numel():
+    #         return self._encode(images)
+
+    #     device = images.device
+    #     dtype = images.dtype
+
+    #     vtoken_kernel_size = self.model.vision_model.embeddings.patch_embedding.kernel_size
+    #     vtoken_stride = self.model.vision_model.embeddings.patch_embedding.stride
+    #     mask_kernal = torch.ones([1, 1, *vtoken_kernel_size], device=device, requires_grad=False).float()
+    #     vtoken_mask = torch.nn.functional.conv2d(masks, mask_kernal, stride=vtoken_stride).flatten(2).transpose(1, 2)
+    #     vtoken_mask = vtoken_mask/np.prod(vtoken_kernel_size)
+
+    #     z = self._encode(images)
+    #     z[:, 1:, :] = z[:, 1:, :] * vtoken_mask.to(dtype)
+    #     z[:, 0, :] = 0
+    #     return z
+
+    def encode(self, images, masks=None):
+        if masks is None:
+            return self._encode(images)
+        else:
+            return self._encode_wmask(images, masks)
+
+@register('clip_image_context_encoder_position_agnostic')
+class CLIPImageContextEncoderPA(CLIPImageContextEncoder):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        import types
+        def customized_embedding_forward(self, pixel_values):
+            batch_size = pixel_values.shape[0]
+            patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
+            patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+
+            class_embeds = self.class_embedding.expand(batch_size, 1, -1)
+            embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+            pembeddings = self.position_embedding(self.position_ids)
+            pembeddings = torch.cat([
+                pembeddings[:, 0:1], 
+                pembeddings[:, 1: ].mean(dim=1, keepdim=True).repeat(1, 256, 1)], dim=1)
+            embeddings = embeddings + pembeddings
+            return embeddings
+
+        self.model.vision_model.embeddings.forward = types.MethodType(
+            customized_embedding_forward, self.model.vision_model.embeddings)
+
+##############
+# from sd2.0 #
+##############
+
+import open_clip
+import torch.nn.functional as F
+
+@register('openclip_text_context_encoder_sdv2')
+class FrozenOpenCLIPTextEmbedderSDv2(AbstractEncoder):
+    """
+    Uses the OpenCLIP transformer encoder for text
+    """
+    LAYERS = [
+        #"pooled",
+        "last",
+        "penultimate"
+    ]
+    def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77,
+                 freeze=True, layer="last"):
+        super().__init__()
+        assert layer in self.LAYERS
+        model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), pretrained=version)
+        del model.visual
+        self.model = model
+
+        self.device = device
+        self.max_length = max_length
+        if freeze:
+            self.freeze()
+        self.layer = layer
+        if self.layer == "last":
+            self.layer_idx = 0
+        elif self.layer == "penultimate":
+            self.layer_idx = 1
+        else:
+            raise NotImplementedError()
+
+    def freeze(self):
+        self.model = self.model.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        tokens = open_clip.tokenize(text)
+        z = self.encode_with_transformer(tokens.to(self.device))
+        return z
+
+    def encode_with_transformer(self, text):
+        x = self.model.token_embedding(text)  # [batch_size, n_ctx, d_model]
+        x = x + self.model.positional_embedding
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.model.ln_final(x)
+        return x
+
+    def text_transformer_forward(self, x: torch.Tensor, attn_mask = None):
+        for i, r in enumerate(self.model.transformer.resblocks):
+            if i == len(self.model.transformer.resblocks) - self.layer_idx:
+                break
+            if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(r, x, attn_mask)
+            else:
+                x = r(x, attn_mask=attn_mask)
+        return x
+
+    def encode(self, text):
+        return self(text)
+
+@register('openclip_text_context_encoder')
+class FrozenOpenCLIPTextEmbedder(AbstractEncoder):
+    """
+    Uses the OpenCLIP transformer encoder for text
+    """
+    def __init__(self, 
+                 arch="ViT-H-14", 
+                 version="laion2b_s32b_b79k", 
+                 max_length=77,
+                 freeze=True,):
+        super().__init__()
+        model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), pretrained=version)
+        del model.visual
+        self.model = model
+        self.max_length = max_length
+        self.device = 'cpu'
+        if freeze:
+            self.freeze()
+
+    def to(self, device):
+        self.device = device
+        super().to(device)
+
+    def freeze(self):
+        self.model = self.model.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        self.device = self.model.ln_final.weight.device # urgly trick
+        tokens = open_clip.tokenize(text)
+        z = self.encode_with_transformer(tokens.to(self.device))
+        return z
+
+    def encode_with_transformer(self, text):
+        x = self.model.token_embedding(text)  # [batch_size, n_ctx, d_model]
+        x = x + self.model.positional_embedding
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.model.transformer(x, attn_mask=self.model.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.model.ln_final(x)
+        x_pool = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.model.text_projection
+        # x_pool_debug = F.normalize(x_pool, dim=-1)
+        x = x @ self.model.text_projection
+        x = x / x_pool.norm(dim=1, keepdim=True).unsqueeze(1)
+        return x
+
+    def encode(self, text):
+        return self(text)
+
+@register('openclip_image_context_encoder')
+class FrozenOpenCLIPImageEmbedder(AbstractEncoder):
+    """
+    Uses the OpenCLIP transformer encoder for text
+    """
+    def __init__(self, 
+                 arch="ViT-H-14", 
+                 version="laion2b_s32b_b79k",
+                 freeze=True,):
+        super().__init__()
+        model, _, preprocess = open_clip.create_model_and_transforms(
+            arch, device=torch.device('cpu'), pretrained=version)
+        self.model = model.visual
+        self.device = 'cpu'
+        import torchvision.transforms as tvtrans
+        # we only need resize & normalization
+        preprocess.transforms[0].size = [224, 224] # make it more precise
+        self.preprocess = tvtrans.Compose([
+            preprocess.transforms[0],
+            preprocess.transforms[4],])
+        if freeze:
+            self.freeze()
+
+    def to(self, device):
+        self.device = device
+        super().to(device)
+
+    def freeze(self):
+        self.model = self.model.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, image):
+        z = self.preprocess(image)
+        z = self.encode_with_transformer(z)
+        return z
+
+    def encode_with_transformer(self, image):
+        x = self.model.conv1(image)
+        x = x.reshape(x.shape[0], x.shape[1], -1)
+        x = x.permute(0, 2, 1)
+        x = torch.cat([
+            self.model.class_embedding.to(x.dtype) 
+            + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
+            x], dim=1)
+        x = x + self.model.positional_embedding.to(x.dtype)
+        x = self.model.ln_pre(x)
+        x = x.permute(1, 0, 2)
+        x = self.model.transformer(x)
+        x = x.permute(1, 0, 2)
+
+        x = self.model.ln_post(x)
+        if self.model.proj is not None:
+            x = x @ self.model.proj
+
+        x_pool = x[:, 0, :]
+        # x_pool_debug = self.model(image)
+        # x_pooln_debug = F.normalize(x_pool_debug, dim=-1)
+        x = x / x_pool.norm(dim=1, keepdim=True).unsqueeze(1)
+        return x
+
+    def _encode(self, image):
+        return self(image)
+
+    def _encode_wmask(self, images, masks):
+        z = self._encode(images)
+        device = z.device
+        vtoken_kernel_size = self.model.conv1.kernel_size
+        vtoken_stride = self.model.conv1.stride
+        mask_kernal = torch.ones([1, 1, *vtoken_kernel_size], device=device, dtype=z.dtype, requires_grad=False)
+        mask_kernal /= np.prod(vtoken_kernel_size)
+
+        assert isinstance(masks, torch.Tensor)
+        assert (len(masks.shape)==4) and (masks.shape[1]==1)
+        masks = torch.clamp(masks, 0, 1)
+        masks = F.interpolate(masks, [224, 224], mode='bilinear')
+
+        vtoken_mask = torch.nn.functional.conv2d(1-masks, mask_kernal, stride=vtoken_stride).flatten(2).transpose(1, 2)
+        z[:, 1:, :] = z[:, 1:, :] * vtoken_mask
+        z[:, 0, :] = 0
+        return z
+
+    def encode(self, images, masks=None):
+        if masks is None:
+            return self._encode(images)
+        else:
+            return self._encode_wmask(images, masks)
+
+############################
+# def customized tokenizer #
+############################
+
+from open_clip import SimpleTokenizer
+
+@register('openclip_text_context_encoder_sdv2_customized_tokenizer_v1')
+class FrozenOpenCLIPEmbedderSDv2CustomizedTokenizerV1(FrozenOpenCLIPTextEmbedderSDv2):
+    """
+    Uses the OpenCLIP transformer encoder for text
+    """
+    def __init__(self, customized_tokens, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if isinstance(customized_tokens, str):
+            customized_tokens = [customized_tokens]
+        self.tokenizer = open_clip.SimpleTokenizer(special_tokens=customized_tokens)
+        self.num_regular_tokens = self.model.token_embedding.weight.shape[0] 
+        self.embedding_dim = self.model.ln_final.weight.shape[0]
+        self.customized_token_embedding = nn.Embedding(
+            len(customized_tokens), embedding_dim=self.embedding_dim)
+        nn.init.normal_(self.customized_token_embedding.weight, std=0.02)
+
+    def tokenize(self, texts):
+        if isinstance(texts, str):
+            texts = [texts]
+        sot_token = self.tokenizer.encoder["<start_of_text>"]
+        eot_token = self.tokenizer.encoder["<end_of_text>"]
+        all_tokens = [[sot_token] + self.tokenizer.encode(text) + [eot_token] for text in texts]
+        maxn = self.num_regular_tokens
+        regular_tokens = [[ti if ti < maxn else 0 for ti in tokens] for tokens in all_tokens]
+        token_mask = [[0 if ti < maxn else 1 for ti in tokens] for tokens in all_tokens]
+        customized_tokens = [[ti-maxn if ti >= maxn else 0 for ti in tokens] for tokens in all_tokens]
+        return regular_tokens, customized_tokens, token_mask
+
+    def pad_to_length(self, tokens, context_length=77, eot_token=None):
+        result = torch.zeros(len(tokens), context_length, dtype=torch.long)
+        eot_token = self.tokenizer.encoder["<end_of_text>"] if eot_token is None else eot_token
+        for i, tokens in enumerate(tokens):
+            if len(tokens) > context_length:
+                tokens = tokens[:context_length]  # Truncate
+                tokens[-1] = eot_token
+            result[i, :len(tokens)] = torch.tensor(tokens)
+        return result
+
+    def forward(self, text):
+        self.device = self.model.ln_final.weight.device # urgly trick
+        regular_tokens, customized_tokens, token_mask = self.tokenize(text)
+        regular_tokens = self.pad_to_length(regular_tokens).to(self.device)
+        customized_tokens = self.pad_to_length(customized_tokens, eot_token=0).to(self.device)
+        token_mask = self.pad_to_length(token_mask, eot_token=0).to(self.device)
+        z0 = self.encode_with_transformer(regular_tokens)
+        z1 = self.customized_token_embedding(customized_tokens)
+        token_mask = token_mask[:, :, None].type(z0.dtype)
+        z = z0 * (1-token_mask) + z1 * token_mask
+        return z
+
+@register('openclip_text_context_encoder_sdv2_customized_tokenizer_v2')
+class FrozenOpenCLIPEmbedderSDv2CustomizedTokenizerV2(FrozenOpenCLIPTextEmbedderSDv2):
+    """
+    Uses the OpenCLIP transformer encoder for text
+    """
+    def __init__(self, customized_tokens, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if isinstance(customized_tokens, str):
+            customized_tokens = [customized_tokens]
+        self.tokenizer = open_clip.SimpleTokenizer(special_tokens=customized_tokens)
+        self.num_regular_tokens = self.model.token_embedding.weight.shape[0] 
+        self.embedding_dim = self.model.token_embedding.weight.shape[1]
+        self.customized_token_embedding = nn.Embedding(
+            len(customized_tokens), embedding_dim=self.embedding_dim)
+        nn.init.normal_(self.customized_token_embedding.weight, std=0.02)
+
+    def tokenize(self, texts):
+        if isinstance(texts, str):
+            texts = [texts]
+        sot_token = self.tokenizer.encoder["<start_of_text>"]
+        eot_token = self.tokenizer.encoder["<end_of_text>"]
+        all_tokens = [[sot_token] + self.tokenizer.encode(text) + [eot_token] for text in texts]
+        maxn = self.num_regular_tokens
+        regular_tokens = [[ti if ti < maxn else 0 for ti in tokens] for tokens in all_tokens]
+        token_mask = [[0 if ti < maxn else 1 for ti in tokens] for tokens in all_tokens]
+        customized_tokens = [[ti-maxn if ti >= maxn else 0 for ti in tokens] for tokens in all_tokens]
+        return regular_tokens, customized_tokens, token_mask
+
+    def pad_to_length(self, tokens, context_length=77, eot_token=None):
+        result = torch.zeros(len(tokens), context_length, dtype=torch.long)
+        eot_token = self.tokenizer.encoder["<end_of_text>"] if eot_token is None else eot_token
+        for i, tokens in enumerate(tokens):
+            if len(tokens) > context_length:
+                tokens = tokens[:context_length]  # Truncate
+                tokens[-1] = eot_token
+            result[i, :len(tokens)] = torch.tensor(tokens)
+        return result
+
+    def forward(self, text):
+        self.device = self.model.token_embedding.weight.device # urgly trick
+        regular_tokens, customized_tokens, token_mask = self.tokenize(text)
+        regular_tokens = self.pad_to_length(regular_tokens).to(self.device)
+        customized_tokens = self.pad_to_length(customized_tokens, eot_token=0).to(self.device)
+        token_mask = self.pad_to_length(token_mask, eot_token=0).to(self.device)
+        z = self.encode_with_transformer(regular_tokens, customized_tokens, token_mask)
+        return z
+
+    def encode_with_transformer(self, token, customized_token, token_mask):
+        x0 = self.model.token_embedding(token)
+        x1 = self.customized_token_embedding(customized_token)
+        token_mask = token_mask[:, :, None].type(x0.dtype)
+        x = x0 * (1-token_mask) + x1 * token_mask        
+        x = x + self.model.positional_embedding
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.model.ln_final(x)
+        return x
+
+class ln_freezed_temp(nn.LayerNorm):
+    def forward(self, x):
+        self.weight.requires_grad = False
+        self.bias.requires_grad = False
+        return super().forward(x)
+
+@register('openclip_text_context_encoder_sdv2_customized_tokenizer_v3')
+class FrozenOpenCLIPEmbedderSDv2CustomizedTokenizerV3(FrozenOpenCLIPEmbedderSDv2CustomizedTokenizerV2):
+    """
+    Uses the OpenCLIP transformer encoder for text
+    """
+    def __init__(self, customized_tokens, texpand=4, lora_rank=None, lora_bias_trainable=True, *args, **kwargs):
+        super().__init__(customized_tokens, *args, **kwargs)
+        if isinstance(customized_tokens, str):
+            customized_tokens = [customized_tokens]
+        self.texpand = texpand
+        self.customized_token_embedding = nn.Embedding(
+            len(customized_tokens)*texpand, embedding_dim=self.embedding_dim)
+        nn.init.normal_(self.customized_token_embedding.weight, std=0.02)
+
+        if lora_rank is not None:
+            from .lora import freeze_param, freeze_module, to_lora
+            def convert_resattnblock(module):
+                module.ln_1.__class__ = ln_freezed_temp
+                # freeze_module(module.ln_1)
+                module.attn = to_lora(module.attn, lora_rank, lora_bias_trainable)
+                module.ln_2.__class__ = ln_freezed_temp
+                # freeze_module(module.ln_2)
+                module.mlp.c_fc = to_lora(module.mlp.c_fc, lora_rank, lora_bias_trainable)
+                module.mlp.c_proj = to_lora(module.mlp.c_proj, lora_rank, lora_bias_trainable)
+            freeze_param(self.model, 'positional_embedding')
+            freeze_param(self.model, 'text_projection')
+            freeze_param(self.model, 'logit_scale')
+            for idx, resattnblock in enumerate(self.model.transformer.resblocks):
+                convert_resattnblock(resattnblock)
+            freeze_module(self.model.token_embedding)
+            self.model.ln_final.__class__ = ln_freezed_temp
+            # freeze_module(self.model.ln_final)
+
+    def tokenize(self, texts):
+        if isinstance(texts, str):
+            texts = [texts]
+        sot_token = self.tokenizer.encoder["<start_of_text>"]
+        eot_token = self.tokenizer.encoder["<end_of_text>"]
+        all_tokens = [[sot_token] + self.tokenizer.encode(text) + [eot_token] for text in texts]
+        maxn = self.num_regular_tokens
+        regular_tokens = [[[ti] if ti < maxn else [0]*self.texpand for ti in tokens] for tokens in all_tokens]
+        token_mask     = [[[ 0] if ti < maxn else [1]*self.texpand for ti in tokens] for tokens in all_tokens]
+        custom_tokens  = [[[ 0] if ti < maxn else [
+            (ti-maxn)*self.texpand+ii for ii in range(self.texpand)]
+                for ti in tokens] for tokens in all_tokens]
+
+        from itertools import chain
+        regular_tokens = [[i for i in chain(*tokens)] for tokens in regular_tokens]
+        token_mask     = [[i for i in chain(*tokens)] for tokens in token_mask]
+        custom_tokens  = [[i for i in chain(*tokens)] for tokens in custom_tokens]
+        return regular_tokens, custom_tokens, token_mask
+
+###################
+# clip expandable #
+###################
+
+@register('clip_text_sdv1_customized_embedding')
+class CLIPTextSD1CE(nn.Module):
+    def __init__(
+            self, 
+            replace_info="text|elon musk",
+            version="openai/clip-vit-large-patch14", 
+            max_length=77):
+        super().__init__()
+
+        self.name = 'clip_text_sdv1_customized_embedding'
+        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.transformer = CLIPTextModel.from_pretrained(version)
+        self.reset_replace_info(replace_info)
+        self.max_length = max_length
+        self.special_token = "<new_token>"
+
+    def reset_replace_info(self, replace_info):
+        rtype, rpara = replace_info.split("|")
+        self.replace_type = rtype
+        if rtype == "token_embedding":
+            ce_num = int(rpara)
+            ce_dim = self.transformer.text_model.embeddings.token_embedding.weight.size(1)
+            self.cembedding = nn.Embedding(ce_num, ce_dim)
+            self.cembedding = self.cembedding.to(self.get_device())
+        elif rtype == "context_embedding":
+            ce_num = int(rpara)
+            ce_dim = self.transformer.text_model.encoder.layers[-1].layer_norm2.weight.size(0)
+            self.cembedding = nn.Embedding(ce_num, ce_dim)
+            self.cembedding = self.cembedding.to(self.get_device())
+        else:
+            assert rtype=="text"
+            self.replace_type = "text"
+            self.replace_string = rpara
+            self.cembedding = None
+
+    def get_device(self):
+        return self.transformer.text_model.embeddings.token_embedding.weight.device
+
+    def position_to_mask(self, tokens, positions):
+        mask = torch.zeros_like(tokens)
+        for idxb, idxs, idxe in zip(*positions):
+            mask[idxb, idxs:idxe] = 1
+        return mask
+
+    def forward(self, text):
+        tokens, positions = self.tokenize(text)
+        mask = self.position_to_mask(tokens, positions)
+        max_token_n = tokens.size(1)
+        positional_ids = torch.arange(max_token_n)[None].to(self.get_device())
+        
+        if self.replace_what == 'token_embedding':
+            cembeds = self.cembedding(tokens * mask)
+
+            def embedding_customized_forward(
+                    self, input_ids=None, position_ids=None, inputs_embeds=None,):
+                seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]
+                if position_ids is None:
+                    position_ids = self.position_ids[:, :seq_length]
+                if inputs_embeds is None:
+                    inputs_embeds = self.token_embedding(input_ids)
+                    inputs_embeds = inputs_embeds * (1-mask.float())[:, :, None]
+                    inputs_embeds = inputs_embeds + cembeds
+                position_embeddings = self.position_embedding(position_ids)
+                embeddings = inputs_embeds + position_embeddings
+                return embeddings
+
+            import types
+            self.transformer.text_model.embeddings.forward = types.MethodType(
+                embedding_customized_forward, self.transformer.text_model.embeddings)
+            
+        else:
+            # TODO: Implement
+            assert False
+
+        outputs = self.transformer(
+            input_ids=tokens, 
+            position_ids=positional_ids, )
+        z = outputs.last_hidden_state
+        return z
+
+    def encode(self, text):
+        return self(text)
+
+    @torch.no_grad()
+    def tokenize(self, text):
+        if isinstance(text, str):
+            text = [text]
+
+        bos_special_text = "<|startoftext|>"
+        text = [ti.replace(self.special_token, bos_special_text) for ti in text]
+
+        batch_encoding = self.tokenizer(
+            text, truncation=True, max_length=self.max_length, return_length=True,
+            return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"]
+
+        bosid = tokens[0,  0]
+        eosid = tokens[0, -1]
+        bs, maxn = tokens.shape
+
+        if self.replace_what in ['token_embedding', 'context_embedding']:
+            newtokens = []
+            ce_num = self.cembedding.weight.size(0)
+            idxi = []; idxstart = []; idxend = [];
+            for idxii, tokeni in enumerate(tokens):
+                newtokeni = []
+                idxjj = 0
+                for ii, tokenii in enumerate(tokeni):
+                    if (tokenii == bosid) and (ii != 0):
+                        newtokeni.extend([i for i in range(ce_num)])
+                        idxi.append(idxii); idxstart.append(idxjj);
+                        idxjj += ce_num
+                        idxjj_record = idxjj if idxjj<=maxn-1 else maxn-1
+                        idxend.append(idxjj_record);
+                    else:
+                        newtokeni.extend([tokenii])
+                        idxjj += 1
+                newtokeni = newtokeni[:maxn]
+                newtokeni[-1] = eosid
+                newtokens.append(newtokeni)
+            return torch.LongTensor(newtokens).to(self.get_device()), (idxi, idxstart, idxend)
+        else:
+            # TODO: Implement
+            assert False

lib/model_zoo/common/get_model.py

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

lib/model_zoo/common/get_optimizer.py

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

lib/model_zoo/common/get_scheduler.py

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

lib/model_zoo/common/utils.py

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

lib/model_zoo/controlnet.py

@@ -0,0 +1,503 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import numpy.random as npr
+import copy
+from functools import partial
+from contextlib import contextmanager
+from lib.model_zoo.common.get_model import get_model, register
+from lib.log_service import print_log
+
+from .openaimodel import \
+    TimestepEmbedSequential, conv_nd, zero_module, \
+    ResBlock, AttentionBlock, SpatialTransformer, \
+    Downsample, timestep_embedding
+
+####################
+# preprocess depth #
+####################
+
+# depth_model = None
+
+# def unload_midas_model():
+#     global depth_model
+#     if depth_model is not None:
+#         depth_model = depth_model.cpu()
+
+# def apply_midas(input_image, a=np.pi*2.0, bg_th=0.1, device='cpu'):
+#     import cv2
+#     from einops import rearrange
+#     from .controlnet_annotators.midas import MiDaSInference
+#     global depth_model
+#     if depth_model is None:
+#         depth_model = MiDaSInference(model_type="dpt_hybrid")
+#         depth_model = depth_model.to(device)
+    
+#     assert input_image.ndim == 3
+#     image_depth = input_image
+#     with torch.no_grad():
+#         image_depth = torch.from_numpy(image_depth).float()
+#         image_depth = image_depth.to(device)
+#         image_depth = image_depth / 127.5 - 1.0
+#         image_depth = rearrange(image_depth, 'h w c -> 1 c h w')
+#         depth = depth_model(image_depth)[0]
+
+#         depth_pt = depth.clone()
+#         depth_pt -= torch.min(depth_pt)
+#         depth_pt /= torch.max(depth_pt)
+#         depth_pt = depth_pt.cpu().numpy()
+#         depth_image = (depth_pt * 255.0).clip(0, 255).astype(np.uint8)
+
+#         depth_np = depth.cpu().numpy()
+#         x = cv2.Sobel(depth_np, cv2.CV_32F, 1, 0, ksize=3)
+#         y = cv2.Sobel(depth_np, cv2.CV_32F, 0, 1, ksize=3)
+#         z = np.ones_like(x) * a
+#         x[depth_pt < bg_th] = 0
+#         y[depth_pt < bg_th] = 0
+#         normal = np.stack([x, y, z], axis=2)
+#         normal /= np.sum(normal ** 2.0, axis=2, keepdims=True) ** 0.5
+#         normal_image = (normal * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
+
+#         return depth_image, normal_image
+
+
+@register('controlnet')
+class ControlNet(nn.Module):
+    def __init__(
+            self,
+            image_size,
+            in_channels,
+            model_channels,
+            hint_channels,
+            num_res_blocks,
+            attention_resolutions,
+            dropout=0,
+            channel_mult=(1, 2, 4, 8),
+            conv_resample=True,
+            dims=2,
+            use_checkpoint=False,
+            use_fp16=False,
+            num_heads=-1,
+            num_head_channels=-1,
+            num_heads_upsample=-1,
+            use_scale_shift_norm=False,
+            resblock_updown=False,
+            use_new_attention_order=False,
+            use_spatial_transformer=False,  # custom transformer support
+            transformer_depth=1,  # custom transformer support
+            context_dim=None,  # custom transformer support
+            n_embed=None,  # custom support for prediction of discrete ids into codebook of first stage vq model
+            legacy=True,
+            disable_self_attentions=None,
+            num_attention_blocks=None,
+            disable_middle_self_attn=False,
+            use_linear_in_transformer=False,
+    ):
+        super().__init__()
+        if use_spatial_transformer:
+            assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
+
+        if context_dim is not None:
+            assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
+            from omegaconf.listconfig import ListConfig
+            if type(context_dim) == ListConfig:
+                context_dim = list(context_dim)
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        if num_heads == -1:
+            assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
+
+        if num_head_channels == -1:
+            assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
+
+        self.dims = dims
+        self.image_size = image_size
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        if isinstance(num_res_blocks, int):
+            self.num_res_blocks = len(channel_mult) * [num_res_blocks]
+        else:
+            if len(num_res_blocks) != len(channel_mult):
+                raise ValueError("provide num_res_blocks either as an int (globally constant) or "
+                                 "as a list/tuple (per-level) with the same length as channel_mult")
+            self.num_res_blocks = num_res_blocks
+        if disable_self_attentions is not None:
+            # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
+            assert len(disable_self_attentions) == len(channel_mult)
+        if num_attention_blocks is not None:
+            assert len(num_attention_blocks) == len(self.num_res_blocks)
+            assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
+            print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
+                  f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
+                  f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
+                  f"attention will still not be set.")
+
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.use_checkpoint = use_checkpoint
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+        self.predict_codebook_ids = n_embed is not None
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            nn.Linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            nn.Linear(time_embed_dim, time_embed_dim),
+        )
+
+        self.input_blocks = nn.ModuleList(
+            [
+                TimestepEmbedSequential(
+                    conv_nd(dims, in_channels, model_channels, 3, padding=1)
+                )
+            ]
+        )
+        self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels)])
+
+        self.input_hint_block = TimestepEmbedSequential(
+            conv_nd(dims, hint_channels, 16, 3, padding=1),
+            nn.SiLU(),
+            conv_nd(dims, 16, 16, 3, padding=1),
+            nn.SiLU(),
+            conv_nd(dims, 16, 32, 3, padding=1, stride=2),
+            nn.SiLU(),
+            conv_nd(dims, 32, 32, 3, padding=1),
+            nn.SiLU(),
+            conv_nd(dims, 32, 96, 3, padding=1, stride=2),
+            nn.SiLU(),
+            conv_nd(dims, 96, 96, 3, padding=1),
+            nn.SiLU(),
+            conv_nd(dims, 96, 256, 3, padding=1, stride=2),
+            nn.SiLU(),
+            zero_module(conv_nd(dims, 256, model_channels, 3, padding=1))
+        )
+
+        self._feature_size = model_channels
+        input_block_chans = [model_channels]
+        ch = model_channels
+        ds = 1
+        for level, mult in enumerate(channel_mult):
+            for nr in range(self.num_res_blocks[level]):
+                layers = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=mult * model_channels,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = mult * model_channels
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads = ch // num_head_channels
+                        dim_head = num_head_channels
+                    if legacy:
+                        # num_heads = 1
+                        dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+                    if disable_self_attentions is not None:
+                        disabled_sa = disable_self_attentions[level]
+                    else:
+                        disabled_sa = False
+
+                    if (num_attention_blocks is None) or nr < num_attention_blocks[level]:
+                        layers.append(
+                            AttentionBlock(
+                                ch,
+                                use_checkpoint=use_checkpoint,
+                                num_heads=num_heads,
+                                num_head_channels=dim_head,
+                                use_new_attention_order=use_new_attention_order,
+                            ) if not use_spatial_transformer else SpatialTransformer(
+                                ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
+                                disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
+                                use_checkpoint=use_checkpoint
+                            )
+                        )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self.zero_convs.append(self.make_zero_conv(ch))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    TimestepEmbedSequential(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True,
+                        )
+                        if resblock_updown
+                        else Downsample(
+                            ch, conv_resample, dims=dims, out_channels=out_ch
+                        )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+                self.zero_convs.append(self.make_zero_conv(ch))
+                ds *= 2
+                self._feature_size += ch
+
+        if num_head_channels == -1:
+            dim_head = ch // num_heads
+        else:
+            num_heads = ch // num_head_channels
+            dim_head = num_head_channels
+        if legacy:
+            # num_heads = 1
+            dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+        self.middle_block = TimestepEmbedSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            AttentionBlock(
+                ch,
+                use_checkpoint=use_checkpoint,
+                num_heads=num_heads,
+                num_head_channels=dim_head,
+                use_new_attention_order=use_new_attention_order,
+            ) if not use_spatial_transformer else SpatialTransformer(  # always uses a self-attn
+                ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
+                disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
+                use_checkpoint=use_checkpoint
+            ),
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self.middle_block_out = self.make_zero_conv(ch)
+        self._feature_size += ch
+
+    def make_zero_conv(self, channels):
+        return TimestepEmbedSequential(zero_module(conv_nd(self.dims, channels, channels, 1, padding=0)))
+
+    def forward(self, x, hint, timesteps, context, **kwargs):
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+        t_emb = t_emb.to(x.dtype)
+        emb = self.time_embed(t_emb)
+
+        guided_hint = self.input_hint_block(hint, emb, context)
+
+        outs = []
+
+        h = x
+        for module, zero_conv in zip(self.input_blocks, self.zero_convs):
+            if guided_hint is not None:
+                h = module(h, emb, context)
+                h += guided_hint
+                guided_hint = None
+            else:
+                h = module(h, emb, context)
+            outs.append(zero_conv(h, emb, context))
+
+        h = self.middle_block(h, emb, context)
+        outs.append(self.middle_block_out(h, emb, context))
+
+        return outs
+
+    def get_device(self):
+        return self.time_embed[0].weight.device
+
+    def get_dtype(self):
+        return self.time_embed[0].weight.dtype
+
+    def preprocess(self, x, type='canny', **kwargs):
+        import torchvision.transforms as tvtrans
+        if isinstance(x, str):
+            import PIL.Image
+            device, dtype = self.get_device(), self.get_dtype()
+            x_list = [PIL.Image.open(x)]
+        elif isinstance(x, torch.Tensor):
+            x_list = [tvtrans.ToPILImage()(xi) for xi in x]
+            device, dtype = x.device, x.dtype
+        else:
+            assert False
+
+        if type == 'none' or type is None:
+            return None
+
+        elif type in ['input', 'shuffle_v11e']:
+            y_torch = torch.stack([tvtrans.ToTensor()(xi) for xi in x_list])
+            y_torch = y_torch.to(device).to(torch.float32)
+            return y_torch
+
+        elif type in ['canny', 'canny_v11p']:
+            low_threshold = kwargs.pop('low_threshold', 100)
+            high_threshold = kwargs.pop('high_threshold', 200)
+            from .controlnet_annotator.canny import apply_canny
+            y_list = [apply_canny(np.array(xi), low_threshold, high_threshold) for xi in x_list]
+            y_torch = torch.stack([tvtrans.ToTensor()(yi) for yi in y_list])
+            y_torch = y_torch.repeat(1, 3, 1, 1) # Make is RGB
+            y_torch = y_torch.to(device).to(torch.float32)
+            return y_torch
+
+        elif type == 'depth':
+            from .controlnet_annotator.midas import apply_midas
+            y_list, _ = zip(*[apply_midas(input_image=np.array(xi), a=np.pi*2.0, device=device) for xi in x_list])
+            y_torch = torch.stack([tvtrans.ToTensor()(yi) for yi in y_list])
+            y_torch = y_torch.repeat(1, 3, 1, 1) # Make is RGB
+            y_torch = y_torch.to(device).to(torch.float32)
+            return y_torch
+
+        elif type in ['hed', 'softedge_v11p']:
+            from .controlnet_annotator.hed import apply_hed
+            y_list = [apply_hed(np.array(xi), device=device) for xi in x_list]
+            y_torch = torch.stack([tvtrans.ToTensor()(yi) for yi in y_list])
+            y_torch = y_torch.repeat(1, 3, 1, 1) # Make is RGB
+            y_torch = y_torch.to(device).to(torch.float32)
+            return y_torch
+
+        elif type in ['mlsd', 'mlsd_v11p']:
+            thr_v = kwargs.pop('thr_v', 0.1)
+            thr_d = kwargs.pop('thr_d', 0.1)
+            from .controlnet_annotator.mlsd import apply_mlsd
+            y_list = [apply_mlsd(np.array(xi), thr_v=thr_v, thr_d=thr_d, device=device) for xi in x_list]
+            y_torch = torch.stack([tvtrans.ToTensor()(yi) for yi in y_list])
+            y_torch = y_torch.repeat(1, 3, 1, 1) # Make is RGB
+            y_torch = y_torch.to(device).to(torch.float32)
+            return y_torch
+
+        elif type == 'normal':
+            bg_th = kwargs.pop('bg_th', 0.4)
+            from .controlnet_annotator.midas import apply_midas
+            _, y_list = zip(*[apply_midas(input_image=np.array(xi), a=np.pi*2.0, bg_th=bg_th, device=device) for xi in x_list])
+            y_torch = torch.stack([tvtrans.ToTensor()(yi.copy()) for yi in y_list])
+            y_torch = y_torch.to(device).to(torch.float32)
+            return y_torch
+
+        elif type in ['openpose', 'openpose_v11p']:
+            from .controlnet_annotator.openpose import OpenposeModel
+            from functools import partial
+            wrapper = OpenposeModel()
+            apply_openpose = partial(
+                wrapper.run_model, include_body=True, include_hand=False, include_face=False, 
+                json_pose_callback=None, device=device)
+            y_list = [apply_openpose(np.array(xi)) for xi in x_list]
+            y_torch = torch.stack([tvtrans.ToTensor()(yi.copy()) for yi in y_list])
+            y_torch = y_torch.to(device).to(torch.float32)
+            return y_torch
+
+        elif type in ['openpose_withface', 'openpose_withface_v11p']:
+            from .controlnet_annotator.openpose import OpenposeModel
+            from functools import partial
+            wrapper = OpenposeModel()
+            apply_openpose = partial(
+                wrapper.run_model, include_body=True, include_hand=False, include_face=True, 
+                json_pose_callback=None, device=device)
+            y_list = [apply_openpose(np.array(xi)) for xi in x_list]
+            y_torch = torch.stack([tvtrans.ToTensor()(yi.copy()) for yi in y_list])
+            y_torch = y_torch.to(device).to(torch.float32)
+            return y_torch
+
+        elif type in ['openpose_withfacehand', 'openpose_withfacehand_v11p']:
+            from .controlnet_annotator.openpose import OpenposeModel
+            from functools import partial
+            wrapper = OpenposeModel()
+            apply_openpose = partial(
+                wrapper.run_model, include_body=True, include_hand=True, include_face=True, 
+                json_pose_callback=None, device=device)
+            y_list = [apply_openpose(np.array(xi)) for xi in x_list]
+            y_torch = torch.stack([tvtrans.ToTensor()(yi.copy()) for yi in y_list])
+            y_torch = y_torch.to(device).to(torch.float32)
+            return y_torch
+
+        elif type == 'scribble':
+            method = kwargs.pop('method', 'pidinet')
+
+            import cv2
+            def nms(x, t, s):
+                x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
+                f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
+                f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
+                f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
+                f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
+                y = np.zeros_like(x)
+                for f in [f1, f2, f3, f4]:
+                    np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
+                z = np.zeros_like(y, dtype=np.uint8)
+                z[y > t] = 255
+                return z
+
+            def make_scribble(result):
+                result = nms(result, 127, 3.0)
+                result = cv2.GaussianBlur(result, (0, 0), 3.0)
+                result[result > 4] = 255
+                result[result < 255] = 0
+                return result
+
+            if method == 'hed':
+                from .controlnet_annotator.hed import apply_hed
+                y_list = [apply_hed(np.array(xi), device=device) for xi in x_list]
+                y_list = [make_scribble(yi) for yi in y_list]
+                y_torch = torch.stack([tvtrans.ToTensor()(yi) for yi in y_list])
+                y_torch = y_torch.repeat(1, 3, 1, 1) # Make is RGB
+                y_torch = y_torch.to(device).to(torch.float32)
+                return y_torch
+            
+            elif method == 'pidinet':
+                from .controlnet_annotator.pidinet import apply_pidinet
+                y_list = [apply_pidinet(np.array(xi), device=device) for xi in x_list]
+                y_list = [make_scribble(yi) for yi in y_list]
+                y_torch = torch.stack([tvtrans.ToTensor()(yi) for yi in y_list])
+                y_torch = y_torch.repeat(1, 3, 1, 1) # Make is RGB
+                y_torch = y_torch.to(device).to(torch.float32)
+                return y_torch
+
+            elif method == 'xdog':
+                threshold = kwargs.pop('threshold', 32)
+                def apply_scribble_xdog(img):
+                    g1 = cv2.GaussianBlur(img.astype(np.float32), (0, 0), 0.5)
+                    g2 = cv2.GaussianBlur(img.astype(np.float32), (0, 0), 5.0)
+                    dog = (255 - np.min(g2 - g1, axis=2)).clip(0, 255).astype(np.uint8)
+                    result = np.zeros_like(img, dtype=np.uint8)
+                    result[2 * (255 - dog) > threshold] = 255
+                    return result
+
+                y_list = [apply_scribble_xdog(np.array(xi), device=device) for xi in x_list]
+                y_torch = torch.stack([tvtrans.ToTensor()(yi) for yi in y_list])
+                y_torch = y_torch.repeat(1, 3, 1, 1) # Make is RGB
+                y_torch = y_torch.to(device).to(torch.float32)
+                return y_torch
+
+            else:
+                raise ValueError
+
+        elif type == 'seg':
+            method = kwargs.pop('method', 'ufade20k')
+            if method == 'ufade20k':
+                from .controlnet_annotator.uniformer import apply_uniformer
+                y_list = [apply_uniformer(np.array(xi), palette='ade20k', device=device) for xi in x_list]
+                y_torch = torch.stack([tvtrans.ToTensor()(yi) for yi in y_list])
+                y_torch = y_torch.to(device).to(torch.float32)
+                return y_torch
+
+            else:
+                raise ValueError

lib/model_zoo/controlnet_annotator/canny/__init__.py

@@ -0,0 +1,5 @@
+import cv2
+
+
+def apply_canny(img, low_threshold, high_threshold):
+    return cv2.Canny(img, low_threshold, high_threshold)

lib/model_zoo/controlnet_annotator/hed/__init__.py

@@ -0,0 +1,134 @@
+# This is an improved version and model of HED edge detection with Apache License, Version 2.0.
+# Please use this implementation in your products
+# This implementation may produce slightly different results from Saining Xie's official implementations,
+# but it generates smoother edges and is more suitable for ControlNet as well as other image-to-image translations.
+# Different from official models and other implementations, this is an RGB-input model (rather than BGR)
+# and in this way it works better for gradio's RGB protocol
+
+import os
+import cv2
+import torch
+import numpy as np
+
+from einops import rearrange
+import os
+
+models_path = 'pretrained/controlnet/preprocess'
+
+def safe_step(x, step=2):
+    y = x.astype(np.float32) * float(step + 1)
+    y = y.astype(np.int32).astype(np.float32) / float(step)
+    return y
+
+class DoubleConvBlock(torch.nn.Module):
+    def __init__(self, input_channel, output_channel, layer_number):
+        super().__init__()
+        self.convs = torch.nn.Sequential()
+        self.convs.append(torch.nn.Conv2d(in_channels=input_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1))
+        for i in range(1, layer_number):
+            self.convs.append(torch.nn.Conv2d(in_channels=output_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1))
+        self.projection = torch.nn.Conv2d(in_channels=output_channel, out_channels=1, kernel_size=(1, 1), stride=(1, 1), padding=0)
+
+    def __call__(self, x, down_sampling=False):
+        h = x
+        if down_sampling:
+            h = torch.nn.functional.max_pool2d(h, kernel_size=(2, 2), stride=(2, 2))
+        for conv in self.convs:
+            h = conv(h)
+            h = torch.nn.functional.relu(h)
+        return h, self.projection(h)
+
+
+class ControlNetHED_Apache2(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.norm = torch.nn.Parameter(torch.zeros(size=(1, 3, 1, 1)))
+        self.block1 = DoubleConvBlock(input_channel=3, output_channel=64, layer_number=2)
+        self.block2 = DoubleConvBlock(input_channel=64, output_channel=128, layer_number=2)
+        self.block3 = DoubleConvBlock(input_channel=128, output_channel=256, layer_number=3)
+        self.block4 = DoubleConvBlock(input_channel=256, output_channel=512, layer_number=3)
+        self.block5 = DoubleConvBlock(input_channel=512, output_channel=512, layer_number=3)
+
+    def __call__(self, x):
+        h = x - self.norm
+        h, projection1 = self.block1(h)
+        h, projection2 = self.block2(h, down_sampling=True)
+        h, projection3 = self.block3(h, down_sampling=True)
+        h, projection4 = self.block4(h, down_sampling=True)
+        h, projection5 = self.block5(h, down_sampling=True)
+        return projection1, projection2, projection3, projection4, projection5
+
+
+netNetwork = None
+remote_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetHED.pth"
+modeldir = os.path.join(models_path, "hed")
+old_modeldir = os.path.dirname(os.path.realpath(__file__))
+
+
+def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
+    """Load file form http url, will download models if necessary.
+
+    Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
+
+    Args:
+        url (str): URL to be downloaded.
+        model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
+            Default: None.
+        progress (bool): Whether to show the download progress. Default: True.
+        file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
+
+    Returns:
+        str: The path to the downloaded file.
+    """
+    from torch.hub import download_url_to_file, get_dir
+    from urllib.parse import urlparse
+    if model_dir is None:  # use the pytorch hub_dir
+        hub_dir = get_dir()
+        model_dir = os.path.join(hub_dir, 'checkpoints')
+
+    os.makedirs(model_dir, exist_ok=True)
+
+    parts = urlparse(url)
+    filename = os.path.basename(parts.path)
+    if file_name is not None:
+        filename = file_name
+    cached_file = os.path.abspath(os.path.join(model_dir, filename))
+    if not os.path.exists(cached_file):
+        print(f'Downloading: "{url}" to {cached_file}\n')
+        download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
+    return cached_file
+
+
+def apply_hed(input_image, is_safe=False, device='cpu'):
+    global netNetwork
+    if netNetwork is None:
+        modelpath = os.path.join(modeldir, "ControlNetHED.pth")
+        old_modelpath = os.path.join(old_modeldir, "ControlNetHED.pth")
+        if os.path.exists(old_modelpath):
+            modelpath = old_modelpath
+        elif not os.path.exists(modelpath):
+            load_file_from_url(remote_model_path, model_dir=modeldir)
+        netNetwork = ControlNetHED_Apache2().to(device)
+        netNetwork.load_state_dict(torch.load(modelpath, map_location='cpu'))
+    netNetwork.to(device).float().eval()
+
+    assert input_image.ndim == 3
+    H, W, C = input_image.shape
+    with torch.no_grad():
+        image_hed = torch.from_numpy(input_image.copy()).float().to(device)
+        image_hed = rearrange(image_hed, 'h w c -> 1 c h w')
+        edges = netNetwork(image_hed)
+        edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges]
+        edges = [cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges]
+        edges = np.stack(edges, axis=2)
+        edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
+        if is_safe:
+            edge = safe_step(edge)
+        edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
+        return edge
+
+    
+def unload_hed_model():
+    global netNetwork
+    if netNetwork is not None:
+        netNetwork.cpu()