fixes for ckpt-only validations when EMA is enabled, and, token length limit application for sd3/auraflow/omnigen fix

Author: bghira

helpers/models/auraflow/model.py

@@ -179,13 +179,15 @@ def check_user_config(self):
                 f"{self.NAME} does not support fp8-quanto. Please use fp8-torchao or int8 precision level instead."
             )
         t5_max_length = 120
-        if (
-            self.config.tokenizer_max_length is None
-            or int(self.config.tokenizer_max_length) > t5_max_length
-        ):
+        if self.config.tokenizer_max_length is None or self.config.tokenizer_max_length == 0:
+            logger.warning(
+                f"Setting T5 XXL tokeniser max length to {t5_max_length} for {self.NAME}."
+            )
+            self.config.tokenizer_max_length = t5_max_length
+        if int(self.config.tokenizer_max_length) > t5_max_length:
             if not self.config.i_know_what_i_am_doing:
                 logger.warning(
-                    f"Updating T5 XXL tokeniser max length to {t5_max_length} for {self.NAME}."
+                    f"Overriding T5 XXL tokeniser max length to {t5_max_length} for {self.NAME} because `--i_know_what_i_am_doing` has not been set."
                 )
                 self.config.tokenizer_max_length = t5_max_length
             else:

helpers/models/auraflow/pipeline.py

@@ -658,7 +658,7 @@ def encode_prompt(
         negative_prompt_embeds: Optional[torch.Tensor] = None,
         prompt_attention_mask: Optional[torch.Tensor] = None,
         negative_prompt_attention_mask: Optional[torch.Tensor] = None,
-        max_sequence_length: int = 256,
+        max_sequence_length: int = 120,
         lora_scale: Optional[float] = None,
     ):
         r"""

helpers/models/hidream/model.py

@@ -306,14 +306,15 @@ def check_user_config(self):
             raise ValueError(
                 f"{self.NAME} does not support fp8-quanto. Please use fp8-torchao or int8 precision level instead."
             )
-        t5_max_length = 128
-        if (
-            self.config.tokenizer_max_length is None
-            or int(self.config.tokenizer_max_length) > t5_max_length
-        ):
+        if self.config.tokenizer_max_length is None or self.config.tokenizer_max_length == 0:
+            logger.warning(
+                f"Setting T5 XXL tokeniser max length to {t5_max_length} for {self.NAME}."
+            )
+            self.config.tokenizer_max_length = t5_max_length
+        if int(self.config.tokenizer_max_length) > t5_max_length:
             if not self.config.i_know_what_i_am_doing:
                 logger.warning(
-                    f"Updating T5 XXL tokeniser max length to {t5_max_length} for {self.NAME}."
+                    f"Overriding T5 XXL tokeniser max length to {t5_max_length} for {self.NAME} because `--i_know_what_i_am_doing` has not been set."
                 )
                 self.config.tokenizer_max_length = t5_max_length
             else:

helpers/training/validation.py

@@ -944,7 +944,7 @@ def validate_prompt(
     ):
         """Generate validation images for a single prompt."""
         # Placeholder for actual image generation and logging
-        logger.debug(f"Validating prompt: {prompt}")
+        logger.debug(f"Validating ({validation_shortname}) prompt: {prompt}")
         # benchmarked / stitched validation images
         stitched_validation_images = {}
         # untouched / un-stitched validation images
@@ -1179,9 +1179,15 @@ def validate_prompt(
                     validation_image_results
                 )
                 if self.config.use_ema:
-                    ema_validation_images[validation_shortname].extend(
-                        ema_image_results
-                    )
+                    if validation_shortname in ema_validation_images and ema_image_results is not None:
+                        if ema_validation_images[validation_shortname] is None:
+                            # init the value
+                            ema_validation_images[validation_shortname] = []
+                        if isinstance(ema_validation_images[validation_shortname], list):
+                            # if we have a list of images, we can stitch them.
+                            ema_validation_images[validation_shortname].extend(
+                                ema_image_results
+                            )
 
             except Exception as e:
                 import traceback