refactor: improve image response handling, adjust aspect ratio defaults, and fix filename construction

- In PollinationsAI.py, modified get_image method to initialize responses set and manage concurrent image fetches with asyncio tasks, adding a while loop to yield responses as they complete
- Changed response index in get_image from 1 to 0 to align with zero-based indexing
- Introduced 'responses' set and 'finished' counter outside inner get_image function for proper progress tracking
- Updated gather() usage to run all get_image tasks concurrently after loop
- In __init__.py, enhanced use_aspect_ratio function: added checks if width and height are None before assigning aspect ratio-based defaults
- Assigned default width and height values for aspect ratios "1:1", "16:9", and "9:16" if not already specified in extra_body
- In copy_images.py, corrected get_filename function to convert tags to strings before joining with '+', ensuring proper filename formatting
- In response.py, refined is_content function to exclude Reasoning objects where is_thinking and token are both None
- Removed __eq__ method from Reasoning class to prevent comparison issues
- In web_search.py, simplified import by removing unused datetime and date modules
This commit is contained in:
hlohaus 2025-05-18 04:51:46 +02:00
parent 57cbd55d74
commit 362c2f0f1a
6 changed files with 46 additions and 33 deletions

View file

@ -1,5 +1,6 @@
from __future__ import annotations
import time
import json
import random
import requests
@ -350,24 +351,34 @@ class PollinationsAI(AsyncGeneratorProvider, ProviderModelMixin):
prompt = quote_plus(prompt)[:2048-len(cls.image_api_endpoint)-len(query)-8]
url = f"{cls.image_api_endpoint}prompt/{prompt}?{query}"
def get_image_url(i: int, seed: Optional[int] = None):
if i == 1:
if i == 0:
if not cache and seed is None:
seed = random.randint(0, 2**32)
else:
seed = random.randint(0, 2**32)
return f"{url}&seed={seed}" if seed else url
async with ClientSession(headers=DEFAULT_HEADERS, connector=get_connector(proxy=proxy)) as session:
async def get_image(i: int, seed: Optional[int] = None):
responses = set()
finished = 0
async def get_image(responses: set, i: int, seed: Optional[int] = None):
nonlocal finished
start = time.time()
async with session.get(get_image_url(i, seed), allow_redirects=False, headers={"referer": referrer}) as response:
try:
await raise_for_status(response)
except Exception as e:
debug.error(f"Error fetching image: {e}")
return str(response.url)
return str(response.url)
yield ImageResponse(await asyncio.gather(*[
get_image(i, seed) for i in range(int(n))
]), prompt)
responses.add(Reasoning(status=f"Image #{i+1} generated in {time.time() - start:.2f}s"))
responses.add(ImageResponse(str(response.url), prompt))
finished += 1
tasks = []
for i in range(int(n)):
tasks.append(asyncio.create_task(get_image(responses, i, seed)))
while finished < n or len(responses) > 0:
while len(responses) > 0:
yield responses.pop()
await asyncio.sleep(0.1)
await asyncio.gather(*tasks)
@classmethod
async def _generate_text(

View file

@ -287,24 +287,25 @@ def to_input_audio(audio: ImageType, filename: str = None) -> str:
def use_aspect_ratio(extra_body: dict, aspect_ratio: str) -> Image:
extra_body = {key: value for key, value in extra_body.items() if value is not None}
if aspect_ratio == "1:1":
extra_body = {
"width": 1024,
"height": 1024,
**extra_body
}
elif aspect_ratio == "16:9":
extra_body = {
"width": 832,
"height": 480,
**extra_body
}
elif aspect_ratio == "9:16":
extra_body = {
"width": 480,
"height": 832,
**extra_body
}
if extra_body.get("width") is None or extra_body.get("height") is None:
if aspect_ratio == "1:1":
extra_body = {
"width": extra_body.get("width", 1024),
"height": extra_body.get("height", 1024),
**extra_body
}
elif aspect_ratio == "16:9":
extra_body = {
"width": extra_body.get("width", 832),
"height": extra_body.get("height", 480),
**extra_body
}
elif aspect_ratio == "9:16":
extra_body = {
"width": extra_body.get("width", 480),
"height": extra_body.get("height", 832),
**extra_body
}
return extra_body
class ImageDataResponse():

View file

@ -79,7 +79,7 @@ async def save_response_media(response: StreamResponse, prompt: str, tags: list[
def get_filename(tags: list[str], alt: str, extension: str, image: str) -> str:
return "".join((
f"{int(time.time())}_",
f"{secure_filename('+'.join([tag for tag in tags if tag]))}+" if tags else "",
f"{secure_filename('+'.join([str(tag) for tag in tags if tag]))}+" if tags else "",
f"{secure_filename(alt)}_",
hashlib.sha256(image.encode()).hexdigest()[:16],
extension

View file

@ -61,6 +61,8 @@ PARAMETER_EXAMPLES = {
"conversation": {"conversation_id": "550e8400-e29b-11d4-a716-...", "message_id": "550e8400-e29b-11d4-a716-..."},
"seed": 42,
"tools": [],
"width": 1024,
"height": 1024,
}
class AbstractProvider(BaseProvider):

View file

@ -7,7 +7,11 @@ from abc import abstractmethod
from urllib.parse import quote_plus, unquote_plus
def is_content(chunk):
return isinstance(chunk, (str, MediaResponse, AudioResponse, Reasoning, ToolCalls))
if isinstance(chunk, Reasoning):
if chunk.is_thinking is None and chunk.token is None:
return False
return True
return isinstance(chunk, (str, MediaResponse, AudioResponse, ToolCalls))
def quote_url(url: str) -> str:
"""
@ -203,11 +207,6 @@ class Reasoning(ResponseType):
return f"{self.status}\n"
return ""
def __eq__(self, other: Reasoning):
return (self.token == other.token and
self.status == other.status and
self.is_thinking == other.is_thinking)
def get_dict(self) -> Dict:
"""Return a dictionary representation of the reasoning."""
if self.label is not None:

View file

@ -5,7 +5,7 @@ import json
import hashlib
from pathlib import Path
from urllib.parse import urlparse, quote_plus
from datetime import datetime, date
from datetime import date
import asyncio
# Optional dependencies