app.py

from gradio_client import Client, file
import gradio as gr
from PIL import Image
import requests
import io

# Configuration for Hugging Face Spaces
CAPTION_SPACE = "gokaygokay/SD3-Long-Captioner"
LLM_SPACE = "hysts/zephyr-7b"
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:
"""

# Initialize Gradio client for captioning and language model
captioning_client = Client(CAPTION_SPACE)
llm_client = Client(LLM_SPACE)

def generate_compliment(image):
    try:
        # Convert PIL image to bytes
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        image_bytes = buffered.getvalue()

        # Get caption for the image using Gradio client
        caption_response = captioning_client.predict("/create_captions_rich", { "image": file(image_bytes) })
        caption_text = caption_response.data[0]

        # Generate compliment based on the caption using language model
        llm_response = llm_client.predict(SYSTEM_PROMPT, f"Caption: {caption_text}\nCompliment: ")
        compliment_text = llm_response.data[0]

    except Exception as e:
        caption_text = f"Error: Failed to get caption. Exception: {str(e)}"
        compliment_text = "Error: Failed to generate compliment."

    return caption_text, compliment_text

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

iface.launch()