Update app.py

app.py

@@ -1,51 +1,64 @@
-import os, io, requests
+import os
+import io
+import requests
 import gradio as gr
 from PIL import Image
 import torch
 from transformers import Blip2Processor, Blip2ForConditionalGeneration
 
-# 1) Загружаем BLIP-2
-model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16).to("cuda")
+# 1) Определяем устройство: CUDA, если доступно, иначе CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# 2) Загружаем BLIP-2 в float32 и сразу в нужное устройство
+model = Blip2ForConditionalGeneration.from_pretrained(
+    "Salesforce/blip2-opt-2.7b", 
+    torch_dtype=torch.float32,      # float16 убираем
+    low_cpu_mem_usage=True          # помогает снизить пиковую нагрузку CPU
+).to(device)
 processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
 
-# 2) TMDb API
+# 3) TMDb API
 TMDB_KEY = os.environ.get("TMDB_API_KEY", "")
 TMDB_SEARCH_URL = "https://api.themoviedb.org/3/search/movie"
 
 def caption_and_search(image: Image.Image, dummy):
     """
-    1) Генерируем описательную подпись BLIP-2
-    2) По этой подписи ищем в TMDb title + ссылку
+    1) Генерируем подпись BLIP-2 на CPU
+    2) По этой подписи ищем в TMDb top-3 фильма
     """
-    # ——— Генерация подписи ———
-    inputs = processor(images=image, return_tensors="pt").to(model.device)
-    gen = model.generate(**inputs, max_new_tokens=50)
+    # — Конвертируем PIL.Image в тензор
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        gen = model.generate(**inputs, max_new_tokens=50)
     caption = processor.decode(gen[0], skip_special_tokens=True)
 
-    # ——— Поиск по TMDb ———
+    # — Поиск по TMDb
     params = {"api_key": TMDB_KEY, "query": caption}
     resp = requests.get(TMDB_SEARCH_URL, params=params).json()
     results = []
     for m in resp.get("results", [])[:3]:
-        title = m.get("title")
-        url   = f"https://www.themoviedb.org/movie/{m['id']}"
-        results.append({"title": title, "url": url})
+        results.append({
+            "title": m.get("title"),
+            "url": f"https://www.themoviedb.org/movie/{m['id']}"
+        })
+
     return {"caption": caption, "results": results}
 
-# 3) Интерфейс Gradio
 iface = gr.Interface(
     fn=caption_and_search,
     inputs=[
         gr.Image(type="pil", label="Постер/кадр фильма"),
-        gr.Textbox(visible=False)  # второй аргумент не нужен, но Gradio требует
+        gr.Textbox(visible=False)  # второй аргумент обязателен, но не используется
     ],
     outputs=[
         gr.Textbox(label="Generated Caption"),
         gr.JSON(label="Top 3 Matches (title + link)")
     ],
-    title="Movie Poster Caption & Search",
-    description="BLIP-2 → TMDb search: получаем описание и ссылки на фильмы"
+    title="Movie Poster Caption & Search (CPU‑Only)",
+    description="BLIP‑2 на CPU + TMDb → описание и ссылки на фильмы"
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    # Для публичной ссылки
+    iface.launch(share=True)