Update app.py

app.py

@@ -4,18 +4,17 @@ import re
 import json
 import tempfile
 import zipfile
-from pathlib import Path
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
-from llama_cpp.llama_chat_format import Qwen2VLChatHandler
+from llama_cpp.llama_chat_format import Llava15ChatHandler
 import base64
 from PIL import Image
 from io import BytesIO
 
 # Константы
 REPO_ID = "mradermacher/VisualQuality-R1-7B-GGUF"
-MODEL_FILE = "VisualQuality-R1-7B.Q4_K_M.gguf"  # 4.68 GB - баланс качества/размера
-MMPROJ_FILE = "VisualQuality-R1-7B.mmproj-Q8_0.gguf"  # 853 MB
+MODEL_FILE = "VisualQuality-R1-7B.Q4_K_M.gguf"
+MMPROJ_FILE = "VisualQuality-R1-7B.mmproj-Q8_0.gguf"
 
 # Промпты
 PROMPT = (
@@ -29,24 +28,21 @@ QUESTION_TEMPLATE_NO_THINKING = "{Question} Please only output the final answer
 
 # Глобальные переменные
 llm = None
-chat_handler = None
 
 
 def download_models():
-    """Скачивание моделей из HuggingFace"""
+    """Скачивание моделей"""
     print("Downloading model files...")
     
     model_path = hf_hub_download(
         repo_id=REPO_ID,
         filename=MODEL_FILE,
-        resume_download=True,
     )
     print(f"Model downloaded: {model_path}")
     
     mmproj_path = hf_hub_download(
         repo_id=REPO_ID,
         filename=MMPROJ_FILE,
-        resume_download=True,
     )
     print(f"MMProj downloaded: {mmproj_path}")
     
@@ -55,45 +51,43 @@ def download_models():
 
 def load_model():
     """Загрузка модели"""
-    global llm, chat_handler
+    global llm
     
     if llm is not None:
         return
     
     model_path, mmproj_path = download_models()
     
-    print("Loading model into memory...")
+    print("Loading model...")
     
-    # Создаём chat handler для Qwen2-VL
-    chat_handler = Qwen2VLChatHandler(
+    # Используем Llava15ChatHandler для vision моделей
+    chat_handler = Llava15ChatHandler(
         clip_model_path=mmproj_path,
         verbose=False
     )
     
-    # Загружаем основную модель
     llm = Llama(
         model_path=model_path,
         chat_handler=chat_handler,
-        n_ctx=4096,  # Контекст
-        n_threads=4,  # Потоки CPU
-        n_gpu_layers=0,  # CPU only
+        n_ctx=4096,
+        n_threads=4,
+        n_gpu_layers=0,
         verbose=False,
     )
     
-    print("Model loaded successfully!")
+    print("Model loaded!")
 
 
-def image_to_base64_uri(image):
+def image_to_data_uri(image):
     """Конвертация PIL Image в data URI"""
     if image is None:
         return None
     
-    # Конвертируем в RGB если нужно
     if image.mode != "RGB":
         image = image.convert("RGB")
     
     # Сжимаем для ускорения
-    max_size = 1024
+    max_size = 768
     if max(image.size) > max_size:
         ratio = max_size / max(image.size)
         new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio))
@@ -107,14 +101,14 @@ def image_to_base64_uri(image):
 
 
 def extract_score(text):
-    """Извлечение оценки из текста"""
+    """Извлечение оценки"""
     try:
-        model_output_matches = re.findall(r'<answer>(.*?)</answer>', text, re.DOTALL)
-        if model_output_matches:
-            model_answer = model_output_matches[-1].strip()
+        matches = re.findall(r'<answer>(.*?)</answer>', text, re.DOTALL)
+        if matches:
+            answer = matches[-1].strip()
         else:
-            model_answer = text.strip()
-        score_match = re.search(r'\d+(\.\d+)?', model_answer)
+            answer = text.strip()
+        score_match = re.search(r'\d+(\.\d+)?', answer)
         if score_match:
             score = float(score_match.group())
             return min(max(score, 1.0), 5.0)
@@ -124,10 +118,10 @@ def extract_score(text):
 
 
 def extract_thinking(text):
-    """Извлечение процесса мышления"""
-    thinking_matches = re.findall(r'<think>(.*?)</think>', text, re.DOTALL)
-    if thinking_matches:
-        return thinking_matches[-1].strip()
+    """Извлечение мышления"""
+    matches = re.findall(r'<think>(.*?)</think>', text, re.DOTALL)
+    if matches:
+        return matches[-1].strip()
     return ""
 
 
@@ -138,17 +132,13 @@ def score_single_image(image, use_thinking=True):
     load_model()
     
     if image is None:
-        yield "❌ Please upload an image first.", "", ""
-        return
+        return "❌ Upload an image first", "", ""
     
-    # Выбор шаблона
     template = QUESTION_TEMPLATE_THINKING if use_thinking else QUESTION_TEMPLATE_NO_THINKING
     prompt_text = template.format(Question=PROMPT)
     
-    # Конвертируем изображение
-    image_uri = image_to_base64_uri(image)
+    image_uri = image_to_data_uri(image)
     
-    # Формируем сообщение
     messages = [
         {
             "role": "user",
@@ -159,11 +149,9 @@ def score_single_image(image, use_thinking=True):
         }
     ]
     
-    # Генерация со стримингом
+    # Стриминг
     generated_text = ""
     
-    yield "⏳ Processing...", "", "*Analyzing image...*"
-    
     try:
         response = llm.create_chat_completion(
             messages=messages,
@@ -183,7 +171,7 @@ def score_single_image(image, use_thinking=True):
                 score = extract_score(generated_text)
                 
                 if score is not None:
-                    score_display = f"⭐ **Quality Score: {score:.2f} / 5.00**"
+                    score_display = f"⭐ **Score: {score:.2f} / 5.00**"
                 else:
                     score_display = "*Analyzing...*"
                 
@@ -196,7 +184,7 @@ def score_single_image(image, use_thinking=True):
         if final_score is not None:
             score_display = f"⭐ **Quality Score: {final_score:.2f} / 5.00**\n\n📊 **For Leaderboard:** `{final_score:.2f}`"
         else:
-            score_display = "❌ Could not extract score. Please try again."
+            score_display = "❌ Could not extract score"
         
         yield generated_text, final_thinking, score_display
         
@@ -205,23 +193,20 @@ def score_single_image(image, use_thinking=True):
 
 
 def process_batch(files, use_thinking=True, progress=gr.Progress()):
-    """Обработка пакета изображений"""
+    """Batch processing"""
     global llm
     
     load_model()
     
     if not files:
-        return "❌ No files uploaded", None
+        return "❌ No files", None
     
     results = []
     template = QUESTION_TEMPLATE_THINKING if use_thinking else QUESTION_TEMPLATE_NO_THINKING
     prompt_text = template.format(Question=PROMPT)
     
-    progress(0, desc="Starting batch processing...")
-    
     for i, file in enumerate(files):
         try:
-            # Загружаем изображение
             if hasattr(file, 'name'):
                 image = Image.open(file.name)
                 filename = os.path.basename(file.name)
@@ -229,7 +214,7 @@ def process_batch(files, use_thinking=True, progress=gr.Progress()):
                 image = Image.open(file)
                 filename = f"image_{i+1}.jpg"
             
-            image_uri = image_to_base64_uri(image)
+            image_uri = image_to_data_uri(image)
             
             messages = [
                 {
@@ -241,7 +226,6 @@ def process_batch(files, use_thinking=True, progress=gr.Progress()):
                 }
             ]
             
-            # Генерация
             response = llm.create_chat_completion(
                 messages=messages,
                 max_tokens=2048 if use_thinking else 256,
@@ -260,7 +244,7 @@ def process_batch(files, use_thinking=True, progress=gr.Progress()):
                 "raw_output": generated_text
             })
             
-            progress((i + 1) / len(files), desc=f"Processed {i+1}/{len(files)}: {filename}")
+            progress((i + 1) / len(files), desc=f"Processed {i+1}/{len(files)}")
             
         except Exception as e:
             results.append({
@@ -270,168 +254,103 @@ def process_batch(files, use_thinking=True, progress=gr.Progress()):
                 "raw_output": str(e)
             })
     
-    # Создаём файлы результатов
+    # Создаём файлы
     with tempfile.TemporaryDirectory() as tmpdir:
-        # Текстовый файл для лидерборда
-        leaderboard_file = os.path.join(tmpdir, "leaderboard_scores.txt")
-        with open(leaderboard_file, "w") as f:
+        # TXT для лидерборда
+        txt_file = os.path.join(tmpdir, "scores.txt")
+        with open(txt_file, "w") as f:
             for r in results:
-                score_str = f"{r['score']:.2f}" if isinstance(r['score'], float) else r['score']
+                score_str = f"{r['score']:.2f}" if isinstance(r['score'], float) else str(r['score'])
                 f.write(f"{r['filename']}\t{score_str}\n")
         
-        # JSON с полными результатами
-        json_file = os.path.join(tmpdir, "full_results.json")
+        # JSON
+        json_file = os.path.join(tmpdir, "results.json")
         with open(json_file, "w") as f:
             json.dump(results, f, indent=2, ensure_ascii=False)
         
-        # CSV файл
+        # CSV
         csv_file = os.path.join(tmpdir, "scores.csv")
         with open(csv_file, "w") as f:
             f.write("filename,score\n")
             for r in results:
-                score_str = f"{r['score']:.2f}" if isinstance(r['score'], float) else r['score']
+                score_str = f"{r['score']:.2f}" if isinstance(r['score'], float) else str(r['score'])
                 f.write(f"{r['filename']},{score_str}\n")
         
-        # Создаём ZIP архив
+        # ZIP
         zip_path = os.path.join(tmpdir, "results.zip")
         with zipfile.ZipFile(zip_path, 'w') as zipf:
-            zipf.write(leaderboard_file, "leaderboard_scores.txt")
-            zipf.write(json_file, "full_results.json")
+            zipf.write(txt_file, "scores.txt")
+            zipf.write(json_file, "results.json")
             zipf.write(csv_file, "scores.csv")
         
-        # Копируем ZIP во временную папку, которая не удалится
+        # Копируем
         final_zip = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
         with open(zip_path, 'rb') as f:
             final_zip.write(f.read())
         final_zip.close()
     
-    # Формируем summary
+    # Summary
     valid_scores = [r['score'] for r in results if isinstance(r['score'], float)]
-    summary = f"""
-## ✅ Batch Processing Complete!
+    avg = sum(valid_scores)/len(valid_scores) if valid_scores else 0
+    
+    summary = f"""## ✅ Done!
 
 **Processed:** {len(results)} images
-**Successful:** {len(valid_scores)} images
-**Failed:** {len(results) - len(valid_scores)} images
+**Success:** {len(valid_scores)}
+**Failed:** {len(results) - len(valid_scores)}
 
-### Statistics:
-- **Average Score:** {sum(valid_scores)/len(valid_scores):.2f} (if valid scores exist)
-- **Min Score:** {min(valid_scores):.2f if valid_scores else 'N/A'}
-- **Max Score:** {max(valid_scores):.2f if valid_scores else 'N/A'}
+**Average:** {avg:.2f}
+**Min:** {min(valid_scores):.2f if valid_scores else 'N/A'}
+**Max:** {max(valid_scores):.2f if valid_scores else 'N/A'}
 
-### Preview (first 10):
-| Filename | Score |
-|----------|-------|
-""" + "\n".join([f"| {r['filename']} | {r['score']:.2f if isinstance(r['score'], float) else r['score']} |" for r in results[:10]])
+### Preview:
+| File | Score |
+|------|-------|
+""" + "\n".join([f"| {r['filename'][:30]} | {r['score']:.2f if isinstance(r['score'], float) else r['score']} |" for r in results[:10]])
     
     return summary, final_zip.name
 
 
-def create_interface():
-    """Создание интерфейса Gradio"""
+# Интерфейс
+with gr.Blocks(title="VisualQuality-R1") as demo:
+    gr.Markdown("""
+    # 🎨 VisualQuality-R1 (GGUF/CPU)
     
-    with gr.Blocks(title="VisualQuality-R1 GGUF") as demo:
-        
-        gr.Markdown("""
-        # 🎨 VisualQuality-R1: Image Quality Assessment (GGUF/CPU)
-        
-        **Reasoning-Induced Image Quality Assessment** | Running on CPU with GGUF quantization
-        
-        [![Paper](https://img.shields.io/badge/arXiv-Paper-red)](https://arxiv.org/abs/2505.14460)
-        [![Model](https://img.shields.io/badge/🤗-Model-yellow)](https://huggingface.co/TianheWu/VisualQuality-R1-7B)
-        
-        > ⚠️ **CPU Mode**: Processing is slower but works without GPU. ~30-60 sec per image.
-        """)
-        
-        with gr.Tabs():
-            # Вкладка для одного изображения
-            with gr.TabItem("📷 Single Image"):
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        image_input = gr.Image(
-                            label="Upload Image",
-                            type="pil",
-                            height=350
-                        )
-                        
-                        thinking_checkbox = gr.Checkbox(
-                            label="🧠 Enable Thinking Mode",
-                            value=True
-                        )
-                        
-                        submit_btn = gr.Button(
-                            "🔍 Analyze Quality",
-                            variant="primary",
-                            size="lg"
-                        )
-                    
-                    with gr.Column(scale=1):
-                        score_output = gr.Markdown(value="*Upload an image to see the score*")
-                        thinking_output = gr.Textbox(label="🧠 Thinking", lines=6, interactive=False)
-                        raw_output = gr.Textbox(label="📝 Full Output", lines=8, interactive=False)
+    Image Quality Assessment | CPU Mode (~30-60 sec/image)
+    
+    [![Paper](https://img.shields.io/badge/arXiv-Paper-red)](https://arxiv.org/abs/2505.14460)
+    """)
+    
+    with gr.Tabs():
+        with gr.TabItem("📷 Single Image"):
+            with gr.Row():
+                with gr.Column():
+                    img_input = gr.Image(label="Upload", type="pil", height=350)
+                    thinking_cb = gr.Checkbox(label="🧠 Thinking Mode", value=True)
+                    btn = gr.Button("🔍 Analyze", variant="primary", size="lg")
                 
-                submit_btn.click(
-                    fn=score_single_image,
-                    inputs=[image_input, thinking_checkbox],
-                    outputs=[raw_output, thinking_output, score_output],
-                )
+                with gr.Column():
+                    score_out = gr.Markdown("*Upload image*")
+                    thinking_out = gr.Textbox(label="Thinking", lines=6)
+                    raw_out = gr.Textbox(label="Output", lines=8)
             
-            # Вкладка для batch processing
-            with gr.TabItem("📁 Batch Processing (1000+ images)"):
-                gr.Markdown("""
-                ### Batch Processing for Leaderboard
-                Upload multiple images (ZIP or individual files) to process them all at once.
-                Results will be saved in a format ready for leaderboard submission.
-                """)
-                
-                with gr.Row():
-                    with gr.Column():
-                        batch_files = gr.File(
-                            label="Upload Images",
-                            file_count="multiple",
-                            file_types=["image"],
-                        )
-                        
-                        batch_thinking = gr.Checkbox(
-                            label="🧠 Enable Thinking Mode (slower but more detailed)",
-                            value=False  # По умолчанию выключено для скорости
-                        )
-                        
-                        batch_btn = gr.Button(
-                            "🚀 Process All Images",
-                            variant="primary",
-                            size="lg"
-                        )
-                    
-                    with gr.Column():
-                        batch_summary = gr.Markdown(value="*Upload images and click Process*")
-                        batch_download = gr.File(label="📥 Download Results")
-                
-                batch_btn.click(
-                    fn=process_batch,
-                    inputs=[batch_files, batch_thinking],
-                    outputs=[batch_summary, batch_download],
-                )
+            btn.click(score_single_image, [img_input, thinking_cb], [raw_out, thinking_out, score_out])
         
-        gr.Markdown("""
-        ---
-        ### 📊 Score Scale
-        | Score | Quality |
-        |-------|---------|
-        | 1.0 | Very poor |
-        | 2.0 | Poor |
-        | 3.0 | Fair |
-        | 4.0 | Good |
-        | 5.0 | Excellent |
-        """)
-    
-    return demo
-
+        with gr.TabItem("📁 Batch (1000+ images)"):
+            gr.Markdown("### Upload multiple images for leaderboard submission")
+            
+            with gr.Row():
+                with gr.Column():
+                    batch_files = gr.File(label="Images", file_count="multiple", file_types=["image"])
+                    batch_thinking = gr.Checkbox(label="🧠 Thinking (slower)", value=False)
+                    batch_btn = gr.Button("🚀 Process All", variant="primary", size="lg")
+                
+                with gr.Column():
+                    batch_summary = gr.Markdown("*Upload and click Process*")
+                    batch_download = gr.File(label="📥 Download Results")
+            
+            batch_btn.click(process_batch, [batch_files, batch_thinking], [batch_summary, batch_download])
 
 if __name__ == "__main__":
-    demo = create_interface()
     demo.queue(max_size=5)
-    demo.launch(
-        show_error=True,
-        ssr_mode=False,
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860)