Method _process_temporal_chunks

looping_sampler.py:310–523

Process all temporal chunks for a single spatial tile.

(
        self,
        tile_config: TileConfig,
        sampling_config: SamplingConfig,
        model_config: ModelConfig,
    )

        )

    def _process_temporal_chunks(
        self,
        tile_config: TileConfig,
        sampling_config: SamplingConfig,
        model_config: ModelConfig,
    ):
        """Process all temporal chunks for a single spatial tile."""
        chunk_index = 0
        tile_out_latents = None
        first_tile_out_latents = None

        for i_temporal_tile, (start_index, end_index) in enumerate(
            zip(
                range(
                    0,
                    tile_config.tile_latents["samples"].shape[2]
                    + sampling_config.temporal_tile_size
                    - sampling_config.temporal_overlap,
                    sampling_config.temporal_tile_size
                    - sampling_config.temporal_overlap,
                ),
                range(
                    sampling_config.temporal_tile_size,
                    tile_config.tile_latents["samples"].shape[2]
                    + sampling_config.temporal_tile_size
                    - sampling_config.temporal_overlap,
                    sampling_config.temporal_tile_size
                    - sampling_config.temporal_overlap,
                ),
            )
        ):
            if tile_config.tile_guiding_latents is not None:
                guiding_latent_chunk = LTXVSelectLatents().select_latents(
                    tile_config.tile_guiding_latents,
                    start_index,
                    min(
                        end_index - 1,
                        tile_config.tile_guiding_latents["samples"].shape[2] - 1,
                    ),
                )[0]
            else:
                guiding_latent_chunk = None

            latent_chunk = LTXVSelectLatents().select_latents(
                tile_config.tile_latents,
                start_index,
                min(end_index - 1, tile_config.tile_latents["samples"].shape[2] - 1),
            )[0]

            print(
                "Processing temporal chunk at index",
                start_index,
                "to",
                min(end_index - 1, tile_config.tile_latents["samples"].shape[2] - 1),
            )

            seed_offset = self._get_per_tile_value(
                sampling_config.per_tile_seed_offsets, i_temporal_tile