Skip to content

Commit 2607aea

Browse files
stikvesstikves
authored
Add configurable VLM image preprocessing strategy (#108)
* Add configurable VLM image preprocessing strategy Support three image preprocessing strategies for VLM vision encoders: - stretch: resize directly to target (default, backward compatible) - center_crop: shortest-edge resize then center crop (for CLIP-based models) - pad: longest-edge resize with zero-padding (preserves geometry) The strategy is declared in metadata.json and inferred from the model at export time. Runtime CLI overrides via --image-strategy. Also adds --image-info flag to optionally inject original image dimensions into the text prompt (useful for models trained with resolution awareness). Closes #100. * Fix Swift format * Address review: add default comment, add padPortrait test
1 parent aeb6ae3 commit 2607aea

7 files changed

Lines changed: 274 additions & 31 deletions

File tree

‎models/vlm/README.md‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,31 @@ with asset roles consumed by the Swift runner's `ModelBundle`:
4343
Add a `VLMSpec(...)` entry to `SUPPORTED_MODELS` in
4444
[`vlm/export.py`](../../python/src/coreai_models/vlm/export.py) with the
4545
HuggingFace ID, output name, image token id, and vision geometry (resolution,
46-
patch/merge sizes, CLIP normalization stats). Models whose text decoder needs a
46+
patch/merge sizes, normalization stats). Models whose text decoder needs a
4747
new architecture also require a class registered in
4848
[`models/registry.py`](../../python/src/coreai_models/models/registry.py).
49+
50+
## Image preprocessing
51+
52+
The vision encoder expects a fixed-size square input. How an arbitrary image
53+
reaches that square is controlled by `image_strategy` in `metadata.json`:
54+
55+
| Strategy | Behavior | Use when |
56+
|---------------|------------------------------------------|--------------------------------------|
57+
| `stretch` | Resize directly to target size | Default. Works for most models. |
58+
| `center_crop` | Shortest-edge resize, then center crop | CLIP-based vision towers (FastVLM) |
59+
| `pad` | Longest-edge resize, zero-pad remainder | Models expecting preserved geometry |
60+
61+
The strategy is inferred from the model's `preprocessor_config.json` at export
62+
time and written into the bundle's `metadata.json`. Override at runtime:
63+
64+
```bash
65+
llm-runner --model vlm_bundle --image photo.jpg --image-strategy center_crop
66+
```
67+
68+
### Original resolution in prompt
69+
70+
Some models (Qwen-VL family) benefit from knowing the original image
71+
dimensions. When `include_image_info` is set in the bundle metadata (or
72+
overridden via `--image-info on`), the original `W×H` is prepended to the
73+
text prompt before tokenization.

‎python/src/coreai_models/vlm/export.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@ class VLMSpec:
6464
image_mean: tuple[float, float, float]
6565
image_std: tuple[float, float, float]
6666
rescale_factor: float
67+
image_strategy: str = "stretch"
68+
include_image_info: bool = False
6769

6870
@property
6971
def num_visual_tokens(self) -> int:
@@ -84,6 +86,8 @@ def num_visual_tokens(self) -> int:
8486
image_mean=(0.5, 0.5, 0.5),
8587
image_std=(0.5, 0.5, 0.5),
8688
rescale_factor=1.0,
89+
image_strategy="stretch",
90+
include_image_info=True,
8791
),
8892
}
8993

@@ -378,6 +382,8 @@ async def export_text_bundle(
378382
"image_mean": list(spec.image_mean),
379383
"image_std": list(spec.image_std),
380384
"rescale_factor": spec.rescale_factor,
385+
"image_strategy": spec.image_strategy,
386+
"include_image_info": spec.include_image_info,
381387
},
382388
"source": {
383389
"hf_model_id": spec.hf_model_id,

‎swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,12 @@ public struct VisionConfig: Codable, Sendable, Equatable {
169169
/// Pixel rescale factor applied before normalization. Defaults to 1.0 when omitted.
170170
public let rescaleFactor: Double
171171

172+
/// Image preprocessing strategy. Defaults to stretch when omitted.
173+
public let imageStrategy: ImageStrategy
174+
175+
/// Whether to include original image dimensions in the text prompt. Defaults to false.
176+
public let includeImageInfo: Bool
177+
172178
/// CLIP normalization (Qwen VL, Pixtral, InternVL, Phi-3.5-vision).
173179
public static let clipMean = [0.48145466, 0.4578275, 0.40821073]
174180
public static let clipStd = [0.26862954, 0.26130258, 0.27577711]
@@ -180,7 +186,9 @@ public struct VisionConfig: Codable, Sendable, Equatable {
180186
imageTokenId: Int32,
181187
imageMean: [Double]? = nil,
182188
imageStd: [Double]? = nil,
183-
rescaleFactor: Double? = nil
189+
rescaleFactor: Double? = nil,
190+
imageStrategy: ImageStrategy? = nil,
191+
includeImageInfo: Bool? = nil
184192
) {
185193
self.imageSize = imageSize
186194
self.patchSize = patchSize
@@ -189,6 +197,8 @@ public struct VisionConfig: Codable, Sendable, Equatable {
189197
self.imageMean = imageMean ?? Self.clipMean
190198
self.imageStd = imageStd ?? Self.clipStd
191199
self.rescaleFactor = rescaleFactor ?? 1.0
200+
self.imageStrategy = imageStrategy ?? .stretch
201+
self.includeImageInfo = includeImageInfo ?? false
192202
}
193203

194204
enum CodingKeys: String, CodingKey {
@@ -199,6 +209,8 @@ public struct VisionConfig: Codable, Sendable, Equatable {
199209
case imageMean = "image_mean"
200210
case imageStd = "image_std"
201211
case rescaleFactor = "rescale_factor"
212+
case imageStrategy = "image_strategy"
213+
case includeImageInfo = "include_image_info"
202214
}
203215

204216
public init(from decoder: Swift.Decoder) throws {
@@ -210,5 +222,7 @@ public struct VisionConfig: Codable, Sendable, Equatable {
210222
self.imageMean = try c.decodeIfPresent([Double].self, forKey: .imageMean) ?? Self.clipMean
211223
self.imageStd = try c.decodeIfPresent([Double].self, forKey: .imageStd) ?? Self.clipStd
212224
self.rescaleFactor = try c.decodeIfPresent(Double.self, forKey: .rescaleFactor) ?? 1.0
225+
self.imageStrategy = try c.decodeIfPresent(ImageStrategy.self, forKey: .imageStrategy) ?? .stretch
226+
self.includeImageInfo = try c.decodeIfPresent(Bool.self, forKey: .includeImageInfo) ?? false
213227
}
214228
}

‎swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -362,7 +362,8 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec
362362
)
363363

364364
// Step 1: Preprocess image to CHW Float32
365-
let chwPixels = try imagePreprocessor.preprocessCHW(cgImage: cgImage)
365+
let chwPixels = try imagePreprocessor.preprocessCHW(
366+
cgImage: cgImage, strategy: config.visionConfig.imageStrategy)
366367

367368
// Step 2: Run encode_image
368369
let encoderOutput = try await runVisionEncoder(pixels: chwPixels)

‎swift/Sources/CoreAIShared/Image/ImagePreprocessor.swift‎

Lines changed: 126 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,14 @@ import CoreImage
99
import Foundation
1010
import ImageIO
1111

12+
// MARK: - Image Strategy
13+
14+
public enum ImageStrategy: String, Codable, Sendable {
15+
case stretch
16+
case centerCrop = "center_crop"
17+
case pad
18+
}
19+
1220
// MARK: - ImagePreprocessor
1321

1422
/// Resizes an image to a target size and applies per-channel normalization
@@ -92,7 +100,125 @@ public struct ImagePreprocessor: Sendable {
92100
public func preprocess(cgImage: CGImage) throws -> (Data, Int, Int) {
93101
let w = Int(targetSize.width)
94102
let h = Int(targetSize.height)
103+
let resized = try renderToContext(
104+
cgImage: cgImage, canvasWidth: w, canvasHeight: h,
105+
drawRect: CGRect(x: 0, y: 0, width: w, height: h))
106+
return try normalize(pixels: resized, width: w, height: h)
107+
}
108+
109+
/// Preprocess a CGImage and return a flat CHW `[C, H, W]` Float32 array.
110+
///
111+
/// Convenience wrapper over ``preprocess(cgImage:)`` that transposes the
112+
/// NHWC RGBA output into the planar `[3, H, W]` layout expected by most
113+
/// vision encoder inputs.
114+
public func preprocessCHW(cgImage: CGImage) throws -> [Float] {
115+
let (data, w, h) = try preprocess(cgImage: cgImage)
116+
let pixelCount = w * h
117+
var chw = [Float](repeating: 0, count: 3 * pixelCount)
118+
data.withUnsafeBytes { rawSrc in
119+
let src = rawSrc.bindMemory(to: Float.self)
120+
for c in 0..<3 {
121+
for i in 0..<pixelCount {
122+
chw[c * pixelCount + i] = src[i * 4 + c]
123+
}
124+
}
125+
}
126+
return chw
127+
}
128+
129+
/// Preprocess with center-crop strategy: resize shortest edge to target,
130+
/// then center-crop to square. Returns flat CHW `[3, H, W]` Float32.
131+
public func preprocessCHWCenterCrop(cgImage: CGImage) throws -> [Float] {
132+
let targetW = Int(targetSize.width)
133+
let targetH = Int(targetSize.height)
134+
let srcW = cgImage.width
135+
let srcH = cgImage.height
136+
137+
let scale =
138+
srcW < srcH
139+
? CGFloat(targetW) / CGFloat(srcW)
140+
: CGFloat(targetH) / CGFloat(srcH)
141+
let resizedW = Int(round(CGFloat(srcW) * scale))
142+
let resizedH = Int(round(CGFloat(srcH) * scale))
143+
144+
let resized = try renderToContext(
145+
cgImage: cgImage, canvasWidth: resizedW, canvasHeight: resizedH,
146+
drawRect: CGRect(x: 0, y: 0, width: resizedW, height: resizedH))
147+
148+
let cropX = (resizedW - targetW) / 2
149+
let cropY = (resizedH - targetH) / 2
150+
guard let cropped = resized.cropping(to: CGRect(x: cropX, y: cropY, width: targetW, height: targetH)) else {
151+
throw ImagePreprocessorError.renderFailed
152+
}
153+
154+
return try preprocessCHW(cgImage: cropped)
155+
}
156+
157+
/// Preprocess with pad strategy: resize longest edge to target, zero-pad
158+
/// the shorter dimension. Returns flat CHW `[3, H, W]` Float32.
159+
public func preprocessCHWPad(cgImage: CGImage) throws -> [Float] {
160+
let targetW = Int(targetSize.width)
161+
let targetH = Int(targetSize.height)
162+
let srcW = cgImage.width
163+
let srcH = cgImage.height
164+
165+
let scale =
166+
srcW > srcH
167+
? CGFloat(targetW) / CGFloat(srcW)
168+
: CGFloat(targetH) / CGFloat(srcH)
169+
let resizedW = Int(round(CGFloat(srcW) * scale))
170+
let resizedH = Int(round(CGFloat(srcH) * scale))
171+
172+
let offsetX = (targetW - resizedW) / 2
173+
let offsetY = (targetH - resizedH) / 2
174+
175+
let padded = try renderToContext(
176+
cgImage: cgImage, canvasWidth: targetW, canvasHeight: targetH,
177+
drawRect: CGRect(x: offsetX, y: offsetY, width: resizedW, height: resizedH))
178+
179+
return try preprocessCHW(cgImage: padded)
180+
}
181+
182+
/// Dispatch preprocessing based on strategy. Returns flat CHW `[3, H, W]`.
183+
public func preprocessCHW(cgImage: CGImage, strategy: ImageStrategy) throws -> [Float] {
184+
switch strategy {
185+
case .stretch: return try preprocessCHW(cgImage: cgImage)
186+
case .centerCrop: return try preprocessCHWCenterCrop(cgImage: cgImage)
187+
case .pad: return try preprocessCHWPad(cgImage: cgImage)
188+
}
189+
}
95190

191+
// MARK: - Private Helpers
192+
193+
private func renderToContext(
194+
cgImage: CGImage, canvasWidth: Int, canvasHeight: Int,
195+
drawRect: CGRect
196+
) throws -> CGImage {
197+
guard let colorSpace = CGColorSpace(name: CGColorSpace.sRGB) else {
198+
throw ImagePreprocessorError.renderFailed
199+
}
200+
guard
201+
let ctx = CGContext(
202+
data: nil,
203+
width: canvasWidth,
204+
height: canvasHeight,
205+
bitsPerComponent: 8,
206+
bytesPerRow: canvasWidth * 4,
207+
space: colorSpace,
208+
bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue
209+
)
210+
else {
211+
throw ImagePreprocessorError.renderFailed
212+
}
213+
ctx.interpolationQuality = .high
214+
ctx.draw(cgImage, in: drawRect)
215+
guard let result = ctx.makeImage() else {
216+
throw ImagePreprocessorError.renderFailed
217+
}
218+
return result
219+
}
220+
221+
private func normalize(pixels cgImage: CGImage, width w: Int, height h: Int) throws -> (Data, Int, Int) {
96222
guard let colorSpace = CGColorSpace(name: CGColorSpace.sRGB) else {
97223
throw ImagePreprocessorError.renderFailed
98224
}
@@ -109,7 +235,6 @@ public struct ImagePreprocessor: Sendable {
109235
else {
110236
throw ImagePreprocessorError.renderFailed
111237
}
112-
ctx.interpolationQuality = .high
113238
ctx.draw(cgImage, in: CGRect(x: 0, y: 0, width: w, height: h))
114239

115240
guard let pixelData = ctx.data else {
@@ -123,10 +248,6 @@ public struct ImagePreprocessor: Sendable {
123248
data.withUnsafeMutableBytes { dstPtr in
124249
guard let dstBase = dstPtr.bindMemory(to: Float.self).baseAddress else { return }
125250

126-
// Combine `(x*scale - mean) / std` into a single vDSP_vsmsa call:
127-
// y = x * (scale/std) + (-mean/std)
128-
// For each color channel, read UInt8 with stride 4, write Float
129-
// back to the interleaved RGBA destination at stride 4.
130251
let scale = Float(rescaleFactor) / 255.0
131252
let means: [Float] = [Float(mean.0), Float(mean.1), Float(mean.2)]
132253
let stds: [Float] = [Float(std.0), Float(std.1), Float(std.2)]
@@ -140,33 +261,12 @@ public struct ImagePreprocessor: Sendable {
140261
vDSP_vsmsa(channel, 1, &a, &b, dstBase.advanced(by: c), 4, n)
141262
}
142263

143-
// Alpha unused — write 0 at stride 4.
144264
var zero: Float = 0
145265
vDSP_vfill(&zero, dstBase.advanced(by: 3), 4, n)
146266
}
147267

148268
return (data, w, h)
149269
}
150-
151-
/// Preprocess a CGImage and return a flat CHW `[C, H, W]` Float32 array.
152-
///
153-
/// Convenience wrapper over ``preprocess(cgImage:)`` that transposes the
154-
/// NHWC RGBA output into the planar `[3, H, W]` layout expected by most
155-
/// vision encoder inputs.
156-
public func preprocessCHW(cgImage: CGImage) throws -> [Float] {
157-
let (data, w, h) = try preprocess(cgImage: cgImage)
158-
let pixelCount = w * h
159-
var chw = [Float](repeating: 0, count: 3 * pixelCount)
160-
data.withUnsafeBytes { rawSrc in
161-
let src = rawSrc.bindMemory(to: Float.self)
162-
for c in 0..<3 {
163-
for i in 0..<pixelCount {
164-
chw[c * pixelCount + i] = src[i * 4 + c]
165-
}
166-
}
167-
}
168-
return chw
169-
}
170270
}
171271

172272
// MARK: - Errors

‎swift/Sources/Tools/llm-runner/LLMRunnerMain.swift‎

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import ArgumentParser
77
import CoreAI
88
import CoreAILanguageModels
99
import CoreAIShared
10+
import CoreImage
1011
import Darwin
1112
import Foundation
1213
import Tokenizers
@@ -41,6 +42,14 @@ enum WarmupMode: ExpressibleByArgument, Equatable {
4142
}
4243
}
4344

45+
extension ImageStrategy: ExpressibleByArgument {}
46+
47+
enum ImageInfoMode: String, ExpressibleByArgument, Sendable {
48+
case on
49+
case off
50+
case auto
51+
}
52+
4453
@main
4554
struct Main {
4655
static func main() async throws {
@@ -179,6 +188,16 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
179188
@Option(name: .customLong("image"), help: "Path to an image file for vision-language models")
180189
var imagePath: String?
181190

191+
@Option(
192+
name: .customLong("image-strategy"),
193+
help: "Image preprocessing: stretch, center_crop, or pad (default: from model metadata)")
194+
var imageStrategy: ImageStrategy?
195+
196+
@Option(
197+
name: .customLong("image-info"),
198+
help: "Include original image resolution in prompt: on, off, auto (default: auto)")
199+
var imageInfo: ImageInfoMode = .auto
200+
182201
@Flag(help: "Enable verbose logging")
183202
var verbose: Bool = false
184203

@@ -839,6 +858,24 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
839858
let embeddedInput = try await vlmEngine.encodeImage(at: imageURL)
840859
CLILogger.log("Image encoded: \(embeddedInput.tokenCount) visual tokens", component: "VLM")
841860

861+
let shouldIncludeInfo =
862+
switch imageInfo {
863+
case .on: true
864+
case .off: false
865+
case .auto: visionConfig.includeImageInfo
866+
}
867+
868+
var effectivePrompt = displayPrompt
869+
if shouldIncludeInfo {
870+
let imageURL = URL(fileURLWithPath: imagePath)
871+
if let ciImage = CIImage(contentsOf: imageURL) {
872+
let w = Int(ciImage.extent.width)
873+
let h = Int(ciImage.extent.height)
874+
effectivePrompt = "Image: \(w)x\(h)\n\(displayPrompt)"
875+
CLILogger.log("Injected image info: \(w)x\(h)", component: "VLM")
876+
}
877+
}
878+
842879
// Build VLM prompt with image placeholder tokens.
843880
// Try using the tokenizer's chat template if available; fall back to
844881
// generic "USER: <image>×N \n prompt \nASSISTANT:" format.
@@ -847,7 +884,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
847884
let vlmTokens: [Int32]
848885

849886
if let chatTemplateTokens = try? buildVLMPromptFromChatTemplate(
850-
prompt: displayPrompt,
887+
prompt: effectivePrompt,
851888
imageTokenCount: imageTokenCount,
852889
imageTokenId: imageTokenId,
853890
tokenizer: tokenizer
@@ -860,7 +897,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
860897
component: "VLM")
861898
var tokens = tokenizer.encode(text: "USER: ", addSpecialTokens: true).map { Int32($0) }
862899
tokens.append(contentsOf: [Int32](repeating: imageTokenId, count: imageTokenCount))
863-
let suffix = "\n" + displayPrompt + "\nASSISTANT:"
900+
let suffix = "\n" + effectivePrompt + "\nASSISTANT:"
864901
tokens.append(
865902
contentsOf: tokenizer.encode(text: suffix, addSpecialTokens: false).map { Int32($0) })
866903
vlmTokens = tokens

0 commit comments

Comments
 (0)