client.py

import random
import time
from io import BytesIO
from os import getenv
from urllib.parse import quote

import requests
from PIL import Image

class Inference:
    def __init__(self):
        self.session = requests.Session()

    def _fetch(self, *args, retries=5, **kwargs):
        for retry in range(retries + 1):
            try:
                response = self.session.get(*args, **kwargs)
                response.raise_for_status()
                return response.content
            except requests.RequestException:
                if e.response.status_code in {429, 500, 502, 503, 504}:
                    delay = min(2 ** retry, 60)
                    time.sleep(delay)
                else:
                    return None
        return None

    def __call__(
        self,
        prompt,
        seed=42,
        width=1024,
        height=1024
    ):
        url = f"https://image.pollinations.ai/prompt/{quote(prompt)}"
        params = {
            "prompt": prompt,
            "seed": seed,
            "width": width,
            "height": height,
            "enhance": "true",
            "nologo": "true",
            "private": "true",
            "referrer": "cherry-ghosts",
            "token": getenv("TOKEN")
        }
        image = self._fetch(url, params=params)
        if image:
            return Image.open(BytesIO(image))
        return None