Added support for Flux Schnell Image generation using Fal.ai (#110)

Author: Ronuhz

Sources/AIProxy/Fal/FalFluxSchnellInputSchema.swift

@@ -0,0 +1,85 @@
+//
+//  FalFluxSchnellInputSchema.swift
+//
+//
+//  Created by Hunor Zoltáni on 01.03.2025.
+//
+
+import Foundation
+
+/// Docstrings taken from the tooltips here: https://fal.ai/models/fal-ai/flux/schnell
+public struct FalFluxSchnellInputSchema: Encodable {
+    // Required
+
+    /// The prompt to generate an image from.
+    public let prompt: String
+
+    // Optional
+
+    /// If set to true, the safety checker will be enabled. Default value: true
+    public let enableSafetyChecker: Bool?
+
+    /// The size of the generated image.
+    /// Default value: `.landscape4x3`
+    public let imageSize: ImageSize?
+
+    /// The number of images to generate.
+    /// Default value: `1`
+    public let numImages: Int?
+
+    /// The number of inference steps to perform.
+    /// Default value: `4`
+    public let numInferenceSteps: Int?
+
+    /// The same seed and the same prompt given to the same version of the model will output the same image every time.
+    public let seed: Int?
+
+    /// If set to true, the function will wait for the image to be generated and uploaded
+    /// before returning the response. This will increase the latency of the function but it
+    /// allows you to get the image directly in the response without going through the CDN.
+    public let syncMode: Bool?
+
+    private enum CodingKeys: String, CodingKey {
+        case enableSafetyChecker = "enable_safety_checker"
+        case imageSize = "image_size"
+        case numImages = "num_images"
+        case numInferenceSteps = "num_inference_steps"
+        case prompt
+        case seed
+        case syncMode = "sync_mode"
+    }
+
+    // This memberwise initializer is autogenerated.
+    // To regenerate, use `cmd-shift-a` > Generate Memberwise Initializer
+    // To format, place the cursor in the initializer's parameter list and use `ctrl-m`
+    public init(
+        prompt: String,
+        enableSafetyChecker: Bool? = nil,
+        imageSize: FalFluxSchnellInputSchema.ImageSize? = nil,
+        numImages: Int? = nil,
+        numInferenceSteps: Int? = nil,
+        seed: Int? = nil,
+        syncMode: Bool? = nil
+    ) {
+        self.prompt = prompt
+        self.enableSafetyChecker = enableSafetyChecker
+        self.imageSize = imageSize
+        self.numImages = numImages
+        self.numInferenceSteps = numInferenceSteps
+        self.seed = seed
+        self.syncMode = syncMode
+    }
+
+}
+
+// MARK: - InputSchema.ImageSize
+extension FalFluxSchnellInputSchema {
+    public enum ImageSize: String, Encodable {
+        case landscape16x9 = "landscape_16_9"
+        case landscape4x3 = "landscape_4_3"
+        case portrait16x9 = "portrait_16_9"
+        case portrait4x3 = "portrait_4_3"
+        case square
+        case squareHD = "square_hd"
+    }
+}

Sources/AIProxy/Fal/FalFluxSchnellOutputSchema.swift

@@ -0,0 +1,24 @@
+//
+//  FalFluxSchnellOutputSchema.swift
+//
+//
+//  Created by Hunor Zoltáni on 01.03.2025.
+//
+
+import Foundation
+
+public struct FalFluxSchnellOutputSchema: Decodable {
+    public let hasNSFWConcepts: [Bool]?
+    public let images: [FalOutputImage]?
+    public let prompt: String?
+    public let seed: UInt64?
+    public let timings: FalTimings?
+
+    private enum CodingKeys: String, CodingKey {
+        case hasNSFWConcepts = "has_nsfw_concepts"
+        case images
+        case prompt
+        case seed
+        case timings
+    }
+}

Sources/AIProxy/Fal/FalService+Convenience.swift

@@ -8,6 +8,25 @@
 import Foundation
 
 extension FalService {
+    /// Convenience method for creating a `fal-ai/flux/schnell` image.
+    ///
+    /// - Parameter input: The input schema. See `FalFluxSchnellInputSchema.swift` for the range of controls that you
+    ///                    can use to adjust the image generation.
+    ///
+    /// - Returns: The inference result. The `images` property of the returned value contains a list of
+    ///            generated images. Each image has a `url` that you can use to fetch the image contents
+    ///            (or use with AsyncImage)
+    public func createFluxSchnellImage(
+        input: FalFluxSchnellInputSchema
+    ) async throws -> FalFluxSchnellOutputSchema {
+        return try await self.createInferenceAndPollForResult(
+            model: "fal-ai/flux/schnell",
+            input: input,
+            pollAttempts: 60,
+            secondsBetweenPollAttempts: 2
+        )
+    }
+    
     /// Convenience method for creating a `fal-ai/fast-sdxl` image.
     ///
     /// - Parameter input: The input schema. See `FalFastSDXLInputSchema.swift` for the range of controls that you

Tests/AIProxyTests/FalFluxSchnellResponseTests.swift

@@ -0,0 +1,41 @@
+//
+//  FalFluxSchnellResponseTests.swift
+//
+//
+//  Created by Hunor Zoltáni on 01.03.2025.
+//
+
+import XCTest
+import Foundation
+@testable import AIProxy
+
+final class FalFluxSchnellResponseTests: XCTestCase {
+
+    func testResponseIsDecodable() throws {
+        let sampleResponse = """
+            {
+              "images": [
+                {
+                  "url": "https://fal.media/files/zebra/_9NDmaWO5okNq9idPY7Il.jpeg",
+                  "width": 1024,
+                  "height": 1024,
+                  "content_type": "image/jpeg"
+                }
+              ],
+              "timings": {
+                "inference": 2.1141920797526836
+              },
+              "seed": 2062765390712234200,
+              "has_nsfw_concepts": [
+                false
+              ],
+              "prompt": "winter wonderland"
+            }
+        """
+        let res = try FalFluxSchnellOutputSchema.deserialize(from: sampleResponse)
+        XCTAssertEqual(
+            "https://fal.media/files/zebra/_9NDmaWO5okNq9idPY7Il.jpeg",
+            res.images?.first?.url?.absoluteString
+        )
+    }
+}