import gradio as gr
from transformers import AutoProcessor, AutoModelForCausalLM  # یا AutoModelForVision2Seq اگر vision2seq باشه
import torch
from PIL import Image

# repo_id دقیق
MODEL_ID = "erfanasghariyan/mobilew-v11-convnext-tiny-6layer-radimagenet"

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID)   # اگر causalLM نباشه، عوض کن به AutoModelForVision2Seq
model.eval()

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

def generate_report(image: Image.Image, prompt: str = "Describe this radiology image in detail:"):
    # پردازش تصویر + متن prompt
    inputs = processor(images=image, text=prompt, return_tensors="pt").to(device)
    
    # generation
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=200,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            num_beams=4,  # اگر deterministic می‌خوای beam search
            repetition_penalty=1.2
        )
    
    # decode خروجی
    generated_text = processor.decode(outputs[0], skip_special_tokens=True)
    
    # اگر prompt در خروجی تکرار شده، تمیز کن
    if generated_text.startswith(prompt):
        generated_text = generated_text[len(prompt):].strip()
    
    return generated_text, image

demo = gr.Interface(
    fn=generate_report,
    inputs=[
        gr.Image(type="pil", label="تصویر رادیولوژی آپلود کن (X-ray, CT, MRI و ...)"),
        gr.Textbox(label="پرامپت (اختیاری)", value="Generate a detailed radiology report for this image:")
    ],
    outputs=[
        gr.Textbox(label="گزارش / توصیف تولید شده"),
        gr.Image(label="تصویر ورودی")
    ],
    title="MobileW-v11 – ConvNeXt Tiny + 6-Layer Decoder for Radiology",
    description=(
        "مدل سبک multimodal برای تحلیل تصاویر پزشکی.\n"
        "encoder: ConvNeXt-Tiny (فریز شده)\n"
        "decoder: 6 لایه برای تولید متن بهتر\n"
        "مثال پرامپت‌ها: 'Describe findings', 'Write a radiology report', 'Is there pneumonia?'"
    ),
    examples=[
        # اگر مثال تصویر داری، آپلود کن به Space و مسیر بده
        # [ "example_chest_xray.jpg", "Write a structured report" ]
    ],
    allow_flagging="never"
)

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