mirror of
https://github.com/modelscope/DiffSynth-Studio.git
synced 2026-03-19 06:48:12 +00:00
45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
'''
|
|
* Adapted from BLIP (https://github.com/salesforce/BLIP)
|
|
'''
|
|
|
|
import transformers
|
|
transformers.logging.set_verbosity_error()
|
|
|
|
from torch import nn
|
|
import os
|
|
from .med import BertConfig, BertModel
|
|
from .blip import create_vit, init_tokenizer
|
|
|
|
class BLIP_Pretrain(nn.Module):
|
|
def __init__(self,
|
|
med_config = "med_config.json",
|
|
image_size = 224,
|
|
vit = 'base',
|
|
vit_grad_ckpt = False,
|
|
vit_ckpt_layer = 0,
|
|
embed_dim = 256,
|
|
queue_size = 57600,
|
|
momentum = 0.995,
|
|
bert_model_path = ""
|
|
):
|
|
"""
|
|
Args:
|
|
med_config (str): path for the mixture of encoder-decoder model's configuration file
|
|
image_size (int): input image size
|
|
vit (str): model size of vision transformer
|
|
"""
|
|
super().__init__()
|
|
|
|
self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer, 0)
|
|
|
|
self.tokenizer = init_tokenizer(bert_model_path)
|
|
encoder_config = BertConfig.from_json_file(med_config)
|
|
encoder_config.encoder_width = vision_width
|
|
self.text_encoder = BertModel(config=encoder_config, add_pooling_layer=False)
|
|
|
|
text_width = self.text_encoder.config.hidden_size
|
|
|
|
self.vision_proj = nn.Linear(vision_width, embed_dim)
|
|
self.text_proj = nn.Linear(text_width, embed_dim)
|
|
|