From 0a4c49c57ca792535174d27aea87449257527b2f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 23 Jul 2024 15:35:28 -0400 Subject: [PATCH] Support MT5. --- comfy/text_encoders/t5.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/text_encoders/t5.py b/comfy/text_encoders/t5.py index 448c5aa..ba62847 100644 --- a/comfy/text_encoders/t5.py +++ b/comfy/text_encoders/t5.py @@ -223,7 +223,7 @@ class T5(torch.nn.Module): self.num_layers = config_dict["num_layers"] model_dim = config_dict["d_model"] - self.encoder = T5Stack(self.num_layers, model_dim, model_dim, config_dict["d_ff"], config_dict["dense_act_fn"], config_dict["is_gated_act"], config_dict["num_heads"], config_dict["model_type"] == "t5", dtype, device, operations) + self.encoder = T5Stack(self.num_layers, model_dim, model_dim, config_dict["d_ff"], config_dict["dense_act_fn"], config_dict["is_gated_act"], config_dict["num_heads"], config_dict["model_type"] != "umt5", dtype, device, operations) self.dtype = dtype self.shared = torch.nn.Embedding(config_dict["vocab_size"], model_dim, device=device)