import gradio as gr
import requests
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration, pipeline
import time

# Carregando o modelo BLIP para geração de legendas
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
model_blip = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")

# Carregando um modelo de geração de texto (exemplo: GPT-2)
generator = pipeline('text-generation', model='gpt2')

def caption(img, min_len, max_len):
    raw_image = Image.open(img).convert('RGB')
    inputs = processor(raw_image, return_tensors="pt")
    out = model_blip.generate(**inputs, min_length=min_len, max_length=max_len)
    return processor.decode(out[0], skip_special_tokens=True)

def generate_nutritional_info(food_description):
    # Gerando informações nutricionais com base na descrição do alimento
    prompt = f"Provide detailed nutritional information about {food_description}."
    result = generator(prompt, max_length=150, num_return_sequences=1)
    return result[0]['generated_text']

def greet(img, min_len, max_len):
    start = time.time()
    
    # Passo 1: Gerar legenda para a imagem
    food_description = caption(img, min_len, max_len)
    
    # Passo 2: Gerar informações nutricionais com base na legenda
    nutritional_info = generate_nutritional_info(food_description)
    
    end = time.time()
    total_time = str(end - start)
    
    # Combinando resultados
    result = f"Food Description: {food_description}\n\nNutritional Information:\n{nutritional_info}\n\nGenerated in {total_time} seconds."
    return result

# Interface Gradio
iface = gr.Interface(
    fn=greet,
    title='Nutritionist Agent with BLIP and GPT-2',
    description="Upload an image of food, and the agent will describe it and provide nutritional information.",
    inputs=[
        gr.Image(type='filepath', label='Image'),
        gr.Slider(label='Minimum Length', minimum=1, maximum=1000, value=30),
        gr.Slider(label='Maximum Length', minimum=1, maximum=1000, value=100)
    ],
    outputs=gr.Textbox(label='Result'),
    theme=gr.themes.Base(primary_hue="teal", secondary_hue="teal", neutral_hue="slate"),
)

iface.launch()