Update app.py

app.py

@@ -4,8 +4,8 @@ from PIL import Image
 import gradio as gr
 import os
 
-# Assuming your API tokens are correctly set in environment variables
-ZEPHYR_API_TOKEN = os.getenv("HF_API_TOKEN")
+# Assuming you have your API tokens set in environment variables
+ZEPHYR_API_TOKEN = os.getenv("ZEPHYR_API_TOKEN")
 SD_API_TOKEN = os.getenv("HF_API_TOKEN")
 
 if not ZEPHYR_API_TOKEN or not SD_API_TOKEN:
@@ -17,16 +17,9 @@ 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})
-    # 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")
+    return response.json()
 
-def generate_image_from_prompt(prompt, negative_prompt="", guidance_scale=7.5, width=1024, height=768, num_inference_steps=30):
+def generate_image_from_prompt(prompt, negative_prompt, guidance_scale, width, height, num_inference_steps):
     headers = {"Authorization": f"Bearer {SD_API_TOKEN}"}
     payload = {
         "inputs": prompt,
@@ -37,7 +30,7 @@ def generate_image_from_prompt(prompt, negative_prompt="", guidance_scale=7.5, w
             "num_inference_steps": num_inference_steps,
         },
     }
-    if negative_prompt:  # Add negative_prompt to the payload if provided
+    if negative_prompt:  # Add negative prompt if provided
         payload["parameters"]["negative_prompt"] = negative_prompt
 
     response = requests.post(SD_API_URL, headers=headers, json=payload)
@@ -46,22 +39,20 @@ def generate_image_from_prompt(prompt, negative_prompt="", guidance_scale=7.5, w
     return image
 
 def generate_image_from_linkedin_text(linkedin_text, negative_prompt, guidance_scale, width, height, num_inference_steps):
+    # Generate a prompt from the LinkedIn text using Zephyr
     zephyr_response = query_zephyr(linkedin_text)
-    
     if zephyr_response and isinstance(zephyr_response, list):
         generated_prompt = zephyr_response[0].get("generated_text", "")
     else:
         raise ValueError("Unexpected response format from Zephyr model.")
-    
-    print("Generated Prompt:", generated_prompt)  # For backend logs
-    
+
+    # Use the generated prompt to create an image with Stable Diffusion
     if generated_prompt:
         image = generate_image_from_prompt(generated_prompt, negative_prompt, guidance_scale, width, height, num_inference_steps)
         return image, generated_prompt
     else:
         raise ValueError("Failed to generate a prompt from the LinkedIn text.")
 
-# Update Gradio interface with additional parameters
 iface = gr.Interface(
     fn=generate_image_from_linkedin_text,
     inputs=[
@@ -74,10 +65,10 @@ iface = gr.Interface(
     ],
     outputs=[
         gr.Image(type="pil"),
-        gr.Textbox(label="Generated Prompt", readonly=True)
+        gr.Label(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. Additional parameters allow for fine-tuning of the generated image."
+    description="Enter a LinkedIn message to generate a creative prompt with Zephyr, which is then used to generate an image with Stable Diffusion. Image parameters can be adjusted."
 )
 
 iface.launch()