Update app.py

app.py

@@ -4,7 +4,7 @@ from PIL import Image
 import gradio as gr
 import os
 
-# Assuming you have your API tokens set in environment variables
+# Assuming your API tokens are correctly set in environment variables
 ZEPHYR_API_TOKEN = os.getenv("HF_API_TOKEN")
 SD_API_TOKEN = os.getenv("HF_API_TOKEN")
 
@@ -17,7 +17,14 @@ SD_API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-dif
 def query_zephyr(prompt):
     headers = {"Authorization": f"Bearer {ZEPHYR_API_TOKEN}"}
     response = requests.post(ZEPHYR_API_URL, headers=headers, json={"inputs": prompt})
-    return response.json()
+    # Assuming the response is correctly formatted JSON and the first item contains the desired data
+    response_data = response.json()
+    if isinstance(response_data, list) and len(response_data) > 0:
+        # Extracting the generated text from the first item in the response list
+        generated_text = response_data[0].get("generated_text", "")
+        return generated_text
+    else:
+        raise ValueError("Unexpected response format from Zephyr API")
 
 def generate_image_from_prompt(prompt, guidance_scale=7.5, width=1024, height=768, num_inference_steps=30):
     headers = {"Authorization": f"Bearer {SD_API_TOKEN}"}
@@ -35,23 +42,23 @@ def generate_image_from_prompt(prompt, guidance_scale=7.5, width=1024, height=76
     image = Image.open(io.BytesIO(image_bytes))
     return image
 
-def generate_image_from_linkedin_text(linkedin_text):
-    # Step 1: Generate a prompt from the LinkedIn text using Zephyr
-    zephyr_output = query_zephyr(linkedin_text)
-    generated_prompt = zephyr_output.get("generated_text", "")
-
-    # Step 2: Use the generated prompt to create an image with Stable Diffusion
+def generate_image_and_show_prompt(linkedin_text):
+    # Generate a prompt from the LinkedIn text using Zephyr
+    generated_prompt = query_zephyr(linkedin_text)
+    
+    # Use the generated prompt to create an image with Stable Diffusion
     if generated_prompt:
-        return generate_image_from_prompt(generated_prompt)
+        image = generate_image_from_prompt(generated_prompt)
+        return image, generated_prompt
     else:
         raise ValueError("Failed to generate a prompt from the LinkedIn text.")
 
 iface = gr.Interface(
-    fn=generate_image_from_linkedin_text,
+    fn=generate_image_and_show_prompt,
     inputs=[gr.Textbox(label="LinkedIn Message", placeholder="Enter LinkedIn message here...")],
-    outputs=gr.Image(type="pil"),
+    outputs=[gr.Image(type="pil"), gr.Textbox(label="Generated Prompt")],
     title="Generate Images from LinkedIn Messages",
-    description="Enter a LinkedIn message to generate a creative prompt with Zephyr, which is then used to generate an image with Stable Diffusion."
+    description="Enter a LinkedIn message to generate a creative prompt with Zephyr, which is then used to generate an image with Stable Diffusion. The generated prompt is also displayed."
 )
 
 iface.launch()