| import json
|
| import os
|
| import yaml
|
| import requests
|
| import pathlib
|
| from aiohttp import web
|
| from server import PromptServer
|
| from ..image import tensor2pil, pil2tensor, image2base64, pil2byte
|
| from ..log import log_node_error
|
|
|
|
|
| root_path = pathlib.Path(__file__).parent.parent.parent.parent
|
| config_path = os.path.join(root_path,'config.yaml')
|
| default_key = [{'name':'Default', 'key':''}]
|
|
|
|
|
| class StabilityAPI:
|
| def __init__(self):
|
| self.api_url = "https://api.stability.ai"
|
| self.api_keys = None
|
| self.api_current = 0
|
| self.user_info = {}
|
|
|
| def getErrors(self, code):
|
| errors = {
|
| 400: "Bad Request",
|
| 403: "ApiKey Forbidden",
|
| 413: "Your request was larger than 10MiB.",
|
| 429: "You have made more than 150 requests in 10 seconds.",
|
| 500: "Internal Server Error",
|
| }
|
| return errors.get(code, "Unknown Error")
|
|
|
| def getAPIKeys(self):
|
| if os.path.isfile(config_path):
|
| with open(config_path, 'r') as f:
|
| data = yaml.load(f, Loader=yaml.FullLoader)
|
| if not data:
|
| data = {'STABILITY_API_KEY': default_key, 'STABILITY_API_DEFAULT':0}
|
| with open(config_path, 'w') as f:
|
| yaml.dump(data, f)
|
| if 'STABILITY_API_KEY' not in data:
|
| data['STABILITY_API_KEY'] = default_key
|
| data['STABILITY_API_DEFAULT'] = 0
|
| with open(config_path, 'w') as f:
|
| yaml.dump(data, f)
|
| api_keys = data['STABILITY_API_KEY']
|
| self.api_current = data['STABILITY_API_DEFAULT']
|
| self.api_keys = api_keys
|
| return api_keys
|
| else:
|
|
|
| with open(config_path, 'w') as f:
|
| data = {'STABILITY_API_KEY': default_key, 'STABILITY_API_DEFAULT':0}
|
| yaml.dump(data, f)
|
| return data['STABILITY_API_KEY']
|
| pass
|
|
|
| def setAPIKeys(self, api_keys):
|
| if len(api_keys) > 0:
|
| self.api_keys = api_keys
|
|
|
| with open(config_path, 'r') as f:
|
| data = yaml.load(f, Loader=yaml.FullLoader)
|
| data['STABILITY_API_KEY'] = api_keys
|
| with open(config_path, 'w') as f:
|
| yaml.dump(data, f)
|
| return True
|
|
|
| def setAPIDefault(self, current):
|
| if current is not None:
|
| self.api_current = current
|
|
|
| with open(config_path, 'r') as f:
|
| data = yaml.load(f, Loader=yaml.FullLoader)
|
| data['STABILITY_API_DEFAULT'] = current
|
| with open(config_path, 'w') as f:
|
| yaml.dump(data, f)
|
| return True
|
|
|
| def generate_sd3_image(self, prompt, negative_prompt, aspect_ratio, model, seed, mode='text-to-image', image=None, strength=1, output_format='png', node_name='easy stableDiffusion3API'):
|
| url = f"{self.api_url}/v2beta/stable-image/generate/sd3"
|
| api_key = self.api_keys[self.api_current]['key']
|
| files = None
|
| data = {
|
| "prompt": prompt,
|
| "mode": mode,
|
| "model": model,
|
| "seed": seed,
|
| "output_format": output_format,
|
| }
|
| if model == 'sd3':
|
| data['negative_prompt'] = negative_prompt
|
|
|
| if mode == 'text-to-image':
|
| files = {"none": ''}
|
| data['aspect_ratio'] = aspect_ratio
|
| elif mode == 'image-to-image':
|
| pil_image = tensor2pil(image)
|
| image_byte = pil2byte(pil_image)
|
| files = {"image": ("output.png", image_byte, 'image/png')}
|
| data['strength'] = strength
|
|
|
| response = requests.post(url,
|
| headers={"authorization": f"{api_key}", "accept": "application/json"},
|
| files=files,
|
| data=data,
|
| )
|
| if response.status_code == 200:
|
| PromptServer.instance.send_sync('stable-diffusion-api-generate-succeed',{"model":model})
|
| json_data = response.json()
|
| image_base64 = json_data['image']
|
| image_data = image2base64(image_base64)
|
| output_t = pil2tensor(image_data)
|
| return output_t
|
| else:
|
| if 'application/json' in response.headers['Content-Type']:
|
| error_info = response.json()
|
| log_node_error(node_name, error_info.get('name', 'No name provided'))
|
| log_node_error(node_name, error_info.get('errors', ['No details provided']))
|
| error_status_text = self.getErrors(response.status_code)
|
| PromptServer.instance.send_sync('easyuse-toast',{"type": "error", "content": error_status_text})
|
| raise Exception(f"Failed to generate image: {error_status_text}")
|
|
|
|
|
| async def getUserAccount(self, cache=True):
|
| url = f"{self.api_url}/v1/user/account"
|
| api_key = self.api_keys[self.api_current]['key']
|
| name = self.api_keys[self.api_current]['name']
|
| if cache and name in self.user_info:
|
| return self.user_info[name]
|
| else:
|
| response = requests.get(url, headers={"Authorization": f"Bearer {api_key}"})
|
| if response.status_code == 200:
|
| user_info = response.json()
|
| self.user_info[name] = user_info
|
| return user_info
|
| else:
|
| PromptServer.instance.send_sync('easyuse-toast',{'type': 'error', 'content': self.getErrors(response.status_code)})
|
| return None
|
|
|
|
|
| async def getUserBalance(self):
|
| url = f"{self.api_url}/v1/user/balance"
|
| api_key = self.api_keys[self.api_current]['key']
|
| response = requests.get(url, headers={
|
| "Authorization": f"Bearer {api_key}"
|
| })
|
| if response.status_code == 200:
|
| return response.json()
|
| else:
|
| PromptServer.instance.send_sync('easyuse-toast', {'type': 'error', 'content': self.getErrors(response.status_code)})
|
| return None
|
|
|
| stableAPI = StabilityAPI()
|
|
|
| @PromptServer.instance.routes.get("/easyuse/stability/api_keys")
|
| async def get_stability_api_keys(request):
|
| stableAPI.getAPIKeys()
|
| return web.json_response({"keys": stableAPI.api_keys, "current": stableAPI.api_current})
|
|
|
| @PromptServer.instance.routes.post("/easyuse/stability/set_api_keys")
|
| async def set_stability_api_keys(request):
|
| post = await request.post()
|
| api_keys = post.get("api_keys")
|
| current = post.get('current')
|
| if api_keys is not None:
|
| api_keys = json.loads(api_keys)
|
| stableAPI.setAPIKeys(api_keys)
|
| if current is not None:
|
| print(current)
|
| stableAPI.setAPIDefault(int(current))
|
| account = await stableAPI.getUserAccount()
|
| balance = await stableAPI.getUserBalance()
|
| return web.json_response({'account': account, 'balance': balance})
|
| else:
|
| return web.json_response({'status': 'ok'})
|
| else:
|
| return web.Response(status=400)
|
|
|
| @PromptServer.instance.routes.post("/easyuse/stability/set_apikey_default")
|
| async def set_stability_api_default(request):
|
| post = await request.post()
|
| current = post.get("current")
|
| if current is not None and current < len(stableAPI.api_keys):
|
| stableAPI.api_current = current
|
| return web.json_response({'status': 'ok'})
|
| else:
|
| return web.Response(status=400)
|
|
|
| @PromptServer.instance.routes.get("/easyuse/stability/user_info")
|
| async def get_account_info(request):
|
| account = await stableAPI.getUserAccount()
|
| balance = await stableAPI.getUserBalance()
|
| return web.json_response({'account': account, 'balance': balance})
|
|
|
| @PromptServer.instance.routes.get("/easyuse/stability/balance")
|
| async def get_balance_info(request):
|
| balance = await stableAPI.getUserBalance()
|
| return web.json_response({'balance': balance})
|
|
|
