Update app.py

app.py

@@ -2,8 +2,6 @@ import io
 import uvicorn
 from PIL import Image
 from fastapi import FastAPI, UploadFile, File, Response
-
-# Импорты
 import torch
 from transformers import AutoModelForImageTextToText, AutoTokenizer, AutoImageProcessor
 
@@ -14,17 +12,18 @@ image_processor = None
 device = "cpu"
 
 try:
-    print(">>> Инициализация загрузки LightOnOCR-1B...")
+    print(">>> Инициализация загрузки LightOnOCR-1B (Fixed VLM pipeline)...")
     
     device = "cuda" if torch.cuda.is_available() else "cpu"
     print(f">>> Устройство: {device}")
 
     repo_id = "lightonai/LightOnOCR-1B-1025"
 
-    # 1. Загружаем токенизатор (для текста)
+    # 1. Загружаем токенизатор
     tokenizer = AutoTokenizer.from_pretrained(repo_id)
     
-    # 2. Загружаем обработчик изображений (для картинок)
+    # 2. Загружаем обработчик изображений
+    # Используем AutoImageProcessor, он должен вернуть правильный класс
     image_processor = AutoImageProcessor.from_pretrained(repo_id)
     
     # 3. Загружаем модель
@@ -40,60 +39,68 @@ try:
 except Exception as e:
     print(f"КРИТИЧЕСКАЯ ОШИБКА загрузки: {e}")
 
-# Инициализация FastAPI
-app = FastAPI(title="LightOnOCR Manual API", version="3.0.0")
+app = FastAPI(title="LightOnOCR Robust API", version="4.0.0")
 
 @app.post("/api/ocr")
 async def run_ocr(file: UploadFile = File(...)):
     if model is None:
-        return Response(content="Сервер не готов. Модель не загружена.", status_code=503)
+        return Response(content="Сервер не готов.", status_code=503)
 
     try:
-        # 1. Чтение и конвертация картинки
+        # 1. Загрузка картинки
         contents = await file.read()
         image = Image.open(io.BytesIO(contents)).convert("RGB")
         
-        # 2. Обработка изображения (получаем тензоры пикселей)
-        # image_processor вернет словарь с ключом 'pixel_values'
+        # 2. Подготовка визуальных данных
+        # ВАЖНО: Мы не просто берем pixel_values, мы берем ВСЕ, что вернет процессор.
+        # Современные модели требуют 'image_sizes' или 'aspect_ratio_ids'.
         vision_outputs = image_processor(images=image, return_tensors="pt")
-        pixel_values = vision_outputs["pixel_values"].to(device)
         
-        # 3. Подготовка текста (Промпта)
-        # Критически важно добавить <image>, чтобы модель знала контекст
+        # Переносим тензоры на устройство (GPU/CPU)
+        # Создаем словарь аргументов для генерации
+        gen_kwargs = {
+            "max_new_tokens": 1024,
+            "do_sample": False,
+            "pad_token_id": tokenizer.pad_token_id
+        }
+
+        # Автоматически добавляем все выходы процессора (pixel_values, image_sizes и т.д.)
+        for key, value in vision_outputs.items():
+            if isinstance(value, torch.Tensor):
+                gen_kwargs[key] = value.to(device)
+            else:
+                gen_kwargs[key] = value
+        
+        # 3. Подготовка текста
+        # Стандартный формат промпта для LLaVA-подобных моделей
         prompt = "<image>\nTranscribe the text in this image."
         
         text_inputs = tokenizer(prompt, return_tensors="pt")
-        input_ids = text_inputs["input_ids"].to(device)
-        attention_mask = text_inputs["attention_mask"].to(device)
+        gen_kwargs["input_ids"] = text_inputs["input_ids"].to(device)
+        gen_kwargs["attention_mask"] = text_inputs["attention_mask"].to(device)
         
         # 4. Генерация
-        # Передаем все компоненты в generate
-        generated_ids = model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            pixel_values=pixel_values,
-            max_new_tokens=1024,
-            do_sample=False, # Детерминированный результат (лучше для OCR)
-            pad_token_id=tokenizer.pad_token_id
-        )
+        # Теперь gen_kwargs содержит и pixel_values, и image_sizes (если они нужны модели)
+        with torch.inference_mode():
+            generated_ids = model.generate(**gen_kwargs)
         
         # 5. Декодирование
         generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
         
-        # Очистка: иногда модель возвращает сам промпт и��и артефакты в начале.
-        # Обычно batch_decode(skiыp_special_tokens=True) убирает <image>, но может оставить текст промпта.
-        # Простая очистка (опционально):
+        # Очистка от артефактов промпта (опционально)
+        # Часто модель возвращает "Transcribe... \n Результат". Уберем промпт.
         clean_text = generated_text.replace("Transcribe the text in this image.", "").strip()
         
         return {"text": clean_text}
 
     except Exception as e:
-        print(f"Ошибка инференса: {e}")
-        return Response(content=f"Error: {str(e)}", status_code=500)
+        import traceback
+        traceback.print_exc() # Печатаем полный лог ошибки в консоль сервера
+        return Response(content=f"Server Error: {str(e)}", status_code=500)
 
 @app.get("/")
 async def home():
-    return {"message": "OCR API Ready. Use POST /api/ocr"}
+    return {"message": "OCR API Ready. POST image to /api/ocr"}
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)