import os
from huggingface_hub import login

# 🔐 Login avec Hugging Face token (optionnel si déjà connecté)
hf_token = os.environ.get("HUGGINGFACE_API_KEY")
if hf_token:
    login(hf_token)

from transformers import AutoModelForCausalLM, AutoProcessor
import torch
import gradio as gr

model_id = "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1"

# ✅ Chargement du modèle et processor
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    torch_dtype=torch.float16,
    device_map="auto"
)

processor = AutoProcessor.from_pretrained(
    model_id,
    trust_remote_code=True
)

# 🧠 Fonction principale
def generate_answer(image, question):
    if not question:
        question = "Please describe this medical image."
    
    prompt = f"### User: {question}\n### Assistant:"
    
    try:
        if image is None:
            inputs = processor(prompt, return_tensors="pt").to(model.device)
        else:
            inputs = processor(prompt, images=image, return_tensors="pt").to(model.device)

        outputs = model.generate(**inputs, max_new_tokens=256)
        response = processor.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response[len(prompt):].strip()
    except Exception as e:
        return f"⚠️ Internal Error: {str(e)}"

# 🎛️ Interface Gradio avec API activée
demo = gr.Interface(
    fn=generate_answer,
    inputs=[
        gr.Image(type="pil", label="Image médicale (optionnelle)"),
        gr.Textbox(label="Question médicale")
    ],
    outputs="text",
    title="🧠 ContactDoctor - Biomedical LLM",
    description="Assistant médical multimodal. Posez une question ou envoyez une image."
)

# 🚀 Lancement avec API activée
demo.launch()