From fbca54957eec4ca49c73eac56159fa92603724b7 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 3 Dec 2025 16:13:07 -0800 Subject: [PATCH] Image generation: Yield partial results for batch count > 1 --- modules/ui_image_generation.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/modules/ui_image_generation.py b/modules/ui_image_generation.py index 06ef3e82..08cf3f64 100644 --- a/modules/ui_image_generation.py +++ b/modules/ui_image_generation.py @@ -677,7 +677,8 @@ def generate(state): if not model_name or model_name == 'None': logger.error("No image model selected. Go to the Model tab and select a model.") - return [] + yield [] + return if shared.image_model is None: result = load_image_model( @@ -690,7 +691,8 @@ def generate(state): ) if result is None: logger.error(f"Failed to load model `{model_name}`.") - return [] + yield [] + return shared.image_model_name = model_name @@ -760,6 +762,7 @@ def generate(state): generator.manual_seed(int(seed + i)) batch_results = shared.image_model(**gen_kwargs).images all_images.extend(batch_results) + yield all_images t1 = time.time() save_generated_images(all_images, state, seed) @@ -768,12 +771,12 @@ def generate(state): total_steps = state["image_steps"] * int(state['image_batch_count']) logger.info(f'Generated {total_images} {"image" if total_images == 1 else "images"} in {(t1 - t0):.2f} seconds ({total_steps / (t1 - t0):.2f} steps/s, seed {seed})') - return all_images + yield all_images except Exception as e: logger.error(f"Image generation failed: {e}") traceback.print_exc() - return [] + yield [] def load_image_model_wrapper(model_name, dtype, attn_backend, cpu_offload, compile_model, quant_method):