drywall-qa-clipseg / src /model /clipseg_wrapper.py
youngPhilosopher's picture
Upload folder using huggingface_hub
b891e61 verified
"""CLIPSeg model loading and freezing utilities."""
from transformers import CLIPSegForImageSegmentation, CLIPSegProcessor
def load_model_and_processor(model_name: str = "CIDAS/clipseg-rd64-refined", freeze_backbone: bool = True):
"""Load CLIPSeg model and processor, optionally freezing the backbone."""
model = CLIPSegForImageSegmentation.from_pretrained(model_name)
processor = CLIPSegProcessor.from_pretrained(model_name)
if freeze_backbone:
trainable, frozen = 0, 0
for name, param in model.named_parameters():
if "decoder" in name:
param.requires_grad = True
trainable += param.numel()
else:
param.requires_grad = False
frozen += param.numel()
print(f"Parameters — trainable (decoder): {trainable:,} | frozen (backbone): {frozen:,}")
else:
trainable = sum(p.numel() for p in model.parameters())
print(f"Parameters — all trainable: {trainable:,}")
return model, processor