inference/inference_runner.py

"""
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()