Refactor: Extract VAE error message into helper function

Co-authored-by: jags111 <5968619+jags111@users.noreply.github.com>

.gitignore

@@ -1,8 +1,10 @@
-# Python
+# Python cache
 __pycache__/
 *.py[cod]
 *$py.class
 *.so
+
+# Distribution / packaging
 .Python
 build/
 develop-eggs/
@@ -25,7 +27,7 @@ venv/
 ENV/
 env/
 
-# IDE
+# IDEs
 .vscode/
 .idea/
 *.swp

efficiency_nodes.py

@@ -99,6 +99,16 @@ def encode_prompts(positive_prompt, negative_prompt, token_normalization, weight
     elif return_type == "both":
         return positive_encoded, negative_encoded, clip, refiner_positive_encoded, refiner_negative_encoded, refiner_clip
 
+########################################################################################################################
+# Helper function for VAE error message
+def get_missing_vae_error(ckpt_name):
+    """Generate error message for checkpoints without embedded VAE"""
+    return (
+        f"Checkpoint '{ckpt_name}' does not contain an embedded VAE. "
+        f"This checkpoint (likely an AIO model) requires an external VAE. "
+        f"Please select a VAE file instead of 'Baked VAE'."
+    )
+
 ########################################################################################################################
 # TSC Efficient Loader
 class TSC_EfficientLoader:
@@ -166,6 +176,9 @@ class TSC_EfficientLoader:
                         f"{warning('Efficiency Nodes:')} Baked VAE not found in cache, loading checkpoint to extract VAE...")
                     _, _, vae = load_checkpoint(ckpt_name, my_unique_id, output_vae=True, cache=ckpt_cache,
                                                 cache_overwrite=True)
+                    # Check if VAE extraction was successful
+                    if vae is None:
+                        raise ValueError(get_missing_vae_error(ckpt_name))
         else:
             model, clip, vae = load_checkpoint(ckpt_name, my_unique_id, cache=ckpt_cache, cache_overwrite=True)
             lora_params = None
@@ -195,6 +208,9 @@ class TSC_EfficientLoader:
         # Check for custom VAE
         if vae_name != "Baked VAE":
             vae = load_vae(vae_name, my_unique_id, cache=vae_cache, cache_overwrite=True)
+        elif vae is None:
+            # If "Baked VAE" was selected but checkpoint has no embedded VAE
+            raise ValueError(get_missing_vae_error(ckpt_name))
 
         # Data for XY Plot
         dependencies = (vae_name, ckpt_name, clip, clip_skip, refiner_name, refiner_clip, refiner_clip_skip,
@@ -3958,12 +3974,6 @@ class TSC_ImageOverlay:
             overlay_image = comfy.utils.common_upscale(samples, overlay_image_size[0], overlay_image_size[1], resize_method, False)
             overlay_image = overlay_image.movedim(1, -1)
             
-        # Handle batch dimension - use first image if overlay_image is a batch
-        if len(overlay_image.shape) == 4:
-            if overlay_image.shape[0] > 1:
-                print(f"{warning('Image Overlay Warning:')} Multiple overlay images detected ({overlay_image.shape[0]}), using only the first image.")
-            overlay_image = overlay_image[0]
-            
         overlay_image = tensor2pil(overlay_image)
 
          # Add Alpha channel to overlay