diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index 8a16cfe55..63eefd420 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -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 diff --git a/comfy/nested_tensor.py b/comfy/nested_tensor.py index b700816fa..08c7133f8 100644 --- a/comfy/nested_tensor.py +++ b/comfy/nested_tensor.py @@ -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)) diff --git a/comfy/samplers.py b/comfy/samplers.py index 29e6bffb3..6fea0913e 100755 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -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() diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 56ef5f526..e81b6328b 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -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: diff --git a/latent_preview.py b/latent_preview.py index a9d777661..6bf2c1869 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -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