image generation

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 from llm_inference import LLMInferenceNode
 import random
+from PIL import Image
+import io
 
 title = """<h1 align="center">Random Prompt Generator</h1>
 <p><center>
@@ -69,6 +71,7 @@ def create_interface():
             generate_button = gr.Button("GENERATE")
         with gr.Row():
             text_output = gr.Textbox(label="LLM Generated Text", lines=10, show_copy_button=True)
+            image_output = gr.Image(label="Generated Image", type="pil")
 
         # Updated Models based on provider
         def update_model_choices(provider):
@@ -157,6 +160,25 @@ def create_interface():
             api_name="generate_random_prompt_with_llm"
         )
 
+        # Add image generation button
+        generate_image_button = gr.Button("Generate Image")
+
+        # Function to generate image
+        def generate_image(text):
+            try:
+                image = llm_node.generate_image(text)
+                return image
+            except Exception as e:
+                print(f"An error occurred while generating the image: {e}")
+                return None
+
+        # Connect the image generation button
+        generate_image_button.click(
+            generate_image,
+            inputs=[text_output],
+            outputs=[image_output]
+        )
+
     return demo
 
 if __name__ == "__main__":

llm_inference.py

@@ -2,6 +2,9 @@ import os
 import random  # Import the random module
 from groq import Groq
 from openai import OpenAI
+import requests
+from PIL import Image
+import io
 
 class LLMInferenceNode:
     def __init__(self):
@@ -237,4 +240,17 @@ Your output is only the caption itself, no comments or extra formatting. The cap
 
         except Exception as e:
             print(f"An error occurred: {e}")
-            return f"Error occurred while processing the request: {str(e)}"
+            return f"Error occurred while processing the request: {str(e)}"
+
+    def generate_image(self, prompt):
+        API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
+        headers = {"Authorization": f"Bearer {self.huggingface_token}"}
+
+        response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
+        
+        if response.status_code != 200:
+            raise Exception(f"Error generating image: {response.text}")
+
+        image_bytes = response.content
+        image = Image.open(io.BytesIO(image_bytes))
+        return image