You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
83 lines
3.0 KiB
83 lines
3.0 KiB
2 years ago
|
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
|