GenAI-Arena / model /models /playground_api.py
DongfuJiang's picture
update
bd0bce1
raw
history blame
1.22 kB
import os
import json
import requests
from PIL import Image
import io
import base64
class PlayGround():
def __init__(self, model_name, model_type=None):
self.model_name = model_name
self.model_type = model_type
self.api_key = os.environ['PlaygroundAPI']
if model_name == "PlayGroundV2":
self._model_name = "Playground_v2"
elif model_name == "PlayGroundV2.5":
self._model_name = "Playground_v2.5"
def __call__(self, prompt):
headers = {
'Content-Type': 'application/json',
'Authorization': "Bearer " + self.api_key,
}
data = json.dumps({"prompt": prompt, "filter_model": self._model_name, "scheduler": "DPMPP_2M_K", "guidance_scale": 3})
response = requests.post('https://playground.com/api/models/external/v1', headers=headers, data=data)
response.raise_for_status()
json_obj = response.json()
image_base64 = json_obj['images'][0]
img = Image.open(io.BytesIO(base64.decodebytes(bytes(image_base64, "utf-8"))))
return img
def load_playground_model(model_name, model_type="generation"):
return PlayGround(model_name, model_type)