app.py

import gradio as gr
import spaces
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:
"""

# Function to generate compliment
def generate_compliment(image):
    # Convert PIL image to bytes
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    image_bytes = buffered.getvalue()
    
    try:
        # Connect to the captioning space
        captioning_space = spaces.connect("gokaygokay/sd3-long-captioner")
        
        # Predict caption for the provided image
        caption_response = captioning_space.predict("/create_captions_rich", { "image": image_bytes })
        caption_text = caption_response.data[0]
        
    except Exception as e:
        return "Error", f"Failed to get caption. Exception: {e}"
    
    try:
        # Connect to the LLM space
        llm_space = spaces.connect("hysts/zephyr-7b")
        
        # Generate compliment using the caption
        llm_payload = {
            "system_prompt": SYSTEM_PROMPT,
            "message": f"Caption: {caption_text}\nCompliment: ",
            "max_new_tokens": 256,
            "temperature": 0.7,
            "top_p": 0.95,
            "top_k": 50,
            "repetition_penalty": 1,
        }
        
        compliment_response = llm_space.run(llm_payload)
        compliment_text = compliment_response.data[0]
        
    except Exception as e:
        return "Error", f"Failed to generate compliment. Exception: {e}"
    
    return caption_text, compliment_text

# Gradio interface setup
# Gradio interface
iface = gr.Interface(
    fn=generate_compliment,
    inputs=gr.Image(type="pil", label="Upload Image"),  # Use gr.inputs.Image for image upload
    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()