Support latent previews for nested latents (#15196)

This commit is contained in:
Jukka Seppänen 2026-08-02 02:57:15 +03:00 committed by GitHub
parent e803f24ea0
commit 49a7422892
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 153 additions and 8 deletions

View File

@ -434,8 +434,138 @@ class LTXV(LatentFormat):
class LTXAV(LTXV):
def __init__(self):
self.latent_rgb_factors = None
self.latent_rgb_factors_bias = None
# video-stream preview factors for the packed AV latent (audio stream is not previewed)
self.latent_rgb_factors = [
[ 0.001135, -0.010555, -0.004925],
[-0.008019, -0.006231, -0.005564],
[ 0.012637, 0.005605, 0.012713],
[ 0.023454, 0.020771, 0.017844],
[-0.011940, -0.000932, 0.009292],
[ 0.018602, 0.011018, 0.013969],
[-0.036369, -0.046631, -0.057898],
[-0.031919, 0.000131, 0.015214],
[ 0.014519, 0.021041, 0.015325],
[ 0.018889, 0.016149, -0.002836],
[-0.003784, -0.006057, -0.008195],
[ 0.013262, 0.030259, 0.029775],
[ 0.050465, 0.050366, 0.025255],
[ 0.018628, 0.007691, 0.002893],
[-0.015698, -0.008451, -0.000676],
[-0.013600, -0.012587, -0.004437],
[ 0.012482, 0.021469, 0.027913],
[-0.018241, -0.013488, -0.010975],
[ 0.013828, 0.012568, 0.021984],
[ 0.017911, 0.006552, 0.005567],
[ 0.026769, 0.006803, -0.009360],
[-0.006794, -0.008447, -0.013921],
[ 0.029708, 0.018671, 0.022811],
[-0.014732, -0.019169, 0.000903],
[ 0.019607, 0.032595, 0.053409],
[-0.003721, 0.003976, 0.010364],
[-0.020193, -0.026076, -0.036068],
[-0.002328, 0.006527, 0.013052],
[ 0.017171, 0.009224, 0.006548],
[ 0.001104, -0.000591, 0.000147],
[-0.000217, 0.011834, 0.017945],
[-0.015329, -0.012463, -0.006178],
[-0.009478, -0.008680, -0.004107],
[-0.005565, -0.006006, -0.001493],
[ 0.009451, 0.008794, 0.013207],
[-0.009989, -0.008027, -0.009568],
[-0.001505, -0.008805, -0.006828],
[ 0.001105, 0.008999, 0.009079],
[ 0.025935, 0.016426, 0.008036],
[ 0.006313, 0.000694, -0.006039],
[-0.001893, -0.006951, -0.009560],
[-0.007082, -0.002566, -0.007152],
[-0.005231, 0.004829, 0.008220],
[-0.004333, 0.001251, -0.004852],
[-0.017024, -0.012730, -0.007457],
[ 0.024988, 0.032963, 0.036556],
[ 0.013697, 0.012278, 0.009979],
[-0.013751, -0.008369, -0.015446],
[-0.009348, -0.001047, 0.007622],
[-0.003135, -0.003350, -0.003766],
[ 0.007436, 0.004957, 0.010480],
[ 0.018315, 0.022066, 0.021104],
[-0.005621, -0.006770, -0.008219],
[-0.007427, 0.001911, -0.001231],
[-0.007413, 0.000486, -0.006039],
[-0.014698, -0.007160, 0.006509],
[ 0.013775, 0.014185, 0.008203],
[ 0.060246, 0.069787, 0.072833],
[ 0.009861, 0.004870, 0.001194],
[-0.003660, 0.003251, 0.008015],
[ 0.003696, -0.003680, -0.008851],
[ 0.014924, 0.006196, 0.005282],
[-0.006740, -0.004319, -0.006729],
[ 0.020635, 0.015163, 0.012385],
[-0.032623, -0.006105, 0.010436],
[-0.058988, -0.030162, -0.037961],
[-0.035614, -0.021929, -0.011062],
[-0.023412, -0.011305, -0.005054],
[-0.002716, -0.005184, -0.004084],
[ 0.014591, 0.015294, 0.014045],
[ 0.008310, 0.002466, -0.003225],
[ 0.005176, 0.001119, 0.000695],
[-0.021569, -0.030886, -0.044732],
[ 0.007517, 0.003891, 0.000551],
[-0.006793, 0.004059, 0.010184],
[-0.086481, -0.082033, -0.083414],
[ 0.004192, 0.000762, -0.008658],
[ 0.010970, 0.009002, 0.007384],
[ 0.004042, -0.006732, -0.011031],
[ 0.012164, 0.006401, 0.007483],
[ 0.029252, 0.013990, 0.011128],
[ 0.048452, 0.034648, 0.016269],
[ 0.024104, 0.012647, 0.011754],
[-0.013216, -0.020192, -0.019752],
[-0.010799, -0.008535, -0.005467],
[ 0.005823, 0.001403, 0.001890],
[ 0.052393, 0.044771, 0.032777],
[ 0.007576, -0.008080, -0.012453],
[ 0.009830, 0.004244, 0.001213],
[-0.025867, -0.013169, -0.010636],
[ 0.008494, 0.003135, 0.000790],
[ 0.003969, -0.002625, -0.010204],
[ 0.006509, 0.008272, 0.020819],
[-0.004943, -0.013424, -0.015351],
[ 0.005541, 0.009136, -0.003666],
[-0.014300, -0.015864, -0.016853],
[ 0.002650, 0.028393, 0.014125],
[-0.027661, -0.045422, -0.064995],
[ 0.009220, 0.015522, 0.010574],
[-0.002236, 0.002915, 0.004557],
[-0.020269, -0.008212, -0.000532],
[ 0.019294, 0.003655, -0.002809],
[ 0.007116, -0.002784, 0.000017],
[ 0.057277, 0.073270, 0.074401],
[-0.002616, -0.001696, -0.000498],
[ 0.007248, 0.009793, 0.022829],
[-0.002590, -0.005601, -0.000436],
[-0.007681, 0.003893, -0.004119],
[-0.057392, -0.045545, -0.025290],
[ 0.045188, 0.047985, 0.054059],
[ 0.000937, -0.008861, -0.038406],
[-0.010192, -0.008036, -0.005385],
[-0.030222, -0.027498, -0.030765],
[-0.008359, 0.013247, 0.010918],
[ 0.004102, 0.002093, 0.006934],
[ 0.039461, 0.027339, 0.008284],
[-0.075747, -0.076340, -0.071625],
[ 0.002692, 0.005096, -0.002247],
[-0.002453, -0.002785, -0.010483],
[ 0.012265, 0.005481, 0.001729],
[ 0.017755, 0.008655, 0.003532],
[ 0.055560, 0.049128, 0.044137],
[-0.025861, -0.023798, -0.018815],
[-0.014876, -0.010770, -0.010713],
[-0.017315, -0.012599, -0.008661],
[-0.008461, -0.006210, -0.007744],
[-0.040175, -0.042255, -0.048119],
[-0.019355, -0.021055, -0.021919],
]
self.latent_rgb_factors_bias = [-0.347892, -0.363814, -0.370287]
class HunyuanVideo(LatentFormat):
latent_channels = 16

View File

@ -51,6 +51,9 @@ class NestedTensor:
def float(self):
return self.to(dtype=torch.float)
def cpu(self):
return self.to(device="cpu")
def chunk(self, *args, **kwargs):
return self.apply_operation(None, lambda x, y: x.chunk(*args, **kwargs))

View File

@ -1284,6 +1284,14 @@ class CFGGuider:
sampler_shapes = [tuple(latent_image.shape)]
detail("Sampler: model=%s latent_shapes=%s", self.model_patcher.model.__class__.__name__, sampler_shapes)
if len(latent_shapes) > 1 and callback is not None:
# samplers run on the flat pack, hand callbacks (previews, x0 output) the nested view
packed_callback = callback
def callback(step, x0, x, total_steps):
x0 = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0, latent_shapes))
x = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x, latent_shapes))
return packed_callback(step, x0, x, total_steps)
if denoise_mask is not None:
if denoise_mask.is_nested:
denoise_masks = denoise_mask.unbind()

View File

@ -775,10 +775,11 @@ class SamplerCustom(io.ComfyNode):
out.pop("downscale_ratio_temporal", None)
out["samples"] = samples
if "x0" in x0_output:
x0_out = model.model.process_latent_out(x0_output["x0"].cpu())
if samples.is_nested:
x0 = x0_output["x0"]
if samples.is_nested and not x0.is_nested:
latent_shapes = [x.shape for x in samples.unbind()]
x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes))
x0 = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0, latent_shapes))
x0_out = model.model.process_latent_out(x0.cpu())
out_denoised = latent.copy()
out_denoised["samples"] = x0_out
else:
@ -1053,10 +1054,11 @@ class SamplerCustomAdvanced(io.ComfyNode):
out.pop("downscale_ratio_temporal", None)
out["samples"] = samples
if "x0" in x0_output:
x0_out = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu())
if samples.is_nested:
x0 = x0_output["x0"]
if samples.is_nested and not x0.is_nested:
latent_shapes = [x.shape for x in samples.unbind()]
x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes))
x0 = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0, latent_shapes))
x0_out = guider.model_patcher.model.process_latent_out(x0.cpu())
out_denoised = latent.copy()
out_denoised["samples"] = x0_out
else:

View File

@ -123,6 +123,8 @@ def prepare_callback(model, steps, x0_output_dict=None):
preview_bytes = None
if previewer:
if x0.is_nested:
x0 = x0.tensors[0]
preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0)
pbar.update_absolute(step + 1, total_steps, preview_bytes)
return callback