Spaces:
Runtime error
Runtime error
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) |