app.py

import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import base64
import io
import re

# Настройки модели
REPO_ID = "mradermacher/VisualQuality-R1-7B-GGUF"
# Ищем файл с квантованием Q8 (наилучшее качество для CPU)
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
            )
            # n_ctx=8192 важно для обработки изображений (визуальные токены занимают место)
            llm = Llama(
                model_path=model_path,
                n_ctx=8192,           # Контекстное окно
                n_gpu_layers=0,       # 0 для CPU (бесплатный Space)
                verbose=True,
                chat_format="chatml-function-calling" # Обычно Qwen совместим с chatml
            )
            print("Модель успешно загружена!")
        except Exception as e:
            print(f"Ошибка загрузки: {e}")
            raise e
    return llm

def image_to_base64(image):
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode('utf-8')

def evaluate_image(image, progress=gr.Progress()):
    if image is None:
        return "Пожалуйста, загрузите изображение.", ""
    
    # Инициализация модели (если еще не загружена)
    model = load_model()
    
    # Системный промпт из официального репозитория VisualQuality-R1
    system_prompt = "You are doing the image quality assessment task."
    
    # Формируем сообщение пользователя по шаблону авторов
    # Модель обучена выдавать <answer>score</answer>
    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."
    )
    
    # Конвертация изображения для Llama-cpp
    base64_image = image_to_base64(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("Начало генерации...")
    
    stream = model.create_chat_completion(
        messages=messages,
        max_tokens=1024,
        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
                # Возвращаем накопленный текст для стриминга в UI
                yield full_response, "Вычисляется..."

    # Парсинг итоговой оценки
    # Ищем текст внутри тегов <answer>...</answer>
    score_match = re.search(r'<answer>\s*([\d\.]+)\s*</answer>', full_response)
    final_score = "Не найдено"
    
    if score_match:
        final_score = score_match.group(1)
    
    yield full_response, final_score

# Интерфейс Gradio
with gr.Blocks(title="VisualQuality-R1 (Q8 GGUF)") as demo:
    gr.Markdown("# 👁️ VisualQuality-R1 (7B Q8)")
    gr.Markdown(
        "Оценка качества изображений с использованием рассуждений (Chain of Thought). "
        "Запущено на CPU (бесплатный Space), поэтому генерация может занять время."
    )
    
    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, show_copy_button=True)

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

# Запуск при старте контейнера, чтобы модель скачалась сразу
if __name__ == "__main__":
    demo.queue().launch()