eve / FLUX.py
Chandima Prabhath
flux & audio reply
b9724da
raw
history blame
5.16 kB
import requests
import time
import io
import os
import re
import json
from PIL import Image, UnidentifiedImageError
from dotenv import load_dotenv
load_dotenv()
# Load the ImgBB API key from the environment variables.
IMGBB_API_KEY = os.getenv("IMGBB_API_KEY")
def upload_to_imgbb(image_path, file_name):
"""
Uploads the image located at image_path to ImgBB.
Returns:
str: URL of the uploaded image on ImgBB or None if failed.
"""
try:
with open(image_path, 'rb') as f:
image_data = f.read()
response = requests.post(
"https://api.imgbb.com/1/upload",
params={"key": IMGBB_API_KEY},
files={"image": (file_name, image_data)}
)
response.raise_for_status()
result = response.json()
if result.get("data") and "url" in result["data"]:
return result["data"]["url"]
else:
print("Failed to upload image to ImgBB.")
return None
except requests.RequestException as e:
print(f"Error uploading image to ImgBB: {e}")
return None
except Exception as e:
print(f"Unexpected error uploading image to ImgBB: {e}")
return None
def generate_image(prompt, request_id, current_request_id, image_dir, attempt=0):
"""
Generate an image using the Pollinations API.
Parameters:
prompt (str): The prompt for image generation.
width (int): Desired image width.
height (int): Desired image height.
request_id (int): The request id for the current operation.
current_request_id (int): The current active request id.
image_dir (str): Directory where image will be saved.
attempt (int): Current attempt count (zero-indexed).
Returns:
tuple: (PIL.Image object, image_path (str), returned_prompt (str), image_url (str))
or None if image fetch fails or request id mismatches.
"""
model = "flux"
width = 1920
height = 1080
enhance_param = "true"
url = f"https://image.pollinations.ai/prompt/{prompt}?nologo=true&safe=false&private=true&model={model}&enhance={enhance_param}&width={width}&height={height}"
print(f"Attempt {attempt + 1}: Fetching image with URL: {url}")
try:
response = requests.get(url, timeout=45)
except Exception as e:
print(f"Error fetching image: {e}")
return None
if response.status_code != 200:
print(f"Failed to fetch image. Status code: {response.status_code}")
return None
if request_id != current_request_id:
print("Request ID mismatch. Operation cancelled.")
return None
print("Image fetched successfully.")
image_data = response.content
try:
image = Image.open(io.BytesIO(image_data))
actual_width, actual_height = image.size
print(f"Actual image dimensions: {actual_width}x{actual_height}")
# Extract metadata from EXIF if available
exif_data = image.info.get('exif', b'')
returned_prompt = prompt
if exif_data:
json_match = re.search(b'{"prompt":.*}', exif_data)
if json_match:
json_str = json_match.group(0).decode('utf-8')
try:
metadata_dict = json.loads(json_str)
returned_prompt = metadata_dict.get('prompt', prompt)
except json.JSONDecodeError as e:
print(f"Failed to parse JSON in metadata: {e}")
else:
print("No JSON data found in EXIF")
if (actual_width, actual_height) != (width, height):
print(f"Warning: Received image dimensions ({actual_width}x{actual_height}) do not match requested dimensions ({width}x{height})")
except UnidentifiedImageError:
print("Error: Received data is not a valid image.")
raise
timestamp = int(time.time())
image_filename = f"background_{timestamp}.png"
image_path = os.path.join(image_dir, image_filename)
# Ensure the image directory exists
os.makedirs(image_dir, exist_ok=True)
try:
image.save(image_path, 'PNG')
print(f"Image saved to {image_path}")
# Upload image to ImgBB
image_url = upload_to_imgbb(image_path, image_filename)
if image_url:
print(f"Image uploaded to ImgBB: {image_url}")
else:
print("Failed to upload image to ImgBB.")
except Exception as e:
print(f"Error saving image: {e}")
return None
return image, image_path, returned_prompt, image_url
if __name__ == "__main__":
# Example usage
prompt = "Beach party, anime style, vibrant colors"
request_id = 1
current_request_id = 1
image_dir = "./images"
image, image_path, returned_prompt, image_url = generate_image(prompt, request_id, current_request_id, image_dir)
if image:
print(f"Image generated and saved at {image_path}")
print(f"Returned prompt: {returned_prompt}")
print(f"Image URL: {image_url}")
else:
print("Failed to generate image.")