Upload SegFace hair segmentation model bundle

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/epoch_010.png filter=lfs diff=lfs merge=lfs -text
+assets/training_curves.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,83 @@
+---
+library_name: pytorch
+pipeline_tag: image-segmentation
+tags:
+- pytorch
+- image-segmentation
+- hair-segmentation
+- segface
+- korean-hairstyle
+- custom-code
+language:
+- ko
+---
+
+# SegFace_k-hair
+
+Personal Hugging Face model repository for a custom PyTorch hair-only segmentation checkpoint trained on a K-Hairstyle based AIHub subset.
+
+## Model Summary
+
+- Backbone: `swin_base`
+- Input size: `512x512`
+- Freeze backbone: `True`
+- LoRA: `rank=8`, `alpha=16.0`, `dropout=0.05`
+- Threshold used during validation: `0.5`
+- Train / Val split used for this run: `50,000` / `5,000`
+
+## Validation Metrics
+
+These are validation metrics from the training run. A separate independent hold-out test split has not been populated yet, so treat these as validation-only results.
+
+| Metric | Value |
+| --- | ---: |
+| Best epoch | 7 |
+| Val IoU | 0.9487 |
+| Val Dice | 0.9736 |
+| Val Precision | 0.9723 |
+| Val Recall | 0.9751 |
+| Epochs completed | 10 |
+| Avg epoch time (sec) | 3546.45 |
+
+## Bundle Contents
+
+- `best.pt`: inference checkpoint
+- `config.json`: training-time model config
+- `training_run_summary.json`: run summary and validation metrics
+- `inference.py`: local / Hub inference example
+- `requirements.txt`: minimal runtime dependencies
+- `hair_mask_dataset/`, `models/`: custom model code required to load the checkpoint
+
+## Inference
+
+Run locally from the root of this model bundle:
+
+```bash
+python inference.py
+  --checkpoint best.pt
+  --input path/to/input.jpg
+  --output-mask output_mask.png
+  --output-overlay output_overlay.png
+```
+
+You can also load directly from the Hugging Face Hub after uploading:
+
+```bash
+python inference.py
+  --repo-id your-username/SegFace_k-hair
+  --input path/to/input.jpg
+  --output-mask output_mask.png
+  --output-overlay output_overlay.png
+```
+
+## Notes
+
+- This repo contains custom code and a raw PyTorch checkpoint, not a Transformers-format model.
+- Preprocessing expects RGB input, resize to `512`, ImageNet normalization, and sigmoid threshold `0.5`.
+- Before making the repository public, verify whether your AIHub / K-Hairstyle data usage terms allow public redistribution of derived model weights.
+
+## Training Artifacts
+
+![Training Curve](assets/training_curves.png)
+
+![Preview](assets/epoch_010.png)

best.pt

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

config.json

@@ -0,0 +1,31 @@
+{
+  "prepared_root": "/workspace/runpod_upload_ready/data/aihub_hairmask_hq_budget_50k",
+  "raw_root": "/workspace/runpod_upload_ready/data/aihub_korean_hairstyle_hq_raw",
+  "image_size": 512,
+  "model_name": "swin_base",
+  "run_dir": "/workspace/runpod_upload_ready/hair_mask_dataset/runs/segface_hair_budget_4090",
+  "epochs": 10,
+  "batch_size": 2,
+  "accumulation_steps": 2,
+  "num_workers": 6,
+  "lr": 0.0001,
+  "weight_decay": 0.0001,
+  "amp": true,
+  "threshold": 0.5,
+  "seed": 42,
+  "save_every": 1,
+  "freeze_backbone": true,
+  "lora_rank": 8,
+  "lora_alpha": 16.0,
+  "lora_dropout": 0.05,
+  "lora_targets": [
+    "attn.qkv",
+    "attn.proj",
+    "mlp.0",
+    "mlp.3"
+  ],
+  "compile_model": true,
+  "channels_last": true,
+  "trainable_params": 5804672,
+  "total_params": 92547896
+}

hair_mask_dataset/__init__.py

@@ -0,0 +1,3 @@
+from .dataset import AIHubHairMaskDataset
+
+__all__ = ["AIHubHairMaskDataset"]

hair_mask_dataset/segface_hair_model.py

@@ -0,0 +1,264 @@
+from __future__ import annotations
+
+import math
+from typing import Dict, Iterable, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+from torchvision.models.swin_transformer import ShiftedWindowAttention
+
+from models.segface.models.segface_celeb import SegFaceCeleb
+
+
+HAIR_CLASS_INDEX = 14
+DEFAULT_LORA_TARGETS = ("attn.qkv", "attn.proj", "mlp.0", "mlp.3")
+
+
+def load_segface_pretrained(model: nn.Module) -> None:
+    ckpt_path = hf_hub_download(
+        repo_id="kartiknarayan/SegFace",
+        filename="swinb_celeba_512/model_299.pt",
+    )
+    checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+    state_dict = checkpoint["state_dict_backbone"] if "state_dict_backbone" in checkpoint else checkpoint
+    model.load_state_dict(state_dict, strict=False)
+
+
+class LoRALinear(nn.Module):
+    def __init__(self, base: nn.Linear, *, rank: int, alpha: float, dropout: float) -> None:
+        super().__init__()
+        if rank <= 0:
+            raise ValueError("LoRA rank must be positive.")
+
+        self.base = base
+        self.rank = rank
+        self.scaling = alpha / rank
+        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        self.lora_down = nn.Linear(base.in_features, rank, bias=False)
+        self.lora_up = nn.Linear(rank, base.out_features, bias=False)
+
+        nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
+        nn.init.zeros_(self.lora_up.weight)
+
+        self.base.weight.requires_grad = False
+        if self.base.bias is not None:
+            self.base.bias.requires_grad = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        base_out = self.base(x)
+        lora_out = self.lora_up(self.lora_down(self.dropout(x))) * self.scaling
+        return base_out + lora_out
+
+
+def linear_bias(module: nn.Module) -> torch.Tensor | None:
+    if isinstance(module, LoRALinear):
+        return module.base.bias
+    if isinstance(module, nn.Linear):
+        return module.bias
+    raise TypeError(f"Unsupported linear module type: {type(module)!r}")
+
+
+def linear_with_lora(module: nn.Module, x: torch.Tensor, *, bias_override: torch.Tensor | None = None) -> torch.Tensor:
+    if isinstance(module, LoRALinear):
+        base_bias = module.base.bias if bias_override is None else bias_override
+        base_out = F.linear(x, module.base.weight, base_bias)
+        lora_hidden = F.linear(module.dropout(x), module.lora_down.weight, None)
+        lora_out = F.linear(lora_hidden, module.lora_up.weight, None) * module.scaling
+        return base_out + lora_out
+    if isinstance(module, nn.Linear):
+        bias = module.bias if bias_override is None else bias_override
+        return F.linear(x, module.weight, bias)
+    raise TypeError(f"Unsupported linear module type: {type(module)!r}")
+
+
+def shifted_window_attention_with_modules(
+    input: torch.Tensor,
+    qkv_module: nn.Module,
+    proj_module: nn.Module,
+    relative_position_bias: torch.Tensor,
+    window_size: list[int],
+    num_heads: int,
+    shift_size: list[int],
+    attention_dropout: float = 0.0,
+    dropout: float = 0.0,
+    logit_scale: torch.Tensor | None = None,
+    training: bool = True,
+) -> torch.Tensor:
+    B, H, W, C = input.shape
+    pad_r = (window_size[1] - W % window_size[1]) % window_size[1]
+    pad_b = (window_size[0] - H % window_size[0]) % window_size[0]
+    x = F.pad(input, (0, 0, 0, pad_r, 0, pad_b))
+    _, pad_H, pad_W, _ = x.shape
+
+    shift_size = shift_size.copy()
+    if window_size[0] >= pad_H:
+        shift_size[0] = 0
+    if window_size[1] >= pad_W:
+        shift_size[1] = 0
+
+    if sum(shift_size) > 0:
+        x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(1, 2))
+
+    num_windows = (pad_H // window_size[0]) * (pad_W // window_size[1])
+    x = x.view(B, pad_H // window_size[0], window_size[0], pad_W // window_size[1], window_size[1], C)
+    x = x.permute(0, 1, 3, 2, 4, 5).reshape(B * num_windows, window_size[0] * window_size[1], C)
+
+    qkv_bias = linear_bias(qkv_module)
+    if logit_scale is not None and qkv_bias is not None:
+        qkv_bias = qkv_bias.clone()
+        length = qkv_bias.numel() // 3
+        qkv_bias[length : 2 * length].zero_()
+
+    qkv = linear_with_lora(qkv_module, x, bias_override=qkv_bias)
+    qkv = qkv.reshape(x.size(0), x.size(1), 3, num_heads, C // num_heads).permute(2, 0, 3, 1, 4)
+    q, k, v = qkv[0], qkv[1], qkv[2]
+
+    if logit_scale is not None:
+        attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)
+        logit_scale = torch.clamp(logit_scale, max=math.log(100.0)).exp()
+        attn = attn * logit_scale
+    else:
+        q = q * (C // num_heads) ** -0.5
+        attn = q.matmul(k.transpose(-2, -1))
+
+    attn = attn + relative_position_bias
+
+    if sum(shift_size) > 0:
+        attn_mask = x.new_zeros((pad_H, pad_W))
+        h_slices = ((0, -window_size[0]), (-window_size[0], -shift_size[0]), (-shift_size[0], None))
+        w_slices = ((0, -window_size[1]), (-window_size[1], -shift_size[1]), (-shift_size[1], None))
+        count = 0
+        for h in h_slices:
+            for w in w_slices:
+                attn_mask[h[0] : h[1], w[0] : w[1]] = count
+                count += 1
+        attn_mask = attn_mask.view(pad_H // window_size[0], window_size[0], pad_W // window_size[1], window_size[1])
+        attn_mask = attn_mask.permute(0, 2, 1, 3).reshape(num_windows, window_size[0] * window_size[1])
+        attn_mask = attn_mask.unsqueeze(1) - attn_mask.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+        attn = attn.view(x.size(0) // num_windows, num_windows, num_heads, x.size(1), x.size(1))
+        attn = attn + attn_mask.unsqueeze(1).unsqueeze(0)
+        attn = attn.view(-1, num_heads, x.size(1), x.size(1))
+
+    attn = F.softmax(attn, dim=-1)
+    attn = F.dropout(attn, p=attention_dropout, training=training)
+
+    x = attn.matmul(v).transpose(1, 2).reshape(x.size(0), x.size(1), C)
+    x = linear_with_lora(proj_module, x)
+    x = F.dropout(x, p=dropout, training=training)
+
+    x = x.view(B, pad_H // window_size[0], pad_W // window_size[1], window_size[0], window_size[1], C)
+    x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, pad_H, pad_W, C)
+
+    if sum(shift_size) > 0:
+        x = torch.roll(x, shifts=(shift_size[0], shift_size[1]), dims=(1, 2))
+
+    return x[:, :H, :W, :].contiguous()
+
+
+def lora_compatible_swin_attention_forward(self: ShiftedWindowAttention, x: torch.Tensor) -> torch.Tensor:
+    relative_position_bias = self.get_relative_position_bias()
+    return shifted_window_attention_with_modules(
+        x,
+        self.qkv,
+        self.proj,
+        relative_position_bias,
+        self.window_size,
+        self.num_heads,
+        shift_size=self.shift_size,
+        attention_dropout=self.attention_dropout,
+        dropout=self.dropout,
+        logit_scale=getattr(self, "logit_scale", None),
+        training=self.training,
+    )
+
+
+def patch_swin_attention_for_lora(module: nn.Module) -> int:
+    patched = 0
+    for child in module.modules():
+        if isinstance(child, ShiftedWindowAttention) and not getattr(child, "_lora_forward_patched", False):
+            child.forward = lora_compatible_swin_attention_forward.__get__(child, type(child))
+            child._lora_forward_patched = True
+            patched += 1
+    return patched
+
+
+def freeze_module(module: nn.Module) -> None:
+    for param in module.parameters():
+        param.requires_grad = False
+
+
+def apply_lora(
+    module: nn.Module,
+    *,
+    rank: int,
+    alpha: float,
+    dropout: float,
+    target_patterns: Iterable[str],
+    prefix: str = "",
+) -> int:
+    replaced = 0
+    patterns = tuple(pattern for pattern in target_patterns if pattern)
+
+    for child_name, child in list(module.named_children()):
+        full_name = f"{prefix}.{child_name}" if prefix else child_name
+        if isinstance(child, nn.Linear) and any(pattern in full_name for pattern in patterns):
+            setattr(module, child_name, LoRALinear(child, rank=rank, alpha=alpha, dropout=dropout))
+            replaced += 1
+            continue
+        replaced += apply_lora(
+            child,
+            rank=rank,
+            alpha=alpha,
+            dropout=dropout,
+            target_patterns=patterns,
+            prefix=full_name,
+        )
+
+    return replaced
+
+
+class SegFaceHairModel(nn.Module):
+    def __init__(
+        self,
+        *,
+        input_resolution: int = 512,
+        model_name: str = "swin_base",
+        load_pretrained: bool = True,
+        freeze_backbone: bool = False,
+        lora_rank: int = 0,
+        lora_alpha: float = 16.0,
+        lora_dropout: float = 0.0,
+        lora_targets: Iterable[str] = DEFAULT_LORA_TARGETS,
+    ) -> None:
+        super().__init__()
+        self.segface = SegFaceCeleb(input_resolution=input_resolution, model=model_name)
+        if load_pretrained:
+            load_segface_pretrained(self.segface)
+
+        if freeze_backbone:
+            freeze_module(self.segface.backbone)
+
+        self.lora_target_patterns: Tuple[str, ...] = tuple(pattern for pattern in lora_targets if pattern)
+        self.lora_replaced = 0
+        self.swin_attention_patched = 0
+        if lora_rank > 0:
+            self.lora_replaced = apply_lora(
+                self.segface.backbone,
+                rank=lora_rank,
+                alpha=lora_alpha,
+                dropout=lora_dropout,
+                target_patterns=self.lora_target_patterns,
+            )
+            if model_name.startswith("swin"):
+                self.swin_attention_patched = patch_swin_attention_for_lora(self.segface.backbone)
+
+    def forward(self, images: torch.Tensor) -> Dict[str, torch.Tensor]:
+        logits = self.segface(images, None, None)
+        hair_logits = logits[:, HAIR_CLASS_INDEX : HAIR_CLASS_INDEX + 1]
+        return {
+            "hair_logits": hair_logits,
+            "all_logits": logits,
+        }

inference.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+from torchvision.transforms import InterpolationMode
+from torchvision.transforms import functional as TF
+
+from hair_mask_dataset.segface_hair_model import SegFaceHairModel
+
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run hair segmentation inference.")
+    parser.add_argument("--input", required=True, help="Path to the input image.")
+    parser.add_argument("--output-mask", required=True, help="Where to save the predicted binary mask.")
+    parser.add_argument("--output-overlay", default="", help="Optional overlay output path.")
+    parser.add_argument("--checkpoint", default="best.pt", help="Local checkpoint path.")
+    parser.add_argument("--config", default="config.json", help="Local config path.")
+    parser.add_argument("--repo-id", default="", help="Optional Hugging Face repo id to download best.pt/config.json from.")
+    parser.add_argument("--revision", default="main", help="Hub revision to download from when using --repo-id.")
+    parser.add_argument("--threshold", type=float, default=None, help="Override sigmoid threshold.")
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu", help="Inference device.")
+    return parser.parse_args()
+
+
+def resolve_artifacts(args: argparse.Namespace) -> tuple[Path, Path]:
+    if args.repo_id:
+        checkpoint_path = Path(
+            hf_hub_download(repo_id=args.repo_id, filename="best.pt", revision=args.revision)
+        )
+        config_path = Path(
+            hf_hub_download(repo_id=args.repo_id, filename="config.json", revision=args.revision)
+        )
+        return checkpoint_path, config_path
+    return Path(args.checkpoint), Path(args.config)
+
+
+def load_model(checkpoint_path: Path, config_path: Path, device: torch.device) -> tuple[torch.nn.Module, dict]:
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    config = checkpoint.get("config")
+    if config is None:
+        config = json.loads(config_path.read_text(encoding="utf-8"))
+
+    model = SegFaceHairModel(
+        input_resolution=config["image_size"],
+        model_name=config["model_name"],
+        load_pretrained=False,
+        freeze_backbone=config["freeze_backbone"],
+        lora_rank=config["lora_rank"],
+        lora_alpha=config["lora_alpha"],
+        lora_dropout=config["lora_dropout"],
+        lora_targets=config["lora_targets"],
+    )
+    model.load_state_dict(checkpoint["model_state"], strict=False)
+    model.to(device)
+    model.eval()
+    return model, config
+
+
+def preprocess(image: Image.Image, image_size: int) -> torch.Tensor:
+    resized = TF.resize(image, [image_size, image_size], interpolation=InterpolationMode.BILINEAR)
+    tensor = TF.to_tensor(resized)
+    tensor = TF.normalize(tensor, IMAGENET_MEAN, IMAGENET_STD)
+    return tensor.unsqueeze(0)
+
+
+def build_overlay(image: Image.Image, mask_u8: np.ndarray) -> Image.Image:
+    image_np = np.asarray(image.convert("RGB"), dtype=np.uint8).copy()
+    overlay = image_np.copy()
+    overlay[mask_u8 > 127] = (overlay[mask_u8 > 127] * 0.4 + np.array([64, 255, 64]) * 0.6).astype(np.uint8)
+    return Image.fromarray(overlay)
+
+
+def main() -> None:
+    args = parse_args()
+    checkpoint_path, config_path = resolve_artifacts(args)
+    device = torch.device(args.device)
+    model, config = load_model(checkpoint_path, config_path, device)
+    threshold = args.threshold if args.threshold is not None else config.get("threshold", 0.5)
+
+    image_path = Path(args.input)
+    output_mask_path = Path(args.output_mask)
+    output_mask_path.parent.mkdir(parents=True, exist_ok=True)
+
+    image = Image.open(image_path).convert("RGB")
+    original_size = image.size
+    inputs = preprocess(image, int(config["image_size"])).to(device)
+
+    with torch.no_grad():
+        logits = model(inputs)["hair_logits"]
+        probs = torch.sigmoid(logits)[0, 0].cpu().numpy()
+
+    mask_small = (probs >= threshold).astype(np.uint8) * 255
+    mask_image = Image.fromarray(mask_small, mode="L").resize(original_size, resample=Image.NEAREST)
+    mask_image.save(output_mask_path)
+
+    if args.output_overlay:
+        output_overlay_path = Path(args.output_overlay)
+        output_overlay_path.parent.mkdir(parents=True, exist_ok=True)
+        overlay = build_overlay(image, np.asarray(mask_image, dtype=np.uint8))
+        overlay.save(output_overlay_path)
+
+    print(f"Saved mask to {output_mask_path}")
+    if args.output_overlay:
+        print(f"Saved overlay to {args.output_overlay}")
+
+
+if __name__ == "__main__":
+    main()

models/segface/__init__.py

@@ -0,0 +1,13 @@
+from .models import SegFaceLapa, SegFaceCeleb, SegFaceHelen
+
+
+def get_model(backbone, input_resolution, model):
+    if backbone == "segface_lapa":
+        model = SegFaceLapa(input_resolution, model)
+    elif backbone == "segface_celeb":
+        model = SegFaceCeleb(input_resolution, model)
+    elif backbone == "segface_helen":
+        model = SegFaceHelen(input_resolution, model)
+    else:
+        raise ValueError("Backbone not implemented")
+    return model

models/segface/models/__init__.py

@@ -0,0 +1,3 @@
+from .segface_lapa import SegFaceLapa
+from .segface_celeb import SegFaceCeleb
+from .segface_helen import SegFaceHelen

models/segface/models/segface_celeb.py

@@ -0,0 +1,348 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+from typing import Any, Optional, Tuple, Type
+from torchvision.models import convnext_large, convnext_base, convnext_small, convnext_tiny, swin_b, swin_v2_b, swin_v2_s, swin_v2_t, mobilenet_v3_large, efficientnet_v2_m
+import pdb
+import numpy as np
+import sys
+import os
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+from models.segface.models.transformer import *
+from models.segface.models.utils_models import *
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        num_layers: int,
+        sigmoid_output: bool = False,
+    ) -> None:
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(
+            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
+        )
+        self.sigmoid_output = sigmoid_output
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        if self.sigmoid_output:
+            x = F.sigmoid(x)
+        return x
+    
+class FaceDecoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        transformer_dim: 256,
+        transformer: nn.Module,
+        activation: Type[nn.Module] = nn.GELU,
+    ) -> None:
+        
+        super().__init__()
+        self.transformer_dim = transformer_dim
+        self.transformer = transformer
+
+        self.background_token = nn.Embedding(1, transformer_dim)
+        self.neck_token = nn.Embedding(1, transformer_dim)
+        self.face_token = nn.Embedding(1, transformer_dim)
+        self.cloth_token = nn.Embedding(1, transformer_dim)
+        self.rightear_token = nn.Embedding(1, transformer_dim)
+        self.leftear_token = nn.Embedding(1, transformer_dim)
+        self.rightbro_token = nn.Embedding(1, transformer_dim)
+        self.leftbro_token = nn.Embedding(1, transformer_dim)
+        self.righteye_token = nn.Embedding(1, transformer_dim)
+        self.lefteye_token = nn.Embedding(1, transformer_dim)
+        self.nose_token = nn.Embedding(1, transformer_dim)
+        self.innermouth_token = nn.Embedding(1, transformer_dim)
+        self.lowerlip_token = nn.Embedding(1, transformer_dim)
+        self.upperlip_token = nn.Embedding(1, transformer_dim)
+        self.hair_token = nn.Embedding(1, transformer_dim)
+        self.glass_token = nn.Embedding(1, transformer_dim)
+        self.hat_token = nn.Embedding(1, transformer_dim)
+        self.earring_token = nn.Embedding(1, transformer_dim)
+        self.necklace_token = nn.Embedding(1, transformer_dim)
+
+
+        self.output_upscaling = nn.Sequential(
+            nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2),
+            LayerNorm2d(transformer_dim // 4),
+            activation(),
+            nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
+            activation(),
+        )
+        
+        self.output_hypernetwork_mlps = MLP(
+            transformer_dim, transformer_dim, transformer_dim // 8, 3
+            )
+
+    def forward(
+        self,
+        image_embeddings: torch.Tensor,
+        image_pe: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        '''
+        image_embeddings - torch.Size([1, 256, 128, 128])
+        image_pe - torch.Size([1, 256, 128, 128])
+        '''
+        output_tokens = torch.cat([
+            self.background_token.weight, self.neck_token.weight, self.face_token.weight, self.cloth_token.weight, 
+            self.rightear_token.weight, self.leftear_token.weight, self.rightbro_token.weight, self.leftbro_token.weight, 
+            self.righteye_token.weight, self.lefteye_token.weight, self.nose_token.weight, self.innermouth_token.weight, 
+            self.lowerlip_token.weight, self.upperlip_token.weight, self.hair_token.weight, self.glass_token.weight, 
+            self.hat_token.weight, self.earring_token.weight, self.necklace_token.weight], dim=0)
+
+        tokens = output_tokens.unsqueeze(0).expand(image_embeddings.size(0), -1, -1) ##### torch.Size([4, 11, 256])
+
+        src = image_embeddings ##### torch.Size([4, 256, 128, 128])
+        pos_src = image_pe.expand(image_embeddings.size(0), -1, -1, -1)
+        b, c, h, w = src.shape
+
+        # Run the transformer
+        hs, src = self.transformer(src, pos_src, tokens)  ####### hs - torch.Size([BS, 11, 256]), src - torch.Size([BS, 16348, 256])
+        mask_token_out =  hs[:, :, :]
+             
+        src = src.transpose(1, 2).view(b, c, h, w) ##### torch.Size([4, 256, 128, 128])
+        upscaled_embedding = self.output_upscaling(src)  ##### torch.Size([4, 32, 512, 512]) 
+        hyper_in = self.output_hypernetwork_mlps(mask_token_out) ##### torch.Size([1, 11, 32])
+        b, c, h, w = upscaled_embedding.shape
+        seg_output = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w)  ##### torch.Size([1, 11, 512, 512])
+        
+        return seg_output
+
+
+class PositionEmbeddingRandom(nn.Module):
+    """
+    Positional encoding using random spatial frequencies.
+    """
+
+    def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:
+        super().__init__()
+        if scale is None or scale <= 0.0:
+            scale = 1.0
+        self.register_buffer(
+            "positional_encoding_gaussian_matrix",
+            scale * torch.randn((2, num_pos_feats)),
+        )
+
+    def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
+        """Positionally encode points that are normalized to [0,1]."""
+        # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
+        coords = 2 * coords - 1
+        coords = coords @ self.positional_encoding_gaussian_matrix
+        coords = 2 * np.pi * coords
+        # outputs d_1 x ... x d_n x C shape
+        return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
+
+    def forward(self, size: Tuple[int, int]) -> torch.Tensor:
+        """Generate positional encoding for a grid of the specified size."""
+        h, w = size
+        device: Any = self.positional_encoding_gaussian_matrix.device
+        grid = torch.ones((h, w), device=device, dtype=torch.float32)
+        y_embed = grid.cumsum(dim=0) - 0.5
+        x_embed = grid.cumsum(dim=1) - 0.5
+        y_embed = y_embed / h
+        x_embed = x_embed / w
+
+        pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
+        return pe.permute(2, 0, 1)  # C x H x W
+
+    def forward_with_coords(
+        self, coords_input: torch.Tensor, image_size: Tuple[int, int]
+    ) -> torch.Tensor:
+        """Positionally encode points that are not normalized to [0,1]."""
+        coords = coords_input.clone()
+        coords[:, :, 0] = coords[:, :, 0] / image_size[1]
+        coords[:, :, 1] = coords[:, :, 1] / image_size[0]
+        return self._pe_encoding(coords.to(torch.float))  # B x N x C
+
+
+class SegfaceMLP(nn.Module):
+    """
+    Linear Embedding.
+    """
+
+    def __init__(self, input_dim):
+        super().__init__()
+        self.proj = nn.Linear(input_dim, 256)
+
+    def forward(self, hidden_states: torch.Tensor):
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        hidden_states = self.proj(hidden_states)
+        return hidden_states
+
+class SegFaceCeleb(nn.Module):
+    def __init__(self, input_resolution, model):
+        super(SegFaceCeleb, self).__init__()
+        self.input_resolution = input_resolution
+        self.model = model
+
+        if self.model == "swin_base":
+            swin_v2 = swin_b(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+        
+        if self.model == "swinv2_base":
+            swin_v2 = swin_v2_b(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "swinv2_small":
+            swin_v2 = swin_v2_s(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "swinv2_tiny":
+            swin_v2 = swin_v2_t(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "convnext_base":
+            convnext = convnext_base(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(convnext.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "convnext_small":
+            convnext = convnext_small(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(convnext.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "convnext_tiny":
+            convnext = convnext_small(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(convnext.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []           
+        
+        if self.model == "resnet":
+            resnet101 = models.resnet101(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(resnet101.children())[:-1]))
+            self.target_layer_names = ['4', '5', '6', '7']
+            self.multi_scale_features = []
+
+        if self.model == "mobilenet":
+            mobilenet = mobilenet_v3_large(pretrained=True).features
+            self.backbone = mobilenet
+            self.target_layer_names = ['3', '6', '12', '16']
+            self.multi_scale_features = []
+        
+        if self.model == "efficientnet":
+            efficientnet = efficientnet_v2_m(pretrained=True).features
+            self.backbone = efficientnet
+            self.target_layer_names = ['2', '3', '5', '8']
+            self.multi_scale_features = []
+
+        embed_dim = 1024
+        out_chans = 256
+        
+        self.pe_layer = PositionEmbeddingRandom(out_chans // 2)   
+
+        for name, module in self.backbone.named_modules():
+            if name in self.target_layer_names:
+                module.register_forward_hook(self.save_features_hook(name))
+        
+        self.face_decoder = FaceDecoder(
+            transformer_dim=256,
+            transformer=TwoWayTransformer(
+                depth=2,
+                embedding_dim=256,
+                mlp_dim=2048,
+                num_heads=8,
+            ))    
+        
+        num_encoder_blocks = 4
+        if self.model in ["swin_base", "swinv2_base", "convnext_base"]:
+            hidden_sizes = [128, 256, 512, 1024] ### Swin Base and ConvNext Base
+        if self.model in ["resnet"]:
+            hidden_sizes = [256, 512, 1024, 2048] ### ResNet
+        if self.model in ["swinv2_small", "swinv2_tiny", "convnext_small", "convnext_tiny"]:
+            hidden_sizes = [96, 192, 384, 768] ### Swin Small/Tiny and ConvNext Small/Tiny
+        if self.model in ["mobilenet"]:
+            hidden_sizes = [24, 40, 112, 960] ### MobileNet
+        if self.model in ["efficientnet"]:
+            hidden_sizes = [48, 80, 176, 1280] ### EfficientNet
+        decoder_hidden_size = 256
+        
+        mlps = []
+        for i in range(num_encoder_blocks):
+            mlp = SegfaceMLP(input_dim=hidden_sizes[i])
+            mlps.append(mlp)
+        self.linear_c = nn.ModuleList(mlps)
+
+        # The following 3 layers implement the ConvModule of the original implementation
+        self.linear_fuse = nn.Conv2d(
+            in_channels=decoder_hidden_size * num_encoder_blocks,
+            out_channels=decoder_hidden_size,
+            kernel_size=1,
+            bias=False,
+        )
+
+    
+    def save_features_hook(self, name):
+        def hook(module, input, output):
+            if self.model in ["swin_base", "swinv2_base", "swinv2_small", "swinv2_tiny"]:
+                self.multi_scale_features.append(output.permute(0,3,1,2).contiguous()) ### Swin, Swinv2
+            if self.model in ["convnext_base", "convnext_small", "convnext_tiny", "mobilenet", "efficientnet"]:
+                self.multi_scale_features.append(output) ### ConvNext, ResNet, EfficientNet, MobileNet
+        return hook
+
+    def forward(self, x, labels, dataset):
+        self.multi_scale_features.clear()
+        
+        _,_,h,w = x.shape
+        features = self.backbone(x).squeeze()
+        
+        batch_size = self.multi_scale_features[-1].shape[0]
+        all_hidden_states = ()
+        for encoder_hidden_state, mlp in zip(self.multi_scale_features, self.linear_c):
+            height, width = encoder_hidden_state.shape[2], encoder_hidden_state.shape[3]
+            encoder_hidden_state = mlp(encoder_hidden_state)
+            encoder_hidden_state = encoder_hidden_state.permute(0, 2, 1)
+            encoder_hidden_state = encoder_hidden_state.reshape(batch_size, -1, height, width)
+            # upsample
+            encoder_hidden_state = nn.functional.interpolate(
+                encoder_hidden_state, size=self.multi_scale_features[0].size()[2:], mode="bilinear", align_corners=False
+            )
+            all_hidden_states += (encoder_hidden_state,)
+        
+        fused_states = self.linear_fuse(torch.cat(all_hidden_states[::-1], dim=1)) #### torch.Size([BS, 256, 128, 128])
+        image_pe = self.pe_layer((fused_states.shape[2], fused_states.shape[3])).unsqueeze(0)
+        seg_output = self.face_decoder(
+                image_embeddings=fused_states,
+                image_pe=image_pe
+            )
+    
+        return seg_output
+
+if __name__ == "__main__":
+    input_resolution = 512
+    model_name = "swin_base"
+    model = SegFaceCeleb(input_resolution, model_name)
+    
+    batch_size = 4
+    num_channels = 3
+    height = 512
+    width = 512
+
+    x = torch.randn(batch_size, num_channels, height, width)
+    
+    labels = {
+        "lnm_seg": torch.randn(batch_size, 5, 2)
+    }
+    
+    dataset = torch.tensor([0,0,0,0])
+
+    seg_output = model(x, labels, dataset)
+    print("Segmentation Output Shape:", seg_output.shape)

models/segface/models/segface_helen.py

@@ -0,0 +1,347 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+from typing import Any, Optional, Tuple, Type
+from torchvision.models import convnext_large, convnext_base, convnext_small, convnext_tiny, swin_b, swin_v2_b, swin_v2_s, swin_v2_t, mobilenet_v3_large, efficientnet_v2_m
+import pdb
+import numpy as np
+import sys
+import os
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+from models.segface.models.transformer import *
+from models.segface.models.utils_models import *
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        num_layers: int,
+        sigmoid_output: bool = False,
+    ) -> None:
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(
+            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
+        )
+        self.sigmoid_output = sigmoid_output
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        if self.sigmoid_output:
+            x = F.sigmoid(x)
+        return x
+    
+class FaceDecoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        transformer_dim: 256,
+        transformer: nn.Module,
+        activation: Type[nn.Module] = nn.GELU,
+    ) -> None:
+        
+        super().__init__()
+        self.transformer_dim = transformer_dim
+        self.transformer = transformer
+
+        self.background_token = nn.Embedding(1, transformer_dim)
+        self.face_token = nn.Embedding(1, transformer_dim)
+        self.leftbro_token = nn.Embedding(1, transformer_dim)
+        self.rightbro_token = nn.Embedding(1, transformer_dim)
+        self.lefteye_token = nn.Embedding(1, transformer_dim)
+        self.righteye_token = nn.Embedding(1, transformer_dim)
+        self.nose_token = nn.Embedding(1, transformer_dim)
+        self.upperlip_token = nn.Embedding(1, transformer_dim)
+        self.innermouth_token = nn.Embedding(1, transformer_dim)
+        self.lowerlip_token = nn.Embedding(1, transformer_dim)
+        self.hair_token = nn.Embedding(1, transformer_dim)
+
+        self.output_upscaling = nn.Sequential(
+            nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2),
+            LayerNorm2d(transformer_dim // 4),
+            activation(),
+            nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
+            activation(),
+        )
+        
+        self.output_hypernetwork_mlps = MLP(
+            transformer_dim, transformer_dim, transformer_dim // 8, 3
+            )
+
+    def forward(
+        self,
+        image_embeddings: torch.Tensor,
+        image_pe: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        '''
+        image_embeddings - torch.Size([1, 256, 128, 128])
+        image_pe - torch.Size([1, 256, 128, 128])
+        '''
+        output_tokens = torch.cat([
+            self.background_token.weight, self.face_token.weight, self.leftbro_token.weight, self.rightbro_token.weight, 
+            self.lefteye_token.weight, self.righteye_token.weight, self.nose_token.weight, self.upperlip_token.weight, 
+            self.innermouth_token.weight, self.lowerlip_token.weight, self.hair_token.weight], dim=0)
+
+        tokens = output_tokens.unsqueeze(0).expand(image_embeddings.size(0), -1, -1) ##### torch.Size([4, 11, 256])
+
+        src = image_embeddings ##### torch.Size([4, 256, 128, 128])
+        pos_src = image_pe.expand(image_embeddings.size(0), -1, -1, -1)
+        b, c, h, w = src.shape
+
+        # Run the transformer
+        hs, src = self.transformer(src, pos_src, tokens)  ####### hs - torch.Size([BS, 11, 256]), src - torch.Size([BS, 16348, 256])
+        mask_token_out =  hs[:, :, :]
+             
+        src = src.transpose(1, 2).view(b, c, h, w) ##### torch.Size([4, 256, 128, 128])
+        upscaled_embedding = self.output_upscaling(src)  ##### torch.Size([4, 32, 512, 512]) 
+        hyper_in = self.output_hypernetwork_mlps(mask_token_out) ##### torch.Size([1, 11, 32])
+        b, c, h, w = upscaled_embedding.shape
+        seg_output = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w)  ##### torch.Size([1, 11, 512, 512])
+        
+        return seg_output
+
+
+
+class PositionEmbeddingRandom(nn.Module):
+    """
+    Positional encoding using random spatial frequencies.
+    """
+
+    def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:
+        super().__init__()
+        if scale is None or scale <= 0.0:
+            scale = 1.0
+        self.register_buffer(
+            "positional_encoding_gaussian_matrix",
+            scale * torch.randn((2, num_pos_feats)),
+        )
+
+    def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
+        """Positionally encode points that are normalized to [0,1]."""
+        # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
+        coords = 2 * coords - 1
+        coords = coords @ self.positional_encoding_gaussian_matrix
+        coords = 2 * np.pi * coords
+        # outputs d_1 x ... x d_n x C shape
+        return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
+
+    def forward(self, size: Tuple[int, int]) -> torch.Tensor:
+        """Generate positional encoding for a grid of the specified size."""
+        h, w = size
+        device: Any = self.positional_encoding_gaussian_matrix.device
+        grid = torch.ones((h, w), device=device, dtype=torch.float32)
+        y_embed = grid.cumsum(dim=0) - 0.5
+        x_embed = grid.cumsum(dim=1) - 0.5
+        y_embed = y_embed / h
+        x_embed = x_embed / w
+
+        pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
+        return pe.permute(2, 0, 1)  # C x H x W
+
+    def forward_with_coords(
+        self, coords_input: torch.Tensor, image_size: Tuple[int, int]
+    ) -> torch.Tensor:
+        """Positionally encode points that are not normalized to [0,1]."""
+        coords = coords_input.clone()
+        coords[:, :, 0] = coords[:, :, 0] / image_size[1]
+        coords[:, :, 1] = coords[:, :, 1] / image_size[0]
+        return self._pe_encoding(coords.to(torch.float))  # B x N x C
+
+
+class SegfaceMLP(nn.Module):
+    """
+    Linear Embedding.
+    """
+    def __init__(self, input_dim):
+        super().__init__()
+        self.proj = nn.Linear(input_dim, 256)
+
+    def forward(self, hidden_states: torch.Tensor):
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        hidden_states = self.proj(hidden_states)
+        return hidden_states
+
+class SegFaceHelen(nn.Module):
+    def __init__(self, input_resolution, model):
+        super(SegFaceHelen, self).__init__()
+        self.input_resolution = input_resolution
+        self.model = model
+
+        if self.model == "swin_base":
+            swin_v2 = swin_b(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+        
+        if self.model == "swinv2_base":
+            swin_v2 = swin_v2_b(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "swinv2_small":
+            swin_v2 = swin_v2_s(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "swinv2_tiny":
+            swin_v2 = swin_v2_t(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "convnext_base":
+            convnext = convnext_base(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(convnext.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "convnext_small":
+            convnext = convnext_small(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(convnext.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "convnext_tiny":
+            convnext = convnext_small(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(convnext.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []           
+        
+        if self.model == "resnet":
+            resnet101 = models.resnet101(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(resnet101.children())[:-1]))
+            self.target_layer_names = ['4', '5', '6', '7']
+            self.multi_scale_features = []
+
+        if self.model == "mobilenet":
+            mobilenet = mobilenet_v3_large(pretrained=True).features
+            self.backbone = mobilenet
+            self.target_layer_names = ['3', '6', '12', '16']
+            self.multi_scale_features = []
+        
+        if self.model == "efficientnet":
+            efficientnet = efficientnet_v2_m(pretrained=True).features
+            self.backbone = efficientnet
+            self.target_layer_names = ['2', '3', '5', '8']
+            self.multi_scale_features = []
+
+        embed_dim = 1024
+        out_chans = 256
+        
+        self.pe_layer = PositionEmbeddingRandom(out_chans // 2) 
+        self.get_matrix_fn = functools.partial(get_face_align_matrix, target_shape=(self.input_resolution, self.input_resolution), target_face_scale=1.0)
+        self.warp_fn = functools.partial(make_tanh_warp_grid, warp_factor=0.8, warped_shape=(self.input_resolution, self.input_resolution))
+        self.inv_warp_fn = functools.partial(make_inverted_tanh_warp_grid, warp_factor=0.8, warped_shape=(self.input_resolution, self.input_resolution))  
+
+        for name, module in self.backbone.named_modules():
+            if name in self.target_layer_names:
+                module.register_forward_hook(self.save_features_hook(name))
+        
+        self.face_decoder = FaceDecoder(
+            transformer_dim=256,
+            transformer=TwoWayTransformer(
+                depth=2,
+                embedding_dim=256,
+                mlp_dim=2048,
+                num_heads=8,
+            ))    
+        
+        num_encoder_blocks = 4
+        if self.model in ["swin_base", "swinv2_base", "convnext_base"]:
+            hidden_sizes = [128, 256, 512, 1024] ### Swin Base and ConvNext Base
+        if self.model in ["resnet"]:
+            hidden_sizes = [256, 512, 1024, 2048] ### ResNet
+        if self.model in ["swinv2_small", "swinv2_tiny", "convnext_small", "convnext_tiny"]:
+            hidden_sizes = [96, 192, 384, 768] ### Swin Small/Tiny and ConvNext Small/Tiny
+        if self.model in ["mobilenet"]:
+            hidden_sizes = [24, 40, 112, 960] ### MobileNet
+        if self.model in ["efficientnet"]:
+            hidden_sizes = [48, 80, 176, 1280] ### EfficientNet
+        decoder_hidden_size = 256
+        
+        mlps = []
+        for i in range(num_encoder_blocks):
+            mlp = SegfaceMLP(input_dim=hidden_sizes[i])
+            mlps.append(mlp)
+        self.linear_c = nn.ModuleList(mlps)
+
+        # The following 3 layers implement the ConvModule of the original implementation
+        self.linear_fuse = nn.Conv2d(
+            in_channels=decoder_hidden_size * num_encoder_blocks,
+            out_channels=decoder_hidden_size,
+            kernel_size=1,
+            bias=False,
+        )
+
+    
+    def save_features_hook(self, name):
+        def hook(module, input, output):
+            if self.model in ["swin_base", "swinv2_base", "swinv2_small", "swinv2_tiny"]:
+                self.multi_scale_features.append(output.permute(0,3,1,2).contiguous()) ### Swin, Swinv2
+            if self.model in ["convnext_base", "convnext_small", "convnext_tiny", "mobilenet", "efficientnet"]:
+                self.multi_scale_features.append(output) ### ConvNext, ResNet, EfficientNet, MobileNet
+        return hook
+
+    def forward(self, x, labels, dataset):
+        self.multi_scale_features.clear()
+        
+        _,_,h,w = x.shape
+        features = self.backbone(x).squeeze()
+        
+        batch_size = self.multi_scale_features[-1].shape[0]
+        all_hidden_states = ()
+        for encoder_hidden_state, mlp in zip(self.multi_scale_features, self.linear_c):
+            height, width = encoder_hidden_state.shape[2], encoder_hidden_state.shape[3]
+            encoder_hidden_state = mlp(encoder_hidden_state)
+            encoder_hidden_state = encoder_hidden_state.permute(0, 2, 1)
+            encoder_hidden_state = encoder_hidden_state.reshape(batch_size, -1, height, width)
+            # upsample
+            encoder_hidden_state = nn.functional.interpolate(
+                encoder_hidden_state, size=self.multi_scale_features[0].size()[2:], mode="bilinear", align_corners=False
+            )
+            all_hidden_states += (encoder_hidden_state,)
+        
+        fused_states = self.linear_fuse(torch.cat(all_hidden_states[::-1], dim=1)) #### torch.Size([BS, 256, 128, 128])
+        image_pe = self.pe_layer((fused_states.shape[2], fused_states.shape[3])).unsqueeze(0)
+        seg_output = self.face_decoder(
+                image_embeddings=fused_states,
+                image_pe=image_pe
+            )
+
+        fused_states = self.linear_fuse(torch.cat(all_hidden_states[::-1], dim=1)) #### torch.Size([bs, 256, 128, 128])
+        image_pe = self.pe_layer((fused_states.shape[2], fused_states.shape[3])).unsqueeze(0)
+        seg_output = self.face_decoder(
+                image_embeddings=fused_states,
+                image_pe=image_pe
+            )
+    
+        return seg_output
+
+if __name__ == "__main__":
+    model_name = "swin_base"
+    input_resolution = 512
+    model = SegFaceHelen(input_resolution, model_name)
+
+    batch_size = 4
+    num_channels = 3
+    height = input_resolution
+    width = input_resolution
+
+    x = torch.randn(batch_size, num_channels, height, width)
+    
+    labels = {
+        "lnm_seg": torch.randn(batch_size, 5, 2)
+    }
+    
+    dataset = torch.tensor([2,2,2,2])
+
+    seg_output = model(x, labels, dataset)
+    print("Segmentation Output Shape:", seg_output.shape)

models/segface/models/segface_lapa.py

@@ -0,0 +1,353 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+from typing import Any, Optional, Tuple, Type
+from torchvision.models import convnext_large, convnext_base, convnext_small, convnext_tiny, swin_b, swin_v2_b, swin_v2_s, swin_v2_t, mobilenet_v3_large, efficientnet_v2_m
+import pdb
+import numpy as np
+import sys
+import os
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+from models.segface.models.transformer import *
+from models.segface.models.utils_models import *
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        num_layers: int,
+        sigmoid_output: bool = False,
+    ) -> None:
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(
+            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
+        )
+        self.sigmoid_output = sigmoid_output
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        if self.sigmoid_output:
+            x = F.sigmoid(x)
+        return x
+    
+class FaceDecoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        transformer_dim: 256,
+        transformer: nn.Module,
+        activation: Type[nn.Module] = nn.GELU,
+    ) -> None:
+        
+        super().__init__()
+        self.transformer_dim = transformer_dim
+        self.transformer = transformer
+
+        self.background_token = nn.Embedding(1, transformer_dim)
+        self.face_token = nn.Embedding(1, transformer_dim)
+        self.leftbro_token = nn.Embedding(1, transformer_dim)
+        self.rightbro_token = nn.Embedding(1, transformer_dim)
+        self.lefteye_token = nn.Embedding(1, transformer_dim)
+        self.righteye_token = nn.Embedding(1, transformer_dim)
+        self.nose_token = nn.Embedding(1, transformer_dim)
+        self.upperlip_token = nn.Embedding(1, transformer_dim)
+        self.innermouth_token = nn.Embedding(1, transformer_dim)
+        self.lowerlip_token = nn.Embedding(1, transformer_dim)
+        self.hair_token = nn.Embedding(1, transformer_dim)
+
+        self.output_upscaling = nn.Sequential(
+            nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2),
+            LayerNorm2d(transformer_dim // 4),
+            activation(),
+            nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
+            activation(),
+        )
+        
+        self.output_hypernetwork_mlps = MLP(
+            transformer_dim, transformer_dim, transformer_dim // 8, 3
+            )
+
+    def forward(
+        self,
+        image_embeddings: torch.Tensor,
+        image_pe: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        '''
+        image_embeddings - torch.Size([1, 256, 128, 128])
+        image_pe - torch.Size([1, 256, 128, 128])
+        '''
+        output_tokens = torch.cat([self.background_token.weight, self.face_token.weight, self.leftbro_token.weight, self.rightbro_token.weight, self.lefteye_token.weight, \
+         self.righteye_token.weight, self.nose_token.weight, self.upperlip_token.weight, self.innermouth_token.weight,self.lowerlip_token.weight, self.hair_token.weight], dim=0) 
+
+        tokens = output_tokens.unsqueeze(0).expand(image_embeddings.size(0), -1, -1) ##### torch.Size([4, 11, 256])
+
+        src = image_embeddings ##### torch.Size([4, 256, 128, 128])
+        pos_src = image_pe.expand(image_embeddings.size(0), -1, -1, -1)
+        b, c, h, w = src.shape
+
+        # Run the transformer
+        hs, src = self.transformer(src, pos_src, tokens)  ####### hs - torch.Size([BS, 11, 256]), src - torch.Size([BS, 16348, 256])
+        mask_token_out =  hs[:, :, :]
+             
+        src = src.transpose(1, 2).view(b, c, h, w) ##### torch.Size([4, 256, 128, 128])
+        upscaled_embedding = self.output_upscaling(src)  ##### torch.Size([4, 32, 512, 512]) 
+        hyper_in = self.output_hypernetwork_mlps(mask_token_out) ##### torch.Size([1, 11, 32])
+        b, c, h, w = upscaled_embedding.shape
+        seg_output = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w)  ##### torch.Size([1, 11, 512, 512])
+        
+        return seg_output
+
+
+
+class PositionEmbeddingRandom(nn.Module):
+    """
+    Positional encoding using random spatial frequencies.
+    """
+
+    def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:
+        super().__init__()
+        if scale is None or scale <= 0.0:
+            scale = 1.0
+        self.register_buffer(
+            "positional_encoding_gaussian_matrix",
+            scale * torch.randn((2, num_pos_feats)),
+        )
+
+    def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
+        """Positionally encode points that are normalized to [0,1]."""
+        # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
+        coords = 2 * coords - 1
+        coords = coords @ self.positional_encoding_gaussian_matrix
+        coords = 2 * np.pi * coords
+        # outputs d_1 x ... x d_n x C shape
+        return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
+
+    def forward(self, size: Tuple[int, int]) -> torch.Tensor:
+        """Generate positional encoding for a grid of the specified size."""
+        h, w = size
+        device: Any = self.positional_encoding_gaussian_matrix.device
+        grid = torch.ones((h, w), device=device, dtype=torch.float32)
+        y_embed = grid.cumsum(dim=0) - 0.5
+        x_embed = grid.cumsum(dim=1) - 0.5
+        y_embed = y_embed / h
+        x_embed = x_embed / w
+
+        pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
+        return pe.permute(2, 0, 1)  # C x H x W
+
+    def forward_with_coords(
+        self, coords_input: torch.Tensor, image_size: Tuple[int, int]
+    ) -> torch.Tensor:
+        """Positionally encode points that are not normalized to [0,1]."""
+        coords = coords_input.clone()
+        coords[:, :, 0] = coords[:, :, 0] / image_size[1]
+        coords[:, :, 1] = coords[:, :, 1] / image_size[0]
+        return self._pe_encoding(coords.to(torch.float))  # B x N x C
+
+
+class SegfaceMLP(nn.Module):
+    """
+    Linear Embedding.
+    """
+
+    def __init__(self, input_dim):
+        super().__init__()
+        self.proj = nn.Linear(input_dim, 256)
+
+    def forward(self, hidden_states: torch.Tensor):
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        hidden_states = self.proj(hidden_states)
+        return hidden_states
+
+class SegFaceLapa(nn.Module):
+    def __init__(self, input_resolution, model):
+        super(SegFaceLapa, self).__init__()
+        self.input_resolution = input_resolution
+        self.model = model
+
+        if self.model == "swin_base":
+            swin_v2 = swin_b(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+        
+        if self.model == "swinv2_base":
+            swin_v2 = swin_v2_b(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "swinv2_small":
+            swin_v2 = swin_v2_s(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "swinv2_tiny":
+            swin_v2 = swin_v2_t(weights='IMAGENET1K_V1')
+            self.backbone = torch.nn.Sequential(*(list(swin_v2.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "convnext_base":
+            convnext = convnext_base(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(convnext.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "convnext_small":
+            convnext = convnext_small(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(convnext.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []
+
+        if self.model == "convnext_tiny":
+            convnext = convnext_small(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(convnext.children())[:-1]))
+            self.target_layer_names = ['0.1', '0.3', '0.5', '0.7']
+            self.multi_scale_features = []           
+        
+        if self.model == "resnet":
+            resnet101 = models.resnet101(pretrained=True)
+            self.backbone = torch.nn.Sequential(*(list(resnet101.children())[:-1]))
+            self.target_layer_names = ['4', '5', '6', '7']
+            self.multi_scale_features = []
+
+        if self.model == "mobilenet":
+            mobilenet = mobilenet_v3_large(pretrained=True).features
+            self.backbone = mobilenet
+            self.target_layer_names = ['3', '6', '12', '16']
+            self.multi_scale_features = []
+        
+        if self.model == "efficientnet":
+            efficientnet = efficientnet_v2_m(pretrained=True).features
+            self.backbone = efficientnet
+            self.target_layer_names = ['2', '3', '5', '8']
+            self.multi_scale_features = []
+
+        embed_dim = 1024
+        out_chans = 256
+        
+        self.pe_layer = PositionEmbeddingRandom(out_chans // 2)   
+        self.get_matrix_fn = functools.partial(get_face_align_matrix, target_shape=(self.input_resolution, self.input_resolution), target_face_scale=1.0)
+        self.warp_fn = functools.partial(make_tanh_warp_grid, warp_factor=0.8, warped_shape=(self.input_resolution, self.input_resolution))
+        self.inv_warp_fn = functools.partial(make_inverted_tanh_warp_grid, warp_factor=0.8, warped_shape=(self.input_resolution, self.input_resolution))
+
+        for name, module in self.backbone.named_modules():
+            if name in self.target_layer_names:
+                module.register_forward_hook(self.save_features_hook(name))
+        
+        self.face_decoder = FaceDecoder(
+            transformer_dim=256,
+            transformer=TwoWayTransformer(
+                depth=2,
+                embedding_dim=256,
+                mlp_dim=2048,
+                num_heads=8,
+            ))    
+        
+        num_encoder_blocks = 4
+        if self.model in ["swin_base", "swinv2_base", "convnext_base"]:
+            hidden_sizes = [128, 256, 512, 1024] ### Swin Base and ConvNext Base
+        if self.model in ["resnet"]:
+            hidden_sizes = [256, 512, 1024, 2048] ### ResNet
+        if self.model in ["swinv2_small", "swinv2_tiny", "convnext_small", "convnext_tiny"]:
+            hidden_sizes = [96, 192, 384, 768] ### Swin Small/Tiny and ConvNext Small/Tiny
+        if self.model in ["mobilenet"]:
+            hidden_sizes = [24, 40, 112, 960] ### MobileNet
+        if self.model in ["efficientnet"]:
+            hidden_sizes = [48, 80, 176, 1280] ### EfficientNet
+        decoder_hidden_size = 256
+        
+        mlps = []
+        for i in range(num_encoder_blocks):
+            mlp = SegfaceMLP(input_dim=hidden_sizes[i])
+            mlps.append(mlp)
+        self.linear_c = nn.ModuleList(mlps)
+
+        # The following 3 layers implement the ConvModule of the original implementation
+        self.linear_fuse = nn.Conv2d(
+            in_channels=decoder_hidden_size * num_encoder_blocks,
+            out_channels=decoder_hidden_size,
+            kernel_size=1,
+            bias=False,
+        )
+
+    
+    def save_features_hook(self, name):
+        def hook(module, input, output):
+            if self.model in ["swin_base", "swinv2_base", "swinv2_small", "swinv2_tiny"]:
+                self.multi_scale_features.append(output.permute(0,3,1,2).contiguous()) ### Swin, Swinv2
+            if self.model in ["convnext_base", "convnext_small", "convnext_tiny", "mobilenet", "efficientnet"]:
+                self.multi_scale_features.append(output) ### ConvNext, ResNet, EfficientNet, MobileNet
+        return hook
+
+    def forward(self, x, labels, dataset):
+        self.multi_scale_features.clear()
+        
+        _,_,h,w = x.shape
+        mask = dataset == 1
+        x_seg = x[mask]  #### torch.Size([4, 3, 512, 512])
+        lnd = labels["lnm_seg"][mask]  #### torch.Size([4, 5, 2])
+        matrix = self.get_matrix_fn(lnd)  #### torch.Size([4, 3, 3])
+        grid = self.warp_fn(matrix=matrix,  orig_shape = (h, w))  #### torch.Size([4, 512, 512, 2])
+        inv_grid = self.inv_warp_fn(matrix=matrix, orig_shape = (h, w)) #### torch.Size([4, 512, 512, 2])
+        w_x_seg = F.grid_sample(x_seg, grid, mode='bilinear', align_corners=False) #### torch.Size([4, 3, 512, 512])
+        x[mask] = w_x_seg
+        features = self.backbone(x).squeeze()
+        
+        batch_size = self.multi_scale_features[-1].shape[0]
+        all_hidden_states = ()
+        for encoder_hidden_state, mlp in zip(self.multi_scale_features, self.linear_c):
+            height, width = encoder_hidden_state.shape[2], encoder_hidden_state.shape[3]
+            encoder_hidden_state = mlp(encoder_hidden_state)
+            encoder_hidden_state = encoder_hidden_state.permute(0, 2, 1)
+            encoder_hidden_state = encoder_hidden_state.reshape(batch_size, -1, height, width)
+            # upsample
+            encoder_hidden_state = nn.functional.interpolate(
+                encoder_hidden_state, size=self.multi_scale_features[0].size()[2:], mode="bilinear", align_corners=False
+            )
+            all_hidden_states += (encoder_hidden_state,)
+        
+        fused_states = self.linear_fuse(torch.cat(all_hidden_states[::-1], dim=1)) #### torch.Size([bs, 256, 128, 128])
+        image_pe = self.pe_layer((fused_states.shape[2], fused_states.shape[3])).unsqueeze(0)
+        seg_output = self.face_decoder(
+                image_embeddings=fused_states,
+                image_pe=image_pe
+            )
+
+        segmentation_indices = dataset == 1
+        seg_lapa_in = seg_output[mask]
+        seg_lapa = F.grid_sample(seg_lapa_in, inv_grid, mode='bilinear', align_corners=False)
+        seg_output[mask] = seg_lapa
+        seg_output = seg_output[segmentation_indices]
+    
+        return seg_output
+
+if __name__ == "__main__":
+    input_resolution = 512
+    model_name = "swin_base"
+    model = SegFaceLapa(input_resolution, model_name)
+    
+    batch_size = 4
+    num_channels = 3
+    height = input_resolution
+    width = input_resolution
+
+    x = torch.randn(batch_size, num_channels, height, width)
+    
+    labels = {
+        "lnm_seg": torch.randn(batch_size, 5, 2)
+    }
+
+    dataset = torch.tensor([0,0,1,1])
+
+    seg_output = model(x, labels, dataset)
+    print("Segmentation Output Shape:", seg_output.shape)

models/segface/models/transformer.py

@@ -0,0 +1,271 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from torch import Tensor, nn
+
+import math
+from typing import Tuple, Type
+
+
+class MLPBlock(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        mlp_dim: int,
+        act: Type[nn.Module] = nn.GELU,
+    ) -> None:
+        super().__init__()
+        self.lin1 = nn.Linear(embedding_dim, mlp_dim)
+        self.lin2 = nn.Linear(mlp_dim, embedding_dim)
+        self.act = act()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lin2(self.act(self.lin1(x)))
+
+
+# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
+# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119  # noqa
+class LayerNorm2d(nn.Module):
+    def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(num_channels))
+        self.bias = nn.Parameter(torch.zeros(num_channels))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        x = self.weight[:, None, None] * x + self.bias[:, None, None]
+        return x
+
+
+class TwoWayTransformer(nn.Module):
+    def __init__(
+        self,
+        depth: int,
+        embedding_dim: int,
+        num_heads: int,
+        mlp_dim: int,
+        activation: Type[nn.Module] = nn.ReLU,
+        attention_downsample_rate: int = 2,
+    ) -> None:
+        """
+        A transformer decoder that attends to an input image using
+        queries whose positional embedding is supplied.
+
+        Args:
+          depth (int): number of layers in the transformer
+          embedding_dim (int): the channel dimension for the input embeddings
+          num_heads (int): the number of heads for multihead attention. Must
+            divide embedding_dim
+          mlp_dim (int): the channel dimension internal to the MLP block
+          activation (nn.Module): the activation to use in the MLP block
+        """
+        super().__init__()
+        self.depth = depth
+        self.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.mlp_dim = mlp_dim
+        self.layers = nn.ModuleList()
+
+        for i in range(depth):
+            self.layers.append(
+                TwoWayAttentionBlock(
+                    embedding_dim=embedding_dim,
+                    num_heads=num_heads,
+                    mlp_dim=mlp_dim,
+                    activation=activation,
+                    attention_downsample_rate=attention_downsample_rate,
+                    skip_first_layer_pe=(i == 0),
+                )
+            )
+
+        self.final_attn_token_to_image = Attention(
+            embedding_dim, num_heads, downsample_rate=attention_downsample_rate
+        )
+        self.norm_final_attn = nn.LayerNorm(embedding_dim)
+
+    def forward(
+        self,
+        image_embedding: Tensor,
+        image_pe: Tensor,
+        point_embedding: Tensor,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Args:
+          image_embedding (torch.Tensor): image to attend to. Should be shape
+            B x embedding_dim x h x w for any h and w.
+          image_pe (torch.Tensor): the positional encoding to add to the image. Must
+            have the same shape as image_embedding.
+          point_embedding (torch.Tensor): the embedding to add to the query points.
+            Must have shape B x N_points x embedding_dim for any N_points.
+
+        Returns:
+          torch.Tensor: the processed point_embedding
+          torch.Tensor: the processed image_embedding
+        """
+        # BxCxHxW -> BxHWxC == B x N_image_tokens x C
+        bs, c, h, w = image_embedding.shape
+        image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
+        image_pe = image_pe.flatten(2).permute(0, 2, 1)
+
+        # Prepare queries
+        queries = point_embedding
+        keys = image_embedding
+
+        # Apply transformer blocks and final layernorm
+        for layer in self.layers:
+            queries, keys = layer(
+                queries=queries,
+                keys=keys,
+                query_pe=point_embedding,
+                key_pe=image_pe,
+            )
+
+        # Apply the final attention layer from the points to the image
+        q = queries + point_embedding
+        k = keys + image_pe
+        attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
+        queries = queries + attn_out
+        queries = self.norm_final_attn(queries)
+
+        return queries, keys
+
+
+class TwoWayAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_heads: int,
+        mlp_dim: int = 2048,
+        activation: Type[nn.Module] = nn.ReLU,
+        attention_downsample_rate: int = 2,
+        skip_first_layer_pe: bool = False,
+    ) -> None:
+        """
+        A transformer block with four layers: (1) self-attention of sparse
+        inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp
+        block on sparse inputs, and (4) cross attention of dense inputs to sparse
+        inputs.
+
+        Arguments:
+          embedding_dim (int): the channel dimension of the embeddings
+          num_heads (int): the number of heads in the attention layers
+          mlp_dim (int): the hidden dimension of the mlp block
+          activation (nn.Module): the activation of the mlp block
+          skip_first_layer_pe (bool): skip the PE on the first layer
+        """
+        super().__init__()
+        self.self_attn = Attention(embedding_dim, num_heads)
+        self.norm1 = nn.LayerNorm(embedding_dim)
+
+        self.cross_attn_token_to_image = Attention(
+            embedding_dim, num_heads, downsample_rate=attention_downsample_rate
+        )
+        self.norm2 = nn.LayerNorm(embedding_dim)
+
+        self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
+        self.norm3 = nn.LayerNorm(embedding_dim)
+
+        self.norm4 = nn.LayerNorm(embedding_dim)
+        self.cross_attn_image_to_token = Attention(
+            embedding_dim, num_heads, downsample_rate=attention_downsample_rate
+        )
+
+        self.skip_first_layer_pe = skip_first_layer_pe
+
+    def forward(
+        self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor
+    ) -> Tuple[Tensor, Tensor]:
+        # Self attention block
+        if self.skip_first_layer_pe:
+            queries = self.self_attn(q=queries, k=queries, v=queries)
+        else:
+            q = queries + query_pe
+            attn_out = self.self_attn(q=q, k=q, v=queries)
+            queries = queries + attn_out
+        queries = self.norm1(queries)
+
+        # Cross attention block, tokens attending to image embedding
+        q = queries + query_pe
+        k = keys + key_pe
+        attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
+        queries = queries + attn_out
+        queries = self.norm2(queries)
+
+        # MLP block
+        mlp_out = self.mlp(queries)
+        queries = queries + mlp_out
+        queries = self.norm3(queries)
+
+        # Cross attention block, image embedding attending to tokens
+        q = queries + query_pe
+        k = keys + key_pe
+        attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)
+        keys = keys + attn_out
+        keys = self.norm4(keys)
+
+        return queries, keys
+
+
+class Attention(nn.Module):
+    """
+    An attention layer that allows for downscaling the size of the embedding
+    after projection to queries, keys, and values.
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_heads: int,
+        downsample_rate: int = 1,
+    ) -> None:
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.internal_dim = embedding_dim // downsample_rate
+        self.num_heads = num_heads
+        assert self.internal_dim % num_heads == 0, "num_heads must divide embedding_dim."
+
+        self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
+        self.k_proj = nn.Linear(embedding_dim, self.internal_dim)
+        self.v_proj = nn.Linear(embedding_dim, self.internal_dim)
+        self.out_proj = nn.Linear(self.internal_dim, embedding_dim)
+
+    def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:
+        b, n, c = x.shape
+        x = x.reshape(b, n, num_heads, c // num_heads)
+        return x.transpose(1, 2)  # B x N_heads x N_tokens x C_per_head
+
+    def _recombine_heads(self, x: Tensor) -> Tensor:
+        b, n_heads, n_tokens, c_per_head = x.shape
+        x = x.transpose(1, 2)
+        return x.reshape(b, n_tokens, n_heads * c_per_head)  # B x N_tokens x C
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        # Input projections
+        q = self.q_proj(q)
+        k = self.k_proj(k)
+        v = self.v_proj(v)
+
+        # Separate into heads
+        q = self._separate_heads(q, self.num_heads)
+        k = self._separate_heads(k, self.num_heads)
+        v = self._separate_heads(v, self.num_heads)
+
+        # Attention
+        _, _, _, c_per_head = q.shape
+        attn = q @ k.permute(0, 1, 3, 2)  # B x N_heads x N_tokens x N_tokens
+        attn = attn / math.sqrt(c_per_head)
+        attn = torch.softmax(attn, dim=-1)
+
+        # Get output
+        out = attn @ v
+        out = self._recombine_heads(out)
+        out = self.out_proj(out)
+
+        return out

models/segface/models/utils_models.py

@@ -0,0 +1,270 @@
+from typing import List, Dict, Callable, Tuple, Optional
+import torch
+import torch.nn.functional as F
+import functools
+import numpy as np
+
+@functools.lru_cache(maxsize=128)
+def _meshgrid(h, w) -> Tuple[torch.Tensor, torch.Tensor]:
+    yy, xx = torch.meshgrid(torch.arange(h).float(),
+                            torch.arange(w).float(),
+                            indexing='ij')
+    return yy, xx
+
+
+def _forge_grid(batch_size: int, device: torch.device,
+                output_shape: Tuple[int, int],
+                fn: Callable[[torch.Tensor], torch.Tensor]
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+    """ Forge transform maps with a given function `fn`.
+
+    Args:
+        output_shape (tuple): (b, h, w, ...).
+        fn (Callable[[torch.Tensor], torch.Tensor]): The function that accepts 
+            a bxnx2 array and outputs the transformed bxnx2 array. Both input 
+            and output store (x, y) coordinates.
+
+    Note: 
+        both input and output arrays of `fn` should store (y, x) coordinates.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Two maps `X` and `Y`, where for each 
+            pixel (y, x) or coordinate (x, y),
+            `(X[y, x], Y[y, x]) = fn([x, y])`
+    """
+    h, w, *_ = output_shape
+    yy, xx = _meshgrid(h, w)  # h x w
+    yy = yy.unsqueeze(0).broadcast_to(batch_size, h, w).to(device)
+    xx = xx.unsqueeze(0).broadcast_to(batch_size, h, w).to(device)
+
+    in_xxyy = torch.stack(
+        [xx, yy], dim=-1).reshape([batch_size, h*w, 2])  # (h x w) x 2
+    out_xxyy: torch.Tensor = fn(in_xxyy)  # (h x w) x 2
+    return out_xxyy.reshape(batch_size, h, w, 2)
+
+def inverted_tanh_warp_transform(coords: torch.Tensor, matrix: torch.Tensor,
+                                 warp_factor: float, warped_shape: Tuple[int, int]):
+    """ Inverted tanh-warp function.
+
+    Args:
+        coords (torch.Tensor): b x n x 2 (x, y). The transformed coordinates.
+        matrix: b x 3 x 3. A matrix that transforms un-normalized coordinates 
+            from the original image to the aligned yet not-warped image.
+        warp_factor (float): The warp factor. 
+            0 means linear transform, 1 means full tanh warp.
+        warped_shape (tuple): [height, width].
+
+    Returns:
+        torch.Tensor: b x n x 2 (x, y). The original coordinates.
+    """
+    h, w, *_ = warped_shape
+    # h -= 1
+    # w -= 1
+
+    w_h = torch.tensor([[w, h]]).to(coords)
+
+    if warp_factor > 0:
+        # normalize coordinates to [-1, +1]
+        coords = coords / w_h * 2 - 1
+
+        nl_part1 = coords > 1.0 - warp_factor
+        nl_part2 = coords < -1.0 + warp_factor
+
+        ret_nl_part1 = _safe_arctanh(
+            (coords - 1.0 + warp_factor) /
+            warp_factor) * warp_factor + \
+            1.0 - warp_factor
+        ret_nl_part2 = _safe_arctanh(
+            (coords + 1.0 - warp_factor) /
+            warp_factor) * warp_factor - \
+            1.0 + warp_factor
+
+        coords = torch.where(nl_part1, ret_nl_part1,
+                             torch.where(nl_part2, ret_nl_part2, coords))
+
+        # denormalize
+        coords = (coords + 1) / 2 * w_h
+
+    coords_homo = torch.cat(
+        [coords, torch.ones_like(coords[:, :, [0]])], dim=-1)  # b x n x 3
+
+    inv_matrix = torch.linalg.inv(matrix)  # b x 3 x 3
+    # inv_matrix = np.linalg.inv(matrix)
+    coords_homo = torch.bmm(
+        coords_homo, inv_matrix.permute(0, 2, 1))  # b x n x 3
+    return coords_homo[:, :, :2] / coords_homo[:, :, [2, 2]]
+
+
+def tanh_warp_transform(
+        coords: torch.Tensor, matrix: torch.Tensor,
+        warp_factor: float, warped_shape: Tuple[int, int]):
+    """ Tanh-warp function.
+
+    Args:
+        coords (torch.Tensor): b x n x 2 (x, y). The original coordinates.
+        matrix: b x 3 x 3. A matrix that transforms un-normalized coordinates 
+            from the original image to the aligned yet not-warped image.
+        warp_factor (float): The warp factor. 
+            0 means linear transform, 1 means full tanh warp.
+        warped_shape (tuple): [height, width].
+
+    Returns:
+        torch.Tensor: b x n x 2 (x, y). The transformed coordinates.
+    """
+    h, w, *_ = warped_shape
+    # h -= 1
+    # w -= 1
+    w_h = torch.tensor([[w, h]]).to(coords)
+
+    coords_homo = torch.cat(
+        [coords, torch.ones_like(coords[:, :, [0]])], dim=-1)  # b x n x 3
+
+    coords_homo = torch.bmm(coords_homo, matrix.transpose(2, 1))  # b x n x 3
+    coords = (coords_homo[:, :, :2] / coords_homo[:, :, [2, 2]])  # b x n x 2
+
+    if warp_factor > 0:
+        # normalize coordinates to [-1, +1]
+        coords = coords / w_h * 2 - 1
+
+        nl_part1 = coords > 1.0 - warp_factor
+        nl_part2 = coords < -1.0 + warp_factor
+
+        ret_nl_part1 = torch.tanh(
+            (coords - 1.0 + warp_factor) /
+            warp_factor) * warp_factor + \
+            1.0 - warp_factor
+        ret_nl_part2 = torch.tanh(
+            (coords + 1.0 - warp_factor) /
+            warp_factor) * warp_factor - \
+            1.0 + warp_factor
+
+        coords = torch.where(nl_part1, ret_nl_part1,
+                             torch.where(nl_part2, ret_nl_part2, coords))
+
+        # denormalize
+        coords = (coords + 1) / 2 * w_h
+
+    return coords
+
+def make_tanh_warp_grid(matrix: torch.Tensor, warp_factor: float,
+                        warped_shape: Tuple[int, int],
+                        orig_shape: Tuple[int, int]):
+    """
+    Args:
+        matrix: bx3x3 matrix.
+        warp_factor: The warping factor. `warp_factor=1.0` represents a vannila Tanh-warping, 
+           `warp_factor=0.0` represents a cropping.
+        warped_shape: The target image shape to transform to.
+
+    Returns:
+        torch.Tensor: b x h x w x 2 (x, y).
+    """
+    orig_h, orig_w, *_ = orig_shape
+    w_h = torch.tensor([orig_w, orig_h]).to(matrix).reshape(1, 1, 1, 2)
+    return _forge_grid(
+        matrix.size(0), matrix.device,
+        warped_shape,
+        functools.partial(inverted_tanh_warp_transform,
+                          matrix=matrix,
+                          warp_factor=warp_factor,
+                          warped_shape=warped_shape)) / w_h*2-1
+
+
+def make_inverted_tanh_warp_grid(matrix: torch.Tensor, warp_factor: float,
+                                 warped_shape: Tuple[int, int],
+                                 orig_shape: Tuple[int, int]):
+    """
+    Args:
+        matrix: bx3x3 matrix.
+        warp_factor: The warping factor. `warp_factor=1.0` represents a vannila Tanh-warping, 
+           `warp_factor=0.0` represents a cropping.
+        warped_shape: The target image shape to transform to.
+        orig_shape: The original image shape that is transformed from.
+
+    Returns:
+        torch.Tensor: b x h x w x 2 (x, y).
+    """
+    h, w, *_ = warped_shape
+    w_h = torch.tensor([w, h]).to(matrix).reshape(1, 1, 1, 2)
+    return _forge_grid(
+        matrix.size(0), matrix.device,
+        orig_shape,
+        functools.partial(tanh_warp_transform,
+                          matrix=matrix,
+                          warp_factor=warp_factor,
+                          warped_shape=warped_shape)) / w_h * 2-1
+
+def _safe_arctanh(x: torch.Tensor, eps: float = 0.001) -> torch.Tensor:
+    return torch.clamp(x, -1+eps, 1-eps).arctanh()
+
+def get_similarity_transform_matrix(
+        from_pts: torch.Tensor, to_pts: torch.Tensor) -> torch.Tensor:
+    """
+    Args:
+        from_pts, to_pts: b x n x 2
+
+    Returns:
+        torch.Tensor: b x 3 x 3
+    """
+    mfrom = from_pts.mean(dim=1, keepdim=True)  # b x 1 x 2
+    mto = to_pts.mean(dim=1, keepdim=True)  # b x 1 x 2
+
+    a1 = (from_pts - mfrom).square().sum([1, 2], keepdim=False)  # b
+    c1 = ((to_pts - mto) * (from_pts - mfrom)).sum([1, 2], keepdim=False)  # b
+
+    to_delta = to_pts - mto
+    from_delta = from_pts - mfrom
+    c2 = (to_delta[:, :, 0] * from_delta[:, :, 1] - to_delta[:,
+          :, 1] * from_delta[:, :, 0]).sum([1], keepdim=False)  # b
+
+    a = c1 / a1
+    b = c2 / a1
+    dx = mto[:, 0, 0] - a * mfrom[:, 0, 0] - b * mfrom[:, 0, 1]  # b
+    dy = mto[:, 0, 1] + b * mfrom[:, 0, 0] - a * mfrom[:, 0, 1]  # b
+
+    ones_pl = torch.ones_like(a1)
+    zeros_pl = torch.zeros_like(a1)
+
+    return torch.stack([
+        a, b, dx,
+        -b, a, dy,
+        zeros_pl, zeros_pl, ones_pl,
+    ], dim=-1).reshape(-1, 3, 3)
+
+
+@functools.lru_cache()
+def _standard_face_pts():
+    pts = torch.tensor([
+        196.0, 226.0,
+        316.0, 226.0,
+        256.0, 286.0,
+        220.0, 360.4,
+        292.0, 360.4], dtype=torch.float32) / 256.0 - 1.0
+    return torch.reshape(pts, (5, 2))
+
+
+def get_face_align_matrix(
+        face_pts: torch.Tensor, target_shape: Tuple[int, int],
+        target_face_scale: float = 1.0, offset_xy: Optional[Tuple[float, float]] = None,
+        target_pts: Optional[torch.Tensor] = None):
+
+    if target_pts is None:
+        with torch.no_grad():
+            std_pts = _standard_face_pts().to(face_pts)  # [-1 1]
+            h, w, *_ = target_shape
+            target_pts = (std_pts * target_face_scale + 1) * \
+                torch.tensor([w-1, h-1]).to(face_pts) / 2.0
+            if offset_xy is not None:
+                target_pts[:, 0] += offset_xy[0]
+                target_pts[:, 1] += offset_xy[1]
+    else:
+        target_pts = target_pts.to(face_pts)
+
+    if target_pts.dim() == 2:
+        target_pts = target_pts.unsqueeze(0)
+    if target_pts.size(0) == 1:
+        target_pts = target_pts.broadcast_to(face_pts.shape)
+
+    assert target_pts.shape == face_pts.shape
+
+    return get_similarity_transform_matrix(face_pts, target_pts)

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.4
+torchvision>=0.19
+numpy>=1.26
+Pillow>=10.0
+huggingface_hub>=0.30

training_run_summary.json

@@ -0,0 +1,47 @@
+{
+  "run_dir": "hair_mask_dataset/runs/segface_hair_budget_4090",
+  "model_name": "swin_base",
+  "prepared_root": "/workspace/runpod_upload_ready/data/aihub_hairmask_hq_budget_50k",
+  "raw_root": "/workspace/runpod_upload_ready/data/aihub_korean_hairstyle_hq_raw",
+  "epochs_completed": 10,
+  "best_epoch": 7,
+  "best_val_iou": 0.9486894006725745,
+  "best_val_dice": 0.9735556454363521,
+  "best_val_precision": 0.9723250788834037,
+  "best_val_recall": 0.9751487422222148,
+  "last_epoch": 10,
+  "last_train_loss": 0.028126267597526313,
+  "last_val_loss": 0.028694584750384094,
+  "last_val_iou": 0.9486362328742441,
+  "last_val_dice": 0.9735264129781702,
+  "last_val_precision": 0.9721686440444964,
+  "last_val_recall": 0.9752568952148782,
+  "avg_epoch_sec": 3546.4520416259766,
+  "train_count": 50000,
+  "val_count": 5000,
+  "test_count": 0,
+  "checkpoint_files": [
+    "best.pt",
+    "epoch_001.pt",
+    "epoch_002.pt",
+    "epoch_003.pt",
+    "epoch_004.pt",
+    "epoch_005.pt",
+    "epoch_006.pt",
+    "epoch_007.pt",
+    "epoch_008.pt",
+    "epoch_009.pt",
+    "epoch_010.pt",
+    "last.pt"
+  ],
+  "plot_path": "hair_mask_dataset/runs/segface_hair_budget_4090/plots/training_curves.png",
+  "latest_preview_path": "hair_mask_dataset/runs/segface_hair_budget_4090/previews/epoch_010.png",
+  "submit_date": "2026-03-17",
+  "github_url": "https://github.com/skn-ai22-251029/SKN22-Final-1Team-AI",
+  "team_members": [
+    "이병재",
+    "장완식",
+    "최정환",
+    "문승준"
+  ]
+}