Update app.py

app.py

@@ -4,13 +4,16 @@ import re
 import json
 import tempfile
 import zipfile
+import traceback
 from huggingface_hub import hf_hub_download
-from llama_cpp import Llama
-from llama_cpp.llama_chat_format import Llava15ChatHandler
 import base64
 from PIL import Image
 from io import BytesIO
 
+print("=" * 50)
+print("Starting VisualQuality-R1 GGUF")
+print("=" * 50)
+
 # Константы
 REPO_ID = "mradermacher/VisualQuality-R1-7B-GGUF"
 MODEL_FILE = "VisualQuality-R1-7B.Q4_K_M.gguf"
@@ -29,10 +32,40 @@ QUESTION_TEMPLATE_NO_THINKING = "{Question} Please only output the final answer
 # Глобальные переменные
 llm = None
 
+print("Importing llama_cpp...")
+try:
+    from llama_cpp import Llama
+    print("llama_cpp imported successfully")
+except Exception as e:
+    print(f"Error importing llama_cpp: {e}")
+    traceback.print_exc()
+
+# Пробуем импортировать chat handlers
+chat_handler_class = None
+try:
+    from llama_cpp.llama_chat_format import Qwen2VLChatHandler
+    chat_handler_class = Qwen2VLChatHandler
+    print("Using Qwen2VLChatHandler")
+except ImportError:
+    print("Qwen2VLChatHandler not found, trying Llava15ChatHandler...")
+    try:
+        from llama_cpp.llama_chat_format import Llava15ChatHandler
+        chat_handler_class = Llava15ChatHandler
+        print("Using Llava15ChatHandler")
+    except ImportError:
+        print("Llava15ChatHandler not found, trying Llava16ChatHandler...")
+        try:
+            from llama_cpp.llama_chat_format import Llava16ChatHandler
+            chat_handler_class = Llava16ChatHandler
+            print("Using Llava16ChatHandler")
+        except ImportError:
+            print("No suitable chat handler found!")
+            chat_handler_class = None
+
 
 def download_models():
     """Скачивание моделей"""
-    print("Downloading model files...")
+    print(f"Downloading {MODEL_FILE}...")
     
     model_path = hf_hub_download(
         repo_id=REPO_ID,
@@ -40,6 +73,7 @@ def download_models():
     )
     print(f"Model downloaded: {model_path}")
     
+    print(f"Downloading {MMPROJ_FILE}...")
     mmproj_path = hf_hub_download(
         repo_id=REPO_ID,
         filename=MMPROJ_FILE,
@@ -51,31 +85,41 @@ def download_models():
 
 def load_model():
     """Загрузка модели"""
-    global llm
+    global llm, chat_handler_class
     
     if llm is not None:
-        return
-    
-    model_path, mmproj_path = download_models()
-    
-    print("Loading model...")
-    
-    # Используем 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,
-        n_gpu_layers=0,
-        verbose=False,
-    )
+        return True
     
-    print("Model loaded!")
+    try:
+        model_path, mmproj_path = download_models()
+        
+        print("Creating chat handler...")
+        if chat_handler_class is not None:
+            chat_handler = chat_handler_class(
+                clip_model_path=mmproj_path,
+                verbose=True
+            )
+            print("Chat handler created")
+        else:
+            print("WARNING: No chat handler, trying without it")
+            chat_handler = None
+        
+        print("Loading LLM...")
+        llm = Llama(
+            model_path=model_path,
+            chat_handler=chat_handler,
+            n_ctx=4096,
+            n_threads=4,
+            n_gpu_layers=0,
+            verbose=True,
+        )
+        print("Model loaded successfully!")
+        return True
+        
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        traceback.print_exc()
+        return False
 
 
 def image_to_data_uri(image):
@@ -86,7 +130,6 @@ def image_to_data_uri(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
     
-    # Сжимаем для ускорения
     max_size = 768
     if max(image.size) > max_size:
         ratio = max_size / max(image.size)
@@ -129,15 +172,20 @@ def score_single_image(image, use_thinking=True):
     """Оценка одного изображения"""
     global llm
     
-    load_model()
+    print(f"score_single_image called, use_thinking={use_thinking}")
     
     if image is None:
         return "❌ Upload an image first", "", ""
     
+    if not load_model():
+        return "❌ Failed to load model. Check logs.", "", ""
+    
     template = QUESTION_TEMPLATE_THINKING if use_thinking else QUESTION_TEMPLATE_NO_THINKING
     prompt_text = template.format(Question=PROMPT)
     
+    print("Converting image to data URI...")
     image_uri = image_to_data_uri(image)
+    print(f"Image URI created, length: {len(image_uri)}")
     
     messages = [
         {
@@ -149,7 +197,7 @@ def score_single_image(image, use_thinking=True):
         }
     ]
     
-    # Стриминг
+    print("Starting generation...")
     generated_text = ""
     
     try:
@@ -170,42 +218,47 @@ def score_single_image(image, use_thinking=True):
                 thinking = extract_thinking(generated_text)
                 score = extract_score(generated_text)
                 
-                if score is not None:
-                    score_display = f"⭐ **Score: {score:.2f} / 5.00**"
-                else:
-                    score_display = "*Analyzing...*"
+                score_display = f"⭐ **Score: {score:.2f} / 5.00**" if score else "*Analyzing...*"
                 
                 yield generated_text, thinking, score_display
         
-        # Финальный результат
+        print(f"Generation complete, output length: {len(generated_text)}")
+        
         final_score = extract_score(generated_text)
         final_thinking = extract_thinking(generated_text) if use_thinking else ""
         
         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"
+            score_display = "❌ Could not extract score. Raw output shown below."
         
         yield generated_text, final_thinking, score_display
         
     except Exception as e:
-        yield f"❌ Error: {str(e)}", "", ""
+        error_msg = f"❌ Error: {str(e)}"
+        print(error_msg)
+        traceback.print_exc()
+        yield error_msg, "", ""
 
 
 def process_batch(files, use_thinking=True, progress=gr.Progress()):
     """Batch processing"""
     global llm
     
-    load_model()
+    print(f"process_batch called with {len(files) if files else 0} files")
     
     if not files:
-        return "❌ No files", None
+        return "❌ No files uploaded", None
+    
+    if not load_model():
+        return "❌ Failed to load model. Check logs.", None
     
     results = []
     template = QUESTION_TEMPLATE_THINKING if use_thinking else QUESTION_TEMPLATE_NO_THINKING
     prompt_text = template.format(Question=PROMPT)
     
     for i, file in enumerate(files):
+        filename = "unknown"
         try:
             if hasattr(file, 'name'):
                 image = Image.open(file.name)
@@ -214,6 +267,8 @@ def process_batch(files, use_thinking=True, progress=gr.Progress()):
                 image = Image.open(file)
                 filename = f"image_{i+1}.jpg"
             
+            print(f"Processing {i+1}/{len(files)}: {filename}")
+            
             image_uri = image_to_data_uri(image)
             
             messages = [
@@ -244,112 +299,154 @@ 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)}")
+            print(f"  Score: {score}")
+            progress((i + 1) / len(files), desc=f"Processed {i+1}/{len(files)}: {filename}")
             
         except Exception as e:
+            print(f"  Error: {e}")
             results.append({
-                "filename": filename if 'filename' in dir() else f"image_{i+1}",
+                "filename": filename,
                 "score": "ERROR",
                 "thinking": "",
                 "raw_output": str(e)
             })
     
-    # Создаём файлы
-    with tempfile.TemporaryDirectory() as tmpdir:
-        # 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 str(r['score'])
-                f.write(f"{r['filename']}\t{score_str}\n")
-        
-        # 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_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 str(r['score'])
-                f.write(f"{r['filename']},{score_str}\n")
-        
-        # ZIP
-        zip_path = os.path.join(tmpdir, "results.zip")
-        with zipfile.ZipFile(zip_path, 'w') as zipf:
-            zipf.write(txt_file, "scores.txt")
-            zipf.write(json_file, "results.json")
-            zipf.write(csv_file, "scores.csv")
-        
-        # Копируем
-        final_zip = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
-        with open(zip_path, 'rb') as f:
-            final_zip.write(f.read())
-        final_zip.close()
+    # Create output files
+    print("Creating output files...")
+    try:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            txt_file = os.path.join(tmpdir, "leaderboard_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 str(r['score'])
+                    f.write(f"{r['filename']}\t{score_str}\n")
+            
+            json_file = os.path.join(tmpdir, "full_results.json")
+            with open(json_file, "w") as f:
+                json.dump(results, f, indent=2, ensure_ascii=False)
+            
+            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 str(r['score'])
+                    f.write(f"{r['filename']},{score_str}\n")
+            
+            zip_path = os.path.join(tmpdir, "results.zip")
+            with zipfile.ZipFile(zip_path, 'w') as zipf:
+                zipf.write(txt_file, "leaderboard_scores.txt")
+                zipf.write(json_file, "full_results.json")
+                zipf.write(csv_file, "scores.csv")
+            
+            final_zip = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
+            with open(zip_path, 'rb') as f:
+                final_zip.write(f.read())
+            final_zip.close()
+            
+            print(f"Results saved to {final_zip.name}")
+    except Exception as e:
+        print(f"Error saving results: {e}")
+        traceback.print_exc()
+        return f"❌ Error saving results: {e}", None
     
     # Summary
     valid_scores = [r['score'] for r in results if isinstance(r['score'], float)]
-    avg = sum(valid_scores)/len(valid_scores) if valid_scores else 0
+    avg = sum(valid_scores) / len(valid_scores) if valid_scores else 0
     
-    summary = f"""## ✅ Done!
+    summary = f"""## ✅ Batch Processing Complete!
 
 **Processed:** {len(results)} images
-**Success:** {len(valid_scores)}
+**Successful:** {len(valid_scores)}
 **Failed:** {len(results) - len(valid_scores)}
 
-**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'}
+### Statistics:
+- **Average Score:** {avg:.2f}
+- **Min Score:** {min(valid_scores):.2f if valid_scores else 'N/A'}
+- **Max Score:** {max(valid_scores):.2f if valid_scores else 'N/A'}
 
-### 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]])
+### Preview (first 10):
+| Filename | Score |
+|----------|-------|
+""" + "\n".join([f"| {r['filename'][:40]} | {r['score']:.2f if isinstance(r['score'], float) else r['score']} |" for r in results[:10]])
     
     return summary, final_zip.name
 
 
-# Интерфейс
-with gr.Blocks(title="VisualQuality-R1") as demo:
+# Gradio Interface
+print("Creating Gradio interface...")
+
+with gr.Blocks(title="VisualQuality-R1 GGUF") as demo:
     gr.Markdown("""
     # 🎨 VisualQuality-R1 (GGUF/CPU)
     
-    Image Quality Assessment | CPU Mode (~30-60 sec/image)
+    **Image Quality Assessment** | CPU Mode (~30-60 sec per image)
     
     [![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)
     """)
     
     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")
+                    img_input = gr.Image(label="📷 Upload Image", type="pil", height=350)
+                    thinking_cb = gr.Checkbox(label="🧠 Enable Thinking Mode", value=True)
+                    analyze_btn = gr.Button("🔍 Analyze Quality", variant="primary", size="lg")
                 
                 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)
+                    score_out = gr.Markdown(value="*Upload an image to see the score*")
+                    thinking_out = gr.Textbox(label="🧠 Thinking Process", lines=6, interactive=False)
+                    raw_out = gr.Textbox(label="📝 Full Output", lines=8, interactive=False)
             
-            btn.click(score_single_image, [img_input, thinking_cb], [raw_out, thinking_out, score_out])
+            analyze_btn.click(
+                score_single_image,
+                inputs=[img_input, thinking_cb],
+                outputs=[raw_out, thinking_out, score_out]
+            )
         
-        with gr.TabItem("📁 Batch (1000+ images)"):
-            gr.Markdown("### Upload multiple images for leaderboard submission")
+        with gr.TabItem("📁 Batch Processing"):
+            gr.Markdown("""
+            ### Batch Processing for Leaderboard
+            Upload multiple images. Results in TXT, CSV, JSON formats.
+            
+            ⚠️ ~30-60 seconds per image on CPU
+            """)
             
             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_files = gr.File(
+                        label="📁 Upload Images",
+                        file_count="multiple",
+                        file_types=["image"]
+                    )
+                    batch_thinking = gr.Checkbox(
+                        label="🧠 Enable 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_summary = gr.Markdown(value="*Upload images and click Process*")
+                    batch_download = gr.File(label="📥 Download Results (ZIP)")
             
-            batch_btn.click(process_batch, [batch_files, batch_thinking], [batch_summary, batch_download])
+            batch_btn.click(
+                process_batch,
+                inputs=[batch_files, batch_thinking],
+                outputs=[batch_summary, batch_download]
+            )
+    
+    gr.Markdown("""
+    ---
+    | Score | Quality |
+    |-------|---------|
+    | 1.0 | Very poor |
+    | 2.0 | Poor |
+    | 3.0 | Fair |
+    | 4.0 | Good |
+    | 5.0 | Excellent |
+    """)
+
+print("Starting server...")
 
 if __name__ == "__main__":
     demo.queue(max_size=5)