zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-24 18:10:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: Support NVIDIA PixelDiT and PiD (CORE-201) (#14103)

Browse Source
This commit is contained in:
Jukka Seppänen
2026-05-27 03:50:14 +03:00
committed by GitHub
parent 921775704c
commit 28f4ef277c
12 changed files with 946 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
comfy/ldm/modules/diffusionmodules/mmdit.py
View File

@@ -211,7 +211,7 @@ class TimestepEmbedder(nn.Module):
Embeds scalar timesteps into vector representations.
"""
def __init__(self, hidden_size, frequency_embedding_size=256, output_size=None, dtype=None, device=None, operations=None):
def __init__(self, hidden_size, frequency_embedding_size=256, output_size=None, dtype=None, device=None, operations=None, max_period=10000):
super().__init__()
if output_size is None:
output_size = hidden_size
@@ -221,9 +221,10 @@ class TimestepEmbedder(nn.Module):
operations.Linear(hidden_size, output_size, bias=True, dtype=dtype, device=device),
)
self.frequency_embedding_size = frequency_embedding_size
self.max_period = max_period
def forward(self, t, dtype, **kwargs):
t_freq = timestep_embedding(t, self.frequency_embedding_size).to(dtype)
t_freq = timestep_embedding(t, self.frequency_embedding_size, max_period=self.max_period).to(dtype)
t_emb = self.mlp(t_freq)
return t_emb