extern/easydiffusion
1
0
Fork 1
mirror of https://github.com/easydiffusion/easydiffusion.git synced 2025-02-16 18:32:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #61 from cmdr2/develop

Browse Source 
Shorter filenames for saved images; Don't crash other images if one fails to save
This commit is contained in:
cmdr2 2022-09-05 16:53:35 +05:30 committed by GitHub
commit 15da928655
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 17 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
ui/sd_internal/runtime.py
View File

@ -1,4 +1,5 @@
import os, re import os, re
import traceback
import torch import torch
import numpy as np import numpy as np
from omegaconf import OmegaConf from omegaconf import OmegaConf
@ -28,7 +29,7 @@ import base64
from io import BytesIO from io import BytesIO
# local # local
session_id = str(uuid.uuid4()) session_id = str(uuid.uuid4())[-8:]
ckpt = None ckpt = None
model = None model = None
@ -188,7 +189,7 @@ def mk_img(req: Request):
print(f"target t_enc is {t_enc} steps") print(f"target t_enc is {t_enc} steps")
if opt_save_to_disk_path is not None: if opt_save_to_disk_path is not None:
session_out_path = os.path.join(opt_save_to_disk_path, 'session-' + session_id) session_out_path = os.path.join(opt_save_to_disk_path, session_id)
os.makedirs(session_out_path, exist_ok=True) os.makedirs(session_out_path, exist_ok=True)
else: else:
session_out_path = None session_out_path = None
@ -198,9 +199,6 @@ def mk_img(req: Request):
for n in trange(opt_n_iter, desc="Sampling"): for n in trange(opt_n_iter, desc="Sampling"):
for prompts in tqdm(data, desc="data"): for prompts in tqdm(data, desc="data"):
if opt_save_to_disk_path is not None:
base_count = len(os.listdir(session_out_path))
with precision_scope("cuda"): with precision_scope("cuda"):
modelCS.to(device) modelCS.to(device)
uc = None uc = None
@ -242,19 +240,23 @@ def mk_img(req: Request):
res.images.append(ResponseImage(data=img_data, seed=opt_seed)) res.images.append(ResponseImage(data=img_data, seed=opt_seed))
if opt_save_to_disk_path is not None: if opt_save_to_disk_path is not None:
prompt_flattened = "_".join(re.split(":| ", prompts[0])) try:
prompt_flattened = prompt_flattened[:150] prompt_flattened = "_".join(re.split(":| ", prompts[0]))
prompt_flattened = prompt_flattened.replace(',', '')
prompt_flattened = prompt_flattened[:50]
file_path = f"sd_{prompt_flattened}_Seed-{opt_seed}_Steps-{opt_ddim_steps}_Guidance-{opt_scale}_{base_count:05}" img_id = str(uuid.uuid4())[-8:]
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}" file_path = f"{prompt_flattened}_{img_id}"
img.save(img_out_path) img_out_path = os.path.join(session_out_path, f"{file_path}.{opt_format}")
with open(meta_out_path, 'w') as f: meta_out_path = os.path.join(session_out_path, f"{file_path}.txt")
f.write(metadata)
base_count += 1 metadata = f"{prompts[0]}\nWidth: {opt_W}\nHeight: {opt_H}\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)
except:
print('could not save the file', traceback.format_exc())
seeds += str(opt_seed) + "," seeds += str(opt_seed) + ","
opt_seed += 1 opt_seed += 1