Method forward

core/drag_embedding.py:106–151

(self, drags_start, drags_end)

            self.channels = self.channels[self.drag_layer_idx:self.drag_layer_idx+1]
        
    def forward(self, drags_start, drags_end):
        scales = self.scales
        channels = self.channels
        
        B, V, N, _ = drags_start.shape  
        
        drags_start = drags_start.view(B*V, N, -1)
        drags_end = drags_end.view(B*V, N, -1)


        multi_scale_merge_start_embeddings = []
        multi_scale_merge_end_embeddings = []
        
        for idx, scale in enumerate(scales):
            drags_start_embeddings = []
            drags_end_embeddings = []
            for i in range(N):
                drag_start_embedding = self.fourier_embedder(drags_start[:, i, :])
                drags_start_embeddings.append(self.linears_drags[idx](drag_start_embedding))
            drags_start_embeddings = torch.stack(drags_start_embeddings, dim=1)
            
            for i in range(N):
                drag_end_embedding = self.fourier_embedder(drags_end[:, i, :])
                drags_end_embeddings.append(self.linears_drags[idx](drag_end_embedding))
            drags_end_embeddings = torch.stack(drags_end_embeddings, dim=1)
        
            merge_start_embeddings = torch.zeros((B*V, channels[idx]//2, scale, scale)).to(drag_start_embedding.device)
            merge_end_embeddings = torch.zeros((B*V, channels[idx]//2, scale, scale)).to(drag_start_embedding.device)
            downsample_ratio = 512 // scale
            
            for i in range(B*V):
                for j in range(N):
                    merge_start_embeddings[i, :, int(drags_start[i, j, 0]) // downsample_ratio, 
                                     int(drags_start[i, j, 1]) // downsample_ratio] += drags_start_embeddings[i,j,:]
                    merge_end_embeddings[i, :, int(drags_end[i, j, 0]) // downsample_ratio, 
                                     int(drags_end[i, j, 1]) // downsample_ratio] += drags_end_embeddings[i,j, :]
                    # merge_end_embeddings[i, :, int(drags_start[i, j, 0]) // downsample_ratio, 
                    #                 int(drags_start[i, j, 1]) // downsample_ratio] += drags_end_embeddings[i,j, :]
            # Add Gaussian Blur
            merge_start_embeddings = self.gaussian_blur(merge_start_embeddings)
            merge_end_embeddings = self.gaussian_blur(merge_end_embeddings)
            
            multi_scale_merge_start_embeddings.append(merge_start_embeddings)
            multi_scale_merge_end_embeddings.append(merge_end_embeddings)
        
        return multi_scale_merge_start_embeddings, multi_scale_merge_end_embeddings