Merge branch 'comfyanonymous:master' into master

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
 

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()

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}

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)