"""Staged freeze/thaw training helpers."""
import torch.nn as nn

def freeze(m):
    for p in m.parameters(): p.requires_grad = False

def unfreeze(m):
    for p in m.parameters(): p.requires_grad = True

def count_trainable(m):
    return sum(p.numel() for p in m.parameters() if p.requires_grad)

def setup_stage1(model):
    unfreeze(model)
    return model

def setup_stage2(model):
    freeze(model)
    for name, p in model.named_parameters():
        if "xattn" in name or "tproj" in name:
            p.requires_grad = True
    return model

def setup_stage3(model):
    unfreeze(model)
    return model