Spaces:
Paused
Paused
File size: 2,376 Bytes
f66643d | 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 | """Base class for all AI models."""
from __future__ import annotations
import logging
from abc import ABC, abstractmethod
from typing import Any
from PIL import Image
logger = logging.getLogger(__name__)
class BaseImageToTextModel(ABC):
"""
Shared interface for AI models with Lazy Loading support.
Pipeline: Call -> [Load Model] -> Prepare -> Predict -> Postprocess.
"""
def __init__(self) -> None:
"""Chỉ khai báo các thuộc tính, KHÔNG tải weights vào VRAM ở đây."""
self.model: Any = None
self.device: Any = None
@abstractmethod
def load_model(self) -> None:
"""
Khởi tạo model và đẩy vào VRAM.
Các class con BẮT BUỘC phải override hàm này.
"""
pass
def unload_model(self) -> None:
"""
Unload model from VRAM
"""
if self.model is not None:
import torch
logger.info("Unloading model from VRAM to free memory...")
del self.model
self.model = None
if torch.cuda.is_available():
torch.cuda.empty_cache()
@abstractmethod
def prepare(self, images: list[Image.Image], *args: Any, **kwargs: Any) -> Any:
"""Preprocess raw images."""
pass
@abstractmethod
def predict(self, *args: Any, **kwargs: Any) -> Any:
"""Run core inference. (Model chắc chắn đã được load khi hàm này chạy)."""
pass
@abstractmethod
def postprocess(self, *args: Any, **kwargs: Any) -> Any:
"""Format raw model outputs."""
pass
def __call__(
self,
images: list[Image.Image],
auto_unload: bool = False,
*args: Any,
**kwargs: Any,
) -> Any:
"""
Hàm trung tâm điều phối toàn bộ Pipeline (Template Method).
"""
if self.model is None:
self.load_model()
try:
prepared_inputs = self.prepare(images, *args, **kwargs)
raw_outputs = self.predict(prepared_inputs, *args, **kwargs)
final_results = self.postprocess(raw_outputs, *args, **kwargs)
return final_results
finally:
# 5. Giải phóng VRAM ngay lập tức nếu auto_unload = True
if auto_unload:
self.unload_model()
|