image generation

huggingface_inference_node.py

@@ -5,6 +5,7 @@ from datetime import datetime
 import anthropic
 from groq import Groq
 from openai import OpenAI
+from gradio_client import Client
 
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 groq_api_key = os.getenv("GROQ_API_KEY")
@@ -22,6 +23,8 @@ class LLMInferenceNode:
             api_key=sambanova_api_key,
             base_url="https://api.sambanova.ai/v1",
         )
+        self.huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
+        self.flux_client = Client("KingNish/Realtime-FLUX", hf_token=self.huggingface_token)
 
     def generate(
         self,
@@ -177,3 +180,18 @@ You are allowed to make up film and branding names, and do them like 80's, 90's
         except Exception as e:
             print(f"An error occurred: {e}")
             return f"Error occurred while processing the request: {str(e)}"
+
+    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"
+            )
+            # Extract the image path from the result tuple
+            image_path = result[0]
+            return image_path
+        except Exception as e:
+            raise Exception(f"Error generating image: {str(e)}")

ui_components.py

@@ -141,6 +141,14 @@ def create_interface():
                 generate_text_button = gr.Button("Generate Prompt with LLM")
                 text_output = gr.Textbox(label="Generated Text", lines=10, show_copy_button=True)
 
+            with gr.Column(scale=2):
+                with gr.Accordion("Image Generation", open=True):
+                    image_output = gr.Image(label="Generated Image", type="filepath")
+                    generate_image_button = gr.Button("Generate Image")
+                    image_seed = gr.Number(label="Image Seed", value=42, step=1)
+                    image_width = gr.Slider(label="Width", minimum=512, maximum=2048, value=1024, step=64)
+                    image_height = gr.Slider(label="Height", minimum=512, maximum=2048, value=1024, step=64)
+
         def create_caption(image, model):
             if image is not None:
                 if model == "Florence-2":
@@ -264,4 +272,21 @@ def create_interface():
             ]
         )
 
+        # Function to generate image
+        def generate_image(text, seed, width, height):
+            try:
+                image_path = llm_node.generate_image(text, seed=seed, width=width, height=height)
+                print(f"Image generated: {image_path}")
+                return image_path
+            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, image_seed, image_width, image_height],
+            outputs=[image_output]
+        )
+
     return demo