feat: add text encoder files (mindspore-lab#1481)

- Add 8 text encoder files: bert, flux, hunyuan_video, llama, qwen_image, qwen_vl, sd3_clip, t5

Signed-off-by: vigo999 <zwiori1982@163.com>

Author: vigo999

mindone/comfyui/comfy/text_encoders/bert.py

@@ -0,0 +1,200 @@
+import comfy.ops
+from comfy.ldm.modules.attention import optimized_attention_for_device
+from mindspore_patch.utils import dtype_to_max
+
+import mindspore
+from mindspore import mint
+
+
+class BertAttention(mindspore.nn.Cell):
+    def __init__(self, embed_dim, heads, dtype, device, operations):
+        super().__init__()
+
+        self.heads = heads
+        self.query = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=None)
+        self.key = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=None)
+        self.value = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=None)
+
+    def construct(self, x, mask=None, optimized_attention=None):
+        q = self.query(x)
+        k = self.key(x)
+        v = self.value(x)
+
+        out = optimized_attention(q, k, v, self.heads, mask)
+        return out
+
+
+class BertOutput(mindspore.nn.Cell):
+    def __init__(self, input_dim, output_dim, layer_norm_eps, dtype, device, operations):
+        super().__init__()
+        self.dense = operations.Linear(input_dim, output_dim, dtype=dtype, device=None)
+        self.LayerNorm = operations.LayerNorm(output_dim, eps=layer_norm_eps, dtype=dtype, device=None)
+        # self.dropout = nn.Dropout(0.0)
+
+    def construct(self, x, y):
+        x = self.dense(x)
+        # hidden_states = self.dropout(hidden_states)
+        x = self.LayerNorm(x + y)
+        return x
+
+
+class BertAttentionBlock(mindspore.nn.Cell):
+    def __init__(self, embed_dim, heads, layer_norm_eps, dtype, device, operations):
+        super().__init__()
+        self.self = BertAttention(embed_dim, heads, dtype, None, operations)
+        self.output = BertOutput(embed_dim, embed_dim, layer_norm_eps, dtype, None, operations)
+
+    def construct(self, x, mask, optimized_attention):
+        y = self.self(x, mask, optimized_attention)
+        return self.output(y, x)
+
+
+class BertIntermediate(mindspore.nn.Cell):
+    def __init__(self, embed_dim, intermediate_dim, dtype, device, operations):
+        super().__init__()
+        self.dense = operations.Linear(embed_dim, intermediate_dim, dtype=dtype, device=None)
+
+    def construct(self, x):
+        x = self.dense(x)
+        return mint.functional.gelu(x)
+
+
+class BertBlock(mindspore.nn.Cell):
+    def __init__(self, embed_dim, intermediate_dim, heads, layer_norm_eps, dtype, device, operations):
+        super().__init__()
+        self.attention = BertAttentionBlock(embed_dim, heads, layer_norm_eps, dtype, None, operations)
+        self.intermediate = BertIntermediate(embed_dim, intermediate_dim, dtype, None, operations)
+        self.output = BertOutput(intermediate_dim, embed_dim, layer_norm_eps, dtype, None, operations)
+
+    def construct(self, x, mask, optimized_attention):
+        x = self.attention(x, mask, optimized_attention)
+        y = self.intermediate(x)
+        return self.output(y, x)
+
+
+class BertEncoder(mindspore.nn.Cell):
+    def __init__(self, num_layers, embed_dim, intermediate_dim, heads, layer_norm_eps, dtype, device, operations):
+        super().__init__()
+        self.layer = mindspore.nn.CellList(
+            [
+                BertBlock(embed_dim, intermediate_dim, heads, layer_norm_eps, dtype, None, operations)
+                for i in range(num_layers)
+            ]
+        )
+
+    def construct(self, x, mask=None, intermediate_output=None):
+        optimized_attention = optimized_attention_for_device(None, mask=mask is not None, small_input=True)
+
+        if intermediate_output is not None:
+            if intermediate_output < 0:
+                intermediate_output = len(self.layer) + intermediate_output
+
+        intermediate = None
+        for i, l in enumerate(self.layer):
+            x = l(x, mask, optimized_attention)
+            if i == intermediate_output:
+                intermediate = x.clone()
+        return x, intermediate
+
+
+class BertEmbeddings(mindspore.nn.Cell):
+    def __init__(
+        self,
+        vocab_size,
+        max_position_embeddings,
+        type_vocab_size,
+        pad_token_id,
+        embed_dim,
+        layer_norm_eps,
+        dtype,
+        device,
+        operations,
+    ):
+        super().__init__()
+        self.word_embeddings = operations.Embedding(
+            vocab_size, embed_dim, padding_idx=pad_token_id, dtype=dtype, device=None
+        )
+        self.position_embeddings = operations.Embedding(max_position_embeddings, embed_dim, dtype=dtype, device=None)
+        self.token_type_embeddings = operations.Embedding(type_vocab_size, embed_dim, dtype=dtype, device=None)
+
+        self.LayerNorm = operations.LayerNorm(embed_dim, eps=layer_norm_eps, dtype=dtype, device=None)
+
+    def construct(self, input_tokens, embeds=None, token_type_ids=None, dtype=None):
+        if embeds is not None:
+            x = embeds
+        else:
+            x = self.word_embeddings(input_tokens, out_dtype=dtype)
+        x += comfy.ops.cast_to_input(self.position_embeddings.weight[: x.shape[1]], x)
+        if token_type_ids is not None:
+            x += self.token_type_embeddings(token_type_ids, out_dtype=x.dtype)
+        else:
+            x += comfy.ops.cast_to_input(self.token_type_embeddings.weight[0], x)
+        x = self.LayerNorm(x)
+        return x
+
+
+class BertModel_(mindspore.nn.Cell):
+    def __init__(self, config_dict, dtype, device, operations):
+        super().__init__()
+        embed_dim = config_dict["hidden_size"]
+        layer_norm_eps = config_dict["layer_norm_eps"]
+
+        self.embeddings = BertEmbeddings(
+            config_dict["vocab_size"],
+            config_dict["max_position_embeddings"],
+            config_dict["type_vocab_size"],
+            config_dict["pad_token_id"],
+            embed_dim,
+            layer_norm_eps,
+            dtype,
+            None,
+            operations,
+        )
+        self.encoder = BertEncoder(
+            config_dict["num_hidden_layers"],
+            embed_dim,
+            config_dict["intermediate_size"],
+            config_dict["num_attention_heads"],
+            layer_norm_eps,
+            dtype,
+            None,
+            operations,
+        )
+
+    def construct(
+        self,
+        input_tokens,
+        attention_mask=None,
+        embeds=None,
+        num_tokens=None,
+        intermediate_output=None,
+        final_layer_norm_intermediate=True,
+        dtype=None,
+        embeds_info=[],
+    ):
+        x = self.embeddings(input_tokens, embeds=embeds, dtype=dtype)
+        mask = None
+        if attention_mask is not None:
+            mask = 1.0 - attention_mask.to(x.dtype).reshape(
+                (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
+            ).expand((attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1]))
+            mask = mask.masked_fill(mask.to(mindspore.bool), -dtype_to_max(x.dtype))
+
+        x, i = self.encoder(x, mask, intermediate_output)
+        return x, i
+
+
+class BertModel(mindspore.nn.Cell):
+    def __init__(self, config_dict, dtype, device, operations):
+        super().__init__()
+        self.bert = BertModel_(config_dict, dtype, None, operations)
+        self.num_layers = config_dict["num_hidden_layers"]
+
+    def get_input_embeddings(self):
+        return self.bert.embeddings.word_embeddings
+
+    def set_input_embeddings(self, embeddings):
+        self.bert.embeddings.word_embeddings = embeddings
+
+    def construct(self, *args, **kwargs):
+        return self.bert(*args, **kwargs)

mindone/comfyui/comfy/text_encoders/flux.py

@@ -0,0 +1,91 @@
+import os
+
+import comfy.model_management
+import comfy.text_encoders.sd3_clip
+import comfy.text_encoders.t5
+from comfy import sd1_clip
+from transformers import T5TokenizerFast
+
+import mindspore
+from mindspore import mint
+
+
+class T5XXLTokenizer(sd1_clip.SDTokenizer):
+    def __init__(self, embedding_directory=None, tokenizer_data={}):
+        tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_tokenizer")
+        super().__init__(
+            tokenizer_path,
+            embedding_directory=embedding_directory,
+            pad_with_end=False,
+            embedding_size=4096,
+            embedding_key="t5xxl",
+            tokenizer_class=T5TokenizerFast,
+            has_start_token=False,
+            pad_to_max_length=False,
+            max_length=99999999,
+            min_length=256,
+            tokenizer_data=tokenizer_data,
+        )
+
+
+class FluxTokenizer:
+    def __init__(self, embedding_directory=None, tokenizer_data={}):
+        self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
+        self.t5xxl = T5XXLTokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
+
+    def tokenize_with_weights(self, text: str, return_word_ids=False, **kwargs):
+        out = {}
+        out["l"] = self.clip_l.tokenize_with_weights(text, return_word_ids, **kwargs)
+        out["t5xxl"] = self.t5xxl.tokenize_with_weights(text, return_word_ids, **kwargs)
+        return out
+
+    def untokenize(self, token_weight_pair):
+        return self.clip_l.untokenize(token_weight_pair)
+
+    def state_dict(self):
+        return {}
+
+
+class FluxClipModel(mindspore.nn.Cell):
+    def __init__(self, dtype_t5=None, device=None, dtype=None, model_options={}):
+        super().__init__()
+        dtype_t5 = comfy.model_management.pick_weight_dtype(dtype_t5, dtype)
+        self.clip_l = sd1_clip.SDClipModel(dtype=dtype, return_projected_pooled=False, model_options=model_options)
+        self.t5xxl = comfy.text_encoders.sd3_clip.T5XXLModel(dtype=dtype_t5, model_options=model_options)
+        self.dtypes = set([dtype, dtype_t5])
+
+    def set_clip_options(self, options):
+        self.clip_l.set_clip_options(options)
+        self.t5xxl.set_clip_options(options)
+
+    def reset_clip_options(self):
+        self.clip_l.reset_clip_options()
+        self.t5xxl.reset_clip_options()
+
+    def encode_token_weights(self, token_weight_pairs):
+        token_weight_pairs_l = token_weight_pairs["l"]
+        token_weight_pairs_t5 = token_weight_pairs["t5xxl"]
+
+        t5_out, t5_pooled = self.t5xxl.encode_token_weights(token_weight_pairs_t5)
+        l_out, l_pooled = self.clip_l.encode_token_weights(token_weight_pairs_l)
+        return t5_out, l_pooled
+
+    def load_sd(self, sd):
+        if "text_model.encoder.layers.1.mlp.fc1.weight" in sd:
+            sd = {f"clip_l.transformer.{k}": v for k, v in sd.items()}
+            return self.clip_l.load_sd(sd)
+        else:
+            sd = {f"t5xxl.transformer.{k}": v for k, v in sd.items()}
+            return self.t5xxl.load_sd(sd)
+
+
+def flux_clip(dtype_t5=None, t5xxl_scaled_fp8=None):
+    class FluxClipModel_(FluxClipModel):
+        def __init__(self, device=None, dtype=None, model_options={}):
+            if t5xxl_scaled_fp8 is not None and "t5xxl_scaled_fp8" not in model_options:
+                # model_options = model_options.copy()
+                # model_options["t5xxl_scaled_fp8"] = t5xxl_scaled_fp8
+                raise NotImplementedError
+            super().__init__(dtype_t5=dtype_t5, device=None, dtype=dtype, model_options=model_options)
+
+    return FluxClipModel_