50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
# Copyright 2023 Qewertyy, MIT License
|
|
import logging
|
|
import asyncio
|
|
from lexica import AsyncClient, Client
|
|
|
|
def ImageModels():
|
|
models = Client().models['models']['image']
|
|
dict_models = {}
|
|
for model in models:
|
|
model_id = model['id']
|
|
model_name = model['name']
|
|
dict_models[model_name] = model_id
|
|
return dict_models
|
|
|
|
async def ImageGeneration(model,prompt):
|
|
try:
|
|
output = await AsyncClient().generate(model, prompt, "")
|
|
if output['code'] != 1:
|
|
return 2
|
|
if output['code'] == 69:
|
|
return output['code']
|
|
task_id, request_id = output['task_id'], output['request_id']
|
|
await asyncio.sleep(20)
|
|
tries = 0
|
|
image_url = None
|
|
resp = await AsyncClient().getImages(task_id, request_id)
|
|
while True:
|
|
if resp['code'] == 2:
|
|
image_url = resp['img_urls']
|
|
break
|
|
if tries > 15:
|
|
break
|
|
await asyncio.sleep(5)
|
|
resp = await AsyncClient().getImages(task_id, request_id)
|
|
tries += 1
|
|
continue
|
|
return image_url
|
|
except Exception as e:
|
|
logging.warning(f"Failed to generate the image:", e)
|
|
finally:
|
|
await AsyncClient().close()
|
|
|
|
async def UpscaleImages(image: bytes) -> str:
|
|
content = await AsyncClient().upscale(image)
|
|
await AsyncClient().close()
|
|
upscaled_file_path = "upscaled.png"
|
|
with open(upscaled_file_path, "wb") as output_file:
|
|
output_file.write(content)
|
|
return upscaled_file_path
|