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

Cube3D: fix graph integration (3D latent, VAE device, fp32 cond, scikit-image)

Browse Source 
Amp-Thread-ID: https://ampcode.com/threads/T-019ec361-addb-70d8-a74b-438ce8a1e096
Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
Jedrzej Kosinski
2026-06-14 22:59:11 -07:00
parent 01a8783bee
commit 871f7bc390
4 changed files with 20 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
comfy/k_diffusion/sampling.py
View File

@@ -2016,12 +2016,16 @@ def sample_cube(model, x, sigmas, extra_args=None, callback=None, disable=None,
bbox = torch.zeros((c.shape[0], 3), device=device, dtype=c.dtype)
return torch.cat([c, cube.bbox_proj(bbox).unsqueeze(1)], dim=1)
with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=autocast_enabled):
cond = add_bbox(cube.encode_text(pos.to(device=device, dtype=weight_dtype)))
if use_cfg:
ucond = add_bbox(cube.encode_text(neg.to(device=device, dtype=weight_dtype)))
cond = torch.cat([cond, ucond], dim=0)
# Conditioning (text_proj + bbox_proj) is computed in the model's weight dtype
# OUTSIDE the bf16 autocast block, matching upstream cube's Engine.prepare_inputs
# (run_clip/encode_text run in full precision). The autocast only covers the
# autoregressive transformer forward, exactly like Engine.run_gpt.
cond = add_bbox(cube.encode_text(pos.to(device=device, dtype=weight_dtype)))
if use_cfg:
ucond = add_bbox(cube.encode_text(neg.to(device=device, dtype=weight_dtype)))
cond = torch.cat([cond, ucond], dim=0)
with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=autocast_enabled):
bos = torch.full((cond.shape[0], 1), cube.shape_bos_id, dtype=torch.long, device=device)
embed = cube.encode_token(bos)
Bp, input_seq_len, dim = embed.shape