diff --git a/python/src/coreai_models/vlm/export.py b/python/src/coreai_models/vlm/export.py index e3713dc3..d2c7bb73 100644 --- a/python/src/coreai_models/vlm/export.py +++ b/python/src/coreai_models/vlm/export.py @@ -412,8 +412,13 @@ class StaticVisionEncoder(nn.Module): produces) and reproduces the Qwen image-processor patchify internally, so the runner needs no Qwen-specific preprocessing beyond resize + normalize. - Input: pixel_values float32 [1, 3, image_size, image_size] (CLIP-normalized, NCHW) - Output: image_features float32 [num_visual_tokens, text_hidden] + Single-image (num_frames=1): + Input: pixel_values float32 [1, 3, image_size, image_size] + Output: image_features float32 [num_visual_tokens, text_hidden] + + Multi-frame (num_frames>1, must be divisible by temporal_patch_size): + Input: pixel_values float32 [1, 3*num_frames, image_size, image_size] + Output: image_features float32 [grid_t * num_visual_tokens, text_hidden] """ def __init__( @@ -424,6 +429,7 @@ def __init__( patch_size: int, spatial_merge_size: int, temporal_patch_size: int, + num_frames: int = 1, ) -> None: super().__init__() self.patch_embed = visual_model.patch_embed @@ -435,49 +441,53 @@ def __init__( self.spatial_merge_size = spatial_merge_size self.temporal_patch_size = temporal_patch_size self.channels = 3 - # Fixed grid for image_size x image_size. - self.grid_t = 1 + self.num_frames = num_frames + + if num_frames % temporal_patch_size != 0: + raise ValueError( + f"num_frames ({num_frames}) must be divisible by " + f"temporal_patch_size ({temporal_patch_size})" + ) + self.grid_t = num_frames // temporal_patch_size self.grid_h = image_size // patch_size self.grid_w = image_size // patch_size - self.num_patches = self.grid_h * self.grid_w + self.num_patches = self.grid_t * self.grid_h * self.grid_w self.patch_dim = temporal_patch_size * self.channels * patch_size * patch_size grid_thw = torch.tensor([[self.grid_t, self.grid_h, self.grid_w]], dtype=torch.int32) with torch.no_grad(): - # Position embeddings [num_patches, vision_hidden] pos_embeds = visual_model.fast_pos_embed_interpolate(grid_thw) self.register_buffer("pos_embeds", pos_embeds) - # Rotary position embeddings - rotary_pos_emb = visual_model.rot_pos_emb(grid_thw) # [num_patches, rot_dim/2] + rotary_pos_emb = visual_model.rot_pos_emb(grid_thw) seq_len = rotary_pos_emb.shape[0] rotary_flat = rotary_pos_emb.reshape(seq_len, -1) - emb = torch.cat([rotary_flat, rotary_flat], dim=-1) # [num_patches, rot_dim] + emb = torch.cat([rotary_flat, rotary_flat], dim=-1) self.register_buffer("rot_cos", emb.cos()) self.register_buffer("rot_sin", emb.sin()) - # cu_seqlens for variable-length attention: [0, num_patches] - # For single image batch: [0, GRID_T * GRID_H * GRID_W] total_patches = self.grid_t * self.grid_h * self.grid_w cu = torch.tensor([0, total_patches], dtype=torch.int32) self.register_buffer("cu_seqlens", cu) def _patchify(self, pixel_values: torch.Tensor) -> torch.Tensor: - """Turn NCHW pixels into Qwen's pre-patchified [num_patches, patch_dim]. + """Turn pixels into Qwen's pre-patchified [num_patches, patch_dim]. - Reproduces the exact reshape/permute of Qwen2/3-VL's image processor - (transpose order ``(0,3,6,4,7,2,1,5,8)``) so the resulting patch order - matches both the precomputed ``pos_embeds`` and the merger's 2×2 - spatial-merge grouping. The single image is duplicated across the - temporal dimension, matching the processor's last-frame repeat. + Single-image: [1, 3, H, W] → duplicate across temporal dim. + Multi-frame: [1, 3*N, H, W] → reshape real frames. """ c, patch, merge = self.channels, self.patch_size, self.spatial_merge_size hw = self.image_size - # [1, 3, H, W] → [3, H, W] → [temporal, 3, H, W] - x = pixel_values.reshape(c, hw, hw) - x = x.unsqueeze(0).repeat(self.temporal_patch_size, 1, 1, 1) - # split H,W into (grid, merge, patch) and T into (grid_t, temporal) + + if self.num_frames == 1: + x = pixel_values.reshape(c, hw, hw) + x = x.unsqueeze(0).repeat(self.temporal_patch_size, 1, 1, 1) + else: + # [1, 3*N, H, W] → [N, 3, H, W] + x = pixel_values.reshape(self.num_frames, c, hw, hw) + + # [N, C, H, W] → split H,W into (grid, merge, patch), T into (grid_t, temporal) x = x.reshape( self.grid_t, self.temporal_patch_size, @@ -594,7 +604,9 @@ def patched(self, grid_thw): vision_model_cls.fast_pos_embed_interpolate = patched -async def export_vision_encoder(spec: VLMSpec, bundle_path: Path, overwrite: bool) -> str: +async def export_vision_encoder( + spec: VLMSpec, bundle_path: Path, overwrite: bool, num_frames: int = 1 +) -> str: """Export the vision encoder as vision.aimodel and patch metadata.json.""" from transformers.models.qwen3_vl.modeling_qwen3_vl import ( Qwen3VLForConditionalGeneration as HFModel, @@ -622,11 +634,16 @@ async def export_vision_encoder(spec: VLMSpec, bundle_path: Path, overwrite: boo patch_size=spec.patch_size, spatial_merge_size=spec.spatial_merge_size, temporal_patch_size=spec.temporal_patch_size, + num_frames=num_frames, ).eval() del hf_model - num_visual_tokens = spec.num_visual_tokens - pixel_shape = (1, 3, spec.image_size, spec.image_size) + grid_t = num_frames // spec.temporal_patch_size + num_visual_tokens = spec.num_visual_tokens * grid_t + if num_frames == 1: + pixel_shape = (1, 3, spec.image_size, spec.image_size) + else: + pixel_shape = (1, 3 * num_frames, spec.image_size, spec.image_size) # ---- 3. Validate output shape before export ---- with torch.no_grad(): @@ -696,6 +713,11 @@ def forward(self, x): with open(bundle_path / "metadata.json") as f: metadata = json.load(f) metadata["assets"]["vision"] = "vision.aimodel" + if num_frames > 1: + metadata["vision"]["image_token_count"] = num_visual_tokens + metadata["vision"]["max_video_frames"] = num_frames + metadata["vision"]["tokens_per_frame"] = spec.num_visual_tokens + metadata["vision"]["temporal_patch_size"] = spec.temporal_patch_size with open(bundle_path / "metadata.json", "w") as f: json.dump(metadata, f, indent=2) @@ -756,6 +778,14 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="Export only the text decoder + embedding (skip the vision encoder)", ) + parser.add_argument( + "--num-frames", + type=int, + default=1, + help="Number of video frames for the vision encoder (default: 1 = single image). " + "Must be divisible by temporal_patch_size (2 for Qwen). " + "Multi-frame exports bake temporal position embeddings for native video support.", + ) parser.add_argument( "--list-models", action="store_true", @@ -786,7 +816,7 @@ async def _run(spec: VLMSpec, args: argparse.Namespace) -> Path: ) if not args.skip_vision: logging.info("Exporting vision encoder...") - await export_vision_encoder(spec, bundle_path, args.overwrite) + await export_vision_encoder(spec, bundle_path, args.overwrite, args.num_frames) return bundle_path diff --git a/swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift b/swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift index 784fa90f..d0f81351 100644 --- a/swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift +++ b/swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift @@ -183,6 +183,15 @@ public struct VisionConfig: Codable, Sendable, Equatable { /// Whether to include original image dimensions in the text prompt. Defaults to false. public let includeImageInfo: Bool + /// Whether this model supports video (multi-frame) input. + public var supportsVideo: Bool { maxVideoFrames != nil } + + /// Maximum number of video frames for multi-frame models. Nil for image-only models. + public let maxVideoFrames: Int? + + /// Visual tokens produced per frame. Nil defaults to `imageTokenCount`. + public let tokensPerFrame: Int? + /// CLIP normalization (Qwen VL, Pixtral, InternVL, Phi-3.5-vision). public static let clipMean = [0.48145466, 0.4578275, 0.40821073] public static let clipStd = [0.26862954, 0.26130258, 0.27577711] @@ -196,7 +205,9 @@ public struct VisionConfig: Codable, Sendable, Equatable { imageStd: [Double]? = nil, rescaleFactor: Double? = nil, imageStrategy: ImageStrategy? = nil, - includeImageInfo: Bool? = nil + includeImageInfo: Bool? = nil, + maxVideoFrames: Int? = nil, + tokensPerFrame: Int? = nil ) { self.imageSize = imageSize self.patchSize = patchSize @@ -207,6 +218,8 @@ public struct VisionConfig: Codable, Sendable, Equatable { self.rescaleFactor = rescaleFactor ?? 1.0 self.imageStrategy = imageStrategy ?? .stretch self.includeImageInfo = includeImageInfo ?? false + self.maxVideoFrames = maxVideoFrames + self.tokensPerFrame = tokensPerFrame } enum CodingKeys: String, CodingKey { @@ -219,6 +232,8 @@ public struct VisionConfig: Codable, Sendable, Equatable { case rescaleFactor = "rescale_factor" case imageStrategy = "image_strategy" case includeImageInfo = "include_image_info" + case maxVideoFrames = "max_video_frames" + case tokensPerFrame = "tokens_per_frame" } public init(from decoder: Swift.Decoder) throws { @@ -232,5 +247,7 @@ public struct VisionConfig: Codable, Sendable, Equatable { self.rescaleFactor = try c.decodeIfPresent(Double.self, forKey: .rescaleFactor) ?? 1.0 self.imageStrategy = try c.decodeIfPresent(ImageStrategy.self, forKey: .imageStrategy) ?? .stretch self.includeImageInfo = try c.decodeIfPresent(Bool.self, forKey: .includeImageInfo) ?? false + self.maxVideoFrames = try c.decodeIfPresent(Int.self, forKey: .maxVideoFrames) + self.tokensPerFrame = try c.decodeIfPresent(Int.self, forKey: .tokensPerFrame) } } diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift index 8c79c83f..e3daf231 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift @@ -64,8 +64,8 @@ public struct VLMModelConfig: InferenceConfiguration, Codable, Sendable { /// /// ## Inference Flow /// -/// 1. `encodeImage(at:)` — preprocess image, run vision encoder + projector, return `EmbeddedInput` -/// 2. `generate(with: EmbeddedInput, tokens:, ...)` — embed tokens, scatter-merge with vision +/// 1. `encodeImage(at:)` — preprocess image, run vision encoder + projector, return `InputEmbeddings` +/// 2. `generate(with: InputEmbeddings, tokens:, ...)` — embed tokens, scatter-merge with vision /// embeddings at placeholder positions, run LLM prefill, then standard autoregressive decode /// /// KV cache is managed identically to `CoreAISequentialEngine`: starts small and grows @@ -341,11 +341,11 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec /// 1. Load image, resize to `visionConfig.imageSize`, normalize channels /// 2. Run vision encoder (`encode_image`) to get patch features /// 3. Run projector (`project`) to map features to LLM hidden dimension - /// 4. Return as `EmbeddedInput` with placeholder token positions + /// 4. Return as `InputEmbeddings` with placeholder token positions /// /// - Parameter url: URL to the image file (JPEG, PNG, HEIC, etc.) - /// - Returns: `EmbeddedInput` containing projected embeddings and token positions - public func encodeImage(at url: URL) async throws -> EmbeddedInput { + /// - Returns: `InputEmbeddings` containing projected embeddings and token positions + public func encodeImage(at url: URL) async throws -> InputEmbeddings { guard let ciImage = CIImage(contentsOf: url) else { throw ImagePreprocessorError.loadFailed(url) } @@ -355,7 +355,7 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec return try await encodeImage(cgImage: cgImage) } - public func encodeImage(cgImage: CGImage) async throws -> EmbeddedInput { + public func encodeImage(cgImage: CGImage) async throws -> InputEmbeddings { let encodeSignpost = InstrumentsProfiler.beginCustomInterval( name: "CoreAIVLM EncodeImage", details: "cgImage" @@ -386,12 +386,89 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec CLILogger.log("VLM encodeImage complete: \(tokenCount) embedding tokens") - return try EmbeddedInput( + return try InputEmbeddings( embeddings: projectedEmbeddings, embeddingPositions: placeholderRange ) } + // MARK: - Video Encoding (MultimodalInferenceEngine) + + /// Encode video frames into concatenated embeddings. + /// + /// Each frame is processed independently through the vision encoder + projector. + /// Embeddings are concatenated along the sequence dimension to produce a single + /// `InputEmbeddings` with shape `[1, N * tokensPerFrame, hidden_dim]`. + /// + /// Frames are encoded sequentially because the GPU vision encoder cannot run + /// multiple inference calls concurrently on the same model function. + public func encodeVideo(_ video: VideoInput) async throws -> InputEmbeddings { + let tokensPerFrame = config.visionConfig.tokensPerFrame ?? config.visionConfig.imageTokenCount + + var frameEmbeddings: [NDArray] = [] + var frameIndex = 0 + + for try await frame in video.frames { + let chwPixels = try imagePreprocessor.preprocessCHW( + cgImage: frame.image, strategy: config.visionConfig.imageStrategy) + let encoderOutput = try await runVisionEncoder(pixels: chwPixels) + let projected = + visionProjectorFused ? encoderOutput : try await runProjector(encoderOutput: encoderOutput) + frameEmbeddings.append(projected) + frameIndex += 1 + CLILogger.log(" - encoded video frame \(frameIndex)", level: 2) + } + + guard !frameEmbeddings.isEmpty else { + throw InferenceRuntimeError.invalidArgument("encodeVideo: no frames in video input") + } + + if frameEmbeddings.count == 1 { + CLILogger.log("VLM encodeVideo complete: 1 frame, \(tokensPerFrame) tokens") + return try InputEmbeddings( + embeddings: frameEmbeddings[0], + embeddingPositions: 0..= imageTokenCount else { + guard imgSeqLen >= expectedCount else { throw InferenceRuntimeError.invalidArgument( - "scatterMerge: image embeddings have \(imgSeqLen) tokens, need \(imageTokenCount)") + "scatterMerge: image embeddings have \(imgSeqLen) tokens, need \(expectedCount)") } // Validate all positions are within bounds @@ -694,7 +771,7 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec /// - tokens: Full token sequence including image placeholder tokens /// - Returns: Logits for the last token (shape: [vocabSize]) private func vlmPrefill( - embeddedInput: EmbeddedInput, + embeddedInput: InputEmbeddings, tokens: [Int32] ) async throws -> [LogitsScalarType] { let prefillSignpost = InstrumentsProfiler.beginCustomInterval( @@ -799,7 +876,7 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec /// image embeddings at placeholder positions, then runs the LLM. Subsequent steps use /// standard token-by-token decode (embed_tokens -> main). public func generate( - with input: EmbeddedInput, + with input: InputEmbeddings, tokens: [TokenId], samplingConfiguration: SamplingConfiguration, inferenceOptions: InferenceOptions @@ -951,7 +1028,7 @@ extension CoreAISequentialVLMEngine { let engine: CoreAISequentialVLMEngine let input: [CoreAISequentialVLMEngine.TokenId] - let embeddedInput: EmbeddedInput? + let embeddedInput: InputEmbeddings? let samplingConfiguration: SamplingConfiguration let inferenceOptions: InferenceOptions let generationToken: GenerationToken @@ -994,7 +1071,7 @@ extension CoreAISequentialVLMEngine.GenerationSequence { private let generationToken: GenerationToken private var inputTokens: [CoreAISequentialVLMEngine.TokenId] - private var embeddedInput: EmbeddedInput? + private var embeddedInput: InputEmbeddings? private var step: Int = 0 private var finished: Bool = false private var prefillDone: Bool = false @@ -1002,7 +1079,7 @@ extension CoreAISequentialVLMEngine.GenerationSequence { init( engine: CoreAISequentialVLMEngine, input: [CoreAISequentialVLMEngine.TokenId], - embeddedInput: EmbeddedInput?, + embeddedInput: InputEmbeddings?, samplingConfiguration: SamplingConfiguration, inferenceOptions: InferenceOptions, stopReasonStore: StopReasonStore, diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift index f9d41616..812196fa 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift @@ -4,6 +4,7 @@ // be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause import CoreAIShared +import CoreGraphics import Foundation import Synchronization @@ -294,27 +295,36 @@ public enum InferenceRuntimeError: Error, LocalizedError { /// /// The typical flow: /// 1. `encodeImage(at:)` — preprocess + run vision encoder, return embeddings -/// 2. `generate(with: EmbeddedInput, ...)` — scatter-merge embeddings into +/// 2. `generate(with: InputEmbeddings, ...)` — scatter-merge embeddings into /// token sequence and run prefill + decode /// /// The caller owns the embeddings and decides caching strategy. public protocol MultimodalInferenceEngine: InferenceEngine { /// Encode an image into embeddings suitable for injection into the VLM. /// Returns the embedded representation — caller decides whether to cache. - func encodeImage(at url: URL) async throws -> EmbeddedInput + func encodeImage(at url: URL) async throws -> InputEmbeddings + + /// Encode a CGImage into embeddings. + func encodeImage(cgImage: CGImage) async throws -> InputEmbeddings + + /// Encode video frames into concatenated embeddings for injection into the VLM. + /// + /// Default implementation encodes each frame independently through `encodeImage(cgImage:)` + /// and concatenates embeddings along the sequence dimension. + func encodeVideo(_ video: VideoInput) async throws -> InputEmbeddings /// Generate tokens from a token sequence with embedded image regions. /// The engine scatter-merges `input.embeddingPositions` with the embedded data /// during prefill, then continues standard autoregressive decode. func generate( - with input: EmbeddedInput, + with input: InputEmbeddings, tokens: [TokenId], samplingConfiguration: SamplingConfiguration, inferenceOptions: InferenceOptions ) async throws -> OutputSequence } -// TODO: Multi-turn — caller can cache EmbeddedInput across turns and pass it +// TODO: Multi-turn — caller can cache InputEmbeddings across turns and pass it // again with the accumulated token context. Engine keeps image in KV cache // via reset(to:) preserving the prefill portion. diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/EmbeddedInput.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InputEmbeddings.swift similarity index 91% rename from swift/Sources/CoreAILanguageModels/InferenceEngines/EmbeddedInput.swift rename to swift/Sources/CoreAILanguageModels/InferenceEngines/InputEmbeddings.swift index 29e10fb0..d7360f8f 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/EmbeddedInput.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/InputEmbeddings.swift @@ -11,7 +11,7 @@ import Foundation /// Used by multimodal engines to pass vision/audio embeddings into the /// language model. The engine performs scatter-merge: replacing placeholder /// token positions with these embeddings before the first forward pass. -public struct EmbeddedInput: Sendable { +public struct InputEmbeddings: Sendable { /// The embedding tensor, shape [batch, seq_len, hidden_dim]. /// Scalar type matches the LLM's expected input (float16, bFloat16, etc.). public let embeddings: NDArray @@ -22,7 +22,7 @@ public struct EmbeddedInput: Sendable { public init(embeddings: NDArray, embeddingPositions: Range) throws { guard embeddings.shape.count == 3 else { throw InferenceRuntimeError.invalidArgument( - "EmbeddedInput requires 3D embeddings [batch, seq_len, hidden_dim], " + "InputEmbeddings requires 3D embeddings [batch, seq_len, hidden_dim], " + "got shape with \(embeddings.shape.count) dimensions") } self.embeddings = embeddings diff --git a/swift/Sources/CoreAIShared/Image/ImagePreprocessor.swift b/swift/Sources/CoreAIShared/Image/ImagePreprocessor.swift index aa6f702f..61b2d55f 100644 --- a/swift/Sources/CoreAIShared/Image/ImagePreprocessor.swift +++ b/swift/Sources/CoreAIShared/Image/ImagePreprocessor.swift @@ -114,13 +114,17 @@ public struct ImagePreprocessor: Sendable { public func preprocessCHW(cgImage: CGImage) throws -> [Float] { let (data, w, h) = try preprocess(cgImage: cgImage) let pixelCount = w * h + guard pixelCount > 0 else { return [] } var chw = [Float](repeating: 0, count: 3 * pixelCount) data.withUnsafeBytes { rawSrc in let src = rawSrc.bindMemory(to: Float.self) + let n = vDSP_Length(pixelCount) for c in 0..<3 { - for i in 0.. FrameSamplingStrategy { + switch self { + case .uniform: .uniform(count: count) + case .fps(let rate, _): .fps(rate: rate, maxFrames: count) + } + } + + /// Compute sample times (in seconds) for a video of the given duration. + /// Each time targets the midpoint of its sampling interval. + /// + /// - Parameters: + /// - duration: Video duration in seconds. + /// - videoFrameRate: Native frame rate of the video (e.g. 30.0). + /// - Returns: Array of sample times in seconds. + public func sampleTimes(forDuration duration: Double, videoFrameRate: Double) -> [Double] { + guard duration > 0, videoFrameRate > 0 else { return [] } + + let totalFrames = Int(duration * videoFrameRate) + guard totalFrames > 0 else { return [] } + + switch self { + case .uniform(let count): + let n = max(1, min(count, totalFrames)) + if n == 1 { + return [duration / 2.0] + } + let step = duration / Double(n) + return (0.. 0, maxFrames > 0 else { return [] } + let interval = 1.0 / rate + var times: [Double] = [] + var t = interval / 2.0 + while t < duration && times.count < maxFrames { + times.append(t) + t += interval + } + if times.isEmpty { + times.append(duration / 2.0) + } + return times + } + } +} diff --git a/swift/Sources/CoreAIShared/Video/VideoFrameExtractor.swift b/swift/Sources/CoreAIShared/Video/VideoFrameExtractor.swift new file mode 100644 index 00000000..7bf59b15 --- /dev/null +++ b/swift/Sources/CoreAIShared/Video/VideoFrameExtractor.swift @@ -0,0 +1,82 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import AVFoundation +import CoreGraphics +import CoreMedia + +/// Extracts frames from video files using AVAssetImageGenerator. +struct VideoFrameExtractor { + /// Extract frames from a video file according to the given sampling strategy. + /// + /// - Parameters: + /// - url: Local file URL to the video. + /// - sampling: How to sample frames from the video. + /// - skipErrors: When true, frames that fail to decode are skipped. + /// When false (default), the first failure throws. + /// - Returns: Tuple of (frame count, duration in seconds, lazy frame sequence). + /// - Throws: ``VideoInputError`` if the video cannot be read. + static func extractFrames( + from url: URL, + sampling: FrameSamplingStrategy, + skipErrors: Bool = false + ) async throws -> (count: Int, duration: Double, frames: VideoFrameSequence) { + let asset = AVURLAsset(url: url) + let duration = try await CMTimeGetSeconds(asset.load(.duration)) + + guard duration > 0, duration.isFinite else { + throw VideoInputError.invalidVideo("Video has zero or invalid duration") + } + + guard let videoTrack = try await asset.loadTracks(withMediaType: .video).first else { + throw VideoInputError.noVideoTrack + } + + let frameRate = try await Double(videoTrack.load(.nominalFrameRate)) + guard frameRate > 0 else { + throw VideoInputError.invalidVideo("Video has invalid frame rate") + } + + let sampleTimes = sampling.sampleTimes(forDuration: duration, videoFrameRate: frameRate) + guard !sampleTimes.isEmpty else { + throw VideoInputError.invalidVideo("No frames to extract") + } + + let count = sampleTimes.count + let cmTimes = sampleTimes.map { + CMTime(seconds: $0, preferredTimescale: 600) + } + + let seq = VideoFrameSequence( + stream: AsyncThrowingStream { continuation in + Task { + let asset = AVURLAsset(url: url) + let generator = AVAssetImageGenerator(asset: asset) + generator.appliesPreferredTrackTransform = true + generator.requestedTimeToleranceBefore = CMTime(seconds: 0.1, preferredTimescale: 600) + generator.requestedTimeToleranceAfter = CMTime(seconds: 0.1, preferredTimescale: 600) + + var frameIndex = 0 + for await result in generator.images(for: cmTimes) { + do { + let image = try result.image + continuation.yield(VideoFrame(image: image, index: frameIndex)) + frameIndex += 1 + } catch { + if skipErrors { + continue + } + continuation.finish( + throwing: VideoInputError.frameExtractionFailed(underlying: error)) + return + } + } + continuation.finish() + } + }) + + return (count: count, duration: duration, frames: seq) + } +} diff --git a/swift/Sources/CoreAIShared/Video/VideoInput.swift b/swift/Sources/CoreAIShared/Video/VideoInput.swift new file mode 100644 index 00000000..a406da7a --- /dev/null +++ b/swift/Sources/CoreAIShared/Video/VideoInput.swift @@ -0,0 +1,121 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import CoreGraphics +import Foundation + +/// A frame produced by video extraction: the decoded image and its position. +public struct VideoFrame: Sendable { + public let image: CGImage + /// Frame index (0-based ordinal within the extraction sequence). + public let index: Int +} + +/// Concrete, Sendable async sequence of video frames. +/// +/// Wraps an `AsyncThrowingStream` so that `VideoInput` can conform to `Sendable` +/// in Swift 6 strict concurrency (existential `any AsyncSequence` cannot). +public struct VideoFrameSequence: AsyncSequence, Sendable { + public typealias Element = VideoFrame + + private let stream: AsyncThrowingStream + + init(stream: AsyncThrowingStream) { + self.stream = stream + } + + public func makeAsyncIterator() -> Iterator { + Iterator(base: stream.makeAsyncIterator()) + } + + public struct Iterator: AsyncIteratorProtocol { + var base: AsyncThrowingStream.AsyncIterator + + public mutating func next() async throws -> VideoFrame? { + try await base.next() + } + } +} + +/// Default number of frames sampled from a video when no explicit count is provided. +public let defaultVideoFrameCount = 8 + +/// Extracted video frames ready for vision encoding. +/// +/// Frames are delivered lazily via `VideoFrameSequence`. The engine processes +/// and releases each frame incrementally, keeping peak memory at 1-2 +/// decoded frames plus accumulated embeddings. +public struct VideoInput: Sendable { + /// Number of frames, if known ahead of time (nil for live streams). + public let frameCount: Int? + /// Video duration in seconds, if known (nil for live streams). + public let duration: Double? + /// Lazy frame sequence. Throws on extraction errors unless `skipErrors` was set. + public let frames: VideoFrameSequence + + public init( + frameCount: Int?, + duration: Double?, + frames: VideoFrameSequence + ) { + self.frameCount = frameCount + self.duration = duration + self.frames = frames + } + + /// Extract frames from a local video file. + /// + /// - Parameters: + /// - url: File URL to a video (MP4, MOV, etc.). + /// - sampling: Frame sampling strategy (default: 8 uniform frames). + /// - skipErrors: When true, frames that fail to decode are skipped + /// instead of throwing. Default is false. + /// - Throws: ``VideoInputError`` if the file cannot be read or has no video track. + public static func fromURL( + _ url: URL, + sampling: FrameSamplingStrategy = .uniform(count: defaultVideoFrameCount), + skipErrors: Bool = false + ) async throws -> VideoInput { + let (count, duration, frames) = try await VideoFrameExtractor.extractFrames( + from: url, sampling: sampling, skipErrors: skipErrors) + return VideoInput(frameCount: count, duration: duration, frames: frames) + } + + /// Wrap pre-extracted frames (e.g. from camera capture). + public static func fromFrames( + _ images: [CGImage] + ) -> VideoInput { + let count = images.count + let stream = AsyncThrowingStream { continuation in + for (i, image) in images.enumerated() { + continuation.yield(VideoFrame(image: image, index: i)) + } + continuation.finish() + } + return VideoInput(frameCount: count, duration: nil, frames: VideoFrameSequence(stream: stream)) + } +} + +// MARK: - Errors + +public enum VideoInputError: Error, LocalizedError { + case invalidVideo(String) + case noVideoTrack + case fileNotFound(URL) + case frameExtractionFailed(underlying: Error) + + public var errorDescription: String? { + switch self { + case .invalidVideo(let reason): + return "Invalid video: \(reason)" + case .noVideoTrack: + return "File has no video track (may be audio-only)" + case .fileNotFound(let url): + return "Video file not found: \(url.path)" + case .frameExtractionFailed(let underlying): + return "Frame extraction failed: \(underlying.localizedDescription)" + } + } +} diff --git a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift index c55727e3..78a15a77 100644 --- a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift +++ b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift @@ -50,6 +50,16 @@ enum ImageInfoMode: String, ExpressibleByArgument, Sendable { case auto } +extension FrameSamplingStrategy: ExpressibleByArgument { + public init?(argument: String) { + switch argument.lowercased() { + case "uniform": self = .uniform(count: defaultVideoFrameCount) + case "fps": self = .fps(rate: 1.0, maxFrames: defaultVideoFrameCount) + default: return nil + } + } +} + @main struct Main { static func main() async throws { @@ -198,6 +208,19 @@ struct LLMRunner: AsyncParsableCommand, Sendable { help: "Include original image resolution in prompt: on, off, auto (default: auto)") var imageInfo: ImageInfoMode = .auto + @Option(name: .customLong("video"), help: "Path to a video file for vision-language models") + var videoPath: String? + + @Option( + name: .customLong("video-frames"), + help: "Number of frames to extract from video (default: \(defaultVideoFrameCount))") + var videoFrames: Int = defaultVideoFrameCount + + @Option( + name: .customLong("video-sampling"), + help: "Video frame sampling: uniform (default) or fps") + var videoSampling: FrameSamplingStrategy = .uniform(count: defaultVideoFrameCount) + @Flag( name: .customLong("clear-coreai-cache"), help: "Clear Core AI cached specialization for this model before loading (forces re-specialization)" @@ -223,6 +246,12 @@ struct LLMRunner: AsyncParsableCommand, Sendable { if let c = chunkSize, c <= 0 { throw ValidationError("--chunk-size must be > 0") } + if imagePath != nil && videoPath != nil { + throw ValidationError("--image and --video cannot be used together") + } + if videoFrames < 1 { + throw ValidationError("--video-frames must be >= 1") + } } func run() async throws { @@ -544,6 +573,21 @@ struct LLMRunner: AsyncParsableCommand, Sendable { return } + // VLM video path: if --video is provided + if let videoPath = videoPath { + try await runVLMVideoInference( + videoPath: videoPath, + inferenceEngine: inferenceEngine, + bundle: bundle, + tokenizer: tokenizer, + samplingConfiguration: samplingConfiguration, + maxTokens: maxTokens, + additionalEosTokenIds: additionalEosTokenIds, + displayPrompt: displayPrompt + ) + return + } + // Build text generator with preloaded inference engine CLILogger.log("Building text generator...", component: "Main") @@ -903,15 +947,90 @@ struct LLMRunner: AsyncParsableCommand, Sendable { } } - // Build VLM prompt with image placeholder tokens. - // Try using the tokenizer's chat template if available; fall back to - // generic "USER: ×N \n prompt \nASSISTANT:" format. + try await runVLMGeneration( + vlmEngine: vlmEngine, + embeddedInput: embeddedInput, + visionConfig: visionConfig, + tokenizer: tokenizer, + samplingConfiguration: samplingConfiguration, + maxTokens: maxTokens, + additionalEosTokenIds: additionalEosTokenIds, + displayPrompt: effectivePrompt + ) + } + + // MARK: - VLM Video Inference + + private func runVLMVideoInference( + videoPath: String, + inferenceEngine: any InferenceEngine, + bundle: LanguageBundle, + tokenizer: any Tokenizer, + samplingConfiguration: SamplingConfiguration, + maxTokens: Int, + additionalEosTokenIds: [Int32], + displayPrompt: String + ) async throws { + guard let vlmEngine = inferenceEngine as? any MultimodalInferenceEngine else { + print("Error: --video requires a vision-language model (engine does not support multimodal)") + throw ExitCode.failure + } + + let videoURL = URL(fileURLWithPath: videoPath) + guard FileManager.default.fileExists(atPath: videoURL.path) else { + print("Error: video not found at \(videoPath)") + throw ExitCode.failure + } + + guard let visionConfig = bundle.visionConfig else { + print("Error: VLM bundle missing 'vision' config in metadata.json") + throw ExitCode.failure + } + + let sampling = videoSampling.withFrameCount(videoFrames) + + if !CLILogger.isVerbose { + print("Generating...") + } + + CLILogger.log("Extracting \(videoFrames) frames from: \(videoPath)", component: "VLM") + let videoInput = try await VideoInput.fromURL(videoURL, sampling: sampling) + CLILogger.log("Encoding video frames...", component: "VLM") + let embeddedInput = try await vlmEngine.encodeVideo(videoInput) + CLILogger.log( + "Video encoded: \(embeddedInput.tokenCount) visual tokens from \(videoInput.frameCount ?? 0) frames", + component: "VLM") + + try await runVLMGeneration( + vlmEngine: vlmEngine, + embeddedInput: embeddedInput, + visionConfig: visionConfig, + tokenizer: tokenizer, + samplingConfiguration: samplingConfiguration, + maxTokens: maxTokens, + additionalEosTokenIds: additionalEosTokenIds, + displayPrompt: displayPrompt + ) + } + + // MARK: - Shared VLM Generation + + private func runVLMGeneration( + vlmEngine: any MultimodalInferenceEngine, + embeddedInput: InputEmbeddings, + visionConfig: VisionConfig, + tokenizer: any Tokenizer, + samplingConfiguration: SamplingConfiguration, + maxTokens: Int, + additionalEosTokenIds: [Int32], + displayPrompt: String + ) async throws { let imageTokenCount = embeddedInput.tokenCount let imageTokenId = visionConfig.imageTokenId let vlmTokens: [Int32] if let chatTemplateTokens = try? buildVLMPromptFromChatTemplate( - prompt: effectivePrompt, + prompt: displayPrompt, imageTokenCount: imageTokenCount, imageTokenId: imageTokenId, tokenizer: tokenizer @@ -924,17 +1043,16 @@ struct LLMRunner: AsyncParsableCommand, Sendable { component: "VLM") var tokens = tokenizer.encode(text: "USER: ", addSpecialTokens: true).map { Int32($0) } tokens.append(contentsOf: [Int32](repeating: imageTokenId, count: imageTokenCount)) - let suffix = "\n" + effectivePrompt + "\nASSISTANT:" + let suffix = "\n" + displayPrompt + "\nASSISTANT:" tokens.append( contentsOf: tokenizer.encode(text: suffix, addSpecialTokens: false).map { Int32($0) }) vlmTokens = tokens } CLILogger.log( - "VLM prompt: \(vlmTokens.count) tokens (\(imageTokenCount) image placeholders)", + "VLM prompt: \(vlmTokens.count) tokens (\(imageTokenCount) visual placeholders)", component: "VLM") - // Build stop token set var eosTokenIds = Set() if let eos = tokenizer.eosTokenId { eosTokenIds.insert(Int32(eos)) } eosTokenIds.formUnion(additionalEosTokenIds) @@ -950,7 +1068,6 @@ struct LLMRunner: AsyncParsableCommand, Sendable { let inferenceID = InstrumentsProfiler.beginInference( promptTokens: vlmTokens.count, maxTokens: maxTokens) - await PerformanceMetrics.shared.recordPromptTokens(vlmTokens.count) let tokenStream = try await vlmEngine.generate( @@ -965,8 +1082,8 @@ struct LLMRunner: AsyncParsableCommand, Sendable { CLILogger.log("VLM generate started, maxTokens=\(maxTokens)", component: "VLM") - // Prompt (prefill) timing — first token latency - var promptSpan: ProfileSpan? = InstrumentsProfiler.beginPrompt(tokens: vlmTokens.count, engine: "CoreAIVLM") + var promptSpan: ProfileSpan? = InstrumentsProfiler.beginPrompt( + tokens: vlmTokens.count, engine: "CoreAIVLM") var extendSpan: ProfileSpan? let needsLogits = printLogits || saveLogits != nil let topKCount = saveLogitsLength.topKForFile ?? 5 @@ -1012,12 +1129,10 @@ struct LLMRunner: AsyncParsableCommand, Sendable { extendSpan?.end() print() - // Save logits to JSON if requested if let path = saveLogits, !allTokenLogits.isEmpty { try LogitsWriter.saveTopKJSON(tokenLogits: allTokenLogits, path: path) } - // Record generation stats InstrumentsProfiler.endInference( generatedTokens: generatedTokens.count, signpostID: inferenceID) await PerformanceMetrics.shared.recordGeneratedTokens(generatedTokens.count) diff --git a/swift/Tests/CoreAISharedTests/Video/VideoInputTests.swift b/swift/Tests/CoreAISharedTests/Video/VideoInputTests.swift new file mode 100644 index 00000000..a303e57e --- /dev/null +++ b/swift/Tests/CoreAISharedTests/Video/VideoInputTests.swift @@ -0,0 +1,158 @@ +// Copyright 2026 Apple Inc. +// +// Use of this source code is governed by a BSD-3-clause license that can +// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +import CoreGraphics +import Foundation +import Testing + +@testable import CoreAIShared + +@Suite("Frame Sampling Strategy") +struct FrameSamplingStrategyTests { + @Test("Uniform sampling with 8 frames from 30fps 10s video") + func uniformSamplingBasic() { + let timestamps = FrameSamplingStrategy.uniform(count: 8) + .sampleTimes(forDuration: 10.0, videoFrameRate: 30.0) + #expect(timestamps.count == 8) + // Frames should be evenly spaced + for i in 1.. timestamps[i - 1]) + } + // All within duration + for t in timestamps { + #expect(t >= 0 && t < 10.0) + } + } + + @Test("Uniform sampling with single frame returns midpoint") + func uniformSamplingSingleFrame() { + let timestamps = FrameSamplingStrategy.uniform(count: 1) + .sampleTimes(forDuration: 10.0, videoFrameRate: 30.0) + #expect(timestamps.count == 1) + #expect(timestamps[0] == 5.0) + } + + @Test("Uniform sampling clamps to available frames") + func uniformSamplingMoreThanAvailable() { + // 5 frames total at 1fps for 5 seconds + let timestamps = FrameSamplingStrategy.uniform(count: 16) + .sampleTimes(forDuration: 5.0, videoFrameRate: 1.0) + #expect(timestamps.count == 5) + } + + @Test("Uniform sampling with zero count gives 1 frame") + func uniformSamplingZeroCount() { + let timestamps = FrameSamplingStrategy.uniform(count: 0) + .sampleTimes(forDuration: 10.0, videoFrameRate: 30.0) + #expect(timestamps.count == 1) + } + + @Test("Uniform sampling with zero duration returns empty") + func uniformSamplingZeroDuration() { + let timestamps = FrameSamplingStrategy.uniform(count: 8) + .sampleTimes(forDuration: 0.0, videoFrameRate: 30.0) + #expect(timestamps.isEmpty) + } + + @Test("Uniform sampling with negative duration returns empty") + func uniformSamplingNegativeDuration() { + let timestamps = FrameSamplingStrategy.uniform(count: 8) + .sampleTimes(forDuration: -5.0, videoFrameRate: 30.0) + #expect(timestamps.isEmpty) + } + + @Test("FPS-based sampling at 1fps from 10s video") + func fpsSampling() { + let timestamps = FrameSamplingStrategy.fps(rate: 1.0, maxFrames: 20) + .sampleTimes(forDuration: 10.0, videoFrameRate: 30.0) + #expect(timestamps.count == 10) + for t in timestamps { + #expect(t >= 0 && t < 10.0) + } + } + + @Test("FPS-based sampling caps at maxFrames") + func fpsSamplingCapped() { + let timestamps = FrameSamplingStrategy.fps(rate: 1.0, maxFrames: 5) + .sampleTimes(forDuration: 60.0, videoFrameRate: 30.0) + #expect(timestamps.count == 5) + } + + @Test("FPS-based sampling at 2fps") + func fpsSamplingHighRate() { + let timestamps = FrameSamplingStrategy.fps(rate: 2.0, maxFrames: 100) + .sampleTimes(forDuration: 5.0, videoFrameRate: 30.0) + #expect(timestamps.count == 10) + } + + @Test("FPS-based sampling with zero rate returns empty") + func fpsSamplingZeroRate() { + let timestamps = FrameSamplingStrategy.fps(rate: 0.0, maxFrames: 10) + .sampleTimes(forDuration: 10.0, videoFrameRate: 30.0) + #expect(timestamps.isEmpty) + } + + @Test("FPS-based sampling with zero maxFrames returns empty") + func fpsSamplingZeroMaxFrames() { + let timestamps = FrameSamplingStrategy.fps(rate: 1.0, maxFrames: 0) + .sampleTimes(forDuration: 10.0, videoFrameRate: 30.0) + #expect(timestamps.isEmpty) + } + + @Test("FPS-based sampling short video returns at least 1 frame") + func fpsSamplingShortVideo() { + let timestamps = FrameSamplingStrategy.fps(rate: 1.0, maxFrames: 10) + .sampleTimes(forDuration: 0.5, videoFrameRate: 30.0) + // 0.5s video at 1fps: interval center at 0.5 is outside, fallback gives 1 frame + #expect(timestamps.count == 1) + } +} + +@Suite("VideoInput") +struct VideoInputTests { + @Test("fromFrames wraps CGImages correctly") + func fromFrames() async throws { + let images = (0..<3).map { _ in makeSolidImage(width: 64, height: 64) } + let input = VideoInput.fromFrames(images) + #expect(input.frameCount == 3) + #expect(input.duration == nil) + + var count = 0 + for try await _ in input.frames { + count += 1 + } + #expect(count == 3) + } + + @Test("fromFrames with empty array gives zero frames") + func fromFramesEmpty() async throws { + let input = VideoInput.fromFrames([]) + #expect(input.frameCount == 0) + + var count = 0 + for try await _ in input.frames { + count += 1 + } + #expect(count == 0) + } +} + +// MARK: - Test Helpers + +private func makeSolidImage(width: Int, height: Int) -> CGImage { + let colorSpace = CGColorSpace(name: CGColorSpace.sRGB)! + let ctx = CGContext( + data: nil, + width: width, + height: height, + bitsPerComponent: 8, + bytesPerRow: width * 4, + space: colorSpace, + bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue + )! + ctx.setFillColor(CGColor(red: 1, green: 0, blue: 0, alpha: 1)) + ctx.fill(CGRect(x: 0, y: 0, width: width, height: height)) + return ctx.makeImage()! +} diff --git a/swift/Tests/LanguageModelsTests/VLMProtocolTests.swift b/swift/Tests/LanguageModelsTests/VLMProtocolTests.swift index 9132269a..ec2e2652 100644 --- a/swift/Tests/LanguageModelsTests/VLMProtocolTests.swift +++ b/swift/Tests/LanguageModelsTests/VLMProtocolTests.swift @@ -15,13 +15,13 @@ import CoreAI @Suite("Multimodal types") struct MultimodalTypeTests { #if canImport(CoreAI) - @Test("EmbeddedInput wraps NDArray with positions") + @Test("InputEmbeddings wraps NDArray with positions") func embeddedInputBasics() throws { let embeddings = NDArray( shape: [1, 256, 2048], scalarType: .float16 ) - let input = try EmbeddedInput( + let input = try InputEmbeddings( embeddings: embeddings, embeddingPositions: 5..<261 ) @@ -75,4 +75,39 @@ struct MultimodalTypeTests { let config = try JSONDecoder().decode(LanguageConfig.self, from: json.data(using: .utf8)!) #expect(config.vision == nil) } + + @Test("VisionConfig decodes with video fields") + func visionConfigWithVideoFields() throws { + let json = """ + { + "image_size": 384, + "patch_size": 14, + "image_token_count": 729, + "image_token_id": 255999, + "max_video_frames": 16, + "tokens_per_frame": 729 + } + """ + let config = try JSONDecoder().decode(VisionConfig.self, from: json.data(using: .utf8)!) + #expect(config.maxVideoFrames == 16) + #expect(config.tokensPerFrame == 729) + #expect(config.imageSize == 384) + } + + @Test("VisionConfig backwards compatible without video fields") + func visionConfigWithoutVideoFields() throws { + let json = """ + { + "image_size": 896, + "patch_size": 14, + "image_token_count": 256, + "image_token_id": 255999 + } + """ + let config = try JSONDecoder().decode(VisionConfig.self, from: json.data(using: .utf8)!) + #expect(config.maxVideoFrames == nil) + #expect(config.tokensPerFrame == nil) + #expect(config.imageSize == 896) + #expect(config.imageTokenCount == 256) + } }