Update app.py

app.py

@@ -3,97 +3,97 @@ import uvicorn
 from PIL import Image
 from fastapi import FastAPI, UploadFile, File, Response
 
-# Импорты для модели
+# Импорты
 import torch
-from transformers import AutoModelForImageTextToText, AutoProcessor
+from transformers import AutoModelForImageTextToText, AutoTokenizer, AutoImageProcessor
 
-# --- 1. Глобальная загрузка модели (СВЕРХБЫСТРЫЙ СТАРТ) ---
-processor = None
+# --- 1. Глобальная загрузка компонентов ---
 model = None
+tokenizer = None
+image_processor = None
 device = "cpu"
 
 try:
-    print(">>> Инициализация загрузки модели LightOnOCR-1B-1025...")
+    print(">>> Инициализация загрузки LightOnOCR-1B...")
     
-    # Определяем устройство
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f">>> Используемое устройство: {device}")
+    print(f">>> Устройство: {device}")
 
-    # Загружаем процессор (он обрабатывает и картинки, и текст)
-    processor = AutoProcessor.from_pretrained("lightonai/LightOnOCR-1B-1025")
+    repo_id = "lightonai/LightOnOCR-1B-1025"
+
+    # 1. Загружаем токенизатор (для текста)
+    tokenizer = AutoTokenizer.from_pretrained(repo_id)
     
-    # Загружаем саму модель
-    # Если есть GPU, используем float16/bfloat16 для скорости. Если CPU - float32 (по умолчанию)
-    dtype = torch.bfloat16 if device == "cuda" else torch.float32
+    # 2. Загружаем обработчик изображений (для картинок)
+    image_processor = AutoImageProcessor.from_pretrained(repo_id)
     
+    # 3. Загружаем модель
+    dtype = torch.bfloat16 if device == "cuda" else torch.float32
     model = AutoModelForImageTextToText.from_pretrained(
-        "lightonai/LightOnOCR-1B-1025",
+        repo_id,
         torch_dtype=dtype,
         low_cpu_mem_usage=True
     ).to(device)
     
-    print(">>> Модель и процессор успешно загружены!")
+    print(">>> Все компоненты успешно загружены!")
     
 except Exception as e:
-    print(f"КРИТИЧЕСКАЯ ОШИБКА при загрузке модели: {e}")
+    print(f"КРИТИЧЕСКАЯ ОШИБКА загрузки: {e}")
 
 # Инициализация FastAPI
-app = FastAPI(
-    title="LightOnOCR API",
-    description="Прямой инференс через Processor + Model",
-    version="2.0.0",
-)
-
-# --- 2. Эндпоинт API ---
+app = FastAPI(title="LightOnOCR Manual API", version="3.0.0")
 
 @app.post("/api/ocr")
 async def run_ocr(file: UploadFile = File(...)):
-    if model is None or processor is None:
-        return Response(content="Сервер не готов. Модель не загрузилась.", status_code=503)
+    if model is None:
+        return Response(content="Сервер не готов. Модель не загружена.", status_code=503)
 
     try:
-        # 1. Чтение файла
+        # 1. Чтение и конвертация картинки
         contents = await file.read()
         image = Image.open(io.BytesIO(contents)).convert("RGB")
         
-        # 2. Подготовка промпта
-        # Эта модель требует текстовую инструкцию для начала работы
-        prompt_text = "Perform Optical Character Recognition (OCR) on the image. Transcribe all the text."
+        # 2. Обработка изображения (получаем тензоры пикселей)
+        # image_processor вернет словарь с ключом 'pixel_values'
+        vision_outputs = image_processor(images=image, return_tensors="pt")
+        pixel_values = vision_outputs["pixel_values"].to(device)
         
-        # 3. Препроцессинг (Самый важный этап, где мы соединяем картинку и текст)
-        inputs = processor(
-            text=prompt_text,
-            images=image,
-            return_tensors="pt"
-        )
+        # 3. Подготовка текста (Промпта)
+        # Критически важно добавить <image>, чтобы модель знала контекст
+        prompt = "<image>\nTranscribe the text in this image."
         
-        # Переносим тензоры на то же устройство, где модель (GPU или CPU)
-        inputs = {k: v.to(device) for k, v in inputs.items()}
+        text_inputs = tokenizer(prompt, return_tensors="pt")
+        input_ids = text_inputs["input_ids"].to(device)
+        attention_mask = text_inputs["attention_mask"].to(device)
         
-        # 4. Генерация (Собственно OCR)
-        # max_new_tokens ограничивает длину ответа, можно увелич��ть при необходимости
+        # 4. Генерация
+        # Передаем все компоненты в generate
         generated_ids = model.generate(
-            **inputs, 
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            pixel_values=pixel_values,
             max_new_tokens=1024,
-            do_sample=False # False делает ответ детерминированным (стабильным)
+            do_sample=False, # Детерминированный результат (лучше для OCR)
+            pad_token_id=tokenizer.pad_token_id
         )
         
-        # 5. Декодирование ответа
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        # 5. Декодирование
+        generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
         
-        # Очистка результата (иногда модель повторяет промпт в ответе, уберем его если нужно)
-        # Обычно batch_decode возвращает чистый текст, но на всякий случай просто вернем результат
+        # Очистка: иногда модель возвращает сам промпт или артефакты в начале.
+        # Обычно batch_decode(skiыp_special_tokens=True) убирает <image>, но может оставить текст промпта.
+        # Простая очистка (опционально):
+        clean_text = generated_text.replace("Transcribe the text in this image.", "").strip()
         
-        return {"text": generated_text}
+        return {"text": clean_text}
 
     except Exception as e:
-        print(f"Ошибка во время обработки: {e}")
-        # Возвращаем JSON с ошибкой для наглядности в тестере
+        print(f"Ошибка инференса: {e}")
         return Response(content=f"Error: {str(e)}", status_code=500)
 
 @app.get("/")
 async def home():
-    return {"message": "OCR API is running via Processor/Model approach. POST image to /api/ocr"}
+    return {"message": "OCR API Ready. Use POST /api/ocr"}
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)