img

app.py

@@ -166,8 +166,9 @@ def create_interface():
         # Function to generate image
         def generate_image(text):
             try:
-                image = llm_node.generate_image(text)
-                return image
+                seed = random.randint(0, 1000000)
+                image_path = llm_node.generate_image(text, seed=seed)
+                return image_path
             except Exception as e:
                 print(f"An error occurred while generating the image: {e}")
                 return None

llm_inference.py

@@ -5,6 +5,7 @@ from openai import OpenAI
 import requests
 from PIL import Image
 import io
+from gradio_client import Client
 
 class LLMInferenceNode:
     def __init__(self):
@@ -21,6 +22,7 @@ class LLMInferenceNode:
             api_key=self.sambanova_api_key,
             base_url="https://api.sambanova.ai/v1",
         )
+        self.flux_client = Client("KingNish/Realtime-FLUX", hf_token=self.huggingface_token)
 
     def generate_prompt(self, dynamic_seed, prompt_type, custom_input):
         """
@@ -242,15 +244,15 @@ Your output is only the caption itself, no comments or extra formatting. The cap
             print(f"An error occurred: {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
+    def generate_image(self, prompt, seed=42, width=1024, height=1024):
+        try:
+            result = self.flux_client.predict(
+                prompt=prompt,
+                seed=seed,
+                width=width,
+                height=height,
+                api_name="/generate_image"
+            )
+            return result
+        except Exception as e:
+            raise Exception(f"Error generating image: {str(e)}")