Update app.py

app.py

@@ -1,8 +1,12 @@
+from gradio_client import Client, file
 import gradio as gr
-import requests
 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.
@@ -17,53 +21,35 @@ Compliment: You are the epitome of elegance and grace, with a style that is as t
 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):
-    # 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"
-    try:
-        caption_response = requests.post(captioning_url, files={"image": ("image.jpg", image_bytes, "image/jpeg")})
-        caption_response.raise_for_status()  # Raise an exception for HTTP errors
-    except requests.exceptions.RequestException as e:
-        return "Error", f"Failed to get caption. Exception: {e}"
-    
-    try:
-        caption = caption_response.json()["data"][0]
-    except Exception as e:
-        return "Error", f"Failed to parse caption response. Error: {str(e)}, Response: {caption_response.text}"
-    
-    # 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,
-    }
+    compliment_text = ""
     try:
-        llm_response = requests.post(llm_url, json=llm_payload)
-        llm_response.raise_for_status()  # Raise an exception for HTTP errors
-    except requests.exceptions.RequestException as e:
-        return "Error", f"Failed to get compliment. Exception: {e}"
-    
-    try:
-        compliment = llm_response.json()["data"][0]
+        # 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:
-        return "Error", f"Failed to parse LLM response. Error: {str(e)}, Response: {llm_response.text}"
+        compliment_text = f"Error: {str(e)}"
+
+    return caption_text, compliment_text
 
-    return caption, compliment
-    
 # Gradio interface
 iface = gr.Interface(
     fn=generate_compliment,
-    inputs=gr.Image(type="pil", label="Upload Image"),  # Use gr.inputs.Image for image upload
+    inputs=gr.Image(type="pil"),
     outputs=[
         gr.Textbox(label="Caption"),
         gr.Textbox(label="Compliment")