Update Main.py

Main.py

@@ -1,167 +1,164 @@
-Import io
-Import streamlit as st
-Import torch
-From transformers import LightOnOcrForConditionalGeneration, LightOnOcrProcessor
-From PIL import Image
-
-# Стильное оформление приложения
-St.set_page_config(
-    Page_title=”LightOnOCR • Распознай текст”,
-    Page_icon=””,
-    Layout=”centered”,
-    Initial_sidebar_state=”expanded”
+import io
+import streamlit as st
+import torch
+from transformers import LightOnOcrForConditionalGeneration, LightOnOcrProcessor
+from PIL import Image
+
+# ==================== Настройки страницы ====================
+st.set_page_config(
+    page_title="LightOnOCR • Распознай текст",
+    page_icon="📄",
+    layout="centered",
+    initial_sidebar_state="expanded"
 )
 
-# Кастомный CSS для современного и стильного вида
-St.markdown(“””
+# ==================== Кастомный CSS ====================
+st.markdown("""
     <style>
-        .main {
-            Background: linear-gradient(180deg, #f8f9fa, #e9f0f7);
-        }
-        .stApp {
-            Max-width: 1200px;
-            Margin: 0 auto;
-        }
-        H1 {
-            Font-family: ‘Segoe UI’, sans-serif;
-            Color: #1e3a8a;
-            Text-align: center;
-            Margin-bottom: 0.2rem;
-        }
+        .main { background: linear-gradient(180deg, #f8f9fa, #e9f0f7); }
+        h1 { color: #1e3a8a; text-align: center; margin-bottom: 0.2rem; }
         .stButton > button {
-            Background: linear-gradient(90deg, #3b82f6, #1e40af);
-            Color: white;
-            Border-radius: 12px;
-            Padding: 12px 32px;
-            Font-weight: 600;
-            Border: none;
-            Box-shadow: 0 4px 15px rgba(59, 130, 246, 0.3);
-            Transition: all 0.3s ease;
+            background: linear-gradient(90deg, #3b82f6, #1e40af);
+            color: white;
+            border-radius: 12px;
+            padding: 12px 32px;
+            font-weight: 600;
+            border: none;
+            box-shadow: 0 4px 15px rgba(59, 130, 246, 0.3);
         }
         .stButton > button:hover {
-            Transform: translateY(-2px);
-            Box-shadow: 0 8px 20px rgba(59, 130, 246, 0.4);
+            transform: translateY(-2px);
+            box-shadow: 0 8px 20px rgba(59, 130, 246, 0.4);
         }
         .result-box {
-            Background: #ffffff;
-            Border-radius: 16px;
-            Padding: 24px;
-            Box-shadow: 0 10px 30px rgba(0, 0, 0, 0.08);
-            Border: 1px solid #e5e7eb;
-        }
-        .header-emoji {
-            Font-size: 3rem;
-            Display: block;
-            Text-align: center;
-            Margin-bottom: 10px;
+            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;
         }
+        .header-emoji { font-size: 3.5rem; text-align: center; margin: 10px 0; }
     </style>
-“””, unsafe_allow_html=True)
+""", unsafe_allow_html=True)
 
-@st.cache_resource(show_spinner=False)
-Def load_model():
-    “””Загрузка модели LightOnOCR-1B-1025 (один раз)”””
-    Model_name = “lightonai/LightOnOCR-1B-1025”
-    
-    # Автоопределение устройства и типа данных
-    If torch.backends.mps.is_available():
-        Device = “mps”
-        Dtype = torch.float32
-    Elif torch.cuda.is_available():
-        Device = “cuda”
-        Dtype = torch.bfloat16
-    Else:
-        Device = “cpu”
-        Dtype = torch.float32
+# ==================== Загрузка модели ====================
+@st.cache_resource(show_spinner="Загрузка модели LightOnOCR-1B-1025...")
+def load_model():
+    model_name = "lightonai/LightOnOCR-1B-1025"
     
-    Model = LightOnOcrForConditionalGeneration.from_pretrained(
-        Model_name,
-        Torch_dtype=dtype,
-        Trust_remote_code=True
+    if torch.backends.mps.is_available():
+        device = "mps"
+        dtype = torch.float32
+    elif torch.cuda.is_available():
+        device = "cuda"
+        dtype = torch.bfloat16
+    else:
+        device = "cpu"
+        dtype = torch.float32
+
+    model = LightOnOcrForConditionalGeneration.from_pretrained(
+        model_name,
+        torch_dtype=dtype,
+        trust_remote_code=True,
+        device_map=None  # загружаем вручную
     ).to(device)
-    
-    Processor = LightOnOcrProcessor.from_pretrained(model_name)
-    
-    Return processor, model, device, dtype
 
-Def load_image():
-    “””Загрузка изображения”””
-    Uploaded_file = st.file_uploader(
-        “ Загрузите изображение (фото, скан, документ)”,
-        Type=[‘png’, ‘jpg’, ‘jpeg’, ‘webp’]
+    processor = LightOnOcrProcessor.from_pretrained(model_name)
+
+    return processor, model, device, dtype
+
+# ==================== Загрузка изображения ====================
+def load_image():
+    uploaded_file = st.file_uploader(
+        "📸 Загрузите изображение (фото, скан, документ)",
+        type=['png', 'jpg', 'jpeg', 'webp']
     )
-    If uploaded_file is not None:
-        Image_data = uploaded_file.getvalue()
-        St.image(image_data, use_container_width=True, caption=”Загруженное изображение”)
-        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(“Современная end-to-end нейросеть LightOnOCR-1B-1025 • Поддерживает документы, чеки, фото, таблицы и сложную вёрстку”)
-
-# Загрузка модели
-Processor, model, device, dtype = load_model()
-
-# Инфо в сайдбаре
-With st.sidebar:
-    St.markdown(“###  О модели”)
-    St.info(“LightOnOCR-1B-1025 — компактная, но очень точная модель для OCR. Отлично работает с английским, латиницей, документами и сложными макетами.”)
-    St.markdown(“**Поддержка:** Английский + 8 других языков (латиница)”)
-    St.markdown(“**Скорость:** до 5+ страниц/сек на GPU”)
-    St.caption(f”Устройство: **{device.upper()}** • dtype: **{dtype}**”)
-
-# Основной интерфейс
-Img = load_image()
-
-If st.button(“ Распознать текст”, use_container_width=True, type=”primary”):
-    If img is None:
-        St.error(“Пожалуйста, загрузите изображение”)
-    Else:
-        With st.spinner(“Распознавание текста… Это может занять несколько секунд (особенно на CPU)”):
-            # Подготовка промпта
-            Prompt = “Extract all the text from this image accurately. Preserve original formatting, layout, tables and line breaks as much as possible.”
-            
-            # Подготовка входных данных
-            Inputs = processor(images=img, text=prompt, 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()
+    if uploaded_file is not None:
+        image_data = uploaded_file.getvalue()
+        st.image(image_data, use_container_width=True, caption="Загруженное изображение")
+        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("Модель LightOnOCR-1B-1025 • Отлично работает с документами, чеками, таблицами и фото")
+
+# Загружаем модель один раз
+processor, model, device, dtype = load_model()
+
+# Сайдбар
+with st.sidebar:
+    st.markdown("### 🚀 О модели")
+    st.info("LightOnOCR-1B-1025 — компактная end-to-end модель для OCR и понимания документов.")
+    st.markdown("**Поддержка:** Английский + латиница, таблицы, сложная вёрстка")
+    st.caption(f"Устройство: **{device.upper()}** • dtype: **{dtype}**")
+
+# Загрузка изображения
+img = load_image()
+
+# Кнопка распознавания
+if st.button("🔍 Распознать текст", use_container_width=True, type="primary"):
+    if img is None:
+        st.error("Пожалуйста, сначала загрузите изображение")
+    else:
+        with st.spinner("Распознавание текста… (на CPU может занять 10–30 секунд)"):
+            # Правильный способ работы с этой моделью (chat template)
+            conversation = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image"},
+                        {"type": "te
+
+
+xt", "text": "Extract all the text from this image accurately. Preserve formatting, tables, and line breaks as much as possible."}
+                    ]
+                }
+            ]
+
+            inputs = processor.apply_chat_template(
+                conversation,
+                add_generation_prompt=True,
+                tokenize=True,
+                return_dict=True,
+                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()
             }
-            
+
             # Генерация
-            Output_ids = model.generate(
+            output_ids = model.generate(
                 **inputs,
-                Max_new_tokens=2048,
-                Do_sample=False,
-                Temperature=0.0
+                max_new_tokens=2048,
+                do_sample=False,
+                temperature=0.0
             )
-            
-            # Берём только сгенерированные токены (убираем промпт)
-            Input_len = inputs[“input_ids”].shape[1]
-            Generated_ids = output_ids[0, input_len:]
-            
-            # Декодирование
-            Generated_text = processor.decode(generated_ids, skip_special_tokens=True)
-            
-            # Результат
-            St.success(“✅ Распознавание завершено!”)
-            St.markdown(‘<div class=”result-box”>’, unsafe_allow_html=True)
-            St.subheader(“ Распознанный текст”)
-            St.markdown(f”```\n{generated_text}\n```”)
-            St.markdown(‘</div>’, unsafe_allow_html=True)
-            
-            # Кнопка копирования
-            St.download_button(
-                Label=” Скачать текст как .txt”,
-                Data=generated_text,
-                File_name=”recognized_text.txt”,
-                Mime=”text/plain”
+
+            # Убираем промпт, оставляем только сгенерированный текст
+            generated_ids = output_ids[0, inputs["input_ids"].shape[1]:]
+            generated_text = processor.decode(generated_ids, skip_special_tokens=True)
+
+            # Вывод результата
+            st.success("✅ Распознавание завершено!")
+            st.markdown('<div class="result-box">', unsafe_allow_html=True)
+            st.subheader("📝 Распознанный текст")
+            st.markdown(f"```\n{generated_text}\n```")
+            st.markdown('</div>', unsafe_allow_html=True)
+
+            # Кнопка скачивания
+            st.download_button(
+                label="💾 Скачать текст (.txt)",
+                data=generated_text,
+                file_name="recognized_text.txt",
+                mime="text/plain"
             )
 
-St.markdown(“---“)
-St.markdown(“**Сделано с ️ на базе LightOnOCR-1B-1025** • [Модель на Hugging Face](https://huggingface.co/lightonai/LightOnOCR-1B-1025)”)
+st.markdown("---")
+st.markdown("**Сделано на базе [lightonai/LightOnOCR-1B-1025](https://huggingface.co/lightonai/LightOnOCR-1B-1025)**")