Update Main.py

Main.py

@@ -5,7 +5,7 @@ import torch
 from transformers import LightOnOcrForConditionalGeneration, LightOnOcrProcessor
 from PIL import Image
 
-# Ускоряем скачивание модели на Hugging Face Spaces
+# Ускоряем скачивание модели
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
 # ==================== Настройки страницы ====================
@@ -35,18 +35,11 @@ st.markdown("""
             border: 1px solid #e5e7eb;
             margin-top: 20px;
         }
-        .stButton > button {
-            background: linear-gradient(90deg, #3b82f6, #1e40af);
-            color: white;
-            border-radius: 12px;
-            padding: 12px 32px;
-            font-weight: 600;
-        }
     </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"
     
@@ -61,7 +54,7 @@ def load_model():
 
     processor = LightOnOcrProcessor.from_pretrained(model_name)
     
-    # Добавляем pad_token, если его нет (важно для генерации)
+    # Исправляем pad_token
     if processor.tokenizer.pad_token is None:
         processor.tokenizer.pad_token = processor.tokenizer.eos_token
 
@@ -83,15 +76,15 @@ def load_image():
 st.markdown('<div class="header-emoji">📄✨</div>', unsafe_allow_html=True)
 st.title("LightOnOCR")
 st.markdown("**Мгновенное распознавание текста с изображений**")
-st.caption("Модель: lightonai/LightOnOCR-1B-1025 • Хорошо работает с английским и документами")
+st.caption("Модель: lightonai/LightOnOCR-1B-1025")
 
-# Загружаем модель
+# Загрузка модели
 processor, model, device, dtype = load_model()
 
-# Информация в боковой панели
+# Сайдбар
 with st.sidebar:
-    st.success(f"✅ Модель успешно загружена на **{device.upper()}**")
-    st.info("LightOnOCR-1B-1025 — современная модель для OCR. Лучше всего работает с английским текстом, документами, чеками и таблицами.")
+    st.success(f"✅ Модель загружена на **{device.upper()}**")
+    st.info("Лучше всего работает с английским текстом, документами и таблицами.")
 
 # Загрузка изображения
 img = load_image()
@@ -103,7 +96,7 @@ if st.button("🔍 Распознать текст", use_container_width=True, t
     else:
         with st.spinner("Распознавание текста... (может занять 5–20 секунд на CPU)"):
             
-            # Правильный формат промпта для LightOnOCR
+            # Правильный промпт
             conversation = [
                 {
                     "role": "user",
@@ -121,16 +114,18 @@ if st.button("🔍 Распознать текст", use_container_width=True, t
                 tokenize=True,
                 return_dict=True,
                 return_tensors="pt",
-                images=img   # Передаём PIL изображение
+                images=img
             )
 
-            # Переносим tensors на правильное устройство
-            inputs = {
+            # Перенос на устройство
+
+
+inputs = {
                 k: (v.to(device=device, dtype=dtype) if v.is_floating_point() else v.to(device))
                 for k, v in inputs.items()
             }
 
-            # Генерация текста
+            # Генерация
             output_ids = model.generate(
                 **inputs,
                 max_new_tokens=2048,
@@ -145,21 +140,19 @@ if st.button("🔍 Распознать текст", use_container_width=True, t
             prompt_length = inputs["input_ids"].shape[1]
             generated_ids = output_ids[0, prompt_length:]
             
-            # Декодируем
             generated_text = processor.decode(
                 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)
             st.markdown('</div>', unsafe_allow_html=True)
 
-            # Кнопка скачивания
             st.download_button(
                 label="💾 Скачать текст как .txt",
                 data=generated_text,
@@ -167,6 +160,6 @@ if st.button("🔍 Распознать текст", use_container_width=True, t
                 mime="text/plain"
             )
 
-# Нижняя подпись
+# Подвал
 st.markdown("---")
 st.caption("Сделано на базе [lightonai/LightOnOCR-1B-1025](https://huggingface.co/lightonai/LightOnOCR-1B-1025)")