client.py

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

import requests
from PIL import Image

MAX_SEED = 2 ** 31 - 1

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

    def _fetch(self, *args, retries=5, **kwargs):
        for retry in range(retries + 1):
            try:
                response = session.get(*args, **kwargs)
                response.raise_for_status()
                return response.content
            except requests.RequestException:
                if 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=None,
        randomize_seed=True,
        width=1024,
        height=1024
    ):
        if seed is None or randomize_seed:
            seed = random.randint(0, MAX_SEED)
        url = f"https://image.pollinations.ai/p/{prompt}"
        params = {
            "prompt": prompt,
            "seed": seed,
            "width": width,
            "height": height
        }
        image = self._fetch(url, params=params)
        if image:
            return Image.open(BytesIO(image))
        return None