| from PIL import Image |
| import torch |
| from transformers import AutoModel, AutoTokenizer, TextIteratorStreamer |
| from transformers.image_utils import load_image |
| from threading import Thread |
| import logging |
| import time |
| import pynvml |
|
|
| class MiniCPM: |
| def __init__(self, model_id): |
| self.model_id = model_id |
| self.model = AutoModel.from_pretrained( |
| model_id, |
| trust_remote_code=True, |
| attn_implementation='sdpa', |
| torch_dtype=torch.bfloat16 |
| ) |
| self.model = self.model.eval().cuda() |
| self.tokenizer = AutoTokenizer.from_pretrained( |
| model_id, trust_remote_code=True |
| ) |
|
|
| self.handle = None |
| if torch.cuda.is_available(): |
| try: |
| pynvml.nvmlInit() |
| device_id = next(self.model.parameters()).device.index |
| self.handle = pynvml.nvmlDeviceGetHandleByIndex(device_id) |
| except Exception as e: |
| logging.error(f"Failed to initialize NVML: {e}") |
|
|
| def __del__(self): |
| if self.handle: |
| try: |
| pynvml.nvmlShutdown() |
| except: |
| pass |
| |
| def generate(self, video, prompt): |
| start_time = time.time() |
|
|
| images = [Image.open(frame).convert('RGB') for frame in video] |
| content = images + [prompt] |
| msgs = [{'role': 'user', 'content': content}] |
|
|
| |
| res = self.model.chat( |
| image=None, |
| msgs=msgs, |
| tokenizer=self.tokenizer, |
| stream=True |
| ) |
| |
| full_response = "" |
| print("Response: ", end="") |
| first_token_time = None |
| for new_text in res: |
| if first_token_time is None: |
| first_token_time = time.time() |
| full_response += new_text |
| print(new_text, end="", flush=True) |
| print() |
| |
| end_time = time.time() |
| |
| if first_token_time is not None: |
| generation_time = end_time - first_token_time |
| else: |
| generation_time = 0 |
| |
| num_generated_tokens = len(self.tokenizer(full_response).input_ids) |
| tokens_per_second = num_generated_tokens / generation_time if generation_time > 0 else 0 |
|
|
| peak_memory_mb = 0 |
| if self.handle: |
| mem_info = pynvml.nvmlDeviceGetMemoryInfo(self.handle) |
| peak_memory_mb = mem_info.used / (1024 * 1024) |
|
|
| return { |
| "response": full_response, |
| "tokens_per_second": tokens_per_second, |
| "peak_gpu_memory_mb": peak_memory_mb, |
| "num_generated_tokens": num_generated_tokens, |
| } |
|
|
