mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-04-23 17:02:38 +08:00
Fix Hunyuan 3D 2.1 multi-GPU worksplit: use cond_or_uncond instead of hardcoded chunk(2)
Amp-Thread-ID: https://ampcode.com/threads/T-019da964-2cc8-77f9-9aae-23f65da233db Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
parent
f0d550bd02
commit
4b8fa98c3c
@ -607,9 +607,14 @@ class HunYuanDiTPlain(nn.Module):
|
|||||||
def forward(self, x, t, context, transformer_options = {}, **kwargs):
|
def forward(self, x, t, context, transformer_options = {}, **kwargs):
|
||||||
|
|
||||||
x = x.movedim(-1, -2)
|
x = x.movedim(-1, -2)
|
||||||
uncond_emb, cond_emb = context.chunk(2, dim = 0)
|
|
||||||
|
|
||||||
context = torch.cat([cond_emb, uncond_emb], dim = 0)
|
cond_or_uncond = transformer_options.get("cond_or_uncond", [])
|
||||||
|
swap_cfg_halves = len(cond_or_uncond) == 2 and set(cond_or_uncond) == {0, 1}
|
||||||
|
|
||||||
|
if swap_cfg_halves:
|
||||||
|
first_half, second_half = context.chunk(2, dim = 0)
|
||||||
|
context = torch.cat([second_half, first_half], dim = 0)
|
||||||
|
|
||||||
main_condition = context
|
main_condition = context
|
||||||
|
|
||||||
t = 1.0 - t
|
t = 1.0 - t
|
||||||
@ -657,5 +662,8 @@ class HunYuanDiTPlain(nn.Module):
|
|||||||
output = self.final_layer(combined)
|
output = self.final_layer(combined)
|
||||||
output = output.movedim(-2, -1) * (-1.0)
|
output = output.movedim(-2, -1) * (-1.0)
|
||||||
|
|
||||||
cond_emb, uncond_emb = output.chunk(2, dim = 0)
|
if swap_cfg_halves:
|
||||||
return torch.cat([uncond_emb, cond_emb])
|
first_half, second_half = output.chunk(2, dim = 0)
|
||||||
|
output = torch.cat([second_half, first_half], dim = 0)
|
||||||
|
|
||||||
|
return output
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user