import logging import torch from .nodes_registry import comfy_node @comfy_node( name="LTXVTiledVAEDecode", ) class LTXVTiledVAEDecode: @classmethod def INPUT_TYPES(s): return { "required": { "vae": ("VAE",), "latents": ("LATENT",), "horizontal_tiles": ("INT", {"default": 1, "min": 1, "max": 6}), "vertical_tiles": ("INT", {"default": 1, "min": 1, "max": 6}), "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",) RETURN_NAMES = ("image",) FUNCTION = "decode" CATEGORY = "latent" def decode( self, vae, latents, horizontal_tiles, vertical_tiles, overlap, last_frame_fix, working_device="auto", working_dtype="auto", ): # Get the latent samples samples = latents["samples"] if last_frame_fix: # Repeat the last frame along dimension 2 (frames) # samples: [batch, channels, frames, height, width] last_frame = samples[ :, :, -1:, :, : ] # shape: [batch, channels, 1, height, width] samples = torch.cat([samples, last_frame], dim=2) 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 # Calculate tile sizes with overlap base_tile_height = (height + (vertical_tiles - 1) * overlap) // vertical_tiles base_tile_width = (width + (horizontal_tiles - 1) * overlap) // horizontal_tiles # Initialize output tensor and weight tensor # VAE decode returns images in format [batch, height, width, channels] 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): # Calculate tile boundaries h_start = h * (base_tile_width - overlap) v_start = v * (base_tile_height - overlap) # Adjust end positions for edge tiles h_end = ( min(h_start + base_tile_width, width) if h < horizontal_tiles - 1 else width ) v_end = ( min(v_start + base_tile_height, height) if v < vertical_tiles - 1 else height ) # Calculate actual tile dimensions tile_height = v_end - v_start tile_width = h_end - h_start 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] # Create tile latents dict tile_latents = {"samples": tile} # Decode the tile decoded_tile = vae.decode(tile_latents["samples"]) # Calculate output tile boundaries out_h_start = v_start * height_scale_factor out_h_end = v_end * height_scale_factor out_w_start = h_start * width_scale_factor out_w_end = h_end * width_scale_factor # Create weight mask for this tile tile_out_height = out_h_end - out_h_start tile_out_width = out_w_end - out_w_start tile_weights = torch.ones( (batch, image_frames, tile_out_height, tile_out_width, 1), device=decoded_tile.device, dtype=decoded_tile.dtype, ) # Calculate overlap regions in output space overlap_out_h = overlap * height_scale_factor overlap_out_w = overlap * width_scale_factor # Apply horizontal blending weights if h > 0: # Left overlap h_blend = torch.linspace( 0, 1, overlap_out_w, device=decoded_tile.device ) tile_weights[:, :, :, :overlap_out_w, :] *= h_blend.view( 1, 1, 1, -1, 1 ) if h < horizontal_tiles - 1: # Right overlap h_blend = torch.linspace( 1, 0, overlap_out_w, device=decoded_tile.device ) tile_weights[:, :, :, -overlap_out_w:, :] *= h_blend.view( 1, 1, 1, -1, 1 ) # Apply vertical blending weights if v > 0: # Top overlap v_blend = torch.linspace( 0, 1, overlap_out_h, device=decoded_tile.device ) tile_weights[:, :, :overlap_out_h, :, :] *= v_blend.view( 1, 1, -1, 1, 1 ) if v < vertical_tiles - 1: # Bottom overlap v_blend = torch.linspace( 1, 0, overlap_out_h, device=decoded_tile.device ) tile_weights[:, :, -overlap_out_h:, :, :] *= v_blend.view( 1, 1, -1, 1, 1 ) # 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.to(target_device, target_dtype) # Normalize by weights output /= weights + 1e-8 # Reshape output to match expected format [batch * frames, height, width, channels] output = output.view( batch * image_frames, output_height, output_width, output.shape[-1] ) if last_frame_fix: 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,)