Created separate page for Flux models

pages/1_Comparing_LLMs.py

@@ -1,185 +0,0 @@
-from dotenv import load_dotenv
-import os
-from PIL import Image
-import streamlit as st
-import fireworks.client
-
-st.set_page_config(page_title="LLM Comparison Tool", page_icon="🎇")
-st.title("LLM-as-a-judge: Comparing LLMs using Fireworks")
-st.write("A light introduction to how easy it is to swap LLMs and how to use the Fireworks Python client")
-
-# Clear the cache before starting
-st.cache_data.clear()
-
-# Specify the path to the .env file in the env/ directory
-dotenv_path = os.path.join(os.path.dirname(__file__), '..', 'env', '.env')
-
-# Load the .env file from the specified path
-load_dotenv(dotenv_path)
-
-# Get the Fireworks API key from the environment variable
-fireworks_api_key = os.getenv("FIREWORKS_API_KEY")
-
-if not fireworks_api_key:
-    raise ValueError("No API key found in the .env file. Please add your FIREWORKS_API_KEY to the .env file.")
-
-# Load the image
-logo_image = Image.open("img/fireworksai_logo.png")
-ash_image = Image.open("img/ash.png")
-bulbasaur_image = Image.open("img/bulbasaur.png")
-squirtel_image = Image.open("img/squirtel.png")
-charmander_image = Image.open("img/charmander.png")
-
-st.divider()
-# Streamlit app
-st.subheader("Fireworks Playground")
-
-st.write("Fireworks AI is a platform that offers serverless and scalable AI models.")
-st.write("👉 Learn more here: [Fireworks Serverless Models](https://fireworks.ai/models?show=Serverless)")
-st.divider()
-
-# Sidebar for selecting models
-with st.sidebar:
-    st.image(logo_image)
-
-    st.write("Select three models to compare their outputs:")
-
-    st.image(bulbasaur_image, width=80)
-    option_1 = st.selectbox("Select Model 1", [
-        "Text: Meta Llama 3.1 Instruct - 70B",
-        "Text: Meta Llama 3.1 Instruct - 8B",
-        "Text: Meta Llama 3.2 Instruct - 3B",
-        "Text: Gemma 2 Instruct - 9B",
-        "Text: Mixtral MoE Instruct - 8x22B",
-        "Text: Mixtral MoE Instruct - 8x7B",
-        "Text: MythoMax L2 - 13B"
-    ], index=2)  # Default to Meta Llama 3.2 Instruct - 3B
-
-    st.image(charmander_image, width=80)
-    option_2 = st.selectbox("Select Model 2", [
-        "Text: Meta Llama 3.1 Instruct - 70B",
-        "Text: Meta Llama 3.1 Instruct - 8B",
-        "Text: Meta Llama 3.2 Instruct - 3B",
-        "Text: Gemma 2 Instruct - 9B",
-        "Text: Mixtral MoE Instruct - 8x22B",
-        "Text: Mixtral MoE Instruct - 8x7B",
-        "Text: MythoMax L2 - 13B"
-    ], index=5)  # Default to Mixtral MoE Instruct - 8x7B
-
-    st.image(squirtel_image, width=80)
-    option_3 = st.selectbox("Select Model 3", [
-        "Text: Meta Llama 3.1 Instruct - 70B",
-        "Text: Meta Llama 3.1 Instruct - 8B",
-        "Text: Meta Llama 3.2 Instruct - 3B",
-        "Text: Gemma 2 Instruct - 9B",
-        "Text: Mixtral MoE Instruct - 8x22B",
-        "Text: Mixtral MoE Instruct - 8x7B",
-        "Text: MythoMax L2 - 13B"
-    ], index=0)  # Default to Gemma 2 Instruct - 9B
-
-    # Dropdown to select the LLM that will perform the comparison
-    st.image(ash_image, width=80)
-    comparison_llm = st.selectbox("Select Comparison Model", [
-        "Text: Meta Llama 3.1 Instruct - 70B",
-        "Text: Meta Llama 3.1 Instruct - 8B",
-        "Text: Meta Llama 3.2 Instruct - 3B",
-        "Text: Gemma 2 Instruct - 9B",
-        "Text: Mixtral MoE Instruct - 8x22B",
-        "Text: Mixtral MoE Instruct - 8x7B",
-        "Text: MythoMax L2 - 13B"
-    ], index=5) # Default to MythoMax L2 - 13B
-
-os.environ["FIREWORKS_API_KEY"] = fireworks_api_key
-
-# Helper text for the prompt
-st.markdown("### Enter your prompt below to generate responses:")
-
-prompt = st.text_input("Prompt", label_visibility="collapsed")
-st.divider()
-
-# Function to generate a response from a text model
-def generate_text_response(model_name, prompt):
-    return fireworks.client.ChatCompletion.create(
-        model=model_name,
-        messages=[{
-            "role": "user",
-            "content": prompt,
-        }]
-    )
-
-# Function to compare the three responses using the selected LLM
-def compare_responses(response_1, response_2, response_3, comparison_model):
-    comparison_prompt = f"Compare the following three responses:\n\nResponse 1: {response_1}\n\nResponse 2: {response_2}\n\nResponse 3: {response_3}\n\nProvide a succinct comparison."
-    
-    comparison_response = fireworks.client.ChatCompletion.create(
-        model=comparison_model,  # Use the selected LLM for comparison
-        messages=[{
-            "role": "user",
-            "content": comparison_prompt,
-        }]
-    )
-    
-    return comparison_response.choices[0].message.content
-
-
-# If Generate button is clicked
-if st.button("Generate"):
-    if not fireworks_api_key.strip() or not prompt.strip():
-        st.error("Please provide the missing fields.")
-    else:
-        try:
-            with st.spinner("Please wait..."):
-                fireworks.client.api_key = fireworks_api_key
-                
-                # Create three columns for side-by-side comparison
-                col1, col2, col3 = st.columns(3)
-                
-                # Model 1
-                with col1:
-                    st.subheader(f"Model 1: {option_1}")
-                    st.image(bulbasaur_image)
-                    if option_1.startswith("Text"):
-                        model_map = {
-                            "Text: Meta Llama 3.1 Instruct - 70B": "accounts/fireworks/models/llama-v3p1-70b-instruct",
-                            "Text: Meta Llama 3.1 Instruct - 8B": "accounts/fireworks/models/llama-v3p1-8b-instruct",
-                            "Text: Meta Llama 3.2 Instruct - 3B": "accounts/fireworks/models/llama-v3p2-3b-instruct",
-                            "Text: Gemma 2 Instruct - 9B": "accounts/fireworks/models/gemma2-9b-it",
-                            "Text: Mixtral MoE Instruct - 8x22B": "accounts/fireworks/models/mixtral-8x22b-instruct",
-                            "Text: Mixtral MoE Instruct - 8x7B": "accounts/fireworks/models/mixtral-8x7b-instruct",
-                            "Text: MythoMax L2 - 13B": "accounts/fireworks/models/mythomax-l2-13b"
-                        }
-                        response_1 = generate_text_response(model_map[option_1], prompt)
-                        st.success(response_1.choices[0].message.content)
-                
-                # Model 2
-                with col2:
-                    st.subheader(f"Model 2: {option_2}")
-                    st.image(charmander_image)
-                    response_2 = generate_text_response(model_map[option_2], prompt)
-                    st.success(response_2.choices[0].message.content)
-                
-                # Model 3
-                with col3:
-                    st.subheader(f"Model 3: {option_3}")
-                    st.image(squirtel_image)
-                    response_3 = generate_text_response(model_map[option_3], prompt)
-                    st.success(response_3.choices[0].message.content)
-
-                # Visual divider between model responses and comparison
-                st.divider()
-
-                # Generate a comparison of the three responses using the selected LLM
-                comparison = compare_responses(
-                    response_1.choices[0].message.content, 
-                    response_2.choices[0].message.content, 
-                    response_3.choices[0].message.content, 
-                    model_map[comparison_llm]
-                )
-                
-                # Display the comparison
-                st.subheader("Comparison of the Three Responses:")
-                st.image(ash_image)
-                st.write(comparison)
-                
-        except Exception as e:
-            st.exception(f"Exception: {e}")

pages/3_Image_Generation.py

@@ -1,28 +1,36 @@
-from dotenv import load_dotenv
 import os
+import requests
+from dotenv import load_dotenv
 from PIL import Image
+from io import BytesIO
 import streamlit as st
 import fireworks.client
 from fireworks.client.image import ImageInference, Answer
 
+# Set page configuration - must be the first Streamlit command
 st.set_page_config(page_title="Image Generation Tool", page_icon="🎇")
-st.title("Image Generation Comparison using Fireworks")
-st.write("An introduction to how easy it is to generate images using the Fireworks Python client.")
 
-# Clear the cache before starting
-st.cache_data.clear()
+# Set the full path to the .env file
+dotenv_path = os.path.join(os.path.dirname(__file__), '.env')
 
-# Specify the path to the .env file in the env/ directory
-dotenv_path = os.path.join(os.path.dirname(__file__), '..', 'env', '.env')
+# Load environment variables from the .env file, overriding existing environment variables if necessary
+load_dotenv(dotenv_path, override=True)
 
-# Load the .env file from the specified path
-load_dotenv(dotenv_path)
-
-# Get the Fireworks API key from the environment variable
+# Get the Fireworks API key from the .env file
 fireworks_api_key = os.getenv("FIREWORKS_API_KEY")
 
+# Debugging check: print the API key to ensure it's being loaded correctly
+st.write(f"API Key loaded in Streamlit: {fireworks_api_key}")
+
+# Ensure the API key is loaded
 if not fireworks_api_key:
-    raise ValueError("No API key found in the .env file. Please add your FIREWORKS_API_KEY to the .env file.")
+    raise ValueError("API key not found. Make sure FIREWORKS_API_KEY is set in the .env file.")
+
+st.title("Image Generation Comparison using Fireworks")
+st.write("An introduction to how easy it is to generate images using the Fireworks Python client.")
+
+# Clear the cache before starting
+st.cache_data.clear()
 
 # Load image for logo
 logo_image = Image.open("img/fireworksai_logo.png")
@@ -41,29 +49,25 @@ with st.sidebar:
     
     st.write("Select three image generation models to compare:")
 
-    # Updated model options with the correct paths and additional FLUX models
+    # Updated model options (FLUX models removed)
     model_options = {
         "Stable Diffusion XL": "stable-diffusion-xl-1024-v1-0",
         "Playground v2 1024": "playground-v2-1024px-aesthetic",
         "Playground v2.5 1024": "playground-v2-5-1024px-aesthetic",
-        "Segmind Stable Diffusion 1B (SSD-1B)": "SSD-1B",
-        "FLUX.1 [dev]": "flux-1-dev",
-        "FLUX.1 [schnell]": "flux-1-schnell"
+        "Segmind Stable Diffusion 1B (SSD-1B)": "SSD-1B"
     }
 
     option_1 = st.selectbox("Select Image Model 1", list(model_options.keys()), index=0)
     option_2 = st.selectbox("Select Image Model 2", list(model_options.keys()), index=1)
     option_3 = st.selectbox("Select Image Model 3", list(model_options.keys()), index=2)
 
-os.environ["FIREWORKS_API_KEY"] = fireworks_api_key
-
 # Helper text for the prompt
 st.markdown("### Enter your prompt below to generate images:")
 
 prompt = st.text_input("Prompt", label_visibility="collapsed")
 st.divider()
 
-# Function to generate a response from an image model
+# Function to generate an image using the Fireworks API models
 def generate_image_response(model_path, prompt):
     # Initialize the ImageInference client
     inference_client = ImageInference(model=model_path)
@@ -117,4 +121,3 @@ if st.button("Generate"):
 
         except Exception as e:
             st.exception(f"Exception: {e}")
-

pages/5_FLUX_image_generation.py

@@ -0,0 +1,77 @@
+import os
+import requests
+from dotenv import load_dotenv
+from PIL import Image
+from io import BytesIO
+import streamlit as st
+
+# Set page configuration
+st.set_page_config(page_title="FLUX Image Generation Tool", page_icon="🎇")
+
+# Correct the path to the .env file to reflect its location
+dotenv_path = os.path.join(os.path.dirname(__file__), '../env/.env')
+
+# Load environment variables from the .env file
+load_dotenv(dotenv_path, override=True)
+
+# Get the Fireworks API key from the .env file
+fireworks_api_key = os.getenv("FIREWORKS_API_KEY")
+
+if not fireworks_api_key:
+    st.error("API key not found. Make sure FIREWORKS_API_KEY is set in the .env file.")
+
+# Function to make requests to the FLUX models
+def generate_flux_image(model_path, prompt, steps):
+    url = f"https://api.fireworks.ai/inference/v1/workflows/accounts/fireworks/models/{model_path}/text_to_image"
+    headers = {
+        "Authorization": f"Bearer {fireworks_api_key}",
+        "Content-Type": "application/json",
+        "Accept": "image/jpeg"
+    }
+    data = {
+        "prompt": prompt,
+        "aspect_ratio": "16:9",
+        "guidance_scale": 3.5,
+        "num_inference_steps": steps,
+        "seed": 0
+    }
+
+    # Send the request
+    response = requests.post(url, headers=headers, json=data)
+    
+    if response.status_code == 200:
+        # Convert the response to an image
+        img_data = response.content
+        img = Image.open(BytesIO(img_data))
+        return img
+    else:
+        raise RuntimeError(f"Error with FLUX model {model_path}: {response.text}")
+
+# Streamlit UI
+st.title("FLUX Image Generation")
+st.write("Generate images using the FLUX models.")
+
+# User input for the prompt
+prompt = st.text_input("Enter your prompt for image generation:")
+
+# Dropdown to select the model
+model_choice = st.selectbox("Select the model:", ["flux-1-dev", "flux-1-schnell"])
+
+# Button to generate images
+if st.button("Generate Image"):
+    if not prompt.strip():
+        st.error("Please provide a prompt.")
+    else:
+        try:
+            with st.spinner("Generating image..."):
+                # Determine steps based on model
+                steps = 30 if model_choice == "flux-1-dev" else 4
+
+                # Generate image
+                generated_image = generate_flux_image(model_choice, prompt, steps)
+
+                # Display the image
+                st.image(generated_image, caption=f"Generated using {model_choice}", use_column_width=True)
+                
+        except Exception as e:
+            st.error(f"An error occurred: {e}")

pages/test_endpoint.py

@@ -0,0 +1,64 @@
+import os
+import requests
+from dotenv import load_dotenv
+from PIL import Image
+from io import BytesIO
+
+# Correct the path to the .env file to reflect its location
+dotenv_path = os.path.join(os.path.dirname(__file__), '../env/.env')
+
+# Load environment variables from the .env file
+load_dotenv(dotenv_path, override=True)
+
+# Get the API key from the .env file
+api_key = os.getenv("FIREWORKS_API_KEY")
+
+if not api_key:
+    raise ValueError("API key not found. Make sure FIREWORKS_API_KEY is set in the .env file.")
+
+# User input for the prompt
+prompt = input("Enter a prompt for image generation: ")
+
+# Validate the prompt input
+if not prompt.strip():
+    raise ValueError("Prompt cannot be empty!")
+
+# Set the model endpoint for either flux-1-dev or flux-1-schnell
+# For dev: "flux-1-dev" (30 steps)
+# For schnell: "flux-1-schnell" (4 steps)
+model_path = "flux-1-schnell"
+# model_path = "flux-1-dev"  # Uncomment if you want to switch to the dev model
+
+# API URL for the model
+url = f"https://api.fireworks.ai/inference/v1/workflows/accounts/fireworks/models/{model_path}/text_to_image"
+
+# Headers for the API request
+headers = {
+    "Authorization": f"Bearer {api_key}",
+    "Content-Type": "application/json",
+    "Accept": "image/jpeg"
+}
+
+# Data payload to send with the request
+data = {
+    "prompt": prompt,  # Use the user-provided prompt
+    "aspect_ratio": "16:9",
+    "guidance_scale": 3.5,
+    "num_inference_steps": 30 if model_path == "flux-1-dev" else 4,  # 30 steps for dev, 4 for schnell
+    "seed": 0
+}
+
+# Make the POST request to the API
+response = requests.post(url, headers=headers, json=data)
+
+# Check the status of the response
+if response.status_code == 200:
+    # If the request is successful, convert the response to an image
+    img_data = response.content
+    img = Image.open(BytesIO(img_data))
+    # Save the image
+    img.save("output_image.jpg")
+    print("Image saved successfully as output_image.jpg.")
+else:
+    # If there's an error, print the status code and response text
+    print(f"Error: {response.status_code}, {response.text}")