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

Fix MageFlow on cards that don't support bf16 (#15081)

Browse Source
This commit is contained in:
Jukka Seppänen
2026-07-26 04:01:51 +03:00
committed by GitHub
parent f966a2b38c
commit 806e092ed4
2 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
comfy/ldm/mage_flow/model.py
View File

@@ -22,7 +22,6 @@ class MageTimestepProjEmbeddings(nn.Module):
) )
def forward(self, timestep, hidden_states): def forward(self, timestep, hidden_states):
timestep = timestep.to(hidden_states.dtype)
half_dim = 128 half_dim = 128
exponent = -math.log(10000) * torch.arange(half_dim, dtype=torch.float32, device=timestep.device) / half_dim exponent = -math.log(10000) * torch.arange(half_dim, dtype=torch.float32, device=timestep.device) / half_dim
emb = torch.exp(exponent).to(timestep.dtype) emb = torch.exp(exponent).to(timestep.dtype)

4
comfy/model_base.py
View File

@@ -2279,6 +2279,10 @@ class MageFlow(QwenImage):
def __init__(self, model_config, model_type=ModelType.FLOW, device=None): def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.mage_flow.model.MageFlowTransformer2DModel) super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.mage_flow.model.MageFlowTransformer2DModel)
def process_timestep(self, timestep, **kwargs):
# Mage runs in bf16 and rounds its timestep frequency table to the timestep dtype, keep that on fp32 devices.
return timestep.to(torch.bfloat16)
def extra_conds_shapes(self, **kwargs): def extra_conds_shapes(self, **kwargs):
out = {} out = {}
ref_latents = kwargs.get("reference_latents", None) ref_latents = kwargs.get("reference_latents", None)