factor out exception handler

This commit is contained in:
Kevin Gibbons 2022-08-25 15:13:07 -07:00
parent 0eba55ddbc
commit a10baf5808

View File

@ -253,23 +253,36 @@ The vast majority of these arguments default to reasonable values.
scope = autocast if self.precision=="autocast" else nullcontext
tic = time.time()
results = list()
def prompt_callback(image, seed):
results.append([image, seed])
if image_callback is not None:
image_callback(image, seed)
try:
if init_img:
assert os.path.exists(init_img),f'{init_img}: File not found'
results = self._img2img(prompt,
self._img2img(prompt,
data=data,precision_scope=scope,
batch_size=batch_size,iterations=iterations,
steps=steps,seed=seed,cfg_scale=cfg_scale,ddim_eta=ddim_eta,
skip_normalize=skip_normalize,
init_img=init_img,strength=strength,variants=variants,
callback=image_callback)
callback=prompt_callback)
else:
results = self._txt2img(prompt,
self._txt2img(prompt,
data=data,precision_scope=scope,
batch_size=batch_size,iterations=iterations,
steps=steps,seed=seed,cfg_scale=cfg_scale,ddim_eta=ddim_eta,
skip_normalize=skip_normalize,
width=width,height=height,
callback=image_callback)
callback=prompt_callback)
except KeyboardInterrupt:
print('*interrupted*')
print('Partial results will be returned; if --grid was requested, nothing will be returned.')
except RuntimeError as e:
print(str(e))
toc = time.time()
print(f'{len(results)} images generated in',"%4.2fs"% (toc-tic))
return results
@ -292,7 +305,6 @@ The vast majority of these arguments default to reasonable values.
image_count = 0
# Gawd. Too many levels of indent here. Need to refactor into smaller routines!
try:
with precision_scope(self.device.type), self.model.ema_scope():
all_samples = list()
for n in trange(iterations, desc="Sampling"):
@ -340,11 +352,6 @@ The vast majority of these arguments default to reasonable values.
callback(image,seed)
seed = self._new_seed()
except KeyboardInterrupt:
print('*interrupted*')
print('Partial results will be returned; if --grid was requested, nothing will be returned.')
except RuntimeError as e:
print(str(e))
return images
@ -379,7 +386,6 @@ The vast majority of these arguments default to reasonable values.
# print(f"target t_enc is {t_enc} steps")
images = list()
try:
with precision_scope(self.device.type), self.model.ema_scope():
all_samples = list()
for n in trange(iterations, desc="Sampling"):
@ -424,12 +430,6 @@ The vast majority of these arguments default to reasonable values.
callback(image,seed)
seed = self._new_seed()
except KeyboardInterrupt:
print('*interrupted*')
print('Partial results will be returned; if --grid was requested, nothing will be returned.')
except RuntimeError as e:
print("Oops! A runtime error has occurred. If this is unexpected, please copy-and-paste this stack trace and post it as an Issue to http://github.com/lstein/stable-diffusion")
traceback.print_exc()
return images
def _new_seed(self):