186 lines
6.0 KiB
Python
186 lines
6.0 KiB
Python
import io
|
|
import time
|
|
import PIL
|
|
import base64
|
|
import re
|
|
import random
|
|
import aiohttp
|
|
import asyncio
|
|
import sentry_sdk
|
|
from aiogoogletrans import Translator
|
|
from typing import Any, Dict
|
|
from discord import File, Asset
|
|
from RUNA.Modules.Tagging import Tagging
|
|
from RUNA.Cogs.Event import model
|
|
|
|
translator = Translator()
|
|
BLOCKTAG = [
|
|
"nsfw",
|
|
"nude",
|
|
"nipples",
|
|
"nipple",
|
|
"pussy",
|
|
"public hair",
|
|
"gay",
|
|
"lesbian",
|
|
"corpse",
|
|
"no panties",
|
|
"no panty",
|
|
"no bra",
|
|
"bra",
|
|
"panty",
|
|
"panties",
|
|
"underwear",
|
|
"undergarment",
|
|
"underpants",
|
|
"underpant",
|
|
"blowjob",
|
|
"sex",
|
|
"sexy",
|
|
"pennis"
|
|
"realistic",
|
|
"open breasts",
|
|
"breasts",
|
|
"bikini",
|
|
"swimsuit",
|
|
"give birth",
|
|
"slave"
|
|
]
|
|
|
|
weight = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.7, 0.8, 0.9]
|
|
cfg = [6, 7, 8, 9, 10, 11]
|
|
|
|
|
|
def is_korean(string):
|
|
pattern = re.compile(r"[ㄱ-ㅎㅏ-ㅣ가-힣]")
|
|
match = pattern.search(string)
|
|
return bool(match)
|
|
|
|
|
|
async def process_prompt(positive: str, negative: str, res: list, isnsfw: bool, style: list, after_process: bool, avatar):
|
|
tags = None
|
|
if avatar is not None:
|
|
try:
|
|
if avatar.url.endswith(".gif"):
|
|
avatar = avatar.replace(format="png")
|
|
avatar = await avatar.read()
|
|
tags = (await Tagging().predict(avatar))[0]
|
|
except Exception:
|
|
tags = None
|
|
if is_korean(positive):
|
|
positive = (await translator.translate(positive, dest="en")).text
|
|
if is_korean(negative):
|
|
negative = (await translator.translate(negative, dest="en")).text
|
|
if not isnsfw:
|
|
for tag in BLOCKTAG:
|
|
if tag in positive:
|
|
positive = positive.replace(tag, "")
|
|
default_negative = "(KHFB, AuroraNegative, easynegative, negative_hand-neg, verybadimagenegative_v1.3:0.8), (Worst Quality, Low Quality:1.4), border, skimpy, grayscale, multiple_girls, 3d, realistic, string, multiple hands, chinese, thick abs, chubby abs, lowres, bad anatomy, asymmetric wings, elf, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, large areolae, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, Multiple legs, multiple feet, tie, (necktie:1.5), several hands, three feet, four legs, three legs,animal ears, nsfw, exposure"
|
|
for styles in style:
|
|
positive += f" <lora:{styles["style"]}:{styles["weight"]}>"
|
|
if negative is not None:
|
|
negative = default_negative + "," + negative
|
|
if tags is not None:
|
|
positive += "," + tags
|
|
payload = {
|
|
"prompt": positive,
|
|
"negative": negative,
|
|
"seed": random.randint(0, 1000000000),
|
|
"width": res[0],
|
|
"height": res[1],
|
|
"cfg_scale": random.choice(cfg),
|
|
"steps": random.randint(20, 25),
|
|
"sampler_index": "DPM++ 2M Karras",
|
|
"refiner_checkpoint": "smooREFINERV2R10_half",
|
|
"refiner_switch_at": 0.45,
|
|
"alwayson_scripts": {
|
|
"ADetailer": {
|
|
'args': [
|
|
{
|
|
'ad_model': 'face_yolov8n.pt',
|
|
'ad_inpaint_only_masked': True
|
|
}]
|
|
},
|
|
"sonar": {
|
|
'args': [
|
|
{
|
|
'sampler': 'Euler a',
|
|
'momentum': 0.95,
|
|
'momentum_hist': 0.75
|
|
}
|
|
]
|
|
}
|
|
}
|
|
}
|
|
if after_process:
|
|
additional = {
|
|
"enable_hr": True,
|
|
"hr_scale": 1.3,
|
|
"hr_upscaler": "R-ESRGAN 4x+ Anime6B",
|
|
"denoising_strength": 0.6
|
|
}
|
|
payload.update(additional)
|
|
return payload
|
|
|
|
|
|
async def post_gpu_server(url: str, payload: Dict[str, Any]):
|
|
async with aiohttp.ClientSession() as session:
|
|
try:
|
|
async with session.post(url, json=payload, timeout=300) as response:
|
|
if response.status == 200:
|
|
return {"status": True, "data": await response.json()}
|
|
else:
|
|
return {"status": False, "data": None}
|
|
except Exception as e:
|
|
if isinstance(e, asyncio.TimeoutError):
|
|
return {"status": False, "data": None}
|
|
|
|
|
|
async def image_to_base64(image) -> str:
|
|
img = PIL.Image.open(image)
|
|
img = img.convert("RGB")
|
|
img_byte_arr = io.BytesIO()
|
|
img.save(img_byte_arr, format='PNG')
|
|
img_byte_arr = img_byte_arr.getvalue()
|
|
img_str = base64.b64encode(img_byte_arr).decode("utf-8")
|
|
return img_str
|
|
|
|
|
|
async def base64_to_image(base642) -> File:
|
|
attachment = File(io.BytesIO(base64.b64decode(base642)), filename="image.png")
|
|
return attachment
|
|
|
|
|
|
async def get_gpuserver_status(url) -> Dict:
|
|
async with aiohttp.ClientSession() as session:
|
|
try:
|
|
async with session.get(f"http://172.30.1.49:7860/sdapi/v1/memory", timeout=10) as response:
|
|
|
|
if response.status == 200:
|
|
result = await response.json()
|
|
memstatus = result["ram"]["used"]
|
|
cudamemstatus = result["cuda"]["system"]["used"]
|
|
oomcount = result["cuda"]["events"]["oom"]
|
|
return {"status": "online", "system_memory_usage": bytes_to_gb(memstatus),
|
|
"cuda_memory_usage": bytes_to_gb(cudamemstatus), "oom_count": oomcount}
|
|
except Exception as e:
|
|
sentry_sdk.capture_exception(e)
|
|
return {"status": "offline"}
|
|
|
|
|
|
def bytes_to_gb(bytes: int) -> float:
|
|
return round(bytes / 1024 / 1024 / 1024, 2)
|
|
|
|
|
|
async def Get_Backend_latency():
|
|
start_time = time.time()
|
|
async with aiohttp.ClientSession() as client:
|
|
try:
|
|
async with client.get("http://172.30.1.49:7860/sdapi/v1/memory", timeout=10) as response:
|
|
if response.status_code == 200 or response.status_code == 404:
|
|
return round(time.time() - start_time, 2)
|
|
else:
|
|
return None
|
|
except Exception as e:
|
|
print(e)
|
|
return None |