cherry-ghosts/dreamer

1import random2import time3from io import BytesIO4from os import getenv5from urllib.parse import quote6 7import requests8from PIL import Image9 10class Inference:11    def __init__(self):12        self.session = requests.Session()13 14    def _fetch(self, *args, retries=5, **kwargs):15        for retry in range(retries + 1):16            try:17                response = self.session.get(*args, **kwargs)18                response.raise_for_status()19                return response.content20            except requests.RequestException:21                if e.response.status_code in {429, 500, 502, 503, 504}:22                    delay = min(2 ** retry, 60)23                    time.sleep(delay)24                else:25                    return None26        return None27 28    def __call__(29        self,30        prompt,31        seed=42,32        width=1024,33        height=102434    ):35        url = f"https://image.pollinations.ai/prompt/{quote(prompt)}"36        params = {37            "prompt": prompt,38            "seed": seed,39            "width": width,40            "height": height,41            "enhance": "true",42            "nologo": "true",43            "private": "true",44            "referrer": "cherry-ghosts",45            "token": getenv("TOKEN")46        }47        image = self._fetch(url, params=params)48        if image:49            return Image.open(BytesIO(image))50        return None