from transformers import TFMobileViTForSemanticSegmentation
import tensorflow as tf

if __name__ == '__main__':
    model = TFMobileViTForSemanticSegmentation.from_pretrained(".")
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.target_spec.supported_ops = [
        tf.lite.OpsSet.TFLITE_BUILTINS,
        tf.lite.OpsSet.SELECT_TF_OPS,
    ]
    tflite_model = converter.convert()
    tflite_filename = "tflite_model.tflite"
    with open(tflite_filename, "wb") as f:
        f.write(tflite_model)