modelscope/DiffSynth-Painter

1import torch2from transformers import T5EncoderModel, T5Config3from .sd_text_encoder import SDTextEncoder4 5 6class FluxTextEncoder1(SDTextEncoder):7    def __init__(self, vocab_size=49408):8        super().__init__(vocab_size=vocab_size)9 10    def forward(self, input_ids, clip_skip=2):11        embeds = self.token_embedding(input_ids) + self.position_embeds12        attn_mask = self.attn_mask.to(device=embeds.device, dtype=embeds.dtype)13        for encoder_id, encoder in enumerate(self.encoders):14            embeds = encoder(embeds, attn_mask=attn_mask)15            if encoder_id + clip_skip == len(self.encoders):16                hidden_states = embeds17        embeds = self.final_layer_norm(embeds)18        pooled_embeds = embeds[torch.arange(embeds.shape[0]), input_ids.to(dtype=torch.int).argmax(dim=-1)]19        return embeds, pooled_embeds20 21    @staticmethod22    def state_dict_converter():23        return FluxTextEncoder1StateDictConverter()24 25 26 27class FluxTextEncoder2(T5EncoderModel):28    def __init__(self, config):29        super().__init__(config)30        self.eval()31 32    def forward(self, input_ids):33        outputs = super().forward(input_ids=input_ids)34        prompt_emb = outputs.last_hidden_state35        return prompt_emb36 37    @staticmethod38    def state_dict_converter():39        return FluxTextEncoder2StateDictConverter()40 41 42 43class FluxTextEncoder1StateDictConverter:44    def __init__(self):45        pass46 47    def from_diffusers(self, state_dict):48        rename_dict = {49            "text_model.embeddings.token_embedding.weight": "token_embedding.weight",50            "text_model.embeddings.position_embedding.weight": "position_embeds",51            "text_model.final_layer_norm.weight": "final_layer_norm.weight",52            "text_model.final_layer_norm.bias": "final_layer_norm.bias"53        }54        attn_rename_dict = {55            "self_attn.q_proj": "attn.to_q",56            "self_attn.k_proj": "attn.to_k",57            "self_attn.v_proj": "attn.to_v",58            "self_attn.out_proj": "attn.to_out",59            "layer_norm1": "layer_norm1",60            "layer_norm2": "layer_norm2",61            "mlp.fc1": "fc1",62            "mlp.fc2": "fc2",63        }64        state_dict_ = {}65        for name in state_dict:66            if name in rename_dict:67                param = state_dict[name]68                if name == "text_model.embeddings.position_embedding.weight":69                    param = param.reshape((1, param.shape[0], param.shape[1]))70                state_dict_[rename_dict[name]] = param71            elif name.startswith("text_model.encoder.layers."):72                param = state_dict[name]73                names = name.split(".")74                layer_id, layer_type, tail = names[3], ".".join(names[4:-1]), names[-1]75                name_ = ".".join(["encoders", layer_id, attn_rename_dict[layer_type], tail])76                state_dict_[name_] = param77        return state_dict_78    79    def from_civitai(self, state_dict):80        return self.from_diffusers(state_dict)81 82 83 84class FluxTextEncoder2StateDictConverter():85    def __init__(self):86        pass87 88    def from_diffusers(self, state_dict):89        state_dict_ = state_dict90        return state_dict_91 92    def from_civitai(self, state_dict):93        return self.from_diffusers(state_dict)94