import gradio as gr
import torch
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
from PIL import Image

# Это точное имя вашей модели, которое мы нашли!
MODEL_ID = "livadies/Russian-Radiologist-ruGPT-ViT"

print("Загрузка модели с Hugging Face (это займет около минуты)...")
# Бесплатные серверы (Spaces) работают на процессоре (CPU)
device = "cpu" 

try:
    model = VisionEncoderDecoderModel.from_pretrained(MODEL_ID).to(device)
    feature_extractor = ViTImageProcessor.from_pretrained(MODEL_ID)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    print("✅ Модель успешно загружена в память!")
except Exception as e:
    print(f"❌ ОШИБКА ЗАГРУЗКИ МОДЕЛИ: {e}")
    print("Возможно, имя модели указано неверно или файлы не догрузились.")

def analyze_xray(image):
    if image is None:
        return "Пожалуйста, загрузите снимок."
    
    try:
        # Подготовка картинки
        image = image.convert("RGB")
        pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
        
        # Генерация текста
        with torch.no_grad():
            generated_ids = model.generate(
                pixel_values,
                max_length=128,
                num_beams=4,
                repetition_penalty=2.0
            )
        
        # Расшифровка текста
        prediction = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
        return prediction
    except Exception as e:
        return f"Произошла ошибка при анализе: {str(e)}"

# Настройка красивого веб-интерфейса
interface = gr.Interface(
    fn=analyze_xray,
    inputs=gr.Image(type="pil", label="Загрузите рентгеновский снимок (X-Ray)"),
    outputs=gr.Textbox(label="Заключение Нейро-Радиолога", lines=6),
    title="🩻 Русский Нейро-Радиолог (AI Assistant)",
    description="Загрузите рентгеновский снимок грудной клетки. Искусственный интеллект, обученный на 6500 снимках, проанализирует его и сгенерирует медицинское заключение на русском языке.\n\n*Внимание: Модель носит исключительно исследовательский характер и не заменяет консультацию врача.*",
    theme="huggingface"
)

if __name__ == "__main__":
    interface.launch()