Update Main.py

Main.py

@@ -5,7 +5,7 @@ import torch
 from transformers import LightOnOcrForConditionalGeneration, LightOnOcrProcessor
 from PIL import Image
 
-# Ускоряем скачивание на HF Spaces
+# Ускоряем скачивание
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
 st.set_page_config(
@@ -15,22 +15,22 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
-# Простой CSS
 st.markdown("""
-<style>
-    .main { background: linear-gradient(180deg, #f8f9fa, #e9f0f7); }
-    .result-box {
-        background: #ffffff;
-        border-radius: 16px;
-        padding: 24px;
-        box-shadow: 0 10px 30px rgba(0,0,0,0.08);
-        margin-top: 20px;
-    }
-    .header-emoji { font-size: 3.5rem; text-align: center; margin: 15px 0; }
-</style>
+    <style>
+        .main { background: linear-gradient(180deg, #f8f9fa, #e9f0f7); }
+        .header-emoji { font-size: 3.5rem; text-align: center; margin: 15px 0; }
+        .result-box {
+            background: #ffffff;
+            border-radius: 16px;
+            padding: 24px;
+            box-shadow: 0 10px 30px rgba(0, 0, 0, 0.08);
+            border: 1px solid #e5e7eb;
+            margin-top: 20px;
+        }
+    </style>
 """, unsafe_allow_html=True)
 
-@st.cache_resource(show_spinner="⏳ Загрузка модели LightOnOCR-1B-1025...\nЭто может занять 2–6 минут при первом запуске на CPU")
+@st.cache_resource(show_spinner="⏳ Загрузка модели LightOnOCR-1B-1025...\n(2–6 минут при первом запуске)")
 def load_model():
     model_name = "lightonai/LightOnOCR-1B-1025"
     
@@ -44,21 +44,12 @@ def load_model():
     ).to(device)
 
     processor = LightOnOcrProcessor.from_pretrained(model_name)
-    
-    return processor, model, device, dtype
-
-# ====================== Заголовок ======================
-st.markdown('<div class="header-emoji">📄✨</div>', unsafe_allow_html=True)
-st.title("LightOnOCR")
-st.markdown("**Распознавание текста с изображений**")
-st.caption("Модель: lightonai/LightOnOCR-1B-1025")
 
-# ====================== Загрузка модели ======================
-processor, model, device, dtype = load_model()
+    if processor.tokenizer.pad_token is None:
+        processor.tokenizer.pad_token = processor.tokenizer.eos_token
 
-st.sidebar.success(f"✅ Модель загружена на **{device.upper()}**")
+    return processor, model, device, dtype
 
-# ====================== Загрузка изображения ======================
 def load_image():
     uploaded_file = st.file_uploader(
         "📸 Загрузите изображение (png, jpg, jpeg, webp)",
@@ -70,69 +61,84 @@ def load_image():
         return Image.open(io.BytesIO(image_data)).convert('RGB')
     return None
 
+# ==================== Интерфейс ====================
+st.markdown('<div class="header-emoji">📄✨</div>', unsafe_allow_html=True)
+st.title("LightOnOCR")
+st.markdown("**Распознавание текста с изображений**")
+st.caption("Модель: lightonai/LightOnOCR-1B-1025")
+
+processor, model, device, dtype = load_model()
+
+with st.sidebar:
+    st.success(f"✅ Модель загружена на **{device.upper()}**")
+
 img = load_image()
 
-# ====================== Распознавание ======================
+# ==================== Распознавание ====================
 if st.button("🔍 Распознать текст", use_container_width=True, type="primary"):
     if img is None:
         st.error("Сначала загрузите изображение")
     else:
-        with st.spinner("Распознавание текста... (может занять 5–20 сек на CPU)"):
+        with st.spinner("Распознавание текста... (5–20 сек на CPU)"):
             
-            # ✅ Правильный формат для LightOnOCR (только изображение + промпт)
+            # Правильный формат разговора (без передачи images здесь)
             conversation = [
                 {
                     "role": "user",
                     "content": [
-                        {"type": "image"},   # изображение передаётся автоматически через processor
-                        {"type": "text", "text": "Extract all the text from this image as accurately as possible. Output clean text with preserved line breaks and formatting."}
+                        {"type": "image"},
+                        {"type": "text", "text": "Extract all the text from this image accurately. Preserve original formatting, tables, and line breaks as much as possible."}
                     ]
                 }
             ]
 
-            # Подготовка inputs (processor обработает и изображение, и текст)
+            # Применяем шаблон чата
             inputs = processor.apply_chat_template(
                 conversation,
                 add_generation_prompt=True,
                 tokenize=True,
                 return_dict=True,
-                return_tensors="pt",
-                # Важно: передаём само PIL-изображение
-                images=img
+                return_tensors="pt"
             )
 
-            # Переносим на устройство
-            inputs = {
-                k: (v.to(device=device, dtype=dtype) if v.is_floating_point() else v.to(device))
-                for k, v in inputs.items()
-            }
+            # Важно: добавляем pixel_values отдельно
+            pixel_values = processor.image_processor(img, return_tensors="pt").pixel_values
+            inputs["pixel_values"] = pixel_values.to(device=device, dtype=dtype)
+
+            # Переносим остальные тензоры
+            for k, v in inputs.items():
+                if isinstance(v, torch.Tensor) and k != "pixel_values":
+                    inputs[k] = v.to(device=device)
 
             # Генерация
             output_ids = model.generate(
                 **inputs,
                 max_new_tokens=2048,
-                do_sample=False,
+                do_sa
+
+
+mple=False,
                 temperature=0.0,
-                num_beams=1,          # для стабильности
+                num_beams=1,
                 pad_token_id=processor.tokenizer.pad_token_id,
                 eos_token_id=processor.tokenizer.eos_token_id,
             )
 
-            # Убираем входной промпт — оставляем только сгенерированный текст
+            # Убираем промпт
             prompt_length = inputs["input_ids"].shape[1]
             generated_ids = output_ids[0, prompt_length:]
             
             generated_text = processor.decode(
-                generated_ids, 
+                generated_ids,
                 skip_special_tokens=True,
                 clean_up_tokenization_spaces=True
             ).strip()
 
-            # Вывод результата
+            # Результат
             st.success("✅ Распознавание завершено!")
             st.markdown('<div class="result-box">', unsafe_allow_html=True)
             st.subheader("📝 Распознанный текст")
-            st.code(generated_text, language=None)   # лучше чем markdown для больших блоков
+            st.code(generated_text, language=None)
             st.markdown('</div>', unsafe_allow_html=True)
 
             st.download_button(