Merge pull request #23 from DropThe8bit/chore/3

[chore] Dockerfile GPU 베이스/HF 캐시/Yolo 모델 관련 설정

Author: hhongyeahh

Dockerfile

@@ -1,15 +1,30 @@
-# 1. Python 3.10 기반 슬림 이미지 사용
-FROM python:3.10-slim
+# 1. 권장 베이스: CUDA 포함 PyTorch 런타임
+FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime
+
+# OpenCV 헤드리스 등에 필요한 OS 패키지
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    libglib2.0-0 libsm6 libxext6 libxrender1 && rm -rf /var/lib/apt/lists/*
+
+# HF 캐시 + PyTorch 메모리 튜닝(선택)
+ENV HF_HOME=/models/hf-cache \
+    HUGGINGFACE_HUB_CACHE=/models/hf-cache \
+    PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
 
 # 2. 작업 디렉토리 생성
 WORKDIR /app
 
-# 3. requirements.txt 먼저 복사하고 설치
+# 3. 의존성만 먼저 복사 → 캐시 최대 활용
 COPY everTale/requirements.txt ./requirements.txt
-RUN pip install --no-cache-dir -r requirements.txt
+RUN pip install --no-cache-dir --upgrade pip \
+ && pip install --no-cache-dir -r /app/requirements.txt \
+ && pip install --no-cache-dir opencv-python-headless ultralytics
+
+# 모델 파일을 이미지에 포함
+COPY models/yolov8n.pt /models/my_yolo_model.pt
+ENV YOLO_MODEL_PATH=/models/my_yolo_model.pt
 
 # 4. 전체 코드 복사
-COPY . .
+COPY . /app
 
 # 5. 컨테이너가 열 포트 설정
 EXPOSE 8000

everTale/app/main.py

@@ -20,3 +20,13 @@
 @app.get("/")
 def root():
     return {"message": "EverTale AI 서버가 정상 실행 중입니다."}
+
+@app.get("/debug/cuda")
+def cuda_health():
+    import torch, os
+    return {
+        "cuda_available": torch.cuda.is_available(),
+        "device_count": torch.cuda.device_count(),
+        "current_device": torch.cuda.current_device() if torch.cuda.is_available() else None,
+        "visible_devices": os.getenv("NVIDIA_VISIBLE_DEVICES", "")
+    }

everTale/app/service/yolo_service.py

@@ -9,14 +9,35 @@
 
 YOLO_MODEL_PATH = os.environ["YOLO_MODEL_PATH"]
 
+from ultralytics import YOLO
+import os, torch
+
+def _resolve_yolo_path() -> str:
+    path = os.getenv("YOLO_MODEL_PATH", "/models/my_yolo_model.pt")
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"YOLO model not found at: {path}")
+    return path
+
+def _require_gpu_for_yolo(stage: str = "YOLO load"):
+    if torch.cuda.is_available():
+        return 0  # device index for CUDA
+    # MPS는 Ultralytics 지원이 제한적이므로 필요한 경우만 허용
+    if torch.backends.mps.is_available() and torch.backends.mps.is_built():
+        return "mps"
+    raise RuntimeError(f"[ERROR] No GPU backend during {stage}. CPU is not allowed for YOLO.")
+
 def load_model() -> YOLO:
+    path = _resolve_yolo_path()
+    device = _require_gpu_for_yolo("YOLO load")
     try:
-        model = YOLO(YOLO_MODEL_PATH)
-        print("모델이 성공적으로 로드되었습니다.")
+        model = YOLO(path)
+        # warm-up(선택): 작은 더미로 한 번 실행해 메모리 로딩
+        model.predict(source=np.zeros((64,64,3), dtype=np.uint8), device=device, imgsz=64, verbose=False)
+        print(f"[INFO] YOLO loaded on device={device} from {path}")
         return model
     except Exception as e:
-        print(f"모델 로드 중 오류가 발생했습니다: {e}")
-        return None
+        raise RuntimeError(f"Failed to load YOLO model at {path}: {e}")
+
 
 def _url_to_bgr(url: str) -> np.ndarray:
     resp = requests.get(url, timeout=10)
@@ -35,13 +56,19 @@ def detect_object(image_paths: List[str]) -> Dict[str, Any]:
           탐지 후보가 전혀 없으면 {"index": None, "url": None, "detection": None}
     """
     model = load_model()
+    device = 0 if torch.cuda.is_available() else "mps"  # 위와 일치
     urls = image_paths[:8]
     candidates: List[Dict[str, Any]] = []
 
     for idx, url in enumerate(urls):
         try:
             img = _url_to_bgr(url)
-            results = model.predict(source=img, verbose=False)
+            results = model.predict(
+                source=img,
+                device=device,
+                half=torch.cuda.is_available(),
+                verbose=False
+            )
             if not results or results[0].boxes is None or results[0].boxes.shape[0] == 0:
                 continue
 

everTale/requirements.txt

@@ -1,27 +1,39 @@
-accelerate==1.8.1
-boto3==1.38.45
-diffusers==0.33.1
-fastapi==0.115.13
-fsspec==2025.5.1
-huggingface-hub==0.33.0
-Jinja2==3.1.6
-numpy==2.0.2
-openai==1.89.0
-pillow==11.2.1
-psutil==7.0.0
-pydantic==2.11.7
-python-dotenv==1.1.0
-python-multipart==0.0.20
+# --- PyTorch stack (이미지에 포함되지만 버전 명시해두면 좋음) ---
+torch==2.3.1
+torchvision==0.18.1
+torchaudio==2.3.1
+
+# --- Diffusers/Transformers/Accelerate (안정 조합) ---
+diffusers==0.29.0
+transformers==4.42.4
+accelerate==0.32.1
+safetensors==0.4.3
+huggingface-hub==0.23.4
+tokenizers==0.19.1
+
+# --- API / 서버 ---
+fastapi==0.115.6        # (너가 쓰는 0.115.x 라인 유지, 너무 최신 patch는 피함)
+uvicorn==0.30.6         # 안정 버전
+starlette==0.38.5       # fastapi 0.115.x와 검증된 조합
+
+# --- 유틸 ---
+numpy==1.26.4           # torch 2.3.x와 널리 쓰이는 안정 버전
+pillow==10.4.0
+requests==2.32.3
+tqdm==4.66.4
+psutil==5.9.8
+python-dotenv==1.0.1
+python-multipart==0.0.9
 PyYAML==6.0.2
-regex==2024.11.6
-requests==2.32.4
-safetensors==0.5.3
-starlette==0.46.2
-tokenizers==0.21.1
-torch==2.7.1
-tqdm==4.67.1
-transformers==4.52.4
-typing-inspection==0.4.1
-typing_extensions==4.14.0
-uvicorn==0.34.3
+regex==2024.5.15
+typing_extensions==4.12.2
+fsspec==2024.6.1
+ultralytics
+
+# --- 선택(LLM LoRA 등에 필요하면) ---
+peft==0.11.1
+
+# --- 필요 시만 추가 (성공/실패 갈릴 수 있으므로 초기엔 제외 권장) ---
+# xformers==0.0.27.post2
+