fixed code v2
This commit is contained in:
@@ -21,7 +21,6 @@ class Settings(BaseSettings):
|
||||
NSFW_THRESHOLD: float = 0.85
|
||||
IMAGE_CLIP_MODEL: str = "openai/clip-vit-base-patch32"
|
||||
CLIP_THRESHOLD: float = 0.75
|
||||
HF_TOKEN: str = "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"
|
||||
MODEL_CACHE_DIR: str = "./models"
|
||||
|
||||
DEVICE: str = "auto"
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
from fastapi import Request
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from app.exceptions.invalid_text_exception import InvalidTextException
|
||||
from app.exceptions.invalid_image_exception import InvalidImageException
|
||||
from app.exceptions.invalid_text_exception import InvalidTextException
|
||||
from app.models.response.error_response import ErrorResponse
|
||||
|
||||
|
||||
|
||||
@@ -1,17 +1,13 @@
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from app.config.logging import logger
|
||||
from app.config.settings import settings
|
||||
from app.ml.model_manager import model_manager
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
if settings.HF_TOKEN:
|
||||
os.environ["HF_TOKEN"] = settings.HF_TOKEN
|
||||
|
||||
logger.info("Loading AI models...")
|
||||
model_manager.initialize()
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import torch
|
||||
import os
|
||||
import torch
|
||||
from transformers import AutoImageProcessor
|
||||
from transformers import AutoModelForImageClassification
|
||||
from transformers import AutoModelForSequenceClassification
|
||||
@@ -17,6 +18,11 @@ class ModelManager:
|
||||
self.image_model = None
|
||||
self.image_processor = None
|
||||
|
||||
def _setup_hf_token(self):
|
||||
if settings.HF_TOKEN:
|
||||
os.environ["HF_TOKEN"] = settings.HF_TOKEN
|
||||
os.environ["HUGGING_FACE_HUB_TOKEN"] = settings.HF_TOKEN
|
||||
|
||||
def load_device(self):
|
||||
if settings.DEVICE == "auto":
|
||||
self.device = torch.device(
|
||||
|
||||
Reference in New Issue
Block a user