fix: match input dtype to model dtype for TrOCR
Browse files🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude <noreply@anthropic.com>
app.py
CHANGED
|
@@ -70,7 +70,8 @@ def extract_text_from_image(image: Image.Image) -> str:
|
|
| 70 |
"""Stage 1: TrOCR로 이미지에서 텍스트 추출 (OCR)"""
|
| 71 |
try:
|
| 72 |
# TrOCR은 이미지 전체를 한 번에 처리
|
| 73 |
-
pixel_values = OCR_PROCESSOR(image, return_tensors="pt").pixel_values
|
|
|
|
| 74 |
|
| 75 |
with torch.no_grad():
|
| 76 |
generated_ids = OCR_MODEL.generate(pixel_values)
|
|
|
|
| 70 |
"""Stage 1: TrOCR로 이미지에서 텍스트 추출 (OCR)"""
|
| 71 |
try:
|
| 72 |
# TrOCR은 이미지 전체를 한 번에 처리
|
| 73 |
+
pixel_values = OCR_PROCESSOR(image, return_tensors="pt").pixel_values
|
| 74 |
+
pixel_values = pixel_values.to(device=OCR_MODEL.device, dtype=OCR_MODEL.dtype)
|
| 75 |
|
| 76 |
with torch.no_grad():
|
| 77 |
generated_ids = OCR_MODEL.generate(pixel_values)
|