Update app.py

app.py

@@ -2,22 +2,20 @@ import gradio as gr
 import torch
 import re
 import random
+import requests
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from torchvision import models, transforms
 from PIL import Image
-import requests
 
 # ==========================================
-# 1. ЗАГРУЗКА ЗРЕНИЯ (ImageNet Classifier)
+# 1. ЗАГРУЗКА ЗРЕНИЯ (~20MB)
 # ==========================================
-# SqueezeNet — весит копейки, работает быстро
+print("--- Загрузка SqueezeNet ---")
 vision_model = models.squeezenet1_1(weights=models.SqueezeNet1_1_Weights.IMAGENET1K_V1).eval()
 
-# Подгружаем названия категорий
 LABELS_URL = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
 labels = requests.get(LABELS_URL).text.splitlines()
 
-# Подготовка картинки
 preprocess = transforms.Compose([
     transforms.Resize(256),
     transforms.CenterCrop(224),
@@ -26,87 +24,82 @@ preprocess = transforms.Compose([
 ])
 
 # ==========================================
-# 2. ТВОИ МОЗГИ (200M Model)
+# 2. ТВОИ МОЗГИ (Лимит 340MB)
 # ==========================================
 MODEL_PATH = "./" 
 TOKENIZER_NAME = "sberbank-ai/rugpt3small_based_on_gpt2"
 
+print("--- Загрузка твоей модели ---")
 tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
+
+# Убираем device_map, чтобы не требовать accelerate, и фиксим dtype
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_PATH,
-    torch_dtype=torch.float32,
-    device_map="cpu",
+    dtype=torch.float32, 
     tie_word_embeddings=False
-)
+).to("cpu") # Явно отправляем на CPU
+
 model.config.max_position_embeddings = 128
 
 # ==========================================
-# 3. ФУНКЦИЯ ПРЕДСКАЗАНИЯ
+# 3. ЛОГИКА
 # ==========================================
 def predict(image, message, history):
-    vision_info = ""
+    vision_info = "ничего не вижу"
     
-    # Если закинули картинку — распознаем
     if image is not None:
-        pil_img = Image.fromarray(image.astype('uint8'), 'RGB')
-        input_tensor = preprocess(pil_img).unsqueeze(0)
-        with torch.no_grad():
-            output = vision_model(input_tensor)
-        
-        # Берем самый вероятный объект
-        _, index = torch.max(output, 1)
-        detected = labels[index[0]].replace("_", " ")
-        vision_info = f"Ты видишь перед собой: {detected}."
-
-    # Собираем промпт. Впихиваем зрение в начало, чтобы модель "прозрела"
-    # Формат: User: (Вижу: банан) Чё это? \n Bot:
+        try:
+            # Gradio может давать массив numpy, переводим в PIL
+            pil_img = Image.fromarray(image.astype('uint8'), 'RGB')
+            input_tensor = preprocess(pil_img).unsqueeze(0)
+            with torch.no_grad():
+                output = vision_model(input_tensor)
+            _, index = torch.max(output, 1)
+            detected = labels[index[0]].replace("_", " ")
+            vision_info = f"вижу {detected}"
+        except Exception:
+            vision_info = "туман"
+
+    # Промпт под твою структуру
     prompt = f"User: ({vision_info}) {message}\nBot:"
     
-    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
     curr_len = inputs.input_ids.shape[1]
     
-    # Лимит 128
     max_to_gen = 128 - curr_len - 1
     if max_to_gen <= 2:
-        return history + [{"role": "user", "content": message}, {"role": "assistant", "content": "Память забита!"}]
+        return history + [{"role": "assistant", "content": "Слишком много инфы, я запутался!"}]
 
     with torch.no_grad():
         output_tokens = model.generate(
             **inputs,
-            max_new_tokens=max_to_gen,
+            max_new_tokens=max_new_tokens,
             do_sample=True,
-            temperature=0.35, # Твоя ТЕМПЕРАТУРА
+            temperature=0.25,
             repetition_penalty=1.8,
             pad_token_id=tokenizer.pad_token_id
         )
 
-    raw_answer = tokenizer.decode(output_tokens[0][curr_len:], skip_special_tokens=True).strip()
-    answer = re.split(r'User:|Bot:|\n', raw_answer)[0].strip()
+    answer = tokenizer.decode(output_tokens[0][curr_len:], skip_special_tokens=True).strip()
+    answer = re.split(r'User:|Bot:|\n', answer)[0].strip()
 
-    if not answer: answer = "Ясно."
+    if not answer: answer = "..."
 
-    # Формат Gradio 6.0
     history.append({"role": "user", "content": message})
     history.append({"role": "assistant", "content": answer})
     return history
 
 # ==========================================
-# 4. ИНТЕРФЕЙС (DARK-YELLOW STYLE)
+# 4. ИНТЕРФЕЙС
 # ==========================================
-with gr.Blocks(theme=gr.themes.Default(primary_hue="yellow", secondary_hue="neutral").set(
-    body_background_fill="#000000", 
-    block_background_fill="#111111",
-    input_background_fill="#222222"
-)) as demo:
-    gr.Markdown("# 🍌 **BananaVision Lite** (Limit: 340MB)")
-    
+with gr.Blocks(theme=gr.themes.Default(primary_hue="yellow")) as demo:
+    gr.Markdown("# 🍌 BananaVision Lite")
     with gr.Row():
-        with gr.Column(scale=1):
-            img_input = gr.Image(label="Глаза бота (Camera/Upload)")
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(type="messages", label="Чат с Хамом")
-            msg = gr.Textbox(placeholder="Спроси чё-нибудь про картинку...")
-            btn = gr.Button("Отправить", variant="primary")
+        img_input = gr.Image(label="Глаза")
+        chatbot = gr.Chatbot(type="messages", label="Чат")
+    
+    msg = gr.Textbox(placeholder="Чё там на картинке?")
+    btn = gr.Button("Спросить")
 
     btn.click(predict, [img_input, msg, chatbot], [chatbot])
     msg.submit(predict, [img_input, msg, chatbot], [chatbot])