Return the image metadata and disk path in the response

This commit is contained in:
cmdr2 2022-09-14 10:15:35 +05:30
parent 0c0c8e503e
commit 885759abc5
2 changed files with 30 additions and 2 deletions

View File

@ -21,6 +21,20 @@ class Request:
use_upscale: str = None # or "RealESRGAN_x4plus" or "RealESRGAN_x4plus_anime_6B"
show_only_filtered_image: bool = False
def json(self):
return {
"prompt": self.prompt,
"num_outputs": self.num_outputs,
"num_inference_steps": self.num_inference_steps,
"guidance_scale": self.guidance_scale,
"width": self.width,
"height": self.height,
"seed": self.seed,
"prompt_strength": self.prompt_strength,
"use_face_correction": self.use_face_correction,
"use_upscale": self.use_upscale,
}
def to_string(self):
return f'''
prompt: {self.prompt}
@ -42,6 +56,7 @@ class Image:
data: str # base64
seed: int
is_nsfw: bool
path_abs: str = None
def __init__(self, data, seed):
self.data = data
@ -51,14 +66,19 @@ class Image:
return {
"data": self.data,
"seed": self.seed,
"path_abs": self.path_abs,
}
class Response:
request: Request
session_id: str
images: list
def json(self):
res = {
"status": 'succeeded',
"session_id": self.session_id,
"request": self.request.json(),
"output": [],
}

View File

@ -191,6 +191,8 @@ def mk_img(req: Request):
stop_processing = False
res = Response()
res.session_id = session_id
res.request = req
res.images = []
model.turbo = req.turbo
@ -373,7 +375,11 @@ def mk_img(req: Request):
if not opt_show_only_filtered:
img_data = img_to_base64_str(img)
res.images.append(ResponseImage(data=img_data, seed=opt_seed))
res_image_orig = ResponseImage(data=img_data, seed=opt_seed)
res.images.append(res_image_orig)
if opt_save_to_disk_path is not None:
res_image_orig.path_abs = img_out_path
if (opt_use_face_correction is not None and opt_use_face_correction.startswith('GFPGAN')) or \
(opt_use_upscale is not None and opt_use_upscale.startswith('RealESRGAN')):
@ -394,13 +400,15 @@ def mk_img(req: Request):
filtered_image = Image.fromarray(x_sample)
filtered_img_data = img_to_base64_str(filtered_image)
res.images.append(ResponseImage(data=filtered_img_data, seed=opt_seed))
res_image_filtered = ResponseImage(data=filtered_img_data, seed=opt_seed)
res.images.append(res_image_filtered)
filters_applied = "_".join(filters_applied)
if opt_save_to_disk_path is not None:
filtered_img_out_path = os.path.join(session_out_path, f"{file_path}_{filters_applied}.{opt_format}")
save_image(filtered_image, filtered_img_out_path)
res_image_filtered.path_abs = filtered_img_out_path
seeds += str(opt_seed) + ","
opt_seed += 1