Method decode_video

diffsynth/models/svd_vae_decoder.py:165–199

(
        self, sample,
        batch_time=8, batch_height=128, batch_width=128,
        stride_time=4, stride_height=32, stride_width=32,
        progress_bar=lambda x:x
    )

    

    def decode_video(
        self, sample,
        batch_time=8, batch_height=128, batch_width=128,
        stride_time=4, stride_height=32, stride_width=32,
        progress_bar=lambda x:x
    ):
        sample = sample.permute(1, 0, 2, 3)
        data_device = sample.device
        computation_device = self.conv_in.weight.device
        torch_dtype = sample.dtype
        _, T, H, W = sample.shape

        weight = torch.zeros((1, T, H*8, W*8), dtype=torch_dtype, device=data_device)
        values = torch.zeros((3, T, H*8, W*8), 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).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*8:hr*8, wl*8:wr*8] += sample_batch * mask
            weight[:, tl:tr, hl*8:hr*8, wl*8:wr*8] += mask
        values /= weight
        return values
    
    
    @staticmethod