Update app.py

app.py

@@ -1,13 +1,10 @@
 import gradio as gr
-import torch
-from transformers import pipeline
-
-# Load the models
-caption_generator = pipeline("image-captioning", model="gokaygokay/SD3-Long-Captioner")
-compliment_generator = pipeline("text-generation", model="hysts/zephyr-7b")
+import requests
+from PIL import Image
+import io
 
 SYSTEM_PROMPT = """
-You are helpful assistant that gives the best compliments to people.
+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.
@@ -21,12 +18,30 @@ Conversation begins below:
 """
 
 def generate_compliment(image):
-    # Generate caption
-    caption = caption_generator(image)[0]['caption']
-    # Generate compliment
-    input_text = f"Caption: {caption}\nCompliment: "
-    generated_compliment = compliment_generator(SYSTEM_PROMPT + input_text, max_new_tokens=50, do_sample=True)[0]['generated_text']
-    compliment = generated_compliment.split('Compliment: ')[-1].strip()
+    # 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