15 changed files with 520 additions and 224 deletions
@ -0,0 +1,59 @@ |
|||||
|
import asyncio |
||||
|
import requests |
||||
|
import json |
||||
|
import logging |
||||
|
|
||||
|
logger = logging.getLogger(__name__) |
||||
|
|
||||
|
class KoboldCppTextWrapper(object): |
||||
|
"""Base Class for koboldcpp""" |
||||
|
|
||||
|
def __init__(self, endpoint_name: str, model_name: str): |
||||
|
self.endpoint_name = endpoint_name |
||||
|
self.model_name = model_name |
||||
|
|
||||
|
def setup(): |
||||
|
os.system("mkdir -p repositories && (cd repositories && git clone https://github.com/LostRuins/koboldcpp.git)") |
||||
|
os.system("apt update && apt-get install libopenblas-dev libclblast-dev libmkl-dev") |
||||
|
os.system("(cd repositories/koboldcpp && make LLAMA_OPENBLAS=1 && cd models && wget https://huggingface.co/concedo/pygmalion-6bv3-ggml-ggjt/resolve/main/pygmalion-6b-v3-ggml-ggjt-q4_0.bin)") |
||||
|
#python3 koboldcpp.py models/pygmalion-6b-v3-ggml-ggjt-q4_0.bin |
||||
|
#python3 koboldcpp.py --smartcontext models/pygmalion-6b-v3-ggml-ggjt-q4_0.bin |
||||
|
|
||||
|
async def generate(self, prompt: str, typing_fn, temperature=0.72, max_new_tokens=200, timeout=180): |
||||
|
# Set the API endpoint URL |
||||
|
endpoint = f"http://{self.endpoint_name}/api/latest/generate" |
||||
|
|
||||
|
# Set the headers for the request |
||||
|
headers = { |
||||
|
"Content-Type": "application/json", |
||||
|
} |
||||
|
|
||||
|
# Define your inputs |
||||
|
input_data = { |
||||
|
"prompt": prompt, |
||||
|
"max_context_length": 2048, |
||||
|
"max_length": max_new_tokens, |
||||
|
"temperature": temperature, |
||||
|
"top_k": 50, |
||||
|
"top_p": 0.85, |
||||
|
"rep_pen": 1.08, |
||||
|
"rep_pen_range": 1024, |
||||
|
"stop_sequence": ['<|endoftext|>'], |
||||
|
} |
||||
|
|
||||
|
logger.info(f"sending request to koboldcpp. endpoint=\"{self.endpoint_name}\"") |
||||
|
|
||||
|
TRIES = 30 |
||||
|
for i in range(TRIES): |
||||
|
r = requests.post(endpoint, json=input_data, headers=headers, timeout=timeout) |
||||
|
r_json = r.json() |
||||
|
logger.info(r_json) |
||||
|
if r.status_code == 200: |
||||
|
output = r_json["results"][0]["text"] |
||||
|
return output |
||||
|
elif r.status_code == 503: |
||||
|
logger.info(f"api is busy. waiting...") |
||||
|
asyncio.sleep(5) |
||||
|
|
||||
|
raise ValueError(f"<ERROR> TIMEOUT / NO OUTOUT") |
||||
|
|
||||
@ -1,42 +0,0 @@ |
|||||
import asyncio |
|
||||
import requests |
|
||||
import os, tempfile |
|
||||
from .runpod import RunpodWrapper |
|
||||
import logging |
|
||||
|
|
||||
logger = logging.getLogger(__name__) |
|
||||
|
|
||||
|
|
||||
class RunpodImageWrapper(RunpodWrapper): |
|
||||
async def download_image(self, url, path): |
|
||||
r = requests.get(url, stream=True) |
|
||||
if r.status_code == 200: |
|
||||
with open(path, 'wb') as f: |
|
||||
for chunk in r: |
|
||||
f.write(chunk) |
|
||||
|
|
||||
async def generate(self, input_prompt: str, negative_prompt: str, endpoint_name: str, typing_fn, timeout=180): |
|
||||
|
|
||||
# Define your inputs |
|
||||
input_data = { |
|
||||
"input": { |
|
||||
"prompt": input_prompt, |
|
||||
"negative_prompt": negative_prompt, |
|
||||
"width": 512, |
|
||||
"height": 768, |
|
||||
"num_outputs": 3, |
|
||||
# "nsfw": True |
|
||||
}, |
|
||||
} |
|
||||
|
|
||||
output = await super().generate(input_data, endpoint_name, typing_fn, timeout) |
|
||||
|
|
||||
os.makedirs("./.data/images", exist_ok=True) |
|
||||
files = [] |
|
||||
for image in output: |
|
||||
temp_name = next(tempfile._get_candidate_names()) |
|
||||
filename = "./.data/images/" + temp_name + ".jpg" |
|
||||
await self.download_image(image["image"], filename) |
|
||||
files.append(filename) |
|
||||
|
|
||||
return files |
|
||||
@ -1,82 +0,0 @@ |
|||||
import asyncio |
|
||||
import requests |
|
||||
import os, tempfile |
|
||||
from .runpod import RunpodWrapper |
|
||||
|
|
||||
import io |
|
||||
import base64 |
|
||||
from PIL import Image, PngImagePlugin |
|
||||
|
|
||||
import logging |
|
||||
|
|
||||
logger = logging.getLogger(__name__) |
|
||||
|
|
||||
|
|
||||
class RunpodImageAutomaticWrapper(RunpodWrapper): |
|
||||
|
|
||||
async def generate(self, input_prompt: str, negative_prompt: str, endpoint_name: str, typing_fn, timeout=180): |
|
||||
|
|
||||
# Define your inputs |
|
||||
input_data = { |
|
||||
"input": { |
|
||||
"prompt": input_prompt, |
|
||||
"nagative_prompt": negative_prompt, |
|
||||
"steps": 25, |
|
||||
"cfg_scale": 7, |
|
||||
"seed": -1, |
|
||||
"width": 512, |
|
||||
"height": 768, |
|
||||
"batch_size": 3, |
|
||||
# "sampler_index": "DPM++ 2M Karras", |
|
||||
# "enable_hr": True, |
|
||||
# "hr_scale": 2, |
|
||||
# "hr_upscaler": "ESRGAN_4x", # "Latent" |
|
||||
# "denoising_strength": 0.5, |
|
||||
# "hr_second_pass_steps": 15, |
|
||||
"restore_faces": True, |
|
||||
# "gfpgan_visibility": 0.5, |
|
||||
# "codeformer_visibility": 0.5, |
|
||||
# "codeformer_weight": 0.5, |
|
||||
## "override_settings": { |
|
||||
## "filter_nsfw": False, |
|
||||
## }, |
|
||||
"api_endpoint": "txt2img", |
|
||||
}, |
|
||||
"cmd": "txt2img" |
|
||||
} |
|
||||
|
|
||||
output = await super().generate(input_data, endpoint_name, typing_fn, timeout) |
|
||||
|
|
||||
upscale = False |
|
||||
if upscale: |
|
||||
count = 0 |
|
||||
for i in output['images']: |
|
||||
payload = { |
|
||||
"init_images": [i], |
|
||||
"prompt": input_prompt, |
|
||||
"nagative_prompt": negative_prompt, |
|
||||
"steps": 20, |
|
||||
"seed": -1, |
|
||||
#"sampler_index": "Euler", |
|
||||
# tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding, upscaler_index, save_upscaled_image, redraw_mode, save_seams_fix_image, seams_fix_mask_blur, seams_fix_type, target_size_type, custom_width, custom_height, custom_scale |
|
||||
# "script_args": ["",512,0,8,32,64,0.275,32,3,False,0,True,8,3,2,1080,1440,1.875], |
|
||||
# "script_name": "Ultimate SD upscale", |
|
||||
} |
|
||||
upscaled_output = await serverless_automatic_request(payload, "img2img", api_url, api_key, typing_fn) |
|
||||
output['images'][count] = upscaled_output['images'][count] |
|
||||
|
|
||||
|
|
||||
os.makedirs("./.data/images", exist_ok=True) |
|
||||
files = [] |
|
||||
for i in output['images']: |
|
||||
temp_name = next(tempfile._get_candidate_names()) |
|
||||
filename = "./.data/images/" + temp_name + ".png" |
|
||||
image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0]))) |
|
||||
info = output['info'] |
|
||||
parameters = output['parameters'] |
|
||||
pnginfo = PngImagePlugin.PngInfo() |
|
||||
pnginfo.add_text("parameters", info) |
|
||||
image.save(filename, pnginfo=pnginfo) |
|
||||
files.append(filename) |
|
||||
|
|
||||
return files |
|
||||
@ -1,31 +0,0 @@ |
|||||
import asyncio |
|
||||
import json |
|
||||
from .runpod import RunpodWrapper |
|
||||
import logging |
|
||||
|
|
||||
logger = logging.getLogger(__name__) |
|
||||
|
|
||||
|
|
||||
class RunpodTextWrapper(RunpodWrapper): |
|
||||
|
|
||||
def __init__(self, api_key, endpoint): |
|
||||
self.api_key = api_key |
|
||||
self.endpoint = endpoint |
|
||||
|
|
||||
async def generate(self, prompt, endpoint_name, typing_fn, temperature=0.72, max_new_tokens=200, timeout=180): |
|
||||
|
|
||||
# Define your inputs |
|
||||
input_data = { |
|
||||
"input": { |
|
||||
"prompt": prompt, |
|
||||
"max_length": min(max_new_tokens, 2048), |
|
||||
"temperature": bot.temperature, |
|
||||
"do_sample": True, |
|
||||
} |
|
||||
} |
|
||||
output = await super().generate(input_data, endpoint_name, api_key, typing_fn, timeout) |
|
||||
output = output.removeprefix(prompt) |
|
||||
return(output) |
|
||||
|
|
||||
async def generate2(self, prompt, typing_fn, temperature=0.72, max_new_tokens=200, timeout=180): |
|
||||
generate(prompt, self.endpoint, typing_fn, temperature, nax_new_tokens, timeout) |
|
||||
@ -0,0 +1,148 @@ |
|||||
|
import asyncio |
||||
|
import requests |
||||
|
import json |
||||
|
import os, tempfile |
||||
|
import logging |
||||
|
|
||||
|
logger = logging.getLogger(__name__) |
||||
|
|
||||
|
class StableHordeWrapper(object): |
||||
|
"""Base Class for stablehorde""" |
||||
|
|
||||
|
def __init__(self, api_key: str, endpoint_name: str, model_name: str): |
||||
|
self.api_key = api_key |
||||
|
self.endpoint_name = endpoint_name |
||||
|
self.model_name = model_name |
||||
|
|
||||
|
async def generate(self, input_data: str, typing_fn, timeout=180): |
||||
|
# Set the API endpoint URL |
||||
|
endpoint = "https://stablehorde.net/api/v2/generate/async" |
||||
|
#endpoint = "https://koboldai.net/api/v2/generate/async" |
||||
|
|
||||
|
# Set the headers for the request |
||||
|
headers = { |
||||
|
"Content-Type": "application/json", |
||||
|
"accept": "application/json", |
||||
|
"apikey": f"{self.api_key}" |
||||
|
} |
||||
|
|
||||
|
logger.info(f"sending request to stablehorde.net. endpoint=\"{self.endpoint_name}\"") |
||||
|
|
||||
|
# Make the request |
||||
|
try: |
||||
|
r = requests.post(endpoint, json=input_data, headers=headers, timeout=timeout) |
||||
|
except requests.exceptions.RequestException as e: |
||||
|
raise ValueError(f"<HTTP ERROR>") |
||||
|
r_json = r.json() |
||||
|
logger.debug(r_json) |
||||
|
|
||||
|
if r.status_code == 202: |
||||
|
#status = r_json["message"] |
||||
|
job_id = r_json["id"] |
||||
|
TIMEOUT = 360 |
||||
|
DELAY = 11 |
||||
|
output = None |
||||
|
for i in range(TIMEOUT//DELAY): |
||||
|
endpoint = f"https://stablehorde.net/api/v2/generate/status/{job_id}" |
||||
|
#endpoint = f"https://koboldai.net/api/v2/generate/text/status/{job_id}" |
||||
|
r = requests.get(endpoint, headers=headers) |
||||
|
r_json = r.json() |
||||
|
logger.info(r_json) |
||||
|
#status = r_json["message"] |
||||
|
if "done" not in r_json: |
||||
|
raise ValueError("<ERROR>") |
||||
|
if "faulted" in r_json and r_json["faulted"] == True: |
||||
|
raise ValueError("<ERROR> Faulted") |
||||
|
if r_json["done"] == True: |
||||
|
output = r_json["generations"] |
||||
|
break |
||||
|
else: |
||||
|
if "processing" in r_json and r_json["processing"] == 1: |
||||
|
await typing_fn() |
||||
|
elif "wait_time" in r_json and r_json["wait_time"] < 20 and r_json["wait_time"] != 0 and r_json["queue_position"] < 100: |
||||
|
await typing_fn() |
||||
|
await asyncio.sleep(DELAY) |
||||
|
else: |
||||
|
raise ValueError(f"<ERROR> HTTP code {r.status_code}") |
||||
|
|
||||
|
if not output: |
||||
|
raise ValueError(f"<ERROR> TIMEOUT / NO OUTOUT") |
||||
|
return output |
||||
|
|
||||
|
|
||||
|
class StableHordeTextWrapper(StableHordeWrapper): |
||||
|
|
||||
|
async def generate(self, prompt, typing_fn, temperature=0.72, max_new_tokens=200, timeout=180): |
||||
|
|
||||
|
# Define your inputs |
||||
|
input_data = { |
||||
|
"prompt": prompt, |
||||
|
"params": { |
||||
|
"n": 1, |
||||
|
# "frmtadsnsp": False, |
||||
|
# "frmtrmblln": False, |
||||
|
# "frmtrmspch": False, |
||||
|
# "frmttriminc": False, |
||||
|
"max_context_length": 1024, |
||||
|
"max_length": 512, |
||||
|
"rep_pen": 1.1, |
||||
|
"rep_pen_range": 1024, |
||||
|
"rep_pen_slope": 0.7, |
||||
|
# "singleline": False, |
||||
|
# "soft_prompt": "", |
||||
|
"temperature": 0.75, |
||||
|
"tfs": 1.0, |
||||
|
"top_a": 0.0, |
||||
|
"top_k": 0, |
||||
|
"top_p": 0.9, |
||||
|
"typical": 1.0, |
||||
|
# "sampler_order": [0], |
||||
|
}, |
||||
|
"softprompts": [], |
||||
|
"trusted_workers": False, |
||||
|
"nsfw": True, |
||||
|
# "workers": [], |
||||
|
"models": [f"{self.endpoint_name}"] |
||||
|
} |
||||
|
|
||||
|
output = await super().generate(input_data, api_key, typing_fn, timeout) |
||||
|
output = output[0]["text"].removeprefix(prompt) |
||||
|
return(output) |
||||
|
|
||||
|
|
||||
|
class StableHordeImageWrapper(StableHordeWrapper): |
||||
|
|
||||
|
async def download_image(self, url, path): |
||||
|
r = requests.get(url, stream=True) |
||||
|
if r.status_code == 200: |
||||
|
with open(path, 'wb') as f: |
||||
|
for chunk in r: |
||||
|
f.write(chunk) |
||||
|
|
||||
|
async def generate(self, input_prompt: str, negative_prompt: str, typing_fn, timeout=180): |
||||
|
|
||||
|
# Define your inputs |
||||
|
input_data = { |
||||
|
"prompt": input_prompt, |
||||
|
"params": { |
||||
|
# "negative_prompt": negative_prompt, |
||||
|
"width": 512, |
||||
|
"height": 512, |
||||
|
}, |
||||
|
"nsfw": True, |
||||
|
"trusted_workers": False, |
||||
|
# "workers": [], |
||||
|
"models": [f"{self.endpoint_name}"] |
||||
|
} |
||||
|
|
||||
|
output = await super().generate(input_data, typing_fn, timeout) |
||||
|
|
||||
|
os.makedirs("./.data/images", exist_ok=True) |
||||
|
files = [] |
||||
|
for image in output: |
||||
|
temp_name = next(tempfile._get_candidate_names()) |
||||
|
filename = "./.data/images/" + temp_name + ".jpg" |
||||
|
await self.download_image(image["img"], filename) |
||||
|
files.append(filename) |
||||
|
|
||||
|
return files |
||||
Loading…
Reference in new issue