Repository: apple/ml-stable-diffusion
File: swift/StableDiffusion/pipeline/Scheduler.swift (Swift)

Code before:
    ///   - timeStep: The current time step in the diffusion chain
    ///   - sample: The current input sample to the diffusion model
    /// - Returns: Predicted de-noised sample at the previous time step
    /// - Postcondition: The scheduler state is updated.
    ///   The state holds the current sample and history of model output noise residuals
    func step(
        output: MLShapedArray<Float32>,
        timeStep t: Int,
        sample s: MLShapedArray<Float32>
    ) -> MLShapedArray<Float32>
}

public extension Scheduler {
    var initNoiseSigma: Float { 1 }
}

public extension Scheduler {
    /// Compute weighted sum of shaped arrays of equal shapes
    ///
    /// - Parameters:
    ///   - weights: The weights each array is multiplied by
    ///   - values: The arrays to be weighted and summed
    /// - Returns: sum_i weights[i]*values[i]
    func weightedSum(_ weights: [Double], _ values: [MLShapedArray<Float32>]) -> MLShapedArray<Float32> {
        assert(weights.count > 1 && values.count == weights.count)
        assert(values.allSatisfy({$0.scalarCount == values.first!.scalarCount}))
        var w = Float(weights.first!)
        var scalars = values.first!.scalars.map({ $0 * w })
        for next in 1 ..< values.count {
            w = Float(weights[next])
            let nextScalars = values[next].scalars
            for i in 0 ..< scalars.count {
                scalars[i] += w * nextScalars[i]
            }
        }
        return MLShapedArray(scalars: scalars, shape: values.first!.shape)
    }
}

/// How to map a beta range to a sequence of betas to step over
public enum BetaSchedule {
    /// Linear stepping between start and end
    case linear
    /// Steps using linspace(sqrt(start),sqrt(end))^2
    case scaledLinear
}


/// A scheduler used to compute a de-noised image
///
///  This implementation matches:

Review comment (suggestion):
```suggestion
        assert(values.allSatisfy({ $0.scalarCount == values.first!.scalarCount }))
```

Code after:
    ///   - timeStep: The current time step in the diffusion chain
    ///   - sample: The current input sample to the diffusion model
    /// - Returns: Predicted de-noised sample at the previous time step
    /// - Postcondition: The scheduler state is updated.
    ///   The state holds the current sample and history of model output noise residuals
    func step(
        output: MLShapedArray<Float32>,
        timeStep t: Int,
        sample s: MLShapedArray<Float32>
    ) -> MLShapedArray<Float32>
}

@available(iOS 16.2, macOS 13.1, *)
public extension Scheduler {
    var initNoiseSigma: Float { 1 }
}

@available(iOS 16.2, macOS 13.1, *)
public extension Scheduler {
    /// Compute weighted sum of shaped arrays of equal shapes
    ///
    /// - Parameters:
    ///   - weights: The weights each array is multiplied by
    ///   - values: The arrays to be weighted and summed
    /// - Returns: sum_i weights[i]*values[i]
    func weightedSum(_ weights: [Double], _ values: [MLShapedArray<Float32>]) -> MLShapedArray<Float32> {
        assert(weights.count > 1 && values.count == weights.count)
        assert(values.allSatisfy({ $0.scalarCount == values.first!.scalarCount }))
        var w = Float(weights.first!)
        var scalars = values.first!.scalars.map({ $0 * w })
        for next in 1 ..< values.count {
            w = Float(weights[next])
            let nextScalars = values[next].scalars
            for i in 0 ..< scalars.count {
                scalars[i] += w * nextScalars[i]
            }
        }
        return MLShapedArray(scalars: scalars, shape: values.first!.shape)
    }
}

/// How to map a beta range to a sequence of betas to step over
@available(iOS 16.2, macOS 13.1, *)
public enum BetaSchedule {
    /// Linear stepping between start and end
    case linear
    /// Steps using linspace(sqrt(start),sqrt(end))^2
    case scaledLinear
}