Update app.py

app.py

@@ -4,15 +4,23 @@ from PIL import Image
 import gradio as gr
 import os
 
-# Ensure your API token is correctly set in your environment variables
-API_TOKEN = os.getenv("HF_API_TOKEN")
-if not API_TOKEN:
-    raise ValueError("Hugging Face API token not found. Please set the HF_API_TOKEN environment variable.")
+# Assuming you have your API tokens set in environment variables
+ZEPHYR_API_TOKEN = os.getenv("HF_API_TOKEN")
+SD_API_TOKEN = os.getenv("HF_API_TOKEN")
 
-API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
-headers = {"Authorization": f"Bearer {API_TOKEN}"}
+if not ZEPHYR_API_TOKEN or not SD_API_TOKEN:
+    raise ValueError("API tokens not found. Please set the ZEPHYR_API_TOKEN and HF_API_TOKEN environment variables.")
 
-def generate_image(prompt, negative_prompt="", guidance_scale=7.5, width=1024, height=768, num_inference_steps=30):
+ZEPHYR_API_URL = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta"
+SD_API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
+
+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()
+
+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}"}
     payload = {
         "inputs": prompt,
         "parameters": {
@@ -22,27 +30,28 @@ def generate_image(prompt, negative_prompt="", guidance_scale=7.5, width=1024, h
             "num_inference_steps": num_inference_steps,
         },
     }
-    if negative_prompt:  # Only add negative_prompt to payload if it's provided
-        payload["parameters"]["negative_prompt"] = negative_prompt
-
-    response = requests.post(API_URL, headers=headers, json=payload)
+    response = requests.post(SD_API_URL, headers=headers, json=payload)
     image_bytes = response.content
     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
+    if generated_prompt:
+        return generate_image_from_prompt(generated_prompt)
+    else:
+        raise ValueError("Failed to generate a prompt from the LinkedIn text.")
+
 iface = gr.Interface(
-    fn=generate_image,
-    inputs=[
-        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
-        gr.Textbox(label="Negative Prompt", placeholder="Enter a negative prompt here (optional)..."),
-        gr.Slider(label="Guidance Scale", minimum=1, maximum=20, step=0.1, value=7.5),
-        gr.Slider(label="Width", minimum=768, maximum=1024, step=1, value=1024),
-        gr.Slider(label="Height", minimum=768, maximum=1024, step=1, value=768),
-        gr.Slider(label="Number of Inference Steps", minimum=20, maximum=50, step=1, value=30)
-    ],
+    fn=generate_image_from_linkedin_text,
+    inputs=[gr.Textbox(label="LinkedIn Message", placeholder="Enter LinkedIn message here...")],
     outputs=gr.Image(type="pil"),
-    title="Stable Diffusion XL Image Generator",
-    description="Generate images with Stable Diffusion XL. Provide a prompt, optionally specify a negative prompt, and adjust other parameters as desired."
+    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."
 )
 
 iface.launch()