app.py

import gradio as gr
import requests
from PIL import Image
import io

SYSTEM_PROMPT = """
You are a helpful assistant that gives the best compliments to people.
You will be given a caption of someone's headshot.
Based on that caption, provide a one sentence compliment to the person in the image.
Make sure you compliment the person in the image and not any objects or scenery.
Do NOT include any hashtags in your compliment or phrases like (emojis: dog, smiling face with heart-eyes, sun).
Here are some examples of the desired behavior:
Caption: a front view of a man who is smiling, there is a lighthouse in the background, there is a grassy area on the left that is green and curved. in the distance you can see the ocean and the shore. there is a grey and cloudy sky above the lighthouse and the trees.
Compliment: Your smile is as bright as a lighthouse, lighting up the world around you. 🌟
Caption: in a close-up, a blonde woman with short, wavy hair, is the focal point of the image. she's dressed in a dark brown turtleneck sweater, paired with a black hat and a black suit jacket. her lips are a vibrant red, and her eyes are a deep brown. in the background, a man with a black hat and a white shirt is visible.
Compliment: You are the epitome of elegance and grace, with a style that is as timeless as your beauty. 💃🎩
Conversation begins below:
"""

def generate_compliment(image):
    # Convert PIL image to bytes
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    image_bytes = buffered.getvalue()
    
    # Connect to the captioning model on Hugging Face Spaces
    captioning_url = "https://gokaygokay-sd3-long-captioner.hf.space/run/create_captions_rich"
    caption_response = requests.post(captioning_url, files={"image": image_bytes})
    caption = caption_response.json()["data"][0]
    
    # Connect to the LLM model on Hugging Face Spaces
    llm_url = "https://hysts-zephyr-7b.hf.space/run/chat"
    llm_payload = {
        "system_prompt": SYSTEM_PROMPT,
        "message": f"Caption: {caption}\nCompliment: ",
        "max_new_tokens": 256,
        "temperature": 0.7,
        "top_p": 0.95,
        "top_k": 50,
        "repetition_penalty": 1,
    }
    llm_response = requests.post(llm_url, json=llm_payload)
    compliment = llm_response.json()["data"][0]

    return caption, compliment

# Gradio interface
iface = gr.Interface(
    fn=generate_compliment,
    inputs=gr.inputs.Image(type="pil"),
    outputs=[
        gr.outputs.Textbox(label="Caption"),
        gr.outputs.Textbox(label="Compliment")
    ],
    title="Compliment Bot 💖",
    description="Upload your headshot and get a personalized compliment!"
)

iface.launch()