feat: enable old py to work

Author: georgeyiasemis

direct/nn/adain/__init__.py

@@ -1,144 +1,3 @@
-from enum import Enum
+from direct.nn.adain.adain import AdaIN2d, AdaIN3d, NormType
 
-import torch
-from torch import nn
-
-__all__ = ["AdaIN2d", "AdaIN3d"]
-
-
-class NormType(str, Enum):
-    INSTANCE = "instance"
-    ADAIN = "adain"
-
-
-import torch
-from torch import nn
-
-
-class AdaIN2d(nn.Module):
-    """
-    Adaptive Instance Normalization for 2D tensors:
-      x: (B, C, H, W)
-      y: (B, F)  auxiliary vector
-    Produces per-sample, per-channel affine params from y.
-    """
-
-    def __init__(
-        self,
-        num_channels: int,
-        aux_in_features: int,
-        hidden_features: int | tuple[int, ...] | None = None,
-        activation: nn.Module | None = None,
-        eps: float = 1e-5,
-        use_one_plus_gamma: bool = True,
-    ):
-        super().__init__()
-        self.num_channels = num_channels
-        self.eps = eps
-        self.use_one_plus_gamma = use_one_plus_gamma
-
-        if activation is None:
-            activation = nn.SiLU()
-
-        # Build an MLP: aux_in_features -> ... -> 2*num_channels (gamma, beta)
-        if hidden_features is None:
-            hidden = []
-        elif isinstance(hidden_features, int):
-            hidden = [hidden_features]
-        else:
-            hidden = list(hidden_features)
-
-        layers: list[nn.Module] = []
-        in_f = aux_in_features
-        for h in hidden:
-            layers += [nn.Linear(in_f, h), activation]
-            in_f = h
-        layers += [nn.Linear(in_f, 2 * num_channels)]
-        self.mlp = nn.Sequential(*layers)
-
-        # Initialize last layer near-zero so AdaIN starts close to plain IN
-        if isinstance(self.mlp[-1], nn.Linear):
-            nn.init.zeros_(self.mlp[-1].weight)
-            nn.init.zeros_(self.mlp[-1].bias)
-
-    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        # Instance-style normalization over spatial dims (H,W), per (B,C)
-        mean = x.mean(dim=(2, 3), keepdim=True)
-        var = x.var(dim=(2, 3), keepdim=True, unbiased=False)
-        x_norm = (x - mean) / torch.sqrt(var + self.eps)
-
-        # Produce gamma/beta from y
-        params = self.mlp(y)  # (B, 2C)
-        gamma, beta = params.chunk(2, 1)  # each (B, C)
-
-        gamma = gamma.view(-1, self.num_channels, 1, 1)
-        beta = beta.view(-1, self.num_channels, 1, 1)
-
-        if self.use_one_plus_gamma:
-            return x_norm * (1.0 + gamma) + beta
-        return x_norm * gamma + beta
-
-
-class AdaIN3d(nn.Module):
-    """
-    Adaptive Instance Normalization for 3D tensors:
-      x: (B, C, Z, H, W)
-      y: (B, F)  auxiliary vector
-    Produces per-sample, per-channel affine params from y.
-    """
-
-    def __init__(
-        self,
-        num_channels: int,
-        aux_in_features: int,
-        hidden_features: int | tuple[int, ...] | None = None,
-        activation: nn.Module | None = None,
-        eps: float = 1e-5,
-        use_one_plus_gamma: bool = True,
-    ):
-        super().__init__()
-        self.num_channels = num_channels
-        self.eps = eps
-        self.use_one_plus_gamma = use_one_plus_gamma
-
-        if activation is None:
-            activation = nn.SiLU()
-
-        # Build an MLP: aux_in_features -> ... -> 2*num_channels (gamma, beta)
-        if hidden_features is None:
-            hidden = []
-        elif isinstance(hidden_features, int):
-            hidden = [hidden_features]
-        else:
-            hidden = list(hidden_features)
-
-        layers: list[nn.Module] = []
-        in_f = aux_in_features
-        for h in hidden:
-            layers += [nn.Linear(in_f, h), activation]
-            in_f = h
-        layers += [nn.Linear(in_f, 2 * num_channels)]
-        self.mlp = nn.Sequential(*layers)
-
-        # Optional: initialize last layer to near-zero so AdaIN starts close to plain IN
-        if isinstance(self.mlp[-1], nn.Linear):
-            nn.init.zeros_(self.mlp[-1].weight)
-            nn.init.zeros_(self.mlp[-1].bias)
-
-    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-
-        # Instance-style normalization over spatial dims (Z,H,W), per (B,C)
-        mean = x.mean(dim=(2, 3, 4), keepdim=True)
-        var = x.var(dim=(2, 3, 4), keepdim=True, unbiased=False)
-        x_norm = (x - mean) / torch.sqrt(var + self.eps)
-
-        # Produce gamma/beta from y
-        params = self.mlp(y)  # (B, 2C)
-        gamma, beta = params.chunk(2, dim=-1)  # each (B, C)
-
-        gamma = gamma.view(-1, self.num_channels, 1, 1, 1)
-        beta = beta.view(-1, self.num_channels, 1, 1, 1)
-
-        if self.use_one_plus_gamma:
-            return x_norm * (1.0 + gamma) + beta
-        return x_norm * gamma + beta
+__all__ = ["AdaIN2d", "AdaIN3d", "NormType"]