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)