extern/easydiffusion
1
0
Fork 1
mirror of https://github.com/easydiffusion/easydiffusion.git synced 2024-11-25 01:34:38 +01:00
Code Issues Packages Projects Releases Wiki Activity

Store images during a session in the same folder; Store the metadata for each image as a txt file next to it

Browse Source
This commit is contained in:
cmdr2 2022-09-04 00:12:48 +05:30
parent e5dc932717
commit 21dc2ece1b
1 changed files with 24 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
ui/sd_internal/runtime.py
View File

@ -15,6 +15,8 @@ from ldm.util import instantiate_from_config
from optimizedSD.optimUtils import split_weighted_subprompts
from transformers import logging
import uuid
logging.set_verbosity_error()
# consts
@ -26,6 +28,8 @@ import base64
from io import BytesIO
# local
session_id = str(uuid.uuid4())
ckpt = None
model = None
modelCS = None
@ -140,7 +144,7 @@ def mk_img(req: Request):
opt_init_img = req.init_image
opt_format = 'png'
print(req.to_string(), 'device', device)
print(req.to_string(), '\n device', device)
seed_everything(opt_seed)
@ -183,15 +187,19 @@ def mk_img(req: Request):
t_enc = int(opt_strength * opt_ddim_steps)
print(f"target t_enc is {t_enc} steps")
if opt_save_to_disk_path is not None:
session_out_path = os.path.join(opt_save_to_disk_path, 'session-' + session_id)
os.makedirs(session_out_path, exist_ok=True)
else:
session_out_path = None
seeds = ""
with torch.no_grad():
for n in trange(opt_n_iter, desc="Sampling"):
for prompts in tqdm(data, desc="data"):
if opt_save_to_disk_path is not None:
sample_path = os.path.join(opt_save_to_disk_path, "_".join(re.split(":| ", prompts[0])))[:150]
os.makedirs(sample_path, exist_ok=True)
base_count = len(os.listdir(sample_path))
base_count = len(os.listdir(session_out_path))
with precision_scope("cuda"):
modelCS.to(device)
@ -234,9 +242,18 @@ def mk_img(req: Request):
res.images.append(ResponseImage(data=img_data, seed=opt_seed))
if opt_save_to_disk_path is not None:
img.save(
os.path.join(sample_path, "seed_" + str(opt_seed) + "_" + f"{base_count:05}.{opt_format}")
)
prompt_flattened = "_".join(re.split(":| ", prompts[0]))
prompt_flattened = prompt_flattened[:150]
file_path = f"sd_{prompt_flattened}_Seed-{opt_seed}_Steps-{opt_ddim_steps}_Guidance-{opt_scale}_{base_count:05}"
img_out_path = os.path.join(session_out_path, f"{file_path}.{opt_format}")
meta_out_path = os.path.join(session_out_path, f"{file_path}.txt")
metadata = f"{prompts[0]}\nSeed: {opt_seed}\nSteps: {opt_ddim_steps}\nGuidance Scale: {opt_scale}"
img.save(img_out_path)
with open(meta_out_path, 'w') as f:
f.write(metadata)
base_count += 1
seeds += str(opt_seed) + ","