Upload app.py

app.py

@@ -12,7 +12,6 @@ from transformers import (
 )
 from datasets import load_dataset
 
-# Загрузка датасета и моделей
 wikiart_dataset = load_dataset("huggan/wikiart", split="train")
 device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
 
@@ -22,11 +21,9 @@ blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image
 clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device).eval()
 clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
-# Загрузка FAISS индексов
 image_index = faiss.read_index("image_index.faiss")
 text_index = faiss.read_index("text_index.faiss")
 
-# Генерация описания через BLIP
 def generate_caption(image: Image.Image):
     inputs = blip_processor(image, return_tensors="pt").to(device)
     with torch.no_grad():
@@ -34,7 +31,6 @@ def generate_caption(image: Image.Image):
     caption = blip_processor.decode(caption_ids[0], skip_special_tokens=True)
     return caption
 
-# Получение CLIP эмбеддинга по тексту
 def get_clip_text_embedding(text):
     inputs = clip_processor(text=[text], return_tensors="pt", padding=True).to(device)
     with torch.no_grad():
@@ -43,7 +39,6 @@ def get_clip_text_embedding(text):
     faiss.normalize_L2(features)
     return features
 
-# Получение CLIP эмбеддинга по изображению
 def get_clip_image_embedding(image):
     inputs = clip_processor(images=image, return_tensors="pt").to(device)
     with torch.no_grad():
@@ -52,13 +47,12 @@ def get_clip_image_embedding(image):
     faiss.normalize_L2(features)
     return features
 
-# Получение похожих изображений по эмбеддингу
 def get_results_with_images(embedding, index, top_k=2):
     D, I = index.search(embedding, top_k)
     results = []
     for idx in I[0]:
         try:
-            item = wikiart_dataset[idx]
+            item = wikiart_dataset[int(idx)]
             img = item["image"]
             title = item.get("title", "Untitled")
             artist = item.get("artist", "Unknown")
@@ -79,7 +73,6 @@ def search_similar_images(image: Image.Image):
 
     return caption, text_results, image_results
 
-# Интерфейс Gradio
 demo = gr.Interface(
     fn=search_similar_images,
     inputs=gr.Image(label="Загрузите изображение", type="pil"),