Automated PR - 2026-01-06
This commit is contained in:
sync-bot
+305
-25
@@ -1,3 +1,5 @@
|
||||
import logging
|
||||
|
||||
import torch
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
@@ -19,6 +21,10 @@ class LTXVTiledVAEDecode:
|
||||
"overlap": ("INT", {"default": 1, "min": 1, "max": 8}),
|
||||
"last_frame_fix": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"optional": {
|
||||
"working_device": (["cpu", "auto"], {"default": "auto"}),
|
||||
"working_dtype": (["float16", "float32", "auto"], {"default": "auto"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
@@ -36,6 +42,8 @@ class LTXVTiledVAEDecode:
|
||||
vertical_tiles,
|
||||
overlap,
|
||||
last_frame_fix,
|
||||
working_device="auto",
|
||||
working_dtype="auto",
|
||||
):
|
||||
# Get the latent samples
|
||||
samples = latents["samples"]
|
||||
@@ -67,6 +75,31 @@ class LTXVTiledVAEDecode:
|
||||
output = None
|
||||
weights = None
|
||||
|
||||
target_device = samples.device if working_device == "auto" else working_device
|
||||
if working_dtype == "auto":
|
||||
target_dtype = samples.dtype
|
||||
elif working_dtype == "float16":
|
||||
target_dtype = torch.float16
|
||||
elif working_dtype == "float32":
|
||||
target_dtype = torch.float32
|
||||
|
||||
output = torch.zeros(
|
||||
(
|
||||
batch,
|
||||
image_frames,
|
||||
output_height,
|
||||
output_width,
|
||||
3,
|
||||
),
|
||||
device=target_device,
|
||||
dtype=target_dtype,
|
||||
)
|
||||
weights = torch.zeros(
|
||||
(batch, image_frames, output_height, output_width, 1),
|
||||
device=target_device,
|
||||
dtype=target_dtype,
|
||||
)
|
||||
|
||||
# Process each tile
|
||||
for v in range(vertical_tiles):
|
||||
for h in range(horizontal_tiles):
|
||||
@@ -90,9 +123,9 @@ class LTXVTiledVAEDecode:
|
||||
tile_height = v_end - v_start
|
||||
tile_width = h_end - h_start
|
||||
|
||||
print(f"Processing VAE decode tile at row {v}, col {h}:")
|
||||
print(f" Position: ({v_start}:{v_end}, {h_start}:{h_end})")
|
||||
print(f" Size: {tile_height}x{tile_width}")
|
||||
logging.info(f"Processing VAE decode tile at row {v}, col {h}:")
|
||||
logging.info(f" Position: ({v_start}:{v_end}, {h_start}:{h_end})")
|
||||
logging.info(f" Size: {tile_height}x{tile_width}")
|
||||
|
||||
# Extract tile
|
||||
tile = samples[:, :, :, v_start:v_end, h_start:h_end]
|
||||
@@ -103,25 +136,6 @@ class LTXVTiledVAEDecode:
|
||||
# Decode the tile
|
||||
decoded_tile = vae.decode(tile_latents["samples"])
|
||||
|
||||
# Initialize output tensors on first tile
|
||||
if output is None:
|
||||
output = torch.zeros(
|
||||
(
|
||||
batch,
|
||||
image_frames,
|
||||
output_height,
|
||||
output_width,
|
||||
decoded_tile.shape[-1],
|
||||
),
|
||||
device=decoded_tile.device,
|
||||
dtype=decoded_tile.dtype,
|
||||
)
|
||||
weights = torch.zeros(
|
||||
(batch, image_frames, output_height, output_width, 1),
|
||||
device=decoded_tile.device,
|
||||
dtype=decoded_tile.dtype,
|
||||
)
|
||||
|
||||
# Calculate output tile boundaries
|
||||
out_h_start = v_start * height_scale_factor
|
||||
out_h_end = v_end * height_scale_factor
|
||||
@@ -176,15 +190,15 @@ class LTXVTiledVAEDecode:
|
||||
# Add weighted tile to output
|
||||
output[:, :, out_h_start:out_h_end, out_w_start:out_w_end, :] += (
|
||||
decoded_tile * tile_weights
|
||||
)
|
||||
).to(target_device, target_dtype)
|
||||
|
||||
# Add weights to weight tensor
|
||||
weights[
|
||||
:, :, out_h_start:out_h_end, out_w_start:out_w_end, :
|
||||
] += tile_weights
|
||||
] += tile_weights.to(target_device, target_dtype)
|
||||
|
||||
# Normalize by weights
|
||||
output = output / (weights + 1e-8)
|
||||
output /= weights + 1e-8
|
||||
|
||||
# Reshape output to match expected format [batch * frames, height, width, channels]
|
||||
output = output.view(
|
||||
@@ -195,3 +209,269 @@ class LTXVTiledVAEDecode:
|
||||
output = output[:-time_scale_factor, :, :]
|
||||
|
||||
return (output,)
|
||||
|
||||
|
||||
def compute_chunk_boundaries(
|
||||
chunk_start: int,
|
||||
temporal_tile_length: int,
|
||||
temporal_overlap: int,
|
||||
total_latent_frames: int,
|
||||
):
|
||||
"""Compute chunk boundaries for temporal tiling.
|
||||
|
||||
Args:
|
||||
chunk_start: Starting frame index for the current chunk
|
||||
temporal_tile_length: Length of each temporal tile
|
||||
temporal_overlap: Number of frames to overlap between chunks
|
||||
total_latent_frames: Total number of latent frames
|
||||
|
||||
Returns:
|
||||
Tuple of (overlap_start, chunk_end)
|
||||
"""
|
||||
if chunk_start == 0:
|
||||
# First chunk: no overlap needed
|
||||
chunk_end = min(chunk_start + temporal_tile_length, total_latent_frames)
|
||||
overlap_start = chunk_start
|
||||
else:
|
||||
# Subsequent chunks: include overlap from previous chunk
|
||||
# -1 because we need one extra frame to overlap, which is decoded to a single frame
|
||||
# never overlap with the first latent frame
|
||||
overlap_start = max(1, chunk_start - temporal_overlap - 1)
|
||||
extra_frames = chunk_start - overlap_start
|
||||
chunk_end = min(
|
||||
chunk_start + temporal_tile_length - extra_frames,
|
||||
total_latent_frames,
|
||||
)
|
||||
|
||||
return overlap_start, chunk_end
|
||||
|
||||
|
||||
def calculate_temporal_output_boundaries(
|
||||
overlap_start: int, time_scale_factor: int, tile_out_frames: int
|
||||
):
|
||||
"""Calculate temporal output boundaries for the decoded tile.
|
||||
|
||||
Args:
|
||||
overlap_start: Starting frame index including overlap
|
||||
time_scale_factor: Time scaling factor from VAE
|
||||
tile_out_frames: Number of frames in the decoded tile
|
||||
|
||||
Returns:
|
||||
Tuple of (out_t_start, out_t_end)
|
||||
"""
|
||||
# +1 for the first frame
|
||||
out_t_start = 1 + overlap_start * time_scale_factor
|
||||
|
||||
# Calculate actual output temporal dimensions
|
||||
out_t_end = out_t_start + tile_out_frames
|
||||
|
||||
return out_t_start, out_t_end
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVSpatioTemporalTiledVAEDecode",
|
||||
)
|
||||
class LTXVSpatioTemporalTiledVAEDecode(LTXVTiledVAEDecode):
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE", {"tooltip": "The VAE to use."}),
|
||||
"latents": ("LATENT", {"tooltip": "The latent samples to decode."}),
|
||||
"spatial_tiles": (
|
||||
"INT",
|
||||
{
|
||||
"default": 4,
|
||||
"min": 1,
|
||||
"max": 8,
|
||||
"tooltip": "The number of spatial tiles to use, horizontal and vertical.",
|
||||
},
|
||||
),
|
||||
"spatial_overlap": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 0,
|
||||
"max": 8,
|
||||
"tooltip": "The overlap between the spatial tiles. (in latent frames)",
|
||||
},
|
||||
),
|
||||
"temporal_tile_length": (
|
||||
"INT",
|
||||
{
|
||||
"default": 16,
|
||||
"min": 2,
|
||||
"max": 1000,
|
||||
"tooltip": "The length of the temporal tile to use for the sampling, in latent frames, including the overlapping region.",
|
||||
},
|
||||
),
|
||||
"temporal_overlap": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 0,
|
||||
"max": 8,
|
||||
"tooltip": "The overlap between the temporal tiles, in latent frames.",
|
||||
},
|
||||
),
|
||||
"last_frame_fix": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "If true, the last frame will be repeated and discarded after the decoding.",
|
||||
},
|
||||
),
|
||||
"working_device": (
|
||||
["cpu", "auto"],
|
||||
{
|
||||
"default": "auto",
|
||||
"tooltip": "The device to use for the decoding. auto->same as the latents.",
|
||||
},
|
||||
),
|
||||
"working_dtype": (
|
||||
["float16", "float32", "auto"],
|
||||
{
|
||||
"default": "auto",
|
||||
"tooltip": "The data type to use for the decoding. auto->same as the latents.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
|
||||
FUNCTION = "decode_spatial_temporal"
|
||||
|
||||
CATEGORY = "latent"
|
||||
|
||||
def decode_spatial_temporal(
|
||||
self,
|
||||
vae,
|
||||
latents,
|
||||
spatial_tiles=4,
|
||||
spatial_overlap=1,
|
||||
temporal_tile_length=16,
|
||||
temporal_overlap=1,
|
||||
last_frame_fix=False,
|
||||
working_device="auto",
|
||||
working_dtype="auto",
|
||||
):
|
||||
if temporal_tile_length < temporal_overlap + 1:
|
||||
raise ValueError(
|
||||
"Temporal tile length must be greater than temporal overlap + 1"
|
||||
)
|
||||
|
||||
# Get the latent samples
|
||||
samples = latents["samples"]
|
||||
|
||||
batch, channels, frames, height, width = samples.shape
|
||||
time_scale_factor, width_scale_factor, height_scale_factor = (
|
||||
vae.downscale_index_formula
|
||||
)
|
||||
image_frames = 1 + (frames - 1) * time_scale_factor
|
||||
|
||||
# Calculate output image dimensions
|
||||
output_height = height * height_scale_factor
|
||||
output_width = width * width_scale_factor
|
||||
|
||||
target_device = samples.device if working_device == "auto" else working_device
|
||||
if working_dtype == "auto":
|
||||
target_dtype = samples.dtype
|
||||
elif working_dtype == "float16":
|
||||
target_dtype = torch.float16
|
||||
elif working_dtype == "float32":
|
||||
target_dtype = torch.float32
|
||||
|
||||
# Initialize output tensor and weight tensor
|
||||
output = torch.empty(
|
||||
(
|
||||
batch,
|
||||
image_frames,
|
||||
output_height,
|
||||
output_width,
|
||||
3,
|
||||
),
|
||||
device=target_device,
|
||||
dtype=target_dtype,
|
||||
)
|
||||
|
||||
# Process temporal chunks similar to reference function
|
||||
total_latent_frames = frames
|
||||
chunk_start = 0
|
||||
|
||||
while chunk_start < total_latent_frames:
|
||||
# Calculate chunk boundaries
|
||||
overlap_start, chunk_end = compute_chunk_boundaries(
|
||||
chunk_start, temporal_tile_length, temporal_overlap, total_latent_frames
|
||||
)
|
||||
|
||||
# units are latent frames
|
||||
chunk_frames = chunk_end - overlap_start
|
||||
logging.info(
|
||||
f"Processing temporal chunk: {overlap_start}:{chunk_end} ({chunk_frames} latent frames)"
|
||||
)
|
||||
|
||||
# Extract tile
|
||||
tile = samples[:, :, overlap_start:chunk_end]
|
||||
|
||||
# Create tile latents dict
|
||||
tile_latents = {"samples": tile}
|
||||
|
||||
# Decode the tile
|
||||
decoded_tile = self.decode(
|
||||
vae=vae,
|
||||
latents=tile_latents,
|
||||
vertical_tiles=spatial_tiles,
|
||||
horizontal_tiles=spatial_tiles,
|
||||
overlap=spatial_overlap,
|
||||
last_frame_fix=last_frame_fix,
|
||||
working_device=working_device,
|
||||
working_dtype=working_dtype,
|
||||
)[0][None]
|
||||
|
||||
if chunk_start == 0:
|
||||
output[:, : decoded_tile.shape[1]] = decoded_tile
|
||||
|
||||
# Drop first frame if needed (overlap)
|
||||
else:
|
||||
if decoded_tile.shape[1] == 1:
|
||||
raise ValueError("Dropping first frame but tile has only 1 frame")
|
||||
decoded_tile = decoded_tile[:, 1:] # Drop first frame
|
||||
|
||||
# Calculate temporal output boundaries
|
||||
out_t_start, out_t_end = calculate_temporal_output_boundaries(
|
||||
overlap_start, time_scale_factor, decoded_tile.shape[1]
|
||||
)
|
||||
|
||||
# Create weight mask for this tile
|
||||
overlap_frames = temporal_overlap * time_scale_factor
|
||||
frame_weights = torch.linspace(
|
||||
0,
|
||||
1,
|
||||
overlap_frames + 2,
|
||||
device=decoded_tile.device,
|
||||
dtype=decoded_tile.dtype,
|
||||
)[1:-1]
|
||||
tile_weights = frame_weights.view(1, -1, 1, 1, 1)
|
||||
after_overlap_frames_start = out_t_start + overlap_frames
|
||||
# Add weighted tile to output
|
||||
overlap_output = decoded_tile[:, :overlap_frames]
|
||||
output[:, out_t_start:after_overlap_frames_start] *= 1 - tile_weights
|
||||
output[:, out_t_start:after_overlap_frames_start] += (
|
||||
tile_weights * overlap_output
|
||||
)
|
||||
output[:, after_overlap_frames_start:out_t_end] = decoded_tile[
|
||||
:, overlap_frames:
|
||||
]
|
||||
|
||||
# Move to next chunk
|
||||
chunk_start = chunk_end
|
||||
|
||||
# Reshape output to match expected format [batch * frames, height, width, channels]
|
||||
output = output.view(
|
||||
batch * image_frames, output_height, output_width, output.shape[-1]
|
||||
)
|
||||
|
||||
return (output,)
|
||||
|
||||
Reference in New Issue
Block a user