Reproduced the fjord example

pages/6_Test_Control_Net_Flux.py

@@ -0,0 +1,143 @@
+import cv2
+import requests
+from io import BytesIO
+from PIL import Image
+import numpy as np
+import os
+from dotenv import load_dotenv
+import streamlit as st
+
+# Set the page title for the Streamlit app
+st.set_page_config(page_title="ControlNet Image Generation")
+
+# 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
+api_key = os.getenv("FIREWORKS_API_KEY")
+
+if not api_key:
+    st.error("API key not found. Make sure FIREWORKS_API_KEY is set in the .env file.")
+    st.stop()
+
+# Load image and apply Canny edge detection
+def process_image(uploaded_image):
+    # Convert the uploaded image into an OpenCV-compatible format (grayscale)
+    image = np.array(Image.open(uploaded_image).convert('L'))
+
+    # Apply Canny edge detection
+    edges = cv2.Canny(image, 100, 200)
+
+    # Convert the single-channel edges image to a 3-channel image (RGB)
+    edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
+
+    # Convert the edges image to a format that can be sent in the HTTP request
+    pil_image = Image.fromarray(edges_rgb)
+    byte_arr = BytesIO()
+    pil_image.save(byte_arr, format='JPEG')
+    byte_arr.seek(0)
+
+    return byte_arr, pil_image
+
+# Function to make the POST request to the API
+def call_control_net_api(uploaded_image, prompt, control_mode=0, aspect_ratio="16:9", 
+                         guidance_scale=3.5, num_inference_steps=30, seed=0, controlnet_conditioning_scale=1.0):
+    # Process the image for control net
+    control_image, processed_image = process_image(uploaded_image)
+
+    # Prepare the payload
+    files = {
+        'control_image': ('control_image.jpg', control_image, 'image/jpeg')
+    }
+    data = {
+        'prompt': prompt,
+        'control_mode': control_mode,  # Control how the control image is used
+        'aspect_ratio': aspect_ratio,
+        'guidance_scale': guidance_scale,
+        'num_inference_steps': num_inference_steps,
+        'seed': seed,
+        'controlnet_conditioning_scale': controlnet_conditioning_scale  # Control how strongly the control image influences
+    }
+
+    headers = {
+        'accept': 'image/jpeg',
+        'authorization': f'Bearer {api_key}',  # Using the API key loaded from the .env file
+    }
+
+    # Send the POST request
+    response = requests.post('https://api.fireworks.ai/inference/v1/workflows/accounts/fireworks/models/flux-1-dev-controlnet-union/control_net', 
+                             files=files, data=data, headers=headers)
+
+    # Handle the response
+    if response.status_code == 200:
+        return Image.open(BytesIO(response.content)), processed_image
+    else:
+        st.error(f"Request failed with status code: {response.status_code}, Response: {response.text}")
+        return None, None
+
+# Streamlit UI
+st.title("ControlNet Image Generation with Fireworks")
+st.write("Upload an image, provide a prompt, and let the model generate an image using Canny edge detection as input.")
+
+# File uploader for image input
+uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
+
+# Text input for prompt
+prompt = st.text_input("Enter your prompt")
+
+# Expander to hide/show additional parameters
+with st.expander("Advanced Parameters"):
+    # Slider for `control_mode`
+    control_mode = st.slider("Control Mode", min_value=0, max_value=2, value=0, help="Control how the edge image affects generation")
+
+    # Slider for `controlnet_conditioning_scale`
+    controlnet_conditioning_scale = st.slider("ControlNet Conditioning Scale", min_value=0.0, max_value=1.0, value=0.5, step=0.1, 
+                                              help="Adjust how strongly the edge-detected image influences the output")
+
+    # Dropdown for `aspect_ratio`
+    aspect_ratio = st.selectbox("Aspect Ratio", options=["16:9", "1:1", "4:3", "3:2", "9:16"], index=0)
+
+    # Slider for `guidance_scale`
+    guidance_scale = st.slider("Guidance Scale", min_value=0.0, max_value=20.0, value=3.5, step=0.1, 
+                               help="Adjust how strongly the model adheres to the prompt")
+
+    # Slider for `num_inference_steps`
+    num_inference_steps = st.slider("Number of Inference Steps", min_value=1, max_value=100, value=30, step=1, 
+                                    help="Number of steps to generate the image")
+
+    # Slider for `seed`
+    seed = st.slider("Random Seed", min_value=0, max_value=1000, value=0, 
+                     help="Set a seed for reproducibility (0 means random)")
+
+# Button to submit
+if st.button("Generate Image"):
+    if uploaded_image is None:
+        st.error("Please upload an image.")
+    elif not prompt.strip():
+        st.error("Please enter a prompt.")
+    else:
+        with st.spinner("Processing..."):
+            # Display the uploaded image first
+            st.subheader("Uploaded Image")
+            st.image(uploaded_image, caption="Original Uploaded Image", use_column_width=True)
+
+            # Call the ControlNet API
+            generated_image, processed_image = call_control_net_api(uploaded_image, prompt, 
+                                                                    control_mode=control_mode, 
+                                                                    aspect_ratio=aspect_ratio,
+                                                                    guidance_scale=guidance_scale,
+                                                                    num_inference_steps=num_inference_steps,
+                                                                    seed=seed,
+                                                                    controlnet_conditioning_scale=controlnet_conditioning_scale)
+
+            if generated_image:
+                # Hide the processed edge-detected image in an expander
+                with st.expander("Edge Detection Result (Input to ControlNet)"):
+                    st.image(processed_image, caption="Processed Edge Detection Image", use_column_width=True)
+
+                # Display the generated image from the API
+                st.subheader("Generated Image")
+                st.image(generated_image, caption="Generated Image from ControlNet", use_column_width=True)

requirements.txt

@@ -135,6 +135,7 @@ numpy==1.26.4
 oauthlib==3.2.2
 onnxruntime==1.19.2
 openai==1.47.0
+opencv-python==4.10.0.84
 opentelemetry-api==1.27.0
 opentelemetry-exporter-otlp-proto-common==1.27.0
 opentelemetry-exporter-otlp-proto-grpc==1.27.0