app.py

import gradio as gr
from gradio_client import Client
import requests
from PIL import Image
import json
import re
import time
from huggingface_hub import InferenceClient

client = InferenceClient(
    "mistralai/Mixtral-8x7B-Instruct-v0.1"
)

def generate(prompt, system_prompt, max_new_tokens):
    generate_kwargs = dict(
        temperature=0.7,
        max_new_tokens=512,
        top_p=0.95,
        repetition_penalty=1.0,
        do_sample=True,
        seed=42,
    )

    formatted_prompt = f" [INST] {system_prompt}, {prompt} [/INST]"
    output = ""

    for response in client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False):
        output += response.token.text

    return translate(output)

def translate(input):
    generate_kwargs = dict(
        temperature=0.7,
        max_new_tokens=512,
        top_p=0.95,
        repetition_penalty=1.0,
        do_sample=True,
        seed=42,
    )

    system_prompt = "Traduza para português."

    formatted_prompt = f" [INST] {system_prompt}, {input} [/INST]"
    output = ""

    for response in client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False):
        output += response.token.text

    return output

def get_caption_from_kosmos(image_in):
    kosmos2_client = Client("https://ydshieh-kosmos-2.hf.space/")

    kosmos2_result = kosmos2_client.predict(
        image_in,	# str (filepath or URL to image) in 'Test Image' Image component
        "Detailed",	# str in 'Description Type' Radio component
        fn_index=4
    )

    print(f"KOSMOS2 RETURNS: {kosmos2_result}")

    with open(kosmos2_result[1], 'r') as f:
        data = json.load(f)
    
    reconstructed_sentence = []
    for sublist in data:
        reconstructed_sentence.append(sublist[0])

    full_sentence = ' '.join(reconstructed_sentence)
    #print(full_sentence)

    # Find the pattern matching the expected format ("Describe this image in detail:" followed by optional space and then the rest)...
    pattern = r'^Describe this image in detail:\s*(.*)$'
    # Apply the regex pattern to extract the description text.
    match = re.search(pattern, full_sentence)
    if match:
        description = match.group(1)
        print(description)
    else:
        print("Unable to locate valid description.")

    # Find the last occurrence of "."
    last_period_index = description.rfind('.')

    # Truncate the string up to the last period
    truncated_caption = description[:last_period_index + 1]

    # print(truncated_caption)
    print(f"\n—\nIMAGE CAPTION: {truncated_caption}")
    
    return truncated_caption

def caption(img, max_len):
  
    img_out = get_caption_from_kosmos(img)
    translated_img_out = translate(img_out)
    
    system_prompt = "Ignore todas as instruções anteriores. Considerando os detalhes fornecidos, crie uma descrição de marketing com um tom atraente/marketing para potenciais clientes, usando palavras que os empoderem. Apresente as informações em um único parágrafo coeso. Se algum caractere especial estiver presente, como '<' ou '>' ou qualquer outro semelhante, você DEVE ignorar e não adicionar à descrição. Qualquer emoji ou emoticon ou qualquer símbolo semelhante **DEVE** ser ignorado, não importa o quê, sem exceções. Abstenha-se de oferecer sugestões sobre o uso ou combinações do produto. Quando uma cor for dada, não implique que várias cores do produto estão disponíveis. Em vez disso, coloque as cores dadas como a cor do produto. Se uma cor não for fornecida, não mencione sobre ela na descrição que você escrever. Empregue uma linguagem propícia a resultados ótimos de mecanismos de busca no Google. Em casos em que o usuário fornece informações aparentemente inalteráveis, reformule ou reescreva o texto sem introduzir detalhes desnecessários. DESCONSIDERE QUALQUER referência a peso (gramas (g), quilogramas (kg), etc.) se mencionado nos detalhes fornecidos. Quando mais de uma cor for dada, não diga que é para produtos diferentes. Em vez disso, insira-as juntas, dizendo que é apenas um produto que tem todas as cores dadas de uma forma natural. Não use palavras clichês no início da frase. Além disso, você DEVE sempre fornecer respostas em português. Adira meticulosamente a estas instruções e abstenha-se de desviar das diretrizes fornecidas. Abstenha-se de fabricar qualquer informação. Todas as instruções dadas **DEVEM** ser respeitadas e seguidas, não importa o quê."


    text_out = generate(translated_img_out, system_prompt, max_len)
    text_out = text_out.replace("</s>", "")
    text_out = text_out.strip()
    return text_out


iface= gr.Interface(fn=caption, 
                     title='Image Captionizer in Marketing Tone', 
                     description="._.",
                     inputs=[gr.Image(type='filepath', label='Image'), gr.Slider(label='Maximum Length', minimum=1, maximum=1000, value=100)], 
                     outputs=gr.Textbox(label='Caption', show_copy_button=True),
                     theme = gr.themes.Base(primary_hue="teal",secondary_hue="teal",neutral_hue="slate"),)
iface.launch()