File size: 7,620 Bytes
7a0c684 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 |
"""
Multi-modal model inference runner using Helium virtual GPU.
"""
import os
import sys
from pathlib import Path
import json
import time
import numpy as np
import cv2
import soundfile as sf
from typing import Dict, List, Optional, Union, Any
# Add parent directory and virtual_gpu_driver to path
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, root_dir)
sys.path.insert(0, os.path.join(root_dir, 'virtual_gpu_driver'))
from inference.app import MultiModalModel
class InferenceRunner:
"""Handles model inference with efficient batching and caching"""
def __init__(
self,
model_id: str = "openai-oss-20b",
device_id: str = "vgpu0",
batch_size: int = 1,
cache_dir: Optional[str] = None
):
# Load model from HuggingFace and store in device DB
self.model = MultiModalModel.from_pretrained(
model_id=model_id,
device_id=device_id,
cache_dir=cache_dir
)
self.model.eval() # Set to inference mode
self.batch_size = batch_size
self.cache_dir = cache_dir
if cache_dir:
os.makedirs(cache_dir, exist_ok=True)
# Batch buffers
self._text_batch = []
self._image_batch = []
self._audio_batch = []
def preprocess_text(self, text: Union[str, List[str]]) -> np.ndarray:
"""Preprocess text input"""
if isinstance(text, str):
text = [text]
# Convert to token IDs (normally would use tokenizer)
# This is just a placeholder - replace with actual tokenization
text_data = np.zeros((len(text), 64), dtype=np.int32)
return text_data
def preprocess_image(self, image_path: str) -> np.ndarray:
"""Preprocess image input"""
# Load and preprocess image
image = cv2.imread(image_path)
if image is None:
raise ValueError(f"Could not load image: {image_path}")
# Resize
image = cv2.resize(image, (224, 224))
# Convert to RGB and normalize
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = image.astype(np.float32) / 127.5 - 1.0
# To NCHW format
image = image.transpose(2, 0, 1)[None]
return image
def preprocess_audio(self, audio_path: str) -> np.ndarray:
"""Preprocess audio input"""
# Load audio file
audio, sr = sf.read(audio_path)
# Convert to mono if stereo
if len(audio.shape) > 1:
audio = audio.mean(axis=1)
# Normalize
audio = audio / np.abs(audio).max()
# Pad or truncate to fixed length (16000 samples = 1 second)
target_len = 16000
if len(audio) > target_len:
audio = audio[:target_len]
else:
audio = np.pad(audio, (0, target_len - len(audio)))
return audio[None, None] # Add batch and channel dims
def add_to_batch(
self,
text: Optional[Union[str, List[str]]] = None,
image_path: Optional[str] = None,
audio_path: Optional[str] = None
) -> None:
"""Add inputs to batch for processing"""
if text is not None:
text_data = self.preprocess_text(text)
self._text_batch.append(text_data)
if image_path is not None:
image_data = self.preprocess_image(image_path)
self._image_batch.append(image_data)
if audio_path is not None:
audio_data = self.preprocess_audio(audio_path)
self._audio_batch.append(audio_data)
# Process batch if full
if len(self._text_batch) >= self.batch_size:
self.process_batch()
def process_batch(self) -> Dict[str, np.ndarray]:
"""Process current batch through model"""
if not any([self._text_batch, self._image_batch, self._audio_batch]):
return {}
# Prepare inputs
inputs = {}
if self._text_batch:
inputs["text"] = np.concatenate(self._text_batch, axis=0)
if self._image_batch:
inputs["image"] = np.concatenate(self._image_batch, axis=0)
if self._audio_batch:
inputs["audio"] = np.concatenate(self._audio_batch, axis=0)
# Run inference
outputs = self.model(inputs)
# Cache results if needed
if self.cache_dir:
timestamp = str(int(time.time()))
cache_path = os.path.join(self.cache_dir, f"batch_{timestamp}.npz")
np.savez(cache_path, **outputs)
# Clear batch buffers
self._text_batch.clear()
self._image_batch.clear()
self._audio_batch.clear()
return outputs
def generate_from_context(
self,
context_text: Optional[str] = None,
context_image: Optional[str] = None,
context_audio: Optional[str] = None,
max_length: int = 100
) -> np.ndarray:
"""Generate sequence using multi-modal context"""
inputs = {}
if context_text:
inputs["text"] = self.preprocess_text(context_text)
if context_image:
inputs["image"] = self.preprocess_image(context_image)
if context_audio:
inputs["audio"] = self.preprocess_audio(context_audio)
return self.model.generate(inputs, max_length=max_length)
def __call__(
self,
text: Optional[Union[str, List[str]]] = None,
image_path: Optional[str] = None,
audio_path: Optional[str] = None,
return_dict: bool = True
) -> Union[np.ndarray, Dict[str, np.ndarray]]:
"""Run inference on inputs"""
self.add_to_batch(text, image_path, audio_path)
return self.process_batch()
def cleanup(self):
"""Clean up resources"""
self.model.cleanup()
if self.cache_dir and os.path.exists(self.cache_dir):
import shutil
shutil.rmtree(self.cache_dir)
def main():
"""Example usage"""
# Initialize runner
runner = InferenceRunner(
model_id="openai-oss-20b",
device_id="vgpu0",
batch_size=4,
cache_dir="inference_cache"
)
# Single inference
outputs = runner(
text="A photo of a cat",
image_path="cat.jpg",
audio_path="meow.wav"
)
print("Single inference outputs:", outputs.keys())
# Batch processing
for i in range(10):
runner.add_to_batch(
text=f"Sample text {i}",
image_path=f"image_{i}.jpg",
audio_path=f"audio_{i}.wav"
)
# Process remaining items in batch
final_outputs = runner.process_batch()
print("Batch processing outputs:", final_outputs.keys())
# Generate from multi-modal context
generated = runner.generate_from_context(
context_text="Describe this image and sound:",
context_image="scene.jpg",
context_audio="ambience.wav",
max_length=50
)
print("Generated sequence shape:", generated.shape)
# Cleanup
runner.cleanup()
if __name__ == "__main__":
main()
|