Is it possible to convert an image to image TensorFlow model to CoreML?

The objective is to convert a TensorFlow model to CoreML using both input and output of type image.

I'm using Python 3.8 and TensorFlow 2.4.1
The model have input and output shape (1, 128, 128, 3)
Pixels are floats between -1 and 1.

I'm using this to create the CoreML module:

import coremltools as ct
from coremltools.proto import FeatureTypes_pb2 as ft

input_name = generator.inputs[0].name.split(':')[0]
image_input = ct.ImageType(name=input_name, shape=(1, 128, 128, 3), scale=2 / 255.0, bias=[-1.0, -1.0, -1.0])
mlModel = ct.convert(generator, inputs=[image_input], source="tensorflow")

# Set Output as Image
spec = mlModel.get_spec()
output = spec.description.output[0]
output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('RGB')
output.type.imageType.width = 128
output.type.imageType.height = 128

# Save
new_model = ct.models.MLModel(spec)
new_model.author = '...'
new_model.license = '...'
new_model.short_description = '...'
new_model.input_description[input_name] = 'Input Image'
new_model.output_description[output.name] = 'Predicted Image'
new_model.save("model.mlmodel")

print("Input name: ", input_name)  # -> input_1
print("Output name: ", output.name)  # -> Identity

I can test predictions like this (the result is correct):

# image shape=(1, 128, 128, 3), dtype=float32)
prediction = generator(image, training=True)
plt.imshow(prediction[0] * 0.5 + 0.5)
plt.show()

The generated model looks like this on Xcode: image to image model

This is what I'm using on IOS:

func generateImage(inputImage: UIImage, completion: @escaping (UIImage?)->()) {
    guard let model = try? VNCoreMLModel(for: Model(configuration: .init()).model) else {
        completion(nil)
        return
    }
    guard let cgImage = inputImage.cgImage else {
        completion(nil)
        return
    }

    let coreMLRequest = VNCoreMLRequest(model: model) { (request, error) in
        guard let results = request.results as? [VNPixelBufferObservation] else {
            print("Can't get results: \(request.results) Error: \(error?.localizedDescription)")
            completion(nil)
            return
        }
        guard let topResult = results.first else {
            print("Can't get top result")
            completion(nil)
            return
        }
        let newImage = UIImage(ciImage: CIImage(cvPixelBuffer: topResult.pixelBuffer))
        completion(newImage)
        return
    }
    do {
        try VNImageRequestHandler(cgImage: cgImage, options: [:]).perform([coreMLRequest])
    } catch {
        print("Failed to perform request. Error: \(error.localizedDescription)")
    }
}

This fails with the error: The VNCoreMLTransform request failed.
And trying the preview option directly form Xcode also fails.

So I believe the input or the output is not properly setup during conversion.
Is it possible to convert from (1, 128, 128, 3) float32 to whatever CoreML expects?