Merge pull request FlagAI-Open#143 from marscrazy/master

Simplify the loading of AltDiffusion models
richardsonjf · Nov 27, 2022 · 46d646c · 46d646c
2 parents db4c59b + 6356cca
commit 46d646c
Show file tree

Hide file tree

Showing 187 changed files with 76 additions and 65 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
diff --git a/CLA.md b/CLA.md
diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md
diff --git a/COMMITTERS.csv b/COMMITTERS.csv
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
diff --git a/GOVERNANCE.md b/GOVERNANCE.md
diff --git a/LICENSE b/LICENSE
diff --git a/SUPPORT.md b/SUPPORT.md
diff --git a/doc_zh/APPENDIX_GLM_IO.md b/doc_zh/APPENDIX_GLM_IO.md
diff --git a/doc_zh/APPENDIX_TASK.md b/doc_zh/APPENDIX_TASK.md
diff --git a/doc_zh/Advanced_Usage.md b/doc_zh/Advanced_Usage.md
diff --git a/doc_zh/DATASET_EXAMPLE.md b/doc_zh/DATASET_EXAMPLE.md
diff --git a/doc_zh/GLM.md b/doc_zh/GLM.md
diff --git a/doc_zh/TUTORIAL_10_METATRON.md b/doc_zh/TUTORIAL_10_METATRON.md
diff --git a/doc_zh/TUTORIAL_11_GLM_BLANK_FILLING_QA.md b/doc_zh/TUTORIAL_11_GLM_BLANK_FILLING_QA.md
diff --git a/doc_zh/TUTORIAL_12_GLM_EXAMPLE_TITLE_GENERATION.md b/doc_zh/TUTORIAL_12_GLM_EXAMPLE_TITLE_GENERATION.md
diff --git a/doc_zh/TUTORIAL_13_GLM_EXAMPLE_PEOTRY_GENERATION.md b/doc_zh/TUTORIAL_13_GLM_EXAMPLE_PEOTRY_GENERATION.md
diff --git a/doc_zh/TUTORIAL_14_HUGGINGFACE_T5.md b/doc_zh/TUTORIAL_14_HUGGINGFACE_T5.md
diff --git a/doc_zh/TUTORIAL_15_BERT_EXAMPLE_TITLE_GENERATION.md b/doc_zh/TUTORIAL_15_BERT_EXAMPLE_TITLE_GENERATION.md
diff --git a/doc_zh/TUTORIAL_16_BERT_EXAMPLE_SEMANTIC_MATCHING.md b/doc_zh/TUTORIAL_16_BERT_EXAMPLE_SEMANTIC_MATCHING.md
diff --git a/doc_zh/TUTORIAL_17_BERT_EXAMPLE_NER.md b/doc_zh/TUTORIAL_17_BERT_EXAMPLE_NER.md
diff --git a/doc_zh/TUTORIAL_18_GPT2_WRITING.md b/doc_zh/TUTORIAL_18_GPT2_WRITING.md
diff --git a/doc_zh/TUTORIAL_19_T5_EXAMPLE_TITLE_GENERATION.md b/doc_zh/TUTORIAL_19_T5_EXAMPLE_TITLE_GENERATION.md
diff --git a/doc_zh/TUTORIAL_1_TOKENIZER.md b/doc_zh/TUTORIAL_1_TOKENIZER.md
diff --git a/doc_zh/TUTORIAL_20_GLM_TNEWS.md b/doc_zh/TUTORIAL_20_GLM_TNEWS.md
diff --git a/doc_zh/TUTORIAL_20_SUPPORTED_TASKS_backup.md b/doc_zh/TUTORIAL_20_SUPPORTED_TASKS_backup.md
diff --git a/doc_zh/TUTORIAL_2_DATASET.md b/doc_zh/TUTORIAL_2_DATASET.md
diff --git a/doc_zh/TUTORIAL_3_MODEL.md b/doc_zh/TUTORIAL_3_MODEL.md
diff --git a/doc_zh/TUTORIAL_4_TRAINER.md b/doc_zh/TUTORIAL_4_TRAINER.md
diff --git a/doc_zh/TUTORIAL_5_INSTRUCTIONS_FOR_AutoLoader.md b/doc_zh/TUTORIAL_5_INSTRUCTIONS_FOR_AutoLoader.md
diff --git a/doc_zh/TUTORIAL_6_INSTRUCTIONS_FOR_PREDICTOR.md b/doc_zh/TUTORIAL_6_INSTRUCTIONS_FOR_PREDICTOR.md
diff --git a/doc_zh/TUTORIAL_7_PROMPT_LEARNING.md b/doc_zh/TUTORIAL_7_PROMPT_LEARNING.md
diff --git a/doc_zh/TUTORIAL_8_ENVIRONMENT_SETUP.md b/doc_zh/TUTORIAL_8_ENVIRONMENT_SETUP.md
diff --git a/doc_zh/TUTORIAL_9_SEQ2SEQ_METHOD.md b/doc_zh/TUTORIAL_9_SEQ2SEQ_METHOD.md
diff --git a/doc_zh/tokenization.md b/doc_zh/tokenization.md
diff --git a/docs/APPENDIX_GLM_IO.md b/docs/APPENDIX_GLM_IO.md
diff --git a/docs/Advanced_Usage.md b/docs/Advanced_Usage.md
diff --git a/docs/DATASET_EXAMPLE.md b/docs/DATASET_EXAMPLE.md
diff --git a/docs/GLM.md b/docs/GLM.md
diff --git a/docs/QuickTour.md b/docs/QuickTour.md
diff --git a/docs/TUTORIAL_10_MEGATRON.md b/docs/TUTORIAL_10_MEGATRON.md
diff --git a/docs/TUTORIAL_11_GLM_BLANK_FILLING_QA.md b/docs/TUTORIAL_11_GLM_BLANK_FILLING_QA.md
diff --git a/docs/TUTORIAL_12_GLM_EXAMPLE_TITLE_GENERATION.md b/docs/TUTORIAL_12_GLM_EXAMPLE_TITLE_GENERATION.md
diff --git a/docs/TUTORIAL_13_GLM_EXAMPLE_PEOTRY_GENERATION.md b/docs/TUTORIAL_13_GLM_EXAMPLE_PEOTRY_GENERATION.md
diff --git a/docs/TUTORIAL_14_HUGGINGFACE_T5.md b/docs/TUTORIAL_14_HUGGINGFACE_T5.md
diff --git a/docs/TUTORIAL_15_BERT_EXAMPLE_TITLE_GENERATION.md b/docs/TUTORIAL_15_BERT_EXAMPLE_TITLE_GENERATION.md
diff --git a/docs/TUTORIAL_16_BERT_EXAMPLE_SEMANTIC_MATCHING.md b/docs/TUTORIAL_16_BERT_EXAMPLE_SEMANTIC_MATCHING.md
diff --git a/docs/TUTORIAL_17_BERT_EXAMPLE_NER.md b/docs/TUTORIAL_17_BERT_EXAMPLE_NER.md
diff --git a/docs/TUTORIAL_18_GPT2_WRITING.md b/docs/TUTORIAL_18_GPT2_WRITING.md
diff --git a/docs/TUTORIAL_19_T5_EXAMPLE_TITLE_GENERATION.md b/docs/TUTORIAL_19_T5_EXAMPLE_TITLE_GENERATION.md
diff --git a/docs/TUTORIAL_1_TOKENIZER.md b/docs/TUTORIAL_1_TOKENIZER.md
diff --git a/docs/TUTORIAL_20_SUPPORTED_TASKS.md b/docs/TUTORIAL_20_SUPPORTED_TASKS.md
diff --git a/docs/TUTORIAL_2_DATASET.md b/docs/TUTORIAL_2_DATASET.md
diff --git a/docs/TUTORIAL_3_MODEL.md b/docs/TUTORIAL_3_MODEL.md
diff --git a/docs/TUTORIAL_4_TRAINER.md b/docs/TUTORIAL_4_TRAINER.md
diff --git a/docs/TUTORIAL_5_INSTRUCTIONS_FOR_AutoLoader.md b/docs/TUTORIAL_5_INSTRUCTIONS_FOR_AutoLoader.md
diff --git a/docs/TUTORIAL_6_INSTRUCTIONS_FOR_PREDICTOR.md b/docs/TUTORIAL_6_INSTRUCTIONS_FOR_PREDICTOR.md
diff --git a/docs/TUTORIAL_7_PROMPT_LEARNING.md b/docs/TUTORIAL_7_PROMPT_LEARNING.md
diff --git a/docs/TUTORIAL_8_ENVIRONMENT_SETUP.md b/docs/TUTORIAL_8_ENVIRONMENT_SETUP.md
diff --git a/docs/TUTORIAL_9_SEQ2SEQ_METHOD.md b/docs/TUTORIAL_9_SEQ2SEQ_METHOD.md
diff --git a/docs/img.png b/docs/img.png
diff --git a/docs/tokenization.md b/docs/tokenization.md
diff --git a/examples/AltDiffusion/README.md b/examples/AltDiffusion/README.md
diff --git a/examples/AltDiffusion/generate.py b/examples/AltDiffusion/generate.py
@@ -9,7 +9,7 @@
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 loader = AutoLoader(task_name="text2img", #contrastive learning
-                    model_name="AltDiffusion",
+                    model_name="AltDiffusion-m9",
                     model_dir="./checkpoints")
 
 model = loader.get_model()
@@ -18,4 +18,4 @@
 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"
-)
+)
diff --git a/examples/bert_title_generation_english/generate.py b/examples/bert_title_generation_english/generate.py
diff --git a/examples/cpm3_finetune/arguments.py b/examples/cpm3_finetune/arguments.py
diff --git a/examples/cpm3_finetune/finetune_cpm3.py b/examples/cpm3_finetune/finetune_cpm3.py
diff --git a/examples/cpm3_finetune/tune_cpm3.sh b/examples/cpm3_finetune/tune_cpm3.sh
diff --git a/examples/cpm3_generation/arguments.py b/examples/cpm3_generation/arguments.py
diff --git a/examples/cpm3_generation/generation.py b/examples/cpm3_generation/generation.py
diff --git a/examples/cpm3_generation/infer.py b/examples/cpm3_generation/infer.py
diff --git a/examples/cpm3_generation/infer.sh b/examples/cpm3_generation/infer.sh
diff --git a/examples/cpm3_pretrain/arguments.py b/examples/cpm3_pretrain/arguments.py
diff --git a/examples/cpm3_pretrain/pretrain_cpm3.py b/examples/cpm3_pretrain/pretrain_cpm3.py
diff --git a/examples/cpm3_pretrain/run_cpm3.sh b/examples/cpm3_pretrain/run_cpm3.sh
diff --git a/examples/glm_blank_filling/glm_generate_samples.py b/examples/glm_blank_filling/glm_generate_samples.py
diff --git a/examples/glm_poetry_generation/generate.py b/examples/glm_poetry_generation/generate.py
diff --git a/examples/glm_title_generation/generate.py b/examples/glm_title_generation/generate.py
diff --git a/examples/gpt2_title_generation/generate.py b/examples/gpt2_title_generation/generate.py
diff --git a/examples/roberta_ner/data/china-people-daily-ner-corpus/example.dev b/examples/roberta_ner/data/china-people-daily-ner-corpus/example.dev
diff --git a/examples/roberta_ner/data/china-people-daily-ner-corpus/example.test b/examples/roberta_ner/data/china-people-daily-ner-corpus/example.test
diff --git a/examples/roberta_ner/data/china-people-daily-ner-corpus/example.train b/examples/roberta_ner/data/china-people-daily-ner-corpus/example.train
diff --git a/examples/roberta_ner/generate.py b/examples/roberta_ner/generate.py
diff --git a/examples/roberta_ner/generate_crf.py b/examples/roberta_ner/generate_crf.py
diff --git a/examples/roberta_ner/generate_global_pointer.py b/examples/roberta_ner/generate_global_pointer.py
diff --git a/examples/roberta_semantic_matching/generate.py b/examples/roberta_semantic_matching/generate.py
diff --git a/examples/roberta_title_generation/generate.py b/examples/roberta_title_generation/generate.py
diff --git a/flagai/auto_model/auto_loader.py b/flagai/auto_model/auto_loader.py
@@ -109,9 +109,9 @@ def __getattr__(self, name):
     "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"],
+    ["flagai.model.mm.diffusion", "LatentDiffusion", "diffusion", "mm","flagai.model.mm.AltCLIP", "AltCLIPProcess"],
     "altdiffusion-m9":
-    ["flagai.model.mm.diffusion", "LatentDiffusion", "diffusion", "mm"],
+    ["flagai.model.mm.diffusion", "LatentDiffusion", "diffusion", "mm","flagai.model.mm.AltCLIP", "AltCLIPProcess"],
     "swinv1-base-patch4-window7-224":
     ["flagai.model.vision.swinv1", "SwinTransformer", "swinv1", "vision"],
     "swinv2-base-patch4-window8-256":
@@ -212,7 +212,10 @@ def __init__(self,
 
         elif model_type == "mm":
             if model_name.startswith("altdiffusion"):
-                self.tokenizer = None
+                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.tokenizer = self.process.tokenizer
+                self.model.tokenizer = self.tokenizer
             elif "altclip" not in model_name:
                 from flagai.data.tokenizer.clip.tokenizer import ClipTokenizer
                 self.tokenizer = ClipTokenizer(bpe_path=os.path.join(download_path, 'bpe_simple_vocab_16e6.txt.gz'))

diff --git a/flagai/data/dataset/cpm3_data/__init__.py b/flagai/data/dataset/cpm3_data/__init__.py
diff --git a/flagai/data/dataset/cpm3_data/dataset.py b/flagai/data/dataset/cpm3_data/dataset.py
diff --git a/flagai/data/dataset/cpm3_data/distributed_indexed.py b/flagai/data/dataset/cpm3_data/distributed_indexed.py
diff --git a/flagai/data/dataset/cpm3_data/indexed.py b/flagai/data/dataset/cpm3_data/indexed.py
diff --git a/flagai/data/tokenizer/cpm_1/cpm1_tokenizer.py b/flagai/data/tokenizer/cpm_1/cpm1_tokenizer.py
diff --git a/flagai/data/tokenizer/cpm_3/__init__.py b/flagai/data/tokenizer/cpm_3/__init__.py
diff --git a/flagai/data/tokenizer/cpm_3/cpm3_tokenizer.py b/flagai/data/tokenizer/cpm_3/cpm3_tokenizer.py
diff --git a/flagai/data/tokenizer/uni_tokenizer/base_tokenizer.py b/flagai/data/tokenizer/uni_tokenizer/base_tokenizer.py
diff --git a/flagai/data/tokenizer/uni_tokenizer/difffusion_bert_tokenizer.py b/flagai/data/tokenizer/uni_tokenizer/difffusion_bert_tokenizer.py
diff --git a/flagai/data/tokenizer/uni_tokenizer/tokenizer.py b/flagai/data/tokenizer/uni_tokenizer/tokenizer.py
diff --git a/flagai/data/tokenizer/uni_tokenizer/wp_tokenizer.py b/flagai/data/tokenizer/uni_tokenizer/wp_tokenizer.py
diff --git a/flagai/model/base_model.py b/flagai/model/base_model.py
@@ -1,12 +1,11 @@
 # Copyright © 2022 BAAI. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License")
-from sklearn.linear_model import HuberRegressor
 from torch.nn import Module
 import torch
 import json
 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
+from flagai.model.file_utils import _get_model_id, _get_checkpoint_path, _get_vocab_path, _get_model_files
 import os
 
 
@@ -46,10 +45,13 @@ def _load_state_dict_into_model(cls,
                                     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 "state_dict" in pl_sd:
+            sd = pl_sd["state_dict"]
+        else:
+            sd = pl_sd
         if "global_step" in pl_sd:
             print(f"Global Step: {pl_sd['global_step']}")
-        m, u = model.load_state_dict(sd, strict=False)
+        m, u = model.load_state_dict(sd, strict=True)
         if len(m) > 0 and verbose:
             print("missing keys:")
             print(m)
@@ -113,7 +115,7 @@ def load_local(checkpoint_path):
                     model.load_weights(checkpoint_path)
             return model
 
-        def load_diffusion_local(yaml_path):
+        def load_diffusion_local(yaml_path, only_download_config=False):
             """
             Now only diffusion models requires yaml
             """
@@ -126,19 +128,19 @@ def load_diffusion_local(yaml_path):
             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,
-            )
+            if not only_download_config:
+                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)
+            return load_diffusion_local(yaml_path, only_download_config=only_download_config)
         elif os.path.exists(config_path):
             """
             It is fine when checkpoint_path does not exist, for the case that only_download_config=True
@@ -202,7 +204,7 @@ def load_diffusion_local(yaml_path):
                         checkpoint_merge,
                         os.path.join(download_path, "pytorch_model.bin"))
         if os.path.exists(yaml_path):
-            return load_diffusion_local(yaml_path)
+            return load_diffusion_local(yaml_path,only_download_config=only_download_config)
         return load_local(checkpoint_path)
 
     @classmethod

diff --git a/flagai/model/bert_model.py b/flagai/model/bert_model.py
diff --git a/flagai/model/blocks/cpm_block.py b/flagai/model/blocks/cpm_block.py
diff --git a/flagai/model/cpm3_model.py b/flagai/model/cpm3_model.py
diff --git a/flagai/model/cpm3_train_model.py b/flagai/model/cpm3_train_model.py
diff --git a/flagai/model/layers/attentions.py b/flagai/model/layers/attentions.py
diff --git a/flagai/model/layers/embeddings.py b/flagai/model/layers/embeddings.py
diff --git a/flagai/model/layers/feedforward.py b/flagai/model/layers/feedforward.py
diff --git a/flagai/model/layers/feedforward_bmt.py b/flagai/model/layers/feedforward_bmt.py
diff --git a/flagai/model/layers/layer_norm.py b/flagai/model/layers/layer_norm.py
diff --git a/flagai/model/layers/linear.py b/flagai/model/layers/linear.py
diff --git a/flagai/model/mm/AltCLIP.py b/flagai/model/mm/AltCLIP.py
@@ -7,7 +7,6 @@
 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,

diff --git a/flagai/model/mm/AltDiffusion.py b/flagai/model/mm/AltDiffusion.py
@@ -14,7 +14,6 @@
 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',
@@ -562,6 +561,8 @@ def __init__(self,
         self.instantiate_first_stage(first_stage_config)
         self.instantiate_cond_stage(cond_stage_config)
         self.cond_stage_forward = cond_stage_forward
+        if self.cond_stage_forward is None:
+            self.set_cond_stage_forward()
         self.clip_denoised = False
         self.bbox_tokenizer = None
 
@@ -619,47 +620,13 @@ def instantiate_first_stage(self, config):
         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
-
+        model = instantiate_from_config(config)
+        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
+
     def _get_denoise_row_from_list(self,
                                    samples,
                                    desc='',
@@ -687,7 +654,23 @@ def get_first_stage_encoding(self, encoder_posterior):
 
             z = encoder_posterior.sample()
         return self.scale_factor * z
-
+    def set_cond_stage_forward(self):
+        def func(c):
+            device = next(self.cond_stage_model.parameters()).device
+            text = self.tokenizer(c,
+                            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.cond_stage_model(**text)
+            return features['projection_state']
+        self.cond_stage_forward = func
     def get_learned_conditioning(self, c):
         # C will be directly returned if it is None
         if c is None:
@@ -701,8 +684,7 @@ def get_learned_conditioning(self, c):
             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)
+            c = self.cond_stage_forward(c)
         return c
 
     def meshgrid(self, h, w):

diff --git a/flagai/model/mm/Sampler.py b/flagai/model/mm/Sampler.py
diff --git a/flagai/model/mm/Unets/Unet.py b/flagai/model/mm/Unets/Unet.py
diff --git a/flagai/model/mm/__init__.py b/flagai/model/mm/__init__.py
diff --git a/flagai/model/mm/attentions/attention.py b/flagai/model/mm/attentions/attention.py
diff --git a/flagai/model/mm/autoencoders.py b/flagai/model/mm/autoencoders.py
diff --git a/flagai/model/mm/clip_guohua/__init__.py b/flagai/model/mm/clip_guohua/__init__.py
diff --git a/flagai/model/mm/clip_guohua/bert_tokenizer.py b/flagai/model/mm/clip_guohua/bert_tokenizer.py
diff --git a/flagai/model/mm/clip_guohua/model.py b/flagai/model/mm/clip_guohua/model.py
diff --git a/flagai/model/mm/clip_guohua/modeling_bert.py b/flagai/model/mm/clip_guohua/modeling_bert.py
diff --git a/flagai/model/mm/clip_model.py b/flagai/model/mm/clip_model.py
diff --git a/flagai/model/mm/lm/clip_guohua.py b/flagai/model/mm/lm/clip_guohua.py
diff --git a/flagai/model/mm/model.py b/flagai/model/mm/model.py
diff --git a/flagai/model/mm/modeling_berts.py b/flagai/model/mm/modeling_berts.py
@@ -27,7 +27,32 @@ class BertSeriesModelWithTransformation(BertPreTrainedModel):
     _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
     config_class = BertSeriesConfig
 
-    def __init__(self, config):
+    def __init__(self, config=None, **kargs):
+        # modify initialization for autoloading 
+        if config is None:
+            config = XLMRobertaConfig()
+            config.attention_probs_dropout_prob= 0.1
+            config.bos_token_id=0
+            config.eos_token_id=2
+            config.hidden_act='gelu'
+            config.hidden_dropout_prob=0.1
+            config.hidden_size=1024
+            config.initializer_range=0.02
+            config.intermediate_size=4096
+            config.layer_norm_eps=1e-05
+            config.max_position_embeddings=514
+
+            config.num_attention_heads=16
+            config.num_hidden_layers=24
+            config.output_past=True
+            config.pad_token_id=1
+            config.position_embedding_type= "absolute"
+
+            config.type_vocab_size= 1
+            config.use_cache=True
+            config.vocab_size= 250002
+            config.project_dim = 768
+            config.learn_encoder = False
         super().__init__(config)
         if config.model_type == 'bert':
             self.bert = BertModel(config)

diff --git a/flagai/model/mm/utils.py b/flagai/model/mm/utils.py
diff --git a/flagai/model/predictor/gpt.py b/flagai/model/predictor/gpt.py
diff --git a/flagai/model/predictor/predictor.py b/flagai/model/predictor/predictor.py
diff --git a/flagai/model/predictor/utils.py b/flagai/model/predictor/utils.py
diff --git a/flagai/model/vision/layers/__init__.py b/flagai/model/vision/layers/__init__.py
diff --git a/flagai/model/vision/layers/activations.py b/flagai/model/vision/layers/activations.py
diff --git a/flagai/model/vision/layers/activations_jit.py b/flagai/model/vision/layers/activations_jit.py
diff --git a/flagai/model/vision/layers/activations_me.py b/flagai/model/vision/layers/activations_me.py
diff --git a/flagai/model/vision/layers/adaptive_avgmax_pool.py b/flagai/model/vision/layers/adaptive_avgmax_pool.py
diff --git a/flagai/model/vision/layers/attention_pool2d.py b/flagai/model/vision/layers/attention_pool2d.py
diff --git a/flagai/model/vision/layers/blur_pool.py b/flagai/model/vision/layers/blur_pool.py
diff --git a/flagai/model/vision/layers/bottleneck_attn.py b/flagai/model/vision/layers/bottleneck_attn.py
diff --git a/flagai/model/vision/layers/cbam.py b/flagai/model/vision/layers/cbam.py
diff --git a/flagai/model/vision/layers/classifier.py b/flagai/model/vision/layers/classifier.py
diff --git a/flagai/model/vision/layers/cond_conv2d.py b/flagai/model/vision/layers/cond_conv2d.py
diff --git a/flagai/model/vision/layers/config.py b/flagai/model/vision/layers/config.py
diff --git a/flagai/model/vision/layers/conv2d_same.py b/flagai/model/vision/layers/conv2d_same.py
diff --git a/flagai/model/vision/layers/conv_bn_act.py b/flagai/model/vision/layers/conv_bn_act.py
diff --git a/flagai/model/vision/layers/create_act.py b/flagai/model/vision/layers/create_act.py
diff --git a/flagai/model/vision/layers/create_attn.py b/flagai/model/vision/layers/create_attn.py
diff --git a/flagai/model/vision/layers/create_conv2d.py b/flagai/model/vision/layers/create_conv2d.py
diff --git a/flagai/model/vision/layers/create_norm_act.py b/flagai/model/vision/layers/create_norm_act.py
diff --git a/flagai/model/vision/layers/drop.py b/flagai/model/vision/layers/drop.py
diff --git a/flagai/model/vision/layers/eca.py b/flagai/model/vision/layers/eca.py
diff --git a/flagai/model/vision/layers/evo_norm.py b/flagai/model/vision/layers/evo_norm.py
diff --git a/flagai/model/vision/layers/filter_response_norm.py b/flagai/model/vision/layers/filter_response_norm.py
diff --git a/flagai/model/vision/layers/gather_excite.py b/flagai/model/vision/layers/gather_excite.py
diff --git a/flagai/model/vision/layers/global_context.py b/flagai/model/vision/layers/global_context.py
diff --git a/flagai/model/vision/layers/halo_attn.py b/flagai/model/vision/layers/halo_attn.py
diff --git a/flagai/model/vision/layers/helpers.py b/flagai/model/vision/layers/helpers.py
diff --git a/flagai/model/vision/layers/inplace_abn.py b/flagai/model/vision/layers/inplace_abn.py
diff --git a/flagai/model/vision/layers/lambda_layer.py b/flagai/model/vision/layers/lambda_layer.py
diff --git a/flagai/model/vision/layers/linear.py b/flagai/model/vision/layers/linear.py
diff --git a/flagai/model/vision/layers/median_pool.py b/flagai/model/vision/layers/median_pool.py
diff --git a/flagai/model/vision/layers/mixed_conv2d.py b/flagai/model/vision/layers/mixed_conv2d.py
diff --git a/flagai/model/vision/layers/ml_decoder.py b/flagai/model/vision/layers/ml_decoder.py
diff --git a/flagai/model/vision/layers/mlp.py b/flagai/model/vision/layers/mlp.py
diff --git a/flagai/model/vision/layers/non_local_attn.py b/flagai/model/vision/layers/non_local_attn.py
diff --git a/flagai/model/vision/layers/norm.py b/flagai/model/vision/layers/norm.py
diff --git a/flagai/model/vision/layers/norm_act.py b/flagai/model/vision/layers/norm_act.py
diff --git a/flagai/model/vision/layers/padding.py b/flagai/model/vision/layers/padding.py
diff --git a/flagai/model/vision/layers/patch_embed.py b/flagai/model/vision/layers/patch_embed.py
diff --git a/flagai/model/vision/layers/pool2d_same.py b/flagai/model/vision/layers/pool2d_same.py
diff --git a/flagai/model/vision/layers/pos_embed.py b/flagai/model/vision/layers/pos_embed.py
diff --git a/flagai/model/vision/layers/selective_kernel.py b/flagai/model/vision/layers/selective_kernel.py
diff --git a/flagai/model/vision/layers/separable_conv.py b/flagai/model/vision/layers/separable_conv.py
diff --git a/flagai/model/vision/layers/space_to_depth.py b/flagai/model/vision/layers/space_to_depth.py
diff --git a/flagai/model/vision/layers/split_attn.py b/flagai/model/vision/layers/split_attn.py
diff --git a/flagai/model/vision/layers/split_batchnorm.py b/flagai/model/vision/layers/split_batchnorm.py
diff --git a/flagai/model/vision/layers/squeeze_excite.py b/flagai/model/vision/layers/squeeze_excite.py
diff --git a/flagai/model/vision/layers/std_conv.py b/flagai/model/vision/layers/std_conv.py
diff --git a/flagai/model/vision/layers/test_time_pool.py b/flagai/model/vision/layers/test_time_pool.py
diff --git a/flagai/model/vision/layers/trace_utils.py b/flagai/model/vision/layers/trace_utils.py
diff --git a/flagai/model/vision/layers/weight_init.py b/flagai/model/vision/layers/weight_init.py
diff --git a/flagai_wechat.png b/flagai_wechat.png
diff --git a/logo.png b/logo.png
diff --git a/prepare_test.sh b/prepare_test.sh
diff --git a/quickstart/glm_title_ch.py b/quickstart/glm_title_ch.py
diff --git a/requirements.txt b/requirements.txt
diff --git a/setup.cfg b/setup.cfg
@@ -1,3 +1,3 @@
 [easy_install]
 
-index_url = https://pypi.tuna.tsinghua.edu.cn/simple
+index_url = https://mirrors.aliyun.com/pypi/simple/
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="flagai",
-    version="v1.4.4",
+    version="v1.4.5",
     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",

diff --git a/test.py b/test.py