Update app.py

app.py

@@ -1,7 +1,9 @@
+# === 1. ПАТЧ ДЛЯ ИСПРАВЛЕНИЯ ОШИБКИ PYTORCHVIDEO (НЕ ДВИГАТЬ!) ===
 import sys
 import torchvision.transforms.functional as F
 sys.modules['torchvision.transforms.functional_tensor'] = F
 
+# === 2. ОСНОВНЫЕ ИМПОРТЫ ===
 import os
 import torch
 import spaces
@@ -9,23 +11,21 @@ import gradio as gr
 from PIL import Image
 from transformers import AutoModelForCausalLM, AutoProcessor
 
+# === 3. НАСТРОЙКИ ===
 MODEL_ID = "zai-org/VisionReward-Video"
 
-# Загружаем процессор
+# Загружаем процессор один раз (он легкий и скачается быстро)
 processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
 
-@spaces.GPU(duration=180)
-def analyze_app_screenshot(image, prompt):
-    if image is None: return "Загрузите скриншот!"
-    
-    # Промпт для глубокого анализа (Checklist)
-    system_prompt = (
-        "Analyze this mobile app screenshot based on: "
-        "1. Visual Quality (Aesthetics), 2. Usability (UI/UX), 3. Text Clarity. "
-        "Provide a score for each."
-    )
-    if not prompt: prompt = system_prompt
+# === 4. ФУНКЦИЯ ИНФЕРЕНСА (ZeroGPU) ===
+@spaces.GPU(duration=180) # 3 минуты на загрузку 25ГБ и расчет
+def predict_ctr(image, prompt):
+    if image is None:
+        return "Ошибка: загрузите изображение"
+    if not prompt:
+        prompt = "high quality mobile app screenshot"
 
+    # Загружаем тяжелую модель прямо в видеокарту
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID, 
         dtype=torch.bfloat16,
@@ -33,50 +33,48 @@ def analyze_app_screenshot(image, prompt):
         device_map="cuda"
     )
 
+    # Подготовка картинки и текста
     inputs = processor(images=image, text=prompt, return_tensors="pt").to("cuda")
     
+    # Расчет Score
     with torch.no_grad():
-        # Генерируем ответ (т.к. это CausalLM, она должна "написать" вердикт)
-        output_ids = model.generate(
-            **inputs, 
-            max_new_tokens=150, 
-            do_sample=False,
-            temperature=0.1
-        )
-        # Убираем входные токены из ответа
-        response = processor.decode(output_ids[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
-        
-        # Также пробуем достать чистый Reward Score (для общего рейтинга)
         try:
-            reward_score = model.get_reward(**inputs).mean().item()
-            final_reward = f"{reward_score:.4f}"
-        except:
-            final_reward = "N/A"
-
-    # Форматируем вывод
-    result = f"### 🏆 Общий Score: {final_reward}\n\n"
-    result += f"### 📝 Детальный разбор:\n{response}"
-    
-    return result
+            # Специфичный метод для VisionReward
+            score = model.get_reward(**inputs)
+        except Exception:
+            # Если get_reward не сработал, пробуем стандартный прогон
+            outputs = model(**inputs)
+            score = outputs.logits if hasattr(outputs, 'logits') else outputs[0]
+        
+    # Пытаемся вернуть красивое число
+    try:
+        return f"Скор (чем выше, тем лучше): {float(score):.4f}"
+    except Exception:
+        return f"Скор (сырой формат): {score}"
 
-# Интерфейс
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 📱 VisionReward: App Store Optimization Analyzer")
-    gr.Markdown("Модель оценивает ваш скриншот по критериям качества, юзабилити и читаемости.")
+# === 5. ИНТЕРФЕЙС GRADIO ===
+with gr.Blocks() as demo:
+    gr.Markdown("# 🚀 VisionReward ASO Predictor")
+    gr.Markdown("Оценка скриншотов для App Store. Модель весит 25 ГБ, первый расчет может занять около минуты.")
     
     with gr.Row():
         with gr.Column():
-            img_input = gr.Image(type="pil", label="Скриншот вашего приложения")
-            custom_prompt = gr.Textbox(
-                label="Доп. инструкции (необязательно)", 
-                placeholder="На чем сфокусироваться?",
-                value=""
+            img_input = gr.Image(type="pil", label="Скриншот приложения")
+            prompt_input = gr.Textbox(
+                label="Что мы оцениваем? (Prompt)", 
+                placeholder="Напр: clean UI, high conversion",
+                value="high quality mobile app screenshot with clean UI"
             )
-            btn = gr.Button("Запустить глубокий анализ", variant="primary")
+            submit_btn = gr.Button("Рассчитать привлекательность", variant="primary")
         
         with gr.Column():
-            output_md = gr.Markdown(label="Анализ нейросети")
+            output_text = gr.Textbox(label="Результат нейросети")
 
-    btn.click(fn=analyze_app_screenshot, inputs=[img_input, custom_prompt], outputs=output_md)
+    submit_btn.click(
+        fn=predict_ctr, 
+        inputs=[img_input, prompt_input], 
+        outputs=output_text
+    )
 
+# Запуск
 demo.launch()