app.py

import os
import sys
import subprocess

# --- УСТАНОВКА LLAMA-CPP-PYTHON (Runtime) ---
try:
    from llama_cpp import Llama
    from llama_cpp.llama_chat_format import Llava15ChatHandler
    print("Библиотека llama-cpp-python проверена.")
except ImportError:
    print("Установка llama-cpp-python с поддержкой CPU...")
    # Устанавливаем версию с официального индекса разработчика
    subprocess.check_call([
        sys.executable, "-m", "pip", "install", 
        "llama-cpp-python>=0.3.2",
        "--extra-index-url", "https://abetlen.github.io/llama-cpp-python/whl/cpu"
    ])
    print("Установка завершена! Импортируем...")
    from llama_cpp import Llama

import gradio as gr
from huggingface_hub import hf_hub_download
import base64
import io
import re
from PIL import Image

# Настройки модели
REPO_ID = "mradermacher/VisualQuality-R1-7B-GGUF"
MODEL_FILENAME = "VisualQuality-R1-7B.Q8_0.gguf" 

llm = None

def load_model():
    global llm
    if llm is None:
        print(f"Загрузка модели {MODEL_FILENAME}...")
        try:
            model_path = hf_hub_download(
                repo_id=REPO_ID,
                filename=MODEL_FILENAME
            )
            
            # ВАЖНО: chat_format="qwen2vl" включает встроенный хендлер для картинок
            llm = Llama(
                model_path=model_path,
                n_ctx=12288,          # 12k контекста (хватит для HD картинки)
                n_gpu_layers=0,       # CPU
                verbose=True,
                chat_format="qwen2vl" # Явно включаем режим Qwen2-VL
            )
            print("Модель успешно загружена!")
        except Exception as e:
            print(f"Критическая ошибка загрузки: {e}")
            raise e
    return llm

def process_image(image):
    # Ресайз слишком больших изображений для экономии памяти и контекста
    max_size = 1024
    if max(image.size) > max_size:
        ratio = max_size / max(image.size)
        new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio))
        image = image.resize(new_size, Image.Resampling.LANCZOS)
    
    # Конвертация в Base64
    buffered = io.BytesIO()
    image = image.convert("RGB")
    image.save(buffered, format="JPEG", quality=90)
    return base64.b64encode(buffered.getvalue()).decode('utf-8')

def evaluate_image(image, progress=gr.Progress()):
    if image is None:
        return "Пожалуйста, загрузите изображение.", ""
    
    progress(0, desc="Инициализация...")
    try:
        model = load_model()
    except Exception as e:
        return f"Ошибка инициализации модели: {str(e)}", "Ошибка"
    
    system_prompt = "You are doing the image quality assessment task."
    user_prompt_text = (
        "What is your overall rating on the quality of this picture? "
        "The rating should be a float between 1 and 5, rounded to two decimal places, "
        "with 1 representing very poor quality and 5 representing excellent quality. "
        "Please only output the final answer with only one score in <answer> </answer> tags."
    )
    
    progress(0.1, desc="Обработка изображения...")
    base64_image = process_image(image)
    image_url = f"data:image/jpeg;base64,{base64_image}"

    messages = [
        {"role": "system", "content": system_prompt},
        {
            "role": "user", 
            "content": [
                {"type": "image_url", "image_url": {"url": image_url}},
                {"type": "text", "text": user_prompt_text}
            ]
        }
    ]

    full_response = ""
    print("Отправка запроса в модель...")
    
    try:
        stream = model.create_chat_completion(
            messages=messages,
            max_tokens=1500,
            temperature=0.6,
            stream=True
        )

        for chunk in stream:
            if "choices" in chunk:
                delta = chunk["choices"][0]["delta"]
                if "content" in delta and delta["content"]:
                    content = delta["content"]
                    full_response += content
                    yield full_response, "Вычисляется..."
                    
    except ValueError as e:
        # Если формат чата не сработал
        err = f"Ошибка формата: {e}. Попробуйте перезагрузить Space."
        print(err)
        yield err, "Error"
        return
    except Exception as e:
        err = f"Внутренняя ошибка: {e}"
        print(err)
        yield err, "Error"
        return

    # Извлечение оценки
    score_match = re.search(r'<answer>\s*([\d\.]+)\s*</answer>', full_response)
    final_score = score_match.group(1) if score_match else "Не найдено"
    
    yield full_response, final_score

with gr.Blocks(title="VisualQuality-R1 (Q8 GGUF)") as demo:
    gr.Markdown("# 👁️ VisualQuality-R1 (7B Q8)")
    gr.Markdown(
        "Оценка качества (IQA) с CoT. Работает на CPU (медленно!).\n"
        "Если видите ошибку 'context window', попробуйте картинку меньшего разрешения."
    )
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Изображение")
            run_btn = gr.Button("Оценить", variant="primary")
        
        with gr.Column():
            output_score = gr.Label(label="Оценка")
            output_text = gr.Textbox(label="Рассуждения (CoT)", lines=15)

    run_btn.click(
        fn=evaluate_image,
        inputs=[input_img],
        outputs=[output_text, output_score]
    )

if __name__ == "__main__":
    demo.queue().launch()