Update app.py

app.py

@@ -10,7 +10,6 @@ from PIL import Image
 # ==========================================
 # 1. ЗАГРУЗКА ЗРЕНИЯ (~20MB)
 # ==========================================
-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"
@@ -29,15 +28,12 @@ preprocess = transforms.Compose([
 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,
     dtype=torch.float32, 
     tie_word_embeddings=False
-).to("cpu") # Явно отправляем на CPU
+).to("cpu")
 
 model.config.max_position_embeddings = 128
 
@@ -49,7 +45,6 @@ def predict(image, message, history):
     
     if image is not None:
         try:
-            # Gradio может давать массив numpy, переводим в PIL
             pil_img = Image.fromarray(image.astype('uint8'), 'RGB')
             input_tensor = preprocess(pil_img).unsqueeze(0)
             with torch.no_grad():
@@ -60,22 +55,25 @@ def predict(image, message, history):
         except Exception:
             vision_info = "туман"
 
-    # Промпт под твою структуру
+    # Собираем промпт
     prompt = f"User: ({vision_info}) {message}\nBot:"
     
     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": "assistant", "content": "Слишком много инфы, я запутался!"}]
+    if max_to_gen <= 5:
+        history.append({"role": "user", "content": message})
+        history.append({"role": "assistant", "content": "Слишком длинно, не влезаю в 128!"})
+        return history
 
     with torch.no_grad():
         output_tokens = model.generate(
             **inputs,
-            max_new_tokens=max_new_tokens,
+            max_new_tokens=max_to_gen,
             do_sample=True,
-            temperature=0.25,
+            temperature=0.35,
             repetition_penalty=1.8,
             pad_token_id=tokenizer.pad_token_id
         )
@@ -85,18 +83,19 @@ def predict(image, message, history):
 
     if not answer: answer = "..."
 
+    # В Gradio 6.0 возвращаем обновленный список сообщений
     history.append({"role": "user", "content": message})
     history.append({"role": "assistant", "content": answer})
     return history
 
 # ==========================================
-# 4. ИНТЕРФЕЙС
+# 4. ИНТЕРФЕЙС (GRADIO 6.0)
 # ==========================================
-with gr.Blocks(theme=gr.themes.Default(primary_hue="yellow")) as demo:
+with gr.Blocks() as demo:
     gr.Markdown("# 🍌 BananaVision Lite")
     with gr.Row():
         img_input = gr.Image(label="Глаза")
-        chatbot = gr.Chatbot(type="messages", label="Чат")
+        chatbot = gr.Chatbot(label="Чат") # БЕЗ type="messages"
     
     msg = gr.Textbox(placeholder="Чё там на картинке?")
     btn = gr.Button("Спросить")
@@ -104,4 +103,5 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="yellow")) as demo:
     btn.click(predict, [img_input, msg, chatbot], [chatbot])
     msg.submit(predict, [img_input, msg, chatbot], [chatbot])
 
-demo.launch()
+# Тема передается здесь
+demo.launch(theme=gr.themes.Default(primary_hue="yellow"))