Update app.py

app.py

@@ -2,94 +2,98 @@ import io
 import uvicorn
 from PIL import Image
 from fastapi import FastAPI, UploadFile, File, Response
-from transformers import pipeline
-import torch 
+
+# Импорты для модели
+import torch
+from transformers import AutoModelForImageTextToText, AutoProcessor
 
 # --- 1. Глобальная загрузка модели (СВЕРХБЫСТРЫЙ СТАРТ) ---
-# Модель загружается только один раз при запуске сервера.
-# Используем пайплайн для простоты и эффективности.
+processor = None
+model = None
+device = "cpu"
+
 try:
-    print(">>> Загрузка модели LightOnOCR-1B-1025... Это может занять несколько минут при первом запуске.")
-    # Этот код теперь будет работать, так как torch импортирован!
-    device_to_use = "cuda" if torch.cuda.is_available() else "cpu"
-    dtype_to_use = torch.bfloat16 if device_to_use == "cuda" else None
+    print(">>> Инициализация загрузки модели LightOnOCR-1B-1025...")
+    
+    # Определяем устройство
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f">>> Используемое устройство: {device}")
+
+    # Загружаем процессор (он обрабатывает и картинки, и текст)
+    processor = AutoProcessor.from_pretrained("lightonai/LightOnOCR-1B-1025")
+    
+    # Загружаем саму модель
+    # Если есть GPU, используем float16/bfloat16 для скорости. Если CPU - float32 (по умолчанию)
+    dtype = torch.bfloat16 if device == "cuda" else torch.float32
+    
+    model = AutoModelForImageTextToText.from_pretrained(
+        "lightonai/LightOnOCR-1B-1025",
+        torch_dtype=dtype,
+        low_cpu_mem_usage=True
+    ).to(device)
+    
+    print(">>> Модель и процессор успешно загружены!")
     
-    ocr_pipeline = pipeline(
-        "image-to-text",
-        model="lightonai/LightOnOCR-1B-1025",
-        device=device_to_use, 
-        torch_dtype=dtype_to_use,
-    )
-    print(">>> Модель успешно загружена!")
 except Exception as e:
-    # Ошибка 503 возникает, потому что код попадает сюда.
-    print(f"Ошибка при загрузке модели: {e}")
-    ocr_pipeline = None
+    print(f"КРИТИЧЕСКАЯ ОШИБКА при загрузке модели: {e}")
 
 # Инициализация FastAPI
 app = FastAPI(
-    title="LightOnOCR Super-Fast API",
-    description="Публичный API для высокоскоростного распознавания текста.",
-    version="1.0.0",
+    title="LightOnOCR API",
+    description="Прямой инференс через Processor + Model",
+    version="2.0.0",
 )
 
 # --- 2. Эндпоинт API ---
 
-# ... (остальной код)
-
-@app.post(
-    "/api/ocr",
-    summary="Запустить OCR по изображению",
-    response_description="Расшифрованный текст",
-)
-async def run_ocr(file: UploadFile = File(..., description="Изображение для расшифровки")):
-    """
-    Принимает изображение (JPG, PNG и т.д.) и возвращает распознанный текст.
-    """
-    if ocr_pipeline is None:
-        return Response(content="Сервер не готов. Модель не загружена.", status_code=503)
+@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)
 
     try:
-        # 1. Чтение файла в память
+        # 1. Чтение файла
         contents = await file.read()
-        
-        # 2. Конвертация байтов в объект PIL Image
-        # Важно: .convert("RGB") гарантирует 3 канала, что часто требуется для VLM.
         image = Image.open(io.BytesIO(contents)).convert("RGB")
         
-        # --- ИСПРАВЛЕНИЕ: ДОБАВЛЕНИЕ ПРОМПТА И ИЗМЕНЕНИЕ ФОРМАТА ВЫЗОВА ---
+        # 2. Подготовка промпта
+        # Эта модель требует текстовую инструкцию для начала работы
+        prompt_text = "Perform Optical Character Recognition (OCR) on the image. Transcribe all the text."
+        
+        # 3. Препроцессинг (Самый важный этап, где мы соединяем картинку и текст)
+        inputs = processor(
+            text=prompt_text,
+            images=image,
+            return_tensors="pt"
+        )
         
-        # Рекомендуемый промпт для OCR с этой моделью 
-        # (взято из документации LightOnOCR)
-        ocr_prompt = "Perform Optical Character Recognition (OCR) on the image. Transcribe all the text."
+        # Переносим тензоры на то же устройство, где модель (GPU или CPU)
+        inputs = {k: v.to(device) for k, v in inputs.items()}
         
-        # 3. Запуск пайплайна с использованием кортежа (изображение, промпт)
-        results = ocr_pipeline(
-            (image, ocr_prompt), 
-            generate_kwargs={"max_new_tokens": 1024}
+        # 4. Генерация (Собственно OCR)
+        # max_new_tokens ограничивает длину ответа, можно увеличить при необходимости
+        generated_ids = model.generate(
+            **inputs, 
+            max_new_tokens=1024,
+            do_sample=False # False делает ответ детерминированным (стабильным)
         )
         
-        # 4. Извлечение текста (остается прежним)
-        if results and isinstance(results, list) and 'generated_text' in results[0]:
-            decoded_text = results[0]['generated_text']
-            return {"text": decoded_text}
-        else:
-            return {"text": "Ошибка: Не удалось распознать текст или результат пуст."}
+        # 5. Декодирование ответа
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        
+        # Очистка результата (иногда модель повторяет промпт в ответе, уберем его если нужно)
+        # Обычно batch_decode возвращает чистый текст, но на всякий случай просто вернем результат
+        
+        return {"text": generated_text}
 
     except Exception as e:
-        # Логирование ошибок для отладки
-        print(f"Ошибка обработки запроса: {e}")
-        return Response(content=f"Внутренняя ошибка сервера: {str(e)}", status_code=500)
-
-# --- 3. Главная страница (для совместимости с HF Space) ---
+        print(f"Ошибка во время обработки: {e}")
+        # Возвращаем JSON с ошибкой для наглядности в тестере
+        return Response(content=f"Error: {str(e)}", status_code=500)
 
-@app.get("/", include_in_schema=False)
+@app.get("/")
 async def home():
-    """Перенаправление на документацию API."""
-    return {"message": "API запущен. Используйте эндпоинт /api/ocr (POST) или посмотрите документацию по /docs"}
+    return {"message": "OCR API is running via Processor/Model approach. POST image to /api/ocr"}
 
-# --- 4. Запуск сервера (для локального тестирования) ---
 if __name__ == "__main__":
-    # На Hugging Face Space этот блок не запускается, 
-    # сервер запускается через команду uvicorn (см. ниже)
     uvicorn.run(app, host="0.0.0.0", port=7860)