Spaces:
Runtime error
Runtime error
#!/usr/bin/env python | |
#pylint: disable=redefined-outer-name | |
""" | |
helper methods that creates HTTP session with managed connection pool | |
provides async HTTP get/post methods and several helper methods | |
""" | |
import io | |
import os | |
import sys | |
import ssl | |
import base64 | |
import asyncio | |
import logging | |
import aiohttp | |
import requests | |
import urllib3 | |
from PIL import Image | |
from util import Map, log | |
from rich import print # pylint: disable=redefined-builtin | |
sd_url = os.environ.get('SDAPI_URL', "http://127.0.0.1:7860") # api url root | |
sd_username = os.environ.get('SDAPI_USR', None) | |
sd_password = os.environ.get('SDAPI_PWD', None) | |
use_session = True | |
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) | |
ssl.create_default_context = ssl._create_unverified_context # pylint: disable=protected-access | |
timeout = aiohttp.ClientTimeout(total = None, sock_connect = 10, sock_read = None) # default value is 5 minutes, we need longer for training | |
sess = None | |
quiet = False | |
BaseThreadPolicy = asyncio.WindowsSelectorEventLoopPolicy if sys.platform == "win32" and hasattr(asyncio, "WindowsSelectorEventLoopPolicy") else asyncio.DefaultEventLoopPolicy | |
class AnyThreadEventLoopPolicy(BaseThreadPolicy): | |
def get_event_loop(self) -> asyncio.AbstractEventLoop: | |
try: | |
return super().get_event_loop() | |
except (RuntimeError, AssertionError): | |
loop = self.new_event_loop() | |
self.set_event_loop(loop) | |
return loop | |
asyncio.set_event_loop_policy(AnyThreadEventLoopPolicy()) | |
def authsync(): | |
if sd_username is not None and sd_password is not None: | |
return requests.auth.HTTPBasicAuth(sd_username, sd_password) | |
return None | |
def auth(): | |
if sd_username is not None and sd_password is not None: | |
return aiohttp.BasicAuth(sd_username, sd_password) | |
return None | |
async def result(req): | |
if req.status != 200: | |
if not quiet: | |
log.error({ 'request error': req.status, 'reason': req.reason, 'url': req.url }) | |
if not use_session and sess is not None: | |
await sess.close() | |
return Map({ 'error': req.status, 'reason': req.reason, 'url': req.url }) | |
else: | |
json = await req.json() | |
if isinstance(json, list): | |
res = json | |
elif json is None: | |
res = {} | |
else: | |
res = Map(json) | |
log.debug({ 'request': req.status, 'url': req.url, 'reason': req.reason }) | |
return res | |
def resultsync(req: requests.Response): | |
if req.status_code != 200: | |
if not quiet: | |
log.error({ 'request error': req.status_code, 'reason': req.reason, 'url': req.url }) | |
return Map({ 'error': req.status_code, 'reason': req.reason, 'url': req.url }) | |
else: | |
json = req.json() | |
if isinstance(json, list): | |
res = json | |
elif json is None: | |
res = {} | |
else: | |
res = Map(json) | |
log.debug({ 'request': req.status_code, 'url': req.url, 'reason': req.reason }) | |
return res | |
async def get(endpoint: str, json: dict = None): | |
global sess # pylint: disable=global-statement | |
sess = sess if sess is not None else await session() | |
try: | |
async with sess.get(url=endpoint, json=json, verify_ssl=False) as req: | |
res = await result(req) | |
return res | |
except Exception as err: | |
log.error({ 'session': err }) | |
return {} | |
def getsync(endpoint: str, json: dict = None): | |
try: | |
req = requests.get(f'{sd_url}{endpoint}', json=json, verify=False, auth=authsync()) # pylint: disable=missing-timeout | |
res = resultsync(req) | |
return res | |
except Exception as err: | |
log.error({ 'session': err }) | |
return {} | |
async def post(endpoint: str, json: dict = None): | |
global sess # pylint: disable=global-statement | |
# sess = sess if sess is not None else await session() | |
if sess and not sess.closed: | |
await sess.close() | |
sess = await session() | |
try: | |
async with sess.post(url=endpoint, json=json, verify_ssl=False) as req: | |
res = await result(req) | |
return res | |
except Exception as err: | |
log.error({ 'session': err }) | |
return {} | |
def postsync(endpoint: str, json: dict = None): | |
req = requests.post(f'{sd_url}{endpoint}', json=json, verify=False, auth=authsync()) # pylint: disable=missing-timeout | |
res = resultsync(req) | |
return res | |
async def interrupt(): | |
res = await get('/sdapi/v1/progress?skip_current_image=true') | |
if 'state' in res and res.state.job_count > 0: | |
log.debug({ 'interrupt': res.state }) | |
res = await post('/sdapi/v1/interrupt') | |
await asyncio.sleep(1) | |
return res | |
else: | |
log.debug({ 'interrupt': 'idle' }) | |
return { 'interrupt': 'idle' } | |
def interruptsync(): | |
res = getsync('/sdapi/v1/progress?skip_current_image=true') | |
if 'state' in res and res.state.job_count > 0: | |
log.debug({ 'interrupt': res.state }) | |
res = postsync('/sdapi/v1/interrupt') | |
return res | |
else: | |
log.debug({ 'interrupt': 'idle' }) | |
return { 'interrupt': 'idle' } | |
async def progress(): | |
res = await get('/sdapi/v1/progress?skip_current_image=false') | |
try: | |
if res is not None and res.get('current_image', None) is not None: | |
res.current_image = Image.open(io.BytesIO(base64.b64decode(res['current_image']))) | |
except Exception: | |
pass | |
log.debug({ 'progress': res }) | |
return res | |
def progresssync(): | |
res = getsync('/sdapi/v1/progress?skip_current_image=true') | |
log.debug({ 'progress': res }) | |
return res | |
def get_log(): | |
res = getsync('/sdapi/v1/log') | |
for line in res: | |
log.debug(line) | |
return res | |
def get_info(): | |
import time | |
t0 = time.time() | |
res = getsync('/sdapi/v1/system-info/status?full=true&refresh=true') | |
t1 = time.time() | |
print({ 'duration': 1000 * round(t1-t0, 3), **res }) | |
return res | |
def options(): | |
opts = getsync('/sdapi/v1/options') | |
flags = getsync('/sdapi/v1/cmd-flags') | |
return { 'options': opts, 'flags': flags } | |
def shutdown(): | |
try: | |
postsync('/sdapi/v1/shutdown') | |
except Exception as e: | |
log.info({ 'shutdown': e }) | |
async def session(): | |
global sess # pylint: disable=global-statement | |
time = aiohttp.ClientTimeout(total = None, sock_connect = 10, sock_read = None) # default value is 5 minutes, we need longer for training | |
sess = aiohttp.ClientSession(timeout = time, base_url = sd_url, auth=auth()) | |
log.debug({ 'sdapi': 'session created', 'endpoint': sd_url }) | |
""" | |
sess = await aiohttp.ClientSession(timeout = timeout).__aenter__() | |
try: | |
async with sess.get(url = f'{sd_url}/') as req: | |
log.debug({ 'sdapi': 'session created', 'endpoint': sd_url }) | |
except Exception as e: | |
log.error({ 'sdapi': e }) | |
await asyncio.sleep(0) | |
await sess.__aexit__(None, None, None) | |
sess = None | |
return sess | |
""" | |
return sess | |
async def close(): | |
if sess is not None: | |
await asyncio.sleep(0) | |
await sess.close() | |
await sess.__aexit__(None, None, None) | |
log.debug({ 'sdapi': 'session closed', 'endpoint': sd_url }) | |
if __name__ == "__main__": | |
sys.argv.pop(0) | |
log.setLevel(logging.DEBUG) | |
if 'interrupt' in sys.argv: | |
asyncio.run(interrupt()) | |
elif 'progress' in sys.argv: | |
asyncio.run(progress()) | |
elif 'progresssync' in sys.argv: | |
progresssync() | |
elif 'options' in sys.argv: | |
opt = options() | |
log.debug({ 'options' }) | |
import json | |
print(json.dumps(opt['options'], indent = 2)) | |
log.debug({ 'cmd-flags' }) | |
print(json.dumps(opt['flags'], indent = 2)) | |
elif 'log' in sys.argv: | |
get_log() | |
elif 'info' in sys.argv: | |
get_info() | |
elif 'shutdown' in sys.argv: | |
shutdown() | |
else: | |
res = getsync(sys.argv[0]) | |
print(res) | |
asyncio.run(close(), debug=True) | |
asyncio.run(asyncio.sleep(0.5)) | |