Add video input support for vision-language models

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
 
 

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)
     }
 }