Burn Fade Animation for USDZ Models in RealityKit using LowLevelMesh, ShaderGraphMaterial & Metal Compute Shader

BurnFadeModelView.swift

import Metal
import RealityKit
import SwiftUI

struct BurnFadeModelView: View {
    @State var lowLevelMesh: LowLevelMesh?
    @State var timer: Timer?
    @State var isForward: Bool = true
    @State var morphProgress: Float = 0.0
    @State var dwellCounter: Int = 0
    @State var isDwelling: Bool = false
    @State var burnScale: Float = 8.0
    @State var hueRotate: Float = 0
    @State var edgeWidth: Float = 0.08
    @State var emberRange: Float = 0.15

    // Store original vertices to preserve colors
    @State var originalVertices: [VertexDataWith4ChannelColor] = []
    
    let modelURL = URL(string: "https://matt54.github.io/Resources/StatueOfBuddha.usdz")!
    
    let timerUpdateDuration: TimeInterval = 1/120.0
    let dwellDuration: Int = 30
    var morphProgressRate: Float = 0.005
    
    let device: MTLDevice
    let commandQueue: MTLCommandQueue
    let computePipelineState: MTLComputePipelineState

    init() {
        self.device = MTLCreateSystemDefaultDevice()!
        self.commandQueue = device.makeCommandQueue()!
        let library = device.makeDefaultLibrary()!
        let kernelFunction = library.makeFunction(name: "burnFadeVertices")!
        self.computePipelineState = try! device.makeComputePipelineState(function: kernelFunction)
    }
    
    var body: some View {
        VStack {
            RealityView { content in
                let model = try! await loadModelEntity(url: modelURL)
                content.add(model)
                let lowLevelMesh = try! createMesh(from: model)
                
                // swap out model mesh with our LowLevelMesh
                model.model?.mesh = try! await MeshResource(from: lowLevelMesh)
                
                let originalMaterial = model.model?.materials.first as! PhysicallyBasedMaterial
                let originalTexture = originalMaterial.baseColor.texture!.resource

                var shaderMaterial = try! await loadMaterial()
                
                try! shaderMaterial.setParameter(name: "ImageTexture", value: .textureResource(originalTexture))
                
                model.model?.materials = [shaderMaterial]
                
                self.lowLevelMesh = lowLevelMesh
            }
            
            VStack {
                HStack {
                    Text("Burn Scale: \(burnScale, specifier: "%.2f")")
                    Spacer()
                    Slider(value: $burnScale, in: 1...40)
                        .frame(width: 300)
                }
                
                HStack {
                    Text("Burn Hue: \(hueRotate, specifier: "%.2f")")
                    Spacer()
                    Slider(value: $hueRotate, in: 0...(.pi*2))
                        .frame(width: 300)
                }
                HStack {
                    Text("Edge Width: \(edgeWidth, specifier: "%.3f")")
                    Spacer()
                    Slider(value: $edgeWidth, in: 0.01...0.2)
                        .frame(width: 300)
                }
                HStack {
                    Text("Ember Range: \(emberRange, specifier: "%.3f")")
                    Spacer()
                    Slider(value: $emberRange, in: 0.05...0.4)
                        .frame(width: 300)
                }
            }
            .frame(width: 500)
            .padding()
            .glassBackgroundEffect()
        }
        .onAppear { startTimer() }
        .onDisappear { stopTimer() }
    }
    
    enum MeshCreationError: Error {
        case modelNotFound, meshPartNotFound
    }
}

// MARK: Animation Loop
extension BurnFadeModelView {
    func startTimer() {
        timer = Timer.scheduledTimer(withTimeInterval: timerUpdateDuration, repeats: true) { timer in
            if isDwelling {
                // Count dwell time
                dwellCounter += 1
                if dwellCounter >= dwellDuration {
                    // Finished dwelling, switch direction and resume morphing
                    isDwelling = false
                    dwellCounter = 0
                    isForward.toggle()
                }
            } else {
                // Update Morph Progress
                if isForward {
                    morphProgress += morphProgressRate
                } else {
                    morphProgress -= morphProgressRate
                }
                
                // Handle bounds
                if morphProgress >= 1.0 {
                    morphProgress = 1.0
                    isDwelling = true
                    dwellCounter = 0
                } else if morphProgress <= 0.0 {
                    morphProgress = 0.0
                    isDwelling = true
                    dwellCounter = 0
                }
            }
            
            updateMesh()
        }
    }
    
    func stopTimer() {
        timer?.invalidate()
        timer = nil
    }
}

// MARK: Get Online Shader Graph Material
extension BurnFadeModelView {
    func loadMaterial() async throws -> ShaderGraphMaterial {
        let baseURL = URL(string: "https://matt54.github.io/Resources/")!
        let fullURL = baseURL.appendingPathComponent("ImageTexturePlus4ChannelColor.usda")
        let data = try Data(contentsOf: fullURL)
        let materialFilenameWithPath: String = "/Root/ImageTexturePlus4ChannelColorMaterial"
        return try await ShaderGraphMaterial(named: materialFilenameWithPath, from: data)
    }
}

// MARK: Download model
extension BurnFadeModelView {
    func loadModelEntity(url: URL) 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
    }
}

// MARK: Mesh functions
extension BurnFadeModelView {
    func createMesh(from modelEntity: ModelEntity) throws -> LowLevelMesh {
        guard let model = modelEntity.model
        else { throw MeshCreationError.modelNotFound }
        
        guard let meshPart = model.mesh.contents.models.first?.parts.first
        else { throw MeshCreationError.meshPartNotFound}
        
        let positions = meshPart[MeshBuffers.positions]?.elements ?? []
        let normals = meshPart[MeshBuffers.normals]?.elements ?? []
        let textureCoordinates = meshPart[MeshBuffers.textureCoordinates]?.elements ?? []
        let triangleIndices = meshPart.triangleIndices?.elements ?? []
        
        let lowLevelMesh = try VertexDataWith4ChannelColor.initializeMesh(vertexCapacity: positions.count,
                                                                          indexCapacity: triangleIndices.count)
        
        // Copy vertex data
        lowLevelMesh.withUnsafeMutableBytes(bufferIndex: 0) { rawBytes in
            let vertexBuffer = rawBytes.bindMemory(to: VertexDataWith4ChannelColor.self)
            for i in 0..<positions.count {
                vertexBuffer[i] = VertexDataWith4ChannelColor(position: positions[i], normal: normals[i], uv: textureCoordinates[i], color: SIMD4(x: 0, y: 0, z: 0, w: 1.0))
                originalVertices.append(vertexBuffer[i])
            }
        }
        
        // Copy index data
        lowLevelMesh.withUnsafeMutableIndices { buffer in
            let indices = buffer.bindMemory(to: UInt32.self)
            for (index, triangleIndex) in triangleIndices.enumerated() {
                indices[index] = UInt32(triangleIndex)
            }
        }
        
        let bounds = model.mesh.bounds
        lowLevelMesh.parts.replaceAll([
            LowLevelMesh.Part(
                indexCount: triangleIndices.count,
                topology: .triangle,
                bounds: bounds
            )
        ])
        
        return lowLevelMesh
    }
    
    func updateMesh() {
        guard let mesh = lowLevelMesh,
              let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder() else { return }
        
        let newVertexBuffer = mesh.replace(bufferIndex: 0, using: commandBuffer)
        
        computeEncoder.setComputePipelineState(computePipelineState)
        let originalVertexBuffer = device.makeBuffer(
            bytes: originalVertices,
            length: originalVertices.count * MemoryLayout<VertexDataWith4ChannelColor>.size,
            options: .storageModeShared
        )
        computeEncoder.setBuffer(originalVertexBuffer, offset: 0, index: 0)     // Read from original
        computeEncoder.setBuffer(newVertexBuffer, offset: 0, index: 1)          // Write to new
        var params = BurnFadeParams(
            progress: morphProgress,
            scale: burnScale,
            hueRotate: hueRotate,
            edgeWidth: edgeWidth,
            emberRange: emberRange
        );
        computeEncoder.setBytes(&params, length: MemoryLayout<BurnFadeParams>.size, index: 2)

        let vertexCount = mesh.vertexCapacity
        let threadsPerGrid = MTLSize(width: vertexCount, height: 1, depth: 1)
        let threadsPerThreadgroup = MTLSize(width: 64, height: 1, depth: 1)
        
        computeEncoder.dispatchThreads(threadsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)
        computeEncoder.endEncoding()
        commandBuffer.commit()
    }
}

#Preview {
    BurnFadeModelView()
}

BurnFadeParams.h

#ifndef BurnFadeParams_h
#define BurnFadeParams_h

struct BurnFadeParams {
    float progress;
    float scale;
    float hueRotate;
    float edgeWidth;
    float emberRange;
};

#endif /* BurnFadeParams_h */

burnFadeVertices.metal

#include <metal_stdlib>
using namespace metal;

#include "HueRotate.metal"
#include "NoiseUtilities.metal"
#include "VertexDataWith4ChannelColor.h"
#include "BurnFadeParams.h"

kernel void burnFadeVertices(device const VertexDataWith4ChannelColor* inputVertices [[buffer(0)]],
                             device VertexDataWith4ChannelColor* outputVertices [[buffer(1)]],
                             constant BurnFadeParams& params [[buffer(2)]],
                             uint vid [[thread_position_in_grid]])
{
    VertexDataWith4ChannelColor inputVertex = inputVertices[vid];
    VertexDataWith4ChannelColor outputVertex = inputVertex;

    float3 scaledPosition = inputVertex.position * params.scale;
    
    // Create noise pattern that determines burn priority
    float primaryNoise = fractalNoise(scaledPosition);
    float secondaryNoise = fractalNoise(scaledPosition * 2.3 + float3(100.0, 50.0, 75.0));
    float detailNoise = noise(scaledPosition * 12.0);
    
    // Combine noise layers - this determines which areas burn first/last
    float burnPriority = primaryNoise * 0.6 + secondaryNoise * 0.3 + detailNoise * 0.1;
    
    // Map burnAmount to affect the full range
    float effectiveBurnAmount = params.progress;
    
    // Calculate how far this vertex is from being burned away
    float burnDistance = burnPriority - effectiveBurnAmount;
    
    // Create smooth transition for alpha
    float edgeWidth = params.edgeWidth;
    float alpha = smoothstep(-edgeWidth, edgeWidth, burnDistance);
    
    // Calculate ember/fire effect at burn edges
    float emberIntensity = 0.0;
    float emberRange = params.emberRange; // How far from the burn edge we show embers
    
    if (burnDistance > -emberRange && burnDistance < emberRange) {
        // We're near the burn edge - calculate ember intensity
        float distanceFromEdge = abs(burnDistance);
        emberIntensity = 1.0 - (distanceFromEdge / emberRange);
        emberIntensity = smoothstep(0.0, 1.0, emberIntensity);
        
        // Add flickering with noise
        float emberFlicker = noise(scaledPosition * 25.0 + float3(params.progress * 10.0, 0.0, 0.0));
        emberIntensity *= 0.7 + 0.3 * emberFlicker;
    }
    
    // Define fire colors
    float3 originalColor = float3(inputVertex.color.x, inputVertex.color.y, inputVertex.color.z);
    float3 emberColor = float3(1.0, 0.4, 0.1);  // Bright orange
    float3 fireColor = float3(1.0, 0.2, 0.0);   // Red-orange
    float3 hotColor = float3(1.0, 0.8, 0.2);    // Yellow-white hot
    
    // Mix between different fire colors based on intensity
    float3 burnColor;
    if (emberIntensity < 0.3) {
        burnColor = mix(emberColor, fireColor, emberIntensity / 0.3);
    } else if (emberIntensity < 0.7) {
        burnColor = mix(fireColor, hotColor, (emberIntensity - 0.3) / 0.4);
    } else {
        burnColor = hotColor;
    }
    
    burnColor = rotateHue(burnColor, params.hueRotate);
    
    // Blend original color with burn effect
    float3 finalColor = mix(originalColor, burnColor, emberIntensity);
    
    // Make burn edges extra bright
    finalColor += emberIntensity * emberIntensity * float3(0.2, 0.2, 0.2);
    
    alpha = clamp(alpha, 0.0, 1.0);
    
    outputVertex.color.rgb = finalColor;
    outputVertex.color.w = alpha;
    outputVertices[vid] = outputVertex;
}

HueRotate.metal

#include <metal_stdlib>
using namespace metal;

inline float3 rotateHue(float3 color, float hueRotation) {
    const float3 k = float3(0.57735, 0.57735, 0.57735);
    float cosAngle = cos(hueRotation);
    return color * cosAngle + cross(k, color) * sin(hueRotation) + k * dot(k, color) * (1.0 - cosAngle);
}

NoiseUtilities.metal

#include <metal_stdlib>
using namespace metal;

inline float hash(float3 p) {
    p = fract(p * 0.3183099 + 0.1);
    p *= 17.0;
    return fract(p.x * p.y * p.z * (p.x + p.y + p.z));
}

inline float noise(float3 p) {
    float3 i = floor(p);
    float3 f = fract(p);
    f = f * f * (3.0 - 2.0 * f);
    
    return mix(mix(mix(hash(i + float3(0,0,0)), hash(i + float3(1,0,0)), f.x),
                   mix(hash(i + float3(0,1,0)), hash(i + float3(1,1,0)), f.x), f.y),
               mix(mix(hash(i + float3(0,0,1)), hash(i + float3(1,0,1)), f.x),
                   mix(hash(i + float3(0,1,1)), hash(i + float3(1,1,1)), f.x), f.y), f.z);
}

inline float fractalNoise(float3 p) {
    float value = 0.0;
    float amplitude = 0.5;
    float frequency = 1.0;
    float maxValue = 0.0;
    
    for (int i = 0; i < 4; i++) {
        value += amplitude * noise(p * frequency);
        maxValue += amplitude;
        amplitude *= 0.5;
        frequency *= 2.0;
    }
    
    return value / maxValue;
}

VertexDataWith4ChannelColor+Extensions.swift

import Foundation
import RealityKit

extension VertexDataWith4ChannelColor {
    static var vertexAttributes: [LowLevelMesh.Attribute] = [
        .init(semantic: .position, format: .float3, offset: MemoryLayout<Self>.offset(of: \.position)!),
        .init(semantic: .normal, format: .float3, offset: MemoryLayout<Self>.offset(of: \.normal)!),
        .init(semantic: .uv0, format: .float2, offset: MemoryLayout<Self>.offset(of: \.uv)!),
        .init(semantic: .uv2, format: .float4, offset: MemoryLayout<Self>.offset(of: \.color)!)
    ]
    
    static var vertexLayouts: [LowLevelMesh.Layout] = [
        .init(bufferIndex: 0, bufferStride: MemoryLayout<Self>.stride)
    ]

    static var descriptor: LowLevelMesh.Descriptor {
        var desc = LowLevelMesh.Descriptor()
        desc.vertexAttributes = VertexDataWith4ChannelColor.vertexAttributes
        desc.vertexLayouts = VertexDataWith4ChannelColor.vertexLayouts
        desc.indexType = .uint32
        return desc
    }
    
    @MainActor static func initializeMesh(vertexCapacity: Int,
                                          indexCapacity: Int) throws -> LowLevelMesh {
        var desc = VertexDataWith4ChannelColor.descriptor
        desc.vertexCapacity = vertexCapacity
        desc.indexCapacity = indexCapacity
        return try LowLevelMesh(descriptor: desc)
    }
}

VertexDataWith4ChannelColor.h

#include <simd/simd.h>

#ifndef VertexDataWith4ChannelColor_h
#define VertexDataWith4ChannelColor_h

struct VertexDataWith4ChannelColor {
    simd_float3 position;
    simd_float3 normal;
    simd_float2 uv;
    simd_float4 color; // rgb + alpha
};

#endif /* VertexDataWith4ChannelColor_h */