import gradio as gr
import openai
import torch
from lavis.models import load_model_and_preprocess
from PIL import Image

openai.api_key = ''


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model, vis_processors, _ = load_model_and_preprocess(name="blip_caption", model_type="base_coco", is_eval=True, device=device)

def generate_caption(raw_image):
    raw_image = raw_image.convert('RGB')
    image = vis_processors["eval"](raw_image).unsqueeze(0).to(device)
    return model.generate({"image": image})[0]

def listing(title):
    response1 = openai.Completion.create(
        model="text-davinci-003",
        prompt='''Given the image caption as: {}.\n\nCan you generate a product title for me?'''.format(title),
        temperature=0.9,
        max_tokens=150,
        top_p=1,
        frequency_penalty=0.0,
        presence_penalty=0.6)
    
    response2 = openai.Completion.create(
        model="text-davinci-003",
        prompt='''Given the product title: {}.\n\nCan you create an Amazon product listing with bullet points and description.'''.format(response1.choices[0].text),
        temperature=0.9,
        max_tokens=300,
        top_p=1,
        frequency_penalty=0.0,
        presence_penalty=0.6)
    
    return response1.choices[0].text, response2.choices[0].text


demo = gr.Blocks()

with demo:

    gr.Markdown("# EcomGenius!😄")

    with gr.Tabs():
        with gr.TabItem("Image Captioning"):
            with gr.Row():
                with gr.Column():
                    input = gr.Image(label="Upload your Image", type = 'pil')
                    caption_button = gr.Button("Caption It", variant="primary")

                with gr.Column():
                    output = gr.components.Textbox(label="Caption!")

        with gr.TabItem("Product Listing Creation"):
            with gr.Row():
                with gr.Column():
                    input2 = gr.Textbox(lines=1, label="Image caption?")
                    generate_button = gr.Button("Create", variant="primary")

                with gr.Column():
                    output2 = gr.components.Textbox(label="Product Title")
                    output3 = gr.components.Textbox(label="Product listing Info")

    caption_button.click(fn=generate_caption, inputs=[input], outputs=[output])
    generate_button.click(fn=listing, inputs=[input2], outputs=[output2, output3])

demo.launch()