Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces
 import requests
 from PIL import Image
 import io
@@ -17,56 +18,56 @@ Compliment: You are the epitome of elegance and grace, with a style that is as t
 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()
     
-    # Connect to the captioning model on Hugging Face Spaces using the correct URL and method
-    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:
+        # Connect to the captioning space
+        captioning_space = spaces.connect("gokaygokay/sd3-long-captioner")
+        
+        # Predict caption for the provided image
+        caption = captioning_space.predict("/create_captions_rich", { "image": image_bytes })
+        
+        # Extract the caption from the response
+        caption_text = caption.data[0]
+        
+    except Exception as e:
         return "Error", f"Failed to get caption. Exception: {e}"
     
     try:
-        caption = caption_response.json()["data"][0]
+        # 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 = llm_space.run(llm_payload)
+        compliment_text = compliment.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,
-    }
-    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}"
+        return "Error", f"Failed to generate compliment. Exception: {e}"
     
-    try:
-        compliment = llm_response.json()["data"][0]
-    except Exception as e:
-        return "Error", f"Failed to parse LLM response. Error: {str(e)}, Response: {llm_response.text}"
-
-    return caption, compliment
+    return caption_text, compliment_text
 
-# Gradio interface
+# Gradio interface setup
 iface = gr.Interface(
     fn=generate_compliment,
-    inputs=gr.Image(type="pil"),
+    inputs=gr.inputs.Image(type="pil", label="Upload Image"),
     outputs=[
-        gr.Textbox(label="Caption"),
-        gr.Textbox(label="Compliment")
+        gr.outputs.Textbox(label="Caption"),
+        gr.outputs.Textbox(label="Compliment")
     ],
     title="Compliment Bot 💖",
     description="Upload your headshot and get a personalized compliment!"