from transformers import PreTrainedModel
from .w_net_3d import WNet3dUNet, WNet3dAttUNet, WNet3dUNetMSS
from .WNetConfigs import WNet3DConfig, AttWNet3DConfig, WNetMSS3DConfig 

class WNet3D(PreTrainedModel):
    config_class = WNet3DConfig
    def __init__(self, config):
        super().__init__(config)
        self.model = WNet3dUNet(
            in_ch=config.in_ch, 
            out_ch=config.out_ch, 
            init_features=config.init_features)
    def forward(self, x):
        return self.model(x)
    
class AttWNet3D(PreTrainedModel):
    config_class = AttWNet3DConfig
    def __init__(self, config):
        super().__init__(config)
        self.model = WNet3dAttUNet(
            in_ch=config.in_ch, 
            out_ch=config.out_ch, 
            init_features=config.init_features)
    def forward(self, x):
        return self.model(x)
    
class WNetMSS3D(PreTrainedModel):
    config_class = WNetMSS3DConfig
    def __init__(self, config):
        super().__init__(config)
        self.model = WNet3dUNetMSS(
            in_ch=config.in_ch, 
            out_ch=config.out_ch, 
            init_features=config.init_features)
    def forward(self, x):
        return self.model(x)