import gradio as gr
from transformers import pipeline
from PIL import Image
import time

print("📥 Chargement de BLIP-large (meilleure qualité)...")
start = time.time()

# BLIP-large au lieu de BLIP-base
captioner = pipeline(
    "image-to-text",  # Ce pipeline accepte image seule
    model="Salesforce/blip-image-captioning-large",
    device=-1
)

print(f"✅ Modèle chargé en {time.time() - start:.1f} secondes")

def generate_caption(image, prompt_text=""):
    """Génère une légende pour l'image"""
    try:
        if hasattr(image, 'size'):
            pil_image = image
        else:
            pil_image = Image.fromarray(image.astype('uint8'))
        
        # BLIP-large ne gère pas les prompts complexes
        # On utilise simplement l'image
        result = captioner(pil_image, max_new_tokens=50)
        
        if isinstance(result, list) and len(result) > 0:
            caption = result[0].get('generated_text', str(result[0]))
        else:
            caption = str(result)
        
        return caption.strip()
    except Exception as e:
        return f"Erreur : {str(e)}"

with gr.Blocks(title="Générateur de titres d'images (BLIP-large)") as demo:
    gr.Markdown("""
    # 🖼️ Générateur de titres d'images
    **Modèle BLIP-large (meilleure qualité)**
    - Temps de réponse : 2-4 minutes par image ⏱️
    - Le prompt est ignoré car ce modèle ne le supporte pas bien
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(label="📸 Uploadez votre image", type="pil")
            btn = gr.Button("🚀 Générer le titre", variant="primary")
        
        with gr.Column(scale=1):
            output = gr.Textbox(label="🏷️ Titre généré", lines=4)
    
    btn.click(
        fn=generate_caption,
        inputs=[image_input],
        outputs=output
    )

demo.launch(server_name="0.0.0.0", server_port=7860)