Update app.py

app.py

@@ -1,47 +1,40 @@
-import os, io, requests
+import os
+import requests
 import gradio as gr
 from PIL import Image
 import torch
-from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
+from transformers import BlipProcessor, BlipForConditionalGeneration
 
-# 1) DEVICE
+# 1) Устройство: CPU (или GPU, если вдруг)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", device)
 
-# 2) Captioning модель
-model = VisionEncoderDecoderModel.from_pretrained(
-    "nlpconnect/vit-gpt2-image-captioning"
-).to(device)
-processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+# 2) Лёгкая BLIP‑модель (~240 MiB)
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+model     = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")\
+              .to(device)
 
-# 3) TMDb
+# 3) TMDb API
 TMDB_KEY = os.environ["TMDB_API_KEY"]
 TMDB_SEARCH_URL = "https://api.themoviedb.org/3/search/movie"
 
-def generate_caption(image: Image.Image):
-    """Генерим краткий caption"""
-    pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device)
-    output_ids = model.generate(pixel_values, max_new_tokens=50)
-    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
-    return caption
+def caption_and_search(image: Image.Image, _):
+    # 4) Генерируем подпись (≈3–5 сек на CPU)
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        out = model.generate(**inputs, max_new_tokens=30)
+    caption = processor.decode(out[0], skip_special_tokens=True).strip()
 
-def caption_to_movies(image: Image.Image, dummy):
-    """
-    1) Caption
-    2) Поиск TMDb по этому caption
-    3) Топ‑3 фильмов с title+url
-    """
-    caption = generate_caption(image)
-
-    # Точный поиск
+    # 5) Делаем поиск в TMDb
     params = {"api_key": TMDB_KEY, "query": caption}
-    resp = requests.get(TMDB_SEARCH_URL, params=params)
+    resp = requests.get(TMDB_SEARCH_URL, params=params, timeout=10)
     if resp.status_code != 200:
-        return {"caption": caption, "results": [{"error": f"TMDb {resp.status_code}"}]}
+        return {"caption": caption,
+                "results": [{"error": f"TMDb API returned {resp.status_code}"}]}
 
-    results_raw = resp.json().get("results", [])[:3]
+    movies = resp.json().get("results", [])[:3]
     results = []
-    for m in results_raw:
+    for m in movies:
         results.append({
             "title": m.get("title", "Unknown"),
             "url":   f"https://www.themoviedb.org/movie/{m['id']}"
@@ -49,19 +42,19 @@ def caption_to_movies(image: Image.Image, dummy):
 
     return {"caption": caption, "results": results}
 
-# 4) Интерфейс
+# 6) Интерфейс Gradio
 iface = gr.Interface(
-    fn=caption_to_movies,
+    fn=caption_and_search,
     inputs=[
         gr.Image(type="pil", label="Постер или кадр фильма"),
-        gr.Textbox(visible=False)  # второй аргумент по сигнатуре
+        gr.Textbox(visible=False)  # второй аргумент для сигнатуры
     ],
     outputs=[
         gr.Textbox(label="Auto‑caption"),
         gr.JSON(label="Top‑3 Movies (title + TMDb URL)")
     ],
-    title="Movie Finder via Caption + TMDb",
-    description="Генерирует подпись кадра и ищет фильмы на TMDb по этому описанию"
+    title="Fast Movie Finder (BLIP‑Base + TMDb)",
+    description="≈240 MiB на CPU даёт caption за 3–5 сек и сразу ищет топ‑3 фильма в TMDb"
 )
 
 if __name__ == "__main__":