refactor: add EarlyReturnException to prevent return value issues

This commit is contained in:
Manuel Schmid 2024-06-08 20:25:50 +02:00
parent e58d913fbb
commit c1f62361ba
No known key found for this signature in database
GPG Key ID: 32C4F7569B40B84B
1 changed files with 42 additions and 43 deletions

View File

@ -109,6 +109,10 @@ class AsyncTask:
async_tasks = [] async_tasks = []
class EarlyReturnException:
pass
def worker(): def worker():
global async_tasks global async_tasks
@ -394,19 +398,20 @@ def worker():
progressbar(async_task, 7, 'Image processing ...') progressbar(async_task, 7, 'Image processing ...')
if 'vary' in goals: if 'vary' in goals:
height, initial_latent, width = apply_vary(async_task, denoising_strength, height, initial_latent, switch, width) height, initial_latent, width = apply_vary(async_task, denoising_strength, switch)
if 'upscale' in goals: if 'upscale' in goals:
denoising_strength, height, initial_latent, tiled, width, direct_return = apply_upscale(async_task, switch) try:
if direct_return: denoising_strength, height, initial_latent, tiled, width = apply_upscale(async_task, switch)
except EarlyReturnException:
return return
if 'inpaint' in goals: if 'inpaint' in goals:
denoising_strength, height, initial_latent, width = apply_inpaint(async_task, initial_latent, try:
inpaint_head_model_path, inpaint_image, denoising_strength, initial_latent, height, width = apply_inpaint(async_task, initial_latent,
inpaint_mask, inpaint_parameterized, inpaint_head_model_path, inpaint_image,
switch) inpaint_mask, inpaint_parameterized,
if async_task.debugging_inpaint_preprocessor: switch)
except EarlyReturnException:
return return
if 'cn' in goals: if 'cn' in goals:
@ -471,7 +476,6 @@ def worker():
processing_time = time.perf_counter() - processing_start_time processing_time = time.perf_counter() - processing_start_time
print(f'Processing time (total): {processing_time:.2f} seconds') print(f'Processing time (total): {processing_time:.2f} seconds')
return
def process_task(all_steps, async_task, callback, controlnet_canny_path, controlnet_cpds_path, current_task_id, def process_task(all_steps, async_task, callback, controlnet_canny_path, controlnet_cpds_path, current_task_id,
denoising_strength, final_scheduler_name, goals, initial_latent, switch, task, tasks, denoising_strength, final_scheduler_name, goals, initial_latent, switch, task, tasks,
@ -521,6 +525,7 @@ def worker():
do_not_show_finished_images=len(tasks) == 1 or async_task.disable_intermediate_results) do_not_show_finished_images=len(tasks) == 1 or async_task.disable_intermediate_results)
return imgs return imgs
def apply_patch_settings(async_task): def apply_patch_settings(async_task):
patch_settings[pid] = PatchSettings( patch_settings[pid] = PatchSettings(
async_task.sharpness, async_task.sharpness,
@ -644,22 +649,22 @@ def worker():
if len(all_ip_tasks) > 0: if len(all_ip_tasks) > 0:
pipeline.final_unet = ip_adapter.patch_model(pipeline.final_unet, all_ip_tasks) pipeline.final_unet = ip_adapter.patch_model(pipeline.final_unet, all_ip_tasks)
def apply_vary(async_task, denoising_strength, height, initial_latent, switch, width): def apply_vary(async_task, uov_input_image, denoising_strength, switch):
if 'subtle' in async_task.uov_method: if 'subtle' in async_task.uov_method:
async_task.denoising_strength = 0.5 async_task.denoising_strength = 0.5
if 'strong' in async_task.uov_method: if 'strong' in async_task.uov_method:
async_task.denoising_strength = 0.85 async_task.denoising_strength = 0.85
if async_task.overwrite_vary_strength > 0: if async_task.overwrite_vary_strength > 0:
async_task.denoising_strength = async_task.overwrite_vary_strength async_task.denoising_strength = async_task.overwrite_vary_strength
shape_ceil = get_image_shape_ceil(async_task.uov_input_image) shape_ceil = get_image_shape_ceil(uov_input_image)
if shape_ceil < 1024: if shape_ceil < 1024:
print(f'[Vary] Image is resized because it is too small.') print(f'[Vary] Image is resized because it is too small.')
shape_ceil = 1024 shape_ceil = 1024
elif shape_ceil > 2048: elif shape_ceil > 2048:
print(f'[Vary] Image is resized because it is too big.') print(f'[Vary] Image is resized because it is too big.')
shape_ceil = 2048 shape_ceil = 2048
async_task.uov_input_image = set_image_shape_ceil(async_task.uov_input_image, shape_ceil) uov_input_image = set_image_shape_ceil(uov_input_image, shape_ceil)
initial_pixels = core.numpy_to_pytorch(async_task.uov_input_image) initial_pixels = core.numpy_to_pytorch(uov_input_image)
progressbar(async_task, 8, 'VAE encoding ...') progressbar(async_task, 8, 'VAE encoding ...')
candidate_vae, _ = pipeline.get_candidate_vae( candidate_vae, _ = pipeline.get_candidate_vae(
steps=async_task.steps, steps=async_task.steps,
@ -671,8 +676,8 @@ def worker():
B, C, H, W = initial_latent['samples'].shape B, C, H, W = initial_latent['samples'].shape
width = W * 8 width = W * 8
height = H * 8 height = H * 8
print(f'Final resolution is {str((height, width))}.') print(f'Final resolution is {str((width, height))}.')
return height, initial_latent, width return initial_latent, width, height
def apply_inpaint(async_task, initial_latent, inpaint_head_model_path, inpaint_image, def apply_inpaint(async_task, initial_latent, inpaint_head_model_path, inpaint_image,
inpaint_mask, inpaint_parameterized, switch): inpaint_mask, inpaint_parameterized, switch):
@ -699,8 +704,8 @@ def worker():
inpaint_image = np.ascontiguousarray(inpaint_image.copy()) inpaint_image = np.ascontiguousarray(inpaint_image.copy())
inpaint_mask = np.ascontiguousarray(inpaint_mask.copy()) inpaint_mask = np.ascontiguousarray(inpaint_mask.copy())
inpaint_strength = 1.0 async_task.inpaint_strength = 1.0
inpaint_respective_field = 1.0 async_task.inpaint_respective_field = 1.0
denoising_strength = async_task.inpaint_strength denoising_strength = async_task.inpaint_strength
inpaint_worker.current_task = inpaint_worker.InpaintWorker( inpaint_worker.current_task = inpaint_worker.InpaintWorker(
image=inpaint_image, image=inpaint_image,
@ -711,7 +716,7 @@ def worker():
if async_task.debugging_inpaint_preprocessor: if async_task.debugging_inpaint_preprocessor:
yield_result(async_task, inpaint_worker.current_task.visualize_mask_processing(), async_task.black_out_nsfw, yield_result(async_task, inpaint_worker.current_task.visualize_mask_processing(), async_task.black_out_nsfw,
do_not_show_finished_images=True) do_not_show_finished_images=True)
return raise EarlyReturnException
progressbar(async_task, 11, 'VAE Inpaint encoding ...') progressbar(async_task, 11, 'VAE Inpaint encoding ...')
inpaint_pixel_fill = core.numpy_to_pytorch(inpaint_worker.current_task.interested_fill) inpaint_pixel_fill = core.numpy_to_pytorch(inpaint_worker.current_task.interested_fill)
@ -751,8 +756,9 @@ def worker():
B, C, H, W = latent_fill.shape B, C, H, W = latent_fill.shape
height, width = H * 8, W * 8 height, width = H * 8, W * 8
final_height, final_width = inpaint_worker.current_task.image.shape[:2] final_height, final_width = inpaint_worker.current_task.image.shape[:2]
print(f'Final resolution is {str((final_height, final_width))}, latent is {str((height, width))}.') print(f'Final resolution is {str((final_height, final_width))}, latent is {str((width, height))}.')
return denoising_strength, height, initial_latent, width
return denoising_strength, initial_latent, width, height
def apply_upscale(async_task, switch): def apply_upscale(async_task, switch):
H, W, C = async_task.uov_input_image.shape H, W, C = async_task.uov_input_image.shape
@ -791,7 +797,7 @@ def worker():
uov_input_image_path = log(async_task.uov_input_image, d, output_format=async_task.output_format) uov_input_image_path = log(async_task.uov_input_image, d, output_format=async_task.output_format)
yield_result(async_task, uov_input_image_path, async_task.black_out_nsfw, False, yield_result(async_task, uov_input_image_path, async_task.black_out_nsfw, False,
do_not_show_finished_images=True) do_not_show_finished_images=True)
return direct_return raise EarlyReturnException
tiled = True tiled = True
denoising_strength = 0.382 denoising_strength = 0.382
@ -811,8 +817,8 @@ def worker():
B, C, H, W = initial_latent['samples'].shape B, C, H, W = initial_latent['samples'].shape
width = W * 8 width = W * 8
height = H * 8 height = H * 8
print(f'Final resolution is {str((height, width))}.') print(f'Final resolution is {str((width, height))}.')
return denoising_strength, height, initial_latent, tiled, width, direct_return return denoising_strength, height, initial_latent, tiled, width
def apply_overrides(async_task, height, width): def apply_overrides(async_task, height, width):
if async_task.overwrite_step > 0: if async_task.overwrite_step > 0:
@ -934,35 +940,28 @@ def worker():
async_task.freeu_s2 async_task.freeu_s2
) )
def patch_discrete(unet, scheduler_name):
return core.opModelSamplingDiscrete.patch(unet, scheduler_name, False)[0]
def patch_edm(unet, scheduler_name):
return core.opModelSamplingContinuousEDM.patch(unet, scheduler_name, 120.0, 0.002)[0]
def patch_samplers(async_task): def patch_samplers(async_task):
final_scheduler_name = async_task.scheduler_name final_scheduler_name = async_task.scheduler_name
if async_task.scheduler_name in ['lcm', 'tcd']: if async_task.scheduler_name in ['lcm', 'tcd']:
final_scheduler_name = 'sgm_uniform' final_scheduler_name = 'sgm_uniform'
def patch_discrete(unet):
return core.opModelSamplingDiscrete.patch(
pipeline.final_unet,
sampling=async_task.scheduler_name,
zsnr=False)[0]
if pipeline.final_unet is not None: if pipeline.final_unet is not None:
pipeline.final_unet = patch_discrete(pipeline.final_unet) pipeline.final_unet = patch_discrete(pipeline.final_unet, async_task.scheduler_name)
if pipeline.final_refiner_unet is not None: if pipeline.final_refiner_unet is not None:
pipeline.final_refiner_unet = patch_discrete(pipeline.final_refiner_unet) pipeline.final_refiner_unet = patch_discrete(pipeline.final_refiner_unet, async_task.scheduler_name)
elif async_task.scheduler_name == 'edm_playground_v2.5': elif async_task.scheduler_name == 'edm_playground_v2.5':
final_scheduler_name = 'karras' final_scheduler_name = 'karras'
def patch_edm(unet):
return core.opModelSamplingContinuousEDM.patch(
unet,
sampling=async_task.scheduler_name,
sigma_max=120.0,
sigma_min=0.002)[0]
if pipeline.final_unet is not None: if pipeline.final_unet is not None:
pipeline.final_unet = patch_edm(pipeline.final_unet) pipeline.final_unet = patch_edm(pipeline.final_unet, async_task.scheduler_name)
if pipeline.final_refiner_unet is not None: if pipeline.final_refiner_unet is not None:
pipeline.final_refiner_unet = patch_edm(pipeline.final_refiner_unet) pipeline.final_refiner_unet = patch_edm(pipeline.final_refiner_unet, async_task.scheduler_name)
return final_scheduler_name return final_scheduler_name