Spaces:
Sleeping
Sleeping
File size: 3,475 Bytes
6349504 |
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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
# image_handler.py
import time
from typing import Dict, Any
from datetime import datetime
import requests
import base64
from requests.exceptions import RequestException
import tempfile
import os
class ImageHandler:
def __init__(self):
self.available_models = ["flux", "flux-realism", "flux-3d", "flux-pro",
"flux-anime", "flux-cablyai", "turbo", "flux-dark"]
self.base_url = "https://image.pollinations.ai/prompt/{prompt}"
def _get_size_dimensions(self, size: str) -> tuple:
size_map = {
"256x256": (256, 256),
"512x512": (512, 512),
"1024x1024": (1024, 1024),
"1792x1024": (1792, 1024),
"1024x1792": (1024, 1792)
}
return size_map.get(size, (768, 768))
def _save_temp_image(self, image_data: bytes) -> str:
"""Save image to temporary file and return its URL."""
temp_dir = tempfile.gettempdir()
temp_file = os.path.join(temp_dir, f"image_{int(time.time())}.jpg")
with open(temp_file, 'wb') as f:
f.write(image_data)
return f"file://{temp_file}"
def generate_images(self, params: Dict[str, Any]) -> Dict[str, Any]:
prompt = params.get("prompt")
n = min(params.get("n", 1), 10)
size = params.get("size", "1024x1024")
model = params.get("model", "flux-pro")
if model not in self.available_models:
model = "flux-pro"
response_format = params.get("response_format", "b64_json")
if response_format != "b64_json":
response_format = "bs64_json"
width, height = self._get_size_dimensions(size)
headers = {
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
"User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:122.0) Gecko/20100101 Firefox/122.0",
}
images = []
for _ in range(n):
try:
url = self.base_url.format(prompt=prompt)
url += f"?width={width}&height={height}&model={model}"
response = requests.get(url, headers=headers, timeout=60)
response.raise_for_status()
# Get the image data
image_data = response.content
if response_format == "b64_json":
# Convert to base64
image_b64 = base64.b64encode(image_data).decode('utf-8')
images.append({"b64_json": image_b64})
else:
# Two options for URL:
# 1. Save locally and return file URL (good for testing)
# 2. Return the pollinations URL (might expire)
# Option 1: Local file
temp_url = self._save_temp_image(image_data)
images.append({"url": temp_url})
# Option 2: Pollinations URL (uncomment if preferred)
# images.append({"url": url})
except RequestException as e:
raise Exception(f"Image generation failed: {str(e)}")
return {
"created": int(datetime.now().timestamp()),
"data": images
}
|