Method tiled_forward

diffsynth/models/svd_unet.py:337–370

(
        self, sample, timestep, encoder_hidden_states, add_time_id,
        batch_time=25, batch_height=128, batch_width=128,
        stride_time=5, stride_height=64, stride_width=64,
        progress_bar=lambda x:x
    )

    

    def tiled_forward(
        self, sample, timestep, encoder_hidden_states, add_time_id,
        batch_time=25, batch_height=128, batch_width=128,
        stride_time=5, stride_height=64, stride_width=64,
        progress_bar=lambda x:x
    ):
        data_device = sample.device
        computation_device = self.conv_in.weight.device
        torch_dtype = sample.dtype
        T, C, H, W = sample.shape

        weight = torch.zeros((T, 1, H, W), dtype=torch_dtype, device=data_device)
        values = torch.zeros((T, 4, H, W), dtype=torch_dtype, device=data_device)

        # Split tasks
        tasks = []
        for t in range(0, T, stride_time):
            for h in range(0, H, stride_height):
                for w in range(0, W, stride_width):
                    if (t-stride_time >= 0 and t-stride_time+batch_time >= T)\
                        or (h-stride_height >= 0 and h-stride_height+batch_height >= H)\
                        or (w-stride_width >= 0 and w-stride_width+batch_width >= W):
                        continue
                    tasks.append((t, t+batch_time, h, h+batch_height, w, w+batch_width))
        
        # Run
        for tl, tr, hl, hr, wl, wr in progress_bar(tasks):
            sample_batch = sample[tl:tr, :, hl:hr, wl:wr].to(computation_device)
            sample_batch = self.forward(sample_batch, timestep, encoder_hidden_states, add_time_id).to(data_device)
            mask = self.build_mask(sample_batch, is_bound=(tl==0, tr>=T, hl==0, hr>=H, wl==0, wr>=W))
            values[tl:tr, :, hl:hr, wl:wr] += sample_batch * mask
            weight[tl:tr, :, hl:hr, wl:wr] += mask
        values /= weight
        return values
    

    def forward(self, sample, timestep, encoder_hidden_states, add_time_id, use_gradient_checkpointing=False, **kwargs):