from transformers import AutoModelForCausalLM, AutoProcessor
from pathlib import Path
import torch
import os
from huggingface_hub import login

HF_TOKEN=os.getenv('HF_TOKEN')
login(HF_TOKEN)

model = AutoModelForCausalLM.from_pretrained("microsoft/maira-2", trust_remote_code=True)
processor = AutoProcessor.from_pretrained("microsoft/maira-2", trust_remote_code=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.is_available():
    print("Using GPU for inference")
model = model.eval()
model = model.to(device)



def generate_report(frontal_image = None, lateral_image = None, indication="", technique="", comparison="" ):
    inputs = processor.format_and_preprocess_reporting_input(
            current_frontal=frontal_image,
            current_lateral=lateral_image,
            prior_frontal=None,
            indication=indication,
            technique=technique,
            comparison=comparison,
            prior_report=None,
            return_tensors="pt",
            get_grounding=False,
        )
    inputs = inputs.to(device)
    with torch.no_grad():
        output_decoding = model.generate(
            **inputs,
            max_new_tokens=300,  # Set to 450 for grounded reporting
            use_cache=True,
        )
    prompt_length = inputs["input_ids"].shape[-1]
    decoded_text = processor.decode(output_decoding[0][prompt_length:], skip_special_tokens=True)
    decoded_text = decoded_text.lstrip() 
    prediction = processor.convert_output_to_plaintext_or_grounded_sequence(decoded_text)
    print("Parsed prediction:", prediction)
    return prediction


print("Maira-2 service is ready to generate reports.")