Update app.py

app.py

@@ -1,21 +1,17 @@
 #!/usr/bin/env python3
 """
-multimodal gpt-oss 120b — Gradio app for Hugging Face Spaces
-
-Функции:
-- Загрузка собственной картинки (type="filepath")
-- Галерея примеров (клик -> подставляет в загрузчик)
-- Автогенерация "More Detailed Caption" через gradio_client Florence-2
-- Streaming ответов от openai/gpt-oss-120b (через NVIDIA integrate / OpenAI-compatible)
-- Кеширование подписи для одной и той же картинки
+multimodal gpt-oss 120b — Gradio app с Florence-2 в браузере (WebGPU)
+
+Что изменилось:
+- Подпись к изображению генерим на стороне пользователя (WebGPU) через Transformers.js.
+- Сервер больше не грузит Florence/torch.
+- LLM остаётся через NVIDIA Integrate (OpenAI-compatible), как и было.
 """
 
 import os
 import traceback
-from typing import Any, Dict, List, Tuple, Optional
-
+from typing import Any, Optional, List
 import gradio as gr
-from gradio_client import Client, handle_file
 from openai import OpenAI
 
 # (опционально) локальный .env при локальном запуске
@@ -34,156 +30,56 @@ if not NV_API_KEY:
         "NV_API_KEY не задан. В Hugging Face Space зайди в Settings → Secrets и добавь NV_API_KEY."
     )
 
-# Florence-2 Gradio wrapper (публичный)
-FLORENCE_WRAPPER = "gokaygokay/Florence-2"
-
-# --------------------- Клиенты ---------------------
-florence = Client(FLORENCE_WRAPPER)
+# OpenAI клиент для LLM
 llm = OpenAI(base_url=NV_BASE_URL, api_key=NV_API_KEY)
 
-# --------------------- Хелперы ---------------------
-def _normalize_florence_result(res: Any) -> str:
-    """
-    Нормализует результат predict от Florence-2: возвращает текстовую подпись.
-    Подстраховываемся на разные форматы (строка, dict, list и т.д.).
-    """
-    try:
-        if res is None:
-            return ""
-        if isinstance(res, str):
-            return res
-        # dict-like
-        if isinstance(res, dict):
-            # часто бывает ключ 'caption' или 'text' или 'generated_text'
-            for k in ("caption", "text", "generated_text", "output", "result"):
-                if k in res and isinstance(res[k], str):
-                    return res[k]
-            # если есть nested fields, попробуем взять первое строковое значение
-            for v in res.values():
-                if isinstance(v, str):
-                    return v
-            # fallback: str()
-            return str(res)
-        # list/tuple: join string elements
-        if isinstance(res, (list, tuple)):
-            pieces = [str(x) for x in res]
-            return "\n".join(pieces)
-        # other types: fallback to str
-        return str(res)
-    except Exception:
-        return f"[Ошибка нормализации подписи: {traceback.format_exc()}]"
-
-def get_caption_for_image(image_path_or_url: str, safety_note: bool = False) -> str:
-    """
-    Запрос к Florence-2: task_prompt="More Detailed Caption".
-    Принимает локальный путь или URL.
-    """
-    try:
-        if not image_path_or_url:
-            return ""
-        # handle_file поддерживает URL и локальные пути
-        res = florence.predict(
-            image=handle_file(image_path_or_url),
-            task_prompt="More Detailed Caption",
-            text_input=None,
-            model_id="microsoft/Florence-2-large",
-            api_name="/process_image"
-        )
-        caption = _normalize_florence_result(res)
-        return caption
-    except Exception as e:
-        # логируем в stdout (HF Spaces покажет лог)
-        print("Ошибка Florence-2 predict:", e)
-        traceback.print_exc()
-        return f"[Ошибка при генерации подписи: {e}]"
 
 def _extract_text_from_stream_chunk(chunk: Any) -> str:
     """
     Универсально извлекает текстовые фрагменты из чанка стриминга LLM.
-    Работает с разными формами chunk (объект SDK, dict, ...)
     """
     try:
-        # объектный стиль: chunk.choices[0].delta.content
         if hasattr(chunk, "choices"):
             choices = getattr(chunk, "choices")
             if choices:
                 c0 = choices[0]
                 delta = getattr(c0, "delta", None)
                 if delta is not None:
-                    # reasoning_content или content
                     txt = getattr(delta, "reasoning_content", None) or getattr(delta, "content", None)
                     if txt:
                         return str(txt)
-                # some SDK might put content in c0.get("text") etc.
                 text_attr = getattr(c0, "text", None)
                 if text_attr:
                     return str(text_attr)
-        # dict-style
         if isinstance(chunk, dict):
             choices = chunk.get("choices") or []
             if choices:
                 delta = choices[0].get("delta") or {}
-                # try common keys
                 return str(delta.get("content") or delta.get("reasoning_content") or choices[0].get("text") or "")
     except Exception:
         pass
     return ""
 
-# --------------------- UI-логика ---------------------
-# Кеш подписи (чтобы не вызывать Florence снова для той же картинки)
-# Храним словарь: {"image_path": "...", "caption": "..."}
-# Будем использовать gr.State для хранения этого словаря в сессии
-def generate_and_cache_caption(image, cache: Optional[Dict[str, str]]):
-    """
-    Вызывается при изменении image_input или при клике по галерее.
-    Возвращает (caption_text, new_cache_dict).
-    """
-    try:
-        if not image:
-            return "", {"image_path": None, "caption": None}
-        # Готовим path/URL
-        img_path = image if isinstance(image, str) else getattr(image, "name", None) or image
-        # Проверка кеша
-        if cache and cache.get("image_path") == img_path and cache.get("caption"):
-            return cache.get("caption"), cache
-        # Иначе генерируем подпись
-        caption = get_caption_for_image(img_path)
-        new_cache = {"image_path": img_path, "caption": caption}
-        return caption, new_cache
-    except Exception as e:
-        print("generate_and_cache_caption exception:", e)
-        traceback.print_exc()
-        return f"[Ошибка генерации подписи: {e}]", {"image_path": None, "caption": None}
 
-def chat_stream(image, user_message, history, cache: Dict[str, str]):
+def chat_stream(image, user_message: str, history: Optional[List[List[str]]], caption_text: str):
     """
-    Основной generator для кнопки Отправить / submit:
-    - Автоматически использует кеш подписи (если есть), иначе генерирует новую
-    - Возвращает по мере стриминга (history, caption) — соответствие outputs=[chatbot, raw_caption]
+    Основной generator для стриминга ответов LLM.
+    Теперь принимает caption_text прямо из браузера (WebGPU).
     """
     history = history or []
-    # Проверки входа
+
     if not user_message:
-        # ничего не делаем, просто возвращаем текущее состояние
-        yield history, (cache.get("caption") if cache else "")
+        yield history, (caption_text or "")
         return
 
     if not image:
-        # если нет картинки — говорим пользователю
         history.append([user_message, "Пожалуйста, загрузите изображение или выберите из галереи."])
-        yield history, (cache.get("caption") if cache else "")
+        yield history, (caption_text or "")
         return
 
-    # получить путь и подпись (используем кеш, если совпадает)
-    img_path = image if isinstance(image, str) else getattr(image, "name", None) or image
-    if cache and cache.get("image_path") == img_path and cache.get("caption"):
-        caption = cache.get("caption")
-    else:
-        caption = get_caption_for_image(img_path)
-        # обновляем кеш локально (не gr.State, а для текущего запроса)
-        cache = {"image_path": img_path, "caption": caption}
+    caption = caption_text or ""
 
-    # система-промпт — даём контекст и просим указывать степень уверенности
+    # Системный промпт с подписью
     system_prompt = (
         "You are 'multimodal gpt-oss 120b', a helpful multimodal assistant. "
         "Use the provided 'More Detailed Caption' as authoritative visual context. "
@@ -195,14 +91,14 @@ def chat_stream(image, user_message, history, cache: Dict[str, str]):
         "Be concise unless asked for details."
     )
 
-    # добавляем пользовательский запрос в историю (пустой ответ пока)
-    history.append([user_message, ""])  # assistant текст будет заполняться по мере стрима
-    # первый yield чтобы UI сразу отобразил user's message и подпись
+    # Добавляем сообщение пользователя
+    history.append([user_message, ""])
+    # Показать подпись справа от чата (как и ��аньше)
     yield history, caption
 
     assistant_accum = ""
     try:
-        # Запускаем стриминг вызов
+        # Стриминг от LLM
         stream = llm.chat.completions.create(
             model="openai/gpt-oss-120b",
             messages=[
@@ -220,15 +116,13 @@ def chat_stream(image, user_message, history, cache: Dict[str, str]):
             if not piece:
                 continue
             assistant_accum += piece
-            # обновляем последний элемент истории (assistant part)
             history[-1][1] = assistant_accum
             yield history, caption
 
     except Exception as e:
-        # Ошибка стриминга: попробуем получить финальный ответ без стрима, либо показать ошибку
-        print("Streaming error:", e)
+        print(f"Streaming error: {e}")
         traceback.print_exc()
-        # Пытаемся сделать не-стриминг вызов (fallback)
+        # Fallback на не-стриминг запрос
         try:
             resp = llm.chat.completions.create(
                 model="openai/gpt-oss-120b",
@@ -241,9 +135,7 @@ def chat_stream(image, user_message, history, cache: Dict[str, str]):
                 max_tokens=1024,
                 stream=False,
             )
-            # нормализуем возможный формат ответа
             final_text = ""
-            # SDK может вернуть object-like resp.choices[0].message.content
             if hasattr(resp, "choices"):
                 try:
                     final_text = getattr(resp.choices[0].message, "content", "") or getattr(resp.choices[0], "text", "") or ""
@@ -264,10 +156,10 @@ def chat_stream(image, user_message, history, cache: Dict[str, str]):
             history[-1][1] = f"[Ошибка LLM: {e2}]"
             yield history, caption
 
-    # финальный yield (гарантируем окончательное состояние)
     yield history, caption
 
-# --------------------- Примеры для галереи (список строк) ---------------------
+
+# --------------------- Примеры для галереи ---------------------
 EXAMPLE_IMAGES = [
     "https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png",
     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.png",
@@ -281,21 +173,108 @@ css = """
 #title { text-align: center; }
 """
 
+# JS-функция: делает caption в браузере через WebGPU (Transformers.js)
+WEBGPU_CAPTION_JS = r"""
+async (image, use_client) => {
+  try {
+    if (!use_client) return null;
+
+    if (!('gpu' in navigator)) {
+      return "[WebGPU недоступен в браузере. Chrome/Edge 113+ (на Linux — chrome://flags/#enable-unsafe-webgpu), Safari TP.]";
+    }
+
+    // Извлекаем источник изображения из значения Gradio Image
+    const toHTMLImage = async (imgVal) => {
+      if (!imgVal) throw new Error("Нет изображения");
+      let src = null;
+      if (typeof imgVal === 'string') {
+        src = imgVal;
+      } else if (imgVal?.image) {
+        src = imgVal.image;
+      } else if (imgVal?.data) {
+        src = imgVal.data;
+      }
+      if (!src) throw new Error("Не удалось прочитать изображение");
+      const im = new Image();
+      im.crossOrigin = 'anonymous';
+      im.src = src;
+      await im.decode();
+      return im;
+    };
+
+    // Подтягиваем Transformers.js
+    const { pipeline, env } = await import("https://cdn.jsdelivr.net/npm/@xenova/[email protected]");
+
+    // Предпочесть WebGPU
+    env.allowRemoteModels = true;
+    env.useBrowserCache = true; // кэш в IndexedDB
+    env.backends.onnx.backend = 'webgpu';
+
+    // Инициализация один раз
+    if (!window.__webgpu_captioner) {
+      const candidates = [
+        'Xenova/Florence-2-large-ft',
+        'Xenova/Florence-2-base-ft'
+      ];
+      let lastErr = null;
+      for (const model of candidates) {
+        try {
+          window.__webgpu_captioner = await pipeline(
+            'image-to-text',
+            model,
+            { device: 'webgpu', dtype: 'fp16', quantized: true }
+          );
+          break;
+        } catch (e) {
+          lastErr = e;
+          console.warn('Failed to load', model, e);
+        }
+      }
+      if (!window.__webgpu_captioner) throw lastErr || new Error("Не удалось инициализировать captioner");
+    }
+
+    const imgEl = await toHTMLImage(image);
+
+    // Для Florence-2 более детальная подпись через специальный токен задачи
+    const out = await window.__webgpu_captioner(imgEl, { text: '<MORE_DETAILED_CAPTION>' });
+
+    const text = Array.isArray(out)
+      ? (out[0]?.generated_text ?? out[0]?.text ?? JSON.stringify(out[0]))
+      : (out?.generated_text ?? out?.text ?? String(out));
+
+    return text;
+  } catch (e) {
+    return `[WebGPU caption error: ${'message' in e ? e.message : e}]`;
+  }
+}
+"""
+
 with gr.Blocks(css=css, analytics_enabled=False) as demo:
-    gr.Markdown("<h2 id='title'>🖼️ multimodal gpt-oss 120b — визуальный чат</h2>")
+    gr.Markdown("<h2 id='title'>🖼️ multimodal gpt-oss 120b — визуальный чат (Florence в браузере / WebGPU)</h2>")
+
     with gr.Row():
         with gr.Column(scale=4):
-            image_input = gr.Image(label="Загрузите картинку (файл / drag-n-drop / камера)", type="filepath")
-            raw_caption = gr.Textbox(label="More Detailed Caption (Florence-2)", interactive=False, lines=6)
-            user_input = gr.Textbox(label="Вопрос по изображению", placeholder="Например: Что происходит на фото?")
-            send_btn = gr.Button("Отправить")
-            clear_btn = gr.Button("Очистить чат")
-            gr.Markdown("**Галерея примеров (клик — подставить в загрузчик и получить подпись)**")
-            # Исправлено: убран .style(), добавлены параметры в конструктор
+            image_input = gr.Image(label="Загрузите картинку", type="filepath")
+            use_webgpu = gr.Checkbox(value=True, label="Генерировать подпись к изображению в браузере (WebGPU)")
+            raw_caption = gr.Textbox(
+                label="More Detailed Caption (WebGPU)", 
+                interactive=True, 
+                lines=6,
+                placeholder="Подпись появится тут (если включён WebGPU-капшенер)"
+            )
+            user_input = gr.Textbox(
+                label="Вопрос по изображению", 
+                placeholder="Например: Что происходит на фото?"
+            )
+            with gr.Row():
+                send_btn = gr.Button("Отправить", variant="primary")
+                clear_btn = gr.Button("Очистить чат")
+
+            gr.Markdown("**Галерея примеров (клик — подставить в загрузчик, подпись посчитается в браузере)**")
             gallery = gr.Gallery(
-                value=EXAMPLE_IMAGES, 
-                label="Примеры", 
-                columns=4, 
+                value=EXAMPLE_IMAGES,
+                label="Примеры",
+                columns=4,
                 rows=1,
                 show_label=False,
                 height="auto",
@@ -305,40 +284,52 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
         with gr.Column(scale=6):
             chatbot = gr.Chatbot(label="Чат с моделью", height=640)
 
-    # gr.State для кеша подписи
-    caption_cache = gr.State(value={"image_path": None, "caption": None})
+    # Клик по галерее: просто подставить изображение и очистить подпись (капшенер сработает на change)
+    def on_gallery_select(evt: gr.SelectData):
+        img = EXAMPLE_IMAGES[evt.index]
+        return img, ""
 
-    # обработчик клика по галерее: сразу подставляет картинку, генерирует подпись и обновляет кеш
-    def on_gallery_select(elem, cache):
-        # elem может быть строкой (URL) или список/tuple в некоторых версиях
-        img = None
-        if isinstance(elem, (list, tuple)):
-            img = elem[0] if elem else None
-        else:
-            img = elem
-        caption, new_cache = generate_and_cache_caption(img, cache)
-        # возвращаем (image_input value, raw_caption, new_cache)
-        return img, caption, new_cache
-
-    gallery.select(on_gallery_select, inputs=[gallery, caption_cache], outputs=[image_input, raw_caption, caption_cache])
+    gallery.select(
+        on_gallery_select,
+        inputs=None,
+        outputs=[image_input, raw_caption]
+    )
 
-    # обработчик изменения image_input (загрузка своей картинки)
-    def on_image_change(image, cache):
-        caption, new_cache = generate_and_cache_caption(image, cache)
-        return caption, new_cache
+    # Изменение картинки: считаем подпись на клиенте (WebGPU)
+    image_input.change(
+        None,
+        inputs=[image_input, use_webgpu],
+        outputs=[raw_caption],
+        js=WEBGPU_CAPTION_JS
+    )
 
-    image_input.change(on_image_change, inputs=[image_input, caption_cache], outputs=[raw_caption, caption_cache])
+    # Отправка сообщения: берём caption прямо из текстбокса (не генерим на сервере)
+    send_btn.click(
+        chat_stream,
+        inputs=[image_input, user_input, chatbot, raw_caption],
+        outputs=[chatbot, raw_caption]
+    )
 
-    # отправка (кнопка или Enter) — стриминг LLM, outputs: (chatbot, raw_caption)
-    send_btn.click(chat_stream, inputs=[image_input, user_input, chatbot, caption_cache], outputs=[chatbot, raw_caption])
-    user_input.submit(chat_stream, inputs=[image_input, user_input, chatbot, caption_cache], outputs=[chatbot, raw_caption])
+    user_input.submit(
+        chat_stream,
+        inputs=[image_input, user_input, chatbot, raw_caption],
+        outputs=[chatbot, raw_caption]
+    )
 
-    # очистка: сбрасываем чат и кеш
+    # Очистка чата + подписи
     def clear_all():
-        return [], {"image_path": None, "caption": None}, ""
+        return [], ""
 
-    clear_btn.click(clear_all, inputs=None, outputs=[chatbot, caption_cache, raw_caption])
+    clear_btn.click(
+        clear_all,
+        inputs=None,
+        outputs=[chatbot, raw_caption]
+    )
 
-# Запуск (в HF Spaces не указывать share=True)
+# Запуск
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=int(os.environ.get("PORT", 7860)),
+        share=False
+    )