Spaces:
Running
Running
| import torch | |
| class StyleTransferController(torch.nn.Module): | |
| def __init__( | |
| self, | |
| num_control_params, | |
| edim, | |
| hidden_dim=256, | |
| agg_method="mlp", | |
| ): | |
| """Plugin parameter controller module to map from input to target style. | |
| Args: | |
| num_control_params (int): Number of plugin parameters to predicted. | |
| edim (int): Size of the encoder representations. | |
| hidden_dim (int, optional): Hidden size of the 3-layer parameter predictor MLP. Default: 256 | |
| agg_method (str, optional): Input/reference embed aggregation method ["conv" or "linear", "mlp"]. Default: "mlp" | |
| """ | |
| super().__init__() | |
| self.num_control_params = num_control_params | |
| self.edim = edim | |
| self.hidden_dim = hidden_dim | |
| self.agg_method = agg_method | |
| if agg_method == "conv": | |
| self.agg = torch.nn.Conv1d( | |
| 2, | |
| 1, | |
| kernel_size=129, | |
| stride=1, | |
| padding="same", | |
| bias=False, | |
| ) | |
| mlp_in_dim = edim | |
| elif agg_method == "linear": | |
| self.agg = torch.nn.Linear(edim * 2, edim) | |
| elif agg_method == "mlp": | |
| self.agg = None | |
| mlp_in_dim = edim * 2 | |
| else: | |
| raise ValueError(f"Invalid agg_method = {self.agg_method}.") | |
| self.mlp = torch.nn.Sequential( | |
| torch.nn.Linear(mlp_in_dim, hidden_dim), | |
| torch.nn.LeakyReLU(0.01), | |
| torch.nn.Linear(hidden_dim, hidden_dim), | |
| torch.nn.LeakyReLU(0.01), | |
| torch.nn.Linear(hidden_dim, num_control_params), | |
| torch.nn.Sigmoid(), # normalize between 0 and 1 | |
| ) | |
| def forward(self, e_x, e_y, z=None): | |
| """Forward pass to generate plugin parameters. | |
| Args: | |
| e_x (tensor): Input signal embedding of shape (batch, edim) | |
| e_y (tensor): Target signal embedding of shape (batch, edim) | |
| Returns: | |
| p (tensor): Estimated control parameters of shape (batch, num_control_params) | |
| """ | |
| # use learnable projection | |
| if self.agg_method == "conv": | |
| e_xy = torch.stack((e_x, e_y), dim=1) # concat on channel dim | |
| e_xy = self.agg(e_xy) | |
| elif self.agg_method == "linear": | |
| e_xy = torch.cat((e_x, e_y), dim=-1) # concat on embed dim | |
| e_xy = self.agg(e_xy) | |
| else: | |
| e_xy = torch.cat((e_x, e_y), dim=-1) # concat on embed dim | |
| # pass through MLP to project to control parametesr | |
| p = self.mlp(e_xy.squeeze(1)) | |
| return p | |