from PIL import Image
import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration

model_id = "Salesforce/blip-image-captioning-base"

model = BlipForConditionalGeneration.from_pretrained(model_id)
processor = BlipProcessor.from_pretrained(model_id)

def launch(input_file):
    # Abrir la imagen subida por el usuario
    image = Image.open(input_file).convert('RGB')
    
    # Procesar la imagen y generar el texto
    inputs = processor(image, return_tensors="pt")
    out = model.generate(**inputs)
    generated_text = processor.decode(out[0], skip_special_tokens=True)
    
    # Agregar prefijo y sufijo al texto generado
    prefijo = "Texture, "
    sufijo = " top view, flat, hq, detailed"
    result = f"{prefijo}{generated_text}{sufijo}"
    
    # Devolver tanto la imagen como el texto
    return image, result

# Cambiar la configuración de Gradio para que la entrada sea un archivo y la salida incluya imagen y texto
iface = gr.Interface(
    fn=launch, 
    inputs="file", 
    outputs=["image", "text"], 
    title="GENERADOR DE TEXTO DEL DR. SERGIO"
)

iface.launch()