(self, fourier_freqs=8, scales=[256, 128, 64, 32, 16, 8], channels=[64, 64, 128, 256, 512, 1024], drag_layer_idx=None)
| 75 | |
| 76 | class DragPositionNetMultiScale(nn.Module): |
| 77 | def __init__(self, fourier_freqs=8, scales=[256, 128, 64, 32, 16, 8], channels=[64, 64, 128, 256, 512, 1024], drag_layer_idx=None): |
| 78 | super().__init__() |
| 79 | self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs) |
| 80 | self.position_dim = fourier_freqs*2*2 # 2 for sin and cos, 2 for 2 dims (x1, y1） or (x2, y2) |
| 81 | |
| 82 | # -------------------------------------------------------------- # |
| 83 | self.linear_drags = [] |
| 84 | for i in range(len(channels)): |
| 85 | if drag_layer_idx is not None and i != drag_layer_idx: |
| 86 | continue |
| 87 | self.linear_drags.append(nn.Sequential( |
| 88 | nn.Linear(self.position_dim, 128), |
| 89 | nn.SiLU(), |
| 90 | nn.Linear(128, 256), |
| 91 | nn.SiLU(), |
| 92 | nn.Linear(256, channels[i]//2), |
| 93 | )) |
| 94 | |
| 95 | self.linears_drags = nn.ModuleList(self.linear_drags) |
| 96 | self.gaussian_blur = GaussianBlur(kernel_size=5, sigma=1.0) |
| 97 | |
| 98 | self.scales = scales |
| 99 | self.channels = channels |
| 100 | |
| 101 | if drag_layer_idx is not None: |
| 102 | self.drag_layer_idx = drag_layer_idx |
| 103 | self.scales = self.scales[self.drag_layer_idx:self.drag_layer_idx+1] |
| 104 | self.channels = self.channels[self.drag_layer_idx:self.drag_layer_idx+1] |
| 105 | |
| 106 | def forward(self, drags_start, drags_end): |
| 107 | scales = self.scales |
nothing calls this directly
no test coverage detected