File size: 2,112 Bytes
e29cdd1 57eccf2 e29cdd1 57eccf2 e29cdd1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
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')
#print("dotenv_path: ", dotenv_path)
# 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-fp8"
model_path = "flux-1-dev-fp8" # 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}")
|