You don’t need to change the image on tap!

Buttons can change their own images when their state changes, so you only need to change the button state:

button.setImage(UIImage(systemName: "square"), for: .normal)
button.setImage(UIImage(systemName: "checkmark.square.fill"), for: .selected)
button.addAction(UIAction { _ in button.isSelected.toggle() }, for: .touchUpInside)

You can also setTitle() and it will all work out-of the box, though you may want to adjust the button.configuration to make it look nice.

For a full (yet short) implementation, see my gist:
CheckboxButton.swift