Update app.py

app.py

@@ -1,81 +1,167 @@
-import io
-import os
-import streamlit as st
-import torch
-from PIL import Image
-from transformers import AutoProcessor, AutoModelForSeq2SeqLM
+Import io
+Import streamlit as st
+Import torch
+From transformers import LightOnOcrForConditionalGeneration, LightOnOcrProcessor
+From PIL import Image
 
-import os
-import tempfile
+# Стильное оформление приложения
+St.set_page_config(
+    Page_title=”LightOnOCR • Распознай текст”,
+    Page_icon=””,
+    Layout=”centered”,
+    Initial_sidebar_state=”expanded”
+)
 
-# Определяем папку для кэша: /data если есть (persistent), иначе /tmp
-if os.path.exists("/data") and os.access("/data", os.W_OK):
-    CACHE_DIR = "/data/.huggingface"
-else:
-    CACHE_DIR = os.path.join(tempfile.gettempdir(), ".huggingface")
+# Кастомный 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;
+        }
+        .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;
+        }
+        .stButton > button:hover {
+            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;
+        }
+    </style>
+“””, unsafe_allow_html=True)
 
-os.makedirs(CACHE_DIR, exist_ok=True)
-os.environ["HF_HOME"] = CACHE_DIR
-os.environ["HF_HUB_CACHE"] = os.path.join(CACHE_DIR, "hub")
-os.environ["TRANSFORMERS_CACHE"] = os.path.join(CACHE_DIR, "transformers")
-
-# --- Функция загрузки модели с кэшированием через Streamlit ---
-@st.cache_resource  # Эта декорация сохраняет модель в памяти между запусками
-def load_model():
-    model_name = "lightonai/LightOnOCR-1B-1025"
+@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
     
-    # Определяем, есть ли GPU (CUDA)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    st.write(f"Using device: {device}")  # Для отладки в логах
+    Model = LightOnOcrForConditionalGeneration.from_pretrained(
+        Model_name,
+        Torch_dtype=dtype,
+        Trust_remote_code=True
+    ).to(device)
     
-    # Загружаем процессор (преобразует изображение в тензоры) и модель
-    processor = AutoProcessor.from_pretrained(model_name, cache_dir=CACHE_DIR)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name, cache_dir=CACHE_DIR)
-    model = model.to(device)  # Перемещаем модель на GPU/CPU
+    Processor = LightOnOcrProcessor.from_pretrained(model_name)
     
-    return processor, model, device
+    Return processor, model, device, dtype
 
-# --- Интерфейс загрузки изображения ---
-def load_image():
-    uploaded_file = st.file_uploader(
-        'Выберите изображение с английским текстом', 
-        type=['png', 'jpg', 'jpeg']
+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_column_width=True)
-        # Конвертируем в RGB (на всякий случай)
-        return Image.open(io.BytesIO(image_data)).convert('RGB')
-    return None
+    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 • Поддерживает документы, чеки, фото, таблицы и сложную вёрстку”)
 
-# --- Заголовок приложения ---
-st.title('🇬🇧 Распознавание английского текста (LightOnOCR)')
+# Загрузка модели
+Processor, model, device, dtype = load_model()
 
-# --- Загружаем модель (один раз) ---
-with st.spinner('Загрузка модели... Это может занять 1-2 минуты при первом запуске'):
-    processor, model, device = 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()
+# Основной интерфейс
+Img = load_image()
 
-# --- Кнопка распознавания ---
-if st.button('Распознать текст') and img is not None:
-    with st.spinner('Распознавание...'):
-        # Преобразуем изображение в формат, понятный модели
-        inputs = processor(images=img, return_tensors="pt").to(device)
-        
-        # Генерируем текст (без вычисления градиентов, чтобы экономить память)
-        with torch.no_grad():
-            generated_ids = model.generate(
+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()
+            }
+            
+            # Генерация
+            Output_ids = model.generate(
                 **inputs,
-                max_new_tokens=512,   # Максимум символов на выходе
-                do_sample=False,      # Детерминированный режим (лучше для OCR)
-                num_beams=1
+                Max_new_tokens=2048,
+                Do_sample=False,
+                Temperature=0.0
             )
-        # Декодируем ID токенов обратно в строку
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        
-        st.success('✅ Распознано!')
-        st.markdown('**📝 Текст на изображении:**')
-        st.markdown(f'`{generated_text}`')
+            
+            # Берём только сгенерированные токены (убираем промпт)
+            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”
+            )
+
+St.markdown(“---“)
+St.markdown(“**Сделано с ️ на базе LightOnOCR-1B-1025** • [Модель на Hugging Face](https://huggingface.co/lightonai/LightOnOCR-1B-1025)”)