Repository: FlagAI-Open/FlagAI (language: python)

Issue:
fix bugs
Signed-off-by: zhaohu xing <920232796@qq.com>


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index ff454454..56df29e1 100755
--- a/README.md
+++ b/README.md
@@ -21,6 +21,8 @@ FlagAI (Fast LArge-scale General AI models) is a fast, easy-to-use and extensibl
 The code is partially based on [GLM](https://github.com/THUDM/GLM), [Transformers](https://github.com/huggingface/transformers) and [DeepSpeedExamples](https://github.com/microsoft/DeepSpeedExamples/tree/master/Megatron-LM).
 
 ## News
+- [10 Nov 2022] release v1.4.0, Added AltDiffusion module in [#90](https://github.com/FlagAI-Open/FlagAI/pull/90)
+- [10 Nov 2022] release v1.4.0, Added AltCLIP module with XLMR-large tokenizer in [#90](https://github.com/FlagAI-Open/FlagAI/pull/90)
 - [29 Aug 2022] release v1.3.0, Added CLIP module and redesigned tokenizer apis in [#81](https://github.com/FlagAI-Open/FlagAI/pull/81)
 - [21 Jul 2022] release v1.2.0, ViTs are supported in [#71](https://github.com/FlagAI-Open/FlagAI/pull/71)
 - [29 Jun 2022] release v1.1.0, support OPTs downloading and inference/finetuning [#63](https://github.com/FlagAI-Open/FlagAI/pull/63)
diff --git a/README_zh.md b/README_zh.md
index 3fa1e38a..0ba89773 100755
--- a/README_zh.md
+++ b/README_zh.md
@@ -21,6 +21,8 @@
 本项目的部分代码基于 [GLM](https://github.com/THUDM/GLM),[Transformers](https://github.com/huggingface/transformers)，[timm](https://github.com/rwightman/pytorch-image-models) 和 [DeepSpeedExamples](https://github.com/microsoft/DeepSpeedExamples/tree/master/Megatron-LM).
 
 ## 动态
+- [10 Nov 2022] release v1.4.0, Added AltDiffusion module in [#90](https://github.com/FlagAI-Open/FlagAI/pull/90)
+- [10 Nov 2022] release v1.4.0, Added AltCLIP module with XLMR-large tokenizer in [#90](https://github.com/FlagAI-Open/FlagAI/pull/90)
 - [29 Aug 2022] release v1.3.0, Added CLIP module and redesigned tokenizer apis in [#81](https://github.com/FlagAI-Open/FlagAI/pull/81)
 - [21 Jul 2022] release v1.2.0, ViTs are supported in [#71](https://github.com/FlagAI-Open/FlagAI/pull/71)
 - [29 Jun 2022] release v1.1.0, support OPTs downloading and inference/finetuning [#63](https://github.com/FlagAI-Open/FlagAI/pull/63)
diff --git a/docs/TUTORIAL_3_MODEL.md b/docs/TUTORIAL_3_MODEL.md
index f6ff7d77..552f097f 100755
--- a/docs/TUTORIAL_3_MODEL.md
+++ b/docs/TUTORIAL_3_MODEL.md
@@ -43,8 +43,10 @@ model = GLMModel.from_pretrain(download_path="./state_dict", model_name="GLM-lar
 
 ## All supported models
 We can use `from_pretrain` to automatically download following models
-| ClassName | ModelName | Language | Model Type |
-|-----------------------------------|------------- ----|----------|------------|
+
+
+| ClassName                                 | ModelName | Language | Model Type  |
+|-------------------------------------------|-----------|----------|-------------|
 | flagai.model.glm_model.GLMModel | ** GLM-10b-ch** | chinese | encoder |
 | flagai.model.glm_model.GLMModel | **GLM-large-ch** | chinese | encoder |
 | flagai.model.bert_model.BertModel | **RoBERTa-base-ch** | chinese | encoder |
diff --git a/examples/AltCLIP/AltCLIP_evaluation.py b/examples/AltCLIP/AltCLIP_evaluation.py
new file mode 100644
index 00000000..d02cea93
--- /dev/null
+++ b/examples/AltCLIP/AltCLIP_evaluation.py
@@ -0,0 +1,89 @@
+# Copyright © 2022 BAAI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License")
+import torch
+from flagai.auto_model.auto_loader import AutoLoader
+import zeroshot_classification
+import json 
+import os 
+from torchvision.datasets import CIFAR10
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+maxlen = 256
+
+dataset_root = "clip_benchmark_datasets"
+dataset_name = "cifar10"
+
+auto_loader = AutoLoader(
+    task_name="txt_img_matching",
+    model_dir="./checkpoints/",
+    model_name="AltCLIP-XLMR-L"   # Load the checkpoints from Modelhub(model.baai.ac.cn/models)
+)
+
+model = auto_loader.get_model()
+model.to(device)
+model.eval()
+tokenizer = auto_loader.get_tokenizer()
+transform = auto_loader.get_transform()
+
+dataset = CIFAR10(root=os.path.join(dataset_root, dataset_name), 
+                transform=transform,   
+                download=True)
+batch_size = 128
+num_workers = 4
+
+template = {"cifar10": [
+        "a photo of a {c}.",
+        "a blurry photo of a {c}.",
+        "a black and white photo of a {c}.",
+        "a low contrast photo of a {c}.",
+        "a high contrast photo of a {c}.",
+        "a bad photo of a {c}.",
+        "a good photo of a {c}.",
+        "a photo of a small {c}.",
+        "a photo of a big {c}.",
+        "a photo of the {c}.",
+        "a blurry photo of the {c}.",
+        "a black and white photo of the {c}.",
+        "a low contrast photo of the {c}.",
+        "a high contrast photo of the {c}.",
+        "a bad photo of the {c}.",
+        "a good photo of the {c}.",
+        "a photo of the small {c}.",
+        "a photo of the big {c}."
+    ],
+}
+def evaluate():
+    if dataset:
+        dataloader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+        )
+
+        zeroshot_templates = template["cifar10"]
+        classnames = dataset.classes if hasattr(dataset, "classes") else None
+
+        metrics = zeroshot_classification.evaluate(
+            model,
+            dataloader,
+            tokenizer,
+            classnames, 
+            zeroshot_templates,
+            device=device,
+            amp=True,
+        )
+       
+        dump = {
+            "dataset": dataset_name,
+            "metrics": metrics
+        }
+
+        print(dump)
+        with open("./result.txt", "w") as f:
+            json.dump(dump, f)
+        return metrics
+
+if __name__ == "__main__":
+    evaluate()
diff --git a/examples/AltCLIP/README.md b/examples/AltCLIP/README.md
new file mode 100644
index 00000000..6ba988cd
--- /dev/null
+++ b/examples/AltCLIP/README.md
@@ -0,0 +1,404 @@
+
+
+我们提出了一个简单高效的方法去训练更加优秀的双语CLIP模型。命名为AltCLIP。
+
+训练共有两个阶段。
+在平行知识蒸馏阶段，我们只是使用平行语料文本来进行蒸馏（平行语料相对于图文对更容易获取且数量更大）。在双语对比学习阶段，我们使用少量的中-英 图像-文本对（一共约2百万）来训练我们的文本编码器以更好地适应图像编码器。
+
+AltCLIP模型可以为本项目中的AltDiffusion模型提供支持，关于AltDiffusion模型的具体信息可查看[此教程](https://github.com/FlagAI-Open/FlagAI/tree/master/examples/AltDiffusion/README.md) 。
+
+模型与权重已经在FlagAI([https://github.com/FlagAI-Open/FlagAI](https://github.com/FlagAI-Open/FlagAI))上开源，我们还提供了微调，推理，验证的脚本，欢迎试用。
+
+We propose a simple and efficient method to train a better bilingual CLIP model. It is named AltCLIP.
+
+There are two phases of training.
+In the parallel knowledge distillation phase, we only use parallel corpus texts for distillation (parallel corpus is easier to obtain and larger in number compared to image text pairs). In the bilingual comparison learning phase, we use a small number of Chinese-English image-text pairs (about 2 million in total) to train our text encoder to better fit the image encoder.
+
+
+The AltCLIP model can provide support for the AltDiffusion model in this project. For specific information about the AltDiffusion model, see [this tutorial](https://github.com/FlagAI-Open/FlagAI/tree/master/examples/AltDiffusion/README.md).
+
+
+
+The model and weights have been open sourced on FlagAI ([https://github.com/FlagAI-Open/FlagAI](https://github.com/FlagAI-Open/FlagAI)), and we also provide scripts for fine-tuning, inference, and evaluation, so feel free to try them out.
+
+# 下游效果 Performance
+
+<table>
+   <tr>
+      <td rowspan=2>Language</td>
+      <td rowspan=2>Method</td>
+      <td colspan=3>Text-to-Image Retrival</td>
+      <td colspan=3>Image-to-Text Retrival</td>
+      <td rowspan=2>MR</td>
+   </tr>
+   <tr>
+      <td>R@1</td>
+      <td>R@5</td>
+      <td>R@10</td>
+      <td>R@1</td>
+      <td>R@5</td>
+      <td>R@10</td>
+   </tr>
+   <tr>
+      <td rowspan=7>English</td>
+      <td>CLIP</td>
+      <td>65.0 </td>
+      <td>87.1 </td>
+      <td>92.2 </td>
+      <td>85.1 </td>
+      <td>97.3 </td>
+      <td>99.2 </td>
+      <td>87.6 </td>
+   </tr>
+   <tr>
+      <td>Taiyi</td>
+      <td>25.3 </td>
+      <td>48.2 </td>
+      <td>59.2 </td>
+      <td>39.3 </td>
+      <td>68.1 </td>
+      <td>79.6 </td>
+      <td>53.3 </td>
+   </tr>
+   <tr>
+      <td>Wukong</td>
+      <td>-</td>
+      <td>-</td>
+      <td>-</td>
+      <td>-</td>
+      <td>-</td>
+      <td>-</td>
+      <td>-</td>
+   </tr>
+   <tr>
+      <td>R2D2</td>
+      <td>-</td>
+      <td>-</td>
+      <td>-</td>
+      <td>-</td>
+      <td>-</td>
+      <td>-</td>
+      <td>-</td>
+   </tr>
+   <tr>
+      <td>CN-CLIP</td>
+      <td>49.5 </td>
+      <td>76.9 </td>
+      <td>83.8 </td>
+      <td>66.5 </td>
+      <td>91.2 </td>
+      <td>96.0 </td>
+      <td>77.3 </td>
+   </tr>
+   <tr>
+      <td>AlterCLIP</td>
+      <td>66.3 </td>
+      <td>87.8 </td>
+      <td>92.7 </td>
+      <td>85.9 </td>
+      <td>97.7 </td>
+      <td>99.1 </td>
+      <td>88.3 </td>
+   </tr>
+   <tr>
+      <td>AlterCLIP∗</td>
+      <td>72.5 </td>
+      <td>91.6 </td>
+      <td>95.4 </td>
+      <td>86.0 </td>
+      <td>98.0 </td>
+      <td>99.1 </td>
+      <td>90.4 </td>
+   </tr>
+   <tr>
+      <td rowspan=7>Chinese</td>
+      <td>CLIP</td>
+      <td>0.0 </td>
+      <td>2.4 </td>
+      <td>4.0 </td>
+      <td>2.3 </td>
+      <td>8.1 </td>
+      <td>12.6 </td>
+      <td>5.0 </td>
+   </tr>
+   <tr>
+      <td>Taiyi</td>
+      <td>53.7 </td>
+      <td>79.8 </td>
+      <td>86.6 </td>
+      <td>63.8 </td>
+      <td>90.5 </td>
+      <td>95.9 </td>
+      <td>78.4 </td>
+   </tr>
+   <tr>
+      <td>Wukong</td>
+      <td>51.7 </td>
+      <td>78.9 </td>
+      <td>86.3 </td>
+      <td>76.1 </td>
+      <td>94.8 </td>
+      <td>97.5 </td>
+      <td>80.9 </td>
+   </tr>
+   <tr>
+      <td>R2D2</td>
+      <td>60.9 </td>
+      <td>86.8 </td>
+      <td>92.7 </td>
+      <td>77.6 </td>
+      <td>96.7 </td>
+      <td>98.9 </td>
+      <td>85.6 </td>
+   </tr>
+   <tr>
+      <td>CN-CLIP</td>
+      <td>68.0 </td>
+      <td>89.7 </td>
+      <td>94.4 </td>
+      <td>80.2 </td>
+      <td>96.6 </td>
+      <td>98.2 </td>
+      <td>87.9 </td>
+   </tr>
+   <tr>
+      <td>AlterCLIP</td>
+      <td>63.7 </td>
+      <td>86.3 </td>
+      <td>92.1 </td>
+      <td>84.7 </td>
+      <td>97.4 </td>
+      <td>98.7 </td>
+      <td>87.2 </td>
+   </tr>
+   <tr>
+      <td>AlterCLIP∗</td>
+      <td>69.8 </td>
+      <td>89.9 </td>
+      <td>94.7 </td>
+      <td>84.8 </td>
+      <td>97.4 </td>
+      <td>98.8 </td>
+      <td>89.2 </td>
+   </tr>
+</table>
+
+
+
+![](imgs/imagenet.png)
+
+# 可视化效果 Visualization effects
+
+基于AltCLIP，我们还开发了AltDiffusion模型，可视化效果如下。
+
+Based on AltCLIP, we have also developed the AltDiffusion model, visualized as follows.
+
+![](https://raw.githubusercontent.com/920232796/test/master/image7.png)
+
+# 模型推理 Inference
+
+```python
+import torch
+from PIL import Image
+from flagai.auto_model.auto_loader import AutoLoader
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+## 一行代码直接自动下载权重到'./checkpoints/clip-xlmr-large'，并自动加载CLIP模型权重
+## modelhub地址: Modelhub(https://model.baai.ac.cn/models)
+loader = AutoLoader(
+    task_name="txt_img_matching",
+    model_dir="./checkpoints",
+    model_name="AltCLIP-XLMR-L"
+)
+## 获取加载好的模型
+model = loader.get_model()
+## 获取tokenizer
+tokenizer = loader.get_tokenizer()
+## 获取transform用来处理图像
+transform = loader.get_transform()
+
+model.eval()
+model.to(device)
+
+## 推理过程,图像与文本匹配
+image = Image.open("./dog.jpeg")
+image = transform(image)
+image = torch.tensor(image["pixel_values"]).to(device)
+text = tokenizer(["a rat", "a dog", "a cat"])["input_ids"]
+
+text = torch.tensor(text).to(device)
+
+with torch.no_grad():
+    image_features = model.get_image_features(image)
+    text_features = model.get_text_features(text)
+    text_probs = (image_features @ text_features.T).softmax(dim=-1)
+
+print(text_probs.cpu().numpy()[0].tolist())
+```
+
+# CLIP微调 Finetuning
+
+微调采用cifar10数据集，并使用FlagAI的Trainer快速开始训练过程。
+
+Fine-tuning was done using the cifar10 dataset and using FlagAI's Trainer to quickly start the training process.
+
+```python
+# Copyright © 2022 BAAI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License")
+import torch
+from flagai.auto_model.auto_loader import AutoLoader
+import os 
+from flagai.trainer import Trainer
+from torchvision.datasets import (
+    CIFAR10
+)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+dataset_root = "./clip_benchmark_datasets"
+dataset_name = "cifar10"
+
+batch_size = 4
+classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
+
+auto_loader = AutoLoader(
+    task_name="txt_img_matching",
+    model_dir="/sharefs/baai-mrnd/xingzhaohu/",
+    model_name="AltCLIP-XLMR-L"   # Load the checkpoints from Modelhub(model.baai.ac.cn/models)
+)
+
+model = auto_loader.get_model()
+model.to(device)
+model.eval()
+tokenizer = auto_loader.get_tokenizer()
+transform = auto_loader.get_transform()
+
+trainer = Trainer(env_type="pytorch",
+                pytorch_device=device,
+                experiment_name="clip_finetuning",
+                batch_size=4,
+                lr=1e-4,
+                epochs=10,
+                log_interval=10)
+
+dataset = CIFAR10(root=os.path.join(dataset_root, dataset_name), 
+                transform=transform,   
+                download=True)
+
+def cifar10_collate_fn(batch):
+    # image shape is (batch, 3, 224, 224)
+    images = torch.tensor([b[0]["pixel_values"][0] for b in batch])
+    # text_id shape is (batch, n)
+    input_ids = torch.tensor([tokenizer(f"a photo of a {b[1]}",padding=True,truncation=True,max_length=77)["input_ids"] for b in batch])    
+
+    return {
+        "pixel_values": images,
+        "input_ids": input_ids
+    }
+    
+if __name__ == "__main__":
+    trainer.train(model=model, train_dataset=dataset, collate_fn=cifar10_collate_fn)
+```
+
+
+
+# 模型验证 Evaluation
+
+我们提供了可以直接运行的验证脚本，在cifar10数据集上进行验证。
+
+期待的输出为：```{'dataset': 'cifar10', 'metrics': {'acc1': 0.95402, 'acc5': 0.99616, 'mean_per_class_recall': 0.9541200000000002}}```
+
+We provide validation scripts that can be run directly on the cifar10 dataset.
+
+```python
+# Copyright © 2022 BAAI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License")
+import torch
+from flagai.auto_model.auto_loader import AutoLoader
+from metrics import zeroshot_classification
+import json 
+import os 
+from torchvision.datasets import CIFAR10
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+maxlen = 256
+
+dataset_root = "./clip_benchmark_datasets"
+dataset_name = "cifar10"
+
+auto_loader = AutoLoader(
+    task_name="txt_img_matching",
+    model_dir="./checkpoints/",
+    model_name="AltCLIP-XLMR-L"
+)
+
+model = auto_loader.get_model()
+model.to(device)
+model.eval()
+tokenizer = auto_loader.get_tokenizer()
+transform = auto_loader.get_transform()
+
+dataset = CIFAR10(root=os.path.join(dataset_root, dataset_name), 
+                transform=transform,   
+                download=True)
+batch_size = 128
+num_workers = 4
+
+template = {"cifar10": [
+        "a photo of a {c}.",
+        "a blurry photo of a {c}.",
+        "a black and white photo of a {c}.",
+        "a low contrast photo of a {c}.",
+        "a high contrast photo of a {c}.",
+        "a bad photo of a {c}.",
+        "a good photo of a {c}.",
+        "a photo of a small {c}.",
+        "a photo of a big {c}.",
+        "a photo of the {c}.",
+        "a blurry photo of the {c}.",
+        "a black and white photo of the {c}.",
+        "a low contrast photo of the {c}.",
+        "a high contrast photo of the {c}.",
+        "a bad photo of the {c}.",
+        "a good photo of the {c}.",
+        "a photo of the small {c}.",
+        "a photo of the big {c}."
+    ],
+}
+def evaluate():
+    if dataset:
+        dataloader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+        )
+        classnames = dataset.classes if hasattr(dataset, "classes") else None
+
+        zeroshot_templates = template["cifar10"]
+        metrics = zeroshot_classification.evaluate(
+            model,
+            dataloader,
+            tokenizer,
+            classnames, 
+            zeroshot_templates,
+            device=device,
+            amp=True,
+        )
+       
+        dump = {
+            "dataset": dataset_name,
+            "metrics": metrics
+        }
+
+        print(dump)
+        with open("./result.txt", "w") as f:
+            json.dump(dump, f)
+        return metrics
+
+if __name__ == "__main__":
+    evaluate()
+
+```
+
diff --git a/examples/AltCLIP/altclip_finetuning.py b/examples/AltCLIP/altclip_finetuning.py
new file mode 100644
index 00000000..8872387d
--- /dev/null
+++ b/examples/AltCLIP/altclip_finetuning.py
@@ -0,0 +1,56 @@
+# Copyright © 2022 BAAI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License")
+import torch
+from flagai.auto_model.auto_loader import AutoLoader
+import os 
+from flagai.trainer import Trainer
+from torchvision.datasets import (
+    CIFAR10
+)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+dataset_root = "./clip_benchmark_datasets"
+dataset_name = "cifar10"
+
+batch_size = 4
+classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
+
+auto_loader = AutoLoader(
+    task_name="txt_img_matching",
+    model_dir="./checkpoints",
+    model_name="AltCLIP-XLMR-L"   # Load the checkpoints from Modelhub(model.baai.ac.cn/models)
+)
+
+model = auto_loader.get_model()
+model.to(device)
+model.eval()
+tokenizer = auto_loader.get_tokenizer()
+transform = auto_loader.get_transform()
+
+trainer = Trainer(env_type="pytorch",
+                pytorch_device=device,
+                experiment_name="clip_finetuning",
+                batch_size=4,
+                lr=1e-4,
+                epochs=10,
+                log_interval=10)
+
+dataset = CIFAR10(root=os.path.join(dataset_root, dataset_name), 
+                transform=transform,   
+                download=True)
+
+def cifar10_collate_fn(batch):
+    # image shape is (batch, 3, 224, 224)
+    images = torch.tensor([b[0]["pixel_values"][0] for b in batch])
+    # text_id shape is (batch, n)
+    input_ids = torch.tensor([tokenizer(f"a photo of a {b[1]}",padding=True,truncation=True,max_length=77)["input_ids"] for b in batch])    
+
+    return {
+        "pixel_values": images,
+        "input_ids": input_ids
+    }
+    
+if __name__ == "__main__":
+    trainer.train(model=model, train_dataset=dataset, collate_fn=cifar10_collate_fn)
\ No newline at end of file
diff --git a/examples/AltCLIP/altclip_inference.py b/examples/AltCLIP/altclip_inference.py
new file mode 100644
index 00000000..42320117
--- /dev/null
+++ b/examples/AltCLIP/altclip_inference.py
@@ -0,0 +1,37 @@
+import torch
+from PIL import Image
+from flagai.auto_model.auto_loader import AutoLoader
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+loader = AutoLoader(
+    task_name="txt_img_matching",
+    model_name="AltCLIP-XLMR-L",   # Load the checkpoints from Modelhub(model.baai.ac.cn/models)
+    model_dir="./checkpoints"
+)
+
+model = loader.get_model()
+tokenizer = loader.get_tokenizer()
+transform = loader.get_transform()
+
+model.eval()
+model.to(device)
+tokenizer = loader.get_tokenizer()
+
+def inference():
+    image = Image.open("./dog.jpeg")
+    image = transform(image)
+    image = torch.tensor(image["pixel_values"]).to(device)
+    text = tokenizer(["a rat", "a dog", "a cat"])["input_ids"]
+
+    text = torch.tensor(text).to(device)
+
+    with torch.no_grad():
+        image_features = model.get_image_features(image)
+        text_features = model.get_text_features(text)
+        text_probs = (image_features @ text_features.T).softmax(dim=-1)
+
+    print(text_probs.cpu().numpy()[0].tolist())
+
+if __name__=="__main__":
+    inference()
diff --git a/examples/AltCLIP/dog.jpeg b/examples/AltCLIP/dog.jpeg
new file mode 100644
index 00000000..3e79ea58
Binary files /dev/null and b/examples/AltCLIP/dog.jpeg differ
diff --git a/examples/AltCLIP/imgs/imagenet.png b/examples/AltCLIP/imgs/imagenet.png
new file mode 100644
index 00000000..8bd87988
Binary files /dev/null and b/examples/AltCLIP/imgs/imagenet.png differ
diff --git a/examples/AltCLIP/zeroshot_classification.py b/examples/AltCLIP/zeroshot_classification.py
new file mode 100644
index 00000000..036167e2
--- /dev/null
+++ b/examples/AltCLIP/zeroshot_classification.py
@@ -0,0 +1,228 @@
+"""
+Code adapated from https://github.com/mlfoundations/open_clip/blob/main/src/training/zero_shot.py
+Thanks to the authors of OpenCLIP
+"""
+import logging
+from contextlib import suppress
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+from sklearn.metrics import classification_report, balanced_accuracy_score
+
+def zero_shot_classifier(model, tokenizer, classnames, templates, device, amp=True):
+    """
+    This function returns zero-shot vectors for each class in order
+    to use it for zero-shot classification.
+    
+
+    model:
+        CLIP-like model with `encode_text`
+    
+    tokenizer:
+        text tokenizer, i.e. convert list of strings to torch.Tensor of integers
+    
+    classnames: list of str
+        name of classes
+    
+    templates: list of str
+        templates to use.
+    
+    Returns
+    -------
+    
+    torch.Tensor of shape (N,C) where N is the number
+    of templates, and C is the number of classes.
+    """
+    autocast = torch.cuda.amp.autocast if amp else suppress
+    with torch.no_grad(), autocast():
+        zeroshot_weights = []
+        for classname in tqdm(classnames):
+            texts = [template.format(c=classname) for template in templates]  # format with class
+            raw_texts = texts
+            if isinstance(tokenizer,Callable):
+                texts = tokenizer(texts,padding=True,truncation=True,max_length=77,return_tensors='pt').to(device)  # tokenize
+            
+            class_embeddings = model.get_text_features(**texts)
+            
+            class_embedding = F.normalize(class_embeddings, dim=-1).mean(dim=0)
+            class_embedding /= class_embedding.norm()
+            zeroshot_weights.append(class_embedding)
+        zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(device)
+    return zeroshot_weights
+
+
+def accuracy(output, target, topk=(1,)):
+    """
+    Compute top-k accuracy
+
+    output: torch.Tensor
+        shape (N, C) where N is the number of examples, C the number of classes.
+        these are the logits.
+    
+    target: torch.Tensor
+        shape (N,) where N is the number of examples. Groundtruth class id of each example.
+    
+    topk: tuple
+        which topk to compute, e.g., topk=(1,5) will compute top-1 and top-5 accuracies
+    
+    Returns
+    -------
+    
+    list of top-k accuracies in the same order as `topk`
+    """
+    # pred = output.topk(max(topk), 1, True, True)[1].t()
+    pred = output.float().topk(max(topk), 1, True, True)[1].t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+    n = len(target)
+    return [float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy()) / n for k in topk]
+
+
+def run_classification(model, classifier, dataloader, device, amp=True):
+    """
+    Run zero-shot classifcation
+
+    model: torch.nn.Module
+        CLIP-like model with `encode_image` and `encode_text`
+    
+    classifier: torch.Tensor
+        obtained from the function `zero_shot_classifier`
+    
+    dataloader: torch.utils.data.Dataloader 
+    
+    Returns
+    -------
+    (pred, true)  where
+        - pred (N, C) are the logits
+        - true (N,) are the actual classes
+    """
+    autocast = torch.cuda.amp.autocast if amp else suppress
+    pred = []
+    true = []
+    nb = 0
+    with torch.no_grad():
+        for images, target in tqdm(dataloader):
+            # print(images)
+            # images = images.to(device)
+            if isinstance(images,torch.Tensor):
+                images = images.to(device)
+            else:
+                images['pixel_values'] = torch.cat(images['pixel_values']).to(device)
+            target = target.to(device)
+
+            with autocast():
+                
+                image_features = model.get_image_features(**images)
+
+                image_features = F.normalize(image_features, dim=-1)
+                logits = 100. * image_features @ classifier
+            
+            true.append(target.cpu())
+            pred.append(logits.cpu())
+
+    pred = torch.cat(pred)
+    true = torch.cat(true)
+    return pred, true
+
+def average_precision_per_class(scores, targets):
+    """
+    Compute average precision  for each class
+    this metric is used for multi-label classification
+    see explanations here https://fangdahan.medium.com/calculate-mean-average-precision-map-for-multi-label-classification-b082679d31be
+    Code is adapted from https://github.com/pytorch/tnt/blob/master/torchnet/meter/meter.py, thanks to the authors of `tnt`.
+
+    Parameters
+    ----------
+
+    scores: torch.Tensor
+        logits, of shape (N,C) where N is the number of examples, C the number of classes
+    
+    targets: torch.Tensor
+        one-hot vectors of groundtruth targets (N, C), where N is the number of examples, C is the
+        number of classes
+    
+    Returns
+    -------
+
+    torch.Tensor of shape (C,) of avereage precision for each class, where C is     
+    the number of classes.
+    
+    """
+    ap = torch.zeros(scores.size(1))
+    rg = torch.arange(1, scores.size(0) + 1).float()
+    # compute average precision for each class
+    for k in range(scores.size(1)):
+        # sort scores
+        scores_k = scores[:, k]
+        targets_k = targets[:, k]
+        _, sortind = torch.sort(scores_k, 0, True)
+        truth = targets_k[sortind]
+        tp = truth.float().cumsum(0)
+        # compute precision curve
+        precision = tp.div(rg)
+        # compute average precision
+        ap[k] = precision[truth.bool()].sum() / max(float(truth.sum()), 1)
+    return ap
+
+
+def evaluate(model, dataloader, tokenizer, classnames, templates, device, amp=True, verbose=False):
+    """
+    Run zero-shot classification and evaluate the metrics
+
+    Parameters
+    ----------
+
+    model: torch.nn.Module
+        CLIP-like model with `encode_image` and `encode_text`
+    
+    dataloader: torch.utils.data.Dataloader
+
+    tokenizer: text tokenizer
+
+    classnames: list of str
+        class names
+    
+    templates: list of str
+        templates to use for zero-shot classification
+    
+    device: cpu/cuda
+
+    amp: whether to use automatic mixed precision
+
+    verbose: whether to use verbose model
+
+    Returns
+    -------
+
+    dict of classification metrics
+    """
+    classifier = zero_shot_classifier(model, tokenizer, classnames, templates, device)
+    logits, target = run_classification(model, classifier, dataloader, device, amp=amp)
+    is_multilabel = (len(target.shape) == 2)
+
+    if is_multilabel:
+        if verbose:
+            print("Detected a multi-label classification dataset")
+        # Multiple labels per image, multiple classes on the dataset
+        ap_per_class = average_precision_per_class(logits, target)
+        if verbose:
+            for class_name, ap in zip(dataloader.dataset.classes, ap_per_class.tolist()):
+                print(f"Class: {class_name}, AveragePrecision: {ap}")
+        return {"mean_average_precision": ap_per_class.mean().item()}
+    else:
+        # Single label per image, multiple classes on the dataset
+        # just compute accuracy and mean_per_class_recall
+
+        pred = logits.argmax(axis=1)
+        # measure accuracy
+        if len(dataloader.dataset.classes) >= 5:
+            acc1, acc5 = accuracy(logits, target, topk=(1, 5))
+        else:
+            acc1, = accuracy(logits, target, topk=(1,))
+            acc5 = float("nan") 
+        mean_per_class_recall = balanced_accuracy_score(target, pred)
+        if verbose:
+            print(classification_report(target, pred, digits=3))
+        return {"acc1": acc1, "acc5": acc5, "mean_per_class_recall": mean_per_class_recall}
diff --git a/examples/AltDiffusion/README.md b/examples/AltDiffusion/README.md
new file mode 100644
index 00000000..35c92438
--- /dev/null
+++ b/examples/AltDiffusion/README.md
@@ -0,0 +1,131 @@
+
+#  模型信息 Model Information
+
+我们使用预训练好的双语语言模型作为我们的text encoder，并使用美学评分在5.5分以上的WuDao数据集(6M)和美学评分在以上的Laion数据(5M)进行微调。
+
+微调时，我们使用stable-diffusion-v1-4作初始化，并冻住双语语言模型，只微调Unet模型中Transformer Block的key模块与vuale模块。
+
+并且在训练时，我们将数据集按图片的长宽比进行分桶，同一个batch内的数据都裁剪到与图片尺寸相近的固定大小，从而克服了原版stable diffusion长图与宽图生成多头的问题。
+
+我们的版本在中英文对齐方面表现非常出色，是目前市面上开源的最强版本，保留了原版stable diffusion的大部分能力，并且在某些例子上比有着比原版模型更出色的能力。
+
+AltDiffusion 模型由名为 AltCLIP 的双语 CLIP 模型支持，该模型也可在本项目中访问。您可以阅读 [此教程](https://github.com/FlagAI-Open/FlagAI/tree/master/examples/AltCLIP/README.md) 了解更多信息。
+
+注意：模型推理要求一张至少10G以上的GPU。
+
+We use the pre-trained bilingual language model as our text encoder and fine-tune it using the WuDao dataset (6M) with an aesthetic score above 5.5 and the Laion data (5M) with an aesthetic score above 5.5.
+
+When fine-tuning, we use stable-diffusion v1-4 as initialization, freeze the bilingual language model, and only fine-tune the key module and vuale module of the Transformer Block in the Unet model.
+
+And during training, we divide the data set into bucketed according to the aspect ratio of the image, and the data in the same batch are cropped to a fixed size similar to the image size, so as to overcome the problem of generating multiple heads for the long and wide images of the original stable diffusion.
+
+AltDiffusion model is backed by a bilingual CLIP model named AltCLIP, which is also accessible in FlagAI. You can read [this tutorial](https://github.com/FlagAI-Open/FlagAI/tree/master/examples/AltCLIP/README.md) for more information. 
+
+Note that the model inference requires a GPU of at least 10G above.
+
+
+
+# 示例 Example
+
+以下示例将为文本输入`Anime portrait of natalie portman as an anime girl by stanley artgerm lau, wlop, rossdraws, james jean, andrei riabovitchev, marc simonetti, and sakimichan, trending on artstation` 在目录`./AltDiffusionOutputs`下生成图片结果。
+
+The following example will generate image results for text input `Anime portrait of natalie portman as an anime girl by stanley artgerm lau, wlop, rossdraws, james jean, andrei riabovitchev, marc simonetti, and sakimichan, trending on artstation` under the default output directory `./AltDiffusionOutputs`
+
+```python
+import torch
+from flagai.auto_model.auto_loader import AutoLoader
+from flagai.model.predictor.predictor import Predictor
+
+# Initialize 
+prompt = "Anime portrait of natalie portman as an anime girl by stanley artgerm lau, wlop, rossdraws, james jean, andrei riabovitchev, marc simonetti, and sakimichan, trending on artstation"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+loader = AutoLoader(task_name="text2img", #contrastive learning
+                    model_name="AltDiffusion",
+                    model_dir="./checkpoints")
+
+model = loader.get_model()
+model.eval()
+model.to(device)
+predictor = Predictor(model)
+predictor.predict_generate_images(prompt)
+```
+
+您可以在`predict_generate_images`函数里通过改变参数来调整设置，具体信息如下:
+
+More parameters of predict_generate_images for you to adjust for `predict_generate_images` are listed below:
+
+
+`prompt: str`: 提示文本; The prompt text
+
+`out_path: str`: 输出路径; The output path to save images
+
+`n_samples: int`: 输出图片数量; Number of images to be generated
+
+`skip_grid: bool`: 如果为True, 会将所有图片拼接在一起，输出一张新的图片; If set to true, image gridding step will be skipped
+
+`ddim_step: int`: DDIM模型的步数; Number of steps in ddim model
+
+`plms: bool`: 如果为True, 则会使用plms模型; If set to true, PLMS Sampler instead of DDIM Sampler will be applied
+
+`scale: float` : 这个值决定了文本在多大程度上影响生成的图片，值越大影响力越强; This value determines how important the prompt incluences generate images
+
+`H: int`: 图片的高度; Height of image
+
+`W: int`: 图片的宽度; Width of image
+
+`C: int`: 图片的channel数; Numeber of channels of generated images
+
+`seed: int`: 随机种子; Random seed number 
+
+# 模型权重 Model Weights
+
+第一次运行AltDiffusion模型时会自动下载下列权重:
+
+The following weights are automatically downloaded when the AltDiffusion model is run for the first time: 
+
+| 模型名称 Model name              | 大小 Size | 描述 Description                                        |
+|------------------------------|---------|-------------------------------------------------------|
+| StableDiffusionSafetyChecker | 1.13G   | 图片的安全检查器；Safety checker for image                     |
+| AltDiffusion                 | 8.0G    | 我们的双语AltDiffusion模型； Our bilingual AltDiffusion model |
+| AltCLIP                      | 3.22G   | 我们的双语AltCLIP模型；Our bilingual AltCLIP model            |
+
+# 更多生成结果 More Results
+
+## 中英文对齐能力
+
+### prompt:dark elf princess, highly detailed, d & d, fantasy, highly detailed, digital painting, trending on artstation, concept art, sharp focus, illustration, art by artgerm and greg rutkowski and fuji choko and viktoria gavrilenko and hoang lap
+### 英文生成结果
+
+![image](./imgs/en_暗黑精灵.png)
+
+### prompt:黑暗精灵公主，非常详细，幻想，非常详细，数字绘画，概念艺术，敏锐的焦点，插图
+### 中文生成结果
+![image](./imgs/cn_暗黑精灵.png)
+
+## 中文表现能力
+
+## prompt:带墨镜的男孩肖像，充满细节，8K高清
+![image](./imgs/小男孩.png)
+
+
+## prompt:带墨镜的中国男孩肖像，充满细节，8K高清
+![image](./imgs/cn_小男孩.png)
+
+## 长图生成能力
+
+### prompt: 一只带着帽子的小狗 
+### 原版 stable diffusion：
+![image](./imgs/多尺度狗（不好）.png)
+
+### Ours:
+![image](./imgs/多尺度狗（好）.png)
+
+注: 此处长图生成技术由右脑科技(RightBrain AI)提供。
+
+Note: Note: The long image generation technology here is provided by Right Brain Technology.
+
+# License
+
+The model is licensed with a [CreativeML Open RAIL-M license](https://huggingface.co/spaces/CompVis/stable-diffusion-license). The authors claim no rights on the outputs you generate, you are free to use them and are accountable for their use which must not go against the provisions set in this license. The license forbids you from sharing any content that violates any laws, produce any harm to a person, disseminate any personal information that would be meant for harm, spread misinformation and target vulnerable groups. You can modify and use the model for commercial purposes, but a copy of the same use restrictions must be included. For the full list of restrictions please [read the license](https://huggingface.co/spaces/CompVis/stable-diffusion-license)
\ No newline at end of file
diff --git a/examples/AltDiffusion/generate.py b/examples/AltDiffusion/generate.py
new file mode 100644
index 00000000..bf66b8eb
--- /dev/null
+++ b/examples/AltDiffusion/generate.py
@@ -0,0 +1,20 @@
+# Copyright © 2022 BAAI. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License")
+import torch
+from flagai.auto_model.auto_loader import AutoLoader
+from flagai.model.predictor.predictor import Predictor
+
+# Initialize 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+loader = AutoLoader(task_name="text2img", #contrastive learning
+                    model_name="AltDiffusion",
+                    model_dir="./checkpoints")
+
+model = loader.get_model()
+model.eval()
+model.to(device)
+predictor = Predictor(model)
+predictor.predict_generate_images("Anime portrait of natalie portman as an anime girl by stanley artgerm lau, wlop, rossdraws, james jean, andrei riabovitchev, marc simonetti, and sakimichan, trending on artstation")
\ No newline at end of file
diff --git "a/examples/AltDiffusion/imgs/cn_\345\260\217\347\224\267\345\255\251.png" "b/examples/AltDiffusion/imgs/cn_\345\260\217\347\224\267\345\255\251.png"
new file mode 100644
index 00000000..7c8dbf50
Binary files /dev/null and "b/examples/AltDiffusion/imgs/cn_\345\260\217\347\224\267\345\255\251.png" differ
diff --git "a/examples/AltDiffusion/imgs/cn_\346\232\227\351\273\221\347\262\276\347\201\265.png" "b/examples/AltDiffusion/imgs/cn_\346\232\227\351\273\221\347\262\276\347\201\265.png"
new file mode 100644
index 00000000..497a2ada
Binary files /dev/null and "b/examples/AltDiffusion/imgs/cn_\346\232\227\351\273\221\347\262\276\347\201\265.png" differ
diff --git "a/examples/AltDiffusion/imgs/en_\346\232\227\351\273\221\347\262\276\347\201\265.png" "b/examples/AltDiffusion/imgs/en_\346\232\227\351\273\221\347\262\276\347\201\265.png"
new file mode 100644
index 00000000..8efff57e
Binary files /dev/null and "b/examples/AltDiffusion/imgs/en_\346\232\227\351\273\221\347\262\276\347\201\265.png" differ
diff --git "a/examples/AltDiffusion/imgs/\345\244\232\345\260\272\345\272\246\347\213\227\357\274\210\344\270\215\345\245\275\357\274\211.png" "b/examples/AltDiffusion/imgs/\345\244\232\345\260\272\345\272\246\347\213\227\357\274\210\344\270\215\345\245\275\357\274\211.png"
new file mode 100644
index 00000000..98c2bbf1
Binary files /dev/null and "b/examples/AltDiffusion/imgs/\345\244\232\345\260\272\345\272\246\347\213\227\357\274\210\344\270\215\345\245\275\357\274\211.png" differ
diff --git "a/examples/AltDiffusion/imgs/\345\244\232\345\260\272\345\272\246\347\213\227\357\274\210\345\245\275\357\274\211.png" "b/examples/AltDiffusion/imgs/\345\244\232\345\260\272\345\272\246\347\213\227\357\274\210\345\245\275\357\274\211.png"
new file mode 100644
index 00000000..99dc4157
Binary files /dev/null and "b/examples/AltDiffusion/imgs/\345\244\232\345\260\272\345\272\246\347\213\227\357\274\210\345\245\275\357\274\211.png" differ
diff --git "a/examples/AltDiffusion/imgs/\345\260\217\347\224\267\345\255\251.png" "b/examples/AltDiffusion/imgs/\345\260\217\347\224\267\345\255\251.png"
new file mode 100644
index 00000000..9896eb79
Binary files /dev/null and "b/examples/AltDiffusion/imgs/\345\260\217\347\224\267\345\255\251.png" differ
diff --git a/examples/bert_title_generation_english/train.py b/examples/bert_title_generation_english/train.py
index bd68fbf6..a7f3423e 100644
--- a/examples/bert_title_generation_english/train.py
+++ b/examples/bert_title_generation_english/train.py
@@ -2,7 +2,6 @@
 #
 # Licensed under the Apache License, Version 2.0 (the "License")
 import sys
-sys.path.append('/mnt/liuguang/FlagAI')
 import os
 import torch
 from torch.utils.data import Dataset
@@ -24,7 +23,7 @@
     weight_decay=1e-3,
     epochs=100,
     log_interval=1,
-    eval_interval=10000,
+    eval_interval=10,
     load_dir=None,
     pytorch_device=device,
     save_dir="checkpoints-bert-title-generation-en",
@@ -33,7 +32,6 @@
     fp16 = False)
 
 model_dir = "../state_dict/"  # download_path for the model 
-
 os.makedirs(model_dir, exist_ok=True)
 maxlen = 256
 
diff --git a/examples/glm_title_generation/generate.py b/examples/glm_title_generation/generate.py
old mode 100644
new mode 100755
index 6e3e354f..51c56b98
--- a/examples/glm_title_generation/generate.py
+++ b/examples/glm_title_generation/generate.py
@@ -4,10 +4,9 @@
 import torch
 from flagai.auto_model.auto_loader import AutoLoader
 from flagai.model.predictor.predictor import Predictor
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Note "./checkpoints_glm_title_generation/{}/mp_rank_00_model_states.pt", {} is a directory in the checkpoints_glm_title_generation.
+
 model_save_path = "./checkpoints_glm_title_generation/8/mp_rank_00_model_states.pt"
 
 auto_loader = AutoLoader("seq2seq",
@@ -18,7 +17,7 @@
 
 model.load_state_dict(
     torch.load(model_save_path, map_location=device)["module"])
-
+    
 model.to(device)
 model.eval()
 predictor = Predictor(model, tokenizer)
diff --git a/flagai/auto_model/auto_loader.py b/flagai/auto_model/auto_loader.py
old mode 100644
new mode 100755
index e5c9f4cc..6344dfdb
--- a/flagai/auto_model/auto_loader.py
+++ b/flagai/auto_model/auto_loader.py
@@ -4,6 +4,8 @@
 import importlib
 import os
 import copy
+from  flagai.model.file_utils import _get_model_id
+
 
 class LazyImport(object):
 
@@ -18,6 +20,7 @@ def __getattr__(self, name):
             self.cache[self.mod_name] = mod
         return getattr(mod, name)
 
+
 # 2 columns : 1-package name,  2-class name
 ALL_TASK = {
     "bert_lm": ["flagai.model.bert_model", "BertModel"],
@@ -25,18 +28,18 @@ def __getattr__(self, name):
     "bert_title-generation": ["flagai.model.bert_model", "BertForSeq2seq"],
     "bert_masklm": ["flagai.model.bert_model", "BertForMaskLM"],
     "bert_sequence-labeling":
-        ["flagai.model.bert_model", "BertForSequenceLabeling"],
+    ["flagai.model.bert_model", "BertForSequenceLabeling"],
     "bert_sequence-labeling-crf":
-        ["flagai.model.bert_model", "BertForSequenceLabeling"],
+    ["flagai.model.bert_model", "BertForSequenceLabeling"],
     "bert_sequence-labeling-gp":
-        ["flagai.model.bert_model", "BertForSequenceLabeling"],
+    ["flagai.model.bert_model", "BertForSequenceLabeling"],
     "bert_ner": ["flagai.model.bert_model", "BertForSequenceLabeling"],
     "bert_ner-crf": ["flagai.model.bert_model", "BertForSequenceLabelingCRF"],
     "bert_ner-gp": ["flagai.model.bert_model", "BertForSequenceLabelingGP"],
     "bert_embedding": ["flagai.model.bert_model", "BertForEmbedding"],
     "bert_classification": ["flagai.model.bert_model", "BertForClsClassifier"],
     "bert_semantic-matching":
-        ["flagai.model.bert_model", "BertForClsClassifier"],
+    ["flagai.model.bert_model", "BertForClsClassifier"],
     "gpt2_seq2seq": ("flagai.model.gpt2_model", "GPT2Model"),
     "gpt2_lm": ("flagai.model.gpt2_model", "GPT2Model"),
     "cpm_seq2seq": ("flagai.model.gpt2_model", "GPT2Model"),
@@ -47,16 +50,21 @@ def __getattr__(self, name):
     "glm_seq2seq": ["flagai.model.glm_model", "GLMForSeq2Seq"],
     "glm_poetry": ["flagai.model.glm_model", "GLMForSeq2Seq"],
     "glm_classification":
-        ["flagai.model.glm_model", "GLMForSequenceClassification"],
+    ["flagai.model.glm_model", "GLMForSequenceClassification"],
     "glm_title-generation": ["flagai.model.glm_model", "GLMForSeq2Seq"],
-    "opt_seq2seq": ("flagai.model.opt_model","OPTModel"),
-    "opt_lm": ("flagai.model.opt_model","OPTModel"),
+    "opt_seq2seq": ("flagai.model.opt_model", "OPTModel"),
+    "opt_lm": ("flagai.model.opt_model", "OPTModel"),
     "vit_classification": ("flagai.model.vision.vit", "VisionTransformer"),
     "clip_txt_img_matching": ("flagai.model.mm.clip_model", "CLIP"),
     "swinv1_classification": ("flagai.model.vision.swinv1", "SwinTransformer"),
-    "swinv2_classification": ("flagai.model.vision.swinv2", "SwinTransformerV2"),
+    "swinv2_classification": ("flagai.model.vision.swinv2",
+                              "SwinTransformerV2"),
     "cpm3_lm": ("flagai.model.cpm3_model", "CPM3"),
     "cpm3_trian": ("flagai.model.cpm3_trian_model", "CPM3"),
+    "diffusion_text2img": ("flagai.model.mm.diffusion", "LatentDiffusion"),
+    # "clipcn_txt_img_matching": ("flagai.model.mm.lm.cn_clip", "CN_CLIP"),
+    "clipcn_txt_img_matching": ("flagai.model.mm.flag_clip", "ChineseCLIP"),
+
 }
 
 # 4 columns : 1-package name,  2-class name, 3-model brief name, 4-model type
@@ -69,41 +77,54 @@ def __getattr__(self, name):
     "glm-large-en": ["flagai.model.glm_model", "GLMModel", "glm", "nlp"],
     "gpt2-base-ch": ["flagai.model.gpt2_model", "GPT2Model", "gpt2", "nlp"],
     "cpm-large-ch": ["flagai.model.gpt2_model", "GPT2Model", "cpm", "nlp"],
-    "opt-125m-en": ["flagai.model.opt_model","OPTModel", "opt", "nlp"],
-    "opt-350m-en": ["flagai.model.opt_model","OPTModel", "opt", "nlp"],
-    "opt-1.3b-en": ["flagai.model.opt_model","OPTModel", "opt", "nlp"],
-    "opt-2.7b-en": ["flagai.model.opt_model","OPTModel", "opt", "nlp"],
-    "opt-6.7b-en": ["flagai.model.opt_model","OPTModel", "opt", "nlp"],
-    "opt-13b-en": ["flagai.model.opt_model","OPTModel", "opt", "nlp"],
-    "opt-30b-en": ["flagai.model.opt_model","OPTModel", "opt", "nlp"],
-    "opt-66b-en": ["flagai.model.opt_model","OPTModel", "opt", "nlp"],
+    "opt-125m-en": ["flagai.model.opt_model", "OPTModel", "opt", "nlp"],
+    "opt-350m-en": ["flagai.model.opt_model", "OPTModel", "opt", "nlp"],
+    "opt-1.3b-en": ["flagai.model.opt_model", "OPTModel", "opt", "nlp"],
+    "opt-2.7b-en": ["flagai.model.opt_model", "OPTModel", "opt", "nlp"],
+    "opt-6.7b-en": ["flagai.model.opt_model", "OPTModel", "opt", "nlp"],
+    "opt-13b-en": ["flagai.model.opt_model", "OPTModel", "opt", "nlp"],
+    "opt-30b-en": ["flagai.model.opt_model", "OPTModel", "opt", "nlp"],
+    "opt-66b-en": ["flagai.model.opt_model", "OPTModel", "opt", "nlp"],
     "glm-10b-ch": ["flagai.model.glm_model", "GLMModel", "glm", "nlp"],
     "cpm3": ["flagai.model.cpm3_model", "CPM3", "cpm3", "nlp"],
     "cpm3-train": ["flagai.model.cpm3_train_model", "CPM3", "cpm3", "nlp"],
-
-    "vit-base-p16-224":["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
-    "vit-base-p16-384":["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
-    "vit-base-p32-224":["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
-    "vit-base-p32-384":["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
-    "vit-large-p16-224":["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
-    "vit-large-p16-384":["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
-    "vit-large-p32-224":["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
-    "vit-large-p32-384":["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
-
-    "clip-base-p32-224":["flagai.model.mm.clip_model", "CLIP", "clip", "mm"],
-    "clip-base-p16-224":["flagai.model.mm.clip_model", "CLIP", "clip", "mm"],
-    "clip-large-p14-224":["flagai.model.mm.clip_model", "CLIP", "clip", "mm"],
-    "clip-large-p14-336":["flagai.model.mm.clip_model", "CLIP", "clip", "mm"],
-
-    "swinv1-base-patch4-window7-224": ["flagai.model.vision.swinv1", "SwinTransformer","swinv1","vision"],
-
-    "swinv2-base-patch4-window8-256": ["flagai.model.vision.swinv2", "SwinTransformerV2", "swinv2", "vision"],
-    "swinv2-base-patch4-window16-256": ["flagai.model.vision.swinv2", "SwinTransformerV2", "swinv2", "vision"],
-    "swinv2-small-patch4-window16-256": ["flagai.model.vision.swinv2", "SwinTransformerV2", "swinv2", "vision"],
+    "vit-base-p16-224":
+    ["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
+    "vit-base-p16-384":
+    ["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
+    "vit-base-p32-224":
+    ["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
+    "vit-base-p32-384":
+    ["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
+    "vit-large-p16-224":
+    ["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
+    "vit-large-p16-384":
+    ["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
+    "vit-large-p32-224":
+    ["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
+    "vit-large-p32-384":
+    ["flagai.model.vision.vit", "VisionTransformer", "vit", "vision"],
+    "clip-base-p32-224": ["flagai.model.mm.clip_model", "CLIP", "clip", "mm"],
+    "clip-base-p16-224": ["flagai.model.mm.clip_model", "CLIP", "clip", "mm"],
+    "clip-large-p14-224": ["flagai.model.mm.clip_model", "CLIP", "clip", "mm"],
+    "clip-large-p14-336": ["flagai.model.mm.clip_model", "CLIP", "clip", "mm"],
+    "clip-large-p14-336": ["flagai.model.mm.clip_model", "CLIP", "clip", "mm"],
+    "altdiffusion":
+    ["flagai.model.mm.diffusion", "LatentDiffusion", "diffusion", "mm"],
+    "swinv1-base-patch4-window7-224":
+    ["flagai.model.vision.swinv1", "SwinTransformer", "swinv1", "vision"],
+    "swinv2-base-patch4-window8-256":
+    ["flagai.model.vision.swinv2", "SwinTransformerV2", "swinv2", "vision"],
+    "swinv2-base-patch4-window16-256":
+    ["flagai.model.vision.swinv2", "SwinTransformerV2", "swinv2", "vision"],
+    "swinv2-small-patch4-window16-256": [
+        "flagai.model.vision.swinv2", "SwinTransformerV2", "swinv2", "vision"
+    ],
+    "clip-cn-b-16": ["flagai.model.mm.lm.cn_clip", "CN_CLIP", "clipcn", "mm"],
+    "altclip-xlmr-l": ["flagai.models.mm.flag_clip", "ChineseCLIP", "clipcn", "mm", "flagai.model.mm.flag_clip", "CHCLIPProcess"],
+    "altclip-bert-base": ["flagai.models.mm.flag_clip", "ChineseCLIP", "clipcn", "mm", "flagai.model.mm.flag_clip", "CHCLIPProcessBert"],
 }
 
-
-
 class AutoLoader:
 
     def __init__(self,
@@ -164,22 +185,42 @@ def __init__(self,
             )
             return
 
+        download_path = os.path.join(model_dir, raw_model_name)
+        print("*" * 20, task_name, model_name)
+
+        model_id = _get_model_id(f"{raw_model_name}-{task_name}")
+        if model_id != 'null':
+            model_name_ = f"{raw_model_name}-{task_name}"
+        else:
+            model_name_ = raw_model_name
+
         self.model = getattr(LazyImport(self.model_name[0]),
                              self.model_name[1]).from_pretrain(
-            download_path=model_dir,
-            model_name=raw_model_name,
-            only_download_config=only_download_config,
-            device=device,
-            **kwargs)
+                                 download_path=model_dir,
+                                 model_name=model_name_,
+                                 only_download_config=only_download_config,
+                                 device=device,
+                                 **kwargs)
 
-        download_path = os.path.join(model_dir, raw_model_name)
-        print("*"*20, task_name, model_name)
-        if model_type == "mm" or model_type == "nlp":
+        if model_type == "nlp":
             tokenizer_class = getattr(LazyImport("flagai.data.tokenizer"),
                                       "Tokenizer")
-            self.tokenizer = tokenizer_class.from_pretrained(model_name, cache_dir=download_path)
+            self.tokenizer = tokenizer_class.from_pretrained(
+                model_name, cache_dir=download_path)
+
+        elif model_type == "mm":
+            if model_name.startswith("altdiffusion"):
+                self.tokenizer = None
+            else:
+                self.process = getattr(LazyImport(MODEL_DICT[model_name][4]),
+                                MODEL_DICT[model_name][5]).from_pretrained(os.path.join(model_dir, raw_model_name))
+                self.transform = self.process.feature_extractor
+                self.tokenizer = self.process.tokenizer
+
         else :
             self.tokenizer = None
+            self.transform = None
+
 
     def get_task_name(self, brief_model_name):
         all_model_task = list(ALL_TASK.keys())
@@ -193,6 +234,9 @@ def get_tokenizer(self):
 
     def get_model(self):
         return self.model
+    
+    def get_transform(self):
+        return self.transform
 
     def load_pretrain_params(self, model_path):
         self.model.load_huggingface_weights(model_path)
diff --git a/flagai/data/tokenizer/uni_tokenizer/base_tokenizer.py b/flagai/data/tokenizer/uni_tokenizer/base_tokenizer.py
index 88706b40..dca6bb82 100755
--- a/flagai/data/tokenizer/uni_tokenizer/base_tokenizer.py
+++ b/flagai/data/tokenizer/uni_tokenizer/base_tokenizer.py
@@ -5,10 +5,11 @@
 
 
 class BaseTokenizer(object):
+
     @classmethod
     def from_pretrained(cls,
                         tokenizer_model_name,
-                        cache_dir=None, 
+                        cache_dir=None,
                         *inputs,
                         **kwargs):
         """
@@ -23,7 +24,7 @@ def from_pretrained(cls,
         """
         if cache_dir is None:
             # cache_dir = os.path.join(os.path.dirname(__file__), 'vocabs', f"{tokenizer_model_name}")
-            cache_dir = './checkpoints/'+tokenizer_model_name
+            cache_dir = './checkpoints/' + tokenizer_model_name
         tokenizer_class = ""
         # search the cache directory for certain files
         if os.path.exists(cache_dir):
@@ -35,38 +36,59 @@ def from_pretrained(cls,
             elif VOCAB_FILE in files:
                 tokenizer_class = "wp"
         if tokenizer_class == "":
-            print("downloading model %s from ModelHub"%tokenizer_model_name)
+            print("downloading model %s from ModelHub" % tokenizer_model_name)
             files = _get_model_files(tokenizer_model_name)
             model_id = _get_model_id(tokenizer_model_name)
             if SP_MODEL_FILE in files:
                 tokenizer_class = "sp"
-                _get_vocab_path(cache_dir + '/', SP_MODEL_FILE, model_id, rank=0)
+                _get_vocab_path(cache_dir + '/',
+                                SP_MODEL_FILE,
+                                model_id,
+                                rank=0)
             elif MERGES_FILE in files:
                 tokenizer_class = "bpe"
                 _get_vocab_path(cache_dir + '/', MERGES_FILE, model_id, rank=0)
                 if VOCAB_JSON_FILE in files:
-                    _get_vocab_path(cache_dir + '/', VOCAB_JSON_FILE, model_id, rank=0)
+                    _get_vocab_path(cache_dir + '/',
+                                    VOCAB_JSON_FILE,
+                                    model_id,
+                                    rank=0)
             elif VOCAB_FILE in files:
                 tokenizer_class = "wp"
                 _get_vocab_path(cache_dir + '/', VOCAB_FILE, model_id, rank=0)
             else:
                 raise FileNotFoundError("Error: no tokenizer files")
-        resolved_vocab_json_file = os.path.join(cache_dir, VOCAB_JSON_FILE) if VOCAB_JSON_FILE in files else None
+        resolved_vocab_json_file = os.path.join(
+            cache_dir, VOCAB_JSON_FILE) if VOCAB_JSON_FILE in files else None
         resolved_vocab_file = os.path.join(cache_dir, VOCAB_FILE)
         resolved_merges_file = os.path.join(cache_dir, MERGES_FILE)
         resolved_sp_file = os.path.join(cache_dir, SP_MODEL_FILE)
         if tokenizer_class == "wp":
-            return cls(vocab_file=resolved_vocab_file, tokenizer_class=tokenizer_class,
-                       tokenizer_model_name=tokenizer_model_name, cache_dir=cache_dir, *inputs, **kwargs)
+            return cls(vocab_file=resolved_vocab_file,
+                       tokenizer_class=tokenizer_class,
+                       tokenizer_model_name=tokenizer_model_name,
+                       cache_dir=cache_dir,
+                       *inputs,
+                       **kwargs)
         elif tokenizer_class == "bpe":
-            return cls(vocab_file=resolved_vocab_json_file, merges_file=resolved_merges_file, tokenizer_class=tokenizer_class,
-                    tokenizer_model_name=tokenizer_model_name, cache_dir=cache_dir, *inputs, **kwargs)
+            return cls(vocab_file=resolved_vocab_json_file,
+                       merges_file=resolved_merges_file,
+                       tokenizer_class=tokenizer_class,
+                       tokenizer_model_name=tokenizer_model_name,
+                       cache_dir=cache_dir,
+                       *inputs,
+                       **kwargs)
         elif tokenizer_class == "sp":
-            return cls(sp_model_file=resolved_sp_file, tokenizer_class=tokenizer_class,
-                       tokenizer_model_name=tokenizer_model_name, cache_dir=cache_dir, *inputs, **kwargs)
+            return cls(sp_model_file=resolved_sp_file,
+                       tokenizer_class=tokenizer_class,
+                       tokenizer_model_name=tokenizer_model_name,
+                       cache_dir=cache_dir,
+                       *inputs,
+                       **kwargs)
         else:
-            raise NotImplementedError("Cannot find a tokenizer class that matches the files settings in the directory or ModelHub")
-
+            raise NotImplementedError(
+                "Cannot find a tokenizer class that matches the files settings in the directory or ModelHub"
+            )
 
     def __init__(self,
                  vocab_file=None,
@@ -84,6 +106,4 @@ def __init__(self,
         self.tokenizer_class = tokenizer_class
         self.tokenizer_model_name = tokenizer_model_name
         self.cache_dir = cache_dir
-        self.deprecation_warnings = (
-            {}
-        )
\ No newline at end of file
+        self.deprecation_warnings = ({})
diff --git a/flagai/data/tokenizer/uni_tokenizer/bpe_tokenizer.py b/flagai/data/tokenizer/uni_tokenizer/bpe_tokenizer.py
index 60952ec4..25883388 100644
--- a/flagai/data/tokenizer/uni_tokenizer/bpe_tokenizer.py
+++ b/flagai/data/tokenizer/uni_tokenizer/bpe_tokenizer.py
@@ -29,6 +29,7 @@
 import ftfy
 from functools import lru_cache
 import regex as re
+
 logger = logging.getLogger(__name__)
 # from flagai.data.tokenizer.glm_10b_en.glm_10b_en_tokenizer import bytes_to_unicode, get_pairs
 import sys
@@ -42,6 +43,7 @@ def lru_cache():
 
 
 class BPETokenizer(object):
+
     def __init__(self,
                  vocab_file,
                  merges_file,
@@ -50,7 +52,7 @@ def __init__(self,
                  **kwargs):
         super().__init__(**kwargs)
         self.max_len = max_len if max_len is not None else int(1e12)
-        
+
         self.errors = errors  # how to handle errors in decoding
         self.byte_encoder = bytes_to_unicode()
         self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
@@ -65,7 +67,7 @@ def __init__(self,
                 vocab.append(''.join(merge))
             self.encoder = dict(zip(vocab, range(len(vocab))))
         else:
-            with open(vocab_file) as file:
+            with open(vocab_file, encoding='utf8') as file:
                 self.encoder = json.load(file)
         self.decoder = {v: k for k, v in self.encoder.items()}
 
@@ -178,9 +180,11 @@ def convert_ids_to_tokens(self, ids, skip_special_tokens=False):
     def convert_tokens_to_string(self, tokens, all_command_token={}):
         """Converts a sequence of tokens (string) in a single string."""
         text = "".join(tokens)
-        text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
+        text = bytearray([self.byte_decoder[c]
+                          for c in text]).decode("utf-8", errors=self.errors)
         return text
 
+
 @lru_cache()
 def bytes_to_unicode():
     """
@@ -208,6 +212,7 @@ def bytes_to_unicode():
     cs = [_chr(n) for n in cs]
     return dict(zip(bs, cs))
 
+
 def get_pairs(word):
     """Return set of symbol pairs in a word.
 
@@ -220,58 +225,67 @@ def get_pairs(word):
         prev_char = char
     return pairs
 
+
 def basic_clean(text):
     text = ftfy.fix_text(text)
     text = html.unescape(html.unescape(text))
     return text.strip()
 
+
 def whitespace_clean(text):
     text = re.sub(r'\s+', ' ', text)
     text = text.strip()
     return text
 
+
 class MMBPETokenizer(BPETokenizer):
+
     def __init__(self,
-                vocab_file,
-                merges_file,
-                errors='replace',
-                max_len=None,
-                special_tokens=None,
-                **kwargs):
+                 vocab_file,
+                 merges_file,
+                 errors='replace',
+                 max_len=None,
+                 special_tokens=None,
+                 **kwargs):
         self.byte_encoder = bytes_to_unicode()
         self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
         merges = open(merges_file).read().split('\n')
-        merges = merges[1:49152-256-2+1]
+        merges = merges[1:49152 - 256 - 2 + 1]
         merges = [tuple(merge.split()) for merge in merges]
         vocab = list(bytes_to_unicode().values())
-        vocab = vocab + [v+'</w>' for v in vocab]
+        vocab = vocab + [v + '</w>' for v in vocab]
         for merge in merges:
             vocab.append(''.join(merge))
         if not special_tokens:
             special_tokens = ['<start_of_text>', '<end_of_text>']
         else:
-            special_tokens = ['<start_of_text>', '<end_of_text>'] + special_tokens
+            special_tokens = ['<start_of_text>', '<end_of_text>'
+                              ] + special_tokens
         vocab.extend(special_tokens)
         self.encoder = dict(zip(vocab, range(len(vocab))))
         self.decoder = {v: k for k, v in self.encoder.items()}
         self.bpe_ranks = dict(zip(merges, range(len(merges))))
-        self.cache = {t:t for t in special_tokens}
+        self.cache = {t: t for t in special_tokens}
         special = "|".join(special_tokens)
-        self.pat = re.compile(special + r"""|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+        self.pat = re.compile(
+            special +
+            r"""|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
+            re.IGNORECASE)
 
         # self.vocab_size = len(self.encoder)
         # self.all_special_ids = [self.encoder[t] for t in special_tokens]
     def bpe(self, token):
         if token in self.cache:
             return self.cache[token]
-        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        word = tuple(token[:-1]) + (token[-1] + '</w>', )
         pairs = get_pairs(word)
 
         if not pairs:
-            return token+'</w>'
+            return token + '</w>'
 
         while True:
-            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            bigram = min(
+                pairs, key=lambda pair: self.bpe_ranks.get(pair, float('inf')))
             if bigram not in self.bpe_ranks:
                 break
             first, second = bigram
@@ -286,8 +300,9 @@ def bpe(self, token):
                     new_word.extend(word[i:])
                     break
 
-                if word[i] == first and i < len(word)-1 and word[i+1] == second:
-                    new_word.append(first+second)
+                if word[i] == first and i < len(word) - 1 and word[
+                        i + 1] == second:
+                    new_word.append(first + second)
                     i += 2
                 else:
                     new_word.append(word[i])
@@ -306,16 +321,24 @@ def encode(self, text):
         bpe_tokens = []
         text = whitespace_clean(basic_clean(text)).lower()
         for token in re.findall(self.pat, text):
-            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
-            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+            token = ''.join(self.byte_encoder[b]
+                            for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token]
+                              for bpe_token in self.bpe(token).split(' '))
         return bpe_tokens
 
     def decode(self, tokens):
         text = ''.join([self.decoder[token] for token in tokens])
-        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        text = bytearray([self.byte_decoder[c] for c in text
+                          ]).decode('utf-8',
+                                    errors="replace").replace('</w>', ' ')
         return text
 
-    def tokenize(self, texts: Union[str, List[str]], sot_token: int, eot_token: int, context_length: int = 77) -> torch.LongTensor:
+    def tokenize(self,
+                 texts: Union[str, List[str]],
+                 sot_token: int,
+                 eot_token: int,
+                 context_length: int = 77) -> torch.LongTensor:
         """
         Returns the tokenized representation of given input string(s)
 
@@ -333,11 +356,12 @@ def tokenize(self, texts: Union[str, List[str]], sot_token: int, eot_token: int,
         if isinstance(texts, str):
             texts = [texts]
 
-        all_tokens = [[sot_token] + self.encode(text) + [eot_token] for text in texts]
+        all_tokens = [[sot_token] + self.encode(text) + [eot_token]
+                      for text in texts]
         result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
 
         for i, tokens in enumerate(all_tokens):
             if len(tokens) > context_length:
                 tokens = tokens[:context_length]  # Truncate
             result[i, :len(tokens)] = torch.tensor(tokens)
-        return result
\ No newline at end of file
+        return result
diff --git a/flagai/data/tokenizer/uni_tokenizer/difffusion_bert_tokenizer.py b/flagai/data/tokenizer/uni_tokenizer/difffusion_bert_tokenizer.py
new file mode 100755
index 00000000..1811facf
--- /dev/null
+++ b/flagai/data/tokenizer/uni_tokenizer/difffusion_bert_tokenizer.py
@@ -0,0 +1,474 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokenization classes."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import sys
+import collections
+import re
+import unicodedata
+import six
+from functools import lru_cache
+import torch
+import os
+from typing import Union, List
+
+unicode = str
+
+
+@lru_cache()
+def default_vocab():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                        "vocab.txt")
+
+
+def validate_case_matches_checkpoint(do_lower_case, init_checkpoint):
+    """Checks whether the casing config is consistent with the checkpoint name."""
+
+    # The casing has to be passed in by the user and there is no explicit check
+    # as to whether it matches the checkpoint. The casing information probably
+    # should have been stored in the bert_config.json file, but it's not, so
+    # we have to heuristically detect it to validate.
+
+    if not init_checkpoint:
+        return
+
+    m = re.match("^.*?([A-Za-z0-9_-]+)/bert_model.ckpt", init_checkpoint)
+    if m is None:
+        return
+
+    model_name = m.group(1)
+
+    lower_models = [
+        "uncased_L-24_H-1024_A-16", "uncased_L-12_H-768_A-12",
+        "multilingual_L-12_H-768_A-12", "chinese_L-12_H-768_A-12"
+    ]
+
+    cased_models = [
+        "cased_L-12_H-768_A-12", "cased_L-24_H-1024_A-16",
+        "multi_cased_L-12_H-768_A-12"
+    ]
+
+    is_bad_config = False
+    if model_name in lower_models and not do_lower_case:
+        is_bad_config = True
+        actual_flag = "False"
+        case_name = "lowercased"
+        opposite_flag = "True"
+
+    if model_name in cased_models and do_lower_case:
+        is_bad_config = True
+        actual_flag = "True"
+        case_name = "cased"
+        opposite_flag = "False"
+
+    if is_bad_config:
+        raise ValueError(
+            "You passed in `--do_lower_case=%s` with `--init_checkpoint=%s`. "
+            "However, `%s` seems to be a %s model, so you "
+            "should pass in `--do_lower_case=%s` so that the fine-tuning matches "
+            "how the model was pre-training. If this error is wrong, please "
+            "just comment out this check." %
+            (actual_flag, init_checkpoint, model_name, case_name,
+             opposite_flag))
+
+
+def convert_to_unicode(text):
+    """Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
+    if six.PY3:
+        if isinstance(text, str):
+            return text
+        elif isinstance(text, bytes):
+            return text.decode("utf-8", "ignore")
+        else:
+            raise ValueError("Unsupported string type: %s" % (type(text)))
+    elif six.PY2:
+        if isinstance(text, str):
+            return text.decode("utf-8", "ignore")
+        elif isinstance(text, unicode):
+            return text
+        else:
+            raise ValueError("Unsupported string type: %s" % (type(text)))
+    else:
+        raise ValueError("Not running on Python2 or Python 3?")
+
+
+def printable_text(text):
+    """Returns text encoded in a way suitable for print or `tf.logging`."""
+
+    # These functions want `str` for both Python2 and Python3, but in one case
+    # it's a Unicode string and in the other it's a byte string.
+    if six.PY3:
+        if isinstance(text, str):
+            return text
+        elif isinstance(text, bytes):
+            return text.decode("utf-8", "ignore")
+        else:
+            raise ValueError("Unsupported string type: %s" % (type(text)))
+    elif six.PY2:
+        if isinstance(text, str):
+            return text
+        elif isinstance(text, unicode):
+            return text.encode("utf-8")
+        else:
+            raise ValueError("Unsupported string type: %s" % (type(text)))
+    else:
+        raise ValueError("Not running on Python2 or Python 3?")
+
+
+def load_vocab(vocab_file):
+    """Loads a vocabulary file into a dictionary."""
+    vocab = collections.OrderedDict()
+    index = 0
+    with open(vocab_file, "r") as reader:
+        while True:
+            token = convert_to_unicode(reader.readline())
+            if not token:
+                break
+            token = token.strip()
+            vocab[token] = index
+            index += 1
+    return vocab
+
+
+def convert_by_vocab(vocab, items):
+    """Converts a sequence of [tokens|ids] using the vocab."""
+    output = []
+    for item in items:
+        output.append(vocab[item])
+    return output
+
+
+def convert_tokens_to_ids(vocab, tokens):
+    return convert_by_vocab(vocab, tokens)
+
+
+def convert_ids_to_tokens(inv_vocab, ids):
+    return convert_by_vocab(inv_vocab, ids)
+
+
+def whitespace_tokenize(text):
+    """Runs basic whitespace cleaning and splitting on a piece of text."""
+    text = text.strip()
+    if not text:
+        return []
+    tokens = text.split()
+    return tokens
+
+
+class FullTokenizer(object):
+    """Runs end-to-end tokenziation."""
+
+    def __init__(self, vocab_file=default_vocab(), do_lower_case=True):
+        self.vocab = load_vocab(vocab_file)
+        self.inv_vocab = {v: k for k, v in self.vocab.items()}
+        self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case)
+        self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab)
+
+    def _tokenize(self, text):
+        split_tokens = []
+        for token in self.basic_tokenizer.tokenize(text):
+            for sub_token in self.wordpiece_tokenizer.tokenize(token):
+                split_tokens.append(sub_token)
+
+        return split_tokens
+
+    def tokenize(self,
+                 texts: Union[str, List[str]],
+                 context_length: int = 64) -> torch.LongTensor:
+        """
+        Returns the tokenized representation of given input string(s)
+        Parameters
+        ----------
+        texts : Union[str, List[str]]
+            An input string or a list of input strings to tokenize
+        context_length : int
+            The context length to use; all baseline models use 24 as the context length
+        Returns
+        -------
+        A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+
+        all_tokens = []
+        for text in texts:
+            all_tokens.append([self.vocab['[CLS]']] +
+                              self.convert_tokens_to_ids(self._tokenize(text))
+                              [:context_length - 2] + [self.vocab['[SEP]']])
+
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+        for i, tokens in enumerate(all_tokens):
+            assert len(tokens) <= context_length
+            result[i, :len(tokens)] = torch.tensor(tokens)
+
+        return result
+
+    def convert_tokens_to_ids(self, tokens):
+        return convert_by_vocab(self.vocab, tokens)
+
+    def convert_token_to_id(self, token):
+        return self.vocab[token]
+
+    def convert_ids_to_tokens(self, ids):
+        return convert_by_vocab(self.inv_vocab, ids)
+
+    @staticmethod
+    def convert_tokens_to_string(tokens, clean_up_tokenization_spaces=True):
+        """ Converts a sequence of tokens (string) in a single string. """
+
+        def clean_up_tokenization(out_string):
+            """ Clean up a list of simple English tokenization artifacts
+            like spaces before punctuations and abreviated forms.
+            """
+            out_string = (out_string.replace(" .", ".").replace(
+                " ?", "?").replace(" !", "!").replace(" ,", ",").replace(
+                    " ' ",
+                    "'").replace(" n't", "n't").replace(" 'm", "'m").replace(
+                        " 's", "'s").replace(" 've",
+                                             "'ve").replace(" 're", "'re"))
+            return out_string
+
+        text = ' '.join(tokens).replace(' ##', '').strip()
+        if clean_up_tokenization_spaces:
+            clean_text = clean_up_tokenization(text)
+            return clean_text
+        else:
+            return text
+
+    @property
+    def vocab_size(self):
+        return len(self.vocab)
+
+
+class BasicTokenizer(object):
+    """Runs basic tokenization (punctuation splitting, lower casing, etc.)."""
+
+    def __init__(self, do_lower_case=True):
+        """Constructs a BasicTokenizer.
+
+        Args:
+          do_lower_case: Whether to lower case the input.
+        """
+        self.do_lower_case = do_lower_case
+
+    def tokenize(self, text):
+        """Tokenizes a piece of text."""
+        text = convert_to_unicode(text)
+        text = self._clean_text(text)
+
+        # This was added on November 1st, 2018 for the multilingual and Chinese
+        # models. This is also applied to the English models now, but it doesn't
+        # matter since the English models were not trained on any Chinese data
+        # and generally don't have any Chinese data in them (there are Chinese
+        # characters in the vocabulary because Wikipedia does have some Chinese
+        # words in the English Wikipedia.).
+        text = self._tokenize_chinese_chars(text)
+
+        orig_tokens = whitespace_tokenize(text)
+        split_tokens = []
+        for token in orig_tokens:
+            if self.do_lower_case:
+                token = token.lower()
+                token = self._run_strip_accents(token)
+            split_tokens.extend(self._run_split_on_punc(token))
+
+        output_tokens = whitespace_tokenize(" ".join(split_tokens))
+        return output_tokens
+
+    def _run_strip_accents(self, text):
+        """Strips accents from a piece of text."""
+        text = unicodedata.normalize("NFD", text)
+        output = []
+        for char in text:
+            cat = unicodedata.category(char)
+            if cat == "Mn":
+                continue
+            output.append(char)
+        return "".join(output)
+
+    def _run_split_on_punc(self, text):
+        """Splits punctuation on a piece of text."""
+        chars = list(text)
+        i = 0
+        start_new_word = True
+        output = []
+        while i < len(chars):
+            char = chars[i]
+            if _is_punctuation(char):
+                output.append([char])
+                start_new_word = True
+            else:
+                if start_new_word:
+                    output.append([])
+                start_new_word = False
+                output[-1].append(char)
+            i += 1
+
+        return ["".join(x) for x in output]
+
+    def _tokenize_chinese_chars(self, text):
+        """Adds whitespace around any CJK character."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if self._is_chinese_char(cp):
+                output.append(" ")
+                output.append(char)
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+    def _is_chinese_char(self, cp):
+        """Checks whether CP is the codepoint of a CJK character."""
+        # This defines a "chinese character" as anything in the CJK Unicode block:
+        #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+        #
+        # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+        # despite its name. The modern Korean Hangul alphabet is a different block,
+        # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+        # space-separated words, so they are not treated specially and handled
+        # like the all of the other languages.
+        if ((cp >= 0x4E00 and cp <= 0x9FFF) or  #
+            (cp >= 0x3400 and cp <= 0x4DBF) or  #
+            (cp >= 0x20000 and cp <= 0x2A6DF) or  #
+            (cp >= 0x2A700 and cp <= 0x2B73F) or  #
+            (cp >= 0x2B740 and cp <= 0x2B81F) or  #
+            (cp >= 0x2B820 and cp <= 0x2CEAF) or
+            (cp >= 0xF900 and cp <= 0xFAFF) or  #
+            (cp >= 0x2F800 and cp <= 0x2FA1F)):  #
+            return True
+
+        return False
+
+    def _clean_text(self, text):
+        """Performs invalid character removal and whitespace cleanup on text."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if cp == 0 or cp == 0xfffd or _is_control(char):
+                continue
+            if _is_whitespace(char):
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+
+class WordpieceTokenizer(object):
+    """Runs WordPiece tokenziation."""
+
+    def __init__(self, vocab, unk_token="[UNK]", max_input_chars_per_word=200):
+        self.vocab = vocab
+        self.unk_token = unk_token
+        self.max_input_chars_per_word = max_input_chars_per_word
+
+    def tokenize(self, text):
+        """Tokenizes a piece of text into its word pieces.
+
+        This uses a greedy longest-match-first algorithm to perform tokenization
+        using the given vocabulary.
+
+        For example:
+          input = "unaffable"
+          output = ["un", "##aff", "##able"]
+
+        Args:
+          text: A single token or whitespace separated tokens. This should have
+            already been passed through `BasicTokenizer.
+
+        Returns:
+          A list of wordpiece tokens.
+        """
+
+        text = convert_to_unicode(text)
+
+        output_tokens = []
+        for token in whitespace_tokenize(text):
+            chars = list(token)
+            if len(chars) > self.max_input_chars_per_word:
+                output_tokens.append(self.unk_token)
+                continue
+
+            is_bad = False
+            start = 0
+            sub_tokens = []
+            while start < len(chars):
+                end = len(chars)
+                cur_substr = None
+                while start < end:
+                    substr = "".join(chars[start:end])
+                    if start > 0:
+                        substr = "##" + substr
+                    if substr in self.vocab:
+                        cur_substr = substr
+                        break
+                    end -= 1
+                if cur_substr is None:
+                    is_bad = True
+                    break
+                sub_tokens.append(cur_substr)
+                start = end
+
+            if is_bad:
+                output_tokens.append(self.unk_token)
+            else:
+                output_tokens.extend(sub_tokens)
+        return output_tokens
+
+
+def _is_whitespace(char):
+    """Checks whether `chars` is a whitespace character."""
+    # \t, \n, and \r are technically contorl characters but we treat them
+    # as whitespace since they are generally considered as such.
+    if char == " " or char == "\t" or char == "\n" or char == "\r":
+        return True
+    cat = unicodedata.category(char)
+    if cat == "Zs":
+        return True
+    return False
+
+
+def _is_control(char):
+    """Checks whether `chars` is a control character."""
+    # These are technically control characters but we count them as whitespace
+    # characters.
+    if char == "\t" or char == "\n" or char == "\r":
+        return False
+    cat = unicodedata.category(char)
+    if cat in ("Cc", "Cf"):
+        return True
+    return False
+
+
+def _is_punctuation(char):
+    """Checks whether `chars` is a punctuation character."""
+    cp = ord(char)
+    # We treat all non-letter/number ASCII as punctuation.
+    # Characters such as "^", "$", and "`" are not in the Unicode
+    # Punctuation class but we treat them as punctuation anyways, for
+    # consistency.
+    if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64)
+            or (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)):
+        return True
+    cat = unicodedata.category(char)
+    if cat.startswith("P"):
+        return True
+    return False
diff --git a/flagai/data/tokenizer/uni_tokenizer/sp_tokenizer.py b/flagai/data/tokenizer/uni_tokenizer/sp_tokenizer.py
index 4426559a..4067e248 100644
--- a/flagai/data/tokenizer/uni_tokenizer/sp_tokenizer.py
+++ b/flagai/data/tokenizer/uni_tokenizer/sp_tokenizer.py
@@ -18,11 +18,13 @@
 """Utilities for using and training tokenizers (char, wordpiece, sentencepiece)"""
 
 import logging
+
 logger = logging.getLogger(__name__)
 import sentencepiece as spm
 
 
 class SentencePieceTokenizer(object):
+
     def __init__(self, model_path):
         self.sp_model = spm.SentencePieceProcessor()
         self.sp_model.Load(model_path)
@@ -35,7 +37,10 @@ def vocab_size(self):
         return self.sp_model.GetPieceSize()
 
     def get_vocab(self):
-        vocab = {self.convert_id_to_token(i): i for i in range(self.vocab_size)}
+        vocab = {
+            self.convert_id_to_token(i): i
+            for i in range(self.vocab_size)
+        }
         # vocab.update(self.added_tokens_encoder)
         return vocab
 
@@ -61,10 +66,10 @@ def convert_tokens_to_string(self, tokens, all_command_token={}):
         for token in tokens:
             # make sure that special tokens are not decoded using sentencepiece model
             if token in all_command_token:
-                out_string += self.sp_model.decode_pieces(current_sub_tokens) + token + " "
+                out_string += self.sp_model.decode_pieces(
+                    current_sub_tokens) + token + " "
                 current_sub_tokens = []
             else:
                 current_sub_tokens.append(token)
         out_string += self.sp_model.decode_pieces(current_sub_tokens)
         return out_string.strip()
-
diff --git a/flagai/data/tokenizer/uni_tokenizer/tokenizer.py b/flagai/data/tokenizer/uni_tokenizer/tokenizer.py
index 4ef31de3..dd264b7f 100755
--- a/flagai/data/tokenizer/uni_tokenizer/tokenizer.py
+++ b/flagai/data/tokenizer/uni_tokenizer/tokenizer.py
@@ -19,18 +19,20 @@
 # from collections import namedtuple
 # import itertools
 
-
 import itertools
 import logging
+
 logger = logging.getLogger(__name__)
 from flagai.data.tokenizer.tokenizer import CommandToken
 from flagai.data.tokenizer.uni_tokenizer.wp_tokenizer import WordpieceTokenizer
 from flagai.data.tokenizer.uni_tokenizer.bpe_tokenizer import BPETokenizer, MMBPETokenizer
 from flagai.data.tokenizer.uni_tokenizer.sp_tokenizer import SentencePieceTokenizer
 from flagai.data.tokenizer.uni_tokenizer.base_tokenizer import BaseTokenizer
+from flagai.data.tokenizer.uni_tokenizer.difffusion_bert_tokenizer import FullTokenizer
 from typing import List, Union, Optional
 import unicodedata
 
+
 def is_control(ch):
     """
     https://en.wikipedia.org/wiki/Control_character
@@ -40,7 +42,9 @@ def is_control(ch):
     """
     return unicodedata.category(ch) in ('Cc', 'Cf')
 
+
 class Tokenizer(BaseTokenizer):
+
     def __init__(self,
                  add_block_symbols=True,
                  add_sentinel_token=0,
@@ -52,14 +56,19 @@ def __init__(self,
 
         if self.tokenizer_class == "wp":
             if self.tokenizer_model_name.lower().endswith("ch"):
-                self.text_tokenizer = WordpieceTokenizer(self.vocab_file, is_ch=True)
+                self.text_tokenizer = WordpieceTokenizer(self.vocab_file,
+                                                         is_ch=True)
+            elif self.tokenizer_model_name.lower().startswith('clip-cn'):
+                self.text_tokenizer = FullTokenizer(self.vocab_file)
             else:
                 self.text_tokenizer = WordpieceTokenizer(self.vocab_file)
         elif self.tokenizer_class == "bpe":
             if self.tokenizer_model_name.lower().startswith('clip'):
-                self.text_tokenizer = MMBPETokenizer(self.vocab_file, self.merges_file)
+                self.text_tokenizer = MMBPETokenizer(self.vocab_file,
+                                                     self.merges_file)
             else:
-                self.text_tokenizer = BPETokenizer(self.vocab_file, self.merges_file)
+                self.text_tokenizer = BPETokenizer(self.vocab_file,
+                                                   self.merges_file)
         elif self.tokenizer_class == "sp":
             self.text_tokenizer = SentencePieceTokenizer(self.sp_model_file)
         else:
@@ -79,34 +88,61 @@ def __init__(self,
             self.token_pad_id = None
             try:
                 self._command_tokens = [
-                    CommandToken('pad', '[PAD]', self.text_tokenizer.convert_token_to_id('[PAD]')),
-                    CommandToken('cls', '[CLS]', self.text_tokenizer.convert_token_to_id('[CLS]')),
-                    CommandToken('MASK', '[MASK]',
-                                 self.text_tokenizer.convert_token_to_id('[MASK]')),
-                    CommandToken('unk', '[UNK]', self.text_tokenizer.convert_token_to_id('[UNK]')),
-                    CommandToken('sep', '[SEP]', self.text_tokenizer.convert_token_to_id('[SEP]')),
-                    CommandToken('eos', '[PAD]', self.text_tokenizer.convert_token_to_id('[PAD]')),
-
+                    CommandToken(
+                        'pad', '[PAD]',
+                        self.text_tokenizer.convert_token_to_id('[PAD]')),
+                    CommandToken(
+                        'cls', '[CLS]',
+                        self.text_tokenizer.convert_token_to_id('[CLS]')),
+                    CommandToken(
+                        'MASK', '[MASK]',
+                        self.text_tokenizer.convert_token_to_id('[MASK]')),
+                    CommandToken(
+                        'unk', '[UNK]',
+                        self.text_tokenizer.convert_token_to_id('[UNK]')),
+                    CommandToken(
+                        'sep', '[SEP]',
+                        self.text_tokenizer.convert_token_to_id('[SEP]')),
+                    CommandToken(
+                        'eos', '[PAD]',
+                        self.text_tokenizer.convert_token_to_id('[PAD]')),
                 ]
-                self.token_start_id = self.text_tokenizer.convert_token_to_id('[CLS]')
-                self.token_end_id = self.text_tokenizer.convert_token_to_id('[SEP]')
-                self.token_pad_id = self.text_tokenizer.convert_token_to_id('[PAD]')
+                self.token_start_id = self.text_tokenizer.convert_token_to_id(
+                    '[CLS]')
+                self.token_end_id = self.text_tokenizer.convert_token_to_id(
+                    '[SEP]')
+                self.token_pad_id = self.text_tokenizer.convert_token_to_id(
+                    '[PAD]')
                 self.text_tokenizer._token_cls = "[CLS]"
                 self.text_tokenizer._token_sep = "[SEP]"
 
             except KeyError:
                 self._command_tokens = [
-                    CommandToken('pad', '[PAD]', self.text_tokenizer.convert_token_to_id('<pad>')),
-                    CommandToken('cls', '[CLS]', self.text_tokenizer.convert_token_to_id('
')),
-                    CommandToken('MASK', '[MASK]',
-                                 self.text_tokenizer.convert_token_to_id('<mask>')),
-                    CommandToken('unk', '[UNK]', self.text_tokenizer.convert_token_to_id('<unk>')),
-                    CommandToken('sep', '[SEP]', self.text_tokenizer.convert_token_to_id('<sep>')),
-                    CommandToken('eos', '[PAD]', self.text_tokenizer.convert_token_to_id('
')),
+                    CommandToken(
+                        'pad', '[PAD]',
+                        self.text_tokenizer.convert_token_to_id('<pad>')),
+                    CommandToken(
+                        'cls', '[CLS]',
+                        self.text_tokenizer.convert_token_to_id('
')),
+                    CommandToken(
+                        'MASK', '[MASK]',
+                        self.text_tokenizer.convert_token_to_id('<mask>')),
+                    CommandToken(
+                        'unk', '[UNK]',
+                        self.text_tokenizer.convert_token_to_id('<unk>')),
+                    CommandToken(
+                        'sep', '[SEP]',
+                        self.text_tokenizer.convert_token_to_id('<sep>')),
+                    CommandToken(
+                        'eos', '[PAD]',
+                        self.text_tokenizer.convert_token_to_id('
')),
                 ]
-                self.token_start_id = self.text_tokenizer.convert_token_to_id('
')
-                self.token_end_id = self.text_tokenizer.convert_token_to_id('
')
-                self.token_pad_id = self.text_tokenizer.convert_token_to_id('<pad>')
+                self.token_start_id = self.text_tokenizer.convert_token_to_id(
+                    '
')
+                self.token_end_id = self.text_tokenizer.convert_token_to_id(
+                    '
')
+                self.token_pad_id = self.text_tokenizer.convert_token_to_id(
+                    '<pad>')
                 self.text_tokenizer._token_cls = "
"
                 self.text_tokenizer._token_sep = "
"
             if add_block_symbols:
@@ -131,7 +167,8 @@ def __init__(self,
             if add_sentinel_token > 0:
                 for i in range(1, add_sentinel_token):
                     self._command_tokens.extend([
-                        CommandToken(f'MASK{i}', f'[MASK{i}]', self.num_tokens),
+                        CommandToken(f'MASK{i}', f'[MASK{i}]',
+                                     self.num_tokens),
                         CommandToken(f'sop{i}', f'<|startofpiece{i}|>',
                                      self.num_tokens + 1)
                     ])
@@ -142,45 +179,59 @@ def __init__(self,
                 self.num_command_tokens = 6
                 self.num_text_tokens = self.num_tokens - 3
                 self._command_tokens = [
-                    CommandToken('pad', '
',
-                                 self.text_tokenizer.convert_token_to_id('
')),
-                    CommandToken('eos', '
',
-                                 self.text_tokenizer.convert_token_to_id('
')),
-                    CommandToken('sep', '[SEP]',
-                                 self.text_tokenizer.convert_token_to_id('
')),
-                    CommandToken('cls', '[CLS]',
-                                 self.text_tokenizer.convert_token_to_id('
')),
-                    CommandToken('MASK',
-                                 '[MASK]',
-                                 self.text_tokenizer.convert_token_to_id('<mask>'),
-                                 lstrip=True),
-                    CommandToken('unk', '[UNK]',
-                                 self.text_tokenizer.convert_token_to_id('<unk>'))
+                    CommandToken(
+                        'pad', '
',
+                        self.text_tokenizer.convert_token_to_id('
')),
+                    CommandToken(
+                        'eos', '
',
+                        self.text_tokenizer.convert_token_to_id('
')),
+                    CommandToken(
+                        'sep', '[SEP]',
+                        self.text_tokenizer.convert_token_to_id('
')),
+                    CommandToken(
+                        'cls', '[CLS]',
+                        self.text_tokenizer.convert_token_to_id('
')),
+                    CommandToken(
+                        'MASK',
+                        '[MASK]',
+                        self.text_tokenizer.convert_token_to_id('<mask>'),
+                        lstrip=True),
+                    CommandToken(
+                        'unk', '[UNK]',
+                        self.text_tokenizer.convert_token_to_id('<unk>'))
                 ]
                 if add_block_symbols:
                     self._command_tokens.extend([
-                        CommandToken('sop', '
', self.num_tokens),
-                        CommandToken('eop', '
', self.num_tokens + 1)
+                        CommandToken('sop', '
',
+                                     self.num_tokens),
+                        CommandToken('eop', '
',
+                                     self.num_tokens + 1)
                     ])
                     self.num_tokens += 2
                     self.num_command_tokens += 2
             elif self.tokenizer_model_name.lower().startswith('clip'):
                 self.num_command_tokens = 2
                 self._command_tokens = [
-                    CommandToken('sot', '<start_of_text>',
-                                 self.text_tokenizer.convert_token_to_id('
')),
-                    CommandToken('eot', '<end_of_text>',
-                                 self.text_tokenizer.convert_token_to_id('
')),
+                    CommandToken(
+                        'sot', '<start_of_text>',
+                        self.text_tokenizer.convert_token_to_id('
')),
+                    CommandToken(
+                        'eot', '<end_of_text>',
+                        self.text_tokenizer.convert_token_to_id('
')),
                 ]
-                self.num_tokens += self.num_command_tokens 
+                self.num_tokens += self.num_command_tokens
             else:
                 self.num_command_tokens = 2
                 self.num_text_tokens = self.num_tokens - 1
                 self._command_tokens = [
-                    CommandToken('pad', '
',
-                                 self.text_tokenizer.convert_token_to_id('
')),
-                    CommandToken('eos', '
',
-                                 self.text_tokenizer.convert_token_to_id('
'))
+                    CommandToken(
+                        'pad', '
',
+                        self.text_tokenizer.convert_token_to_id(
+                            '
')),
+                    CommandToken(
+                        'eos', '
',
+                        self.text_tokenizer.convert_token_to_id(
+                            '
'))
                 ]
                 if add_block_symbols:
                     if self.tokenizer_model_name.lower().startswith('glm'):
@@ -188,12 +239,16 @@ def __init__(self,
                         cls_token_id = self.num_tokens + 2
                         num_tokens_to_add = 5
                     else:
-                        unk_token_id = self.text_tokenizer.convert_token_to_id('
')
-                        cls_token_id = self.text_tokenizer.convert_token_to_id('
')
+                        unk_token_id = self.text_tokenizer.convert_token_to_id(
+                            '
')
+                        cls_token_id = self.text_tokenizer.convert_token_to_id(
+                            '
')
                         num_tokens_to_add = 4
                     self._command_tokens.extend([
-                        CommandToken('sop', '
', self.num_tokens),
-                        CommandToken('eop', '
', self.num_tokens + 1),
+                        CommandToken('sop', '
',
+                                     self.num_tokens),
+                        CommandToken('eop', '
',
+                                     self.num_tokens + 1),
                         CommandToken('cls', '[CLS]', cls_token_id),
                         CommandToken('MASK',
                                      '[MASK]',
@@ -251,7 +306,8 @@ def __init__(self,
             self.num_command_tokens += 6
             if add_block_symbols:
                 self._command_tokens.extend([
-                    CommandToken('sop', '
', self.num_tokens + 1),
+                    CommandToken('sop', '
',
+                                 self.num_tokens + 1),
                     CommandToken('eop', '
', self.num_tokens + 2)
                 ])
                 if fix_command_token:
@@ -390,20 +446,24 @@ def DecodeIds(self, ids):
                 tokens.append(self.command_id_map[id].token)
             else:
                 try:
-                    tokens.extend(self.text_tokenizer.convert_ids_to_tokens([id]))
+                    tokens.extend(
+                        self.text_tokenizer.convert_ids_to_tokens([id]))
                 except KeyError:
                     pass
-        return self.text_tokenizer.convert_tokens_to_string(tokens, self.command_token_map)
+        return self.text_tokenizer.convert_tokens_to_string(
+            tokens, self.command_token_map)
 
     def encode(self, text):
-        return self.text_tokenizer.convert_tokens_to_ids(self.text_tokenizer.tokenize(text))
+        return self.text_tokenizer.convert_tokens_to_ids(
+            self.text_tokenizer.tokenize(text))
 
     def decode(self, ids):
         return self.DecodeIds(ids)
 
     def DecodeTokens(self, tokens):
         """converts wordpiece tokens to a text string"""
-        return self.text_tokenizer.convert_tokens_to_string(tokens, self.command_token_map)
+        return self.text_tokenizer.convert_tokens_to_string(
+            tokens, self.command_token_map)
 
     def EncodeAsIds(self, text, process_fn=None):
         """
@@ -503,12 +563,12 @@ def get_input_ids(text):
         )
 
     def prepare_for_model(
-            self,
-            ids: List[int],
-            pair_ids: Optional[List[int]] = None,
-            add_special_tokens: bool = True,
-            truncation: Union[bool, str] = True,
-            max_length: Optional[int] = None,
+        self,
+        ids: List[int],
+        pair_ids: Optional[List[int]] = None,
+        add_special_tokens: bool = True,
+        truncation: Union[bool, str] = True,
+        max_length: Optional[int] = None,
     ):
 
         pair = bool(pair_ids is not None)
@@ -536,7 +596,8 @@ def prepare_for_model(
                 token_type_ids = [0] * (len(ids) + 2) + [1] * (len(pair_ids) +
                                                                1)
             else:
-                sequence = [self.get_command_id("cls")] + ids + [self.get_command_id("eos")]
+                sequence = [self.get_command_id("cls")
+                            ] + ids + [self.token_end_id]
                 token_type_ids = [0] * (len(ids) + 2)
         else:
             sequence = ids + pair_ids if pair else ids
@@ -548,15 +609,16 @@ def prepare_for_model(
         return encoded_inputs
 
     def encode_plus(  #for Seq2seq
-            self,
-            source_text: str,
-            second_text=None,
-            target_text=None,
-            truncation=True,
-            max_length=None,
+        self,
+        source_text: str,
+        second_text=None,
+        target_text=None,
+        truncation=True,
+        max_length=None,
     ):
         if not self.tokenizer_model_name.lower().startswith("glm"):
-            return self.encode_plus_non_glm(source_text, second_text, truncation, max_length)
+            return self.encode_plus_non_glm(source_text, second_text,
+                                            truncation, max_length)
         sop_id = self.get_command_id('sop')  #start of piece
         eop_id = self.get_command_id('eop')  #end of piece
         sep_id = self.get_command_id('sep')  #seperation
@@ -631,7 +693,9 @@ def tokenize_as_tensor(self, texts):
         """
         sot_token = self.get_command_id('sot')
         eot_token = self.get_command_id('eot')
-        return self.text_tokenizer.tokenize(texts, sot_token=sot_token, eot_token=eot_token)
+        return self.text_tokenizer.tokenize(texts,
+                                            sot_token=sot_token,
+                                            eot_token=eot_token)
 
     def tokenize(self, text, maxlen=None, add_spatial_tokens=False):
         tokens = self.text_tokenizer.tokenize(text)
@@ -643,9 +707,4 @@ def tokenize(self, text, maxlen=None, add_spatial_tokens=False):
         if maxlen is not None:
             index = int(self.get_command_id('sep') is not None) + 1
             self.truncate_sequence(maxlen, tokens, pop_index=-index)
-        return tokens
-
-
-
-
-
+        return tokens
\ No newline at end of file
diff --git a/flagai/data/tokenizer/uni_tokenizer/wp_tokenizer.py b/flagai/data/tokenizer/uni_tokenizer/wp_tokenizer.py
index 4e6e8531..dfbc97b5 100755
--- a/flagai/data/tokenizer/uni_tokenizer/wp_tokenizer.py
+++ b/flagai/data/tokenizer/uni_tokenizer/wp_tokenizer.py
@@ -19,9 +19,8 @@
 # from collections import namedtuple
 # import itertools
 
-
-
 import logging
+
 logger = logging.getLogger(__name__)
 import os
 # from flagai.data.tokenizer.glm_large_en.wordpiece import load_vocab, BasicTokenizer, whitespace_tokenize
@@ -31,15 +30,23 @@
 
 
 class WordpieceTokenizer(object):
-    def __init__(self, vocab_file=None, do_basic_tokenize=True,
-                         do_lower_case=True, max_len=None,
-                         never_split=("[UNK]", "[SEP]", "[PAD]", "[CLS]", "[MASK]"),
-                         unk_token="[UNK]", max_input_chars_per_word=100, is_ch=False, *input, **kwargs):
+
+    def __init__(self,
+                 vocab_file=None,
+                 do_basic_tokenize=True,
+                 do_lower_case=True,
+                 max_len=None,
+                 never_split=("[UNK]", "[SEP]", "[PAD]", "[CLS]", "[MASK]"),
+                 unk_token="[UNK]",
+                 max_input_chars_per_word=100,
+                 is_ch=False,
+                 *input,
+                 **kwargs):
         if not os.path.isfile(vocab_file):
             raise ValueError(
                 "Can't find a vocabulary file at path '{}'. To load the vocabulary from a Google pretrained "
                 "model use `tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
-                    .format(vocab_file))
+                .format(vocab_file))
         self.vocab = load_vocab(vocab_file)
         self.ids_to_tokens = collections.OrderedDict([
             (ids, tok) for tok, ids in self.vocab.items()
@@ -174,6 +181,7 @@ def convert_tokens_to_string(self, tokens, all_command_token={}):
             out_string = " ".join(tokens).replace(" ##", "").strip()
         return out_string
 
+
 def load_vocab(vocab_file):
     """Loads a vocabulary file into a dictionary."""
     vocab = collections.OrderedDict()
@@ -203,7 +211,6 @@ def whitespace_tokenize(text):
 class BasicTokenizer(object):
     """Runs basic tokenization (punctuation splitting, lower casing, etc.)."""
 
-
     def __init__(self,
                  do_lower_case=True,
                  never_split=("[UNK]", "[SEP]", "[PAD]", "[CLS]", "[MASK]")):
@@ -214,7 +221,6 @@ def __init__(self,
         self.do_lower_case = do_lower_case
         self.never_split = never_split
 
-
     def tokenize(self, text):
         """Tokenizes a piece of text."""
         text = self._clean_text(text)
@@ -236,7 +242,6 @@ def tokenize(self, text):
         output_tokens = whitespace_tokenize(" ".join(split_tokens))
         return output_tokens
 
-
     def _run_strip_accents(self, text):
         """Strips accents from a piece of text."""
         text = unicodedata.normalize("NFD", text)
@@ -248,7 +253,6 @@ def _run_strip_accents(self, text):
             output.append(char)
         return "".join(output)
 
-
     def _run_split_on_punc(self, text):
         """Splits punctuation on a piece of text."""
         if text in self.never_split:
@@ -271,7 +275,6 @@ def _run_split_on_punc(self, text):
 
         return ["".join(x) for x in output]
 
-
     def _tokenize_chinese_chars(self, text):
         """Adds whitespace around any CJK character."""
         output = []
@@ -285,7 +288,6 @@ def _tokenize_chinese_chars(self, text):
                 output.append(char)
         return "".join(output)
 
-
     def _is_chinese_char(self, cp):
         """Checks whether CP is the codepoint of a CJK character."""
         # This defines a "chinese character" as anything in the CJK Unicode block:
@@ -297,18 +299,17 @@ def _is_chinese_char(self, cp):
         # space-separated words, so they are not treated specially and handled
         # like the all of the other languages.
         if ((cp >= 0x4E00 and cp <= 0x9FFF) or  #
-                (cp >= 0x3400 and cp <= 0x4DBF) or  #
-                (cp >= 0x20000 and cp <= 0x2A6DF) or  #
-                (cp >= 0x2A700 and cp <= 0x2B73F) or  #
-                (cp >= 0x2B740 and cp <= 0x2B81F) or  #
-                (cp >= 0x2B820 and cp <= 0x2CEAF) or
-                (cp >= 0xF900 and cp <= 0xFAFF) or  #
-                (cp >= 0x2F800 and cp <= 0x2FA1F)):  #
+            (cp >= 0x3400 and cp <= 0x4DBF) or  #
+            (cp >= 0x20000 and cp <= 0x2A6DF) or  #
+            (cp >= 0x2A700 and cp <= 0x2B73F) or  #
+            (cp >= 0x2B740 and cp <= 0x2B81F) or  #
+            (cp >= 0x2B820 and cp <= 0x2CEAF) or
+            (cp >= 0xF900 and cp <= 0xFAFF) or  #
+            (cp >= 0x2F800 and cp <= 0x2FA1F)):  #
             return True
 
         return False
 
-
     def _clean_text(self, text):
         """Performs invalid character removal and whitespace cleanup on text."""
         output = []
@@ -361,7 +362,3 @@ def _is_punctuation(char):
     if cat.startswith("P"):
         return True
     return False
-
-
-
-
diff --git a/flagai/model/base_model.py b/flagai/model/base_model.py
index bd26918b..3e28cb10 100755
--- a/flagai/model/base_model.py
+++ b/flagai/model/base_model.py
@@ -5,10 +5,11 @@
 from torch.nn import Module
 import torch
 import json
-from typing import Union 
+from typing import Union
 from flagai.model.file_utils import _get_model_id, _get_config_path, _get_checkpoint_path, _get_vocab_path, _get_model_files
 import os
 
+
 # The base model for models
 class BaseModel(Module):
 
@@ -39,6 +40,26 @@ def init_from_json(cls, config_file='./config.json', **kwargs):
             args['checkpoint_activations'] = False
         return cls(args, **kwargs)
 
+    @classmethod
+    def _load_state_dict_into_model(cls,
+                                    model,
+                                    pretrained_model_name_or_path,
+                                    verbose=False):
+        pl_sd = torch.load(pretrained_model_name_or_path, map_location="cpu")
+        sd = pl_sd["state_dict"]
+        if "global_step" in pl_sd:
+            print(f"Global Step: {pl_sd['global_step']}")
+        m, u = model.load_state_dict(sd, strict=False)
+        if len(m) > 0 and verbose:
+            print("missing keys:")
+            print(m)
+        if len(u) > 0 and verbose:
+            print("unexpected keys:")
+            print(u)
+        model.cuda()
+        model.eval()
+        return model
+
     @classmethod
     def from_pretrain(cls,
                       download_path='./checkpoints/',
@@ -48,8 +69,10 @@ def from_pretrain(cls,
                       **kwargs):
         model_id = None
 
+        raw_download_path = download_path
         # Try load model from local path
         download_path = os.path.join(download_path, model_name)
+
         config_path = os.path.join(download_path, "config.json")
         checkpoint_path = os.path.join(download_path, "pytorch_model.bin")
 
@@ -72,8 +95,9 @@ def load_local(checkpoint_path):
                     if model_parallel_size > 1 and not check_pytorch_model_mp_size(
                             download_path, model_parallel_size):
                         brief_model_name = MODEL_DICT[model_name.lower()][2]
-                        change_pytorch_model_mp_from_1_to_n_new(brief_model_name,
-                            download_path, model_parallel_size)
+                        change_pytorch_model_mp_from_1_to_n_new(
+                            brief_model_name, download_path,
+                            model_parallel_size)
 
                 from flagai import mpu
                 torch.distributed.barrier(group=mpu.get_model_parallel_group())
@@ -90,7 +114,33 @@ def load_local(checkpoint_path):
                     model.load_weights(checkpoint_path)
             return model
 
-        if os.path.exists(config_path):
+        def load_diffusion_local(yaml_path):
+            """
+            Now only diffusion models requires yaml
+            """
+            from omegaconf import OmegaConf
+            checkpoint_path = os.path.join(
+                download_path, "model.ckpt")  # Specific to diffusion models
+
+            config = OmegaConf.load(f"{yaml_path}")
+            model_config = config.model
+            model_config.params.cond_stage_config.params.download_path = raw_download_path
+
+            model = cls(**model_config.get("params", dict()))
+
+            model = cls._load_state_dict_into_model(
+                model,
+                checkpoint_path,
+            )
+            return model
+
+        yaml_path = os.path.join(download_path, "config.yaml")
+        if os.path.exists(yaml_path):
+            """
+            Now only diffusion models requires yaml
+            """
+            return load_diffusion_local(yaml_path)
+        elif os.path.exists(config_path):
             """
             It is fine when checkpoint_path does not exist, for the case that only_download_config=True
             At that time the model will not be loaded. 
@@ -111,7 +161,8 @@ def load_local(checkpoint_path):
                 if not file_name.endswith("bin"):
                     _get_vocab_path(download_path, file_name, model_id)
 
-            if not only_download_config and os.path.exists(os.path.join(download_path, 'config.json')):
+            if not only_download_config and os.path.exists(
+                    os.path.join(download_path, 'config.json')):
                 if os.getenv('ENV_TYPE') == 'deepspeed+mpu':
                     model_parallel_size = int(os.getenv("MODEL_PARALLEL_SIZE"))
                     if model_parallel_size > 1:
@@ -129,25 +180,48 @@ def load_local(checkpoint_path):
                     # Only to download the model slices(megatron-lm).
                     for file_to_load in model_files:
                         if "pytorch_model_" in file_to_load:
-                            _get_checkpoint_path(download_path,
-                                                 file_to_load,
+                            _get_checkpoint_path(download_path, file_to_load,
                                                  model_id)
 
                 elif 'pytorch_model.bin' in model_files:
-                    checkpoint_path = _get_checkpoint_path(download_path,
-                                                           'pytorch_model.bin',
-                                                           model_id)
-                else :
+                    checkpoint_path = _get_checkpoint_path(
+                        download_path, 'pytorch_model.bin', model_id)
+                else:
                     checkpoint_merge = {}
                     # maybe multi weights files
                     for file_to_load in model_files:
                         if "pytorch_model-0" in file_to_load:
-                            _get_checkpoint_path(download_path,
-                                                 file_to_load,
+                            _get_checkpoint_path(download_path, file_to_load,
                                                  model_id)
-                            checkpoint_to_load = torch.load(os.path.join(download_path, file_to_load), map_location="cpu")
+                            checkpoint_to_load = torch.load(os.path.join(
+                                download_path, file_to_load),
+                                                            map_location="cpu")
                             for k, v in checkpoint_to_load.items():
                                 checkpoint_merge[k] = v
                     # save all parameters
-                    torch.save(checkpoint_merge, os.path.join(download_path, "pytorch_model.bin"))
+                    torch.save(
+                        checkpoint_merge,
+                        os.path.join(download_path, "pytorch_model.bin"))
+        if os.path.exists(yaml_path):
+            return load_diffusion_local(yaml_path)
         return load_local(checkpoint_path)
+
+    @classmethod
+    def download(cls,
+                download_path='./checkpoints/',
+                model_name='RoBERTa-base-ch',
+                **kwargs):
+        try:
+            model_id = _get_model_id(model_name)
+        except:
+            print("Model hub is not reachable!")
+        # prepare the download path
+        # downloading the files
+        if model_id and model_id != "null":
+            model_files = eval(_get_model_files(model_name))
+            print("model files:" + str(model_files))
+            for file_name in model_files:
+                if not file_name.endswith("bin"):
+                    _get_vocab_path(os.path.join(download_path, model_name), file_name, model_id)
+                else :
+                    _get_checkpoint_path(os.path.join(download_path, model_name), file_name, model_id)
\ No newline at end of file
diff --git a/flagai/model/cpm3_model.py b/flagai/model/cpm3_model.py
old mode 100644
new mode 100755
index b53cd9b8..025f662e
--- a/flagai/model/cpm3_model.py
+++ b/flagai/model/cpm3_model.py
@@ -1,4 +1,3 @@
-from selectors import EpollSelector
 import torch
 from flagai.model.blocks.cpm_block import CPM3Block
 from flagai.model.layers.layer_norm import CPM3LayerNorm
@@ -63,7 +62,6 @@ def forward(self, hidden_states : torch.Tensor,
                       attention_mask : torch.Tensor,
                       position_bias : torch.Tensor = None,
                       past_key_values = None,
-                      isinfer = True,
                       ):
         """
         Args:
@@ -76,66 +74,67 @@ def forward(self, hidden_states : torch.Tensor,
 
         """
         # (batch, seq_enc, dim_model)
+        
         present_key_values = []
-        if isinfer:
-            for i, module in enumerate(self.layers):
-                hidden_states, present_key_value = module(hidden_states, attention_mask, position_bias, None, None, None, past_key_values[i])
-                present_key_values.append(present_key_value)
-            # (batch, seq_enc, dim_model)
-            hidden_states = self.output_layernorm(hidden_states)
-        else:
-            hidden_states = self.layers(hidden_states, attention_mask, position_bias, None, None, None)
-            # (batch, seq_enc, dim_model)
-            hidden_states = self.output_layernorm(hidden_states)
+        for i, module in enumerate(self.layers):
+            hidden_states, present_key_value = module(hidden_states, attention_mask, position_bias, None, None, None, past_key_values[i])
+            present_key_values.append(present_key_value)
+        # (batch, seq_enc, dim_model)
+        hidden_states = self.output_layernorm(hidden_states)
         return hidden_states, present_key_values
 
 
 class CPM3(BaseModel):
     def __init__(self, config, **kwargs):
         super().__init__(config, **kwargs)
+        self.config = config
         if type(config) is dict:
             config_cpm3 = CPM3Config(
-                vocab_size = config.vocab_size, 
-                num_layers = config.num_layers,
-                dim_model = config.dim_model, 
-                dim_ff = config.dim_ff,
-                num_heads = config.num_heads,
-                dim_head = config.dim_head,
-                dtype = config.dtype, 
-                int8 = config.int8,
-                norm_eps = config.norm_eps, 
-                norm_init_var = config.norm_init_var,
-                norm_bias = config.norm_bias,
-                att_init_mean = config.att_init_mean, 
-                att_init_std = config.att_init_std,
-                att_bias = config.att_bias,
-                att_mask_value = float(config.att_mask_value),
-                pos_bias_type = config.pos_bias_type,
-                ffn_init_mean = config.ffn_init_mean, 
-                ffn_init_std = config.ffn_init_std,
-                ffn_bias = config.ffn_bias,
-                ffn_activate_fn = config.ffn_activate_fn,
-                length_scale = config.length_scale,
-                attn_scale = config.attn_scale,
-                dropout_p = config.dropout_p,
-                prompt_types = config.prompt_types,
-                prompt_length = config.prompt_length, 
-                segment_types = config.segment_types,
-                position_bias_num_buckets = config.position_bias_num_buckets,
-                position_bias_max_distance = config.position_bias_max_distance,
-                max_exact_rate = config.max_exact_rate,
-                max_distance_rate = config.max_distance_rate,
-                absolute_inner_segment = config.absolute_inner_segment,
-                tied = config.tied,
-                cls_head = config.cls_head,
-                init_mean = config.proj_init_mean,
-                init_std = config.proj_init_std,
-                bias = config.proj_bias,
+                vocab_size = self.config['vocab_size'],
+                        dim_model = self.config['dim_model'],
+                        num_heads = self.config['num_heads'],
+                        dim_head = self.config['dim_head'],
+                        dim_ff = self.config['dim_ff'],
+                        num_layers = self.config['num_layers'],
+                        dropout_p = self.config['dropout_p'],
+                        emb_init_mean = self.config['emb_init_mean'],
+                        emb_init_std = self.config['emb_init_std'],
+                        pos_bias_type = self.config['pos_bias_type'],
+                        position_bias_num_buckets = self.config['position_bias_num_buckets'],
+                        position_bias_max_distance = self.config['position_bias_max_distance'],
+                        pos_init_mean = self.config['pos_init_mean'],
+                        pos_init_std = self.config['pos_init_std'],
+                        norm_init_var = self.config['norm_init_var'],
+                        norm_bias = self.config['norm_bias'],
+                        norm_eps = self.config['norm_eps'],
+                        att_init_mean = self.config['att_init_mean'], 
+                        att_init_std = self.config['att_init_std'],
+                        att_bias = self.config['att_bias'],
+                        att_mask_value = float(self.config['att_mask_value']) if type(self.config['att_mask_value']) == str else self.config['att_mask_value'],
+                        ffn_init_mean = self.config['ffn_init_mean'], 
+                        ffn_init_std = self.config['ffn_init_std'],
+                        ffn_bias = self.config['ffn_bias'],
+                        ffn_activate_fn = self.config['ffn_activate_fn'],
+                        proj_init_mean = self.config['proj_init_mean'],
+                        proj_init_std = self.config['proj_init_std'],
+                        proj_bias = self.config['proj_bias'],
+                        length_scale = self.config['length_scale'],
+                        attn_scale = self.config['attn_scale'],
+                        half = self.config['half'], 
+                        int8 = self.config['int8'],
+                        tied = self.config['tied'],
+                        prompt_types = self.config['prompt_types'],
+                        prompt_length = self.config['prompt_length'], 
+                        segment_types = self.config['segment_types'],
+                        max_exact_rate = self.config['max_exact_rate'],
+                        max_distance_rate = self.config['max_distance_rate'],
+                        absolute_inner_segment = self.config['absolute_inner_segment'],
+                        cls_head = self.config.get('cls_head', None),
+                        post_layer_norm= self.config.get('post_layer_norm', None)
             )
         self.config = config_cpm3
         
         self.encoder = CPM3Stack(self.config)
-
         self.prompt_embedding = CPM3Embedding(
             vocab_size = self.config.prompt_types * self.config.prompt_length, 
             embedding_size = self.config.dim_model,
@@ -154,7 +153,7 @@ def __init__(self, config, **kwargs):
             num_segments = self.config.segment_types,
             num_heads = self.config.num_heads, 
             num_buckets = self.config.position_bias_num_buckets, 
-            max_distance = self.config.position_bias_max_distance,
+            max_distance = self.config.position_bias_max_distance, 
             max_exact_rate = self.config.max_exact_rate,
             max_distance_rate = self.config.max_distance_rate,
             absolute_inner_segment = self.config.absolute_inner_segment,
@@ -194,12 +193,10 @@ def forward(self, input : torch.Tensor, # (batch, seqlen)
                       past_key_values = None,  # num_layers * 2 * (batch, num_heads, seqlen, dim_head)
                       right_ctx_start_idx = None,
                       last_input_idx = None,
-                      cached_attn_mask_pos_bias = None,
-                      isinfer = True,
+                      cached_attn_mask_pos_bias = None
                     ):
 
         batch = input.size(0)
-        seqlen = context.size(1)
 
         if past_key_values is None:
             past_length = 0
@@ -224,7 +221,7 @@ def forward(self, input : torch.Tensor, # (batch, seqlen)
         if cached_attn_mask_pos_bias is None:
             with torch.no_grad():
                 device = input.device
-                # seqlen = context.size(1)
+                seqlen = context.size(1)
                 directional_mask_2d = torch.arange(seqlen, device=device) <= torch.arange(seqlen, device=device).view(-1, 1)
                 attention_mask = context[:, None, :] | (context[:, :, None].logical_not() & directional_mask_2d.view(1, seqlen, seqlen))
                 attention_mask = attention_mask & (span[:, None, :] == span[:, :, None])
@@ -254,8 +251,8 @@ def forward(self, input : torch.Tensor, # (batch, seqlen)
         else:
             attention_mask = attention_mask[:, past_length:valid_len, :valid_len]
             position_bias = position_bias[:, :, past_length:valid_len, :valid_len]
-        # example is infer or finetune & pretrain
-        hidden_states, present_key_values = self.encoder(hidden_states, attention_mask, position_bias, past_key_values, isinfer)
+
+        hidden_states, present_key_values = self.encoder(hidden_states, attention_mask, position_bias, past_key_values)
 
         if self.cls_head:
             logits = self.output_projection(hidden_states)
@@ -265,7 +262,6 @@ def forward(self, input : torch.Tensor, # (batch, seqlen)
             logits = self.input_embedding.projection(hidden_states)
 
         return logits, hidden_states, present_key_values, cached_attn_mask_pos_bias
-
     
     def load_weights(self, checkpoint_path):
         self.load_state_dict(
@@ -311,7 +307,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 class CPM3Config(Config):
 
-    def __init__(self, vocab_size = 30730,
+    def __init__(self, vocab_size = 30720,
                         dim_model = 4096,
                         num_heads = 64,
                         dim_head = 64,
@@ -401,4 +397,4 @@ def __init__(self, vocab_size = 30730,
         self.attn_scale = attn_scale
         self.post_layer_norm = post_layer_norm
         self.is_decoder = is_decoder
-        self.parallel_ffn  = parallel_ffn
+        self.parallel_ffn  = parallel_ffn
\ No newline at end of file
diff --git a/flagai/model/mm/Sampler.py b/flagai/model/mm/Sampler.py
new file mode 100755
index 00000000..0b75d756
--- /dev/null
+++ b/flagai/model/mm/Sampler.py
@@ -0,0 +1,467 @@
+
+import torch
+import numpy as np
+from tqdm import tqdm
+from flagai.model.mm.utils import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like,extract_into_tensor
+
+class DDIMSampler(object):
+    def __init__(self, model, schedule="linear", **kwargs):
+        super().__init__()
+        self.model = model
+        self.ddpm_num_timesteps = model.num_timesteps
+        self.schedule = schedule
+
+    def register_buffer(self, name, attr):
+        if type(attr) == torch.Tensor:
+            if attr.device != torch.device("cuda"):
+                attr = attr.to(torch.device("cuda"))
+        setattr(self, name, attr)
+
+    def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
+        self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
+                                                  num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
+        alphas_cumprod = self.model.alphas_cumprod
+        assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
+        to_torch = lambda x: x.clone().detach().to(torch.float32).to('cuda')
+
+        self.register_buffer('betas', to_torch(self.model.betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
+        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
+
+        # ddim sampling parameters
+        ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
+                                                                                   ddim_timesteps=self.ddim_timesteps,
+                                                                                   eta=ddim_eta,verbose=verbose)
+        self.register_buffer('ddim_sigmas', ddim_sigmas)
+        self.register_buffer('ddim_alphas', ddim_alphas)
+        self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
+        self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
+        sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
+            (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
+                        1 - self.alphas_cumprod / self.alphas_cumprod_prev))
+        self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
+
+    @torch.no_grad()
+    def sample(self,
+               S,
+               batch_size,
+               shape,
+               conditioning=None,
+               callback=None,
+               normals_sequence=None,
+               img_callback=None,
+               quantize_x0=False,
+               eta=0.,
+               mask=None,
+               x0=None,
+               temperature=1.,
+               noise_dropout=0.,
+               score_corrector=None,
+               corrector_kwargs=None,
+               verbose=True,
+               x_T=None,
+               log_every_t=100,
+               unconditional_guidance_scale=1.,
+               unconditional_conditioning=None,
+               # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
+               **kwargs
+               ):
+        if conditioning is not None:
+            if isinstance(conditioning, dict):
+                cbs = conditioning[list(conditioning.keys())[0]].shape[0]
+                if cbs != batch_size:
+                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+            else:
+                if conditioning.shape[0] != batch_size:
+                    print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
+        self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
+        # sampling
+        C, H, W = shape
+        size = (batch_size, C, H, W)
+        print(f'Data shape for DDIM sampling is {size}, eta {eta}')
+        samples, intermediates = self.ddim_sampling(conditioning, size,
+                                                    callback=callback,
+                                                    img_callback=img_callback,
+                                                    quantize_denoised=quantize_x0,
+                                                    mask=mask, x0=x0,
+                                                    ddim_use_original_steps=False,
+                                                    noise_dropout=noise_dropout,
+                                                    temperature=temperature,
+                                                    score_corrector=score_corrector,
+                                                    corrector_kwargs=corrector_kwargs,
+                                                    x_T=x_T,
+                                                    log_every_t=log_every_t,
+                                                    unconditional_guidance_scale=unconditional_guidance_scale,
+                                                    unconditional_conditioning=unconditional_conditioning,
+                                                    )
+        return samples, intermediates
+
+    @torch.no_grad()
+    def ddim_sampling(self, cond, shape,
+                      x_T=None, ddim_use_original_steps=False,
+                      callback=None, timesteps=None, quantize_denoised=False,
+                      mask=None, x0=None, img_callback=None, log_every_t=100,
+                      temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+                      unconditional_guidance_scale=1., unconditional_conditioning=None,):
+        device = self.model.betas.device
+        b = shape[0]
+        if x_T is None:
+            img = torch.randn(shape, device=device)
+        else:
+            img = x_T
+
+        if timesteps is None:
+            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+        elif timesteps is not None and not ddim_use_original_steps:
+            subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
+            timesteps = self.ddim_timesteps[:subset_end]
+
+        intermediates = {'x_inter': [img], 'pred_x0': [img]}
+        time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
+        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+        print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
+
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((b,), step, device=device, dtype=torch.long)
+            if mask is not None:
+                assert x0 is not None
+                img_orig = self.model.q_sample(x0, ts)  # TODO: deterministic forward pass?
+                img = img_orig * mask + (1. - mask) * img
+
+            outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
+                                      quantize_denoised=quantize_denoised, temperature=temperature,
+                                      noise_dropout=noise_dropout, score_corrector=score_corrector,
+                                      corrector_kwargs=corrector_kwargs,
+                                      unconditional_guidance_scale=unconditional_guidance_scale,
+                                      unconditional_conditioning=unconditional_conditioning)
+            img, pred_x0 = outs
+            if callback: callback(i)
+            if img_callback: img_callback(pred_x0, i)
+
+            if index % log_every_t == 0 or index == total_steps - 1:
+                intermediates['x_inter'].append(img)
+                intermediates['pred_x0'].append(pred_x0)
+
+        return img, intermediates
+
+    @torch.no_grad()
+    def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
+                      temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+                      unconditional_guidance_scale=1., unconditional_conditioning=None):
+        b, *_, device = *x.shape, x.device
+
+        if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+            e_t = self.model.apply_model(x, t, c)
+        else:
+            # x_in = torch.cat([x] * 2)
+            # t_in = torch.cat([t] * 2)
+
+            # c_in = torch.cat([unconditional_conditioning, c])
+            # e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
+            x_in = torch.cat([x])
+            t_in = torch.cat([t]) 
+            e_t_uncond = self.model.apply_model(x_in, t_in, unconditional_conditioning)
+            e_t = self.model.apply_model(x_in, t_in, c)
+            e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
+
+        if score_corrector is not None:
+            assert self.model.parameterization == "eps"
+            e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
+
+        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+        alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+        sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+        # select parameters corresponding to the currently considered timestep
+        a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
+        a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
+        sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
+        sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
+
+        # current prediction for x_0
+        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+        if quantize_denoised:
+            pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+        # direction pointing to x_t
+        dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+        noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+        if noise_dropout > 0.:
+            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+        x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+        return x_prev, pred_x0
+
+    @torch.no_grad()
+    def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
+        # fast, but does not allow for exact reconstruction
+        # t serves as an index to gather the correct alphas
+        if use_original_steps:
+            sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
+            sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
+        else:
+            sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
+            sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
+
+        if noise is None:
+            noise = torch.randn_like(x0)
+        return (extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 +
+                extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise)
+
+
+    @torch.no_grad()
+    def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
+               use_original_steps=False):
+
+        timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
+        timesteps = timesteps[:t_start]
+
+        time_range = np.flip(timesteps)
+        total_steps = timesteps.shape[0]
+        print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+        iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
+        x_dec = x_latent
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
+            x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
+                                          unconditional_guidance_scale=unconditional_guidance_scale,
+                                          unconditional_conditioning=unconditional_conditioning)
+        return x_dec
+
+
+class PLMSSampler(object):
+    def __init__(self, model, schedule="linear", **kwargs):
+        super().__init__()
+        self.model = model
+        self.ddpm_num_timesteps = model.num_timesteps
+        self.schedule = schedule
+
+    def register_buffer(self, name, attr):
+        if type(attr) == torch.Tensor:
+            if attr.device != torch.device("cuda"):
+                attr = attr.to(torch.device("cuda"))
+        setattr(self, name, attr)
+
+    def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
+        if ddim_eta != 0:
+            raise ValueError('ddim_eta must be 0 for PLMS')
+        self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
+                                                  num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
+        alphas_cumprod = self.model.alphas_cumprod
+        assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
+        to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
+
+        self.register_buffer('betas', to_torch(self.model.betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
+        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
+
+        # ddim sampling parameters
+        ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
+                                                                                   ddim_timesteps=self.ddim_timesteps,
+                                                                                   eta=ddim_eta,verbose=verbose)
+        self.register_buffer('ddim_sigmas', ddim_sigmas)
+        self.register_buffer('ddim_alphas', ddim_alphas)
+        self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
+        self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
+        sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
+            (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
+                        1 - self.alphas_cumprod / self.alphas_cumprod_prev))
+        self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
+
+    @torch.no_grad()
+    def sample(self,
+               S,
+               batch_size,
+               shape,
+               conditioning=None,
+               callback=None,
+               normals_sequence=None,
+               img_callback=None,
+               quantize_x0=False,
+               eta=0.,
+               mask=None,
+               x0=None,
+               temperature=1.,
+               noise_dropout=0.,
+               score_corrector=None,
+               corrector_kwargs=None,
+               verbose=True,
+               x_T=None,
+               log_every_t=100,
+               unconditional_guidance_scale=1.,
+               unconditional_conditioning=None,
+               # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
+               **kwargs
+               ):
+        if conditioning is not None:
+            if isinstance(conditioning, dict):
+                cbs = conditioning[list(conditioning.keys())[0]].shape[0]
+                if cbs != batch_size:
+                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+            else:
+                if conditioning.shape[0] != batch_size:
+                    print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
+
+        self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
+        # sampling
+        C, H, W = shape
+        size = (batch_size, C, H, W)
+        print(f'Data shape for PLMS sampling is {size}')
+        samples, intermediates = self.plms_sampling(conditioning, size,
+                                                    callback=callback,
+                                                    img_callback=img_callback,
+                                                    quantize_denoised=quantize_x0,
+                                                    mask=mask, x0=x0,
+                                                    ddim_use_original_steps=False,
+                                                    noise_dropout=noise_dropout,
+                                                    temperature=temperature,
+                                                    score_corrector=score_corrector,
+                                                    corrector_kwargs=corrector_kwargs,
+                                                    x_T=x_T,
+                                                    log_every_t=log_every_t,
+                                                    unconditional_guidance_scale=unconditional_guidance_scale,
+                                                    unconditional_conditioning=unconditional_conditioning,
+                                                    )
+        return samples, intermediates
+
+    @torch.no_grad()
+    def plms_sampling(self, cond, shape,
+                      x_T=None, ddim_use_original_steps=False,
+                      callback=None, timesteps=None, quantize_denoised=False,
+                      mask=None, x0=None, img_callback=None, log_every_t=100,
+                      temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+                      unconditional_guidance_scale=1., unconditional_conditioning=None,):
+        device = self.model.betas.device
+        b = shape[0]
+        if x_T is None:
+            img = torch.randn(shape, device=device)
+        else:
+            img = x_T
+
+        if timesteps is None:
+            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+        elif timesteps is not None and not ddim_use_original_steps:
+            subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
+            timesteps = self.ddim_timesteps[:subset_end]
+
+        intermediates = {'x_inter': [img], 'pred_x0': [img]}
+        time_range = list(reversed(range(0,timesteps))) if ddim_use_original_steps else np.flip(timesteps)
+        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+        print(f"Running PLMS Sampling with {total_steps} timesteps")
+
+        iterator = tqdm(time_range, desc='PLMS Sampler', total=total_steps)
+        old_eps = []
+
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((b,), step, device=device, dtype=torch.long)
+            ts_next = torch.full((b,), time_range[min(i + 1, len(time_range) - 1)], device=device, dtype=torch.long)
+
+            if mask is not None:
+                assert x0 is not None
+                img_orig = self.model.q_sample(x0, ts)  # TODO: deterministic forward pass?
+                img = img_orig * mask + (1. - mask) * img
+            outs = self.p_sample_plms(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
+                                      quantize_denoised=quantize_denoised, temperature=temperature,
+                                      noise_dropout=noise_dropout, score_corrector=score_corrector,
+                                      corrector_kwargs=corrector_kwargs,
+                                      unconditional_guidance_scale=unconditional_guidance_scale,
+                                      unconditional_conditioning=unconditional_conditioning,
+                                      old_eps=old_eps, t_next=ts_next)
+            img, pred_x0, e_t = outs
+            old_eps.append(e_t)
+            if len(old_eps) >= 4:
+                old_eps.pop(0)
+            if callback: callback(i)
+            if img_callback: img_callback(pred_x0, i)
+
+            if index % log_every_t == 0 or index == total_steps - 1:
+                intermediates['x_inter'].append(img)
+                intermediates['pred_x0'].append(pred_x0)
+
+        return img, intermediates
+
+    @torch.no_grad()
+    def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
+                      temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+                      unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None):
+        b, *_, device = *x.shape, x.device
+
+        def get_model_output(x, t):
+            if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+                e_t = self.model.apply_model(x, t, c)
+            else:
+                x_in = torch.cat([x] * 2)
+                t_in = torch.cat([t] * 2)
+                #c_in = torch.cat([unconditional_conditioning, c])
+                e_t_uncond = self.model.apply_model(x_in, t_in, unconditional_conditioning)
+                e_t = self.model.apply_model(x_in, t_in, c) 
+                #e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
+                e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
+
+            if score_corrector is not None:
+                assert self.model.parameterization == "eps"
+                e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
+
+            return e_t
+
+        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+        alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+        sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+
+        def get_x_prev_and_pred_x0(e_t, index):
+            # select parameters corresponding to the currently considered timestep
+            a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
+            a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
+            sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
+            sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
+
+            # current prediction for x_0
+            pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+            if quantize_denoised:
+                pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+            # direction pointing to x_t
+            dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+            noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+            if noise_dropout > 0.:
+                noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+            x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+            return x_prev, pred_x0
+
+        e_t = get_model_output(x, t)
+        if len(old_eps) == 0:
+            # Pseudo Improved Euler (2nd order)
+            x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
+            e_t_next = get_model_output(x_prev, t_next)
+            e_t_prime = (e_t + e_t_next) / 2
+        elif len(old_eps) == 1:
+            # 2nd order Pseudo Linear Multistep (Adams-Bashforth)
+            e_t_prime = (3 * e_t - old_eps[-1]) / 2
+        elif len(old_eps) == 2:
+            # 3nd order Pseudo Linear Multistep (Adams-Bashforth)
+            e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
+        elif len(old_eps) >= 3:
+            # 4nd order Pseudo Linear Multistep (Adams-Bashforth)
+            e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
+
+        x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
+
+        return x_prev, pred_x0, e_t
diff --git a/flagai/model/mm/Unets/Unet.py b/flagai/model/mm/Unets/Unet.py
new file mode 100755
index 00000000..2175fbdf
--- /dev/null
+++ b/flagai/model/mm/Unets/Unet.py
@@ -0,0 +1,959 @@
+from abc import abstractmethod
+import math
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+from abc import abstractmethod
+import math
+from flagai.model.mm.utils import (
+    checkpoint,
+    conv_nd,
+    linear,
+    avg_pool_nd,
+    zero_module,
+    normalization,
+    timestep_embedding,
+)
+from flagai.model.mm.attentions.attention import SpatialTransformer
+
+
+# dummy replace
+def convert_module_to_f16(x):
+    pass
+
+def convert_module_to_f32(x):
+    pass
+
+
+## go
+class AttentionPool2d(nn.Module):
+    """
+    Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
+    """
+
+    def __init__(
+        self,
+        spacial_dim: int,
+        embed_dim: int,
+        num_heads_channels: int,
+        output_dim: int = None,
+    ):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5)
+        self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
+        self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
+        self.num_heads = embed_dim // num_heads_channels
+        self.attention = QKVAttention(self.num_heads)
+
+    def forward(self, x):
+        b, c, *_spatial = x.shape
+        x = x.reshape(b, c, -1)  # NC(HW)
+        x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1)  # NC(HW+1)
+        x = x + self.positional_embedding[None, :, :].to(x.dtype)  # NC(HW+1)
+        x = self.qkv_proj(x)
+        x = self.attention(x)
+        x = self.c_proj(x)
+        return x[:, :, 0]
+
+
+class TimestepBlock(nn.Module):
+    """
+    Any module where forward() takes timestep embeddings as a second argument.
+    """
+
+    @abstractmethod
+    def forward(self, x, emb):
+        """
+        Apply the module to `x` given `emb` timestep embeddings.
+        """
+
+
+class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
+    """
+    A sequential module that passes timestep embeddings to the children that
+    support it as an extra input.
+    """
+
+    def forward(self, x, emb, context=None):
+        for layer in self:
+            if isinstance(layer, TimestepBlock):
+                x = layer(x, emb)
+            elif isinstance(layer, SpatialTransformer):
+                x = layer(x, context)
+            else:
+                x = layer(x)
+        return x
+
+
+class Upsample(nn.Module):
+    """
+    An upsampling layer with an optional convolution.
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 upsampling occurs in the inner-two dimensions.
+    """
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        if use_conv:
+            self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        if self.dims == 3:
+            x = F.interpolate(
+                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
+            )
+        else:
+            x = F.interpolate(x, scale_factor=2, mode="nearest")
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+class TransposedUpsample(nn.Module):
+    'Learned 2x upsampling without padding'
+    def __init__(self, channels, out_channels=None, ks=5):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+
+        self.up = nn.ConvTranspose2d(self.channels,self.out_channels,kernel_size=ks,stride=2)
+
+    def forward(self,x):
+        return self.up(x)
+
+
+class Downsample(nn.Module):
+    """
+    A downsampling layer with an optional convolution.
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 downsampling occurs in the inner-two dimensions.
+    """
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None,padding=1):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        stride = 2 if dims != 3 else (1, 2, 2)
+        if use_conv:
+            self.op = conv_nd(
+                dims, self.channels, self.out_channels, 3, stride=stride, padding=padding
+            )
+        else:
+            assert self.channels == self.out_channels
+            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        return self.op(x)
+
+
+class ResBlock(TimestepBlock):
+    """
+    A residual block that can optionally change the number of channels.
+    :param channels: the number of input channels.
+    :param emb_channels: the number of timestep embedding channels.
+    :param dropout: the rate of dropout.
+    :param out_channels: if specified, the number of out channels.
+    :param use_conv: if True and out_channels is specified, use a spatial
+        convolution instead of a smaller 1x1 convolution to change the
+        channels in the skip connection.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param use_checkpoint: if True, use gradient checkpointing on this module.
+    :param up: if True, use this block for upsampling.
+    :param down: if True, use this block for downsampling.
+    """
+
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        dims=2,
+        use_checkpoint=False,
+        up=False,
+        down=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_checkpoint = use_checkpoint
+        self.use_scale_shift_norm = use_scale_shift_norm
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
+        )
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(
+                conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = conv_nd(
+                dims, channels, self.out_channels, 3, padding=1
+            )
+        else:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+    def forward(self, x, emb):
+        """
+        Apply the block to a Tensor, conditioned on a timestep embedding.
+        :param x: an [N x C x ...] Tensor of features.
+        :param emb: an [N x emb_channels] Tensor of timestep embeddings.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        return checkpoint(
+            self._forward, (x, emb), self.parameters(), self.use_checkpoint
+        )
+
+
+    def _forward(self, x, emb):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        emb_out = self.emb_layers(emb).type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        if self.use_scale_shift_norm:
+            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+            scale, shift = th.chunk(emb_out, 2, dim=1)
+            h = out_norm(h) * (1 + scale) + shift
+            h = out_rest(h)
+        else:
+            h = h + emb_out
+            h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    An attention block that allows spatial positions to attend to each other.
+    Originally ported from here, but adapted to the N-d case.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
+    """
+
+    def __init__(
+        self,
+        channels,
+        num_heads=1,
+        num_head_channels=-1,
+        use_checkpoint=False,
+        use_new_attention_order=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        if num_head_channels == -1:
+            self.num_heads = num_heads
+        else:
+            assert (
+                channels % num_head_channels == 0
+            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
+            self.num_heads = channels // num_head_channels
+        self.use_checkpoint = use_checkpoint
+        self.norm = normalization(channels)
+        self.qkv = conv_nd(1, channels, channels * 3, 1)
+        if use_new_attention_order:
+            # split qkv before split heads
+            self.attention = QKVAttention(self.num_heads)
+        else:
+            # split heads before split qkv
+            self.attention = QKVAttentionLegacy(self.num_heads)
+
+        self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
+
+    def forward(self, x):
+        return checkpoint(self._forward, (x,), self.parameters(), True)   # TODO: check checkpoint usage, is True # TODO: fix the .half call!!!
+        #return pt_checkpoint(self._forward, x)  # pytorch
+
+    def _forward(self, x):
+        b, c, *spatial = x.shape
+        x = x.reshape(b, c, -1)
+        qkv = self.qkv(self.norm(x))
+        h = self.attention(qkv)
+        h = self.proj_out(h)
+        return (x + h).reshape(b, c, *spatial)
+
+
+def count_flops_attn(model, _x, y):
+    """
+    A counter for the `thop` package to count the operations in an
+    attention operation.
+    Meant to be used like:
+        macs, params = thop.profile(
+            model,
+            inputs=(inputs, timestamps),
+            custom_ops={QKVAttention: QKVAttention.count_flops},
+        )
+    """
+    b, c, *spatial = y[0].shape
+    num_spatial = int(np.prod(spatial))
+    # We perform two matmuls with the same number of ops.
+    # The first computes the weight matrix, the second computes
+    # the combination of the value vectors.
+    matmul_ops = 2 * b * (num_spatial ** 2) * c
+    model.total_ops += th.DoubleTensor([matmul_ops])
+
+
+class QKVAttentionLegacy(nn.Module):
+    """
+    A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
+    """
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        """
+        Apply QKV attention.
+        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = th.einsum(
+            "bct,bcs->bts", q * scale, k * scale
+        )  # More stable with f16 than dividing afterwards
+        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = th.einsum("bts,bcs->bct", weight, v)
+        return a.reshape(bs, -1, length)
+
+    @staticmethod
+    def count_flops(model, _x, y):
+        return count_flops_attn(model, _x, y)
+
+
+class QKVAttention(nn.Module):
+    """
+    A module which performs QKV attention and splits in a different order.
+    """
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        """
+        Apply QKV attention.
+        :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.chunk(3, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = th.einsum(
+            "bct,bcs->bts",
+            (q * scale).view(bs * self.n_heads, ch, length),
+            (k * scale).view(bs * self.n_heads, ch, length),
+        )  # More stable with f16 than dividing afterwards
+        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
+        return a.reshape(bs, -1, length)
+
+    @staticmethod
+    def count_flops(model, _x, y):
+        return count_flops_attn(model, _x, y)
+
+
+class UNetModel(nn.Module):
+    """
+    The full UNet model with attention and timestep embedding.
+    :param in_channels: channels in the input Tensor.
+    :param model_channels: base channel count for the model.
+    :param out_channels: channels in the output Tensor.
+    :param num_res_blocks: number of residual blocks per downsample.
+    :param attention_resolutions: a collection of downsample rates at which
+        attention will take place. May be a set, list, or tuple.
+        For example, if this contains 4, then at 4x downsampling, attention
+        will be used.
+    :param dropout: the dropout probability.
+    :param channel_mult: channel multiplier for each level of the UNet.
+    :param conv_resample: if True, use learned convolutions for upsampling and
+        downsampling.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param num_classes: if specified (as an int), then this model will be
+        class-conditional with `num_classes` classes.
+    :param use_checkpoint: use gradient checkpointing to reduce memory usage.
+    :param num_heads: the number of attention heads in each attention layer.
+    :param num_heads_channels: if specified, ignore num_heads and instead use
+                               a fixed channel width per attention head.
+    :param num_heads_upsample: works with num_heads to set a different number
+                               of heads for upsampling. Deprecated.
+    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
+    :param resblock_updown: use residual blocks for up/downsampling.
+    :param use_new_attention_order: use a different attention pattern for potentially
+                                    increased efficiency.
+    """
+
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        num_classes=None,
+        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,
+    ):
+        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.image_size = image_size
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.num_classes = num_classes
+        self.use_checkpoint = use_checkpoint
+        self.dtype = th.float16 if use_fp16 else th.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(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        if self.num_classes is not None:
+            self.label_emb = nn.Embedding(num_classes, time_embed_dim)
+
+        self.input_blocks = nn.ModuleList(
+            [
+                TimestepEmbedSequential(
+                    conv_nd(dims, in_channels, 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 _ in range(num_res_blocks):
+                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
+                    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
+                        )
+                    )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                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)
+                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(
+                            ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim
+                        ),
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+
+        self.output_blocks = nn.ModuleList([])
+        for level, mult in list(enumerate(channel_mult))[::-1]:
+            for i in range(num_res_blocks + 1):
+                ich = input_block_chans.pop()
+                layers = [
+                    ResBlock(
+                        ch + ich,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=model_channels * mult,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = model_channels * mult
+                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
+                    layers.append(
+                        AttentionBlock(
+                            ch,
+                            use_checkpoint=use_checkpoint,
+                            num_heads=num_heads_upsample,
+                            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
+                        )
+                    )
+                if level and i == num_res_blocks:
+                    out_ch = ch
+                    layers.append(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            up=True,
+                        )
+                        if resblock_updown
+                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                    )
+                    ds //= 2
+                self.output_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+
+        self.out = nn.Sequential(
+            normalization(ch),
+            nn.SiLU(),
+            zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
+        )
+        if self.predict_codebook_ids:
+            self.id_predictor = nn.Sequential(
+            normalization(ch),
+            conv_nd(dims, model_channels, n_embed, 1),
+            #nn.LogSoftmax(dim=1)  # change to cross_entropy and produce non-normalized logits
+        )
+
+    def convert_to_fp16(self):
+        """
+        Convert the torso of the model to float16.
+        """
+        self.input_blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+        self.output_blocks.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self):
+        """
+        Convert the torso of the model to float32.
+        """
+        self.input_blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+        self.output_blocks.apply(convert_module_to_f32)
+
+    def forward(self, x, timesteps=None, context=None, y=None,**kwargs):
+        """
+        Apply the model to an input batch.
+        :param x: an [N x C x ...] Tensor of inputs.
+        :param timesteps: a 1-D batch of timesteps.
+        :param context: conditioning plugged in via crossattn
+        :param y: an [N] Tensor of labels, if class-conditional.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        assert (y is not None) == (
+            self.num_classes is not None
+        ), "must specify y if and only if the model is class-conditional"
+        hs = []
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+        emb = self.time_embed(t_emb)
+
+        if self.num_classes is not None:
+            assert y.shape == (x.shape[0],)
+            emb = emb + self.label_emb(y)
+
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            h = module(h, emb, context)
+            hs.append(h)
+        h = self.middle_block(h, emb, context)
+        for module in self.output_blocks:
+            h = th.cat([h, hs.pop()], dim=1)
+            h = module(h, emb, context)
+        h = h.type(x.dtype)
+        if self.predict_codebook_ids:
+            return self.id_predictor(h)
+        else:
+            return self.out(h)
+
+
+class EncoderUNetModel(nn.Module):
+    """
+    The half UNet model with attention and timestep embedding.
+    For usage, see UNet.
+    """
+
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_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,
+        pool="adaptive",
+        *args,
+        **kwargs
+    ):
+        super().__init__()
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        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 = th.float16 if use_fp16 else th.float32
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        self.input_blocks = nn.ModuleList(
+            [
+                TimestepEmbedSequential(
+                    conv_nd(dims, in_channels, 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 _ in range(num_res_blocks):
+                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:
+                    layers.append(
+                        AttentionBlock(
+                            ch,
+                            use_checkpoint=use_checkpoint,
+                            num_heads=num_heads,
+                            num_head_channels=num_head_channels,
+                            use_new_attention_order=use_new_attention_order,
+                        )
+                    )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                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)
+                ds *= 2
+                self._feature_size += ch
+
+        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=num_head_channels,
+                use_new_attention_order=use_new_attention_order,
+            ),
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+        self.pool = pool
+        if pool == "adaptive":
+            self.out = nn.Sequential(
+                normalization(ch),
+                nn.SiLU(),
+                nn.AdaptiveAvgPool2d((1, 1)),
+                zero_module(conv_nd(dims, ch, out_channels, 1)),
+                nn.Flatten(),
+            )
+        elif pool == "attention":
+            assert num_head_channels != -1
+            self.out = nn.Sequential(
+                normalization(ch),
+                nn.SiLU(),
+                AttentionPool2d(
+                    (image_size // ds), ch, num_head_channels, out_channels
+                ),
+            )
+        elif pool == "spatial":
+            self.out = nn.Sequential(
+                nn.Linear(self._feature_size, 2048),
+                nn.ReLU(),
+                nn.Linear(2048, self.out_channels),
+            )
+        elif pool == "spatial_v2":
+            self.out = nn.Sequential(
+                nn.Linear(self._feature_size, 2048),
+                normalization(2048),
+                nn.SiLU(),
+                nn.Linear(2048, self.out_channels),
+            )
+        else:
+            raise NotImplementedError(f"Unexpected {pool} pooling")
+
+    def convert_to_fp16(self):
+        """
+        Convert the torso of the model to float16.
+        """
+        self.input_blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self):
+        """
+        Convert the torso of the model to float32.
+        """
+        self.input_blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+
+    def forward(self, x, timesteps):
+        """
+        Apply the model to an input batch.
+        :param x: an [N x C x ...] Tensor of inputs.
+        :param timesteps: a 1-D batch of timesteps.
+        :return: an [N x K] Tensor of outputs.
+        """
+        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
+
+        results = []
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            h = module(h, emb)
+            if self.pool.startswith("spatial"):
+                results.append(h.type(x.dtype).mean(dim=(2, 3)))
+        h = self.middle_block(h, emb)
+        if self.pool.startswith("spatial"):
+            results.append(h.type(x.dtype).mean(dim=(2, 3)))
+            h = th.cat(results, axis=-1)
+            return self.out(h)
+        else:
+            h = h.type(x.dtype)
+            return self.out(h)
+
diff --git a/flagai/model/mm/Unets/__init__.py b/flagai/model/mm/Unets/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/flagai/model/mm/__init__.py b/flagai/model/mm/__init__.py
old mode 100644
new mode 100755
diff --git a/flagai/model/mm/attentions/__init__.py b/flagai/model/mm/attentions/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/flagai/model/mm/attentions/attention.py b/flagai/model/mm/attentions/attention.py
new file mode 100755
index 00000000..b149f1f2
--- /dev/null
+++ b/flagai/model/mm/attentions/attention.py
@@ -0,0 +1,259 @@
+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 flagai.model.mm.utils import checkpoint
+
+def exists(val):
+    return val is not None
+
+
+def uniq(arr):
+    return{el: True for el in arr}.keys()
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+
+def init_(tensor):
+    dim = tensor.shape[-1]
+    std = 1 / math.sqrt(dim)
+    tensor.uniform_(-std, std)
+    return tensor
+
+
+# feedforward
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        project_in = nn.Sequential(
+            nn.Linear(dim, inner_dim),
+            nn.GELU()
+        ) if not glu else GEGLU(dim, inner_dim)
+
+        self.net = nn.Sequential(
+            project_in,
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class LinearAttention(nn.Module):
+    def __init__(self, dim, heads=4, dim_head=32):
+        super().__init__()
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
+        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x)
+        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
+        k = k.softmax(dim=-1)  
+        context = torch.einsum('bhdn,bhen->bhde', k, v)
+        out = torch.einsum('bhde,bhdn->bhen', context, q)
+        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
+        return self.to_out(out)
+
+
+class SpatialSelfAttention(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.k = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.v = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=1,
+                                        stride=1,
+                                        padding=0)
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b,c,h,w = q.shape
+        q = rearrange(q, 'b c h w -> b (h w) c')
+        k = rearrange(k, 'b c h w -> b c (h w)')
+        w_ = torch.einsum('bij,bjk->bik', q, k)
+
+        w_ = w_ * (int(c)**(-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = rearrange(v, 'b c h w -> b c (h w)')
+        w_ = rearrange(w_, 'b i j -> b j i')
+        h_ = torch.einsum('bij,bjk->bik', v, w_)
+        h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
+        h_ = self.proj_out(h_)
+
+        return x+h_
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+
+        context_dim = default(context_dim, query_dim)
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x, context=None, mask=None):
+        h = self.heads
+
+        q = self.to_q(x)
+        context = default(context, x)
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+
+        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+
+        if exists(mask):
+            mask = rearrange(mask, 'b ... -> b (...)')
+            max_neg_value = -torch.finfo(sim.dtype).max
+            mask = repeat(mask, 'b j -> (b h) () j', h=h)
+            sim.masked_fill_(~mask, max_neg_value)
+
+        # attention, what we cannot get enough of
+        attn = sim.softmax(dim=-1)
+
+        out = einsum('b i j, b j d -> b i d', attn, v)
+        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+        return self.to_out(out)
+
+
+class BasicTransformerBlock(nn.Module):
+    def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True):
+        super().__init__()
+        self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout)  # is a self-attention
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
+                                    heads=n_heads, dim_head=d_head, dropout=dropout)  # is self-attn if context is none
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+
+    def forward(self, x, context=None):
+        return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
+
+    def _forward(self, x, context=None):
+        x = self.attn1(self.norm1(x)) + x
+        x = self.attn2(self.norm2(x), context=context) + x
+        x = self.ff(self.norm3(x)) + x
+        return x
+
+
+class SpatialTransformer(nn.Module):
+    """
+    Transformer block for image-like data.
+    First, project the input (aka embedding)
+    and reshape to b, t, d.
+    Then apply standard transformer action.
+    Finally, reshape to image
+    """
+    def __init__(self, in_channels, n_heads, d_head,
+                 depth=1, dropout=0., context_dim=None):
+        super().__init__()
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = Normalize(in_channels)
+
+        self.proj_in = nn.Conv2d(in_channels,
+                                 inner_dim,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+
+        self.transformer_blocks = nn.ModuleList(
+            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim)
+                for d in range(depth)]
+        )
+
+        self.proj_out = zero_module(nn.Conv2d(inner_dim,
+                                              in_channels,
+                                              kernel_size=1,
+                                              stride=1,
+                                              padding=0))
+
+    def forward(self, x, context=None):
+        # note: if no context is given, cross-attention defaults to self-attention
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        x = self.proj_in(x)
+        x = rearrange(x, 'b c h w -> b (h w) c')
+        for block in self.transformer_blocks:
+            x = block(x, context=context)
+        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
+        x = self.proj_out(x)
+        return x + x_in
\ No newline at end of file
diff --git a/flagai/model/mm/autoencoders.py b/flagai/model/mm/autoencoders.py
new file mode 100755
index 00000000..ad0a1e15
--- /dev/null
+++ b/flagai/model/mm/autoencoders.py
@@ -0,0 +1,483 @@
+import torch
+import pytorch_lightning as pl
+import torch.nn.functional as F
+from contextlib import contextmanager
+
+from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
+from torch.optim.lr_scheduler import LambdaLR
+from flagai.model.mm.model import Encoder, Decoder
+import numpy as np
+from flagai.model.mm.utils import instantiate_from_config
+
+
+class VQModel(torch.nn.Module):
+    def __init__(self,
+                 ddconfig,
+                 lossconfig,
+                 n_embed,
+                 embed_dim,
+                 ckpt_path=None,
+                 ignore_keys=[],
+                 image_key="image",
+                 colorize_nlabels=None,
+                 monitor=None,
+                 batch_resize_range=None,
+                 scheduler_config=None,
+                 lr_g_factor=1.0,
+                 remap=None,
+                 sane_index_shape=False, # tell vector quantizer to return indices as bhw
+                 use_ema=False
+                 ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.n_embed = n_embed
+        self.image_key = image_key
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        self.loss = instantiate_from_config(lossconfig)
+        self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25,
+                                        remap=remap,
+                                        sane_index_shape=sane_index_shape)
+        self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+        if colorize_nlabels is not None:
+            assert type(colorize_nlabels)==int
+            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
+        if monitor is not None:
+            self.monitor = monitor
+        self.batch_resize_range = batch_resize_range
+        if self.batch_resize_range is not None:
+            print(f"{self.__class__.__name__}: Using per-batch resizing in range {batch_resize_range}.")
+
+        self.use_ema = use_ema
+
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+        self.scheduler_config = scheduler_config
+        self.lr_g_factor = lr_g_factor
+
+    @contextmanager
+    def ema_scope(self, context=None):
+        if self.use_ema:
+            self.model_ema.store(self.parameters())
+            self.model_ema.copy_to(self)
+            if context is not None:
+                print(f"{context}: Switched to EMA weights")
+        try:
+            yield None
+        finally:
+            if self.use_ema:
+                self.model_ema.restore(self.parameters())
+                if context is not None:
+                    print(f"{context}: Restored training weights")
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        sd = torch.load(path, map_location="cpu")["state_dict"]
+        keys = list(sd.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print("Deleting key {} from state_dict.".format(k))
+                    del sd[k]
+        missing, unexpected = self.load_state_dict(sd, strict=False)
+        print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
+        if len(missing) > 0:
+            print(f"Missing Keys: {missing}")
+            print(f"Unexpected Keys: {unexpected}")
+
+    def on_train_batch_end(self, *args, **kwargs):
+        if self.use_ema:
+            self.model_ema(self)
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        quant, emb_loss, info = self.quantize(h)
+        return quant, emb_loss, info
+
+    def encode_to_prequant(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        return h
+
+    def decode(self, quant):
+        quant = self.post_quant_conv(quant)
+        dec = self.decoder(quant)
+        return dec
+
+    def decode_code(self, code_b):
+        quant_b = self.quantize.embed_code(code_b)
+        dec = self.decode(quant_b)
+        return dec
+
+    def forward(self, input, return_pred_indices=False):
+        quant, diff, (_,_,ind) = self.encode(input)
+        dec = self.decode(quant)
+        if return_pred_indices:
+            return dec, diff, ind
+        return dec, diff
+
+    def get_input(self, batch, k):
+        x = batch[k]
+        if len(x.shape) == 3:
+            x = x[..., None]
+        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
+        if self.batch_resize_range is not None:
+            lower_size = self.batch_resize_range[0]
+            upper_size = self.batch_resize_range[1]
+            if self.global_step <= 4:
+                # do the first few batches with max size to avoid later oom
+                new_resize = upper_size
+            else:
+                new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16))
+            if new_resize != x.shape[2]:
+                x = F.interpolate(x, size=new_resize, mode="bicubic")
+            x = x.detach()
+        return x
+
+    def training_step(self, batch, batch_idx, optimizer_idx):
+        # https://github.com/pytorch/pytorch/issues/37142
+        # try not to fool the heuristics
+        x = self.get_input(batch, self.image_key)
+        xrec, qloss, ind = self(x, return_pred_indices=True)
+
+        if optimizer_idx == 0:
+            # autoencode
+            aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
+                                            last_layer=self.get_last_layer(), split="train",
+                                            predicted_indices=ind)
+
+            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
+            return aeloss
+
+        if optimizer_idx == 1:
+            # discriminator
+            discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
+                                            last_layer=self.get_last_layer(), split="train")
+            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
+            return discloss
+
+    def validation_step(self, batch, batch_idx):
+        log_dict = self._validation_step(batch, batch_idx)
+        with self.ema_scope():
+            log_dict_ema = self._validation_step(batch, batch_idx, suffix="_ema")
+        return log_dict
+
+    def _validation_step(self, batch, batch_idx, suffix=""):
+        x = self.get_input(batch, self.image_key)
+        xrec, qloss, ind = self(x, return_pred_indices=True)
+        aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0,
+                                        self.global_step,
+                                        last_layer=self.get_last_layer(),
+                                        split="val"+suffix,
+                                        predicted_indices=ind
+                                        )
+
+        discloss, log_dict_disc = self.loss(qloss, x, xrec, 1,
+                                            self.global_step,
+                                            last_layer=self.get_last_layer(),
+                                            split="val"+suffix,
+                                            predicted_indices=ind
+                                            )
+        rec_loss = log_dict_ae[f"val{suffix}/rec_loss"]
+        self.log(f"val{suffix}/rec_loss", rec_loss,
+                   prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
+        self.log(f"val{suffix}/aeloss", aeloss,
+                   prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
+
+        self.log_dict(log_dict_ae)
+        self.log_dict(log_dict_disc)
+        return self.log_dict
+
+    def configure_optimizers(self):
+        lr_d = self.learning_rate
+        lr_g = self.lr_g_factor*self.learning_rate
+        print("lr_d", lr_d)
+        print("lr_g", lr_g)
+        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
+                                  list(self.decoder.parameters())+
+                                  list(self.quantize.parameters())+
+                                  list(self.quant_conv.parameters())+
+                                  list(self.post_quant_conv.parameters()),
+                                  lr=lr_g, betas=(0.5, 0.9))
+        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
+                                    lr=lr_d, betas=(0.5, 0.9))
+
+        if self.scheduler_config is not None:
+            scheduler = instantiate_from_config(self.scheduler_config)
+
+            print("Setting up LambdaLR scheduler...")
+            scheduler = [
+                {
+                    'scheduler': LambdaLR(opt_ae, lr_lambda=scheduler.schedule),
+                    'interval': 'step',
+                    'frequency': 1
+                },
+                {
+                    'scheduler': LambdaLR(opt_disc, lr_lambda=scheduler.schedule),
+                    'interval': 'step',
+                    'frequency': 1
+                },
+            ]
+            return [opt_ae, opt_disc], scheduler
+        return [opt_ae, opt_disc], []
+
+    def get_last_layer(self):
+        return self.decoder.conv_out.weight
+
+    def log_images(self, batch, only_inputs=False, plot_ema=False, **kwargs):
+        log = dict()
+        x = self.get_input(batch, self.image_key)
+        x = x.to(self.device)
+        if only_inputs:
+            log["inputs"] = x
+            return log
+        xrec, _ = self(x)
+        if x.shape[1] > 3:
+            # colorize with random projection
+            assert xrec.shape[1] > 3
+            x = self.to_rgb(x)
+            xrec = self.to_rgb(xrec)
+        log["inputs"] = x
+        log["reconstructions"] = xrec
+        if plot_ema:
+            with self.ema_scope():
+                xrec_ema, _ = self(x)
+                if x.shape[1] > 3: xrec_ema = self.to_rgb(xrec_ema)
+                log["reconstructions_ema"] = xrec_ema
+        return log
+
+    def to_rgb(self, x):
+        assert self.image_key == "segmentation"
+        if not hasattr(self, "colorize"):
+            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
+        x = F.conv2d(x, weight=self.colorize)
+        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
+        return x
+
+
+class VQModelInterface(VQModel):
+    def __init__(self, embed_dim, *args, **kwargs):
+        super().__init__(embed_dim=embed_dim, *args, **kwargs)
+        self.embed_dim = embed_dim
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        return h
+
+    def decode(self, h, force_not_quantize=False):
+        # also go through quantization layer
+        if not force_not_quantize:
+            quant, emb_loss, info = self.quantize(h)
+        else:
+            quant = h
+        quant = self.post_quant_conv(quant)
+        dec = self.decoder(quant)
+        return dec
+
+
+class AutoencoderKL(torch.nn.Module):
+    def __init__(self,
+                 ddconfig,
+                 lossconfig,
+                 embed_dim,
+                 ckpt_path=None,
+                 ignore_keys=[],
+                 image_key="image",
+                 colorize_nlabels=None,
+                 monitor=None,
+                 ):
+        super().__init__()
+        self.image_key = image_key
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        self.loss = instantiate_from_config(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
+        if colorize_nlabels is not None:
+            assert type(colorize_nlabels)==int
+            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
+        if monitor is not None:
+            self.monitor = monitor
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        sd = torch.load(path, map_location="cpu")["state_dict"]
+        keys = list(sd.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print("Deleting key {} from state_dict.".format(k))
+                    del sd[k]
+        self.load_state_dict(sd, strict=False)
+        print(f"Restored from {path}")
+
+    def encode(self, x):
+        h = self.encoder(x)
+        moments = self.quant_conv(h)
+        posterior = DiagonalGaussianDistribution(moments)
+        return posterior
+
+    def decode(self, z):
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+        return dec
+
+    def forward(self, input, sample_posterior=True):
+        posterior = self.encode(input)
+        if sample_posterior:
+            z = posterior.sample()
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        return dec, posterior
+
+    def get_input(self, batch, k):
+        x = batch[k]
+        if len(x.shape) == 3:
+            x = x[..., None]
+        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
+        return x
+
+    def training_step(self, batch, batch_idx, optimizer_idx):
+        inputs = self.get_input(batch, self.image_key)
+        reconstructions, posterior = self(inputs)
+
+        if optimizer_idx == 0:
+            # train encoder+decoder+logvar
+            aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
+                                            last_layer=self.get_last_layer(), split="train")
+            self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
+            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False)
+            return aeloss
+
+        if optimizer_idx == 1:
+            # train the discriminator
+            discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
+                                                last_layer=self.get_last_layer(), split="train")
+
+            self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
+            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False)
+            return discloss
+
+    def validation_step(self, batch, batch_idx):
+        inputs = self.get_input(batch, self.image_key)
+        reconstructions, posterior = self(inputs)
+        aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
+                                        last_layer=self.get_last_layer(), split="val")
+
+        discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
+                                            last_layer=self.get_last_layer(), split="val")
+
+        self.log("val/rec_loss", log_dict_ae["val/rec_loss"])
+        self.log_dict(log_dict_ae)
+        self.log_dict(log_dict_disc)
+        return self.log_dict
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
+                                  list(self.decoder.parameters())+
+                                  list(self.quant_conv.parameters())+
+                                  list(self.post_quant_conv.parameters()),
+                                  lr=lr, betas=(0.5, 0.9))
+        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
+                                    lr=lr, betas=(0.5, 0.9))
+        return [opt_ae, opt_disc], []
+
+    def get_last_layer(self):
+        return self.decoder.conv_out.weight
+
+    @torch.no_grad()
+    def log_images(self, batch, only_inputs=False, **kwargs):
+        log = dict()
+        x = self.get_input(batch, self.image_key)
+        x = x.to(self.device)
+        if not only_inputs:
+            xrec, posterior = self(x)
+            if x.shape[1] > 3:
+                # colorize with random projection
+                assert xrec.shape[1] > 3
+                x = self.to_rgb(x)
+                xrec = self.to_rgb(xrec)
+            log["samples"] = self.decode(torch.randn_like(posterior.sample()))
+            log["reconstructions"] = xrec
+        log["inputs"] = x
+        return log
+
+    def to_rgb(self, x):
+        assert self.image_key == "segmentation"
+        if not hasattr(self, "colorize"):
+            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
+        x = F.conv2d(x, weight=self.colorize)
+        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
+        return x
+
+
+class IdentityFirstStage(torch.nn.Module):
+    def __init__(self, *args, vq_interface=False, **kwargs):
+        self.vq_interface = vq_interface  # TODO: Should be true by default but check to not break older stuff
+        super().__init__()
+
+    def encode(self, x, *args, **kwargs):
+        return x
+
+    def decode(self, x, *args, **kwargs):
+        return x
+
+    def quantize(self, x, *args, **kwargs):
+        if self.vq_interface:
+            return x, None, [None, None, None]
+        return x
+
+    def forward(self, x, *args, **kwargs):
+        return x
+
+
+
+class DiagonalGaussianDistribution(object):
+    def __init__(self, parameters, deterministic=False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
+
+    def sample(self):
+        x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
+        return x
+
+    def kl(self, other=None):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        else:
+            if other is None:
+                return 0.5 * torch.sum(torch.pow(self.mean, 2)
+                                       + self.var - 1.0 - self.logvar,
+                                       dim=[1, 2, 3])
+            else:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean - other.mean, 2) / other.var
+                    + self.var / other.var - 1.0 - self.logvar + other.logvar,
+                    dim=[1, 2, 3])
+
+    def nll(self, sample, dims=[1,2,3]):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims)
+
+    def mode(self):
+        return self.mean
+
+
diff --git a/flagai/model/mm/clip_cn/__init__.py b/flagai/model/mm/clip_cn/__init__.py
new file mode 100755
index 00000000..f19f1cf8
--- /dev/null
+++ b/flagai/model/mm/clip_cn/__init__.py
@@ -0,0 +1,2 @@
+# from .bert_tokenizer import FullTokenizer
+# _tokenizer = FullTokenizer()
\ No newline at end of file
diff --git a/flagai/model/mm/clip_cn/bert_tokenizer.py b/flagai/model/mm/clip_cn/bert_tokenizer.py
new file mode 100755
index 00000000..f024285e
--- /dev/null
+++ b/flagai/model/mm/clip_cn/bert_tokenizer.py
@@ -0,0 +1,439 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tokenization classes."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import sys
+import collections
+import re
+import unicodedata
+import six
+from functools import lru_cache
+import os
+unicode = str
+
+
+@lru_cache()
+def default_vocab():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "vocab.txt")
+
+def validate_case_matches_checkpoint(do_lower_case, init_checkpoint):
+    """Checks whether the casing config is consistent with the checkpoint name."""
+
+    # The casing has to be passed in by the user and there is no explicit check
+    # as to whether it matches the checkpoint. The casing information probably
+    # should have been stored in the bert_config.json file, but it's not, so
+    # we have to heuristically detect it to validate.
+
+    if not init_checkpoint:
+        return
+
+    m = re.match("^.*?([A-Za-z0-9_-]+)/bert_model.ckpt", init_checkpoint)
+    if m is None:
+        return
+
+    model_name = m.group(1)
+
+    lower_models = [
+        "uncased_L-24_H-1024_A-16", "uncased_L-12_H-768_A-12",
+        "multilingual_L-12_H-768_A-12", "chinese_L-12_H-768_A-12"
+    ]
+
+    cased_models = [
+        "cased_L-12_H-768_A-12", "cased_L-24_H-1024_A-16",
+        "multi_cased_L-12_H-768_A-12"
+    ]
+
+    is_bad_config = False
+    if model_name in lower_models and not do_lower_case:
+        is_bad_config = True
+        actual_flag = "False"
+        case_name = "lowercased"
+        opposite_flag = "True"
+
+    if model_name in cased_models and do_lower_case:
+        is_bad_config = True
+        actual_flag = "True"
+        case_name = "cased"
+        opposite_flag = "False"
+
+    if is_bad_config:
+        raise ValueError(
+            "You passed in `--do_lower_case=%s` with `--init_checkpoint=%s`. "
+            "However, `%s` seems to be a %s model, so you "
+            "should pass in `--do_lower_case=%s` so that the fine-tuning matches "
+            "how the model was pre-training. If this error is wrong, please "
+            "just comment out this check." % (actual_flag, init_checkpoint,
+                                              model_name, case_name, opposite_flag))
+
+
+def convert_to_unicode(text):
+    """Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
+    if six.PY3:
+        if isinstance(text, str):
+            return text
+        elif isinstance(text, bytes):
+            return text.decode("utf-8", "ignore")
+        else:
+            raise ValueError("Unsupported string type: %s" % (type(text)))
+    elif six.PY2:
+        if isinstance(text, str):
+            return text.decode("utf-8", "ignore")
+        elif isinstance(text, unicode):
+            return text
+        else:
+            raise ValueError("Unsupported string type: %s" % (type(text)))
+    else:
+        raise ValueError("Not running on Python2 or Python 3?")
+
+
+def printable_text(text):
+    """Returns text encoded in a way suitable for print or `tf.logging`."""
+
+    # These functions want `str` for both Python2 and Python3, but in one case
+    # it's a Unicode string and in the other it's a byte string.
+    if six.PY3:
+        if isinstance(text, str):
+            return text
+        elif isinstance(text, bytes):
+            return text.decode("utf-8", "ignore")
+        else:
+            raise ValueError("Unsupported string type: %s" % (type(text)))
+    elif six.PY2:
+        if isinstance(text, str):
+            return text
+        elif isinstance(text, unicode):
+            return text.encode("utf-8")
+        else:
+            raise ValueError("Unsupported string type: %s" % (type(text)))
+    else:
+        raise ValueError("Not running on Python2 or Python 3?")
+
+
+def load_vocab(vocab_file):
+    """Loads a vocabulary file into a dictionary."""
+    vocab = collections.OrderedDict()
+    index = 0
+    with open(vocab_file, "r") as reader:
+        while True:
+            token = convert_to_unicode(reader.readline())
+            if not token:
+                break
+            token = token.strip()
+            vocab[token] = index
+            index += 1
+    return vocab
+
+
+def convert_by_vocab(vocab, items):
+    """Converts a sequence of [tokens|ids] using the vocab."""
+    output = []
+    for item in items:
+        output.append(vocab[item])
+    return output
+
+
+def convert_tokens_to_ids(vocab, tokens):
+    return convert_by_vocab(vocab, tokens)
+
+
+def convert_ids_to_tokens(inv_vocab, ids):
+    return convert_by_vocab(inv_vocab, ids)
+
+
+def whitespace_tokenize(text):
+    """Runs basic whitespace cleaning and splitting on a piece of text."""
+    text = text.strip()
+    if not text:
+        return []
+    tokens = text.split()
+    return tokens
+
+
+class FullTokenizer(object):
+    """Runs end-to-end tokenziation."""
+
+    def __init__(self, vocab_file=default_vocab(), do_lower_case=True):
+        self.vocab = load_vocab(vocab_file)
+        self.inv_vocab = {v: k for k, v in self.vocab.items()}
+        self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case)
+        self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab)
+
+    def tokenize(self, text):
+        split_tokens = []
+        for token in self.basic_tokenizer.tokenize(text):
+            for sub_token in self.wordpiece_tokenizer.tokenize(token):
+                split_tokens.append(sub_token)
+
+        return split_tokens
+
+    def convert_tokens_to_ids(self, tokens):
+        return convert_by_vocab(self.vocab, tokens)
+
+    def convert_ids_to_tokens(self, ids):
+        return convert_by_vocab(self.inv_vocab, ids)
+
+    @staticmethod
+    def convert_tokens_to_string(tokens, clean_up_tokenization_spaces=True):
+        """ Converts a sequence of tokens (string) in a single string. """
+
+        def clean_up_tokenization(out_string):
+            """ Clean up a list of simple English tokenization artifacts
+            like spaces before punctuations and abreviated forms.
+            """
+            out_string = (
+                out_string.replace(" .", ".")
+                    .replace(" ?", "?")
+                    .replace(" !", "!")
+                    .replace(" ,", ",")
+                    .replace(" ' ", "'")
+                    .replace(" n't", "n't")
+                    .replace(" 'm", "'m")
+                    .replace(" 's", "'s")
+                    .replace(" 've", "'ve")
+                    .replace(" 're", "'re")
+            )
+            return out_string
+
+        text = ' '.join(tokens).replace(' ##', '').strip()
+        if clean_up_tokenization_spaces:
+            clean_text = clean_up_tokenization(text)
+            return clean_text
+        else:
+            return text
+
+    def vocab_size(self):
+        return len(self.vocab)
+
+
+class BasicTokenizer(object):
+    """Runs basic tokenization (punctuation splitting, lower casing, etc.)."""
+
+    def __init__(self, do_lower_case=True):
+        """Constructs a BasicTokenizer.
+
+        Args:
+          do_lower_case: Whether to lower case the input.
+        """
+        self.do_lower_case = do_lower_case
+
+    def tokenize(self, text):
+        """Tokenizes a piece of text."""
+        text = convert_to_unicode(text)
+        text = self._clean_text(text)
+
+        # This was added on November 1st, 2018 for the multilingual and Chinese
+        # models. This is also applied to the English models now, but it doesn't
+        # matter since the English models were not trained on any Chinese data
+        # and generally don't have any Chinese data in them (there are Chinese
+        # characters in the vocabulary because Wikipedia does have some Chinese
+        # words in the English Wikipedia.).
+        text = self._tokenize_chinese_chars(text)
+
+        orig_tokens = whitespace_tokenize(text)
+        split_tokens = []
+        for token in orig_tokens:
+            if self.do_lower_case:
+                token = token.lower()
+                token = self._run_strip_accents(token)
+            split_tokens.extend(self._run_split_on_punc(token))
+
+        output_tokens = whitespace_tokenize(" ".join(split_tokens))
+        return output_tokens
+
+    def _run_strip_accents(self, text):
+        """Strips accents from a piece of text."""
+        text = unicodedata.normalize("NFD", text)
+        output = []
+        for char in text:
+            cat = unicodedata.category(char)
+            if cat == "Mn":
+                continue
+            output.append(char)
+        return "".join(output)
+
+    def _run_split_on_punc(self, text):
+        """Splits punctuation on a piece of text."""
+        chars = list(text)
+        i = 0
+        start_new_word = True
+        output = []
+        while i < len(chars):
+            char = chars[i]
+            if _is_punctuation(char):
+                output.append([char])
+                start_new_word = True
+            else:
+                if start_new_word:
+                    output.append([])
+                start_new_word = False
+                output[-1].append(char)
+            i += 1
+
+        return ["".join(x) for x in output]
+
+    def _tokenize_chinese_chars(self, text):
+        """Adds whitespace around any CJK character."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if self._is_chinese_char(cp):
+                output.append(" ")
+                output.append(char)
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+    def _is_chinese_char(self, cp):
+        """Checks whether CP is the codepoint of a CJK character."""
+        # This defines a "chinese character" as anything in the CJK Unicode block:
+        #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+        #
+        # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+        # despite its name. The modern Korean Hangul alphabet is a different block,
+        # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+        # space-separated words, so they are not treated specially and handled
+        # like the all of the other languages.
+        if ((cp >= 0x4E00 and cp <= 0x9FFF) or  #
+            (cp >= 0x3400 and cp <= 0x4DBF) or  #
+            (cp >= 0x20000 and cp <= 0x2A6DF) or  #
+            (cp >= 0x2A700 and cp <= 0x2B73F) or  #
+            (cp >= 0x2B740 and cp <= 0x2B81F) or  #
+            (cp >= 0x2B820 and cp <= 0x2CEAF) or
+            (cp >= 0xF900 and cp <= 0xFAFF) or  #
+                (cp >= 0x2F800 and cp <= 0x2FA1F)):  #
+            return True
+
+        return False
+
+    def _clean_text(self, text):
+        """Performs invalid character removal and whitespace cleanup on text."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if cp == 0 or cp == 0xfffd or _is_control(char):
+                continue
+            if _is_whitespace(char):
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+
+class WordpieceTokenizer(object):
+    """Runs WordPiece tokenziation."""
+
+    def __init__(self, vocab, unk_token="[UNK]", max_input_chars_per_word=200):
+        self.vocab = vocab
+        self.unk_token = unk_token
+        self.max_input_chars_per_word = max_input_chars_per_word
+
+    def tokenize(self, text):
+        """Tokenizes a piece of text into its word pieces.
+
+        This uses a greedy longest-match-first algorithm to perform tokenization
+        using the given vocabulary.
+
+        For example:
+          input = "unaffable"
+          output = ["un", "##aff", "##able"]
+
+        Args:
+          text: A single token or whitespace separated tokens. This should have
+            already been passed through `BasicTokenizer.
+
+        Returns:
+          A list of wordpiece tokens.
+        """
+
+        text = convert_to_unicode(text)
+
+        output_tokens = []
+        for token in whitespace_tokenize(text):
+            chars = list(token)
+            if len(chars) > self.max_input_chars_per_word:
+                output_tokens.append(self.unk_token)
+                continue
+
+            is_bad = False
+            start = 0
+            sub_tokens = []
+            while start < len(chars):
+                end = len(chars)
+                cur_substr = None
+                while start < end:
+                    substr = "".join(chars[start:end])
+                    if start > 0:
+                        substr = "##" + substr
+                    if substr in self.vocab:
+                        cur_substr = substr
+                        break
+                    end -= 1
+                if cur_substr is None:
+                    is_bad = True
+                    break
+                sub_tokens.append(cur_substr)
+                start = end
+
+            if is_bad:
+                output_tokens.append(self.unk_token)
+            else:
+                output_tokens.extend(sub_tokens)
+        return output_tokens
+
+
+def _is_whitespace(char):
+    """Checks whether `chars` is a whitespace character."""
+    # \t, \n, and \r are technically contorl characters but we treat them
+    # as whitespace since they are generally considered as such.
+    if char == " " or char == "\t" or char == "\n" or char == "\r":
+        return True
+    cat = unicodedata.category(char)
+    if cat == "Zs":
+        return True
+    return False
+
+
+def _is_control(char):
+    """Checks whether `chars` is a control character."""
+    # These are technically control characters but we count them as whitespace
+    # characters.
+    if char == "\t" or char == "\n" or char == "\r":
+        return False
+    cat = unicodedata.category(char)
+    if cat in ("Cc", "Cf"):
+        return True
+    return False
+
+
+def _is_punctuation(char):
+    """Checks whether `chars` is a punctuation character."""
+    cp = ord(char)
+    # We treat all non-letter/number ASCII as punctuation.
+    # Characters such as "^", "$", and "`" are not in the Unicode
+    # Punctuation class but we treat them as punctuation anyways, for
+    # consistency.
+    if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or
+            (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)):
+        return True
+    cat = unicodedata.category(char)
+    if cat.startswith("P"):
+        return True
+    return False
diff --git a/flagai/model/mm/clip_cn/clip.py b/flagai/model/mm/clip_cn/clip.py
new file mode 100644
index 00000000..959aaf6e
--- /dev/null
+++ b/flagai/model/mm/clip_cn/clip.py
@@ -0,0 +1,53 @@
+# Code modified from https://github.com/openai/CLIP
+
+from typing import Union, List
+
+import torch
+
+from .bert_tokenizer import FullTokenizer
+_tokenizer = FullTokenizer()
+from .model import restore_model
+
+__all__ = ["load", "tokenize"]
+
+
+def load(model, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", clip_path=None, bert_path=None):
+    """Load CLIP and BERT model weights
+    """
+    bert_state_dict = torch.load(bert_path, map_location="cpu") if bert_path else None
+    clip_state_dict = torch.load(clip_path, map_location="cpu") if clip_path else None
+
+    restore_model(model, clip_state_dict, bert_state_dict).to(device)
+
+    if str(device) == "cpu":
+        model.float()
+    return model
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 64) -> torch.LongTensor:
+    """
+    Returns the tokenized representation of given input string(s)
+    Parameters
+    ----------
+    texts : Union[str, List[str]]
+        An input string or a list of input strings to tokenize
+    context_length : int
+        The context length to use; all baseline models use 24 as the context length
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    all_tokens = []
+    for text in texts:
+        all_tokens.append([_tokenizer.vocab['[CLS]']] + _tokenizer.convert_tokens_to_ids(_tokenizer.tokenize(text))[:context_length - 2] + [_tokenizer.vocab['[SEP]']])
+
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    for i, tokens in enumerate(all_tokens):
+        assert len(tokens) <= context_length
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result
\ No newline at end of file
diff --git a/flagai/model/mm/clip_cn/configuration_bert.py b/flagai/model/mm/clip_cn/configuration_bert.py
new file mode 100644
index 00000000..491909e9
--- /dev/null
+++ b/flagai/model/mm/clip_cn/configuration_bert.py
@@ -0,0 +1,84 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" BERT model configuration """
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class BertConfig(object):
+    r"""
+        :class:`~transformers.BertConfig` is the configuration class to store the configuration of a
+        `BertModel`.
+
+
+        Arguments:
+            vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `BertModel`.
+            hidden_size: Size of the encoder layers and the pooler layer.
+            num_hidden_layers: Number of hidden layers in the Transformer encoder.
+            num_attention_heads: Number of attention heads for each attention layer in
+                the Transformer encoder.
+            intermediate_size: The size of the "intermediate" (i.e., feed-forward)
+                layer in the Transformer encoder.
+            hidden_act: The non-linear activation function (function or string) in the
+                encoder and pooler. If string, "gelu", "relu", "swish" and "gelu_new" are supported.
+            hidden_dropout_prob: The dropout probabilitiy for all fully connected
+                layers in the embeddings, encoder, and pooler.
+            attention_probs_dropout_prob: The dropout ratio for the attention
+                probabilities.
+            max_position_embeddings: The maximum sequence length that this model might
+                ever be used with. Typically set this to something large just in case
+                (e.g., 512 or 1024 or 2048).
+            type_vocab_size: The vocabulary size of the `token_type_ids` passed into
+                `BertModel`.
+            initializer_range: The sttdev of the truncated_normal_initializer for
+                initializing all weight matrices.
+            layer_norm_eps: The epsilon used by LayerNorm.
+    """
+
+    def __init__(self,
+                 vocab_size_or_config_json_file=30522,
+                 hidden_size=768,
+                 num_hidden_layers=12,
+                 num_attention_heads=12,
+                 intermediate_size=3072,
+                 hidden_act="gelu",
+                 hidden_dropout_prob=0.1,
+                 attention_probs_dropout_prob=0.1,
+                 max_position_embeddings=512,
+                 type_vocab_size=2,
+                 initializer_range=0.02,
+                 layer_norm_eps=1e-12,
+                 output_attentions=False,
+                 output_hidden_states=False
+                 ):
+        self.vocab_size = vocab_size_or_config_json_file
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.output_attentions = output_attentions
+        self.output_hidden_states = output_hidden_states
\ No newline at end of file
diff --git a/flagai/model/mm/clip_cn/model.py b/flagai/model/mm/clip_cn/model.py
new file mode 100755
index 00000000..134fc67d
--- /dev/null
+++ b/flagai/model/mm/clip_cn/model.py
@@ -0,0 +1,497 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import os
+import json
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from .modeling_bert import BertModel
+from .configuration_bert import BertConfig
+from flagai.model.base_model import BaseModel
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(
+                OrderedDict([("-1", nn.AvgPool2d(stride)),
+                             ("0",
+                              nn.Conv2d(inplanes,
+                                        planes * self.expansion,
+                                        1,
+                                        stride=1,
+                                        bias=False)),
+                             ("1", nn.BatchNorm2d(planes * self.expansion))]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.relu(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+
+    def __init__(self,
+                 spacial_dim: int,
+                 embed_dim: int,
+                 num_heads: int,
+                 output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(
+            torch.randn(spacial_dim**2 + 1, embed_dim) / embed_dim**0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.reshape(x.shape[0], x.shape[1],
+                      x.shape[2] * x.shape[3]).permute(2, 0,
+                                                       1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x,
+            key=x,
+            value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat(
+                [self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False)
+
+        return x[0]
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self,
+                 layers,
+                 output_dim,
+                 heads,
+                 input_resolution=224,
+                 width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3,
+                               width // 2,
+                               kernel_size=3,
+                               stride=2,
+                               padding=1,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.conv2 = nn.Conv2d(width // 2,
+                               width // 2,
+                               kernel_size=3,
+                               padding=1,
+                               bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.conv3 = nn.Conv2d(width // 2,
+                               width,
+                               kernel_size=3,
+                               padding=1,
+                               bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.avgpool = nn.AvgPool2d(2)
+        self.relu = nn.ReLU(inplace=True)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim,
+                                        heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+
+        def stem(x):
+            for conv, bn in [(self.conv1, self.bn1), (self.conv2, self.bn2),
+                             (self.conv3, self.bn3)]:
+                x = self.relu(bn(conv(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+
+    def __init__(self,
+                 d_model: int,
+                 n_head: int,
+                 attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(
+            OrderedDict([("c_fc", nn.Linear(d_model, d_model * 4)),
+                         ("gelu", QuickGELU()),
+                         ("c_proj", nn.Linear(d_model * 4, d_model))]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(
+            dtype=x.dtype,
+            device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False,
+                         attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+
+    def __init__(self,
+                 width: int,
+                 layers: int,
+                 heads: int,
+                 attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[
+            ResidualAttentionBlock(width, heads, attn_mask)
+            for _ in range(layers)
+        ])
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+
+class VisualTransformer(nn.Module):
+
+    def __init__(self, input_resolution: int, patch_size: int, width: int,
+                 layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3,
+                               out_channels=width,
+                               kernel_size=patch_size,
+                               stride=patch_size,
+                               bias=False)
+
+        scale = width**-0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn(
+            (input_resolution // patch_size)**2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1],
+                      -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([
+            self.class_embedding.to(x.dtype) + torch.zeros(
+                x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x
+        ],
+                      dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.ln_post(x[:, 0, :])
+
+        if self.proj is not None:
+            x = x @ self.proj
+
+        return x
+
+
+class CLIP(BaseModel):
+
+    def __init__(self, config, **kwargs):
+        #  embed_dim: int,
+        #  # vision
+        #  image_resolution: int,
+        #  vision_layers: Union[Tuple[int, int, int, int], int],
+        #  vision_width: int,
+        #  vision_patch_size: int,
+        #  # text
+        #  vocab_size: int,
+        #  text_attention_probs_dropout_prob: float,
+        #  text_hidden_act: str,
+        #  text_hidden_dropout_prob: float,
+        #  text_hidden_size: int,
+        #  text_initializer_range: float,
+        #  text_intermediate_size: int,
+        #  text_max_position_embeddings: int,
+        #  text_num_attention_heads: int,
+        #  text_num_hidden_layers: int,
+        #  text_type_vocab_size: int,
+        #  tokenizer = _tokenizer,
+        #  ):
+        super().__init__(config, **kwargs)
+        self.config = config
+        embed_dim = config["embed_dim"]
+        # vision
+        image_resolution = config["image_resolution"]
+        vision_layers = config["vision_layers"]
+        vision_width = config["vision_width"]
+        vision_patch_size = config["vision_patch_size"]
+        # text
+        vocab_size = config["vocab_size"]
+        text_attention_probs_dropout_prob = config[
+            "text_attention_probs_dropout_prob"]
+        text_hidden_act = config["text_hidden_act"]
+        text_hidden_dropout_prob = config["text_hidden_dropout_prob"]
+        checkpoint_activations = config[
+            "checkpoint_activations"]  # TODO: implement checkpoint activations
+        text_hidden_size = config["text_hidden_size"]
+        text_initializer_range = config["text_initializer_range"]
+        text_intermediate_size = config["text_intermediate_size"]
+        text_max_position_embeddings = config["text_max_position_embeddings"]
+        text_num_attention_heads = config["text_num_attention_heads"]
+        text_num_hidden_layers = config.get("text_num_hidden_layers", None)
+        text_type_vocab_size = config["text_type_vocab_size"]
+        tokenizer = config["tokenizer"]
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(layers=vision_layers,
+                                         output_dim=embed_dim,
+                                         heads=vision_heads,
+                                         input_resolution=image_resolution,
+                                         width=vision_width)
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisualTransformer(input_resolution=image_resolution,
+                                            patch_size=vision_patch_size,
+                                            width=vision_width,
+                                            layers=vision_layers,
+                                            heads=vision_heads,
+                                            output_dim=embed_dim)
+
+        self.bert_config = BertConfig(
+            vocab_size_or_config_json_file=vocab_size,
+            hidden_size=text_hidden_size,
+            num_hidden_layers=text_num_hidden_layers,
+            num_attention_heads=text_num_attention_heads,
+            intermediate_size=text_intermediate_size,
+            hidden_act=text_hidden_act,
+            hidden_dropout_prob=text_hidden_dropout_prob,
+            attention_probs_dropout_prob=text_attention_probs_dropout_prob,
+            max_position_embeddings=text_max_position_embeddings,
+            type_vocab_size=text_type_vocab_size,
+            initializer_range=text_initializer_range,
+            layer_norm_eps=1e-12,
+        )
+        self.bert = BertModel(self.bert_config)
+
+        self.text_projection = nn.Parameter(
+            torch.empty(text_hidden_size, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.tokenizer = tokenizer
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features**-0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [
+                    self.visual.layer1, self.visual.layer2, self.visual.layer3,
+                    self.visual.layer4
+            ]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection,
+                            std=self.bert_config.hidden_size**-0.5)
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+        pad_index = self.tokenizer.get_command_id('pad')
+        attn_mask = text.ne(pad_index).type(self.dtype)
+        x = self.bert(text, attention_mask=attn_mask)[0].type(
+            self.dtype)  # [batch_size, seq_length, hidden_size]
+        #return x[:, 0, :] @ self.text_projection
+        return x
+
+    def forward(self, image, text):
+        assert image is not None or text is not None, "text and image cannot both be None!"
+
+        if image is None:
+            return self.encode_text(text)
+        elif text is None:
+            return self.encode_image(image)
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        image_features = image_features / image_features.norm(dim=-1,
+                                                              keepdim=True)
+        text_features = text_features / text_features.norm(dim=-1,
+                                                           keepdim=True)
+
+        return image_features, text_features, self.logit_scale.exp()
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [
+                    *[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]],
+                    "in_proj_bias", "bias_k", "bias_v"
+            ]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        if isinstance(l, BertModel):
+            l.to(torch.half)
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def restore_model(model, clip_state_dict: dict, bert_state_dict: dict):
+    merged_state_dict = {}
+
+    # use clip_state_dict to initialize the image encoder & logit scale
+    if clip_state_dict is not None:
+        for k, v in clip_state_dict.items():
+            if k.startswith("visual") or k == "logit_scale":
+                merged_state_dict[k] = v
+
+    # use bert_state_dict to initialize the text encoder
+    if bert_state_dict is not None:
+        for k, v in bert_state_dict.items():
+            if k.startswith("bert"):
+                merged_state_dict[k] = v
+
+    convert_weights(model)
+    model.load_state_dict(merged_state_dict, strict=False)
+    return model.eval()
diff --git a/flagai/model/mm/clip_cn/modeling_bert.py b/flagai/model/mm/clip_cn/modeling_bert.py
new file mode 100755
index 00000000..9a476ec1
--- /dev/null
+++ b/flagai/model/mm/clip_cn/modeling_bert.py
@@ -0,0 +1,442 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch BERT model. """
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import json
+import logging
+import math
+import os
+import sys
+from io import open
+
+import torch
+from torch import nn
+
+from .configuration_bert import BertConfig
+if sys.version_info[0] >= 3:
+    unicode = str
+logger = logging.getLogger(__name__)
+
+def gelu(x):
+    """ Original Implementation of the gelu activation function in Google Bert repo when initially created.
+        For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
+        0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
+        Also see https://arxiv.org/abs/1606.08415
+    """
+    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
+
+def gelu_new(x):
+    """ Implementation of the gelu activation function currently in Google Bert repo (identical to OpenAI GPT).
+        Also see https://arxiv.org/abs/1606.08415
+    """
+    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
+
+def swish(x):
+    return x * torch.sigmoid(x)
+
+
+ACT2FN = {"gelu": gelu, "relu": torch.nn.functional.relu, "swish": swish, "gelu_new": gelu_new}
+
+
+BertLayerNorm = torch.nn.LayerNorm
+
+class BertEmbeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings.
+    """
+    def __init__(self, config):
+        super(BertEmbeddings, self).__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+        # any TensorFlow checkpoint file
+        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, input_ids, token_type_ids=None, position_ids=None):
+        seq_length = input_ids.size(1)
+        if position_ids is None:
+            position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
+            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+
+        words_embeddings = self.word_embeddings(input_ids)
+        position_embeddings = self.position_embeddings(position_ids)
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        embeddings = words_embeddings + position_embeddings + token_type_embeddings
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class BertSelfAttention(nn.Module):
+    def __init__(self, config):
+        super(BertSelfAttention, self).__init__()
+        if config.hidden_size % config.num_attention_heads != 0:
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention "
+                "heads (%d)" % (config.hidden_size, config.num_attention_heads))
+        self.output_attentions = config.output_attentions
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(self, hidden_states, attention_mask=None, head_mask=None):
+        mixed_query_layer = self.query(hidden_states)
+        mixed_key_layer = self.key(hidden_states)
+        mixed_value_layer = self.value(hidden_states)
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+        key_layer = self.transpose_for_scores(mixed_key_layer)
+        value_layer = self.transpose_for_scores(mixed_value_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
+            attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs = self.dropout(attention_probs)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attention_probs = attention_probs * head_mask
+
+        context_layer = torch.matmul(attention_probs, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
+        return outputs
+
+
+class BertSelfOutput(nn.Module):
+    def __init__(self, config):
+        super(BertSelfOutput, self).__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertAttention(nn.Module):
+    def __init__(self, config):
+        super(BertAttention, self).__init__()
+        self.self = BertSelfAttention(config)
+        self.output = BertSelfOutput(config)
+        self.pruned_heads = set()
+
+    def forward(self, input_tensor, attention_mask=None, head_mask=None):
+        self_outputs = self.self(input_tensor, attention_mask, head_mask)
+        attention_output = self.output(self_outputs[0], input_tensor)
+        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+class BertIntermediate(nn.Module):
+    def __init__(self, config):
+        super(BertIntermediate, self).__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str) or (sys.version_info[0] == 2 and isinstance(config.hidden_act, unicode)):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+class BertOutput(nn.Module):
+    def __init__(self, config):
+        super(BertOutput, self).__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertLayer(nn.Module):
+    def __init__(self, config):
+        super(BertLayer, self).__init__()
+        self.attention = BertAttention(config)
+        self.intermediate = BertIntermediate(config)
+        self.output = BertOutput(config)
+
+    def forward(self, hidden_states, attention_mask=None, head_mask=None):
+        attention_outputs = self.attention(hidden_states, attention_mask, head_mask)
+        attention_output = attention_outputs[0]
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        outputs = (layer_output,) + attention_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+class BertEncoder(nn.Module):
+    def __init__(self, config):
+        super(BertEncoder, self).__init__()
+        self.output_attentions = config.output_attentions
+        self.output_hidden_states = config.output_hidden_states
+        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
+
+    def forward(self, hidden_states, attention_mask=None, head_mask=None):
+        all_hidden_states = ()
+        all_attentions = ()
+        for i, layer_module in enumerate(self.layer):
+            if self.output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_outputs = layer_module(hidden_states, attention_mask, head_mask[i])
+            hidden_states = layer_outputs[0]
+
+            if self.output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        # Add last layer
+        if self.output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        outputs = (hidden_states,)
+        if self.output_hidden_states:
+            outputs = outputs + (all_hidden_states,)
+        if self.output_attentions:
+            outputs = outputs + (all_attentions,)
+        return outputs  # last-layer hidden state, (all hidden states), (all attentions)
+
+
+class BertPooler(nn.Module):
+    def __init__(self, config):
+        super(BertPooler, self).__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class BertPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super(BertPredictionHeadTransform, self).__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str) or (sys.version_info[0] == 2 and isinstance(config.hidden_act, unicode)):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
+
+class BertLMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super(BertLMPredictionHead, self).__init__()
+        self.transform = BertPredictionHeadTransform(config)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(config.hidden_size,
+                                 config.vocab_size,
+                                 bias=False)
+
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states) + self.bias
+        return hidden_states
+
+
+class BertOnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super(BertOnlyMLMHead, self).__init__()
+        self.predictions = BertLMPredictionHead(config)
+
+    def forward(self, sequence_output):
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
+class BertOnlyNSPHead(nn.Module):
+    def __init__(self, config):
+        super(BertOnlyNSPHead, self).__init__()
+        self.seq_relationship = nn.Linear(config.hidden_size, 2)
+
+    def forward(self, pooled_output):
+        seq_relationship_score = self.seq_relationship(pooled_output)
+        return seq_relationship_score
+
+
+class BertPreTrainingHeads(nn.Module):
+    def __init__(self, config):
+        super(BertPreTrainingHeads, self).__init__()
+        self.predictions = BertLMPredictionHead(config)
+        self.seq_relationship = nn.Linear(config.hidden_size, 2)
+
+    def forward(self, sequence_output, pooled_output):
+        prediction_scores = self.predictions(sequence_output)
+        seq_relationship_score = self.seq_relationship(pooled_output)
+        return prediction_scores, seq_relationship_score
+
+
+class BertPreTrainedModel(nn.Module):
+    config_class = BertConfig
+    base_model_prefix = "bert"
+
+    def __init__(self, config):
+        super(BertPreTrainedModel, self).__init__() 
+        self.config = config   
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, BertLayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+class BertModel(BertPreTrainedModel):
+    r"""
+    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
+        **last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``
+            Sequence of hidden-states at the output of the last layer of the model.
+        **pooler_output**: ``torch.FloatTensor`` of shape ``(batch_size, hidden_size)``
+            Last layer hidden-state of the first token of the sequence (classification token)
+            further processed by a Linear layer and a Tanh activation function. The Linear
+            layer weights are trained from the next sentence prediction (classification)
+            objective during Bert pretraining. This output is usually *not* a good summary
+            of the semantic content of the input, you're often better with averaging or pooling
+            the sequence of hidden-states for the whole input sequence.
+        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
+            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
+            of shape ``(batch_size, sequence_length, hidden_size)``:
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
+            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
+
+    Examples::
+
+        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        model = BertModel.from_pretrained('bert-base-uncased')
+        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
+        outputs = model(input_ids)
+        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple
+
+    """
+    def __init__(self, config):
+        super(BertModel, self).__init__(config)
+
+        self.embeddings = BertEmbeddings(config)
+        self.encoder = BertEncoder(config)
+        self.pooler = BertPooler(config)
+
+        self.apply(self._init_weights)
+
+    def forward(self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+
+        # We create a 3D attention mask from a 2D tensor mask.
+        # Sizes are [batch_size, 1, 1, to_seq_length]
+        # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
+        # this attention mask is more simple than the triangular masking of causal attention
+        # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
+        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+
+        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+        # masked positions, this operation will create a tensor which is 0.0 for
+        # positions we want to attend and -10000.0 for masked positions.
+        # Since we are adding it to the raw scores before the softmax, this is
+        # effectively the same as removing these entirely.
+        extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        if head_mask is not None:
+            if head_mask.dim() == 1:
+                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+                head_mask = head_mask.expand(self.config.num_hidden_layers, -1, -1, -1, -1)
+            elif head_mask.dim() == 2:
+                head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
+            head_mask = head_mask.to(dtype=next(self.parameters()).dtype) # switch to fload if need + fp16 compatibility
+        else:
+            head_mask = [None] * self.config.num_hidden_layers
+
+        embedding_output = self.embeddings(input_ids, position_ids=position_ids, token_type_ids=token_type_ids)
+        encoder_outputs = self.encoder(embedding_output,
+                                       extended_attention_mask,
+                                       head_mask=head_mask)
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output)
+
+        outputs = (sequence_output, pooled_output,) + encoder_outputs[1:]  # add hidden_states and attentions if they are here
+        return outputs  # sequence_output, pooled_output, (hidden_states), (attentions)
\ No newline at end of file
diff --git a/flagai/model/mm/clip_model.py b/flagai/model/mm/clip_model.py
old mode 100644
new mode 100755
diff --git a/flagai/model/mm/diffusion.py b/flagai/model/mm/diffusion.py
new file mode 100755
index 00000000..fa813e75
--- /dev/null
+++ b/flagai/model/mm/diffusion.py
@@ -0,0 +1,1950 @@
+import sys
+import torch
+import torch.nn as nn
+import numpy as np
+from torch.optim.lr_scheduler import LambdaLR
+from einops import rearrange, repeat
+from contextlib import contextmanager
+from functools import partial
+from tqdm import tqdm
+from torchvision.utils import make_grid
+from pytorch_lightning.utilities.distributed import rank_zero_only
+from flagai.model.mm.utils import exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config, log_txt_as_img
+from flagai.model.mm.autoencoders import VQModelInterface, IdentityFirstStage, AutoencoderKL
+from flagai.model.mm.utils import make_beta_schedule, extract_into_tensor, noise_like
+from flagai.model.mm.Sampler import DDIMSampler
+from flagai.model.base_model import BaseModel
+from flagai.auto_model.auto_loader import AutoLoader
+
+__conditioning_keys__ = {
+    'concat': 'c_concat',
+    'crossattn': 'c_crossattn',
+    'adm': 'y'
+}
+
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+
+def uniform_on_device(r1, r2, shape, device):
+    return (r1 - r2) * torch.rand(*shape, device=device) + r2
+
+
+class DDPM(BaseModel):
+    # classic DDPM with Gaussian diffusion, in image space
+    def __init__(
+        self,
+        unet_config,
+        timesteps=1000,
+        beta_schedule="linear",
+        loss_type="l2",
+        ckpt_path=None,
+        ignore_keys=[],
+        load_only_unet=False,
+        monitor="val/loss",
+        use_ema=False,
+        first_stage_key="image",
+        image_size=256,
+        channels=3,
+        log_every_t=100,
+        clip_denoised=True,
+        linear_start=1e-4,
+        linear_end=2e-2,
+        cosine_s=8e-3,
+        given_betas=None,
+        original_elbo_weight=0.,
+        v_posterior=0.,  # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta
+        l_simple_weight=1.,
+        conditioning_key=None,
+        parameterization="eps",  # all assuming fixed variance schedules
+        scheduler_config=None,
+        use_positional_encodings=False,
+        learn_logvar=False,
+        logvar_init=0.,
+        **kwargs,
+    ):
+        super(DDPM, self).__init__(unet_config, **kwargs)
+        self.device = torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu")
+        assert parameterization in [
+            "eps", "x0"
+        ], 'currently only supporting "eps" and "x0"'
+        self.parameterization = parameterization
+        print(
+            f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode"
+        )
+        self.cond_stage_model = None
+        self.clip_denoised = clip_denoised
+        self.log_every_t = log_every_t
+        self.first_stage_key = first_stage_key
+        self.image_size = image_size  # try conv?
+        self.channels = channels
+        self.use_positional_encodings = use_positional_encodings
+        self.model = DiffusionWrapper(unet_config, conditioning_key)
+        self.use_ema = use_ema
+        count_params(self.model, verbose=True)
+
+        self.use_scheduler = scheduler_config is not None
+        if self.use_scheduler:
+            self.scheduler_config = scheduler_config
+
+        self.v_posterior = v_posterior
+        self.original_elbo_weight = original_elbo_weight
+        self.l_simple_weight = l_simple_weight
+
+        if monitor is not None:
+            self.monitor = monitor
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path,
+                                ignore_keys=ignore_keys,
+                                only_model=load_only_unet)
+
+        self.register_schedule(given_betas=given_betas,
+                               beta_schedule=beta_schedule,
+                               timesteps=timesteps,
+                               linear_start=linear_start,
+                               linear_end=linear_end,
+                               cosine_s=cosine_s)
+
+        self.loss_type = loss_type
+
+        self.learn_logvar = learn_logvar
+        self.logvar = torch.full(fill_value=logvar_init,
+                                 size=(self.num_timesteps, ))
+        if self.learn_logvar:
+            self.logvar = nn.Parameter(self.logvar, requires_grad=True)
+
+    def register_schedule(self,
+                          given_betas=None,
+                          beta_schedule="linear",
+                          timesteps=1000,
+                          linear_start=1e-4,
+                          linear_end=2e-2,
+                          cosine_s=8e-3):
+        if exists(given_betas):
+            betas = given_betas
+        else:
+            betas = make_beta_schedule(beta_schedule,
+                                       timesteps,
+                                       linear_start=linear_start,
+                                       linear_end=linear_end,
+                                       cosine_s=cosine_s)
+        alphas = 1. - betas
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+
+        timesteps, = betas.shape
+        self.num_timesteps = int(timesteps)
+        self.linear_start = linear_start
+        self.linear_end = linear_end
+        assert alphas_cumprod.shape[
+            0] == self.num_timesteps, 'alphas have to be defined for each timestep'
+
+        to_torch = partial(torch.tensor, dtype=torch.float32)
+
+        self.register_buffer('betas', to_torch(betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev',
+                             to_torch(alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod',
+                             to_torch(np.sqrt(alphas_cumprod)))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod',
+                             to_torch(np.sqrt(1. - alphas_cumprod)))
+        self.register_buffer('log_one_minus_alphas_cumprod',
+                             to_torch(np.log(1. - alphas_cumprod)))
+        self.register_buffer('sqrt_recip_alphas_cumprod',
+                             to_torch(np.sqrt(1. / alphas_cumprod)))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod',
+                             to_torch(np.sqrt(1. / alphas_cumprod - 1)))
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        posterior_variance = (1 - self.v_posterior) * betas * (
+            1. - alphas_cumprod_prev) / (
+                1. - alphas_cumprod) + self.v_posterior * betas
+        # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
+        self.register_buffer('posterior_variance',
+                             to_torch(posterior_variance))
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.register_buffer(
+            'posterior_log_variance_clipped',
+            to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
+        self.register_buffer(
+            'posterior_mean_coef1',
+            to_torch(betas * np.sqrt(alphas_cumprod_prev) /
+                     (1. - alphas_cumprod)))
+        self.register_buffer(
+            'posterior_mean_coef2',
+            to_torch((1. - alphas_cumprod_prev) * np.sqrt(alphas) /
+                     (1. - alphas_cumprod)))
+
+        if self.parameterization == "eps":
+            lvlb_weights = self.betas**2 / (2 * self.posterior_variance *
+                                            to_torch(alphas) *
+                                            (1 - self.alphas_cumprod))
+        elif self.parameterization == "x0":
+            lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (
+                2. * 1 - torch.Tensor(alphas_cumprod))
+        else:
+            raise NotImplementedError("mu not supported")
+        # TODO how to choose this term
+        lvlb_weights[0] = lvlb_weights[1]
+        self.register_buffer('lvlb_weights', lvlb_weights, persistent=False)
+        assert not torch.isnan(self.lvlb_weights).all()
+
+    @contextmanager
+    def ema_scope(self, context=None):
+        if self.use_ema:
+            self.model_ema.store(self.model.parameters())
+            self.model_ema.copy_to(self.model)
+            if context is not None:
+                print(f"{context}: Switched to EMA weights")
+        try:
+            yield None
+        finally:
+            if self.use_ema:
+                self.model_ema.restore(self.model.parameters())
+                if context is not None:
+                    print(f"{context}: Restored training weights")
+
+    def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
+        sd = torch.load(path, map_location="cpu")
+        if "state_dict" in list(sd.keys()):
+            sd = sd["state_dict"]
+        keys = list(sd.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print("Deleting key {} from state_dict.".format(k))
+                    del sd[k]
+        missing, unexpected = self.load_state_dict(
+            sd,
+            strict=False) if not only_model else self.model.load_state_dict(
+                sd, strict=False)
+        print(
+            f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys"
+        )
+        if len(missing) > 0:
+            print(f"Missing Keys: {missing}")
+        if len(unexpected) > 0:
+            print(f"Unexpected Keys: {unexpected}")
+
+    def q_mean_variance(self, x_start, t):
+        """
+        Get the distribution q(x_t | x_0).
+        :param x_start: the [N x C x ...] tensor of noiseless inputs.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+        """
+        mean = (
+            extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) *
+            x_start)
+        variance = extract_into_tensor(1.0 - self.alphas_cumprod, t,
+                                       x_start.shape)
+        log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod,
+                                           t, x_start.shape)
+        return mean, variance, log_variance
+
+    def predict_start_from_noise(self, x_t, t, noise):
+        return (
+            extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) *
+            x_t - extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t,
+                                      x_t.shape) * noise)
+
+    def q_posterior(self, x_start, x_t, t):
+        posterior_mean = (
+            extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) *
+            x_start +
+            extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t)
+        posterior_variance = extract_into_tensor(self.posterior_variance, t,
+                                                 x_t.shape)
+        posterior_log_variance_clipped = extract_into_tensor(
+            self.posterior_log_variance_clipped, t, x_t.shape)
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(self, x, t, clip_denoised: bool):
+        model_out = self.model(x, t)
+        if self.parameterization == "eps":
+            x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
+        elif self.parameterization == "x0":
+            x_recon = model_out
+        if clip_denoised:
+            x_recon.clamp_(-1., 1.)
+
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(
+            x_start=x_recon, x_t=x, t=t)
+        return model_mean, posterior_variance, posterior_log_variance
+
+    @torch.no_grad()
+    def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
+        b, *_, device = *x.shape, x.device
+        model_mean, _, model_log_variance = self.p_mean_variance(
+            x=x, t=t, clip_denoised=clip_denoised)
+        noise = noise_like(x.shape, device, repeat_noise)
+        # no noise when t == 0
+        nonzero_mask = (1 - (t == 0).float()).reshape(
+            b, *((1, ) * (len(x.shape) - 1)))
+        return model_mean + nonzero_mask * (0.5 *
+                                            model_log_variance).exp() * noise
+
+    @torch.no_grad()
+    def p_sample_loop(self, shape, return_intermediates=False):
+        device = self.betas.device
+        b = shape[0]
+        img = torch.randn(shape, device=device)
+        intermediates = [img]
+        for i in tqdm(reversed(range(0, self.num_timesteps)),
+                      desc='Sampling t',
+                      total=self.num_timesteps):
+            img = self.p_sample(img,
+                                torch.full((b, ),
+                                           i,
+                                           device=device,
+                                           dtype=torch.long),
+                                clip_denoised=self.clip_denoised)
+            if i % self.log_every_t == 0 or i == self.num_timesteps - 1:
+                intermediates.append(img)
+        if return_intermediates:
+            return img, intermediates
+        return img
+
+    @torch.no_grad()
+    def sample(self, batch_size=16, return_intermediates=False):
+        image_size = self.image_size
+        channels = self.channels
+        return self.p_sample_loop(
+            (batch_size, channels, image_size, image_size),
+            return_intermediates=return_intermediates)
+
+    def q_sample(self, x_start, t, noise=None):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        return (
+            extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) *
+            x_start + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod,
+                                          t, x_start.shape) * noise)
+
+    def get_loss(self, pred, target, mean=True):
+        if self.loss_type == 'l1':
+            loss = (target - pred).abs()
+            if mean:
+                loss = loss.mean()
+        elif self.loss_type == 'l2':
+            if mean:
+                loss = torch.nn.functional.mse_loss(target, pred)
+            else:
+                loss = torch.nn.functional.mse_loss(target,
+                                                    pred,
+                                                    reduction='none')
+        else:
+            raise NotImplementedError("unknown loss type '{loss_type}'")
+
+        return loss
+
+    def p_losses(self, x_start, t, noise=None):
+
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        # print(f"after get noise: {torch.cuda.memory_allocated()/1024/1024/1024:.2f}GB")
+        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+        # print(f"after sampled noise: {torch.cuda.memory_allocated()/1024/1024/1024:.2f}GB")
+        model_out = self.model(x_noisy, t)
+        # print(f"after predict noise: {torch.cuda.memory_allocated()/1024/1024/1024:.2f}GB")
+
+        loss_dict = {}
+        if self.parameterization == "eps":
+            target = noise
+        elif self.parameterization == "x0":
+            target = x_start
+        else:
+            raise NotImplementedError(
+                f"Paramterization {self.parameterization} not yet supported")
+
+        loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3])
+        # print(f"after get loss: {torch.cuda.memory_allocated()/1024/1024/1024:.2f}GB")
+
+        log_prefix = 'train' if self.training else 'val'
+
+        loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()})
+        loss_simple = loss.mean() * self.l_simple_weight
+
+        loss_vlb = (self.lvlb_weights[t] * loss).mean()
+        loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb})
+
+        loss = loss_simple + self.original_elbo_weight * loss_vlb
+
+        loss_dict.update({f'{log_prefix}/loss': loss})
+
+        return loss, loss_dict
+
+    def forward(self, x, *args, **kwargs):
+        # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
+        # assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
+        t = torch.randint(0,
+                          self.num_timesteps, (x.shape[0], ),
+                          device=self.device).long()
+        return self.p_losses(x, t, *args, **kwargs)
+
+    def get_input(self, batch, k):
+        # The k here stands for key, and is primarily used to handle the case that batch contains both image and caption
+        x = batch[k]
+        # In unsupervised learning, k is not used
+        #x = batch
+        if len(x.shape) == 3:
+            x = x[..., None]
+        # dataseet中已经处理了图片顺序，不需要再做rearrange
+        # x = rearrange(x, 'b h w c -> b c h w')
+        x = x.to(memory_format=torch.contiguous_format).float()
+        return x
+
+    def shared_step(self, batch):
+        x = self.get_input(batch, self.first_stage_key)
+        loss, loss_dict = self(x)
+        return loss, loss_dict
+
+    def training_step(self, batch, batch_idx):
+
+        loss, loss_dict = self.shared_step(batch)
+        self.log_dict(loss_dict,
+                      prog_bar=True,
+                      logger=True,
+                      on_step=True,
+                      on_epoch=True)
+        try:
+            self.log("global_step",
+                     self.global_step,
+                     prog_bar=True,
+                     logger=True,
+                     on_step=True,
+                     on_epoch=False)
+
+            if self.use_scheduler:
+                lr = self.optimizers().param_groups[0]['lr']
+                self.log('lr_abs',
+                         lr,
+                         prog_bar=True,
+                         logger=True,
+                         on_step=True,
+                         on_epoch=False)
+        except:
+            return loss
+
+        return loss
+
+    @torch.no_grad()
+    def validation_step(self, batch, batch_idx):
+        _, loss_dict_no_ema = self.shared_step(batch)
+        with self.ema_scope():
+            _, loss_dict_ema = self.shared_step(batch)
+            loss_dict_ema = {
+                key + '_ema': loss_dict_ema[key]
+                for key in loss_dict_ema
+            }
+        self.log_dict(loss_dict_no_ema,
+                      prog_bar=False,
+                      logger=True,
+                      on_step=False,
+                      on_epoch=True)
+        self.log_dict(loss_dict_ema,
+                      prog_bar=False,
+                      logger=True,
+                      on_step=False,
+                      on_epoch=True)
+
+    def on_train_batch_end(self, *args, **kwargs):
+        if self.use_ema:
+            self.model_ema(self.model)
+
+    def _get_rows_from_list(self, samples):
+        n_imgs_per_row = len(samples)
+        denoise_grid = rearrange(samples, 'n b c h w -> b n c h w')
+        denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
+        denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
+        return denoise_grid
+
+    @torch.no_grad()
+    def log_images(self,
+                   batch,
+                   N=8,
+                   n_row=2,
+                   sample=True,
+                   return_keys=None,
+                   **kwargs):
+        log = dict()
+        x = self.get_input(batch, self.first_stage_key)
+        N = min(x.shape[0], N)
+        n_row = min(x.shape[0], n_row)
+        x = x.to(self.device)[:N]
+        log["inputs"] = x
+
+        # get diffusion row
+        diffusion_row = list()
+        x_start = x[:n_row]
+
+        for t in range(self.num_timesteps):
+            if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+                t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
+                t = t.to(self.device).long()
+                noise = torch.randn_like(x_start)
+                x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+                diffusion_row.append(x_noisy)
+
+        log["diffusion_row"] = self._get_rows_from_list(diffusion_row)
+
+        if sample:
+            # get denoise row
+            with self.ema_scope("Plotting"):
+                samples, denoise_row = self.sample(batch_size=N,
+                                                   return_intermediates=True)
+
+            log["samples"] = samples
+            log["denoise_row"] = self._get_rows_from_list(denoise_row)
+
+        if return_keys:
+            if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
+                return log
+            else:
+                return {key: log[key] for key in return_keys}
+        return log
+
+
+class LatentDiffusion(DDPM):
+    """main class"""
+
+    def __init__(self,
+                 first_stage_config,
+                 cond_stage_config,
+                 num_timesteps_cond=None,
+                 cond_stage_key="image",
+                 cond_stage_trainable=False,
+                 concat_mode=True,
+                 cond_stage_forward=None,
+                 conditioning_key=None,
+                 scale_factor=1.0,
+                 scale_by_std=False,
+                 tokenizer=None,
+                 *args,
+                 **kwargs):
+        self.tokenizer = tokenizer
+        self.num_timesteps_cond = default(num_timesteps_cond, 1)
+        self.scale_by_std = scale_by_std
+        assert self.num_timesteps_cond <= kwargs['timesteps']
+        # for backwards compatibility after implementation of DiffusionWrapper
+        if conditioning_key is None:
+            conditioning_key = 'concat' if concat_mode else 'crossattn'
+        if cond_stage_config == '__is_unconditional__':
+            conditioning_key = None
+
+        ckpt_path = kwargs.pop("ckpt_path", None)
+        ignore_keys = kwargs.pop("ignore_keys", [])
+        super().__init__(conditioning_key=conditioning_key, *args, **kwargs)
+        self.concat_mode = concat_mode
+        self.cond_stage_trainable = cond_stage_trainable
+        self.cond_stage_key = cond_stage_key
+        self.conditioning_key = conditioning_key
+        try:
+            self.num_downs = len(
+                first_stage_config.params.ddconfig.ch_mult) - 1
+        except:
+            self.num_downs = 0
+        if not scale_by_std:
+            self.scale_factor = scale_factor
+        else:
+            self.register_buffer('scale_factor', torch.tensor(scale_factor))
+        self.instantiate_first_stage(first_stage_config)
+        self.instantiate_cond_stage(cond_stage_config)
+        self.cond_stage_forward = cond_stage_forward
+        self.clip_denoised = False
+        self.bbox_tokenizer = None
+
+        self.restarted_from_ckpt = False
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys)
+            self.restarted_from_ckpt = True
+
+    def make_cond_schedule(self, ):
+        self.cond_ids = torch.full(size=(self.num_timesteps, ),
+                                   fill_value=self.num_timesteps - 1,
+                                   dtype=torch.long)
+        ids = torch.round(
+            torch.linspace(0, self.num_timesteps - 1,
+                           self.num_timesteps_cond)).long()
+        self.cond_ids[:self.num_timesteps_cond] = ids
+
+    # In the preprocessing before the start of training, an error about positional parameter will appear
+    @rank_zero_only
+    @torch.no_grad()
+    #def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
+    def on_train_batch_start(self, batch, batch_idx):
+        # only for very first batch
+        if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt:
+            assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously'
+            # set rescale weight to 1./std of encodings
+            print("### USING STD-RESCALING ###")
+            x = super().get_input(batch, self.first_stage_key)
+            x = x.to(self.device)
+            # There will be an compression of x
+            encoder_posterior = self.encode_first_stage(x)
+            z = self.get_first_stage_encoding(encoder_posterior).detach()
+            del self.scale_factor
+            self.register_buffer('scale_factor', 1. / z.flatten().std())
+            print(f"setting self.scale_factor to {self.scale_factor}")
+            print("### USING STD-RESCALING ###")
+
+    def register_schedule(self,
+                          given_betas=None,
+                          beta_schedule="linear",
+                          timesteps=1000,
+                          linear_start=1e-4,
+                          linear_end=2e-2,
+                          cosine_s=8e-3):
+        super().register_schedule(given_betas, beta_schedule, timesteps,
+                                  linear_start, linear_end, cosine_s)
+
+        self.shorten_cond_schedule = self.num_timesteps_cond > 1
+        if self.shorten_cond_schedule:
+            self.make_cond_schedule()
+
+    def instantiate_first_stage(self, config):
+        model = instantiate_from_config(config)
+        self.first_stage_model = model.eval()
+        self.first_stage_model.train = disabled_train
+        for param in self.first_stage_model.parameters():
+            param.requires_grad = False
+
+    def instantiate_cond_stage(self, config):
+        dct = config.get("params", dict())
+        model_dir = dct.get("model_dir", None)
+        if not model_dir:
+            model_dir = dct.get("download_path", None)
+        if not self.cond_stage_trainable:
+            if config == "__is_first_stage__":
+                print("Using first stage also as cond stage.")
+                self.cond_stage_model = self.first_stage_model
+            elif config == "__is_unconditional__":
+                print(
+                    f"Training {self.__class__.__name__} as an unconditional model."
+                )
+                self.cond_stage_model = None
+            else:
+                loader = AutoLoader(
+                    task_name="txt_img_matching",  #contrastive learning
+                    model_name=dct["model_name"],
+                    model_dir=model_dir)
+                model = loader.get_model()
+                tokenizer = loader.get_tokenizer()
+                self.tokenizer = tokenizer
+                model.to(self.device)
+                self.cond_stage_model = model.eval()
+                self.cond_stage_model.train = disabled_train
+                for param in self.cond_stage_model.parameters():
+                    param.requires_grad = False
+        else:
+            assert config != '__is_first_stage__'
+            assert config != '__is_unconditional__'
+            loader = AutoLoader(
+                task_name="txt_img_matching",  #contrastive learning
+                model_name=dct["model_name"],
+                model_dir=model_dir)
+            tokenizer = loader.get_tokenizer()
+            self.tokenizer = tokenizer
+            model = loader.get_model()
+            model.to(self.device)
+            self.cond_stage_model = model
+
+    def _get_denoise_row_from_list(self,
+                                   samples,
+                                   desc='',
+                                   force_no_decoder_quantization=False):
+        denoise_row = []
+        for zd in tqdm(samples, desc=desc):
+            denoise_row.append(
+                self.decode_first_stage(
+                    zd.to(self.device),
+                    force_not_quantize=force_no_decoder_quantization))
+        n_imgs_per_row = len(denoise_row)
+        denoise_row = torch.stack(denoise_row)  # n_log_step, n_row, C, H, W
+        denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w')
+        denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
+        denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
+        return denoise_grid
+
+    def get_first_stage_encoding(self, encoder_posterior):
+        if isinstance(encoder_posterior, DiagonalGaussianDistribution):
+            z = encoder_posterior.sample()
+        elif isinstance(encoder_posterior, torch.Tensor):
+            z = encoder_posterior
+        else:
+            # raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented")
+
+            z = encoder_posterior.sample()
+        return self.scale_factor * z
+
+    def get_learned_conditioning(self, c):
+        # C will be directly returned if it is None
+        if c is None:
+            return None
+        if self.cond_stage_forward is None:
+            if hasattr(self.cond_stage_model, 'encode') and callable(
+                    self.cond_stage_model.encode):
+                c = self.cond_stage_model.encode(c, self.tokenizer)
+                if isinstance(c, DiagonalGaussianDistribution):
+                    c = c.mode()
+            else:
+                c = self.cond_stage_model(c)
+        else:
+            assert hasattr(self.cond_stage_model, self.cond_stage_forward)
+            c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
+        return c
+
+    def meshgrid(self, h, w):
+        y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1)
+        x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1)
+
+        arr = torch.cat([y, x], dim=-1)
+        return arr
+
+    def delta_border(self, h, w):
+        """
+        :param h: height
+        :param w: width
+        :return: normalized distance to image border,
+         wtith min distance = 0 at border and max dist = 0.5 at image center
+        """
+        lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2)
+        arr = self.meshgrid(h, w) / lower_right_corner
+        dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0]
+        dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0]
+        edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down],
+                                        dim=-1),
+                              dim=-1)[0]
+        return edge_dist
+
+    def get_weighting(self, h, w, Ly, Lx, device):
+        weighting = self.delta_border(h, w)
+        weighting = torch.clip(
+            weighting,
+            self.split_input_params["clip_min_weight"],
+            self.split_input_params["clip_max_weight"],
+        )
+        weighting = weighting.view(1, h * w, 1).repeat(1, 1,
+                                                       Ly * Lx).to(device)
+
+        if self.split_input_params["tie_braker"]:
+            L_weighting = self.delta_border(Ly, Lx)
+            L_weighting = torch.clip(
+                L_weighting, self.split_input_params["clip_min_tie_weight"],
+                self.split_input_params["clip_max_tie_weight"])
+
+            L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device)
+            weighting = weighting * L_weighting
+        return weighting
+
+    def get_fold_unfold(self,
+                        x,
+                        kernel_size,
+                        stride,
+                        uf=1,
+                        df=1):  # todo load once not every time, shorten code
+        """
+        :param x: img of size (bs, c, h, w)
+        :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1])
+        """
+        bs, nc, h, w = x.shape
+
+        # number of crops in image
+        Ly = (h - kernel_size[0]) // stride[0] + 1
+        Lx = (w - kernel_size[1]) // stride[1] + 1
+
+        if uf == 1 and df == 1:
+            fold_params = dict(kernel_size=kernel_size,
+                               dilation=1,
+                               padding=0,
+                               stride=stride)
+            unfold = torch.nn.Unfold(**fold_params)
+
+            fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params)
+
+            weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly,
+                                           Lx, x.device).to(x.dtype)
+            normalization = fold(weighting).view(1, 1, h,
+                                                 w)  # normalizes the overlap
+            weighting = weighting.view(
+                (1, 1, kernel_size[0], kernel_size[1], Ly * Lx))
+
+        elif uf > 1 and df == 1:
+            fold_params = dict(kernel_size=kernel_size,
+                               dilation=1,
+                               padding=0,
+                               stride=stride)
+            unfold = torch.nn.Unfold(**fold_params)
+
+            fold_params2 = dict(kernel_size=(kernel_size[0] * uf,
+                                             kernel_size[0] * uf),
+                                dilation=1,
+                                padding=0,
+                                stride=(stride[0] * uf, stride[1] * uf))
+            fold = torch.nn.Fold(output_size=(x.shape[2] * uf,
+                                              x.shape[3] * uf),
+                                 **fold_params2)
+
+            weighting = self.get_weighting(kernel_size[0] * uf,
+                                           kernel_size[1] * uf, Ly, Lx,
+                                           x.device).to(x.dtype)
+            normalization = fold(weighting).view(1, 1, h * uf, w *
+                                                 uf)  # normalizes the overlap
+            weighting = weighting.view(
+                (1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx))
+
+        elif df > 1 and uf == 1:
+            fold_params = dict(kernel_size=kernel_size,
+                               dilation=1,
+                               padding=0,
+                               stride=stride)
+            unfold = torch.nn.Unfold(**fold_params)
+
+            fold_params2 = dict(kernel_size=(kernel_size[0] // df,
+                                             kernel_size[0] // df),
+                                dilation=1,
+                                padding=0,
+                                stride=(stride[0] // df, stride[1] // df))
+            fold = torch.nn.Fold(output_size=(x.shape[2] // df,
+                                              x.shape[3] // df),
+                                 **fold_params2)
+
+            weighting = self.get_weighting(kernel_size[0] // df,
+                                           kernel_size[1] // df, Ly, Lx,
+                                           x.device).to(x.dtype)
+            normalization = fold(weighting).view(1, 1, h // df, w //
+                                                 df)  # normalizes the overlap
+            weighting = weighting.view(
+                (1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx))
+
+        else:
+            raise NotImplementedError
+
+        return fold, unfold, normalization, weighting
+
+    @torch.no_grad()
+    def get_input(self,
+                  batch,
+                  k,
+                  return_first_stage_outputs=False,
+                  force_c_encode=False,
+                  cond_key=None,
+                  return_original_cond=False,
+                  bs=None):
+        x = super().get_input(batch, k)
+        if bs is not None:
+            x = x[:bs]
+        x = x.to(self.device)
+        # VAE-style compression by default
+        encoder_posterior = self.encode_first_stage(x)
+
+        z = self.get_first_stage_encoding(encoder_posterior).detach()
+        # self.model.conditioning_key is of type str by default. Note that setting it to None in config file will return string 'None'
+        # Comment it in yaml if you do not want to set it
+
+        if self.model.conditioning_key is not None:
+            if cond_key is None:
+                cond_key = self.cond_stage_key
+
+            if cond_key != self.first_stage_key:
+                if cond_key in ['caption', 'coordinates_bbox']:
+                    xc = batch[cond_key]
+                elif cond_key == 'class_label':
+                    xc = batch
+                # Add the condition for text-image modification
+                elif cond_key == "img_and_caption":
+                    xc = batch["img_and_caption"]
+                else:
+                    xc = super().get_input(batch, cond_key).to(self.device)
+            else:
+                # Obtain xc according to the dictionary
+                xc = x
+
+            if not self.cond_stage_trainable or force_c_encode:
+                if isinstance(xc, dict) or isinstance(xc, list):
+                    # import pudb; pudb.set_trace()
+                    # Determine if learning the cond, otherwise it will return
+                    # if the image editing is driven by texts, we will process the init_img and caption
+                    if cond_key == "img_and_caption":
+                        # To prevent xc["init_img"] saved after being encoded, we fetch it from batch at each time
+                        c = {}
+                        init_img = batch["img_and_caption"]["init_img"]
+                        caption = batch["img_and_caption"]["caption"]
+                        init_img_posterior = self.encode_first_stage(init_img)
+                        init_img_posterior = self.get_first_stage_encoding(
+                            init_img_posterior).detach()
+
+                        c["init_img"] = self.get_first_stage_encoding(
+                            init_img_posterior)
+                        c["caption"] = self.get_learned_conditioning(caption)
+
+                    else:
+                        c = self.get_learned_conditioning(xc)
+                else:
+                    #
+                    c = self.get_learned_conditioning(xc.to(self.device))
+            else:
+                c = xc
+            if bs is not None:
+                if cond_key == "img_and_caption":
+                    c["init_img"] = c["init_img"][:bs]
+                    c["caption"] = c["caption"][:bs]
+                else:
+                    c = c[:bs]
+
+            if self.use_positional_encodings:
+                pos_x, pos_y = self.compute_latent_shifts(batch)
+                ckey = __conditioning_keys__[self.model.conditioning_key]
+                c = {ckey: c, 'pos_x': pos_x, 'pos_y': pos_y}
+
+        else:
+            c = None
+            xc = None
+            if self.use_positional_encodings:
+                pos_x, pos_y = self.compute_latent_shifts(batch)
+                c = {'pos_x': pos_x, 'pos_y': pos_y}
+        out = [z, c]
+        if return_first_stage_outputs:
+            xrec = self.decode_first_stage(z)
+            out.extend([x, xrec])
+        if return_original_cond:
+            out.append(xc)
+        return out
+
+    @torch.no_grad()
+    def decode_first_stage(self,
+                           z,
+                           predict_cids=False,
+                           force_not_quantize=False):
+        if predict_cids:
+            if z.dim() == 4:
+                z = torch.argmax(z.exp(), dim=1).long()
+            z = self.first_stage_model.quantize.get_codebook_entry(z,
+                                                                   shape=None)
+            z = rearrange(z, 'b h w c -> b c h w').contiguous()
+
+        z = 1. / self.scale_factor * z
+
+        if hasattr(self, "split_input_params"):
+            if self.split_input_params["patch_distributed_vq"]:
+                ks = self.split_input_params["ks"]  # eg. (128, 128)
+                stride = self.split_input_params["stride"]  # eg. (64, 64)
+                uf = self.split_input_params["vqf"]
+                bs, nc, h, w = z.shape
+                if ks[0] > h or ks[1] > w:
+                    ks = (min(ks[0], h), min(ks[1], w))
+                    print("reducing Kernel")
+
+                if stride[0] > h or stride[1] > w:
+                    stride = (min(stride[0], h), min(stride[1], w))
+                    print("reducing stride")
+
+                fold, unfold, normalization, weighting = self.get_fold_unfold(
+                    z, ks, stride, uf=uf)
+
+                z = unfold(z)  # (bn, nc * prod(**ks), L)
+                # 1. Reshape to img shape
+                z = z.view((z.shape[0], -1, ks[0], ks[1],
+                            z.shape[-1]))  # (bn, nc, ks[0], ks[1], L )
+
+                # 2. apply model loop over last dim
+                if isinstance(self.first_stage_model, VQModelInterface):
+                    output_list = [
+                        self.first_stage_model.decode(
+                            z[:, :, :, :, i],
+                            force_not_quantize=predict_cids
+                            or force_not_quantize) for i in range(z.shape[-1])
+                    ]
+                else:
+
+                    output_list = [
+                        self.first_stage_model.decode(z[:, :, :, :, i])
+                        for i in range(z.shape[-1])
+                    ]
+
+                o = torch.stack(output_list,
+                                axis=-1)  # # (bn, nc, ks[0], ks[1], L)
+                o = o * weighting
+                # Reverse 1. reshape to img shape
+                o = o.view((o.shape[0], -1,
+                            o.shape[-1]))  # (bn, nc * ks[0] * ks[1], L)
+                # stitch crops together
+                decoded = fold(o)
+                decoded = decoded / normalization  # norm is shape (1, 1, h, w)
+                return decoded
+            else:
+                if isinstance(self.first_stage_model, VQModelInterface):
+                    return self.first_stage_model.decode(
+                        z,
+                        force_not_quantize=predict_cids or force_not_quantize)
+                else:
+                    return self.first_stage_model.decode(z)
+
+        else:
+            if isinstance(self.first_stage_model, VQModelInterface):
+                return self.first_stage_model.decode(
+                    z, force_not_quantize=predict_cids or force_not_quantize)
+            else:
+                return self.first_stage_model.decode(z)
+
+    # same as above but without decorator
+    def differentiable_decode_first_stage(self,
+                                          z,
+                                          predict_cids=False,
+                                          force_not_quantize=False):
+        if predict_cids:
+            if z.dim() == 4:
+                z = torch.argmax(z.exp(), dim=1).long()
+            z = self.first_stage_model.quantize.get_codebook_entry(z,
+                                                                   shape=None)
+            z = rearrange(z, 'b h w c -> b c h w').contiguous()
+
+        z = 1. / self.scale_factor * z
+
+        if hasattr(self, "split_input_params"):
+            if self.split_input_params["patch_distributed_vq"]:
+                ks = self.split_input_params["ks"]  # eg. (128, 128)
+                stride = self.split_input_params["stride"]  # eg. (64, 64)
+                uf = self.split_input_params["vqf"]
+                bs, nc, h, w = z.shape
+                if ks[0] > h or ks[1] > w:
+                    ks = (min(ks[0], h), min(ks[1], w))
+                    print("reducing Kernel")
+
+                if stride[0] > h or stride[1] > w:
+                    stride = (min(stride[0], h), min(stride[1], w))
+                    print("reducing stride")
+
+                fold, unfold, normalization, weighting = self.get_fold_unfold(
+                    z, ks, stride, uf=uf)
+
+                z = unfold(z)  # (bn, nc * prod(**ks), L)
+                # 1. Reshape to img shape
+                z = z.view((z.shape[0], -1, ks[0], ks[1],
+                            z.shape[-1]))  # (bn, nc, ks[0], ks[1], L )
+
+                # 2. apply model loop over last dim
+                if isinstance(self.first_stage_model, VQModelInterface):
+                    output_list = [
+                        self.first_stage_model.decode(
+                            z[:, :, :, :, i],
+                            force_not_quantize=predict_cids
+                            or force_not_quantize) for i in range(z.shape[-1])
+                    ]
+                else:
+
+                    output_list = [
+                        self.first_stage_model.decode(z[:, :, :, :, i])
+                        for i in range(z.shape[-1])
+                    ]
+
+                o = torch.stack(output_list,
+                                axis=-1)  # # (bn, nc, ks[0], ks[1], L)
+                o = o * weighting
+                # Reverse 1. reshape to img shape
+                o = o.view((o.shape[0], -1,
+                            o.shape[-1]))  # (bn, nc * ks[0] * ks[1], L)
+                # stitch crops together
+                decoded = fold(o)
+                decoded = decoded / normalization  # norm is shape (1, 1, h, w)
+                return decoded
+            else:
+                if isinstance(self.first_stage_model, VQModelInterface):
+                    return self.first_stage_model.decode(
+                        z,
+                        force_not_quantize=predict_cids or force_not_quantize)
+                else:
+                    return self.first_stage_model.decode(z)
+
+        else:
+            if isinstance(self.first_stage_model, VQModelInterface):
+                return self.first_stage_model.decode(
+                    z, force_not_quantize=predict_cids or force_not_quantize)
+            else:
+                return self.first_stage_model.decode(z)
+
+    @torch.no_grad()
+    def encode_first_stage(self, x):
+        if hasattr(self, "split_input_params"):
+            if self.split_input_params["patch_distributed_vq"]:
+                ks = self.split_input_params["ks"]  # eg. (128, 128)
+                stride = self.split_input_params["stride"]  # eg. (64, 64)
+                df = self.split_input_params["vqf"]
+                self.split_input_params['original_image_size'] = x.shape[-2:]
+                bs, nc, h, w = x.shape
+                if ks[0] > h or ks[1] > w:
+                    ks = (min(ks[0], h), min(ks[1], w))
+                    print("reducing Kernel")
+
+                if stride[0] > h or stride[1] > w:
+                    stride = (min(stride[0], h), min(stride[1], w))
+                    print("reducing stride")
+
+                fold, unfold, normalization, weighting = self.get_fold_unfold(
+                    x, ks, stride, df=df)
+                z = unfold(x)  # (bn, nc * prod(**ks), L)
+                # Reshape to img shape
+                z = z.view((z.shape[0], -1, ks[0], ks[1],
+                            z.shape[-1]))  # (bn, nc, ks[0], ks[1], L )
+
+                output_list = [
+                    self.first_stage_model.encode(z[:, :, :, :, i])
+                    for i in range(z.shape[-1])
+                ]
+
+                o = torch.stack(output_list, axis=-1)
+                o = o * weighting
+
+                # Reverse reshape to img shape
+                o = o.view((o.shape[0], -1,
+                            o.shape[-1]))  # (bn, nc * ks[0] * ks[1], L)
+                # stitch crops together
+                decoded = fold(o)
+                decoded = decoded / normalization
+                return decoded
+
+            else:
+                return self.first_stage_model.encode(x)
+        else:
+            return self.first_stage_model.encode(x)
+
+    def shared_step(self, batch, **kwargs):
+        x, c = self.get_input(batch, self.first_stage_key)
+        loss = self(x, c)
+        return loss
+
+    def forward(self, x, c, *args, **kwargs):
+        t = torch.randint(0,
+                          self.num_timesteps, (x.shape[0], ),
+                          device=self.device).long()
+        if self.model.conditioning_key is not None:
+            assert c is not None
+            if self.cond_stage_trainable:
+                if self.cond_stage_key == "img_and_caption":
+                    init_img_posterior = self.encode_first_stage(c["init_img"])
+                    c["init_img"] = self.get_first_stage_encoding(
+                        init_img_posterior).detach()
+                    c["caption"] = self.get_learned_conditioning(c["caption"])
+                else:
+                    c = self.get_learned_conditioning(c)
+
+            if self.shorten_cond_schedule:  # TODO: drop this option
+                tc = self.cond_ids[t].to(self.device)
+                c = self.q_sample(x_start=c,
+                                  t=tc,
+                                  noise=torch.randn_like(c.float()))
+        return self.p_losses(x, c, t, *args, **kwargs)
+
+    def _rescale_annotations(self, bboxes,
+                             crop_coordinates):  # TODO: move to dataset
+
+        def rescale_bbox(bbox):
+            x0 = torch.clamp(
+                (bbox[0] - crop_coordinates[0]) / crop_coordinates[2])
+            y0 = torch.clamp(
+                (bbox[1] - crop_coordinates[1]) / crop_coordinates[3])
+            w = min(bbox[2] / crop_coordinates[2], 1 - x0)
+            h = min(bbox[3] / crop_coordinates[3], 1 - y0)
+            return x0, y0, w, h
+
+        return [rescale_bbox(b) for b in bboxes]
+
+    def apply_model(self, x_noisy, t, cond, return_ids=False):
+
+        if self.conditioning_key == "hybrid":
+            new_cond = {
+                'c_concat': [cond["init_img"]],
+                'c_crossattn': [cond["caption"]]
+            }
+            cond = new_cond
+        else:
+            if not isinstance(cond, list):
+                cond = [cond]
+            key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn'
+            cond = {key: cond}
+
+        if hasattr(self, "split_input_params"):
+            assert len(
+                cond) == 1  # todo can only deal with one conditioning atm
+            assert not return_ids
+            ks = self.split_input_params["ks"]  # eg. (128, 128)
+            stride = self.split_input_params["stride"]  # eg. (64, 64)
+
+            h, w = x_noisy.shape[-2:]
+
+            fold, unfold, normalization, weighting = self.get_fold_unfold(
+                x_noisy, ks, stride)
+
+            z = unfold(x_noisy)  # (bn, nc * prod(**ks), L)
+            # Reshape to img shape
+            z = z.view((z.shape[0], -1, ks[0], ks[1],
+                        z.shape[-1]))  # (bn, nc, ks[0], ks[1], L )
+            z_list = [z[:, :, :, :, i] for i in range(z.shape[-1])]
+
+            if self.cond_stage_key in [
+                    "image", "LR_image", "segmentation", 'bbox_img'
+            ] and self.model.conditioning_key:  # todo check for completeness
+                c_key = next(iter(cond.keys()))  # get key
+                c = next(iter(cond.values()))  # get value
+                assert (len(c) == 1
+                        )  # todo extend to list with more than one elem
+                c = c[0]  # get element
+
+                c = unfold(c)
+                c = c.view((c.shape[0], -1, ks[0], ks[1],
+                            c.shape[-1]))  # (bn, nc, ks[0], ks[1], L )
+
+                cond_list = [{
+                    c_key: [c[:, :, :, :, i]]
+                } for i in range(c.shape[-1])]
+
+            elif self.cond_stage_key == 'coordinates_bbox':
+                assert 'original_image_size' in self.split_input_params, 'BoudingBoxRescaling is missing original_image_size'
+
+                # assuming padding of unfold is always 0 and its dilation is always 1
+                n_patches_per_row = int((w - ks[0]) / stride[0] + 1)
+                full_img_h, full_img_w = self.split_input_params[
+                    'original_image_size']
+                # as we are operating on latents, we need the factor from the original image size to the
+                # spatial latent size to properly rescale the crops for regenerating the bbox annotations
+                num_downs = self.first_stage_model.encoder.num_resolutions - 1
+                rescale_latent = 2**(num_downs)
+
+                # get top left postions of patches as conforming for the bbbox tokenizer, therefore we
+                # need to rescale the tl patch coordinates to be in between (0,1)
+                tl_patch_coordinates = [
+                    (rescale_latent * stride[0] *
+                     (patch_nr % n_patches_per_row) / full_img_w,
+                     rescale_latent * stride[1] *
+                     (patch_nr // n_patches_per_row) / full_img_h)
+                    for patch_nr in range(z.shape[-1])
+                ]
+
+                # patch_limits are tl_coord, width and height coordinates as (x_tl, y_tl, h, w)
+                patch_limits = [
+                    (x_tl, y_tl, rescale_latent * ks[0] / full_img_w,
+                     rescale_latent * ks[1] / full_img_h)
+                    for x_tl, y_tl in tl_patch_coordinates
+                ]
+                # patch_values = [(np.arange(x_tl,min(x_tl+ks, 1.)),np.arange(y_tl,min(y_tl+ks, 1.))) for x_tl, y_tl in tl_patch_coordinates]
+
+                # tokenize crop coordinates for the bounding boxes of the respective patches
+                patch_limits_tknzd = [
+                    torch.LongTensor(
+                        self.bbox_tokenizer._crop_encoder(bbox))[None].to(
+                            self.device) for bbox in patch_limits
+                ]  # list of length l with tensors of shape (1, 2)
+                print(patch_limits_tknzd[0].shape)
+                # cut tknzd crop position from conditioning
+                assert isinstance(
+                    cond, dict), 'cond must be dict to be fed into model'
+                cut_cond = cond['c_crossattn'][0][..., :-2].to(self.device)
+                print(cut_cond.shape)
+
+                adapted_cond = torch.stack([
+                    torch.cat([cut_cond, p], dim=1) for p in patch_limits_tknzd
+                ])
+                adapted_cond = rearrange(adapted_cond, 'l b n -> (l b) n')
+                print(adapted_cond.shape)
+                adapted_cond = self.get_learned_conditioning(adapted_cond)
+                print(adapted_cond.shape)
+                adapted_cond = rearrange(adapted_cond,
+                                         '(l b) n d -> l b n d',
+                                         l=z.shape[-1])
+                print(adapted_cond.shape)
+
+                cond_list = [{'c_crossattn': [e]} for e in adapted_cond]
+
+            else:
+                cond_list = [cond for i in range(z.shape[-1])
+                             ]  # Todo make this more efficient
+
+            # apply model by loop over crops
+            output_list = [
+                self.model(z_list[i], t, **cond_list[i])
+                for i in range(z.shape[-1])
+            ]
+            assert not isinstance(
+                output_list[0], tuple
+            )  # todo cant deal with multiple model outputs check this never happens
+
+            o = torch.stack(output_list, axis=-1)
+            o = o * weighting
+            # Reverse reshape to img shape
+            o = o.view(
+                (o.shape[0], -1, o.shape[-1]))  # (bn, nc * ks[0] * ks[1], L)
+            # stitch crops together
+            x_recon = fold(o) / normalization
+
+        else:
+            x_recon = self.model(x_noisy, t, **cond)
+
+        if isinstance(x_recon, tuple) and not return_ids:
+            return x_recon[0]
+        else:
+            return x_recon
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \
+               extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def _prior_bpd(self, x_start):
+        """
+        Get the prior KL term for the variational lower-bound, measured in
+        bits-per-dim.
+        This term can't be optimized, as it only depends on the encoder.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :return: a batch of [N] KL values (in bits), one per batch element.
+        """
+        batch_size = x_start.shape[0]
+        t = torch.tensor([self.num_timesteps - 1] * batch_size,
+                         device=x_start.device)
+        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+        kl_prior = normal_kl(mean1=qt_mean,
+                             logvar1=qt_log_variance,
+                             mean2=0.0,
+                             logvar2=0.0)
+        return mean_flat(kl_prior) / np.log(2.0)
+
+    def p_losses(self, x_start, cond, t, noise=None):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+        # 条件图像也要变成噪声图像
+        # init_img_noise = self.q_sample(x_start=cond["init_img"], t=t, noise=noise)
+        # cond["init_img"] = init_img_noise
+        model_output = self.apply_model(x_noisy, t, cond)
+
+        loss_dict = {}
+        prefix = 'train' if self.training else 'val'
+
+        if self.parameterization == "x0":
+            target = x_start
+        elif self.parameterization == "eps":
+            target = noise
+        else:
+            raise NotImplementedError()
+        loss_simple = self.get_loss(model_output, target,
+                                    mean=False).mean([1, 2, 3])
+        #print(loss_simple)
+
+        loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})
+
+        logvar_t = self.logvar[t].to(self.device)
+        loss = loss_simple / torch.exp(logvar_t) + logvar_t
+        # loss = loss_simple / torch.exp(self.logvar) + self.logvar
+        if self.learn_logvar:
+            loss_dict.update({f'{prefix}/loss_gamma': loss.mean()})
+            loss_dict.update({'logvar': self.logvar.data.mean()})
+
+        loss = self.l_simple_weight * loss.mean()
+
+        loss_vlb = self.get_loss(model_output, target,
+                                 mean=False).mean(dim=(1, 2, 3))
+        loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean()
+        loss_dict.update({f'{prefix}/loss_vlb': loss_vlb})
+        loss += (self.original_elbo_weight * loss_vlb)
+        loss_dict.update({f'{prefix}/loss': loss})
+
+        return loss, loss_dict
+
+    def p_mean_variance(self,
+                        x,
+                        c,
+                        t,
+                        clip_denoised: bool,
+                        return_codebook_ids=False,
+                        quantize_denoised=False,
+                        return_x0=False,
+                        score_corrector=None,
+                        corrector_kwargs=None):
+        t_in = t
+        model_out = self.apply_model(x,
+                                     t_in,
+                                     c,
+                                     return_ids=return_codebook_ids)
+
+        if score_corrector is not None:
+            assert self.parameterization == "eps"
+            model_out = score_corrector.modify_score(self, model_out, x, t, c,
+                                                     **corrector_kwargs)
+
+        if return_codebook_ids:
+            model_out, logits = model_out
+
+        if self.parameterization == "eps":
+            x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
+        elif self.parameterization == "x0":
+            x_recon = model_out
+        else:
+            raise NotImplementedError()
+
+        if clip_denoised:
+            x_recon.clamp_(-1., 1.)
+        if quantize_denoised:
+            x_recon, _, [_, _,
+                         indices] = self.first_stage_model.quantize(x_recon)
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(
+            x_start=x_recon, x_t=x, t=t)
+        if return_codebook_ids:
+            return model_mean, posterior_variance, posterior_log_variance, logits
+        elif return_x0:
+            return model_mean, posterior_variance, posterior_log_variance, x_recon
+        else:
+            return model_mean, posterior_variance, posterior_log_variance
+
+    @torch.no_grad()
+    def p_sample(self,
+                 x,
+                 c,
+                 t,
+                 clip_denoised=False,
+                 repeat_noise=False,
+                 return_codebook_ids=False,
+                 quantize_denoised=False,
+                 return_x0=False,
+                 temperature=1.,
+                 noise_dropout=0.,
+                 score_corrector=None,
+                 corrector_kwargs=None):
+        b, *_, device = *x.shape, x.device
+        outputs = self.p_mean_variance(x=x,
+                                       c=c,
+                                       t=t,
+                                       clip_denoised=clip_denoised,
+                                       return_codebook_ids=return_codebook_ids,
+                                       quantize_denoised=quantize_denoised,
+                                       return_x0=return_x0,
+                                       score_corrector=score_corrector,
+                                       corrector_kwargs=corrector_kwargs)
+        if return_codebook_ids:
+            raise DeprecationWarning("Support dropped.")
+            model_mean, _, model_log_variance, logits = outputs
+        elif return_x0:
+            model_mean, _, model_log_variance, x0 = outputs
+        else:
+            model_mean, _, model_log_variance = outputs
+
+        noise = noise_like(x.shape, device, repeat_noise) * temperature
+        if noise_dropout > 0.:
+            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+        # no noise when t == 0
+        nonzero_mask = (1 - (t == 0).float()).reshape(
+            b, *((1, ) * (len(x.shape) - 1)))
+
+        if return_codebook_ids:
+            return model_mean + nonzero_mask * (
+                0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1)
+        if return_x0:
+            return model_mean + nonzero_mask * (
+                0.5 * model_log_variance).exp() * noise, x0
+        else:
+            return model_mean + nonzero_mask * (
+                0.5 * model_log_variance).exp() * noise
+
+    @torch.no_grad()
+    def progressive_denoising(self,
+                              cond,
+                              shape,
+                              verbose=True,
+                              callback=None,
+                              quantize_denoised=False,
+                              img_callback=None,
+                              mask=None,
+                              x0=None,
+                              temperature=1.,
+                              noise_dropout=0.,
+                              score_corrector=None,
+                              corrector_kwargs=None,
+                              batch_size=None,
+                              x_T=None,
+                              start_T=None,
+                              log_every_t=None):
+        if not log_every_t:
+            log_every_t = self.log_every_t
+        timesteps = self.num_timesteps
+        if batch_size is not None:
+            b = batch_size if batch_size is not None else shape[0]
+            shape = [batch_size] + list(shape)
+        else:
+            b = batch_size = shape[0]
+        if x_T is None:
+            img = torch.randn(shape, device=self.device)
+        else:
+            img = x_T
+        intermediates = []
+        if cond is not None:
+            if isinstance(cond, dict):
+                cond = {
+                    key:
+                    cond[key][:batch_size] if not isinstance(cond[key], list)
+                    else list(map(lambda x: x[:batch_size], cond[key]))
+                    for key in cond
+                }
+            else:
+                cond = [c[:batch_size] for c in cond] if isinstance(
+                    cond, list) else cond[:batch_size]
+
+        if start_T is not None:
+            timesteps = min(timesteps, start_T)
+        iterator = tqdm(reversed(range(0, timesteps)),
+                        desc='Progressive Generation',
+                        total=timesteps) if verbose else reversed(
+                            range(0, timesteps))
+        if type(temperature) == float:
+            temperature = [temperature] * timesteps
+
+        for i in iterator:
+            ts = torch.full((b, ), i, device=self.device, dtype=torch.long)
+            if self.shorten_cond_schedule:
+                assert self.model.conditioning_key != 'hybrid'
+                tc = self.cond_ids[ts].to(cond.device)
+                cond = self.q_sample(x_start=cond,
+                                     t=tc,
+                                     noise=torch.randn_like(cond))
+
+            img, x0_partial = self.p_sample(
+                img,
+                cond,
+                ts,
+                clip_denoised=self.clip_denoised,
+                quantize_denoised=quantize_denoised,
+                return_x0=True,
+                temperature=temperature[i],
+                noise_dropout=noise_dropout,
+                score_corrector=score_corrector,
+                corrector_kwargs=corrector_kwargs)
+            if mask is not None:
+                assert x0 is not None
+                img_orig = self.q_sample(x0, ts)
+                img = img_orig * mask + (1. - mask) * img
+
+            if i % log_every_t == 0 or i == timesteps - 1:
+                intermediates.append(x0_partial)
+            if callback: callback(i)
+            if img_callback: img_callback(img, i)
+        return img, intermediates
+
+    @torch.no_grad()
+    def p_sample_loop(self,
+                      cond,
+                      shape,
+                      return_intermediates=False,
+                      x_T=None,
+                      verbose=True,
+                      callback=None,
+                      timesteps=None,
+                      quantize_denoised=False,
+                      mask=None,
+                      x0=None,
+                      img_callback=None,
+                      start_T=None,
+                      log_every_t=None):
+
+        if not log_every_t:
+            log_every_t = self.log_every_t
+        device = self.betas.device
+        b = shape[0]
+        if x_T is None:
+            img = torch.randn(shape, device=device)
+        else:
+            img = x_T
+
+        intermediates = [img]
+        if timesteps is None:
+            timesteps = self.num_timesteps
+
+        if start_T is not None:
+            timesteps = min(timesteps, start_T)
+        iterator = tqdm(
+            reversed(range(0, timesteps)), desc='Sampling t',
+            total=timesteps) if verbose else reversed(range(0, timesteps))
+
+        if mask is not None:
+            assert x0 is not None
+            assert x0.shape[2:3] == mask.shape[2:
+                                               3]  # spatial size has to match
+
+        for i in iterator:
+            ts = torch.full((b, ), i, device=device, dtype=torch.long)
+            if self.shorten_cond_schedule:
+                assert self.model.conditioning_key != 'hybrid'
+                tc = self.cond_ids[ts].to(cond.device)
+                cond = self.q_sample(x_start=cond,
+                                     t=tc,
+                                     noise=torch.randn_like(cond))
+
+            img = self.p_sample(img,
+                                cond,
+                                ts,
+                                clip_denoised=self.clip_denoised,
+                                quantize_denoised=quantize_denoised)
+            if mask is not None:
+                img_orig = self.q_sample(x0, ts)
+                img = img_orig * mask + (1. - mask) * img
+
+            if i % log_every_t == 0 or i == timesteps - 1:
+                intermediates.append(img)
+            if callback: callback(i)
+            if img_callback: img_callback(img, i)
+
+        if return_intermediates:
+            return img, intermediates
+        return img
+
+    @torch.no_grad()
+    def sample(self,
+               cond,
+               batch_size=16,
+               return_intermediates=False,
+               x_T=None,
+               verbose=True,
+               timesteps=None,
+               quantize_denoised=False,
+               mask=None,
+               x0=None,
+               shape=None,
+               **kwargs):
+        if shape is None:
+            shape = (batch_size, self.channels, self.image_size,
+                     self.image_size)
+        if cond is not None:
+            if isinstance(cond, dict):
+                cond = {
+                    key:
+                    cond[key][:batch_size] if not isinstance(cond[key], list)
+                    else list(map(lambda x: x[:batch_size], cond[key]))
+                    for key in cond
+                }
+            else:
+                cond = [c[:batch_size] for c in cond] if isinstance(
+                    cond, list) else cond[:batch_size]
+        return self.p_sample_loop(cond,
+                                  shape,
+                                  return_intermediates=return_intermediates,
+                                  x_T=x_T,
+                                  verbose=verbose,
+                                  timesteps=timesteps,
+                                  quantize_denoised=quantize_denoised,
+                                  mask=mask,
+                                  x0=x0)
+
+    @torch.no_grad()
+    def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
+
+        if ddim:
+            ddim_sampler = DDIMSampler(self)
+            shape = (self.channels, self.image_size, self.image_size)
+            samples, intermediates = ddim_sampler.sample(ddim_steps,
+                                                         batch_size,
+                                                         shape,
+                                                         cond,
+                                                         verbose=False,
+                                                         **kwargs)
+
+        else:
+            samples, intermediates = self.sample(cond=cond,
+                                                 batch_size=batch_size,
+                                                 return_intermediates=True,
+                                                 **kwargs)
+
+        return samples, intermediates
+
+    @torch.no_grad()
+    def log_images(self,
+                   batch,
+                   N=8,
+                   n_row=4,
+                   sample=True,
+                   ddim_steps=200,
+                   ddim_eta=1.,
+                   return_keys=None,
+                   quantize_denoised=True,
+                   inpaint=True,
+                   plot_denoise_rows=False,
+                   plot_progressive_rows=True,
+                   plot_diffusion_rows=True,
+                   **kwargs):
+
+        use_ddim = ddim_steps is not None
+
+        log = dict()
+        z, c, x, xrec, xc = self.get_input(batch,
+                                           self.first_stage_key,
+                                           return_first_stage_outputs=True,
+                                           force_c_encode=True,
+                                           return_original_cond=True,
+                                           bs=N)
+        N = min(x.shape[0], N)
+        n_row = min(x.shape[0], n_row)
+        log["inputs"] = x
+        #log["reconstruction"] = xrec
+        if self.conditioning_key == 'hybrid':
+            # log文本和初始图像
+            log["init_img"] = batch["img_and_caption"]["init_img"]
+            xc = log_txt_as_img((x.shape[2], x.shape[3]),
+                                batch["img_and_caption"]["caption"])
+            log["conditioning"] = xc
+
+        if self.model.conditioning_key is not None:
+            if hasattr(self.cond_stage_model, "decode"):
+                xc = self.cond_stage_model.decode(c)
+                log["conditioning"] = xc
+            elif self.cond_stage_key in ["caption"]:
+                xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["caption"])
+                log["conditioning"] = xc
+            elif self.cond_stage_key == 'class_label':
+                xc = log_txt_as_img((x.shape[2], x.shape[3]),
+                                    batch["human_label"])
+                log['conditioning'] = xc
+            elif isimage(xc):
+                log["conditioning"] = xc
+            if ismap(xc):
+                log["original_conditioning"] = self.to_rgb(xc)
+
+        if plot_diffusion_rows:
+            # get diffusion row
+            diffusion_row = list()
+            z_start = z[:n_row]
+            for t in range(self.num_timesteps):
+                if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+                    t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
+                    t = t.to(self.device).long()
+                    noise = torch.randn_like(z_start)
+                    z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
+                    diffusion_row.append(self.decode_first_stage(z_noisy))
+
+            diffusion_row = torch.stack(
+                diffusion_row)  # n_log_step, n_row, C, H, W
+            diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
+            diffusion_grid = rearrange(diffusion_grid,
+                                       'b n c h w -> (b n) c h w')
+            diffusion_grid = make_grid(diffusion_grid,
+                                       nrow=diffusion_row.shape[0])
+            #log["diffusion_row"] = diffusion_grid
+
+        if sample:
+            # get denoise row
+            with self.ema_scope("Plotting"):
+                samples, z_denoise_row = self.sample_log(cond=c,
+                                                         batch_size=N,
+                                                         ddim=use_ddim,
+                                                         ddim_steps=ddim_steps,
+                                                         eta=ddim_eta)
+                # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
+            x_samples = self.decode_first_stage(samples)
+            log["samples"] = x_samples
+            if plot_denoise_rows:
+                denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
+                #log["denoise_row"] = denoise_grid
+
+            if quantize_denoised and not isinstance(
+                    self.first_stage_model, AutoencoderKL) and not isinstance(
+                        self.first_stage_model, IdentityFirstStage):
+                # also display when quantizing x0 while sampling
+                with self.ema_scope("Plotting Quantized Denoised"):
+                    samples, z_denoise_row = self.sample_log(
+                        cond=c,
+                        batch_size=N,
+                        ddim=use_ddim,
+                        ddim_steps=ddim_steps,
+                        eta=ddim_eta,
+                        quantize_denoised=True)
+                    # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True,
+                    #                                      quantize_denoised=True)
+                x_samples = self.decode_first_stage(samples.to(self.device))
+                log["samples_x0_quantized"] = x_samples
+
+            if inpaint:
+                # make a simple center square
+                b, h, w = z.shape[0], z.shape[2], z.shape[3]
+                mask = torch.ones(N, h, w).to(self.device)
+                # zeros will be filled in
+                mask[:, h // 4:3 * h // 4, w // 4:3 * w // 4] = 0.
+                mask = mask[:, None, ...]
+                with self.ema_scope("Plotting Inpaint"):
+
+                    samples, _ = self.sample_log(cond=c,
+                                                 batch_size=N,
+                                                 ddim=use_ddim,
+                                                 eta=ddim_eta,
+                                                 ddim_steps=ddim_steps,
+                                                 x0=z[:N],
+                                                 mask=mask)
+                x_samples = self.decode_first_stage(samples.to(self.device))
+                #log["samples_inpainting"] = x_samples
+                #log["mask"] = mask
+
+                # outpaint
+                with self.ema_scope("Plotting Outpaint"):
+                    samples, _ = self.sample_log(cond=c,
+                                                 batch_size=N,
+                                                 ddim=use_ddim,
+                                                 eta=ddim_eta,
+                                                 ddim_steps=ddim_steps,
+                                                 x0=z[:N],
+                                                 mask=mask)
+                x_samples = self.decode_first_stage(samples.to(self.device))
+                #log["samples_outpainting"] = x_samples
+
+        if plot_progressive_rows:
+            with self.ema_scope("Plotting Progressives"):
+                img, progressives = self.progressive_denoising(
+                    c,
+                    shape=(self.channels, self.image_size, self.image_size),
+                    batch_size=N)
+            prog_row = self._get_denoise_row_from_list(
+                progressives, desc="Progressive Generation")
+            log["progressive_row"] = prog_row
+
+        if return_keys:
+            if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
+                return log
+            else:
+                return {key: log[key] for key in return_keys}
+        return log
+
+    def configure_optimizers(self):
+
+        lr = self.learning_rate
+        params = list(self.model.parameters())
+
+        if self.cond_stage_trainable:
+            print(
+                f"{self.__class__.__name__}: Also optimizing conditioner params!"
+            )
+            params = params + list(self.cond_stage_model.parameters())
+        if self.learn_logvar:
+            print('Diffusion model optimizing logvar')
+            params.append(self.logvar)
+        #opt = torch.optim.AdamW(params, lr=lr)
+        opt = torch.optim.Adam(params, lr=lr)
+        if self.use_scheduler:
+            assert 'target' in self.scheduler_config
+            scheduler = instantiate_from_config(self.scheduler_config)
+            print("Setting up LambdaLR scheduler...")
+            scheduler = [{
+                'scheduler':
+                LambdaLR(opt, lr_lambda=scheduler.schedule),
+                'interval':
+                'step',
+                'frequency':
+                1
+            }]
+            return [opt], scheduler
+        return opt
+
+    @torch.no_grad()
+    def to_rgb(self, x):
+        x = x.float()
+        if not hasattr(self, "colorize"):
+            self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x)
+        x = nn.functional.conv2d(x, weight=self.colorize)
+        x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
+        return x
+
+
+class DiffusionWrapper(torch.nn.Module):
+
+    def __init__(self, diff_model_config, conditioning_key):
+        super().__init__()
+        self.diffusion_model = instantiate_from_config(diff_model_config)
+        self.conditioning_key = conditioning_key
+        assert self.conditioning_key in [
+            None, 'concat', 'crossattn', 'hybrid', 'adm'
+        ]
+
+    def forward(self, x, t, c_concat: list = None, c_crossattn: list = None):
+        if self.conditioning_key is None:
+            out = self.diffusion_model(x, t)
+        elif self.conditioning_key == 'concat':
+            xc = torch.cat([x] + c_concat, dim=1)
+            out = self.diffusion_model(xc, t)
+        elif self.conditioning_key == 'crossattn':
+            cc = torch.cat(c_crossattn, 1)
+            out = self.diffusion_model(x, t, context=cc)
+        elif self.conditioning_key == 'hybrid':
+            xc = torch.cat([x] + c_concat, dim=1)
+            #xc = x + c_concat[0]
+            cc = torch.cat(c_crossattn, 1)
+            out = self.diffusion_model(xc, t, context=cc)
+        elif self.conditioning_key == 'adm':
+            cc = c_crossattn[0]
+            out = self.diffusion_model(x, t, y=cc)
+        else:
+            raise NotImplementedError()
+
+        return out
+
+
+class DiagonalGaussianDistribution(object):
+
+    def __init__(self, parameters, deterministic=False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(
+                self.mean).to(device=self.parameters.device)
+
+    def sample(self):
+        x = self.mean + self.std * torch.randn(
+            self.mean.shape).to(device=self.parameters.device)
+        return x
+
+    def kl(self, other=None):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        else:
+            if other is None:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+                    dim=[1, 2, 3])
+            else:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean - other.mean, 2) / other.var +
+                    self.var / other.var - 1.0 - self.logvar + other.logvar,
+                    dim=[1, 2, 3])
+
+    def nll(self, sample, dims=[1, 2, 3]):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(logtwopi + self.logvar +
+                               torch.pow(sample - self.mean, 2) / self.var,
+                               dim=dims)
+
+    def mode(self):
+        return self.mean
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12
+    Compute the KL divergence between two gaussians.
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, torch.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for torch.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor)
+        for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (-1.0 + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) +
+                  ((mean1 - mean2)**2) * torch.exp(-logvar2))
diff --git a/flagai/model/mm/flag_clip.py b/flagai/model/mm/flag_clip.py
new file mode 100755
index 00000000..cc94c58d
--- /dev/null
+++ b/flagai/model/mm/flag_clip.py
@@ -0,0 +1,454 @@
+from time import clock_settime
+from transformers.models.clip.modeling_clip import *
+import torch.nn as nn
+import torch
+from transformers.models.clip.modeling_clip import CLIPOutput
+from transformers import CLIPProcessor
+import os
+from flagai.model.base_model import BaseModel
+
+from .modeling_berts import BertSeriesConfig, RobertaSeriesConfig, BertSeriesModelWithTransformation, RobertaSeriesModelWithTransformation
+from transformers.models.bert.tokenization_bert import BertTokenizer
+
+STUDENT_CONFIG_DICT = {
+    'hfl/chinese-roberta-wwm-ext': BertSeriesConfig,
+    'hfl/chinese-roberta-wwm-ext-large': BertSeriesConfig,
+    'xlm-roberta-large': RobertaSeriesConfig,
+    'xlm-roberta-base': RobertaSeriesConfig,
+    'bert-base-uncased': BertSeriesConfig,
+    'bert': BertSeriesConfig,
+    'xlm-roberta': RobertaSeriesConfig,
+    'clip_text_model': BertSeriesConfig,
+}
+
+STUDENT_MODEL_DICT = {
+    'hfl/chinese-roberta-wwm-ext': BertSeriesModelWithTransformation,
+    'hfl/chinese-roberta-wwm-ext-large': BertSeriesModelWithTransformation,
+    'xlm-roberta-large': RobertaSeriesModelWithTransformation,
+    'xlm-roberta-base': RobertaSeriesModelWithTransformation,
+    'bert': BertSeriesModelWithTransformation,
+    'xlm-roberta': RobertaSeriesModelWithTransformation,
+    'clip_text_model': BertSeriesModelWithTransformation,
+}
+
+
+@dataclass
+class OursCLIPOutput(CLIPOutput):
+
+    def to_tuple(self) -> Tuple[Any]:
+        return tuple(
+            self[k] if k not in ["text_model_output", "vision_model_output"
+                                 ] else getattr(self, k).to_tuple()
+            for k in self.keys())
+
+
+class CHCLIPProcess(CLIPProcessor):
+    tokenizer_class = ("XLMRobertaTokenizer", "XLMRobertaTokenizerFast")
+
+    # tokenizer_class = ("BertTokenizer", "BertTokenizerFast")
+    def __init__(self, feature_extractor, tokenizer):
+        super().__init__(feature_extractor, tokenizer)
+        # ## in some cases, we need to switch to different tokenizer
+        # if tokenizer.vocab_size!=250002:
+        #     print("tokenizer is not XLM-R, switched to BertTokenizer.")
+        #     self.tokenizer = BertTokenizer.from_pretrained(tokenizer.name_or_path)
+
+
+class CHCLIPProcessBert(CLIPProcessor):
+    # tokenizer_class = ("XLMRobertaTokenizer","XLMRobertaTokenizerFast")
+    tokenizer_class = ("BertTokenizer", "BertTokenizerFast")
+
+    def __init__(self, feature_extractor, tokenizer):
+        super().__init__(feature_extractor, tokenizer)
+        # ## in some cases, we need to switch to different tokenizer
+        # if tokenizer.vocab_size!=250002:
+        #     print("tokenizer is not XLM-R, switched to BertTokenizer.")
+        #     self.tokenizer = BertTokenizer.from_pretrained(tokenizer.name_or_path)
+
+
+class ChCLIPConfig(CLIPConfig):
+
+    def __init__(self,
+                 text_model_name=None,
+                 vision_model_name=None,
+                 text_config_dict=None,
+                 vision_config_dict=None,
+                 projection_dim=512,
+                 logit_scale_init_value=2.6592,
+                 num_layers=3,
+                 variant='invert',
+                 **kwargs):
+        super().__init__(text_config_dict, vision_config_dict, projection_dim,
+                         logit_scale_init_value, **kwargs)
+        if text_config_dict is None:
+            text_config_dict = {}
+        # when reload the config from local, we need name to select which class should be instanced.
+        self.text_config = STUDENT_CONFIG_DICT[
+            kwargs['text_config']['model_type']](**kwargs.pop('text_config'))
+        self.num_layers = num_layers
+        self.text_model_name = text_model_name
+        self.vision_model_name = vision_model_name
+        self.variant = variant
+
+
+class ChineseCLIPHF(CLIPPreTrainedModel):
+    config_class = ChCLIPConfig
+
+    def __init__(self, config: ChCLIPConfig, clip_model=None):
+        super().__init__(config)
+
+        if not isinstance(config.vision_config, CLIPVisionConfig):
+            raise ValueError(
+                "config.vision_config is expected to be of type CLIPVisionConfig but is of type"
+                f" {type(config.vision_config)}.")
+
+        text_config = config.text_config
+        vision_config = config.vision_config
+
+        self.projection_dim = config.projection_dim
+        self.text_embed_dim = text_config.project_dim
+        self.vision_embed_dim = vision_config.hidden_size
+
+        self.text_model = STUDENT_MODEL_DICT[text_config.model_type](
+            text_config)
+
+        self.vision_model = CLIPVisionTransformer(vision_config)
+
+        self.visual_projection = nn.Linear(self.vision_embed_dim,
+                                           self.projection_dim,
+                                           bias=False)
+        self.text_projection = nn.Linear(self.text_embed_dim,
+                                         self.projection_dim,
+                                         bias=False)
+        self.logit_scale = nn.Parameter(
+            torch.ones([]) * self.config.logit_scale_init_value)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_text_features(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        token_type_ids=None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> torch.FloatTensor:
+        r"""
+        Returns:
+            text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by
+            applying the projection layer to the pooled output of [`CLIPTextModel`].
+
+        Examples:
+
+        ```python
+        >>> from transformers import CLIPTokenizer, CLIPModel
+
+        >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+        >>> tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+
+        >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
+        >>> text_features = model.get_text_features(**inputs)
+        ```"""
+        # Use CLIP model's config for some fields (if specified) instead of those of vision & text components.
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else
+                                self.config.output_hidden_states)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        pooled_output = text_outputs['pooler_output']
+        text_features = self.text_projection(pooled_output)
+
+        return text_features
+
+    def get_text_features_diffusion(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        token_type_ids=None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> torch.FloatTensor:
+        r"""
+        Returns:
+            text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by
+            applying the projection layer to the pooled output of [`CLIPTextModel`].
+
+        Examples:
+
+        ```python
+        >>> from transformers import CLIPTokenizer, CLIPModel
+
+        >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+        >>> tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+
+        >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
+        >>> text_features = model.get_text_features(**inputs)
+        ```"""
+        # Use CLIP model's config for some fields (if specified) instead of those of vision & text components.
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else
+                                self.config.output_hidden_states)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        text_features = text_outputs['projection_state']
+        # text_features = self.text_projection(pooled_output)
+
+        return text_features
+
+    def tokenize(self, texts, tokenizer, context_length: int = 77):
+
+        if isinstance(texts, str):
+            texts = [texts]
+
+        all_tokens = []
+        for text in texts:
+            all_tokens.append([tokenizer.vocab['
']] +
+                              tokenizer.convert_tokens_to_ids(
+                                  tokenizer.tokenize(text))[:context_length -
+                                                            2] +
+                              [tokenizer.vocab['
']])
+
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+        for i, tokens in enumerate(all_tokens):
+            assert len(tokens) <= context_length
+            result[i, :len(tokens)] = torch.tensor(tokens)
+
+        return result
+
+    def encode(self,
+               text,
+               tokenizer,
+               padding="max_length",
+               truncation=True,
+               max_length=77):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        text = tokenizer(text,
+                         truncation=True,
+                         max_length=77,
+                         return_length=False,
+                         return_overflowing_tokens=False,
+                         padding="max_length",
+                         return_tensors="pt")
+        text["input_ids"] = torch.tensor(text["input_ids"]).to(device)
+        text["attention_mask"] = torch.tensor(
+            text['attention_mask']).to(device)
+        # text = torch.tensor(text).to(device
+        features = self.get_text_features_diffusion(**text)
+        return features
+
+    def get_text_encoder_out(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        token_type_ids=None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> torch.FloatTensor:
+        # Use CLIP model's config for some fields (if specified) instead of those of vision & text components.
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else
+                                self.config.output_hidden_states)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_out = text_outputs["sequence_out"]
+
+        return sequence_out
+
+    @add_start_docstrings_to_model_forward(CLIP_VISION_INPUTS_DOCSTRING)
+    def get_image_features(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> torch.FloatTensor:
+        r"""
+        Returns:
+            image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by
+            applying the projection layer to the pooled output of [`CLIPVisionModel`].
+
+        Examples:
+
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import CLIPProcessor, CLIPModel
+
+        >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+        >>> processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> inputs = processor(images=image, return_tensors="pt")
+
+        >>> image_features = model.get_image_features(**inputs)
+        ```"""
+        # Use CLIP model's config for some fields (if specified) instead of those of vision & text components.
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else
+                                self.config.output_hidden_states)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        pooled_output = vision_outputs[1]  # pooled_output
+        image_features = self.visual_projection(pooled_output)
+
+        return image_features
+
+    @add_start_docstrings_to_model_forward(CLIP_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=CLIPOutput, config_class=CLIPConfig)
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        token_type_ids=None,
+        return_loss: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, CLIPOutput]:
+        r"""
+        Returns:
+
+        Examples:
+
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import CLIPProcessor, CLIPModel
+
+        >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+        >>> processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> inputs = processor(
+        ...     text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True
+        ... )
+
+        >>> outputs = model(**inputs)
+        >>> logits_per_image = outputs.logits_per_image  # this is the image-text similarity score
+        >>> probs = logits_per_image.softmax(dim=1)  # we can take the softmax to get the label probabilities
+        ```"""
+        # Use CLIP model's config for some fields (if specified) instead of those of vision & text components.
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else
+                                self.config.output_hidden_states)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        image_embeds = vision_outputs[1]
+        image_embeds = self.visual_projection(image_embeds)
+
+        # text_embeds = self.text_model.get_text_embeds(text_outputs['pooler_output'],clip_outputs[1])
+        text_embeds = text_outputs['pooler_output']
+        text_embeds = self.text_projection(text_embeds)
+
+        # normalized features
+        image_embeds = image_embeds / image_embeds.norm(
+            p=2, dim=-1, keepdim=True)
+        text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_text = torch.matmul(text_embeds,
+                                       image_embeds.t()) * logit_scale
+        logits_per_image = logits_per_text.T
+
+        loss = clip_loss(logits_per_text)
+
+        return {
+            "loss": loss,
+            "logits_per_image": logits_per_image,
+            "logits_per_text": logits_per_text,
+            "text_embeds": text_embeds,
+            "image_embeds": image_embeds,
+            "text_model_output": text_outputs,
+            "vision_model_output": vision_outputs,
+            "logits": torch.cat([image_embeds, text_embeds], dim=-1)
+        }
+
+
+class ChineseCLIP(BaseModel):
+
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
+
+    @classmethod
+    def from_pretrain(cls,
+                      download_path='./checkpoints/',
+                      model_name='RoBERTa-base-ch',
+                      only_download_config=False,
+                      device="cpu",
+                      **kwargs):
+        super().download(download_path, model_name)
+        pretrained_model_name_or_path = os.path.join(download_path, model_name)
+        print(pretrained_model_name_or_path)
+        return ChineseCLIPHF.from_pretrained(pretrained_model_name_or_path)
diff --git a/flagai/model/mm/lm/__init__.py b/flagai/model/mm/lm/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/flagai/model/mm/lm/cn_clip.py b/flagai/model/mm/lm/cn_clip.py
new file mode 100755
index 00000000..85bafe6d
--- /dev/null
+++ b/flagai/model/mm/lm/cn_clip.py
@@ -0,0 +1,68 @@
+import torch
+import torch.nn as nn
+from flagai.model.mm.clip_cn.model import CLIP
+from flagai.model.base_model import BaseModel
+
+
+class LayerNorm(nn.Module):
+
+    def __init__(self, hidden_size, eps=1e-12):
+        """Construct a layernorm module in the TF style (epsilon inside the square root).
+        """
+        super(LayerNorm, self).__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.bias = nn.Parameter(torch.zeros(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x):
+        u = x.mean(-1, keepdim=True)
+        s = (x - u).pow(2).mean(-1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
+        return self.weight * x + self.bias
+
+
+class CN_CLIP(BaseModel):
+
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
+        max_length = config["max_length"]
+        normalize = config["normalize"]
+        self.tokenizer = kwargs.get("tokenizer")
+
+        self.device = torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu")
+        self.max_length = max_length
+
+        self.model = CLIP(config, **kwargs)
+
+        # 冻结权重
+        # for param in self.model.parameters():
+        #     param.requires_grad = False
+
+        self.normalize = normalize
+        self.layer_norm1 = LayerNorm(768)
+        self.layer_norm2 = LayerNorm(768)
+
+        self.proj = nn.Linear(768, 768)
+
+    def forward(self, text):
+        text = self.tokenizer.tokenize(text).to(self.device)
+        z = self.model.encode_text(text)
+        z = self.layer_norm1(z)
+        z = self.proj(z)
+        z = self.layer_norm2(z)
+
+        return z
+
+    def encode(self, text):
+        z = self(text)
+        if z.ndim == 2:
+            z = z[:, None, :]
+        return z
+
+    def load_weights(self, checkpoint_path):
+        checkpoint = torch.load(checkpoint_path)
+        sd = checkpoint["state_dict"]
+        if next(iter(sd.items()))[0].startswith('module'):
+            sd = {k[len('module.'):]: v for k, v in sd.items()}
+        self.model.load_state_dict(sd)
diff --git a/flagai/model/mm/lm/dcn_clip.py b/flagai/model/mm/lm/dcn_clip.py
new file mode 100755
index 00000000..253ba4bd
--- /dev/null
+++ b/flagai/model/mm/lm/dcn_clip.py
@@ -0,0 +1,155 @@
+import torch
+import torch.nn as nn
+from functools import partial
+import clip
+from einops import rearrange, repeat
+import transformers
+from models.lm.lgs.modeling_chclip_new import ChineseCLIP
+from models.lm.lgs.modeling_chclip_new import CHCLIPProcess
+import json
+
+class AbstractEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+
+class LayerNorm(nn.Module):
+
+    def __init__(self, hidden_size, eps=1e-12):
+        """Construct a layernorm module in the TF style (epsilon inside the square root).
+        """
+        super(LayerNorm, self).__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.bias = nn.Parameter(torch.zeros(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x):
+        u = x.mean(-1, keepdim=True)
+        s = (x - u).pow(2).mean(-1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
+        return self.weight * x + self.bias
+
+class DCNCLIP3M768(AbstractEncoder):
+    def __init__(self, version='ViT-b/14', device="cuda", max_length=77, n_repeat=1, normalize=True):
+        super().__init__()
+        self.device = device
+        self.max_length = max_length
+        self.ch_clip_model = ChineseCLIP.from_pretrained("/sharefs/baai-mrnd/zch/bs-bs-3m")
+        self.ch_clip_model = self.ch_clip_model.eval()
+        for param in self.ch_clip_model.parameters():
+            param.requires_grad = False
+        
+        self.processor = CHCLIPProcess.from_pretrained("/sharefs/baai-mrnd/zch/bs-bs-3m")
+        self.tokenizer = self.processor.tokenizer
+        self.text_encoder = self.ch_clip_model.text_model
+        self.layer_norm = LayerNorm(768)
+
+    def forward(self, text):
+        tokens = self.tokenizer(text,return_tensors='pt',padding=True).to(self.device)
+        #tokens = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=False,
+        #                                 return_overflowing_tokens=False, padding="max_length", return_tensors="pt").to(self.device)
+        z = self.text_encoder(**tokens)['hidden_states'][-1]
+        z = self.layer_norm(z)
+        
+        return z
+
+    def encode(self, text):
+        z = self(text)
+        if z.ndim==2:
+            z = z[:, None, :]
+        return z
+    
+
+class DCNCLIP30M1024(AbstractEncoder):
+
+    def __init__(self, version='ViT-L/14', device="cuda", max_length=50, n_repeat=1, normalize=True):
+        super().__init__()
+        self.device = device
+        self.max_length = max_length
+        self.ch_clip_model = ChineseCLIP.from_pretrained("/sharefs/baai-mrnd/czz/lg30M")
+        self.ch_clip_model = self.ch_clip_model.eval()
+
+        for param in self.ch_clip_model.parameters():
+            param.requires_grad = False
+        
+        self.processor = CHCLIPProcess.from_pretrained("/sharefs/baai-mrnd/czz/lg30M")
+        self.tokenizer = self.processor.tokenizer
+
+        self.text_encoder = self.ch_clip_model.text_model
+        # self.normalize = normalize
+        # self.layer_norm1 = LayerNorm(768)
+        # self.layer_norm2 = LayerNorm(768)
+
+        # self.proj = nn.Linear(1024,768)
+        
+
+
+    def forward(self, text):
+        tokens = self.tokenizer(text,return_tensors='pt',truncation=True, 
+        max_length=self.max_length, padding=True).to(self.device)
+        #tokens = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=False,
+        #                                 return_overflowing_tokens=False, padding="max_length", return_tensors="pt").to(self.device)
+
+        z = self.text_encoder(**tokens)['mimic_states']
+        #  进行维度映射
+        # z = self.layer_norm1(z)
+
+        # z = self.proj(z)
+        # z = self.layer_norm2(z)
+
+        return z
+
+    def encode(self, text):
+        z = self(text)
+        if z.ndim==2:
+            z = z[:, None, :]
+        return z
+
+class DCNCLIP3M1024(AbstractEncoder):
+
+    def __init__(self, version='ViT-L/14', device="cuda", max_length=77, n_repeat=1, normalize=True):
+        super().__init__()
+        self.device = device
+        self.max_length = max_length
+        self.ch_clip_model = ChineseCLIP.from_pretrained("/sharefs/baai-mrnd/zch/my_stable_diffusion/lm/lgs/lg-lg-8m")
+        self.ch_clip_model = self.ch_clip_model.eval()
+
+        for param in self.ch_clip_model.parameters():
+            param.requires_grad = False
+        
+        self.processor = CHCLIPProcess.from_pretrained("/sharefs/baai-mrnd/zch/my_stable_diffusion/lm/lgs/lg-lg-8m")
+        self.tokenizer = self.processor.tokenizer
+
+        self.text_encoder = self.ch_clip_model.text_model
+        self.normalize = normalize
+        self.layer_norm1 = LayerNorm(1024)
+        self.layer_norm2 = LayerNorm(768)
+
+        self.proj = nn.Linear(1024,768)
+        
+
+
+    def forward(self, text):
+        tokens = self.tokenizer(text,return_tensors='pt',padding=True).to(self.device)
+        #tokens = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=False,
+        #                                 return_overflowing_tokens=False, padding="max_length", return_tensors="pt").to(self.device)
+
+        z = self.text_encoder(**tokens)['hidden_states'][-1]
+        #  进行维度映射
+        z = self.layer_norm1(z)
+        
+
+        z = self.proj(z)
+        z = self.layer_norm2(z)
+
+        return z
+
+    def encode(self, text):
+        z = self(text)
+        if z.ndim==2:
+            z = z[:, None, :]
+        return z
diff --git a/flagai/model/mm/lm/en_clip.py b/flagai/model/mm/lm/en_clip.py
new file mode 100644
index 00000000..a82584c7
--- /dev/null
+++ b/flagai/model/mm/lm/en_clip.py
@@ -0,0 +1,47 @@
+import torch
+import torch.nn as nn
+from functools import partial
+import clip
+from einops import rearrange, repeat
+import transformers
+from transformers import CLIPTokenizer, CLIPTextModel
+import json
+
+
+
+class AbstractEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+
+class FrozenCLIPEmbedder(AbstractEncoder):
+    """Uses the CLIP transformer encoder for text (from Hugging Face)"""
+    def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77):
+        super().__init__()
+        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.transformer = CLIPTextModel.from_pretrained(version)
+        self.device = device
+        self.max_length = max_length
+        self.freeze()
+
+    def freeze(self):
+        self.transformer = self.transformer.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=False,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"].to(self.device)
+        outputs = self.transformer(input_ids=tokens)
+
+        z = outputs.last_hidden_state
+        return z
+
+    def encode(self, text):
+        return self(text)
diff --git a/flagai/model/mm/lm/x_transformer.py b/flagai/model/mm/lm/x_transformer.py
new file mode 100644
index 00000000..5fc15bf9
--- /dev/null
+++ b/flagai/model/mm/lm/x_transformer.py
@@ -0,0 +1,641 @@
+"""shout-out to https://github.com/lucidrains/x-transformers/tree/main/x_transformers"""
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+from functools import partial
+from inspect import isfunction
+from collections import namedtuple
+from einops import rearrange, repeat, reduce
+
+# constants
+
+DEFAULT_DIM_HEAD = 64
+
+Intermediates = namedtuple('Intermediates', [
+    'pre_softmax_attn',
+    'post_softmax_attn'
+])
+
+LayerIntermediates = namedtuple('Intermediates', [
+    'hiddens',
+    'attn_intermediates'
+])
+
+
+class AbsolutePositionalEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len):
+        super().__init__()
+        self.emb = nn.Embedding(max_seq_len, dim)
+        self.init_()
+
+    def init_(self):
+        nn.init.normal_(self.emb.weight, std=0.02)
+
+    def forward(self, x):
+        n = torch.arange(x.shape[1], device=x.device)
+        return self.emb(n)[None, :, :]
+
+
+class FixedPositionalEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+    def forward(self, x, seq_dim=1, offset=0):
+        t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq) + offset
+        sinusoid_inp = torch.einsum('i , j -> i j', t, self.inv_freq)
+        emb = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1)
+        return emb[None, :, :]
+
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def always(val):
+    def inner(*args, **kwargs):
+        return val
+    return inner
+
+
+def not_equals(val):
+    def inner(x):
+        return x != val
+    return inner
+
+
+def equals(val):
+    def inner(x):
+        return x == val
+    return inner
+
+
+def max_neg_value(tensor):
+    return -torch.finfo(tensor.dtype).max
+
+
+# keyword argument helpers
+
+def pick_and_pop(keys, d):
+    values = list(map(lambda key: d.pop(key), keys))
+    return dict(zip(keys, values))
+
+
+def group_dict_by_key(cond, d):
+    return_val = [dict(), dict()]
+    for key in d.keys():
+        match = bool(cond(key))
+        ind = int(not match)
+        return_val[ind][key] = d[key]
+    return (*return_val,)
+
+
+def string_begins_with(prefix, str):
+    return str.startswith(prefix)
+
+
+def group_by_key_prefix(prefix, d):
+    return group_dict_by_key(partial(string_begins_with, prefix), d)
+
+
+def groupby_prefix_and_trim(prefix, d):
+    kwargs_with_prefix, kwargs = group_dict_by_key(partial(string_begins_with, prefix), d)
+    kwargs_without_prefix = dict(map(lambda x: (x[0][len(prefix):], x[1]), tuple(kwargs_with_prefix.items())))
+    return kwargs_without_prefix, kwargs
+
+
+# classes
+class Scale(nn.Module):
+    def __init__(self, value, fn):
+        super().__init__()
+        self.value = value
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        x, *rest = self.fn(x, **kwargs)
+        return (x * self.value, *rest)
+
+
+class Rezero(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+        self.g = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, **kwargs):
+        x, *rest = self.fn(x, **kwargs)
+        return (x * self.g, *rest)
+
+
+class ScaleNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(1))
+
+    def forward(self, x):
+        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+        return x / norm.clamp(min=self.eps) * self.g
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-8):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+        return x / norm.clamp(min=self.eps) * self.g
+
+
+class Residual(nn.Module):
+    def forward(self, x, residual):
+        return x + residual
+
+
+class GRUGating(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.gru = nn.GRUCell(dim, dim)
+
+    def forward(self, x, residual):
+        gated_output = self.gru(
+            rearrange(x, 'b n d -> (b n) d'),
+            rearrange(residual, 'b n d -> (b n) d')
+        )
+
+        return gated_output.reshape_as(x)
+
+
+# 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)
+
+
+# attention.
+class Attention(nn.Module):
+    def __init__(
+            self,
+            dim,
+            dim_head=DEFAULT_DIM_HEAD,
+            heads=8,
+            causal=False,
+            mask=None,
+            talking_heads=False,
+            sparse_topk=None,
+            use_entmax15=False,
+            num_mem_kv=0,
+            dropout=0.,
+            on_attn=False
+    ):
+        super().__init__()
+        if use_entmax15:
+            raise NotImplementedError("Check out entmax activation instead of softmax activation!")
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        self.causal = causal
+        self.mask = mask
+
+        inner_dim = dim_head * heads
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(dim, inner_dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+        # talking heads
+        self.talking_heads = talking_heads
+        if talking_heads:
+            self.pre_softmax_proj = nn.Parameter(torch.randn(heads, heads))
+            self.post_softmax_proj = nn.Parameter(torch.randn(heads, heads))
+
+        # explicit topk sparse attention
+        self.sparse_topk = sparse_topk
+
+        # entmax
+        #self.attn_fn = entmax15 if use_entmax15 else F.softmax
+        self.attn_fn = F.softmax
+
+        # add memory key / values
+        self.num_mem_kv = num_mem_kv
+        if num_mem_kv > 0:
+            self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
+            self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
+
+        # attention on attention
+        self.attn_on_attn = on_attn
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim * 2), nn.GLU()) if on_attn else nn.Linear(inner_dim, dim)
+
+    def forward(
+            self,
+            x,
+            context=None,
+            mask=None,
+            context_mask=None,
+            rel_pos=None,
+            sinusoidal_emb=None,
+            prev_attn=None,
+            mem=None
+    ):
+        b, n, _, h, talking_heads, device = *x.shape, self.heads, self.talking_heads, x.device
+        kv_input = default(context, x)
+
+        q_input = x
+        k_input = kv_input
+        v_input = kv_input
+
+        if exists(mem):
+            k_input = torch.cat((mem, k_input), dim=-2)
+            v_input = torch.cat((mem, v_input), dim=-2)
+
+        if exists(sinusoidal_emb):
+            # in shortformer, the query would start at a position offset depending on the past cached memory
+            offset = k_input.shape[-2] - q_input.shape[-2]
+            q_input = q_input + sinusoidal_emb(q_input, offset=offset)
+            k_input = k_input + sinusoidal_emb(k_input)
+
+        q = self.to_q(q_input)
+        k = self.to_k(k_input)
+        v = self.to_v(v_input)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v))
+
+        input_mask = None
+        if any(map(exists, (mask, context_mask))):
+            q_mask = default(mask, lambda: torch.ones((b, n), device=device).bool())
+            k_mask = q_mask if not exists(context) else context_mask
+            k_mask = default(k_mask, lambda: torch.ones((b, k.shape[-2]), device=device).bool())
+            q_mask = rearrange(q_mask, 'b i -> b () i ()')
+            k_mask = rearrange(k_mask, 'b j -> b () () j')
+            input_mask = q_mask * k_mask
+
+        if self.num_mem_kv > 0:
+            mem_k, mem_v = map(lambda t: repeat(t, 'h n d -> b h n d', b=b), (self.mem_k, self.mem_v))
+            k = torch.cat((mem_k, k), dim=-2)
+            v = torch.cat((mem_v, v), dim=-2)
+            if exists(input_mask):
+                input_mask = F.pad(input_mask, (self.num_mem_kv, 0), value=True)
+
+        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
+        mask_value = max_neg_value(dots)
+
+        if exists(prev_attn):
+            dots = dots + prev_attn
+
+        pre_softmax_attn = dots
+
+        if talking_heads:
+            dots = einsum('b h i j, h k -> b k i j', dots, self.pre_softmax_proj).contiguous()
+
+        if exists(rel_pos):
+            dots = rel_pos(dots)
+
+        if exists(input_mask):
+            dots.masked_fill_(~input_mask, mask_value)
+            del input_mask
+
+        if self.causal:
+            i, j = dots.shape[-2:]
+            r = torch.arange(i, device=device)
+            mask = rearrange(r, 'i -> () () i ()') < rearrange(r, 'j -> () () () j')
+            mask = F.pad(mask, (j - i, 0), value=False)
+            dots.masked_fill_(mask, mask_value)
+            del mask
+
+        if exists(self.sparse_topk) and self.sparse_topk < dots.shape[-1]:
+            top, _ = dots.topk(self.sparse_topk, dim=-1)
+            vk = top[..., -1].unsqueeze(-1).expand_as(dots)
+            mask = dots < vk
+            dots.masked_fill_(mask, mask_value)
+            del mask
+
+        attn = self.attn_fn(dots, dim=-1)
+        post_softmax_attn = attn
+
+        attn = self.dropout(attn)
+
+        if talking_heads:
+            attn = einsum('b h i j, h k -> b k i j', attn, self.post_softmax_proj).contiguous()
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+
+        intermediates = Intermediates(
+            pre_softmax_attn=pre_softmax_attn,
+            post_softmax_attn=post_softmax_attn
+        )
+
+        return self.to_out(out), intermediates
+
+
+class AttentionLayers(nn.Module):
+    def __init__(
+            self,
+            dim,
+            depth,
+            heads=8,
+            causal=False,
+            cross_attend=False,
+            only_cross=False,
+            use_scalenorm=False,
+            use_rmsnorm=False,
+            use_rezero=False,
+            rel_pos_num_buckets=32,
+            rel_pos_max_distance=128,
+            position_infused_attn=False,
+            custom_layers=None,
+            sandwich_coef=None,
+            par_ratio=None,
+            residual_attn=False,
+            cross_residual_attn=False,
+            macaron=False,
+            pre_norm=True,
+            gate_residual=False,
+            **kwargs
+    ):
+        super().__init__()
+        ff_kwargs, kwargs = groupby_prefix_and_trim('ff_', kwargs)
+        attn_kwargs, _ = groupby_prefix_and_trim('attn_', kwargs)
+
+        dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD)
+
+        self.dim = dim
+        self.depth = depth
+        self.layers = nn.ModuleList([])
+
+        self.has_pos_emb = position_infused_attn
+        self.pia_pos_emb = FixedPositionalEmbedding(dim) if position_infused_attn else None
+        self.rotary_pos_emb = always(None)
+
+        assert rel_pos_num_buckets <= rel_pos_max_distance, 'number of relative position buckets must be less than the relative position max distance'
+        self.rel_pos = None
+
+        self.pre_norm = pre_norm
+
+        self.residual_attn = residual_attn
+        self.cross_residual_attn = cross_residual_attn
+
+        norm_class = ScaleNorm if use_scalenorm else nn.LayerNorm
+        norm_class = RMSNorm if use_rmsnorm else norm_class
+        norm_fn = partial(norm_class, dim)
+
+        norm_fn = nn.Identity if use_rezero else norm_fn
+        branch_fn = Rezero if use_rezero else None
+
+        if cross_attend and not only_cross:
+            default_block = ('a', 'c', 'f')
+        elif cross_attend and only_cross:
+            default_block = ('c', 'f')
+        else:
+            default_block = ('a', 'f')
+
+        if macaron:
+            default_block = ('f',) + default_block
+
+        if exists(custom_layers):
+            layer_types = custom_layers
+        elif exists(par_ratio):
+            par_depth = depth * len(default_block)
+            assert 1 < par_ratio <= par_depth, 'par ratio out of range'
+            default_block = tuple(filter(not_equals('f'), default_block))
+            par_attn = par_depth // par_ratio
+            depth_cut = par_depth * 2 // 3  # 2 / 3 attention layer cutoff suggested by PAR paper
+            par_width = (depth_cut + depth_cut // par_attn) // par_attn
+            assert len(default_block) <= par_width, 'default block is too large for par_ratio'
+            par_block = default_block + ('f',) * (par_width - len(default_block))
+            par_head = par_block * par_attn
+            layer_types = par_head + ('f',) * (par_depth - len(par_head))
+        elif exists(sandwich_coef):
+            assert sandwich_coef > 0 and sandwich_coef <= depth, 'sandwich coefficient should be less than the depth'
+            layer_types = ('a',) * sandwich_coef + default_block * (depth - sandwich_coef) + ('f',) * sandwich_coef
+        else:
+            layer_types = default_block * depth
+
+        self.layer_types = layer_types
+        self.num_attn_layers = len(list(filter(equals('a'), layer_types)))
+
+        for layer_type in self.layer_types:
+            if layer_type == 'a':
+                layer = Attention(dim, heads=heads, causal=causal, **attn_kwargs)
+            elif layer_type == 'c':
+                layer = Attention(dim, heads=heads, **attn_kwargs)
+            elif layer_type == 'f':
+                layer = FeedForward(dim, **ff_kwargs)
+                layer = layer if not macaron else Scale(0.5, layer)
+            else:
+                raise Exception(f'invalid layer type {layer_type}')
+
+            if isinstance(layer, Attention) and exists(branch_fn):
+                layer = branch_fn(layer)
+
+            if gate_residual:
+                residual_fn = GRUGating(dim)
+            else:
+                residual_fn = Residual()
+
+            self.layers.append(nn.ModuleList([
+                norm_fn(),
+                layer,
+                residual_fn
+            ]))
+
+    def forward(
+            self,
+            x,
+            context=None,
+            mask=None,
+            context_mask=None,
+            mems=None,
+            return_hiddens=False
+    ):
+        hiddens = []
+        intermediates = []
+        prev_attn = None
+        prev_cross_attn = None
+
+        mems = mems.copy() if exists(mems) else [None] * self.num_attn_layers
+
+        for ind, (layer_type, (norm, block, residual_fn)) in enumerate(zip(self.layer_types, self.layers)):
+            is_last = ind == (len(self.layers) - 1)
+
+            if layer_type == 'a':
+                hiddens.append(x)
+                layer_mem = mems.pop(0)
+
+            residual = x
+
+            if self.pre_norm:
+                x = norm(x)
+
+            if layer_type == 'a':
+                out, inter = block(x, mask=mask, sinusoidal_emb=self.pia_pos_emb, rel_pos=self.rel_pos,
+                                   prev_attn=prev_attn, mem=layer_mem)
+            elif layer_type == 'c':
+                out, inter = block(x, context=context, mask=mask, context_mask=context_mask, prev_attn=prev_cross_attn)
+            elif layer_type == 'f':
+                out = block(x)
+
+            x = residual_fn(out, residual)
+
+            if layer_type in ('a', 'c'):
+                intermediates.append(inter)
+
+            if layer_type == 'a' and self.residual_attn:
+                prev_attn = inter.pre_softmax_attn
+            elif layer_type == 'c' and self.cross_residual_attn:
+                prev_cross_attn = inter.pre_softmax_attn
+
+            if not self.pre_norm and not is_last:
+                x = norm(x)
+
+        if return_hiddens:
+            intermediates = LayerIntermediates(
+                hiddens=hiddens,
+                attn_intermediates=intermediates
+            )
+
+            return x, intermediates
+
+        return x
+
+
+class Encoder(AttentionLayers):
+    def __init__(self, **kwargs):
+        assert 'causal' not in kwargs, 'cannot set causality on encoder'
+        super().__init__(causal=False, **kwargs)
+
+
+
+class TransformerWrapper(nn.Module):
+    def __init__(
+            self,
+            *,
+            num_tokens,
+            max_seq_len,
+            attn_layers,
+            emb_dim=None,
+            max_mem_len=0.,
+            emb_dropout=0.,
+            num_memory_tokens=None,
+            tie_embedding=False,
+            use_pos_emb=True
+    ):
+        super().__init__()
+        assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder'
+
+        dim = attn_layers.dim
+        emb_dim = default(emb_dim, dim)
+
+        self.max_seq_len = max_seq_len
+        self.max_mem_len = max_mem_len
+        self.num_tokens = num_tokens
+
+        self.token_emb = nn.Embedding(num_tokens, emb_dim)
+        self.pos_emb = AbsolutePositionalEmbedding(emb_dim, max_seq_len) if (
+                    use_pos_emb and not attn_layers.has_pos_emb) else always(0)
+        self.emb_dropout = nn.Dropout(emb_dropout)
+
+        self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity()
+        self.attn_layers = attn_layers
+        self.norm = nn.LayerNorm(dim)
+
+        self.init_()
+
+        self.to_logits = nn.Linear(dim, num_tokens) if not tie_embedding else lambda t: t @ self.token_emb.weight.t()
+
+        # memory tokens (like [cls]) from Memory Transformers paper
+        num_memory_tokens = default(num_memory_tokens, 0)
+        self.num_memory_tokens = num_memory_tokens
+        if num_memory_tokens > 0:
+            self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
+
+            # let funnel encoder know number of memory tokens, if specified
+            if hasattr(attn_layers, 'num_memory_tokens'):
+                attn_layers.num_memory_tokens = num_memory_tokens
+
+    def init_(self):
+        nn.init.normal_(self.token_emb.weight, std=0.02)
+
+    def forward(
+            self,
+            x,
+            return_embeddings=False,
+            mask=None,
+            return_mems=False,
+            return_attn=False,
+            mems=None,
+            **kwargs
+    ):
+        b, n, device, num_mem = *x.shape, x.device, self.num_memory_tokens
+        x = self.token_emb(x)
+        x += self.pos_emb(x)
+        x = self.emb_dropout(x)
+
+        x = self.project_emb(x)
+
+        if num_mem > 0:
+            mem = repeat(self.memory_tokens, 'n d -> b n d', b=b)
+            x = torch.cat((mem, x), dim=1)
+
+            # auto-handle masking after appending memory tokens
+            if exists(mask):
+                mask = F.pad(mask, (num_mem, 0), value=True)
+
+        x, intermediates = self.attn_layers(x, mask=mask, mems=mems, return_hiddens=True, **kwargs)
+        x = self.norm(x)
+
+        mem, x = x[:, :num_mem], x[:, num_mem:]
+
+        out = self.to_logits(x) if not return_embeddings else x
+
+        if return_mems:
+            hiddens = intermediates.hiddens
+            new_mems = list(map(lambda pair: torch.cat(pair, dim=-2), zip(mems, hiddens))) if exists(mems) else hiddens
+            new_mems = list(map(lambda t: t[..., -self.max_mem_len:, :].detach(), new_mems))
+            return out, new_mems
+
+        if return_attn:
+            attn_maps = list(map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates))
+            return out, attn_maps
+
+        return out
+
diff --git a/flagai/model/mm/model.py b/flagai/model/mm/model.py
new file mode 100755
index 00000000..09396431
--- /dev/null
+++ b/flagai/model/mm/model.py
@@ -0,0 +1,1059 @@
+# pytorch_diffusion + derived encoder decoder
+import math
+import torch
+import torch.nn as nn
+import numpy as np
+from einops import rearrange
+from flagai.model.mm.utils import instantiate_from_config
+
+
+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)
+
+
+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
+
+
+class DiagonalGaussianDistribution(object):
+
+    def __init__(self, parameters, deterministic=False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(
+                self.mean).to(device=self.parameters.device)
+
+    def sample(self):
+        x = self.mean + self.std * torch.randn(
+            self.mean.shape).to(device=self.parameters.device)
+        return x
+
+    def kl(self, other=None):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        else:
+            if other is None:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+                    dim=[1, 2, 3])
+            else:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean - other.mean, 2) / other.var +
+                    self.var / other.var - 1.0 - self.logvar + other.logvar,
+                    dim=[1, 2, 3])
+
+    def nll(self, sample, dims=[1, 2, 3]):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(logtwopi + self.logvar +
+                               torch.pow(sample - self.mean, 2) / self.var,
+                               dim=dims)
+
+    def mode(self):
+        return self.mean
diff --git a/flagai/model/mm/modeling_berts.py b/flagai/model/mm/modeling_berts.py
new file mode 100644
index 00000000..43d7dda7
--- /dev/null
+++ b/flagai/model/mm/modeling_berts.py
@@ -0,0 +1,105 @@
+from transformers import BertPreTrainedModel,BertModel,BertConfig
+import torch.nn as nn
+import torch
+from transformers.models.xlm_roberta.configuration_xlm_roberta import XLMRobertaConfig
+from transformers import XLMRobertaModel
+from typing import Optional
+
+class BertSeriesConfig(BertConfig):
+    def __init__(self, vocab_size=30522, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act="gelu", hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1, max_position_embeddings=512, type_vocab_size=2, initializer_range=0.02, layer_norm_eps=1e-12, pad_token_id=0, position_embedding_type="absolute", use_cache=True, classifier_dropout=None,project_dim=512, pooler_fn="average",learn_encoder=False,model_type='bert',**kwargs):
+    
+        super().__init__(vocab_size, hidden_size, num_hidden_layers, num_attention_heads, intermediate_size, hidden_act, hidden_dropout_prob, attention_probs_dropout_prob, max_position_embeddings, type_vocab_size, initializer_range, layer_norm_eps, pad_token_id, position_embedding_type, use_cache, classifier_dropout, **kwargs)
+        self.project_dim = project_dim
+        self.pooler_fn = pooler_fn
+        self.learn_encoder = learn_encoder
+
+class RobertaSeriesConfig(XLMRobertaConfig):
+    def __init__(self, pad_token_id=1, bos_token_id=0, eos_token_id=2,project_dim=512,pooler_fn='cls',learn_encoder=False, **kwargs):
+        super().__init__(pad_token_id, bos_token_id, eos_token_id, **kwargs)
+        self.project_dim = project_dim
+        self.pooler_fn = pooler_fn
+        self.learn_encoder = learn_encoder
+
+    
+
+class BertSeriesModelWithTransformation(BertPreTrainedModel):
+
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+    config_class = BertSeriesConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        if config.model_type == 'bert':
+            self.bert = BertModel(config)
+        elif config.model_type == 'xlm-roberta':
+            self.roberta = XLMRobertaModel(config)
+        self.transformation = nn.Linear(config.hidden_size,config.project_dim)
+        self.learn_encoder = config.learn_encoder
+        if config.learn_encoder:
+            self.encoder_tfm = nn.Linear(config.hidden_size,config.project_dim)
+        self.pre_LN=nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        
+        self.pooler = lambda x: x[:,0]
+        self.post_init()
+        
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+    ) :
+        r"""
+        """
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        
+        outputs = self.base_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=True,
+            return_dict=return_dict,
+        )
+        
+        # last module outputs
+        sequence_output = outputs[0]
+        
+
+        # project every module
+        sequence_output_ln = self.pre_LN(sequence_output)
+            
+        # pooler
+        pooler_output = self.pooler(sequence_output_ln)
+        pooler_output = self.transformation(pooler_output)
+        projection_state = self.transformation(outputs.last_hidden_state)
+        
+        return {
+            'pooler_output':pooler_output,
+            'last_hidden_state':outputs.last_hidden_state,
+            'hidden_states':outputs.hidden_states,
+            'attentions':outputs.attentions,
+            'projection_state':projection_state,
+            'sequence_out': sequence_output
+        }
+
+
+class RobertaSeriesModelWithTransformation(BertSeriesModelWithTransformation):
+    base_model_prefix = 'roberta'
+    config_class= RobertaSeriesConfig
diff --git a/flagai/model/mm/utils.py b/flagai/model/mm/utils.py
old mode 100644
new mode 100755
index bba444f2..de5b6618
--- a/flagai/model/mm/utils.py
+++ b/flagai/model/mm/utils.py
@@ -1,13 +1,517 @@
-
 from itertools import repeat
 import collections.abc
+import importlib
+
+import torch
+import numpy as np
+from collections import abc
+from einops import rearrange
+from functools import partial
+
+import multiprocessing as mp
+from threading import Thread
+from queue import Queue
+
+from inspect import isfunction
+from PIL import Image, ImageDraw, ImageFont
+
+import os
+import math
+import torch.nn as nn
+import numpy as np
+from einops import repeat
+
+
 # From PyTorch internals
 def _ntuple(n):
+
     def parse(x):
         if isinstance(x, collections.abc.Iterable):
             return x
         return tuple(repeat(x, n))
+
     return parse
 
 
-to_2tuple = _ntuple(2) #for clip
\ No newline at end of file
+to_2tuple = _ntuple(2)  #for clip
+
+
+def log_txt_as_img(wh, xc, size=10):
+    # wh a tuple of (width, height)
+    # xc a list of captions to plot
+    b = len(xc)
+    txts = list()
+    for bi in range(b):
+
+        #font = ImageFont.truetype('data/DejaVuSans.ttf', size=size)
+        # 用中文字体
+        try:
+            txt = Image.new("RGB", wh, color="white")
+            draw = ImageDraw.Draw(txt)
+            font = ImageFont.truetype("simsun.ttc", size=size, encoding="unic")
+            nc = int(40 * (wh[0] / 256))
+            lines = "\n".join(xc[bi][start:start + nc]
+                              for start in range(0, len(xc[bi]), nc))
+            draw.text((0, 0), lines, fill="black", font=font)
+        except:
+            print("Cant encode string for logging. Skipping.")
+
+        txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
+        txts.append(txt)
+    txts = np.stack(txts)
+    txts = torch.tensor(txts)
+    return txts
+
+
+def ismap(x):
+    if not isinstance(x, torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] > 3)
+
+
+def isimage(x):
+    if not isinstance(x, torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
+
+
+def exists(x):
+    return x is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def mean_flat(tensor):
+    """
+    https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+def count_params(model, verbose=False):
+    total_params = sum(p.numel() for p in model.parameters())
+    if verbose:
+        print(
+            f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params."
+        )
+    return total_params
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        if config == '__is_first_stage__':
+            return None
+        elif config == "__is_unconditional__":
+            return None
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+def _do_parallel_data_prefetch(func, Q, data, idx, idx_to_fn=False):
+    # create dummy dataset instance
+
+    # run prefetching
+    if idx_to_fn:
+        res = func(data, worker_id=idx)
+    else:
+        res = func(data)
+    Q.put([idx, res])
+    Q.put("Done")
+
+
+def parallel_data_prefetch(func: callable,
+                           data,
+                           n_proc,
+                           target_data_type="ndarray",
+                           cpu_intensive=True,
+                           use_worker_id=False):
+    # if target_data_type not in ["ndarray", "list"]:
+    #     raise ValueError(
+    #         "Data, which is passed to parallel_data_prefetch has to be either of type list or ndarray."
+    #     )
+    if isinstance(data, np.ndarray) and target_data_type == "list":
+        raise ValueError("list expected but function got ndarray.")
+    elif isinstance(data, abc.Iterable):
+        if isinstance(data, dict):
+            print(
+                f'WARNING:"data" argument passed to parallel_data_prefetch is a dict: Using only its values and disregarding keys.'
+            )
+            data = list(data.values())
+        if target_data_type == "ndarray":
+            data = np.asarray(data)
+        else:
+            data = list(data)
+    else:
+        raise TypeError(
+            f"The data, that shall be processed parallel has to be either an np.ndarray or an Iterable, but is actually {type(data)}."
+        )
+
+    if cpu_intensive:
+        Q = mp.Queue(1000)
+        proc = mp.Process
+    else:
+        Q = Queue(1000)
+        proc = Thread
+    # spawn processes
+    if target_data_type == "ndarray":
+        arguments = [[func, Q, part, i, use_worker_id]
+                     for i, part in enumerate(np.array_split(data, n_proc))]
+    else:
+        step = (int(len(data) / n_proc +
+                    1) if len(data) % n_proc != 0 else int(len(data) / n_proc))
+        arguments = [[func, Q, part, i, use_worker_id]
+                     for i, part in enumerate(
+                         [data[i:i + step]
+                          for i in range(0, len(data), step)])]
+    processes = []
+    for i in range(n_proc):
+        p = proc(target=_do_parallel_data_prefetch, args=arguments[i])
+        processes += [p]
+
+    # start processes
+    print(f"Start prefetching...")
+    import time
+
+    start = time.time()
+    gather_res = [[] for _ in range(n_proc)]
+    try:
+        for p in processes:
+            p.start()
+
+        k = 0
+        while k < n_proc:
+            # get result
+            res = Q.get()
+            if res == "Done":
+                k += 1
+            else:
+                gather_res[res[0]] = res[1]
+
+    except Exception as e:
+        print("Exception: ", e)
+        for p in processes:
+            p.terminate()
+
+        raise e
+    finally:
+        for p in processes:
+            p.join()
+        print(f"Prefetching complete. [{time.time() - start} sec.]")
+
+    if target_data_type == 'ndarray':
+        if not isinstance(gather_res[0], np.ndarray):
+            return np.concatenate([np.asarray(r) for r in gather_res], axis=0)
+
+        # order outputs
+        return np.concatenate(gather_res, axis=0)
+    elif target_data_type == 'list':
+        out = []
+        for r in gather_res:
+            out.extend(r)
+        return out
+    else:
+        return gather_res
+
+
+def make_beta_schedule(schedule,
+                       n_timestep,
+                       linear_start=1e-4,
+                       linear_end=2e-2,
+                       cosine_s=8e-3):
+    if schedule == "linear":
+        betas = (torch.linspace(linear_start**0.5,
+                                linear_end**0.5,
+                                n_timestep,
+                                dtype=torch.float64)**2)
+
+    elif schedule == "cosine":
+        timesteps = (
+            torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep +
+            cosine_s)
+        alphas = timesteps / (1 + cosine_s) * np.pi / 2
+        alphas = torch.cos(alphas).pow(2)
+        alphas = alphas / alphas[0]
+        betas = 1 - alphas[1:] / alphas[:-1]
+        betas = np.clip(betas, a_min=0, a_max=0.999)
+
+    elif schedule == "sqrt_linear":
+        betas = torch.linspace(linear_start,
+                               linear_end,
+                               n_timestep,
+                               dtype=torch.float64)
+    elif schedule == "sqrt":
+        betas = torch.linspace(linear_start,
+                               linear_end,
+                               n_timestep,
+                               dtype=torch.float64)**0.5
+    else:
+        raise ValueError(f"schedule '{schedule}' unknown.")
+    return betas.numpy()
+
+
+def make_ddim_timesteps(ddim_discr_method,
+                        num_ddim_timesteps,
+                        num_ddpm_timesteps,
+                        verbose=True):
+    if ddim_discr_method == 'uniform':
+        c = num_ddpm_timesteps // num_ddim_timesteps
+        ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
+    elif ddim_discr_method == 'quad':
+        ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8),
+                                       num_ddim_timesteps))**2).astype(int)
+    else:
+        raise NotImplementedError(
+            f'There is no ddim discretization method called "{ddim_discr_method}"'
+        )
+
+    # assert ddim_timesteps.shape[0] == num_ddim_timesteps
+    # add one to get the final alpha values right (the ones from first scale to data during sampling)
+    steps_out = ddim_timesteps + 1
+    if verbose:
+        print(f'Selected timesteps for ddim sampler: {steps_out}')
+    return steps_out
+
+
+def make_ddim_sampling_parameters(alphacums,
+                                  ddim_timesteps,
+                                  eta,
+                                  verbose=True):
+    # select alphas for computing the variance schedule
+    alphas = alphacums[ddim_timesteps]
+    alphas_prev = np.asarray([alphacums[0]] +
+                             alphacums[ddim_timesteps[:-1]].tolist())
+
+    # according the the formula provided in https://arxiv.org/abs/2010.02502
+    sigmas = eta * np.sqrt(
+        (1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
+    if verbose:
+        print(
+            f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}'
+        )
+        print(
+            f'For the chosen value of eta, which is {eta}, '
+            f'this results in the following sigma_t schedule for ddim sampler {sigmas}'
+        )
+    return sigmas, alphas, alphas_prev
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+
+def extract_into_tensor(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1, ) * (len(x_shape) - 1)))
+
+
+def checkpoint(func, inputs, params, flag):
+    """
+    Evaluate a function without caching intermediate activations, allowing for
+    reduced memory at the expense of extra compute in the backward pass.
+    :param func: the function to evaluate.
+    :param inputs: the argument sequence to pass to `func`.
+    :param params: a sequence of parameters `func` depends on but does not
+                   explicitly take as arguments.
+    :param flag: if False, disable gradient checkpointing.
+    """
+    if flag:
+        args = tuple(inputs) + tuple(params)
+        return CheckpointFunction.apply(func, len(inputs), *args)
+    else:
+        return func(*inputs)
+
+
+class CheckpointFunction(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, run_function, length, *args):
+        ctx.run_function = run_function
+        ctx.input_tensors = list(args[:length])
+        ctx.input_params = list(args[length:])
+
+        with torch.no_grad():
+            output_tensors = ctx.run_function(*ctx.input_tensors)
+        return output_tensors
+
+    @staticmethod
+    def backward(ctx, *output_grads):
+        ctx.input_tensors = [
+            x.detach().requires_grad_(True) for x in ctx.input_tensors
+        ]
+        with torch.enable_grad():
+            # Fixes a bug where the first op in run_function modifies the
+            # Tensor storage in place, which is not allowed for detach()'d
+            # Tensors.
+            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
+            output_tensors = ctx.run_function(*shallow_copies)
+        input_grads = torch.autograd.grad(
+            output_tensors,
+            ctx.input_tensors + ctx.input_params,
+            output_grads,
+            allow_unused=True,
+        )
+        del ctx.input_tensors
+        del ctx.input_params
+        del output_tensors
+        return (None, None) + input_grads
+
+
+def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+    """
+    Create sinusoidal timestep embeddings.
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    if not repeat_only:
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) *
+            torch.arange(start=0, end=half, dtype=torch.float32) /
+            half).to(device=timesteps.device)
+        args = timesteps[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat(
+                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    else:
+        embedding = repeat(timesteps, 'b -> b d', d=dim)
+    return embedding
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+def scale_module(module, scale):
+    """
+    Scale the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().mul_(scale)
+    return module
+
+
+def mean_flat(tensor):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+def normalization(channels):
+    """
+    Make a standard normalization layer.
+    :param channels: number of input channels.
+    :return: an nn.Module for normalization.
+    """
+    return GroupNorm32(32, channels)
+
+
+# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
+class SiLU(nn.Module):
+
+    def forward(self, x):
+        return x * torch.sigmoid(x)
+
+
+class GroupNorm32(nn.GroupNorm):
+
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+
+def conv_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D convolution module.
+    """
+    if dims == 1:
+        return nn.Conv1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.Conv3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def linear(*args, **kwargs):
+    """
+    Create a linear module.
+    """
+    return nn.Linear(*args, **kwargs)
+
+
+def avg_pool_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D average pooling module.
+    """
+    if dims == 1:
+        return nn.AvgPool1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.AvgPool2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.AvgPool3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+class HybridConditioner(nn.Module):
+
+    def __init__(self, c_concat_config, c_crossattn_config):
+        super().__init__()
+        self.concat_conditioner = instantiate_from_config(c_concat_config)
+        self.crossattn_conditioner = instantiate_from_config(
+            c_crossattn_config)
+
+    def forward(self, c_concat, c_crossattn):
+        c_concat = self.concat_conditioner(c_concat)
+        c_crossattn = self.crossattn_conditioner(c_crossattn)
+        return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
+
+
+def noise_like(shape, device, repeat=False):
+    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(
+        shape[0], *((1, ) * (len(shape) - 1)))
+    noise = lambda: torch.randn(shape, device=device)
+    return repeat_noise() if repeat else noise()
+
+
+def extract_into_tensor(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1, ) * (len(x_shape) - 1)))
diff --git a/flagai/model/predictor/predictor.py b/flagai/model/predictor/predictor.py
index 65062625..9e37e673 100755
--- a/flagai/model/predictor/predictor.py
+++ b/flagai/model/predictor/predictor.py
@@ -9,13 +9,17 @@
     t5_beamsearch, gpt_beamsearch, bert_random_sample, glm_beamsearch, glm_random_sample, cpm_beamsearch
 from typing import List, Union, Dict, Tuple, Any
 from flagai.model.predictor.gpt import gpt_random_sample_use_cache
-
+from flagai.model.mm.Sampler import DDIMSampler, PLMSSampler
+import os
+from PIL import Image
+from tqdm import trange, tqdm
+import time
+from contextlib import contextmanager, nullcontext
+from einops import rearrange
 
 class Predictor:
 
-    def __init__(self,
-                 model,
-                 tokenizer):
+    def __init__(self, model, tokenizer=None):
         """
         Args:
             model: The model loaded by the AutoLoader class.
@@ -42,7 +46,8 @@ def __init__(self,
         if getattr(self.tokenizer, "token_start_id", None) is None:
             if word2idx is not None:
                 if word2idx.get("[CLS]", None) is not None:
-                    setattr(self.tokenizer, "token_start_id", word2idx["[CLS]"])
+                    setattr(self.tokenizer, "token_start_id",
+                            word2idx["[CLS]"])
                 elif word2idx.get("
", None) is not None:
                     setattr(self.tokenizer, "token_start_id", word2idx["
"])
                 else:
@@ -66,7 +71,6 @@ def __init__(self,
                 else:
                     setattr(self.tokenizer, "token_pad_id", 0)
 
-
         self.model = model
         self.model.eval()
         self.class_name = type(model).__name__
@@ -172,7 +176,7 @@ def predict_ner(self,
         tokens = tokenizer.tokenize(text,
                                     maxlen=maxlen,
                                     add_spatial_tokens=add_spatial_token)
-        
+
         mapping = tokenizer.rematch(text, tokens)
         token_ids = tokenizer.text_tokenizer.convert_tokens_to_ids(tokens)
         token_ids = torch.tensor([token_ids], dtype=torch.long, device=device)
@@ -288,11 +292,13 @@ def predict_generate_randomsample(self,
                                     top_p, repetition_penalty, temperature,
                                     device)
 
-        elif "gpt" in self.class_name.lower() or "opt" in self.class_name.lower():
-            return gpt_random_sample_use_cache(self.model, self.tokenizer, text,
-                                     input_max_length, out_max_length, top_k,
-                                     top_p, repetition_penalty, temperature,
-                                     device)
+        elif "gpt" in self.class_name.lower(
+        ) or "opt" in self.class_name.lower():
+            return gpt_random_sample_use_cache(self.model, self.tokenizer,
+                                               text, input_max_length,
+                                               out_max_length, top_k, top_p,
+                                               repetition_penalty, temperature,
+                                               device)
         elif "glm" in self.class_name.lower():
             return glm_random_sample(self.model, self.tokenizer, text,
                                      out_max_length, top_k, top_p,
@@ -310,3 +316,155 @@ def predict_generate_randomsample(self,
             print("Unsupported decoding mode")
             import os
             os._exit(0)
+
+    def predict_generate_images(self,
+                                prompt: str,
+                                outpath: str = "AltDiffusionOutputs",
+                                n_samples: int = 4,
+                                n_rows: int = 0,
+                                skip_grid: bool = False,
+                                skip_save: bool = False,
+                                ddim_steps: int = 50,
+                                n_iter: int = 1,
+                                plms: bool = False,
+                                fixed_code: bool = False,
+                                ddim_eta: float = 0.0,
+                                H: int = 512,
+                                W: int = 512,
+                                C: int = 4,
+                                f: int = 8,
+                                scale: float = 7.5,
+                                from_file: str = None,
+                                seed: int = 34234):
+        from torchvision.utils import make_grid
+        from pytorch_lightning import seed_everything
+        from flagai.model.predictor.utils import chunk, check_safety, get_safety_checker
+        safety_checker, safety_feature_extractor = get_safety_checker()
+        """
+        Args:
+        prompt: the prompt text
+        out_path: the output path
+        n_samples: how many images to be generated
+        skip_grid: not to grid images
+        skip_save: do not save images
+        ddim_step: number of steps in ddim model
+        n_iter: number of iterations
+        plms: use PLMSSampler
+        fixed_code: sampled from a initial start code
+        seed: Random seed
+        H: height of image
+        W: width of image
+        C: channels of images, 4 for colored images
+        """
+        seed_everything(seed)
+
+        assert "diffusion" in self.class_name.lower()
+        device = next(self.model.parameters()).device
+        if plms:
+            sampler = PLMSSampler(self.model)
+        else:
+            sampler = DDIMSampler(self.model)
+
+        os.makedirs(outpath, exist_ok=True)
+
+        batch_size = n_samples
+        n_rows = n_rows if n_rows > 0 else batch_size
+        if not from_file:
+            assert prompt is not None
+            data = [batch_size * [prompt]]
+
+        else:
+            print(f"reading prompts from {from_file}")
+            with open(from_file, "r") as f:
+                data = f.read().splitlines()
+                data = list(chunk(data, batch_size))
+
+        sample_path = os.path.join(outpath, "samples")
+        os.makedirs(sample_path, exist_ok=True)
+        base_count = len(os.listdir(sample_path))
+        grid_count = len(os.listdir(outpath)) - 1
+
+        start_code = None
+        if fixed_code:
+            start_code = torch.randn([n_samples, C, H // f, W // f],
+                                     device=device)
+
+        precision_scope = nullcontext
+        with torch.no_grad():
+            with precision_scope("cuda"):
+                with self.model.ema_scope():
+                    tic = time.time()
+                    all_samples = list()
+                    for n in trange(n_iter, desc="Sampling"):
+                        for prompts in tqdm(data, desc="data"):
+                            uc = None
+                            if scale != 1.0:
+                                uc = self.model.get_learned_conditioning(
+                                    batch_size * [""])
+                            if isinstance(prompts, tuple):
+                                prompts = list(prompts)
+                            c = self.model.get_learned_conditioning(prompts)
+                            shape = [C, H // f, W // f]
+                            samples_ddim, _ = sampler.sample(
+                                S=ddim_steps,
+                                conditioning=c,
+                                batch_size=n_samples,
+                                shape=shape,
+                                verbose=False,
+                                unconditional_guidance_scale=scale,
+                                unconditional_conditioning=uc,
+                                eta=ddim_eta,
+                                x_T=start_code)
+
+                            x_samples_ddim = self.model.decode_first_stage(
+                                samples_ddim)
+                            x_samples_ddim = torch.clamp(
+                                (x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
+                            x_samples_ddim = x_samples_ddim.cpu().permute(
+                                0, 2, 3, 1).numpy()
+
+                            x_checked_image, has_nsfw_concept = check_safety(safety_checker, safety_feature_extractor, x_samples_ddim)
+
+                            x_checked_image_torch = torch.from_numpy(
+                                x_samples_ddim).permute(0, 3, 1, 2)
+
+                            prompt_count = 0
+                            if not skip_save:
+                                for x_sample in x_checked_image_torch:
+                                    x_sample = 255. * rearrange(
+                                        x_sample.cpu().numpy(),
+                                        'c h w -> h w c')
+                                    img = Image.fromarray(
+                                        x_sample.astype(np.uint8))
+                                    img.save(
+                                        os.path.join(sample_path,
+                                                     f"{base_count:05}.png"))
+                                    #img.save(os.path.join(sample_path, f"{prompts[prompt_count]}.png"))
+
+                                    base_count += 1
+                                    prompt_count = prompt_count % batch_size
+                                    prompt_count += 1
+
+                            if not skip_grid:
+                                all_samples.append(x_checked_image_torch)
+
+                    if not skip_grid:
+                        # additionally, save as grid
+                        grid = torch.stack(all_samples, 0)
+                        grid = rearrange(grid, 'n b c h w -> (n b) c h w')
+                        grid = make_grid(grid, nrow=n_rows)
+
+                        # to image
+                        grid = 255. * rearrange(
+                            grid, 'c h w -> h w c').cpu().numpy()
+                        img = Image.fromarray(grid.astype(np.uint8))
+                        # img = put_watermark(img, wm_encoder)
+                        img.save(
+                            os.path.join(outpath, f'grid-{grid_count:04}.png'))
+                        grid_count += 1
+
+                    toc = time.time()
+
+        print(
+            f"Your samples are ready and waiting for you here: \n{outpath} \n"
+            f" \nEnjoy.")
diff --git a/flagai/model/predictor/utils.py b/flagai/model/predictor/utils.py
old mode 100644
new mode 100755
index cd201a72..45f05f4c
--- a/flagai/model/predictor/utils.py
+++ b/flagai/model/predictor/utils.py
@@ -8,11 +8,50 @@
 import numpy as np
 import torch.nn.functional as F
 import time
-
+from PIL import Image
+from itertools import islice
+from transformers import AutoFeatureExtractor
 import math
 
 join = os.path.join
 
+def get_safety_checker():
+    # load safety model
+    from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+    safety_model_id = "CompVis/stable-diffusion-safety-checker"
+    safety_feature_extractor = AutoFeatureExtractor.from_pretrained(
+        safety_model_id)
+    safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id)
+    return safety_checker, safety_feature_extractor
+
+def chunk(it, size):
+    it = iter(it)
+    return iter(lambda: tuple(islice(it, size)), ())
+
+
+def numpy_to_pil(images):
+    """
+    Convert a numpy image or a batch of images to a PIL image.
+    """
+    if images.ndim == 3:
+        images = images[None, ...]
+    images = (images * 255).round().astype("uint8")
+    pil_images = [Image.fromarray(image) for image in images]
+
+    return pil_images
+
+
+def load_replacement(x):
+    try:
+        hwc = x.shape
+        y = Image.open("assets/rick.jpeg").convert("RGB").resize(
+            (hwc[1], hwc[0]))
+        y = (np.array(y) / 255.0).astype(x.dtype)
+        assert y.shape == x.shape
+        return y
+    except Exception:
+        return x
+
 
 def load_config(config_path):
     with open(config_path) as f:
@@ -21,6 +60,30 @@ def load_config(config_path):
     return j
 
 
+def load_replacement(x):
+    try:
+        hwc = x.shape
+        y = Image.open("assets/rick.jpeg").convert("RGB").resize(
+            (hwc[1], hwc[0]))
+        y = (np.array(y) / 255.0).astype(x.dtype)
+        assert y.shape == x.shape
+        return y
+    except Exception:
+        return x
+
+
+def check_safety(safety_checker, safety_feature_extractor, x_image):
+    safety_checker_input = safety_feature_extractor(numpy_to_pil(x_image),
+                                                    return_tensors="pt")
+    x_checked_image, has_nsfw_concept = safety_checker(
+        images=x_image, clip_input=safety_checker_input.pixel_values)
+    assert x_checked_image.shape[0] == len(has_nsfw_concept)
+    for i in range(len(has_nsfw_concept)):
+        if has_nsfw_concept[i]:
+            x_checked_image[i] = load_replacement(x_checked_image[i])
+    return x_checked_image, has_nsfw_concept
+
+
 class LogitsProcessor:
     """Abstract base class for all logit processors that can be applied during generation."""
 
@@ -188,7 +251,12 @@ def __call__(self, input_ids: torch.LongTensor,
 
 class BeamHypotheses(object):
 
-    def __init__(self, n_hyp, max_len, length_penalty, early_stopping, tokenizer=None):
+    def __init__(self,
+                 n_hyp,
+                 max_len,
+                 length_penalty,
+                 early_stopping,
+                 tokenizer=None):
         """
         Initialize n-best list of hypotheses.
         """
@@ -210,17 +278,19 @@ def add(self, hyp, sum_logprobs):
         """
         Add a new hypothesis to the list.
         """
-        score = sum_logprobs / len(hyp) ** self.length_penalty
+        score = sum_logprobs / len(hyp)**self.length_penalty
 
         if len(self) < self.n_hyp or score > self.worst_score:
             self.hyp.append((score, hyp))
             if len(self) > self.n_hyp:
-                sorted_scores = sorted([(s, idx) for idx, (s, _) in enumerate(self.hyp)])
+                sorted_scores = sorted([
+                    (s, idx) for idx, (s, _) in enumerate(self.hyp)
+                ])
                 del self.hyp[sorted_scores[0][1]]
                 self.worst_score = sorted_scores[1][0]
             else:
                 self.worst_score = min(score, self.worst_score)
-        
+
     def is_done(self, best_sum_logprobs, cur_len):
         """
         If there are enough hypotheses and that none of the hypotheses being generated
@@ -231,8 +301,7 @@ def is_done(self, best_sum_logprobs, cur_len):
         elif self.early_stopping:
             return True
         else:
-            return self.worst_score >= best_sum_logprobs / cur_len ** self.length_penalty
-
+            return self.worst_score >= best_sum_logprobs / cur_len**self.length_penalty
 
 
 def viterbi_decode(nodes, trans):
@@ -763,6 +832,7 @@ def gpt_beamsearch(model, tokenizer, text, input_max_length, out_max_length,
     output = tokenizer.decode(out_puts_ids)
     return output
 
+
 def cpm_beamsearch(model, tokenizer, text, input_max_length, out_max_length,
                    beam_size):
     # tokenizer_out = tokenizer.encode_plus(text, max_length=input_max_length)
@@ -1055,6 +1125,7 @@ def convert_to_ids(tokenizer, text):
     ids = [j for j in ids if j != tokenizer.unk_id]
     return ids
 
+
 def get_control(control, tokenizer, task):
     sep_id1 = 30665
     sep_id2 = 30666
@@ -1067,11 +1138,16 @@ def get_control(control, tokenizer, task):
                 keywords += convert_to_ids(tokenizer, keyword)
             if i != len(control['keywords']) - 1:
                 keywords += [sep_id1]
-        keywords = [tokenizer.begin_of_keyword_id] + keywords +[tokenizer.end_of_keyword_id]
+        keywords = [tokenizer.begin_of_keyword_id
+                    ] + keywords + [tokenizer.end_of_keyword_id]
 
-    if 'genre' in control and control['genre'] != "" and control['genre'] not in ['新闻', '学术', '公文', '诗歌', '武侠仙侠', '玄幻奇幻', '科幻灵异', '军事历史', '言情']:
+    if 'genre' in control and control['genre'] != "" and control[
+            'genre'] not in [
+                '新闻', '学术', '公文', '诗歌', '武侠仙侠', '玄幻奇幻', '科幻灵异', '军事历史', '言情'
+            ]:
         style = convert_to_ids(tokenizer, control['genre'])
-        style = [tokenizer.begin_of_style_id] + style +[tokenizer.end_of_style_id]
+        style = [tokenizer.begin_of_style_id
+                 ] + style + [tokenizer.end_of_style_id]
     else:
         style = []
 
@@ -1087,22 +1163,27 @@ def get_control(control, tokenizer, task):
                     relation += convert_to_ids(tokenizer, item)
                     if i != len(items) - 1:
                         relation += [sep_id1]
-                ids = [tokenizer.begin_of_relation_id] + relation + [tokenizer.end_of_relation_id]
+                ids = [tokenizer.begin_of_relation_id
+                       ] + relation + [tokenizer.end_of_relation_id]
                 relations = relations + ids
 
     events = []
     if 'events' in control and control['events'] != []:
         events_set = set()
         for i in control['events']:
-            event_join = "/".join([":".join(x) for x in sorted(i.items(), key=lambda x:x[0])])
+            event_join = "/".join(
+                [":".join(x) for x in sorted(i.items(), key=lambda x: x[0])])
             if event_join not in events_set:
                 events_set.add(event_join)
                 event = []
                 for idx, j in enumerate(i):
-                    event += convert_to_ids(tokenizer, j) + [sep_id2] + convert_to_ids(tokenizer, i[j])
+                    event += convert_to_ids(tokenizer,
+                                            j) + [sep_id2] + convert_to_ids(
+                                                tokenizer, i[j])
                     if idx != len(i) - 1:
                         event += [sep_id1]
-                ids = [tokenizer.begin_of_event_id] + event + [tokenizer.end_of_event_id]
+                ids = [tokenizer.begin_of_event_id
+                       ] + event + [tokenizer.end_of_event_id]
                 events = events + ids
 
     if task == 0:
@@ -1128,19 +1209,19 @@ def get_control(control, tokenizer, task):
         res = keywords + relations + events
     else:
         raise ValueError("task id error")
-    
+
     return res
 
 
 def encode(tokenizer, i, target_span_len=100, use_target=False):
     task_ids = {
-        'lm':0,
-        'compress':1,
-        'expand':2,
-        'rewrite':3,
-        'rewrite_s':4,
-        'compress_para':5,
-        'expand_para':6,
+        'lm': 0,
+        'compress': 1,
+        'expand': 2,
+        'rewrite': 3,
+        'rewrite_s': 4,
+        'compress_para': 5,
+        'expand_para': 6,
     }
 
     task = task_ids[i['mode']]
@@ -1155,14 +1236,14 @@ def encode(tokenizer, i, target_span_len=100, use_target=False):
 
     ids += control
     info.append(len(control))
-    
+
     assert len(i['source']) <= 2
     src = i['source'][0]
 
     src_ids = convert_to_ids(tokenizer, src)
     src_ids = [tokenizer.bos_id] + src_ids
     if task != 0:
-         src_ids += [tokenizer.eos_id]
+        src_ids += [tokenizer.eos_id]
     ids += src_ids
     info.append(len(src_ids))
     if not use_target:
@@ -1199,64 +1280,61 @@ def encode(tokenizer, i, target_span_len=100, use_target=False):
         ids += tgt_ids
         # 不控制长度，eos自己生成
         info.extend([len(tgt_ids), 0])
-        
+
     info = info[:1] + np.cumsum(info[1:]).tolist()
 
     assert len(ids) == info[-1]
-    assert len(info) % 2 == 1 # task, control, src, tgt, src, tgt, ...., end
+    assert len(info) % 2 == 1  # task, control, src, tgt, src, tgt, ...., end
     return ids, info
 
 
 def make_input_cpm3(ctx, info, prompt_length):
     task = info[0]
     len_ctx = len(ctx)
-    inp = np.arange((prompt_length+len_ctx), dtype = np.int64) + prompt_length * task
+    inp = np.arange(
+        (prompt_length + len_ctx), dtype=np.int64) + prompt_length * task
     inp[prompt_length:] = ctx[:len_ctx]
     len_inp = len(inp)
 
     info = [x + prompt_length for x in info[1:]]
     context_inp = np.full(len_inp, True)
     # 保证end一定能看见
-    for i in range(1, len(info)-1, 2):
-        context_inp[info[i]:info[i+1]] = False
-    
-    tgt = np.full((len_inp), -100, dtype = np.int64)
-    tgt[:-1] = np.where(
-        context_inp[1:],
-        -100,
-        inp[1:]
-    )
+    for i in range(1, len(info) - 1, 2):
+        context_inp[info[i]:info[i + 1]] = False
+
+    tgt = np.full((len_inp), -100, dtype=np.int64)
+    tgt[:-1] = np.where(context_inp[1:], -100, inp[1:])
 
-    position_inp = np.arange((len_inp), dtype = np.float32) / prompt_length
-    segment_inp = np.zeros((len_inp), dtype = np.int64)
+    position_inp = np.arange((len_inp), dtype=np.float32) / prompt_length
+    segment_inp = np.zeros((len_inp), dtype=np.int64)
 
     if task == 0:
         arr = [(2, info[0]), (1, 0), (1, info[-1])]
     else:
-        arr = [(2, info[0]), (2+task, info[1]), (1, info[-1])]
-    
+        arr = [(2, info[0]), (2 + task, info[1]), (1, info[-1])]
+
     last = prompt_length
     for (typ, end) in arr:
         if end > last:
             segment_inp[last:end] = typ
-            position_inp[last:end] = np.arange(end-last) / (end-last)
+            position_inp[last:end] = np.arange(end - last) / (end - last)
             last = end
     assert last == len_inp
 
     max_length = (len_inp + 2 - 1) // 2 * 2
 
-    _ctx = torch.zeros((max_length,), dtype=torch.long)
+    _ctx = torch.zeros((max_length, ), dtype=torch.long)
     _ctx[:len_inp] = torch.from_numpy(inp)[:len_inp].long()
-    _context = torch.full((max_length,), False, dtype=torch.bool)
+    _context = torch.full((max_length, ), False, dtype=torch.bool)
     _context[:len_inp] = torch.from_numpy(context_inp)[:len_inp].bool()
-    _position = torch.full((max_length,), False, dtype=torch.float)
+    _position = torch.full((max_length, ), False, dtype=torch.float)
     _position[:len_inp] = torch.from_numpy(position_inp)[:len_inp].float()
-    _segment = torch.full((max_length,), False, dtype=torch.long)
+    _segment = torch.full((max_length, ), False, dtype=torch.long)
     _segment[:len_inp] = torch.from_numpy(segment_inp)[:len_inp].long()
-    _tgt = torch.full((max_length,), -100, dtype=torch.long)
+    _tgt = torch.full((max_length, ), -100, dtype=torch.long)
     _tgt[:len_inp] = torch.from_numpy(tgt)[:len_inp].long()
 
-    _span = torch.zeros((max_length + 1,), dtype=torch.long)
+    _span = torch.zeros((max_length + 1, ), dtype=torch.long)
     _span[len_inp] = 1  # 每个拼接的句子结尾的后一位是1
     _span = torch.cumsum(_span, dim=-1)[:-1]
 
@@ -1272,34 +1350,38 @@ def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float("inf")):
 
     if top_k > 0:
         # Remove all tokens with a probability less than the last token of the top-k
-        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1,
+                                                                  None]
         logits[indices_to_remove] = filter_value
 
     batch_size = logits.size()[0]
     if top_p > 0.0:
-        logits=logits.view(batch_size, -1).contiguous()
+        logits = logits.view(batch_size, -1).contiguous()
         for index in range(len(logits)):
 
-            sorted_logits, sorted_indices = torch.sort(logits[index].view(-1), descending=True)
-            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+            sorted_logits, sorted_indices = torch.sort(logits[index].view(-1),
+                                                       descending=True)
+            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1),
+                                            dim=-1)
             # Remove tokens with cumulative probability above the threshold
             sorted_indices_to_remove = cumulative_probs > top_p
             # Shift the indices to the right to keep also the first token above the threshold
-            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+                ..., :-1].clone()
             sorted_indices_to_remove[..., 0] = 0
             indices_to_remove = sorted_indices[sorted_indices_to_remove]
             logits[index][indices_to_remove] = filter_value
 
-        logits=logits.view(batch_size, -1).contiguous()
+        logits = logits.view(batch_size, -1).contiguous()
 
     return logits
 
 
-def enforce_repetition_penalty_(tokenizer, 
-                                lprobs, 
-                                batch_size, 
-                                num_beams, 
-                                prev_output_tokens, 
+def enforce_repetition_penalty_(tokenizer,
+                                lprobs,
+                                batch_size,
+                                num_beams,
+                                prev_output_tokens,
                                 repetition_penalty,
                                 start_idx=None,
                                 end_idx=None,
@@ -1312,9 +1394,12 @@ def enforce_repetition_penalty_(tokenizer,
         else:
             if end_idx >= start_idx:
                 if window_size:
-                    output_tokens = prev_output_tokens[i][max(start_idx, end_idx + 1 - window_size): end_idx+1].tolist()
+                    output_tokens = prev_output_tokens[
+                        i][max(start_idx, end_idx + 1 - window_size):end_idx +
+                           1].tolist()
                 else:
-                    output_tokens = prev_output_tokens[i][start_idx: end_idx+1].tolist()
+                    output_tokens = prev_output_tokens[i][start_idx:end_idx +
+                                                          1].tolist()
             else:
                 output_tokens = []
         #print(output_tokens)
@@ -1333,7 +1418,8 @@ def _get_ngrams(ngram_size: int, prev_input_ids: torch.Tensor, num_hypos: int):
         generated_ngram = generated_ngrams[idx]
         for ngram in zip(*[gen_tokens[i:] for i in range(ngram_size)]):
             prev_ngram_tuple = tuple(ngram[:-1])
-            generated_ngram[prev_ngram_tuple] = generated_ngram.get(prev_ngram_tuple, []) + [ngram[-1]]
+            generated_ngram[prev_ngram_tuple] = generated_ngram.get(
+                prev_ngram_tuple, []) + [ngram[-1]]
     return generated_ngrams
 
 
@@ -1345,17 +1431,24 @@ def _get_generated_ngrams(banned_ngrams, prev_input_ids, ngram_size, cur_len):
     return banned_ngrams.get(ngram_idx, [])
 
 
-def calc_banned_ngram_tokens(
-    prev_input_ids: torch.Tensor, num_hypos: int, ngram_size: int, start_idx=None, end_idx=None, window_size=None, tokenizer=None):
+def calc_banned_ngram_tokens(prev_input_ids: torch.Tensor,
+                             num_hypos: int,
+                             ngram_size: int,
+                             start_idx=None,
+                             end_idx=None,
+                             window_size=None,
+                             tokenizer=None):
     """Copied from fairseq for no_repeat_ngram in beam_search"""
     if start_idx is not None and end_idx is not None:
         if window_size:
-            prev_input_ids = prev_input_ids[:, max(start_idx, end_idx + 1 - window_size): end_idx+1]
+            prev_input_ids = prev_input_ids[:,
+                                            max(start_idx, end_idx + 1 -
+                                                window_size):end_idx + 1]
         else:
-            prev_input_ids = prev_input_ids[:, start_idx: end_idx+1]
-        
+            prev_input_ids = prev_input_ids[:, start_idx:end_idx + 1]
+
     cur_len = prev_input_ids.size(1)
-    
+
     if cur_len + 1 < ngram_size:
         # return no banned tokens if we haven't generated no_repeat_ngram_size tokens yet
         return [[] for _ in range(num_hypos)]
@@ -1363,16 +1456,21 @@ def calc_banned_ngram_tokens(
     generated_ngrams = _get_ngrams(ngram_size, prev_input_ids, num_hypos)
 
     banned_tokens = [
-        _get_generated_ngrams(generated_ngrams[hypo_idx], prev_input_ids[hypo_idx], ngram_size, cur_len)
+        _get_generated_ngrams(generated_ngrams[hypo_idx],
+                              prev_input_ids[hypo_idx], ngram_size, cur_len)
         for hypo_idx in range(num_hypos)
     ]
 
     return banned_tokens
 
-def calc_banned_bad_words_ids(prev_input_ids, bad_words_ids, start_idx=None,  end_idx=None):
+
+def calc_banned_bad_words_ids(prev_input_ids,
+                              bad_words_ids,
+                              start_idx=None,
+                              end_idx=None):
     if start_idx is not None and end_idx is not None:
-        prev_input_ids = prev_input_ids[:, start_idx: end_idx+1]
-        
+        prev_input_ids = prev_input_ids[:, start_idx:end_idx + 1]
+
     banned_tokens = []
 
     def _tokens_match(prev_tokens, tokens):
@@ -1383,7 +1481,7 @@ def _tokens_match(prev_tokens, tokens):
             # if bad word tokens are longer then prev input_ids they can't be equal
             return False
 
-        if prev_tokens[-len(tokens) :] == tokens:
+        if prev_tokens[-len(tokens):] == tokens:
             # if tokens match
             return True
         else:
@@ -1393,11 +1491,13 @@ def _tokens_match(prev_tokens, tokens):
         banned_tokens_slice = []
 
         for banned_token_seq in bad_words_ids:
-            assert len(banned_token_seq) > 0, "Banned words token sequences {} cannot have an empty list".format(
-                bad_words_ids
-            )
+            assert len(
+                banned_token_seq
+            ) > 0, "Banned words token sequences {} cannot have an empty list".format(
+                bad_words_ids)
 
-            if _tokens_match(prev_input_ids_slice.tolist(), banned_token_seq[:-1]) is False:
+            if _tokens_match(prev_input_ids_slice.tolist(),
+                             banned_token_seq[:-1]) is False:
                 # if tokens do not match continue
                 continue
             banned_tokens_slice.append(banned_token_seq[-1])
@@ -1428,21 +1528,26 @@ def postprocess_next_token_scores(tokenizer,
 
     # repetition penalty (from CTRL paper https://arxiv.org/abs/1909.05858)
     if repetition_penalty != 1.0:
-        enforce_repetition_penalty_(
-            tokenizer, scores, batch_size, num_beams, input_ids, repetition_penalty, start_idx, end_idx, window_size
-        )
+        enforce_repetition_penalty_(tokenizer, scores, batch_size, num_beams,
+                                    input_ids, repetition_penalty, start_idx,
+                                    end_idx, window_size)
 
     if no_repeat_ngram_size > 0:
         # calculate a list of banned tokens to prevent repetitively generating the same ngrams
         num_batch_hypotheses = batch_size * num_beams
         # from fairseq: https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345
-        banned_batch_tokens = calc_banned_ngram_tokens(input_ids, num_batch_hypotheses, no_repeat_ngram_size, start_idx, end_idx, window_size, tokenizer)
+        banned_batch_tokens = calc_banned_ngram_tokens(input_ids,
+                                                       num_batch_hypotheses,
+                                                       no_repeat_ngram_size,
+                                                       start_idx, end_idx,
+                                                       window_size, tokenizer)
         for i, banned_tokens in enumerate(banned_batch_tokens):
             scores[i, banned_tokens] = -float("inf")
 
     if bad_words_ids is not None:
         # calculate a list of banned tokens according to bad words
-        banned_tokens = calc_banned_bad_words_ids(input_ids, bad_words_ids, start_idx, end_idx)
+        banned_tokens = calc_banned_bad_words_ids(input_ids, bad_words_ids,
+                                                  start_idx, end_idx)
 
         for i, banned_tokens in enumerate(banned_tokens):
             scores[i, banned_tokens] = -float("inf")
@@ -1451,40 +1556,62 @@ def postprocess_next_token_scores(tokenizer,
     scores[:, [0, 1, 2, 3, 4, 5] + [x for x in range(8, 20)]] = -float("inf")
 
     if start_idx is not None and end_idx is not None and min_len is not None:
-        min_length_constraint(scores, end_idx - start_idx + 2, min_len, tokenizer)
+        min_length_constraint(scores, end_idx - start_idx + 2, min_len,
+                              tokenizer)
 
     return scores
 
 
-def cpm_beam_search(model, tokenizer, instance, target_span_len = 100, beam_size = 3,
-                     temperature = .9, top_k = 0, top_p = 0.9,
-                     no_repeat_ngram_size = 0, repetition_penalty = 1.2, random_sample=False, min_len=None, **kwags):
+def cpm_beam_search(model,
+                    tokenizer,
+                    instance,
+                    target_span_len=100,
+                    beam_size=3,
+                    temperature=.9,
+                    top_k=0,
+                    top_p=0.9,
+                    no_repeat_ngram_size=0,
+                    repetition_penalty=1.2,
+                    random_sample=False,
+                    min_len=None,
+                    **kwags):
     print('tokenizer is', tokenizer)
     vocab_size = tokenizer.vocab_size
 
     ids, info = encode(tokenizer, instance, target_span_len)
 
     prompt_length = 64
-    input_tokens, input_length, context_input, position_input, segment_input, span_input, _ = make_input_cpm3(ids, info, prompt_length)
+    input_tokens, input_length, context_input, position_input, segment_input, span_input, _ = make_input_cpm3(
+        ids, info, prompt_length)
 
-    # (batch, max_length) 
+    # (batch, max_length)
     max_length = input_tokens.size(-1)
     batch_size = input_tokens.size(0)
 
-    # (batch, beam_size, max_length)    
-    input_tokens = input_tokens.unsqueeze(1).expand(batch_size, beam_size, max_length)
+    # (batch, beam_size, max_length)
+    input_tokens = input_tokens.unsqueeze(1).expand(batch_size, beam_size,
+                                                    max_length)
     input_length = input_length.unsqueeze(1).expand(batch_size, beam_size)
-    span_input = span_input.unsqueeze(1).expand(batch_size, beam_size, max_length)
-    context_input = context_input.unsqueeze(1).expand(batch_size, beam_size, max_length)
-    position_input = position_input.unsqueeze(1).expand(batch_size, beam_size, max_length)
-    segment_input = segment_input.unsqueeze(1).expand(batch_size, beam_size, max_length)
-    # (batch * beam_size, max_length)    
-    input_tokens = input_tokens.contiguous().view(batch_size * beam_size, max_length)
-    input_length = input_length.contiguous().view(batch_size * beam_size,)
-    span_input = span_input.contiguous().view(batch_size * beam_size, max_length)
-    context_input = context_input.contiguous().view(batch_size * beam_size, max_length)
-    position_input = position_input.contiguous().view(batch_size * beam_size, max_length)
-    segment_input = segment_input.contiguous().view(batch_size * beam_size, max_length)
+    span_input = span_input.unsqueeze(1).expand(batch_size, beam_size,
+                                                max_length)
+    context_input = context_input.unsqueeze(1).expand(batch_size, beam_size,
+                                                      max_length)
+    position_input = position_input.unsqueeze(1).expand(
+        batch_size, beam_size, max_length)
+    segment_input = segment_input.unsqueeze(1).expand(batch_size, beam_size,
+                                                      max_length)
+    # (batch * beam_size, max_length)
+    input_tokens = input_tokens.contiguous().view(batch_size * beam_size,
+                                                  max_length)
+    input_length = input_length.contiguous().view(batch_size * beam_size, )
+    span_input = span_input.contiguous().view(batch_size * beam_size,
+                                              max_length)
+    context_input = context_input.contiguous().view(batch_size * beam_size,
+                                                    max_length)
+    position_input = position_input.contiguous().view(batch_size * beam_size,
+                                                      max_length)
+    segment_input = segment_input.contiguous().view(batch_size * beam_size,
+                                                    max_length)
 
     input_tokens = input_tokens.int().cuda()
     input_length = input_length.int().cuda()
@@ -1495,9 +1622,11 @@ def cpm_beam_search(model, tokenizer, instance, target_span_len = 100, beam_size
 
     done = [False for _ in range(batch_size)]
     # (batch_size * beam_size, 0)
-    
-    beam_scores = torch.zeros((batch_size, beam_size), dtype=torch.float, device=input_tokens.device)
-    beam_scores[:, 1:] = -1e9 # 确保第一次只在一个vocab大小里选取
+
+    beam_scores = torch.zeros((batch_size, beam_size),
+                              dtype=torch.float,
+                              device=input_tokens.device)
+    beam_scores[:, 1:] = -1e9  # 确保第一次只在一个vocab大小里选取
     beam_scores = beam_scores.view(-1)
 
     # current position
@@ -1510,24 +1639,32 @@ def cpm_beam_search(model, tokenizer, instance, target_span_len = 100, beam_size
 
     # generated hypotheses
     generated_hyps = [
-        BeamHypotheses(beam_size, span_length, length_penalty=1 , early_stopping=False, tokenizer=tokenizer)
-        for _ in range(batch_size)
+        BeamHypotheses(beam_size,
+                       span_length,
+                       length_penalty=1,
+                       early_stopping=False,
+                       tokenizer=tokenizer) for _ in range(batch_size)
     ]
 
-
     with torch.inference_mode():
         past_key_values = None
         cached_attn_mask_pos_bias = None
         for i in range(lef - 1, rig):
             # skip all steps when we are done with each sentence
             if all(done):
-                break # Note: break not supports multi-GPUs
+                break  # Note: break not supports multi-GPUs
 
             if i == lef - 1:
                 # for the first time step, we will move the right context to the beginning inside model
-                logits, _, past_key_values, cached_attn_mask_pos_bias = model(input_tokens, input_length, context_input, position_input, segment_input, span_input, past_key_values, rig, i, cached_attn_mask_pos_bias)
+                logits, _, past_key_values, cached_attn_mask_pos_bias = model(
+                    input_tokens, input_length, context_input, position_input,
+                    segment_input, span_input, past_key_values, rig, i,
+                    cached_attn_mask_pos_bias)
             else:
-                logits, _, past_key_values, cached_attn_mask_pos_bias = model(input_tokens[:, i:i+1], input_length, context_input, position_input, segment_input, span_input, past_key_values, rig, i, cached_attn_mask_pos_bias)
+                logits, _, past_key_values, cached_attn_mask_pos_bias = model(
+                    input_tokens[:, i:i + 1], input_length, context_input,
+                    position_input, segment_input, span_input, past_key_values,
+                    rig, i, cached_attn_mask_pos_bias)
 
             logits = logits[:, -1, :]
 
@@ -1543,52 +1680,69 @@ def cpm_beam_search(model, tokenizer, instance, target_span_len = 100, beam_size
                 start_idx=lef,
                 end_idx=i,
                 window_size=None,
-                min_len=min_len
-            )
+                min_len=min_len)
             scores = F.log_softmax(logits, dim=-1)
 
             if random_sample:
                 assert temperature != 0, "temperature should not be zero!"
                 scores = scores - math.log(temperature)
                 _scores = scores + beam_scores[:, None].expand_as(scores)
-                             
+
                 _scores = top_k_logits(_scores, top_k=top_k, top_p=top_p)
-                _scores = _scores.contiguous().view(batch_size, beam_size * vocab_size)
+                _scores = _scores.contiguous().view(batch_size,
+                                                    beam_size * vocab_size)
                 # Sample 2 next tokens for each beam (so we have some spare tokens and match output of greedy beam search)
                 probs = F.softmax(_scores, dim=-1)
-                next_words = torch.multinomial(probs, num_samples=2 * beam_size)  # (batch_size, beam_size * 2)
+                next_words = torch.multinomial(
+                    probs,
+                    num_samples=2 * beam_size)  # (batch_size, beam_size * 2)
                 # Compute next scores
-                next_scores = torch.gather(_scores, -1, next_words)  # (batch_size, beam_size * 2)
+                next_scores = torch.gather(
+                    _scores, -1, next_words)  # (batch_size, beam_size * 2)
                 # sort the sampled vector to make sure that the first beam_size samples are the best
-                next_scores, next_scores_indices = torch.sort(next_scores, descending=True, dim=1)
-                next_words = torch.gather(next_words, -1, next_scores_indices)  # (batch_size, beam_size * 2)            
+                next_scores, next_scores_indices = torch.sort(next_scores,
+                                                              descending=True,
+                                                              dim=1)
+                next_words = torch.gather(
+                    next_words, -1,
+                    next_scores_indices)  # (batch_size, beam_size * 2)
             else:
-                next_scores = scores + beam_scores[:, None].expand_as(scores)  # (batch_size * beam_size, vocab_size)
+                next_scores = scores + beam_scores[:, None].expand_as(
+                    scores)  # (batch_size * beam_size, vocab_size)
 
                 # re-organize to group the beam together (we are keeping top hypothesis accross beams)
                 next_scores = next_scores.view(
-                    batch_size, beam_size * vocab_size
-                )  # (batch_size, beam_size * vocab_size)
+                    batch_size, beam_size *
+                    vocab_size)  # (batch_size, beam_size * vocab_size)
 
-                next_scores, next_words = torch.topk(next_scores, 2 * beam_size, dim=1, largest=True, sorted=True)
+                next_scores, next_words = torch.topk(next_scores,
+                                                     2 * beam_size,
+                                                     dim=1,
+                                                     largest=True,
+                                                     sorted=True)
 
-            assert next_scores.size() == next_words.size() == (batch_size, 2 * beam_size)
+            assert next_scores.size() == next_words.size() == (batch_size,
+                                                               2 * beam_size)
             # next batch beam content
             next_batch_beam = []
 
             for sent_id in range(batch_size):
 
-                 # if we are done with this sentence
-                done[sent_id] = done[sent_id] or generated_hyps[sent_id].is_done(next_scores[sent_id].max().item(), cur_len)
+                # if we are done with this sentence
+                done[sent_id] = done[
+                    sent_id] or generated_hyps[sent_id].is_done(
+                        next_scores[sent_id].max().item(), cur_len)
                 if done[sent_id]:
-                    next_batch_beam.extend([(0, tokenizer.pad_id, 0)] * beam_size)  # pad the batch
+                    next_batch_beam.extend([(0, tokenizer.pad_id, 0)] *
+                                           beam_size)  # pad the batch
                     continue
 
                 # next sentence beam content
                 next_sent_beam = []
 
                 # next words for this sentence
-                for idx, value in zip(next_words[sent_id], next_scores[sent_id]):
+                for idx, value in zip(next_words[sent_id],
+                                      next_scores[sent_id]):
 
                     # get beam and word IDs
                     beam_id = idx // vocab_size
@@ -1597,25 +1751,31 @@ def cpm_beam_search(model, tokenizer, instance, target_span_len = 100, beam_size
                     # end of sentence, or next word
                     if word_id == tokenizer.eos_id or cur_len == span_length:
                         if cur_len > 0:
-                            generated_hyps[sent_id].add(input_tokens[sent_id * beam_size + beam_id, lef:lef+cur_len].clone(), value.item())
+                            generated_hyps[sent_id].add(
+                                input_tokens[sent_id * beam_size + beam_id,
+                                             lef:lef + cur_len].clone(),
+                                value.item())
                     else:
-                        next_sent_beam.append((value, word_id, sent_id * beam_size + beam_id))
+                        next_sent_beam.append(
+                            (value, word_id, sent_id * beam_size + beam_id))
 
                     # the beam for next step is full
                     if len(next_sent_beam) == beam_size:
                         break
 
                 # update next beam content
-                assert len(next_sent_beam) == 0 if cur_len == span_length else beam_size
+                assert len(next_sent_beam
+                           ) == 0 if cur_len == span_length else beam_size
                 if len(next_sent_beam) == 0:
-                    next_sent_beam = [(0, tokenizer.pad_id, 0)] * beam_size  # pad the batch
+                    next_sent_beam = [(0, tokenizer.pad_id, 0)
+                                      ] * beam_size  # pad the batch
                 next_batch_beam.extend(next_sent_beam)
                 assert len(next_batch_beam) == beam_size * (sent_id + 1)
 
             # At the last step, we should not add the token to the next position
             if i == rig - 1:
                 break
-            
+
             # sanity check / prepare next batch
             assert len(next_batch_beam) == batch_size * beam_size
             beam_scores = beam_scores.new([x[0] for x in next_batch_beam])
@@ -1625,7 +1785,7 @@ def cpm_beam_search(model, tokenizer, instance, target_span_len = 100, beam_size
             # re-order batch and internal states
             input_tokens = input_tokens[beam_idx, :]
             input_tokens[:, lef + cur_len] = beam_words
-            
+
             for key_value_layer in past_key_values:
                 key_value_layer[0] = key_value_layer[0][beam_idx]
                 key_value_layer[1] = key_value_layer[1][beam_idx]
@@ -1639,13 +1799,13 @@ def cpm_beam_search(model, tokenizer, instance, target_span_len = 100, beam_size
         for i, hypotheses in enumerate(generated_hyps):
             best_hyp = max(hypotheses.hyp, key=lambda x: x[0])[1]
             best.append(best_hyp)
-        
+
         if instance['mode'] == 'lm':
             decode_start_idx = 0
         else:
-            decode_start_idx = 1  
+            decode_start_idx = 1
         return best[0][decode_start_idx:].cpu().numpy()
         # for id in best[0][decode_start_idx:].cpu().numpy():
         #     token = tokenizer.decode([id])
 
-        #     yield token
\ No newline at end of file
+        #     yield token
diff --git a/quickstart/glm_title_ch.py b/quickstart/glm_title_ch.py
old mode 100644
new mode 100755
index 5719f555..8717ff07
--- a/quickstart/glm_title_ch.py
+++ b/quickstart/glm_title_ch.py
@@ -5,6 +5,7 @@
 from flagai.auto_model.auto_loader import AutoLoader
 from flagai.model.predictor.predictor import Predictor
 
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 auto_loader = AutoLoader("title-generation",
diff --git a/quickstart/ner_ch.py b/quickstart/ner_ch.py
index fbf2ed5b..0d16e821 100644
--- a/quickstart/ner_ch.py
+++ b/quickstart/ner_ch.py
@@ -3,6 +3,9 @@
 # Licensed under the Apache License, Version 2.0 (the "License")
 from flagai.auto_model.auto_loader import AutoLoader
 from flagai.model.predictor.predictor import Predictor
+import torch 
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 task_name = "ner"
 model_name = "RoBERTa-base-ch"   # Load the checkpoints from Modelhub(model.baai.ac.cn/models)
@@ -14,6 +17,7 @@
                          class_num=len(target))
 
 model = auto_loader.get_model()
+model.to(device)
 tokenizer = auto_loader.get_tokenizer()
 
 predictor = Predictor(model, tokenizer)
diff --git a/quickstart/title_ch.py b/quickstart/title_ch.py
index 62e91cf6..7e0f141c 100644
--- a/quickstart/title_ch.py
+++ b/quickstart/title_ch.py
@@ -13,6 +13,7 @@
     model_name="RoBERTa-base-ch"   # Load the checkpoints from Modelhub(model.baai.ac.cn/models)
 )
 model = auto_loader.get_model()
+model.to(device)
 tokenizer = auto_loader.get_tokenizer()
 
 predictor = Predictor(model, tokenizer)
diff --git a/quickstart/title_en.py b/quickstart/title_en.py
index 2feb1281..77550df4 100644
--- a/quickstart/title_en.py
+++ b/quickstart/title_en.py
@@ -13,6 +13,7 @@
     model_name="BERT-base-en"
 )
 model = auto_loader.get_model()
+model.to(device)
 tokenizer = auto_loader.get_tokenizer()
 
 predictor = Predictor(model, tokenizer)
diff --git a/quickstart/writing_ch.py b/quickstart/writing_ch.py
index a289fb02..44432475 100644
--- a/quickstart/writing_ch.py
+++ b/quickstart/writing_ch.py
@@ -3,12 +3,16 @@
 # Licensed under the Apache License, Version 2.0 (the "License")
 from flagai.auto_model.auto_loader import AutoLoader
 from flagai.model.predictor.predictor import Predictor
+import torch 
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 if __name__ == '__main__':
     loader = AutoLoader("seq2seq",
                         "GPT2-base-ch",
                         model_dir="./state_dict/")
     model = loader.get_model()
+    model.to(device)
     tokenizer = loader.get_tokenizer()
     predictor = Predictor(model, tokenizer)
 
diff --git a/requirements.txt b/requirements.txt
index a797f8f6..e3bc8738 100755
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,4 +10,8 @@ datasets == 2.0.0
 setuptools == 59.5.0
 protobuf == 3.20.1
 ftfy == 6.1.1
-bmtrain
\ No newline at end of file
+Pillow == 9.2.0
+einops == 0.3.0
+diffusers == 0.7.2
+pytorch-lightning == 1.6.5
+taming-transformers-rom1504 == 0.0.6
diff --git a/setup.py b/setup.py
index f52cd405..fc3432ba 100755
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="flagai",
-    version="v1.3.2",
+    version="v1.4.0",
     description="FlagAI aims to help researchers and developers to freely train and test large-scale models for NLP/CV/VL tasks.",
     long_description=open("README.md", encoding="utf-8").read(),
     long_description_content_type="text/markdown",
@@ -29,6 +29,10 @@
         'setuptools==59.5.0',
         'protobuf==3.20.1',
         'ftfy == 6.1.1',
-        'bmtrain',
+        'Pillow == 9.2.0',
+        'einops == 0.3.0',
+        'diffusers == 0.7.2',
+        'pytorch-lightning == 1.6.5',
+        'taming-transformers-rom1504 == 0.0.6',
     ]
 )


Test patch:
diff --git a/examples/cpm3_generation/test.py b/examples/cpm3_generation/test.py
new file mode 100644
index 00000000..6596cd44
--- /dev/null
+++ b/examples/cpm3_generation/test.py
@@ -0,0 +1,21 @@
+from flagai.data.tokenizer.cpm_3.cpm3_tokenizer import CPM3Tokenizer
+from flagai.auto_model.auto_loader import AutoLoader
+from flagai.model.predictor.predictor import Predictor
+import torch
+import json
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+loader = AutoLoader("lm", "cpm3", model_dir="/sharefs/baai-mrnd/xw/")
+model = loader.get_model()
+tokenizer = CPM3Tokenizer('/sharefs/baai-mrnd/xw/cpm3/vocab.txt', space_token = '</_>', line_token = '</n>',)
+predictor = Predictor(model, tokenizer)
+model.to(device)
+text = { "mode": "lm", "source": ["半坡饰族斜挎包牛皮女包包女宽肩带撞色条纹小方包运动风单肩包红色，条纹风格逆袭，柔软牛皮，\
+    实用型收纳，时毗宽肩带，匠心工艺，深蓝，牛皮拉牌，大红，触感柔软，不同风格搭配的冲突感正是时毗的秘诀所在，于休闲中诠释自我。轻奢光泽坚韧耐用，1拉链暗袋，让代表性条纹装饰成为手袋的主角，实力与颜值集于一身。伴随着运动成为时尚生活方式，荔枝"],\
+         "targets": "", "control": { "keywords": [], "genre": "", "relations": [], "events": [] } }
+out_1 = predictor.predict_generate_beamsearch(text,
+                                            beam_size=3,
+                                            input_max_length=512,
+                                            out_max_length=100)
+print(out_1)
\ No newline at end of file