wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-18 20:38:15 +08:00
Code Issues Actions 27 Packages Projects Releases Wiki Activity

Create PixArt dropout masks on the input device

Browse Source
This commit is contained in:
野生の男 2026-07-13 14:39:17 +09:00
parent e6f26fa2cc
commit e451fbbac4
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
comfy/ldm/pixart/blocks.py
View File

@ -296,7 +296,7 @@ class LabelEmbedder(nn.Module):
Drops labels to enable classifier-free guidance.
"""
if force_drop_ids is None:
drop_ids = torch.rand(labels.shape[0]).cuda() < self.dropout_prob
drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
else:
drop_ids = force_drop_ids == 1
labels = torch.where(drop_ids, self.num_classes, labels)
@ -328,7 +328,7 @@ class CaptionEmbedder(nn.Module):
Drops labels to enable classifier-free guidance.
"""
if force_drop_ids is None:
drop_ids = torch.rand(caption.shape[0]).cuda() < self.uncond_prob
drop_ids = torch.rand(caption.shape[0], device=caption.device) < self.uncond_prob
else:
drop_ids = force_drop_ids == 1
caption = torch.where(drop_ids[:, None, None, None], self.y_embedding, caption)
@ -363,7 +363,7 @@ class CaptionEmbedderDoubleBr(nn.Module):
Drops labels to enable classifier-free guidance.
"""
if force_drop_ids is None:
drop_ids = torch.rand(global_caption.shape[0]).cuda() < self.uncond_prob
drop_ids = torch.rand(global_caption.shape[0], device=global_caption.device) < self.uncond_prob
else:
drop_ids = force_drop_ids == 1
global_caption = torch.where(drop_ids[:, None], self.embedding, global_caption)