Upload folder using huggingface_hub

README.md

@@ -0,0 +1,75 @@
+---
+pipeline_tag: image-segmentation
+library_name: transformers
+tags:
+- ultrasound
+- medical-image-segmentation
+- attention-unet
+- custom-pipeline
+---
+
+# Cond-UNet Attention for Ultrasound Segmentation
+
+Cond-UNet Attention is a binary ultrasound segmentation model based on an
+attention-conditioned U-Net. It was trained to predict a foreground mask from
+an RGB ultrasound image.
+
+## Model Details
+
+- Architecture: attention-conditioned U-Net
+- U-Net depth: 5
+- Base channels: 16
+- Input resolution: 512 x 512
+- Attention embedding dimension: 768
+- Attention patch size: 8
+- Output: one foreground logit per pixel
+- Organ conditioning: optional; omitted IDs use the unknown-organ token (`-1`)
+
+The model does not use DWT features or the optional shape-conditioning branch.
+
+## Usage
+
+This repository contains custom Transformers code. Pass `trust_remote_code=True`
+when loading it.
+
+```python
+from transformers import pipeline
+
+segmenter = pipeline(
+    "image-segmentation",
+    model="YOUR_ORG/YOUR_MODEL",
+    trust_remote_code=True,
+)
+
+result = segmenter("ultrasound.png")
+mask = result["mask"]
+```
+
+When organ metadata is known, pass its integer class ID:
+
+```python
+result = segmenter("ultrasound.png", organ_id=3)
+```
+
+If `organ_id` is not provided, the model automatically uses `-1`, matching the
+unknown-organ conditioning used in training.
+
+## Output
+
+The pipeline applies a sigmoid to the foreground logit and returns a binary
+PIL mask thresholded at 0.5. Change the threshold if needed:
+
+```python
+result = segmenter("ultrasound.png", threshold=0.4)
+```
+
+## Preprocessing
+
+Images are converted to RGB, resized directly to 512 x 512, and normalized in
+the 0-255 value range with mean `[123.675, 116.28, 103.53]` and standard
+deviation `[58.395, 57.12, 57.375]`.
+
+## Limitations
+
+This model is intended for research use only. It is not a clinical diagnostic
+device and must not be used as the sole basis for medical decisions.

config.json

@@ -0,0 +1,50 @@
+{
+  "architectures": [
+    "CondUNetForSemanticSegmentation"
+  ],
+  "attn_start": 0,
+  "auto_map": {
+    "AutoConfig": "configuration_cond_unet.CondUNetConfig",
+    "AutoImageProcessor": "image_processing_cond_unet.CondUNetImageProcessor",
+    "AutoModelForImageSegmentation": "modeling_cond_unet.CondUNetForSemanticSegmentation"
+  },
+  "custom_pipelines": {
+    "image-segmentation": {
+      "impl": "pipeline.CondUNetImageSegmentationPipeline",
+      "pt": [
+        "AutoModelForImageSegmentation"
+      ],
+      "type": "image"
+    }
+  },
+  "depth": 5,
+  "dtype": "float32",
+  "dwt_bands": [
+    "LL",
+    "LH",
+    "HL",
+    "HH"
+  ],
+  "emb_dim": 768,
+  "id2label": {
+    "0": "foreground"
+  },
+  "image_size": 512,
+  "in_channels": 3,
+  "keep_aspect_ratio": false,
+  "label2id": {
+    "foreground": 0
+  },
+  "model_type": "cond_unet",
+  "n_heads": 8,
+  "n_organs": 10,
+  "patch_size": 8,
+  "shape_res": 32,
+  "size": 16,
+  "transformers_version": "5.16.1",
+  "unknown_organ_id": -1,
+  "use_attn": true,
+  "use_dwt": false,
+  "use_shape": false,
+  "wavelet": "haar"
+}

configuration_cond_unet.py

@@ -0,0 +1,49 @@
+from transformers import PretrainedConfig
+
+
+class CondUNetConfig(PretrainedConfig):
+    model_type = "cond_unet"
+
+    def __init__(
+        self,
+        in_channels=3,
+        num_labels=1,
+        n_organs=10,
+        size=32,
+        depth=5,
+        attn_start=0,
+        use_attn=True,
+        image_size=512,
+        patch_size=8,
+        emb_dim=768,
+        n_heads=8,
+        use_dwt=False,
+        wavelet="haar",
+        dwt_bands=None,
+        use_shape=False,
+        shape_res=32,
+        unknown_organ_id=-1,
+        keep_aspect_ratio=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.num_labels = num_labels
+        self.n_organs = n_organs
+        self.size = size
+        self.depth = depth
+        self.attn_start = attn_start
+        self.use_attn = use_attn
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.emb_dim = emb_dim
+        self.n_heads = n_heads
+        self.use_dwt = use_dwt
+        self.wavelet = wavelet
+        self.dwt_bands = dwt_bands or ["LL", "LH", "HL", "HH"]
+        self.use_shape = use_shape
+        self.shape_res = shape_res
+        self.unknown_organ_id = unknown_organ_id
+        self.keep_aspect_ratio = keep_aspect_ratio
+        self.id2label = {0: "foreground"}
+        self.label2id = {"foreground": 0}

image_processing_cond_unet.py

@@ -0,0 +1,79 @@
+from typing import Optional, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+
+
+class CondUNetImageProcessor(BaseImageProcessor):
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        image_size=512,
+        keep_aspect_ratio=False,
+        mean=None,
+        std=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.image_size = image_size
+        self.keep_aspect_ratio = keep_aspect_ratio
+        self.mean = mean or [123.675, 116.28, 103.53]
+        self.std = std or [58.395, 57.12, 57.375]
+
+    def preprocess(
+        self,
+        images: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        return_tensors: Optional[Union[str, torch.Tensor]] = None,
+        **kwargs,
+    ):
+        if not isinstance(images, (list, tuple)):
+            images = [images]
+        pixel_values = [self._preprocess_image(image) for image in images]
+        return BatchFeature(data={"pixel_values": pixel_values}, tensor_type=return_tensors)
+
+    def _preprocess_image(self, image):
+        if isinstance(image, Image.Image):
+            image = np.array(image.convert("RGB"), copy=True)
+        if isinstance(image, np.ndarray):
+            image = torch.from_numpy(image)
+        if image.ndim != 3:
+            raise ValueError("Expected an HWC or CHW RGB image.")
+        if image.shape[-1] in (1, 3):
+            image = image.permute(2, 0, 1)
+        if image.shape[0] == 1:
+            image = image.expand(3, -1, -1)
+        if image.shape[0] != 3:
+            raise ValueError("Cond-UNet requires one or three input channels.")
+        image = image.to(dtype=torch.float32)
+        if image.max() <= 1:
+            image = image * 255.0
+        height, width = image.shape[-2:]
+        if self.keep_aspect_ratio:
+            scale = self.image_size / max(height, width)
+            new_height = int(height * scale)
+            new_width = int(width * scale)
+            new_height += new_height % 2
+            new_width += new_width % 2
+            image = F.interpolate(
+                image.unsqueeze(0),
+                size=(new_height, new_width),
+                mode="bilinear",
+                align_corners=False,
+            ).squeeze(0)
+            pad_left = (self.image_size - new_width) // 2
+            pad_top = (self.image_size - new_height) // 2
+            image = F.pad(image, (pad_left, self.image_size - new_width - pad_left, pad_top, self.image_size - new_height - pad_top))
+        else:
+            image = F.interpolate(
+                image.unsqueeze(0),
+                size=(self.image_size, self.image_size),
+                mode="bilinear",
+                align_corners=False,
+            ).squeeze(0)
+        mean = torch.tensor(self.mean, dtype=image.dtype).view(-1, 1, 1)
+        std = torch.tensor(self.std, dtype=image.dtype).view(-1, 1, 1)
+        return (image - mean) / std

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7dabb114d6c30c42b81eab8354cade8fd9f8335cc70dc05070b4b65e4ac7f747
+size 191410692

modeling_cond_unet.py

@@ -0,0 +1,74 @@
+from typing import Optional
+
+import torch
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import SemanticSegmenterOutput
+
+from .configuration_cond_unet import CondUNetConfig
+
+try:
+    from .unet_attn import UNet2DAttn
+except ModuleNotFoundError:
+    # The release exporter bundles these modules into the Hub repository root.
+    from nets.unet_attn import UNet2DAttn
+
+
+class CondUNetForSemanticSegmentation(PreTrainedModel):
+    config_class = CondUNetConfig
+    main_input_name = "pixel_values"
+
+    def __init__(self, config: CondUNetConfig):
+        super().__init__(config)
+        self.unet = UNet2DAttn(
+            in_channels=config.in_channels,
+            num_classes=config.num_labels,
+            n_organs=config.n_organs,
+            size=config.size,
+            depth=config.depth,
+            attn_start=config.attn_start,
+            use_attn=config.use_attn,
+            img_size=config.image_size,
+            patch_size=config.patch_size,
+            emb_dim=config.emb_dim,
+            n_heads=config.n_heads,
+            distill=False,
+            distill_unet=False,
+            use_dwt=config.use_dwt,
+            wavelet=config.wavelet,
+            dwt_bands=config.dwt_bands,
+            use_shape=config.use_shape,
+            shape_res=config.shape_res,
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        organ_id: Optional[torch.LongTensor] = None,
+        labels: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ):
+        if organ_id is None:
+            organ_id = torch.full(
+                (pixel_values.shape[0],),
+                self.config.unknown_organ_id,
+                device=pixel_values.device,
+                dtype=torch.long,
+            )
+        else:
+            organ_id = organ_id.to(device=pixel_values.device, dtype=torch.long)
+
+        outputs = self.unet(
+            pixel_values=pixel_values,
+            organ_id=organ_id,
+            masks=labels,
+            **kwargs,
+        )
+        logits = outputs["logits"]
+        if logits.ndim == 3:
+            logits = logits.unsqueeze(1)
+        loss = outputs["loss"] if labels is not None else None
+        if return_dict is False:
+            return (loss, logits) if loss is not None else (logits,)
+        return SemanticSegmenterOutput(loss=loss, logits=logits)

pipeline.py

@@ -0,0 +1,50 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from transformers.pipelines.base import Pipeline
+
+from .image_processing_cond_unet import CondUNetImageProcessor
+
+
+class CondUNetImageSegmentationPipeline(Pipeline):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.image_processor is None:
+            self.image_processor = CondUNetImageProcessor(
+                image_size=self.model.config.image_size,
+                keep_aspect_ratio=self.model.config.keep_aspect_ratio,
+            )
+
+    def _sanitize_parameters(self, organ_id=None, threshold=None, **kwargs):
+        preprocess_kwargs = {}
+        postprocess_kwargs = {}
+        if organ_id is not None:
+            preprocess_kwargs["organ_id"] = organ_id
+        if threshold is not None:
+            postprocess_kwargs["threshold"] = threshold
+        return preprocess_kwargs, {}, postprocess_kwargs
+
+    def preprocess(self, image, organ_id=None, **kwargs):
+        if not isinstance(image, Image.Image):
+            image = Image.open(image).convert("RGB")
+        else:
+            image = image.convert("RGB")
+        width, height = image.size
+        inputs = self.image_processor(images=image, return_tensors="pt")
+        inputs["original_size"] = (height, width)
+        if organ_id is not None:
+            inputs["organ_id"] = torch.tensor([organ_id], dtype=torch.long)
+        return inputs
+
+    def _forward(self, model_inputs, **kwargs):
+        original_size = model_inputs.pop("original_size")
+        outputs = self.model(**model_inputs)
+        return {"logits": outputs.logits, "original_size": original_size}
+
+    def postprocess(self, model_outputs, threshold=0.5, **kwargs):
+        logits = model_outputs["logits"]
+        height, width = model_outputs["original_size"]
+        probabilities = torch.sigmoid(F.interpolate(logits, size=(height, width), mode="bilinear", align_corners=False))[0, 0]
+        mask = (probabilities >= threshold).to(torch.uint8).cpu().numpy() * 255
+        return {"label": "foreground", "mask": Image.fromarray(mask), "score": float(probabilities.mean())}

preprocessor_config.json

@@ -0,0 +1,15 @@
+{
+  "image_processor_type": "CondUNetImageProcessor",
+  "image_size": 512,
+  "keep_aspect_ratio": false,
+  "mean": [
+    123.675,
+    116.28,
+    103.53
+  ],
+  "std": [
+    58.395,
+    57.12,
+    57.375
+  ]
+}

requirements.txt

@@ -0,0 +1,8 @@
+transformers>=4.40
+torch
+torchvision
+safetensors
+einops
+ptwt
+wandb
+peft

segm_net.py

@@ -0,0 +1,1051 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+import numpy as np
+import wandb
+from peft import LoraConfig, get_peft_model
+from torchvision.transforms import v2
+from .unet_base import BaseUnet
+
+
+class SplitQKVLinear(nn.Module):
+    def __init__(self, base_layer: nn.Linear):
+        super().__init__()
+        if base_layer.out_features % 3 != 0:
+            raise ValueError("Expected fused qkv projection with out_features divisible by 3.")
+
+        self.proj_dim = base_layer.out_features // 3
+        self.q_proj = nn.Linear(base_layer.in_features, self.proj_dim, bias=base_layer.bias is not None)
+        self.k_proj = nn.Linear(base_layer.in_features, self.proj_dim, bias=base_layer.bias is not None)
+        self.v_proj = nn.Linear(base_layer.in_features, self.proj_dim, bias=base_layer.bias is not None)
+
+        with torch.no_grad():
+            self.q_proj.weight.copy_(base_layer.weight[: self.proj_dim])
+            self.k_proj.weight.copy_(base_layer.weight[self.proj_dim : 2 * self.proj_dim])
+            self.v_proj.weight.copy_(base_layer.weight[2 * self.proj_dim :])
+
+            if base_layer.bias is not None:
+                self.q_proj.bias.copy_(base_layer.bias[: self.proj_dim])
+                self.k_proj.bias.copy_(base_layer.bias[self.proj_dim : 2 * self.proj_dim])
+                self.v_proj.bias.copy_(base_layer.bias[2 * self.proj_dim :])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+        return torch.cat((q, k, v), dim=-1)
+
+
+def replace_qkv_with_split_projections(module: nn.Module):
+    replaced_layers = 0
+    for child_name, child_module in list(module.named_children()):
+        if (
+            child_name == "qkv"
+            and isinstance(child_module, nn.Linear)
+            and child_module.out_features == child_module.in_features * 3
+        ):
+            setattr(module, child_name, SplitQKVLinear(child_module))
+            replaced_layers += 1
+            continue
+
+        replaced_layers += replace_qkv_with_split_projections(child_module)
+
+    return replaced_layers
+
+
+def apply_lora_to_qv_projections(module: nn.Module, rank: int, alpha: float | None = None):
+    replaced_layers = replace_qkv_with_split_projections(module)
+    if replaced_layers == 0:
+        raise ValueError("No qkv projections found in MedSAM image encoder for LoRA injection.")
+
+    lora_config = LoraConfig(
+        r=rank,
+        lora_alpha=(alpha if alpha is not None else rank),
+        target_modules=["q_proj", "v_proj"],
+        bias="none",
+    )
+    peft_module = get_peft_model(module, lora_config)
+    return peft_module, replaced_layers
+
+def pad_to_2d(x: torch.Tensor, stride: int):
+    h, w = x.shape[-2:]
+
+    new_h = h if h % stride == 0 else h + stride - (h % stride)
+    new_w = w if w % stride == 0 else w + stride - (w % stride)
+
+    top = (new_h - h) // 2
+    bottom = (new_h - h) - top
+    left = (new_w - w) // 2
+    right = (new_w - w) - left
+
+    pads = (left, right, top, bottom)
+    x_pad = F.pad(x, pads, mode="constant", value=0)
+    return x_pad, pads
+
+
+def unpad_2d(x: torch.Tensor, pads):
+    left, right, top, bottom = pads
+
+    if top or bottom:
+        end_h = -bottom if bottom > 0 else None
+        x = x[:, :, top:end_h, :]
+
+    if left or right:
+        end_w = -right if right > 0 else None
+        x = x[:, :, :, left:end_w]
+
+    return x
+
+
+class ConvBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        conv_kwargs={"kernel_size": 3, "stride": 1, "padding": 1},
+    ):
+        super(ConvBlock, self).__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, **conv_kwargs),
+            nn.InstanceNorm2d(out_channels),
+            nn.LeakyReLU(),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+# Single encoder block
+class DownConvBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: list,
+        out_channels: list,
+        conv_kwargs={"kernel_size": 3, "stride": 1, "padding": 1},
+    ):
+        super(DownConvBlock, self).__init__()
+
+        assert len(in_channels) == len(
+            out_channels
+        ), f"in_channels length is {len(in_channels)} while out_channels is {len(out_channels)}"
+
+        # Variable number of convolutional block in each layer, based on the in_channels and out_channels length
+        self.conv_blocks = nn.ModuleList(
+            [
+                ConvBlock(in_ch, out_ch, conv_kwargs)
+                for in_ch, out_ch in zip(in_channels, out_channels)
+            ]
+        )
+
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        for block in self.conv_blocks:
+            x = block(x)
+        return self.pool(x), x
+
+
+# Single decoder block
+class UpConvBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: list,
+        out_channels: list,
+        up_conv=True,
+        conv_kwargs={"kernel_size": 3, "stride": 1, "padding": 1},
+        upconv_kwargs={"kernel_size": 2, "stride": 2},
+    ):
+        super(UpConvBlock, self).__init__()
+
+        assert len(in_channels) == len(
+            out_channels
+        ), f"in_channels length is {len(in_channels)} while out_channels is {len(out_channels)}"
+
+        # Variable number of convolutional block in each layer, based on the in_channels and out_channels length
+        self.conv_blocks = nn.ModuleList(
+            [
+                ConvBlock(in_ch, out_ch, conv_kwargs)
+                for in_ch, out_ch in zip(in_channels, out_channels)
+            ]
+        )
+
+        self.up_conv = up_conv
+        if self.up_conv:
+            self.up_conv_op = nn.ConvTranspose2d(
+                out_channels[-1], out_channels[-1], **upconv_kwargs
+            )
+
+    def forward(self, x):
+        for block in self.conv_blocks:
+            x = block(x)
+
+        if self.up_conv:
+            return self.up_conv_op(x)
+        else:
+            return x
+
+
+class FiLM2d(nn.Module):
+    """
+    Feature-wise Linear Modulation for a 2-D feature map.
+    (γ, β) are generated from a learned embedding of organ_id.
+    """
+
+    def __init__(
+        self,
+        n_organs: int,
+        in_channels: int,
+        emb_dim: int | None = None,
+        hidden: int | None = None,
+    ):
+        super().__init__()
+        hidden = hidden or 2 * in_channels
+        self.embed = nn.Embedding(n_organs + 1, emb_dim)
+
+        self.mlp = nn.Sequential(
+            nn.Linear(emb_dim, hidden),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden, 2 * in_channels),  # → [β‖γ]
+        )
+
+        # initialise so that FiLM starts as identity: γ≈1, β≈0
+        nn.init.zeros_(self.mlp[-1].weight)
+        nn.init.constant_(self.mlp[-1].bias[:in_channels], 0)  # β
+        nn.init.constant_(self.mlp[-1].bias[in_channels:], 1)  # γ
+
+    def compute_modulation(self, organ_id: torch.Tensor):
+        """Return FiLM coefficients and conditioning embedding for a batch."""
+        B = organ_id.shape[0]
+        device = organ_id.device
+        mask = organ_id >= 0  # [B]
+        q_org = self.embed(organ_id.clamp(min=0))  # [B, D] (dummy for unknown)
+        q_img = self.embed(
+            torch.tensor([self.embed.weight.shape[0] - 1], device=device).expand(B)
+        )
+        q = torch.where(mask[:, None], q_org, q_img)
+        beta_gamma = self.mlp(q)  # (B, 2C)
+        beta, gamma = beta_gamma.chunk(2, dim=-1)  # each (B, C)
+        beta = beta.unsqueeze(-1).unsqueeze(-1)
+        gamma = gamma.unsqueeze(-1).unsqueeze(-1)
+        return gamma, beta, q
+
+    def forward(self, x: torch.Tensor, organ_id: torch.Tensor):
+        """
+        x : (B, C, H, W)
+        organ_id : (B,) integer 0…n_organs-1
+        """
+        gamma, beta, _ = self.compute_modulation(organ_id)
+        return gamma * x + beta
+
+
+class SharedFiLMModulator:
+    """
+    Compatibility helper that exposes the same collection interface used by
+    UNet2DAttn notebooks: compute all per-layer FiLM coefficients in one call.
+    """
+
+    def __init__(self, model: "UNet2DFiLM"):
+        self.model = model
+
+    def compute_all_modulations(
+        self,
+        original_img: torch.Tensor,
+        organ_id: torch.Tensor,
+        layer_configs: list[tuple[int, int]],
+        return_compact: bool = False,
+    ):
+        del original_img  # FiLM modulation depends only on organ_id in this model.
+
+        modulations = {}
+        compact_modulations = {} if return_compact else None
+        film_layers = self.model._get_film_layers_in_order()
+
+        if len(film_layers) != len(layer_configs):
+            raise RuntimeError(
+                f"FiLM layer/config mismatch: {len(film_layers)} layers vs "
+                f"{len(layer_configs)} configs."
+            )
+
+        for (layer_id, n_channels), film_layer in zip(layer_configs, film_layers):
+            gamma, beta, film_context = film_layer.compute_modulation(organ_id)
+            modulations[layer_id] = (gamma, beta)
+
+            if not return_compact:
+                continue
+
+            gamma_flat = gamma.squeeze(-1).squeeze(-1)
+            beta_flat = beta.squeeze(-1).squeeze(-1)
+
+            if n_channels != self.model.max_channels:
+                gamma_compact = F.interpolate(
+                    gamma_flat.unsqueeze(1),
+                    size=self.model.max_channels,
+                    mode="linear",
+                    align_corners=False,
+                ).squeeze(1)
+                beta_compact = F.interpolate(
+                    beta_flat.unsqueeze(1),
+                    size=self.model.max_channels,
+                    mode="linear",
+                    align_corners=False,
+                ).squeeze(1)
+            else:
+                gamma_compact = gamma_flat
+                beta_compact = beta_flat
+
+            compact_modulations[layer_id] = (
+                gamma_compact,
+                beta_compact,
+                film_context,
+            )
+
+        projected_shapes = torch.empty(0, device=organ_id.device)
+        if return_compact:
+            return modulations, projected_shapes, compact_modulations
+        return modulations, projected_shapes
+
+
+class DownConvBlockFiLM(nn.Module):
+    """
+    Conv → FiLM → Conv → FiLM → Pool.
+    Except for FiLM, the API and behaviour remain the same.
+    """
+
+    def __init__(
+        self,
+        in_channels: list[int],
+        out_channels: list[int],
+        n_organs: int,
+        conv_kwargs={"kernel_size": 3, "stride": 1, "padding": 1},
+        emb_dim: int = 64,
+    ):
+        super().__init__()
+        assert len(in_channels) == len(
+            out_channels
+        ), f"in_channels length is {len(in_channels)} while out_channels is {len(out_channels)}"
+
+        self.conv_blocks = nn.ModuleList(
+            [
+                ConvBlock(in_ch, out_ch, conv_kwargs)
+                for in_ch, out_ch in zip(in_channels, out_channels)
+            ]
+        )
+
+        self.film_blocks = nn.ModuleList(
+            [
+                FiLM2d(
+                    n_organs=n_organs,
+                    in_channels=out_ch,
+                    emb_dim=emb_dim,
+                )
+                for out_ch in out_channels
+            ]
+        )
+
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        gammas: list[torch.Tensor] | None = None,
+        betas: list[torch.Tensor] | None = None,
+    ):
+        """
+        organ_id : (B,) torch.long   (0 = breast, 1 = thyroid, …)
+        """
+        if (gammas is None) != (betas is None):
+            raise ValueError("gammas and betas must be provided together.")
+
+        for idx, (conv, film) in enumerate(zip(self.conv_blocks, self.film_blocks)):
+            x = conv(x)  # usual conv-norm-ReLU
+            if gammas is None:
+                if organ_id is None:
+                    raise ValueError("organ_id is required when FiLM modulations are not precomputed.")
+                x = film(x, organ_id)  # FiLM modulation
+            else:
+                x = gammas[idx] * x + betas[idx]
+        return self.pool(x), x  # (downsampled, skip-connection)
+
+
+class UpConvBlockFiLM(nn.Module):
+    """
+    Up-sampling block with FiLM conditioning.
+
+    Conv → FiLM → Conv → FiLM → (optional) ConvTranspose2d
+    """
+
+    def __init__(
+        self,
+        in_channels: list[int],
+        out_channels: list[int],
+        n_organs: int,
+        up_conv: bool = True,
+        conv_kwargs: dict = {"kernel_size": 3, "stride": 1, "padding": 1},
+        upconv_kwargs: dict = {"kernel_size": 2, "stride": 2},
+        emb_dim: int = 64,
+    ):
+        super().__init__()
+        assert len(in_channels) == len(
+            out_channels
+        ), f"in_channels length is {len(in_channels)} while out_channels is {len(out_channels)}"
+
+        self.conv_blocks = nn.ModuleList(
+            [
+                ConvBlock(in_ch, out_ch, conv_kwargs)
+                for in_ch, out_ch in zip(in_channels, out_channels)
+            ]
+        )
+
+        self.film_blocks = nn.ModuleList(
+            [
+                FiLM2d(
+                    n_organs=n_organs,
+                    in_channels=out_ch,
+                    emb_dim=emb_dim,
+                )
+                for out_ch in out_channels
+            ]
+        )
+
+        self.up_conv = up_conv
+        if self.up_conv:
+            self.up_conv_op = nn.ConvTranspose2d(
+                out_channels[-1], out_channels[-1], **upconv_kwargs
+            )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        gammas: list[torch.Tensor] | None = None,
+        betas: list[torch.Tensor] | None = None,
+    ):
+        """
+        x : (B, C, H, W)
+        organ_id : (B,) long tensor - 0=breast, 1=thyroid, …
+        """
+        if (gammas is None) != (betas is None):
+            raise ValueError("gammas and betas must be provided together.")
+
+        for idx, (conv, film) in enumerate(zip(self.conv_blocks, self.film_blocks)):
+            x = conv(x)
+            if gammas is None:
+                if organ_id is None:
+                    raise ValueError("organ_id is required when FiLM modulations are not precomputed.")
+                x = film(x, organ_id)
+            else:
+                x = gammas[idx] * x + betas[idx]
+
+        if self.up_conv:
+            x = self.up_conv_op(x)
+
+        return x
+
+
+class UNet2DFiLM(BaseUnet):
+    def __init__(
+        self,
+        in_channels: int,
+        num_classes: int,
+        n_organs: int,
+        size: int = 32,
+        depth: int = 3,
+        *,
+        film_start: int = 0,
+        use_film: bool = True,
+        film_embed: int = 64,
+        distill: bool = False,
+        distill_unet: bool = False,
+        medsam_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/medsam_unfreezed/model.safetensors",
+        unet_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/unet5_attn/model.safetensors",
+        unet_teacher_kwargs: dict | None = None,
+    ):
+        """
+        UNet with symmetric FiLM conditioning in encoder and decoder.
+        """
+        super().__init__(
+            in_channels=in_channels,
+            num_classes=num_classes,
+            n_organs=n_organs,
+            size=size,
+            depth=depth,
+            film_start=film_start,
+            use_film=use_film,
+            film_embed=film_embed,
+            distill=distill,
+            distill_unet=distill_unet,
+            medsam_teacher_ckpt=medsam_teacher_ckpt,
+            unet_teacher_ckpt=unet_teacher_ckpt,
+            unet_teacher_kwargs=unet_teacher_kwargs,
+        )
+
+    def _build_model(
+        self,
+        *,
+        film_start: int = 0,
+        use_film: bool = True,
+        film_embed: int = 64,
+        distill: bool = False,
+        distill_unet: bool = False,
+        medsam_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/medsam_unfreezed/model.safetensors",
+        unet_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/unet5_attn/model.safetensors",
+        unet_teacher_kwargs: dict | None = None,
+        **kwargs,
+    ):
+        if kwargs:
+            unknown = ", ".join(sorted(kwargs.keys()))
+            raise TypeError(f"Unexpected UNet2DFiLM kwargs: {unknown}")
+
+        self.film_start = max(0, int(film_start))
+        self.use_film = bool(use_film)
+        self.use_attn = self.use_film
+        self.attn_start = self.film_start
+        self.film_embed = int(film_embed)
+        self.max_channels = self.size * (2 ** (self.depth + 1))
+        self.criterion = DiceBCELoss()
+        self.shared_attn = SharedFiLMModulator(self) if self.use_film else None
+
+        # ---------------- Encoder ----------------
+        self.encoder = nn.ModuleDict()
+
+        if self.use_film and 0 >= self.film_start:
+            self.encoder["0"] = DownConvBlockFiLM(
+                [self.in_channels, self.size],
+                [self.size, self.size * 2],
+                n_organs=self.n_organs,
+                emb_dim=self.film_embed,
+            )
+        else:
+            self.encoder["0"] = DownConvBlock(
+                [self.in_channels, self.size], [self.size, self.size * 2]
+            )
+
+        for i in range(1, self.depth):
+            in_ch = [self.size * (2**i), self.size * (2**i)]
+            out_ch = [self.size * (2**i), self.size * (2 ** (i + 1))]
+            key = str(i)
+
+            if self.use_film and i >= self.film_start:
+                self.encoder[key] = DownConvBlockFiLM(
+                    in_ch,
+                    out_ch,
+                    n_organs=self.n_organs,
+                    emb_dim=self.film_embed,
+                )
+            else:
+                self.encoder[key] = DownConvBlock(in_ch, out_ch)
+
+        # ---------------- Bottleneck ----------------
+        if self.use_film:
+            self.bottleneck = UpConvBlockFiLM(
+                [self.size * (2**self.depth), self.size * (2**self.depth)],
+                [self.size * (2**self.depth), self.size * (2 ** (self.depth + 1))],
+                n_organs=self.n_organs,
+                emb_dim=self.film_embed,
+            )
+        else:
+            self.bottleneck = UpConvBlock(
+                [self.size * (2**self.depth), self.size * (2**self.depth)],
+                [self.size * (2**self.depth), self.size * (2 ** (self.depth + 1))],
+            )
+
+        # ---------------- Decoder ----------------
+        self.decoder = nn.ModuleDict()
+
+        for i in range(self.depth, 1, -1):
+            use_film_at_level = self.use_film and (i - 1) >= self.film_start
+
+            if use_film_at_level:
+                self.decoder[str(i - 1)] = UpConvBlockFiLM(
+                    [
+                        self.size * (2 ** (i + 1)) + self.size * (2**i),
+                        self.size * (2**i),
+                    ],
+                    [self.size * (2**i), self.size * (2**i)],
+                    n_organs=self.n_organs,
+                    emb_dim=self.film_embed,
+                )
+            else:
+                self.decoder[str(i - 1)] = UpConvBlock(
+                    [
+                        self.size * (2 ** (i + 1)) + self.size * (2**i),
+                        self.size * (2**i),
+                    ],
+                    [self.size * (2**i), self.size * (2**i)],
+                )
+
+        if self.use_film and 0 >= self.film_start:
+            self.decoder["0"] = UpConvBlockFiLM(
+                [self.size * 4 + self.size * 2, self.size * 2],
+                [self.size * 2, self.size * 2],
+                n_organs=self.n_organs,
+                up_conv=False,
+                emb_dim=self.film_embed,
+            )
+        else:
+            self.decoder["0"] = UpConvBlock(
+                [self.size * 4 + self.size * 2, self.size * 2],
+                [self.size * 2, self.size * 2],
+                up_conv=False,
+            )
+
+        self.out_layer = ConvBlock(
+            self.size * 2,
+            self.out_channels,
+            conv_kwargs={"kernel_size": 1, "stride": 1, "padding": 0},
+        )
+
+        self._init_distillation(
+            distill=distill,
+            distill_unet=distill_unet,
+            medsam_teacher_ckpt=medsam_teacher_ckpt,
+            unet_teacher_ckpt=unet_teacher_ckpt,
+            unet_teacher_kwargs=unet_teacher_kwargs,
+        )
+
+    def _build_layer_configs(self):
+        """Match UNet2DAttn layer ordering for notebook-side FiLM collection."""
+        configs = []
+        layer_id = 0
+
+        for i in range(self.depth):
+            if self.use_film and i >= self.film_start:
+                configs.append((layer_id, self.size * (2**i)))
+                layer_id += 1
+                configs.append((layer_id, self.size * (2 ** (i + 1))))
+                layer_id += 1
+
+        if self.use_film:
+            configs.append((layer_id, self.size * (2**self.depth)))
+            layer_id += 1
+            configs.append((layer_id, self.size * (2 ** (self.depth + 1))))
+            layer_id += 1
+
+        for i in range(self.depth - 1, -1, -1):
+            if self.use_film and i >= self.film_start:
+                configs.append((layer_id, self.size * (2 ** (i + 1))))
+                layer_id += 1
+                configs.append((layer_id, self.size * (2 ** (i + 1))))
+                layer_id += 1
+
+        return configs
+
+    def _get_film_layers_in_order(self):
+        """Return FiLM modules in the same order as `_build_layer_configs`."""
+        film_layers = []
+
+        for i in range(self.depth):
+            layer = self.encoder[str(i)]
+            if isinstance(layer, DownConvBlockFiLM):
+                film_layers.extend(layer.film_blocks)
+
+        if isinstance(self.bottleneck, UpConvBlockFiLM):
+            film_layers.extend(self.bottleneck.film_blocks)
+
+        for i in range(self.depth - 1, -1, -1):
+            layer = self.decoder[str(i)]
+            if isinstance(layer, UpConvBlockFiLM):
+                film_layers.extend(layer.film_blocks)
+
+        return film_layers
+
+    def _prepare_forward(
+        self,
+        *,
+        pixel_values: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        **kwargs,
+    ) -> dict:
+        forward_ctx = {
+            "mod_list": None,
+            "mod_idx": 0,
+        }
+        if self.use_film:
+            if organ_id is None:
+                raise ValueError("organ_id is required when FiLM is enabled.")
+            layer_configs = self._build_layer_configs()
+            film_layers = self._get_film_layers_in_order()
+            if len(film_layers) != len(layer_configs):
+                raise RuntimeError(
+                    f"FiLM layer/config mismatch: {len(film_layers)} layers vs "
+                    f"{len(layer_configs)} configs."
+                )
+            modulations = {
+                layer_id: film_layer.compute_modulation(organ_id)[:2]
+                for (layer_id, _), film_layer in zip(layer_configs, film_layers)
+            }
+            forward_ctx["mod_list"] = [modulations[i] for i in range(len(layer_configs))]
+        return forward_ctx
+
+    def _next_modulation(self, forward_ctx: dict):
+        mod_idx = forward_ctx["mod_idx"]
+        mod_list = forward_ctx["mod_list"]
+        gammas = [mod_list[mod_idx][0], mod_list[mod_idx + 1][0]]
+        betas = [mod_list[mod_idx][1], mod_list[mod_idx + 1][1]]
+        forward_ctx["mod_idx"] = mod_idx + 2
+        return gammas, betas
+
+    def _encode(
+        self,
+        layer,
+        x,
+        organ_id=None,
+        forward_ctx=None,
+    ):
+        if isinstance(layer, DownConvBlockFiLM):
+            if forward_ctx is not None and forward_ctx.get("mod_list") is not None:
+                gammas, betas = self._next_modulation(forward_ctx)
+                return layer(x, gammas=gammas, betas=betas)
+            return layer(x, organ_id=organ_id)
+        else:
+            return layer(x)
+
+    def _decode(
+        self,
+        layer,
+        x,
+        organ_id=None,
+        forward_ctx=None,
+    ):
+        if isinstance(layer, UpConvBlockFiLM):
+            if forward_ctx is not None and forward_ctx.get("mod_list") is not None:
+                gammas, betas = self._next_modulation(forward_ctx)
+                return layer(x, gammas=gammas, betas=betas)
+            return layer(x, organ_id=organ_id)
+        else:
+            return layer(x)
+
+    def _bottleneck(
+        self,
+        x,
+        organ_id=None,
+        forward_ctx=None,
+    ):
+        if isinstance(self.bottleneck, UpConvBlockFiLM):
+            if forward_ctx is not None and forward_ctx.get("mod_list") is not None:
+                gammas, betas = self._next_modulation(forward_ctx)
+                return self.bottleneck(x, gammas=gammas, betas=betas)
+            return self.bottleneck(x, organ_id=organ_id)
+        else:
+            return self.bottleneck(x)
+
+    def __str__(self):
+        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
+        params = sum([np.prod(p.size()) for p in model_parameters])
+        film_status = "enabled" if self.use_film else "disabled"
+        film_range = f"from level {self.film_start}" if self.use_film else "N/A"
+        return (
+            super().__str__() + f"\nTrainable parameters: {params}"
+            f"\nFiLM: {film_status} ({film_range})"
+        )
+
+
+class DiceBCELoss(nn.Module):
+    def __init__(self, dice_weight: float = 1.0, bce_weight: float = 1.0):
+        super().__init__()
+        self.dice_weight = dice_weight
+        self.bce_weight = bce_weight
+        self.eps = 1e-6
+
+    def forward(self, logits: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
+
+        gt = gt.float()
+
+        bce = F.binary_cross_entropy_with_logits(
+            logits.squeeze(), gt.squeeze(), reduction="mean"
+        )
+
+        # Soft Dice loss
+        probs = torch.sigmoid(logits)
+        dims = tuple(range(2, probs.dim()))  # (H, W)  or (D,H,W)
+
+        # per‑class Dice, per‑sample
+        inter = (probs * gt).sum(dims) * 2
+        union = probs.sum(dims) + gt.sum(dims)
+        dice = 1 - (inter + self.eps) / (union + self.eps)  # [B, C]
+
+        dice = dice.mean()
+
+        loss = self.dice_weight * dice + self.bce_weight * bce
+        return loss
+
+
+class MedSAM(nn.Module):
+    def __init__(
+        self,
+        image_encoder,
+        mask_decoder,
+        prompt_encoder,
+        freeze_image_encoder=True,
+        predict_bboxes=False,
+        image_encoder_lora_rank=0,
+    ):
+        super().__init__()
+        self.image_encoder = image_encoder
+        self.mask_decoder = mask_decoder
+        self.prompt_encoder = prompt_encoder
+        self.criterion = DiceBCELoss()
+        self.freeze_image_encoder = freeze_image_encoder
+        self.predict_bboxes = predict_bboxes
+
+        if self.freeze_image_encoder:
+            for param in self.image_encoder.parameters():
+                param.requires_grad = False
+
+        self.image_encoder_lora_rank = int(image_encoder_lora_rank)
+        self.image_encoder_lora_layers = 0
+        if self.image_encoder_lora_rank > 0:
+            self.image_encoder, self.image_encoder_lora_layers = apply_lora_to_qv_projections(
+                self.image_encoder,
+                rank=self.image_encoder_lora_rank,
+            )
+
+        # Classification head
+        self.multi_cls = nn.Sequential(nn.Linear(256 * 64 * 64, 10))
+
+        # Bounding box regression head
+        self.bbox_regr = nn.Sequential(nn.Linear(256 * 64 * 64, 4))
+
+        # Learnable prompt embeddings (no input required)
+        self.learned_sparse_embeddings = nn.Parameter(
+            torch.randn(1, 2, 256)  # (batch, num_tokens, embed_dim)
+        )
+        self.learned_dense_embeddings = nn.Parameter(
+            torch.randn(1, 256, 64, 64)  # (batch, embed_dim, H, W)
+        )
+
+    def forward(
+        self,
+        pixel_values,
+        organ_id=None,
+        labels=None,
+        masks=None,
+        bbox_coords=None,
+        organ_id_metric=None,
+        pixel_values_medsam=None,
+        **kwargs,  # ignored, for peft compatibility
+    ):
+        batch_size = pixel_values.shape[0]
+
+        # Get image embeddings
+        image_embedding = self.image_encoder(pixel_values_medsam)  # (B, 256, 64, 64)
+
+        # Classification output
+        emb_flattened = torch.flatten(image_embedding, 1)
+        multi_cls_out = self.multi_cls(emb_flattened)
+
+        # Bounding box output (if enabled)
+        if self.predict_bboxes:
+            bbox_out = self.bbox_regr(emb_flattened)
+        else:
+            bbox_out = None
+
+        # Expand learned embeddings to batch size
+        sparse_embeddings = self.learned_sparse_embeddings
+        dense_embeddings = self.learned_dense_embeddings
+
+        # Get positional encoding
+        image_pe = self.prompt_encoder.get_dense_pe()  # (1, 256, 64, 64)
+
+        # Decode mask
+        low_res_masks, iou_predictions = self.mask_decoder(
+            image_embeddings=image_embedding,  # (B, 256, 64, 64)
+            image_pe=image_pe,  # (1, 256, 64, 64)
+            sparse_prompt_embeddings=sparse_embeddings,  # (B, 2, 256)
+            dense_prompt_embeddings=dense_embeddings,  # (B, 256, 64, 64)
+            multimask_output=False,
+        )  # (B, 1, 256, 256)
+        if masks is not None:
+            loss = self.criterion(
+                low_res_masks.squeeze(1), v2.functional.resize(masks, (256, 256))
+            )
+        else:
+            loss = 0.0
+
+        return {
+            "loss": loss,
+            "logits": low_res_masks.squeeze(1),
+            "labels": masks,
+            "organ_id": organ_id,
+            "organ_id_metric": organ_id_metric,
+        }
+
+
+class DistillationLoss(nn.Module):
+    """
+    Improved distillation loss with multiple components:
+    - Cosine similarity loss (direction alignment)
+    - MSE loss (magnitude alignment)
+    - Optional L1 loss (sparsity)
+    """
+
+    def __init__(self, temperature=3.0, alpha=0.5, use_l1=False):
+        super().__init__()
+        self.temperature = temperature
+        self.alpha = alpha  # Weight between cosine and MSE
+        self.use_l1 = use_l1
+
+    def forward(self, student_logits, teacher_logits, tau=0.7):
+        """
+        Args:
+            student_logits: (B, D) - student features
+            teacher_logits: (B, D) - teacher features
+        """
+        # # Normalize features for stable training
+        # student_norm = F.normalize(student_logits, p=2, dim=1)
+        # teacher_norm = F.normalize(teacher_logits, p=2, dim=1)
+
+        # # Cosine similarity loss (encourages directional alignment)
+        # cosine_loss = (
+        #     1 - F.cosine_similarity(student_logits, teacher_logits, dim=1).mean()
+        # )
+
+        # # MSE loss on normalized features (encourages magnitude alignment)
+        # mse_loss = F.mse_loss(student_norm, teacher_norm)
+
+        # # Combined loss
+        # loss = self.alpha * cosine_loss + (1 - self.alpha) * mse_loss
+
+        # # Optional L1 for sparsity
+        # if self.use_l1:
+        #     l1_loss = F.l1_loss(student_logits, teacher_logits)
+        #     loss = loss + 0.1 * l1_loss
+
+        B = student_logits.size(0)
+
+        s = student_logits.flatten(start_dim=1)
+        t = teacher_logits.flatten(start_dim=1).detach()
+
+        s = F.normalize(s, dim=-1)
+        t = F.normalize(t, dim=-1)
+
+        d2 = (s[:, None, :] - t[None, :, :]).pow(2).sum(dim=-1)
+        logits = -d2 / tau
+        labels = torch.arange(B, device=s.device)
+
+        loss = F.cross_entropy(logits, labels)
+        return {
+            "loss": loss,
+        }
+
+
+class MedSAMPrompt(nn.Module):
+    def __init__(
+        self,
+        image_encoder,
+        mask_decoder,
+        prompt_encoder,
+        freeze_image_encoder=True,
+        predict_bboxes=False,
+        n_organs=1,
+        image_encoder_lora_rank=0,
+    ):
+        super().__init__()
+        self.image_encoder = image_encoder
+        self.mask_decoder = mask_decoder
+        self.prompt_encoder = prompt_encoder
+        self.criterion = DiceBCELoss()
+        self.freeze_image_encoder = freeze_image_encoder
+        self.predict_bboxes = predict_bboxes
+
+        if self.freeze_image_encoder:
+            for param in self.image_encoder.parameters():
+                param.requires_grad = False
+
+        self.image_encoder_lora_rank = int(image_encoder_lora_rank)
+        self.image_encoder_lora_layers = 0
+        if self.image_encoder_lora_rank > 0:
+            self.image_encoder, self.image_encoder_lora_layers = apply_lora_to_qv_projections(
+                self.image_encoder,
+                rank=self.image_encoder_lora_rank,
+            )
+
+        # Classification head
+        self.multi_cls = nn.Sequential(nn.Linear(256 * 64 * 64, 10))
+
+        # Bounding box regression head
+        self.bbox_regr = nn.Sequential(nn.Linear(256 * 64 * 64, 4))
+
+        # Learnable prompt embeddings (no input required)
+        self.sparse_embeddings = nn.Parameter(torch.randn(1, 2, 256))
+        self.dense_embeddings = nn.Embedding(n_organs + 1, 1 * 256 * 64 * 64)
+
+        # self.sparse_embeddings = nn.ParameterDict(
+        #     {
+        #         str(id_): nn.Parameter(torch.randn(1, 2, 256))
+        #         for id_ in set(organ_to_class_dict.values())
+        #     }
+        # )
+        # self.dense_embeddings = nn.ParameterDict(
+        #     {
+        #         str(id_): nn.Parameter(torch.randn(1, 256, 64, 64))
+        #         for id_ in set(organ_to_class_dict.values())
+        #     }
+        # )
+
+    def forward(
+        self,
+        pixel_values,
+        organ_id=None,
+        labels=None,
+        masks=None,
+        bbox_coords=None,
+        organ_id_metric=None,
+        pixel_values_medsam=None,
+        **kwargs,  # ignored, for peft compatibility
+    ):
+        batch_size = pixel_values.shape[0]
+        B = pixel_values.shape[0]
+
+        # Get image embeddings
+        image_embedding = self.image_encoder(pixel_values_medsam)  # (B, 256, 64, 64)
+
+        # Classification output
+        emb_flattened = torch.flatten(image_embedding, 1)
+        multi_cls_out = self.multi_cls(emb_flattened)
+
+        # Bounding box output (if enabled)
+        if self.predict_bboxes:
+            bbox_out = self.bbox_regr(emb_flattened)
+        else:
+            bbox_out = None
+
+        if organ_id is None:
+            raise ValueError("organ_id must be provided for selecting embeddings.")
+
+        sparse_emb_list, dense_emb_list = [], []
+
+        mask = organ_id >= 0  # [B]
+        idx = torch.where(mask, organ_id, self.dense_embeddings.weight.shape[0] - 1)
+        dense_embeddings = self.dense_embeddings(idx).view(B, 256, 64, 64)
+
+        # Get positional encoding
+        image_pe = self.prompt_encoder.get_dense_pe()  # (1, 256, 64, 64)
+
+        # Decode mask
+        low_res_masks, iou_predictions = self.mask_decoder(
+            image_embeddings=image_embedding,  # (B, 256, 64, 64)
+            image_pe=image_pe,  # (1, 256, 64, 64)
+            sparse_prompt_embeddings=self.sparse_embeddings,  # (B, 2, 256)
+            dense_prompt_embeddings=dense_embeddings,  # (B, 256, 64, 64)
+            multimask_output=False,
+        )  # (B, 1, 256, 256)
+        if masks is not None:
+            loss = self.criterion(
+                low_res_masks.squeeze(1), v2.functional.resize(masks, (256, 256))
+            )
+        else:
+            loss = 0.0
+
+        return {
+            "loss": loss,
+            "logits": low_res_masks.squeeze(1),
+            "labels": masks,
+            "organ_id": organ_id,
+            "organ_id_metric": organ_id_metric,
+        }
+
+
+
+# config = {"in_channels": 3,"num_classes": 1,"n_organs": 8,"size": 32,"depth": 5,"film_start": 0,"use_film": 1}

unet_attn.py

@@ -0,0 +1,971 @@
+import torch, wandb
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from .segm_net import (
+    ConvBlock,
+    DownConvBlock,
+    UpConvBlock,
+    DiceBCELoss,
+)
+import ptwt
+import numpy as np
+from .unet_base import BaseUnet
+
+
+def canonicalize_mask_batch(gt_masks: torch.Tensor, canon_res: int = 32, pad: int = 4):
+    """
+    Args:
+        gt_masks: (B, H, W) binary float tensor (0/1)
+        canon_res: output resolution (int)
+        pad: pixels of padding around bbox (optional)
+    Returns:
+        canon_masks: (B, canon_res, canon_res) float tensor in [0,1]
+    """
+    B, H, W = gt_masks.shape
+    device = gt_masks.device
+    canon_masks = torch.zeros(
+        (B, canon_res, canon_res), device=device, dtype=gt_masks.dtype
+    )
+
+    for i in range(B):
+        m = gt_masks[i]
+        nz = torch.nonzero(m, as_tuple=False)
+        if nz.numel() == 0:
+            # empty mask -> keep zeros
+            continue
+        y_min = int(nz[:, 0].min().clamp(0, H - 1).item())
+        y_max = int(nz[:, 0].max().clamp(0, H - 1).item())
+        x_min = int(nz[:, 1].min().clamp(0, W - 1).item())
+        x_max = int(nz[:, 1].max().clamp(0, W - 1).item())
+
+        # pad bbox
+        y0 = max(0, y_min - pad)
+        y1 = min(H, y_max + pad + 1)
+        x0 = max(0, x_min - pad)
+        x1 = min(W, x_max + pad + 1)
+
+        crop = m[y0:y1, x0:x1].unsqueeze(0).unsqueeze(0)  # (1,1,hc,wc)
+        # Resize to canonical resolution
+        crop_resized = F.interpolate(
+            crop, size=(canon_res, canon_res), mode="bilinear", align_corners=False
+        )
+        canon_masks[i] = crop_resized[0, 0]
+
+    return canon_masks  # (B, canon_res, canon_res)
+
+
+def canonicalize_mask_batch_normalized(
+    gt_masks: torch.Tensor,
+    canon_res: int = 32,
+    target_scale: float = 0.7,  # Target mask to fill 70% of canonical space
+):
+    """
+    Canonical masks with consistent scale normalization.
+    """
+    B, H, W = gt_masks.shape
+    device = gt_masks.device
+    canon_masks = torch.zeros(
+        (B, canon_res, canon_res), device=device, dtype=gt_masks.dtype
+    )
+
+    for i in range(B):
+        m = gt_masks[i]
+        nz = torch.nonzero(m, as_tuple=False)
+        if nz.numel() == 0:
+            continue
+
+        y_coords = nz[:, 0].float()
+        x_coords = nz[:, 1].float()
+
+        # Center of mass
+        com_y = y_coords.mean()
+        com_x = x_coords.mean()
+
+        # Bbox dimensions
+        y_min, y_max = y_coords.min(), y_coords.max()
+        x_min, x_max = x_coords.min(), x_coords.max()
+        bbox_h = (y_max - y_min + 1).item()
+        bbox_w = (x_max - x_min + 1).item()
+
+        # Compute crop size to achieve target scale
+        max_dim = max(bbox_h, bbox_w)
+        crop_size = int(
+            max_dim / target_scale
+        )  # Scale up to make mask fill target_scale
+
+        # Create square crop centered on COM
+        half_crop = crop_size // 2
+        y0 = int(com_y.item()) - half_crop
+        y1 = y0 + crop_size
+        x0 = int(com_x.item()) - half_crop
+        x1 = x0 + crop_size
+
+        # Clamp and pad
+        y0_clamped = max(0, y0)
+        y1_clamped = min(H, y1)
+        x0_clamped = max(0, x0)
+        x1_clamped = min(W, x1)
+
+        crop = m[y0_clamped:y1_clamped, x0_clamped:x1_clamped]
+
+        # Padding
+        pad_top = y0_clamped - y0
+        pad_bottom = crop_size - (y1_clamped - y0_clamped) - pad_top
+        pad_left = x0_clamped - x0
+        pad_right = crop_size - (x1_clamped - x0_clamped) - pad_left
+
+        crop = F.pad(
+            crop.unsqueeze(0).unsqueeze(0),
+            (pad_left, pad_right, pad_top, pad_bottom),
+            mode="constant",
+            value=0,
+        )
+
+        # Resize
+        crop_resized = F.interpolate(
+            crop, size=(canon_res, canon_res), mode="bilinear", align_corners=False
+        )
+        canon_masks[i] = crop_resized[0, 0]
+
+    return canon_masks
+
+
+class PatchEmbed(nn.Module):
+    """Convert image to patches and embed them with positional encoding"""
+
+    def __init__(self, img_size=256, patch_size=16, in_channels=3, embed_dim=256):
+        super().__init__()
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.n_patches = (img_size // patch_size) ** 2
+
+        self.proj = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+        # Positional encoding - CRITICAL for spatial understanding
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.n_patches, embed_dim))
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+
+    def forward(self, x):
+        """
+        x: (B, C, H, W)
+        returns: (B, N, D) where N = num_patches
+        """
+        x = self.proj(x)  # (B, embed_dim, H/P, W/P)
+        x = rearrange(x, "b c h w -> b (h w) c")  # (B, N, D)
+        x = x + self.pos_embed  # Add positional encoding
+        return x
+
+
+class DWTPatchEmbed(nn.Module):
+    """
+    Apply DWT and create patch embeddings from subbands.
+    Interface matches PatchEmbed to minimize changes to SharedAttnModulator.
+
+    DWT decomposes (B, C, H, W) into:
+    - LL: Low-freq approximation (B, C, H/2, W/2)
+    - LH, HL, HH: High-freq details (B, C, H/2, W/2 each)
+    """
+
+    def __init__(
+        self,
+        img_size=256,
+        patch_size=8,
+        in_channels=3,
+        embed_dim=256,
+        wavelet="haar",
+        mode="zero",
+        bands: list[str] | None = None,
+    ):
+        super().__init__()
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.wavelet = wavelet
+        self.mode = mode
+        self.available_bands = ["LL", "LH", "HL", "HH"]
+        if not bands:
+            self.bands = list(self.available_bands)
+        else:
+            normalized = []
+            seen = set()
+            for band in bands:
+                band = band.upper()
+                if band in seen:
+                    continue
+                if band not in self.available_bands:
+                    raise ValueError(
+                        f"Invalid DWT band '{band}'. Valid bands: {self.available_bands}"
+                    )
+                normalized.append(band)
+                seen.add(band)
+            self.bands = normalized
+
+        # After 1-level DWT, each subband is img_size/2
+        dwt_size = img_size // 2  # 128 for img_size=256
+
+        # Number of patches per subband
+        patches_per_dim = dwt_size // patch_size  # 128/8 = 16
+        n_patches_per_band = patches_per_dim**2  # 256
+
+        # Total patches: selected subbands × n_patches_per_band
+        self.n_patches = len(self.bands) * n_patches_per_band
+        self.n_patches_per_band = n_patches_per_band
+
+        # Separate projections for each subband
+        self.proj_LL = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+        self.proj_LH = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+        self.proj_HL = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+        self.proj_HH = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+        # Shared positional encoding for all patches
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.n_patches, embed_dim))
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+
+        # Learnable subband type embeddings
+        self.subband_type_embed = nn.Parameter(torch.zeros(4, embed_dim))
+        nn.init.trunc_normal_(self.subband_type_embed, std=0.02)
+
+    def forward(self, x):
+        """
+        Args:
+            x: (B, C, H, W) - input image
+        Returns:
+            patches: (B, N, D) where N = len(bands) * (H/2/patch_size)^2
+        """
+        B, C, H, W = x.shape
+
+        # Apply DWT using ptwt - correct API
+        # wavedec2 returns coefficients in format: [LL, (LH, HL, HH)]
+        coeffs = ptwt.wavedec2(x, wavelet=self.wavelet, mode=self.mode, level=1)
+
+        # Extract subbands
+        # coeffs[0] is LL (approximation)
+        # coeffs[1] is tuple of (LH, HL, HH) detail coefficients
+        LL = coeffs[0]  # (B, C, H/2, W/2)
+        LH, HL, HH = coeffs[1]  # Each (B, C, H/2, W/2)
+
+        subband_tensors = {
+            "LL": (LL, self.proj_LL, 0),
+            "LH": (LH, self.proj_LH, 1),
+            "HL": (HL, self.proj_HL, 2),
+            "HH": (HH, self.proj_HH, 3),
+        }
+        patches_list = []
+        for band in self.bands:
+            band_tensor, proj, band_idx = subband_tensors[band]
+            band_patches = proj(band_tensor)  # (B, embed_dim, H/2/P, W/2/P)
+            band_patches = rearrange(band_patches, "b d h w -> b (h w) d")
+            band_patches = band_patches + self.subband_type_embed[band_idx]
+            patches_list.append(band_patches)
+
+        # Concatenate selected subbands
+        patches = torch.cat(patches_list, dim=1)
+
+        # Add positional encoding
+        patches = patches + self.pos_embed
+
+        return patches
+
+
+class FiLMLayer(nn.Module):
+    """Simple FiLM layer that applies gamma * x + beta modulation"""
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor):
+        """
+        Args:
+            x: (B, C, H, W) - input features
+            gamma: (B, C, 1, 1) - scaling factors
+            beta: (B, C, 1, 1) - bias terms
+        Returns:
+            (B, C, H, W) - modulated features
+        """
+        return gamma * x + beta
+
+
+class SharedAttnModulator(nn.Module):
+    """
+    Shared attention modulator that computes gamma/beta for all layers at once.
+    Uses layer_id embeddings to differentiate between layers.
+    """
+
+    def __init__(
+        self,
+        n_organs: int,
+        n_layers: int,
+        max_channels: int,
+        img_size: int = 256,
+        patch_size: int = 16,
+        img_channels: int = 3,
+        emb_dim: int = 256,
+        n_heads: int = 8,
+        dropout: float = 0.1,
+        use_dwt: bool = True,
+        wavelet: str = "haar",
+        dwt_bands: list[str] | None = None,
+        use_shape: bool = False,
+        shape_res=32,
+    ) -> None:
+        super().__init__()
+        self.emb_dim = emb_dim
+        self.n_layers = n_layers
+        self.max_channels = max_channels
+        self.use_shape = use_shape
+        self.shape_res = shape_res
+
+        if use_dwt:
+            self.patch_embed = DWTPatchEmbed(
+                img_size=img_size,
+                patch_size=patch_size,
+                in_channels=img_channels,
+                embed_dim=emb_dim,
+                wavelet=wavelet,
+                bands=dwt_bands,
+            )
+        else:
+            self.patch_embed = PatchEmbed(
+                img_size=img_size,
+                patch_size=patch_size,
+                in_channels=img_channels,
+                embed_dim=emb_dim,
+            )
+
+        # Organ embedding
+        self.organ_embed = nn.Embedding(
+            n_organs + 1, emb_dim
+        )  # the +1 is for the unknown organ
+        nn.init.normal_(self.organ_embed.weight, mean=0, std=0.02)
+
+        # Layer embedding
+        self.layer_embed = nn.Embedding(n_layers, emb_dim)
+        nn.init.normal_(self.layer_embed.weight, mean=0, std=0.02)
+
+        if self.use_shape:
+            # Shape embedding
+            self.shape_embed = nn.Embedding(
+                n_organs + 1, emb_dim
+            )  # the +1 is for the unknown organ
+            nn.init.normal_(self.shape_embed.weight, mean=0, std=0.02)
+            self.shape_proj = nn.Linear(self.emb_dim, self.shape_res**2)
+
+        # Learnable influence token
+        self.influence_token = nn.Parameter(torch.zeros(1, 1, emb_dim))
+        nn.init.trunc_normal_(self.influence_token, std=0.02)
+
+        # Attention with dropout
+        self.attn = nn.MultiheadAttention(
+            embed_dim=emb_dim, num_heads=n_heads, dropout=dropout, batch_first=True
+        )
+
+        # Shared components for all layers
+        self.shared_norm = nn.LayerNorm(emb_dim)
+
+        # Per-layer linear projections
+        self.layer_linears = nn.ModuleList(
+            [nn.Linear(emb_dim, emb_dim) for _ in range(n_layers)]
+        )
+
+        # Shared components after layer-specific linear
+        self.shared_gelu = nn.GELU()
+        self.shared_dropout = nn.Dropout(dropout)
+        self.shared_output = nn.Linear(emb_dim, 2 * max_channels)
+
+        # Initialize to identity (gamma≈1, beta≈0)
+        nn.init.zeros_(self.shared_output.weight)
+        nn.init.constant_(self.shared_output.bias[:max_channels], 0)  # β
+        nn.init.constant_(self.shared_output.bias[max_channels:], 1)  # γ
+
+    def compute_all_modulations(
+        self,
+        original_img: torch.Tensor,
+        organ_id: torch.Tensor,
+        layer_configs: list[tuple[int, int]],
+        return_compact: bool = False,
+    ):
+        """
+        Compute gamma and beta for all layers at once.
+
+        Args:
+            original_img: (B, 3, H, W) - original input image
+            organ_id: (B,) - organ type IDs (if >= 0, use organ embedding)
+            layer_configs: list of (layer_id, n_channels) tuples
+
+        Returns:
+            dict: {layer_id: (gamma, beta)} where gamma/beta are (B, C, 1, 1)
+            If return_compact=True, also returns raw (gamma_max, beta_max) per layer
+            before channel pooling, each shaped (B, max_channels).
+
+        Performance note: all L layer queries are stacked into a single attention
+        call (B, L*n_q, D) so the (B, N, D) key/value projections are computed
+        only once instead of L times. For N=4096, D=768, L=22 this is ~20x faster.
+        """
+        B = original_img.shape[0]
+        L = len(layer_configs)
+        D = self.emb_dim
+        device = original_img.device
+
+        # Patch embedding — shared across all layers, computed once.
+        patches = self.patch_embed(original_img)  # (B, N, D)
+
+        # Organ query base — same for all layers, select by organ_id validity.
+        mask = organ_id >= 0  # (B,)
+        unk_idx = torch.tensor(
+            [self.organ_embed.weight.shape[0] - 1], device=device
+        ).expand(B)
+        q_org_base = torch.where(
+            mask[:, None],
+            self.organ_embed(organ_id.clamp(min=0)),
+            self.organ_embed(unk_idx),
+        )  # (B, D)
+
+        # All layer embeddings at once.
+        all_layer_ids = torch.tensor(
+            [lid for lid, _ in layer_configs], device=device
+        )  # (L,)
+        all_layer_embs = self.layer_embed(all_layer_ids)  # (L, D)
+
+        # Per-layer organ query: q_org_base + layer_emb  →  (B, L, D)
+        q_per_layer = q_org_base.unsqueeze(1) + all_layer_embs.unsqueeze(0)
+
+        # Influence token tiled to (B, L, D).
+        influence_all = self.influence_token.expand(B, L, -1)  # (B, L, D)
+
+        if self.use_shape:
+            unk_shape_idx = torch.tensor(
+                [self.shape_embed.weight.shape[0] - 1], device=device
+            ).expand(B)
+            q_shape = torch.where(
+                mask[:, None],
+                self.shape_embed(organ_id.clamp(min=0)),
+                self.shape_embed(unk_shape_idx),
+            )  # (B, D) — keep un-detached for shape_proj gradient
+
+            # projected_shape is identical for every layer (depends only on organ).
+            flattened_shape = self.shape_proj(q_shape)  # (B, shape_res²)
+            projected_shape = flattened_shape.reshape(B, self.shape_res, self.shape_res)
+            projected_shapes = torch.stack([projected_shape] * L)  # (L, B, sr, sr)
+
+            # Detach for use as attention query to isolate gradient path.
+            q_shape_attn = q_shape.detach().unsqueeze(1).expand(-1, L, -1)  # (B, L, D)
+
+            # Stack: [influence, organ+layer, shape]  →  (B, L, 3, D)
+            queries_all = torch.stack(
+                [influence_all, q_per_layer, q_shape_attn], dim=2
+            )
+            n_q = 3
+        else:
+            projected_shapes = torch.Tensor([])
+            # Stack: [influence, organ+layer]  →  (B, L, 2, D)
+            queries_all = torch.stack([influence_all, q_per_layer], dim=2)
+            n_q = 2
+
+        # Single attention call: K/V projections on (B, N, D) happen only once
+        # instead of L times — the dominant cost for large N.
+        queries_flat = queries_all.reshape(B, L * n_q, D)  # (B, L*n_q, D)
+        attn_out_flat, _ = self.attn(
+            query=queries_flat,  # (B, L*n_q, D)
+            key=patches,         # (B, N, D)
+            value=patches,
+        )  # → (B, L*n_q, D)
+
+        # Mean-pool each layer's n_q tokens  →  (B, L, D)
+        attn_out_all = attn_out_flat.reshape(B, L, n_q, D).mean(dim=2)
+
+        # Per-layer post-processing (cheap: norm + linear + gelu + dropout + output).
+        modulations = {}
+        compact_modulations = {} if return_compact else None
+        for li, (layer_id, n_channels) in enumerate(layer_configs):
+            attn_out = attn_out_all[:, li, :]  # (B, D)
+
+            x = self.shared_norm(attn_out)
+            x = self.layer_linears[layer_id](x)
+            x = self.shared_gelu(x)
+            x = self.shared_dropout(x)
+            gamma_beta = self.shared_output(x)  # (B, 2 * max_channels)
+
+            beta_max, gamma_max = gamma_beta.chunk(2, dim=-1)  # each (B, max_channels)
+
+            if return_compact:
+                compact_modulations[layer_id] = (gamma_max, beta_max, attn_out)
+
+            beta = F.adaptive_avg_pool1d(beta_max, n_channels).unsqueeze(-1).unsqueeze(-1)
+            gamma = F.adaptive_avg_pool1d(gamma_max, n_channels).unsqueeze(-1).unsqueeze(-1)
+
+            modulations[layer_id] = (gamma, beta)
+
+        if return_compact:
+            return modulations, projected_shapes.to(device), compact_modulations
+        return modulations, projected_shapes.to(device)
+
+
+class DownConvBlockFiLM(nn.Module):
+    """
+    Encoder block with FiLM modulation.
+    Conv → FiLM → Conv → FiLM → Pool
+    """
+
+    def __init__(
+        self,
+        in_channels: list[int],
+        out_channels: list[int],
+        conv_kwargs={"kernel_size": 3, "stride": 1, "padding": 1},
+    ):
+        super().__init__()
+        assert len(in_channels) == len(out_channels)
+
+        self.conv_blocks = nn.ModuleList(
+            [
+                ConvBlock(in_ch, out_ch, conv_kwargs)
+                for in_ch, out_ch in zip(in_channels, out_channels)
+            ]
+        )
+
+        self.film_layers = nn.ModuleList([FiLMLayer() for _ in out_channels])
+
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x: torch.Tensor, gammas: list, betas: list):
+        """
+        Args:
+            x: (B, C, H, W) - features from previous layer
+            gammas: list of gamma tensors for each conv block
+            betas: list of beta tensors for each conv block
+        """
+        for conv, film, gamma, beta in zip(
+            self.conv_blocks, self.film_layers, gammas, betas
+        ):
+            x = conv(x)
+            x = film(x, gamma, beta)
+        return self.pool(x), x
+
+
+class UpConvBlockFiLM(nn.Module):
+    """
+    Decoder block with FiLM modulation.
+    Conv → FiLM → Conv → FiLM → (optional) ConvTranspose2d
+    """
+
+    def __init__(
+        self,
+        in_channels: list[int],
+        out_channels: list[int],
+        up_conv: bool = True,
+        conv_kwargs: dict = {"kernel_size": 3, "stride": 1, "padding": 1},
+        upconv_kwargs: dict = {"kernel_size": 2, "stride": 2},
+    ):
+        super().__init__()
+        assert len(in_channels) == len(out_channels)
+
+        self.conv_blocks = nn.ModuleList(
+            [
+                ConvBlock(in_ch, out_ch, conv_kwargs)
+                for in_ch, out_ch in zip(in_channels, out_channels)
+            ]
+        )
+
+        self.film_layers = nn.ModuleList([FiLMLayer() for _ in out_channels])
+
+        self.up_conv = up_conv
+        if self.up_conv:
+            self.up_conv_op = nn.ConvTranspose2d(
+                out_channels[-1], out_channels[-1], **upconv_kwargs
+            )
+
+    def forward(self, x: torch.Tensor, gammas: list, betas: list):
+        """
+        Args:
+            x: (B, C, H, W)
+            gammas: list of gamma tensors for each conv block
+            betas: list of beta tensors for each conv block
+        """
+        for conv, film, gamma, beta in zip(
+            self.conv_blocks, self.film_layers, gammas, betas
+        ):
+            x = conv(x)
+            x = film(x, gamma, beta)
+
+        if self.up_conv:
+            x = self.up_conv_op(x)
+
+        return x
+
+
+class UNet2DAttn(BaseUnet):
+    """
+    UNet with shared attention-based conditioning.
+    Computes all gamma/beta at the start of forward pass.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        num_classes: int,
+        n_organs: int,
+        size: int = 32,
+        depth: int = 3,
+        *,
+        attn_start: int = 0,
+        use_attn: bool = True,
+        img_size: int = 512,
+        patch_size: int = 8,
+        emb_dim: int = 768,
+        n_heads: int = 8,
+        distill: bool = False,
+        distill_unet: bool = False,
+        medsam_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/medsam_unfreezed/model.safetensors",
+        unet_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/unet5_attn/model.safetensors",
+        unet_teacher_kwargs: dict | None = None,
+        use_dwt: bool = True,
+        wavelet: str = "haar",
+        dwt_bands: list[str] | None = None,
+        use_shape: bool = False,
+        shape_res: int = 64,
+    ):
+        """
+        Args:
+            in_channels: Number of input channels
+            num_classes: Number of output classes
+            n_organs: Number of organ types
+            size: Base number of channels
+            depth: Number of encoder/decoder levels
+            attn_start: Level where attention starts (0-based)
+            use_attn: If False, use plain Conv blocks
+            img_size: Size of input images (assumed square)
+            patch_size: Patch size for attention
+            emb_dim: Embedding dimension for attention
+            n_heads: Number of attention heads
+        """
+        super().__init__(
+            in_channels=in_channels,
+            num_classes=num_classes,
+            n_organs=n_organs,
+            size=size,
+            depth=depth,
+            attn_start=attn_start,
+            use_attn=use_attn,
+            img_size=img_size,
+            patch_size=patch_size,
+            emb_dim=emb_dim,
+            n_heads=n_heads,
+            distill=distill,
+            distill_unet=distill_unet,
+            medsam_teacher_ckpt=medsam_teacher_ckpt,
+            unet_teacher_ckpt=unet_teacher_ckpt,
+            unet_teacher_kwargs=unet_teacher_kwargs,
+            use_dwt=use_dwt,
+            wavelet=wavelet,
+            dwt_bands=dwt_bands,
+            use_shape=use_shape,
+            shape_res=shape_res,
+        )
+
+    def _build_model(
+        self,
+        *,
+        attn_start: int = 0,
+        use_attn: bool = True,
+        img_size: int = 512,
+        patch_size: int = 8,
+        emb_dim: int = 768,
+        n_heads: int = 8,
+        distill: bool = False,
+        distill_unet: bool = False,
+        medsam_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/medsam_unfreezed/model.safetensors",
+        unet_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/unet5_attn/model.safetensors",
+        unet_teacher_kwargs: dict | None = None,
+        use_dwt: bool = True,
+        wavelet: str = "haar",
+        dwt_bands: list[str] | None = None,
+        use_shape: bool = False,
+        shape_res: int = 64,
+        **kwargs,
+    ):
+        if kwargs:
+            unknown = ", ".join(sorted(kwargs.keys()))
+            raise TypeError(f"Unexpected UNet2DAttn kwargs: {unknown}")
+
+        self.attn_start = max(0, int(attn_start))
+        self.use_attn = bool(use_attn)
+        self.use_shape = bool(use_shape)
+        self.img_size = int(img_size)
+        self.shape_res = int(shape_res)
+        self.criterion = DiceBCELoss()
+        self.steps_counter = 0
+
+        # Compute max channels needed
+        max_channels = self.size * (2 ** (self.depth + 1))
+
+        # Create shared attention modulator
+        if self.use_attn:
+            # Count total number of layers that will use attention
+            n_attn_layers = 0
+            for i in range(self.depth):
+                if i >= self.attn_start:
+                    n_attn_layers += 2  # encoder has 2 conv blocks per level
+            n_attn_layers += 2  # bottleneck
+            for i in range(self.depth):
+                if i >= self.attn_start:
+                    n_attn_layers += 2  # decoder has 2 conv blocks per level
+
+            self.shared_attn = SharedAttnModulator(
+                n_organs=self.n_organs,
+                n_layers=n_attn_layers,
+                max_channels=max_channels,
+                img_size=self.img_size,
+                patch_size=patch_size,
+                emb_dim=emb_dim,
+                n_heads=n_heads,
+                use_dwt=use_dwt,
+                wavelet=wavelet,
+                dwt_bands=dwt_bands,
+                use_shape=use_shape,
+                shape_res=shape_res,
+            )
+
+        # ---------------- Encoder ----------------
+        self.encoder = nn.ModuleDict()
+
+        # First encoder block
+        if self.use_attn and 0 >= self.attn_start:
+            self.encoder["0"] = DownConvBlockFiLM(
+                [self.in_channels, self.size],
+                [self.size, self.size * 2],
+            )
+        else:
+            self.encoder["0"] = DownConvBlock(
+                [self.in_channels, self.size], [self.size, self.size * 2]
+            )
+
+        # Remaining encoder blocks
+        for i in range(1, self.depth):
+            in_ch = [self.size * (2**i), self.size * (2**i)]
+            out_ch = [self.size * (2**i), self.size * (2 ** (i + 1))]
+            key = str(i)
+
+            if self.use_attn and i >= self.attn_start:
+                self.encoder[key] = DownConvBlockFiLM(in_ch, out_ch)
+            else:
+                self.encoder[key] = DownConvBlock(in_ch, out_ch)
+
+        # ---------------- Bottleneck ----------------
+        if self.use_attn:
+            self.bottleneck = UpConvBlockFiLM(
+                [self.size * (2**self.depth), self.size * (2**self.depth)],
+                [self.size * (2**self.depth), self.size * (2 ** (self.depth + 1))],
+            )
+        else:
+            self.bottleneck = UpConvBlock(
+                [self.size * (2**self.depth), self.size * (2**self.depth)],
+                [self.size * (2**self.depth), self.size * (2 ** (self.depth + 1))],
+            )
+
+        # ---------------- Decoder ----------------
+        self.decoder = nn.ModuleDict()
+
+        for i in range(self.depth, 1, -1):
+            use_attn_at_level = self.use_attn and (i - 1) >= self.attn_start
+
+            if use_attn_at_level:
+                self.decoder[str(i - 1)] = UpConvBlockFiLM(
+                    [
+                        self.size * (2 ** (i + 1)) + self.size * (2**i),
+                        self.size * (2**i),
+                    ],
+                    [self.size * (2**i), self.size * (2**i)],
+                )
+            else:
+                self.decoder[str(i - 1)] = UpConvBlock(
+                    [
+                        self.size * (2 ** (i + 1)) + self.size * (2**i),
+                        self.size * (2**i),
+                    ],
+                    [self.size * (2**i), self.size * (2**i)],
+                )
+
+        # Final decoder block
+        if self.use_attn and 0 >= self.attn_start:
+            self.decoder["0"] = UpConvBlockFiLM(
+                [self.size * 4 + self.size * 2, self.size * 2],
+                [self.size * 2, self.size * 2],
+                up_conv=False,
+            )
+        else:
+            self.decoder["0"] = UpConvBlock(
+                [self.size * 4 + self.size * 2, self.size * 2],
+                [self.size * 2, self.size * 2],
+                up_conv=False,
+            )
+
+        self.out_layer = ConvBlock(
+            self.size * 2,
+            self.out_channels,
+            conv_kwargs={"kernel_size": 1, "stride": 1, "padding": 0},
+        )
+        self._init_distillation(
+            distill=distill,
+            distill_unet=distill_unet,
+            medsam_teacher_ckpt=medsam_teacher_ckpt,
+            unet_teacher_ckpt=unet_teacher_ckpt,
+            unet_teacher_kwargs=unet_teacher_kwargs,
+        )
+
+    def _build_layer_configs(self):
+        """
+        Build list of (layer_id, n_channels) for all layers that use attention.
+        Layer IDs are sequential: 0, 1, 2, ...
+        """
+        configs = []
+        layer_id = 0
+
+        # Encoder layers
+        for i in range(self.depth):
+            if self.use_attn and i >= self.attn_start:
+                # Each encoder block has 2 conv outputs
+                configs.append((layer_id, self.size * (2**i)))
+                layer_id += 1
+                configs.append((layer_id, self.size * (2 ** (i + 1))))
+                layer_id += 1
+
+        # Bottleneck (2 conv blocks)
+        if self.use_attn:
+            configs.append((layer_id, self.size * (2**self.depth)))
+            layer_id += 1
+            configs.append((layer_id, self.size * (2 ** (self.depth + 1))))
+            layer_id += 1
+
+        # Decoder layers
+        for i in range(self.depth - 1, -1, -1):
+            if self.use_attn and i >= self.attn_start:
+                # Each decoder block has 2 conv outputs
+                configs.append((layer_id, self.size * (2 ** (i + 1))))
+                layer_id += 1
+                configs.append((layer_id, self.size * (2 ** (i + 1))))
+                layer_id += 1
+
+        return configs
+
+    def _prepare_forward(
+        self,
+        *,
+        pixel_values: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        **kwargs,
+    ) -> dict:
+        forward_ctx = {
+            "mod_list": None,
+            "mod_idx": 0,
+            "projected_shapes": None,
+        }
+        if self.use_attn:
+            layer_configs = self._build_layer_configs()
+            modulations, projected_shapes = self.shared_attn.compute_all_modulations(
+                pixel_values, organ_id, layer_configs
+            )
+            forward_ctx["mod_list"] = [modulations[i] for i in range(len(layer_configs))]
+            forward_ctx["projected_shapes"] = projected_shapes
+        return forward_ctx
+
+    def _next_modulation(self, forward_ctx: dict):
+        mod_idx = forward_ctx["mod_idx"]
+        mod_list = forward_ctx["mod_list"]
+        gammas = [mod_list[mod_idx][0], mod_list[mod_idx + 1][0]]
+        betas = [mod_list[mod_idx][1], mod_list[mod_idx + 1][1]]
+        forward_ctx["mod_idx"] = mod_idx + 2
+        return gammas, betas
+
+    def _encode(
+        self,
+        layer: nn.Module,
+        x: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        forward_ctx: dict | None = None,
+    ):
+        if isinstance(layer, DownConvBlockFiLM):
+            gammas, betas = self._next_modulation(forward_ctx)
+            return layer(x, gammas, betas)
+        return layer(x)
+
+    def _bottleneck(
+        self,
+        x: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        forward_ctx: dict | None = None,
+    ):
+        if isinstance(self.bottleneck, UpConvBlockFiLM):
+            gammas, betas = self._next_modulation(forward_ctx)
+            return self.bottleneck(x, gammas, betas)
+        return self.bottleneck(x)
+
+    def _decode(
+        self,
+        layer: nn.Module,
+        x: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        forward_ctx: dict | None = None,
+    ):
+        if isinstance(layer, UpConvBlockFiLM):
+            gammas, betas = self._next_modulation(forward_ctx)
+            return layer(x, gammas, betas)
+        return layer(x)
+
+    def _apply_auxiliary_losses(
+        self,
+        *,
+        loss: torch.Tensor | float,
+        logits: torch.Tensor,
+        masks: torch.Tensor | None,
+        organ_id: torch.Tensor | None = None,
+        forward_ctx: dict | None = None,
+        **kwargs,
+    ):
+        self.steps_counter += 1
+        if not self.use_shape or masks is None:
+            return loss
+
+        projected_shapes = forward_ctx.get("projected_shapes", None)
+        if projected_shapes is None or projected_shapes.numel() == 0:
+            return loss
+
+        reduced_masks = canonicalize_mask_batch_normalized(
+            masks, canon_res=self.shape_res
+        )
+        loss_shape = torch.tensor([0.0], requires_grad=True).to(logits.device)
+        for projected_shape in projected_shapes:
+            loss_shape = loss_shape + self.criterion(projected_shape, reduced_masks)
+
+        loss_shape = loss_shape / projected_shapes.shape[0]
+        if wandb.run is not None:
+            wandb.log({"shape_loss": loss_shape.item()}, commit=False)
+        loss = loss + loss_shape
+
+        if self.steps_counter % 200 == 0:
+            with torch.no_grad():
+                proj_tokens = torch.sigmoid(
+                    self.shared_attn.shape_proj(
+                        self.shared_attn.shape_embed.weight[:8].clone().detach()
+                    )
+                    .reshape(8, self.shape_res, self.shape_res)
+                    .detach()
+                )
+            proj_tokens = (proj_tokens > 0.5).to(torch.float32)
+            proj_vis = (
+                proj_tokens.view(2, 4, self.shape_res, self.shape_res)
+                .permute(0, 2, 1, 3)
+                .reshape(self.shape_res * 2, self.shape_res * 4)
+            )
+            wandb.log({"projected_tokens": wandb.Image(proj_vis.unsqueeze(0))})
+            print("log")
+
+        return loss

unet_base.py

@@ -0,0 +1,386 @@
+from copy import deepcopy
+from abc import ABC, abstractmethod
+import torch
+import torch.nn.functional as F
+from safetensors.torch import load_file
+from torch import nn
+from torchvision.transforms import v2
+import wandb
+
+
+class BaseUnet(nn.Module, ABC):
+    def __init__(
+        self,
+        in_channels: int,
+        num_classes: int,
+        n_organs: int,
+        size: int = 32,
+        depth: int = 3,
+        **kwargs
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = num_classes
+        self.n_organs = n_organs
+        self.size = size
+        self.depth = depth
+        self.distill = False
+        self.distill_unet = False
+        self.distill_model = None
+        self.distill_adapter = None
+        self.distill_loss = None
+        self._build_model(**kwargs)
+
+    @staticmethod
+    def _pad_to_2d(x: torch.Tensor, stride: int):
+        h, w = x.shape[-2:]
+
+        new_h = h if h % stride == 0 else h + stride - (h % stride)
+        new_w = w if w % stride == 0 else w + stride - (w % stride)
+
+        top = (new_h - h) // 2
+        bottom = (new_h - h) - top
+        left = (new_w - w) // 2
+        right = (new_w - w) - left
+
+        pads = (left, right, top, bottom)
+        x_pad = F.pad(x, pads, mode="constant", value=0)
+        return x_pad, pads
+
+    @staticmethod
+    def _unpad_2d(x: torch.Tensor, pads):
+        left, right, top, bottom = pads
+
+        if top or bottom:
+            end_h = -bottom if bottom > 0 else None
+            x = x[:, :, top:end_h, :]
+
+        if left or right:
+            end_w = -right if right > 0 else None
+            x = x[:, :, :, left:end_w]
+
+        return x
+
+    @abstractmethod
+    def _build_model(self, **kwargs):
+        pass
+
+    def _prepare_forward(
+        self,
+        *,
+        pixel_values: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        **kwargs,
+    ) -> dict:
+        return {}
+
+    @abstractmethod
+    def _encode(
+        self,
+        layer: nn.Module,
+        x: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        forward_ctx: dict | None = None,
+    ):
+        pass
+
+    @abstractmethod
+    def _bottleneck(
+        self,
+        x: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        forward_ctx: dict | None = None,
+    ):
+        pass
+
+    @abstractmethod
+    def _decode(
+        self,
+        layer: nn.Module,
+        x: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        forward_ctx: dict | None = None,
+    ):
+        pass
+
+    def encode(
+        self,
+        x: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        forward_ctx: dict | None = None,
+    ):
+        feat_list = []
+        pads = None
+
+        pre_padding = (
+            (x.size(-1) % 2**self.depth != 0)
+            or (x.size(-2) % 2**self.depth != 0)
+            or (x.size(-3) % 2**self.depth != 0)
+        )
+        if pre_padding:
+            x, pads = self._pad_to_2d(x, 2**self.depth)
+
+        out, feat = self._encode(
+            self.encoder["0"], x, organ_id=organ_id, forward_ctx=forward_ctx
+        )
+        feat_list.append(feat)
+
+        for key in list(self.encoder.keys())[1:]:
+            out, feat = self._encode(
+                self.encoder[key], out, organ_id=organ_id, forward_ctx=forward_ctx
+            )
+            feat_list.append(feat)
+
+        out = self._bottleneck(out, organ_id=organ_id, forward_ctx=forward_ctx)
+        return out, feat_list, pads
+
+    def decode(
+        self,
+        out: torch.Tensor,
+        feat_list: list[torch.Tensor],
+        pads,
+        organ_id: torch.Tensor | None = None,
+        forward_ctx: dict | None = None,
+    ):
+        for key in self.decoder:
+            out = self._decode(
+                self.decoder[key],
+                torch.cat((out, feat_list[int(key)]), dim=1),
+                organ_id=organ_id,
+                forward_ctx=forward_ctx,
+            )
+            del feat_list[int(key)]
+
+        out = self.out_layer(out)
+        if pads is not None:
+            out = self._unpad_2d(out, pads).squeeze(1)
+        return out
+
+    def _apply_auxiliary_losses(
+        self,
+        *,
+        loss: torch.Tensor | float,
+        logits: torch.Tensor,
+        masks: torch.Tensor | None,
+        organ_id: torch.Tensor | None = None,
+        forward_ctx: dict | None = None,
+        **kwargs,
+    ):
+        return loss
+
+    def _init_distillation(
+        self,
+        *,
+        distill: bool = False,
+        distill_unet: bool = False,
+        medsam_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/medsam_unfreezed/model.safetensors",
+        unet_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/unet5_attn_distilled/model.safetensors",
+        unet_teacher_kwargs: dict | None = None,
+    ):
+        if distill and distill_unet:
+            raise ValueError("distill and distill_unet cannot both be enabled.")
+
+        self.distill = bool(distill)
+        self.distill_unet = bool(distill_unet)
+        self.distill_model = None
+        self.distill_adapter = None
+        self.distill_loss = None
+
+        if not self.distill and not self.distill_unet:
+            return
+
+        from .segm_net import DistillationLoss, MedSAM
+
+        student_channels = (2048 // (32 // self.size)) // (2 ** (5 - self.depth))
+
+        if self.distill:
+            self.distill_adapter = nn.Conv2d(student_channels, 256, kernel_size=1)
+        else:
+            self.distill_adapter = nn.Conv2d(student_channels, 2048, kernel_size=1)
+
+        if self.distill:
+            from segment_anything import sam_model_registry
+            from utils.paths import MEDSAM_BASE_WEIGHTS
+
+            sam_model = sam_model_registry["vit_b"](checkpoint=MEDSAM_BASE_WEIGHTS)
+            self.distill_model = MedSAM(
+                image_encoder=deepcopy(sam_model.image_encoder),
+                mask_decoder=deepcopy(sam_model.mask_decoder),
+                prompt_encoder=deepcopy(sam_model.prompt_encoder),
+                predict_bboxes=True,
+                freeze_image_encoder=0,
+            )
+            state_dict = load_file(medsam_teacher_ckpt)
+            load_result = self.distill_model.load_state_dict(state_dict)
+            print(f"Loaded MedSam teacher model and loaded weights:\n{load_result}")
+        else:
+            from .unet_attn import UNet2DAttn
+
+            teacher_kwargs = {
+                "in_channels": 3,
+                "num_classes": 1,
+                "n_organs": 10,
+                "size": 32,
+                "depth": 5,
+                "attn_start": 0,
+                "use_attn": True,
+                "img_size": 512,
+                "patch_size": 8,
+                "emb_dim": 768,
+                "n_heads": 8,
+                "distill": False,
+                "distill_unet": False,
+                "use_dwt": False,
+                "wavelet": "haar",
+                "use_shape": False,
+                "shape_res": 64,
+            }
+            if unet_teacher_kwargs is not None:
+                teacher_kwargs.update(unet_teacher_kwargs)
+
+            self.distill_model = UNet2DAttn(**teacher_kwargs)
+            state_dict = load_file(unet_teacher_ckpt)
+            state_dict = {k: v for k, v in state_dict.items() if "distill" not in k}
+            load_result = self.distill_model.load_state_dict(state_dict)
+            print(f"Loaded UNet teacher model and loaded weights:\n{load_result}")
+
+        for p in self.distill_model.parameters():
+            p.requires_grad = False
+        self.distill_model.eval()
+        self.distill_loss = DistillationLoss()
+
+    def _forward_distillation(
+        self,
+        *,
+        student_logits: torch.Tensor,
+        student_bottleneck: torch.Tensor,
+        pixel_values: torch.Tensor,
+        organ_id: torch.Tensor | None = None,
+        pixel_values_medsam: torch.Tensor | None = None,
+    ) -> torch.Tensor | None:
+        if not self.distill and not self.distill_unet:
+            return None
+
+        if self.distill:
+            if pixel_values_medsam is None:
+                raise ValueError("pixel_values_medsam is required when distill=True.")
+
+            with torch.no_grad():
+                up_pixel_values = v2.functional.resize(
+                    pixel_values_medsam, 1024, v2.InterpolationMode.BICUBIC
+                )
+                image_embedding = self.distill_model.image_encoder(up_pixel_values)
+
+            student_resized = F.interpolate(
+                student_bottleneck,
+                size=(image_embedding.shape[-1], image_embedding.shape[-1]),
+                mode="bilinear",
+                align_corners=False,
+            )
+            up_feat = self.distill_adapter(student_resized)
+            distill_loss_emb = self.distill_loss(
+                student_logits=up_feat,
+                teacher_logits=image_embedding.detach(),
+            )
+
+            if wandb.run is not None:
+                wandb.log(
+                    {
+                        "distill_loss_emb": distill_loss_emb["loss"].item(),
+                    },
+                    commit=False,
+                )
+
+            return distill_loss_emb["loss"]
+
+        with torch.no_grad():
+            forward_ctx = self._prepare_forward(
+                pixel_values=pixel_values,
+                organ_id=organ_id,
+            )
+
+            teacher_embedding, _, _ = self.encode(
+                pixel_values, organ_id=organ_id, forward_ctx=forward_ctx
+            )
+
+        student_resized = F.interpolate(
+                student_bottleneck,
+                size=(teacher_embedding.shape[-1], teacher_embedding.shape[-1]),
+                mode="bilinear",
+                align_corners=False,
+            )
+        up_feat = self.distill_adapter(student_resized)
+        distill_loss_emb = self.distill_loss(
+            student_logits=student_bottleneck,
+            teacher_logits=teacher_embedding.detach(),
+        )
+        if wandb.run is not None:
+            wandb.log(
+                {"distill_loss_logits": distill_loss_emb["loss"].item()},
+                commit=False,
+            )
+        return distill_loss_emb["loss"]
+
+    def forward(
+        self,
+        pixel_values,
+        organ_id=None,
+        labels=None,
+        masks=None,
+        bbox_coords=None,
+        organ_id_metric=None,
+        teacher_embedding=None,
+        teacher_mask=None,
+        pixel_values_medsam=None,
+        **kwargs,
+    ):
+        forward_ctx = self._prepare_forward(
+            pixel_values=pixel_values,
+            organ_id=organ_id,
+            masks=masks,
+            **kwargs,
+        )
+
+        out_bottleneck, feat_list, pads = self.encode(
+            pixel_values, organ_id=organ_id, forward_ctx=forward_ctx
+        )
+        out = self.decode(
+            out_bottleneck,
+            feat_list,
+            pads,
+            organ_id=organ_id,
+            forward_ctx=forward_ctx,
+        )
+
+        if masks is not None:
+            loss = self.criterion(out, masks)
+        else:
+            loss = 0.0
+
+        distill_loss = self._forward_distillation(
+            student_logits=out,
+            student_bottleneck=out_bottleneck,
+            pixel_values=pixel_values,
+            organ_id=organ_id,
+            pixel_values_medsam=pixel_values_medsam,
+        )
+        if distill_loss is not None:
+            loss = loss + distill_loss
+
+        loss = self._apply_auxiliary_losses(
+            loss=loss,
+            logits=out,
+            masks=masks,
+            organ_id=organ_id,
+            forward_ctx=forward_ctx,
+            pixel_values=pixel_values,
+            **kwargs,
+        )
+
+        return {
+            "loss": loss,
+            "logits": out,
+            "labels": masks,
+            "organ_id": organ_id,
+            "organ_id_metric": organ_id_metric,
+        }