RealityView loading a USDZ model and visualizing the UV coordinates with a LowLevelTexture

ModelUVTextureRealityView.swift

import SwiftUI
import RealityKit
import MetalKit

struct ModelUVTextureRealityView: View {
    @State var entity: ModelEntity?
    @State var texture: LowLevelTexture?
    let device: MTLDevice
    let commandQueue: MTLCommandQueue
    let textureComputePipeline: MTLComputePipelineState
    
    init() {
        self.device = MTLCreateSystemDefaultDevice()!
        self.commandQueue = device.makeCommandQueue()!
        let library = device.makeDefaultLibrary()!
        let updateFunction = library.makeFunction(name: "uvVisualizationKernel")!
        self.textureComputePipeline = try! device.makeComputePipelineState(function: updateFunction)
    }
    
    var body: some View {
        RealityView { content in
            let model = try! await loadModelEntity()
            self.entity = model
            var material = model.model?.materials.first as! PhysicallyBasedMaterial
            let texture = try! LowLevelTexture(descriptor: textureDescriptor)
            let textureResource = try! await TextureResource(from: texture)
            material.baseColor.texture = .init(textureResource)
            updateTexture(texture)
            model.model?.materials = [material]
            content.add(model)
            self.texture = texture
        }
    }
    
    var textureDescriptor: LowLevelTexture.Descriptor {
        var descriptor = LowLevelTexture.Descriptor()
        descriptor.textureType = .type2D
        descriptor.pixelFormat = .rgba16Float
        descriptor.width = 2048
        descriptor.height = 2048
        descriptor.mipmapLevelCount = 1
        descriptor.textureUsage = [.shaderRead, .shaderWrite]
        return descriptor
    }
    
    func updateTexture(_ texture: LowLevelTexture) {
        let commandBuffer = commandQueue.makeCommandBuffer()!
        let computeEncoder = commandBuffer.makeComputeCommandEncoder()!

        computeEncoder.setComputePipelineState(textureComputePipeline)
        computeEncoder.setTexture(texture.replace(using: commandBuffer), index: 0)

        let threadGroupSize = MTLSizeMake(8, 8, 1)
        let threadGroups = MTLSizeMake(
            (texture.descriptor.width + threadGroupSize.width - 1) / threadGroupSize.width,
            (texture.descriptor.height + threadGroupSize.height - 1) / threadGroupSize.height,
            1
        )
        
        computeEncoder.dispatchThreadgroups(threadGroups, threadsPerThreadgroup: threadGroupSize)
        computeEncoder.endEncoding()
        
        commandBuffer.commit()
    }
    
    func loadModelEntity(url: URL = URL(string: "https://matt54.github.io/Resources/Anubis_Statue_1.usdz")!) async throws -> ModelEntity {
        let (downloadedURL, _) = try await URLSession.shared.download(from: url)
        let documentsDirectory = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first!
        let destinationURL = documentsDirectory.appendingPathComponent("downloadedModel.usdz")
        if FileManager.default.fileExists(atPath: destinationURL.path) {
            try FileManager.default.removeItem(at: destinationURL)
        }
        try FileManager.default.moveItem(at: downloadedURL, to: destinationURL)
        let entity = try await ModelEntity.init(contentsOf: destinationURL)
        try FileManager.default.removeItem(at: destinationURL)
        return entity
    }
}

#Preview {
    ModelUVTextureRealityView()
}

uvVisualizationKernel.metal

#include <metal_stdlib>
using namespace metal;

kernel void uvVisualizationKernel(texture2d<float, access::write> outTexture [[texture(0)]],
                                  uint2 gid [[thread_position_in_grid]])
{
    float2 uv = float2(gid) / float2(outTexture.get_width(), outTexture.get_height());
    
    float4 color;
    color.r = uv.x;
    color.g = uv.y;
    color.b = 0.5;
    color.a = 1.0;
    
    outTexture.write(color, gid);
}