From 334ba48cea2961994e92c2fb25de9417b19897ed Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 23 Jul 2024 14:13:32 -0400 Subject: [PATCH 1/4] More generic unet prefix detection code. --- comfy/model_detection.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index c62e2b822..ae88eeb97 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -261,13 +261,22 @@ def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=Fal return model_config def unet_prefix_from_state_dict(state_dict): - if "model.model.postprocess_conv.weight" in state_dict: #audio models - unet_key_prefix = "model.model." - elif "model.double_layers.0.attn.w1q.weight" in state_dict: #aura flow - unet_key_prefix = "model." + candidates = ["model.diffusion_model.", #ldm/sgm models + "model.model.", #audio models + ] + counts = {k: 0 for k in candidates} + for k in state_dict: + for c in candidates: + if k.startswith(c): + counts[c] += 1 + break + + top = max(counts, key=counts.get) + if counts[top] > 5: + return top else: - unet_key_prefix = "model.diffusion_model." - return unet_key_prefix + return "model." #aura flow and others + def convert_config(unet_config): new_config = unet_config.copy() From 88ed89303463ab030654ef87f7f54d72ee3979bb Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 23 Jul 2024 14:17:42 -0400 Subject: [PATCH 2/4] Allow SPieceTokenizer to load model from a byte string. --- comfy/text_encoders/spiece_tokenizer.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/comfy/text_encoders/spiece_tokenizer.py b/comfy/text_encoders/spiece_tokenizer.py index d611d5bb7..fa303da0b 100644 --- a/comfy/text_encoders/spiece_tokenizer.py +++ b/comfy/text_encoders/spiece_tokenizer.py @@ -1,14 +1,18 @@ import os class SPieceTokenizer: + add_eos = True + @staticmethod def from_pretrained(path): return SPieceTokenizer(path) def __init__(self, tokenizer_path): import sentencepiece - self.tokenizer = sentencepiece.SentencePieceProcessor(model_file=tokenizer_path) - self.end = self.tokenizer.eos_id() + if isinstance(tokenizer_path, bytes): + self.tokenizer = sentencepiece.SentencePieceProcessor(model_proto=tokenizer_path, add_eos=self.add_eos) + else: + self.tokenizer = sentencepiece.SentencePieceProcessor(model_file=tokenizer_path, add_eos=self.add_eos) def get_vocab(self): out = {} @@ -18,5 +22,4 @@ class SPieceTokenizer: def __call__(self, string): out = self.tokenizer.encode(string) - out += [self.end] return {"input_ids": out} From 0a4c49c57ca792535174d27aea87449257527b2f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 23 Jul 2024 15:35:28 -0400 Subject: [PATCH 3/4] 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 448c5aad3..ba628478d 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) From 10b43ceea52d86395d10a7029eda2822fd65dfd1 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 24 Jul 2024 01:12:59 -0400 Subject: [PATCH 4/4] Remove duplicate code. --- comfy/ldm/modules/diffusionmodules/mmdit.py | 29 ++------------------- 1 file changed, 2 insertions(+), 27 deletions(-) diff --git a/comfy/ldm/modules/diffusionmodules/mmdit.py b/comfy/ldm/modules/diffusionmodules/mmdit.py index 927451534..f37f7ff77 100644 --- a/comfy/ldm/modules/diffusionmodules/mmdit.py +++ b/comfy/ldm/modules/diffusionmodules/mmdit.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn from .. import attention from einops import rearrange, repeat +from .util import timestep_embedding def default(x, y): if x is not None: @@ -230,34 +231,8 @@ class TimestepEmbedder(nn.Module): ) self.frequency_embedding_size = frequency_embedding_size - @staticmethod - def timestep_embedding(t, dim, max_period=10000): - """ - Create sinusoidal timestep embeddings. - :param t: a 1-D Tensor of N indices, one per batch element. - These may be fractional. - :param dim: the dimension of the output. - :param max_period: controls the minimum frequency of the embeddings. - :return: an (N, D) Tensor of positional embeddings. - """ - half = dim // 2 - freqs = torch.exp( - -math.log(max_period) - * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) - / half - ) - args = t[:, None].float() * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat( - [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 - ) - if torch.is_floating_point(t): - embedding = embedding.to(dtype=t.dtype) - return embedding - def forward(self, t, dtype, **kwargs): - t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype) + t_freq = timestep_embedding(t, self.frequency_embedding_size).to(dtype) t_emb = self.mlp(t_freq) return t_emb