model.py

import torch
import torch.nn as nn
from .networks import StemBlock3D, ResBlock3D, ASPP3D, DecoderBlock3D
from .normalizing_flow import CyLNormalizingFlow

class UnifiedFoundationModel(nn.Module):
    """
    3D V-Net Encoder-Decoder Architecture.
    Acts as the Foundation Model (FM) trained purely on the Source Domain.
    """
    def __init__(self, in_channels=1, num_classes=6):
        super().__init__()
        # Unified Encoder
        self.c1 = StemBlock3D(in_channels, 16, stride=1)
        self.c2 = ResBlock3D(16, 32, stride=2)
        self.c3 = ResBlock3D(32, 64, stride=2)
        self.c4 = ResBlock3D(64, 128, stride=2)
        self.b1 = ASPP3D(128, 256) # Bottleneck Latent Space (Z)
        
        # Unified Decoder
        self.d1 = DecoderBlock3D(gate_channels=256, skip_channels=64, out_channels=128)
        self.d2 = DecoderBlock3D(gate_channels=128, skip_channels=32, out_channels=64)
        self.d3 = DecoderBlock3D(gate_channels=64, skip_channels=16, out_channels=32)
        self.aspp_out = ASPP3D(32, 16) # Output Topological Space
        
        # Final Classification Head
        self.head = nn.Conv3d(16, num_classes, kernel_size=1)

    def encode(self, x):
        s1 = self.c1(x)
        s2 = self.c2(s1)
        s3 = self.c3(s2)
        s4 = self.c4(s3)
        z = self.b1(s4)
        return s1, s2, s3, s4, z

    def decode(self, s1, s2, s3, z):
        x = self.d1(z, s3)
        x = self.d2(x, s2)
        x = self.d3(x, s1)
        features = self.aspp_out(x)
        logits = self.head(features)
        return logits, features

class CyLAdapterModel(nn.Module):
    """
    Cross-Modality Latent Adapter.
    Wraps the Frozen Foundation Model and aligns out-of-distribution 
    Target latent spaces into the Source latent space via Normalizing Flows.
    """
    def __init__(self, foundation_model):
        super().__init__()
        self.foundation = foundation_model
        
        # Forward Latent Normalizing Flow at the Bottleneck (256 channels)
        self.T_flow = CyLNormalizingFlow(channels=256, num_layers=3)

    def forward_foundation(self, x):
        """Phase 1: Run standard inference on Source Data without adaptation."""
        s1, s2, s3, s4, z = self.foundation.encode(x)
        logits, features = self.foundation.decode(s1, s2, s3, z)
        return logits, z

    def forward_adapter(self, x):
        """
        Phase 2: Run Target Data inference. 
        Target latent space is mapped to Source latent space using T_flow.
        The aligned features are then processed by the frozen Source Decoder/Head.
        """
        s1, s2, s3, s4, z_target = self.foundation.encode(x)
        
        z_source_aligned = self.T_flow(z_target, reverse=False)
        logits, features = self.foundation.decode(s1, s2, s3, z_source_aligned)
        
        return logits, z_target, z_source_aligned