Upload folder using huggingface_hub

.gitattributes

@@ -35,3 +35,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 PixelGen-XL-16-256/demo.png filter=lfs diff=lfs merge=lfs -text
 PixelGen-XXL-16-512-t2i/tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+PixelGen-XXL-16-512-t2i/demo.png filter=lfs diff=lfs merge=lfs -text

PixelGen-XXL-16-512-t2i/transformer/transformer_jit_t2i.py

@@ -107,19 +107,19 @@ class JiTT2ITimestepEmbedder(nn.Module):
         self.frequency_embedding_size = frequency_embedding_size
 
     @staticmethod
-    def timestep_embedding(t, dim, max_period=10000):
+    def timestep_embedding(t, dim, max_period=10):
         half = dim // 2
-        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
-            device=t.device
+        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]
+        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)
-        return embedding
+        return embedding.to(dtype=t.dtype)
 
     def forward(self, t, dtype=None):
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_freq = self.timestep_embedding(t.reshape(-1), self.frequency_embedding_size)
         if dtype is not None:
             t_freq = t_freq.to(dtype=dtype)
         return self.mlp(t_freq)