feat: --image-source flag to add noise generation and custom dirs#710
Open
ajcasagrande wants to merge 1 commit intomainfrom
Open
feat: --image-source flag to add noise generation and custom dirs#710ajcasagrande wants to merge 1 commit intomainfrom
ajcasagrande wants to merge 1 commit intomainfrom
Conversation
…lasses Move AudioDefaults, ImageDefaults, and VideoDefaults out of the shared config_defaults module and into their respective config classes to improve cohesion. Add ImageSource enum and --image-source CLI option to support noise-based and custom directory image generation modes alongside the existing assets-based approach. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
Try out this PRQuick install: pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@8a9f290d5281dbcdfca264bdfa44b3e7bc869a0bRecommended with virtual environment (using uv): uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@8a9f290d5281dbcdfca264bdfa44b3e7bc869a0bLast updated for commit: |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
ajcasagrande
commented
Feb 25, 2026
| from aiperf.common.config.base_config import BaseConfig | ||
| from aiperf.common.config.cli_parameter import CLIParameter, DisableCLI | ||
| from aiperf.common.config.config_defaults import ( | ||
| AudioDefaults, |
Contributor
Author
There was a problem hiding this comment.
I removed the defaults file and inline the defaults directly to reduce cognitive load
Contributor
There was a problem hiding this comment.
I thought we generally like relegating magic numbers into these default classes?
WalkthroughThis pull request refactors media generation configuration by removing centralized defaults classes and embedding defaults directly into config fields. It introduces validation methods ensuring generation is only enabled when required options are provided, adds an Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Tip Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs). Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (2)
tests/unit/dataset/generator/test_image_generator.py (1)
356-362: Verify thatnot hasattr(generator, "config")is the intended contract for disabled generators.This test asserts that
ImageGeneratordoes not setself.configwhen images are disabled (early return in__init__). This is a somewhat fragile assertion—ifBaseGenerator.__init__or any future refactor setsconfig, this test breaks. Consider asserting a more explicit disabled-state indicator if one exists, or adding a brief comment documenting why this assertion is appropriate.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/dataset/generator/test_image_generator.py` around lines 356 - 362, The test currently relies on the absence of an attribute by asserting not hasattr(generator, "config"), which is fragile if BaseGenerator.__init__ or future refactors add that attribute; update the test in TestImageGeneratorDisabled.test_disabled_images_skips_init to assert an explicit disabled-state indicator instead (e.g., assert generator.disabled is True or assert generator.enabled is False or check generator.initialized is False) using whatever explicit flag exists on ImageGenerator/ImageConfig, and if no such flag exists either add a short comment explaining why the hasattr check is the intended contract or add a small explicit property (e.g., disabled or initialized) in ImageGenerator.__init__ to signal the early-return state and assert that instead.src/aiperf/dataset/composer/synthetic.py (1)
30-39: Pre-existing: validation doesn't account for video-only payloads.The guard at lines 30-34 raises if prompt, image, and audio are all disabled—but doesn't check
self.include_video. If video is the only enabled input, this would incorrectly raiseValueError. This is pre-existing (not introduced by this PR), but worth noting sinceinclude_videowas just wired up.Proposed fix
if ( not self.include_prompt and not self.include_image and not self.include_audio + and not self.include_video ): raise ValueError( "All synthetic data are disabled. " - "Please enable at least one of prompt, image, or audio by " + "Please enable at least one of prompt, image, audio, or video by " "setting the mean to a positive value." )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/aiperf/dataset/composer/synthetic.py` around lines 30 - 39, The validation that raises ValueError when all synthetic inputs are disabled currently checks only self.include_prompt, self.include_image, and self.include_audio; update that guard in the synthetic data validation block to also consider self.include_video (i.e., require at least one of self.include_prompt, self.include_image, self.include_audio, or self.include_video to be True) so a video-only payload does not incorrectly trigger the error; modify the conditional that constructs the exception to include self.include_video or refactor to assert any([self.include_prompt, self.include_image, self.include_audio, self.include_video]) before raising.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/aiperf/common/config/audio_config.py`:
- Around line 155-158: The current check raises whenever audio options were
provided while audio_enabled() is False, which wrongly rejects an explicit
disable via batch_size==0; update the logic around the raise to allow the
explicit disable case by: compute whether only batch_size was provided (e.g.,
only_batch_size_provided = audio_options_set and all other audio option fields
like audio_length_mean are unset/None), then change the condition to raise only
when not self.audio_enabled() and audio_options_set and not (self.batch_size ==
0 and only_batch_size_provided); reference the existing audio_enabled(),
audio_options_set, and self.batch_size symbols when making this change.
In `@src/aiperf/common/config/image_config.py`:
- Around line 167-170: The validator is rejecting an explicit disable via
"--image-batch-size 0"; update the check around images_enabled() so providing
image options is allowed when the user explicitly set batch size to 0.
Concretely, in the block that raises when not self.images_enabled() and
image_options_set, add a condition to skip the error if the batch-size flag is
explicitly zero (e.g. check self.image_batch_size == 0 or the equivalent
attribute/flag used to store --image-batch-size); alternatively implement
images_enabled() to return False when batch size == 0 and ensure the raise only
happens if images are truly enabled. Reference images_enabled(),
image_options_set and the image-batch-size flag/attribute when making the
change.
In `@src/aiperf/common/config/video_config.py`:
- Around line 35-38: The current check raises when video generation is disabled
but any video option was provided, which incorrectly rejects the explicit
disable flag `--video-batch-size 0`; update the conditional in the validator so
it does not raise if the user explicitly set video_batch_size to 0. Concretely,
change the guard using videos_enabled() and video_options_set to also allow the
case self.video_batch_size == 0 (e.g., only raise if not self.videos_enabled()
and video_options_set and self.video_batch_size != 0), referencing
videos_enabled(), video_options_set, and the video_batch_size attribute.
In `@src/aiperf/common/random_generator.py`:
- Around line 122-134: The integers method is missing type annotations for size,
dtype and the return value; update the integers signature in class
RandomGenerator (method integers) to annotate size as Optional[Union[int,
Tuple[int, ...]]], dtype as Optional[np.typing.DTypeLike] (or np.dtype | None),
and the return type as Union[int, np.ndarray]; also add the necessary imports
(Optional, Union, Tuple from typing and np.typing as needed) so the function has
full parameter and return type hints.
In `@src/aiperf/dataset/generator/image.py`:
- Around line 67-77: The code currently collects all directory entries into
image_paths and attempts to open each with Image.open causing failures on
subdirectories or non-image files; update the logic in the generator so
image_paths only includes regular files with image-like extensions or add a safe
open loop that checks Path(path).is_file() and wraps Image.open(path) in a
try/except (catching OSError/UnidentifiedImageError) to skip non-image files,
appending only successfully opened images into the images list; reference
image_paths, source_path, Image.open, and images to locate where to apply the
filtering and error-handling change.
- Around line 34-37: In the ImageGenerator constructor (__init__) do not return
before required instance state is set; instead always initialize self.config and
any source/factory attributes (e.g., image source factory) and set an enabled
flag (e.g., self.images_enabled = config.images_enabled()) so generate() can
safely inspect state; if images are disabled, log via self.debug and early-exit
behavior should rely on the flag rather than leaving members uninitialized so
later calls to generate() won't raise AttributeError.
- Around line 60-79: The synchronous disk I/O in _load_source_images_from_disk
(glob + Image.open + img.copy) must be moved out of __init__ into an
async-friendly lifecycle: create an async method (e.g., async_load_source_images
or load_source_images_async) that performs the work on a threadpool (use
asyncio.to_thread to call the existing _load_source_images_from_disk logic or
move the per-file open/copy into the thread) and update any startup code to
await that method from the appropriate async initialization hook; ensure the
debug(...) call remains after loading and that exceptions
(FileNotFoundError/NotADirectoryError/ValueError) are preserved and raised from
the async loader so callers can handle them.
---
Nitpick comments:
In `@src/aiperf/dataset/composer/synthetic.py`:
- Around line 30-39: The validation that raises ValueError when all synthetic
inputs are disabled currently checks only self.include_prompt,
self.include_image, and self.include_audio; update that guard in the synthetic
data validation block to also consider self.include_video (i.e., require at
least one of self.include_prompt, self.include_image, self.include_audio, or
self.include_video to be True) so a video-only payload does not incorrectly
trigger the error; modify the conditional that constructs the exception to
include self.include_video or refactor to assert any([self.include_prompt,
self.include_image, self.include_audio, self.include_video]) before raising.
In `@tests/unit/dataset/generator/test_image_generator.py`:
- Around line 356-362: The test currently relies on the absence of an attribute
by asserting not hasattr(generator, "config"), which is fragile if
BaseGenerator.__init__ or future refactors add that attribute; update the test
in TestImageGeneratorDisabled.test_disabled_images_skips_init to assert an
explicit disabled-state indicator instead (e.g., assert generator.disabled is
True or assert generator.enabled is False or check generator.initialized is
False) using whatever explicit flag exists on ImageGenerator/ImageConfig, and if
no such flag exists either add a short comment explaining why the hasattr check
is the intended contract or add a small explicit property (e.g., disabled or
initialized) in ImageGenerator.__init__ to signal the early-return state and
assert that instead.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
docs/cli_options.mdsrc/aiperf/common/config/__init__.pysrc/aiperf/common/config/audio_config.pysrc/aiperf/common/config/config_defaults.pysrc/aiperf/common/config/image_config.pysrc/aiperf/common/config/video_config.pysrc/aiperf/common/enums/__init__.pysrc/aiperf/common/enums/enums.pysrc/aiperf/common/random_generator.pysrc/aiperf/dataset/composer/synthetic.pysrc/aiperf/dataset/generator/image.pytests/unit/common/config/test_audio_config.pytests/unit/common/config/test_image_config.pytests/unit/dataset/composer/test_synthetic_composer.pytests/unit/dataset/generator/test_image_generator.pytests/unit/dataset/generator/test_video_generator.py
💤 Files with no reviewable changes (3)
- tests/unit/dataset/composer/test_synthetic_composer.py
- src/aiperf/common/config/init.py
- src/aiperf/common/config/config_defaults.py
Comment on lines
+155
to
+158
| if not self.audio_enabled() and audio_options_set: | ||
| raise ValueError( | ||
| "Audio generation is disabled but audio options were provided. Please set `--audio-batch-size` and `--audio-length-mean` to enable audio generation." | ||
| ) |
There was a problem hiding this comment.
batch_size=0 should be accepted as an explicit disabled state.
Line 155 treats any explicitly set audio option as invalid when disabled, which also rejects batch_size=0. Since audio_enabled() uses batch_size > 0, this should be a valid explicit disable path.
Proposed fix
`@model_validator`(mode="after")
def _validate_audio_options(self) -> Self:
"""Validate the audio options."""
+ if self.batch_size == 0:
+ return self
+
audio_options_set = {*self.model_fields_set, *self.length.model_fields_set}
if not self.audio_enabled() and audio_options_set:
raise ValueError(
"Audio generation is disabled but audio options were provided. Please set `--audio-batch-size` and `--audio-length-mean` to enable audio generation."
)
return self🧰 Tools
🪛 Ruff (0.15.2)
[warning] 156-158: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/aiperf/common/config/audio_config.py` around lines 155 - 158, The current
check raises whenever audio options were provided while audio_enabled() is
False, which wrongly rejects an explicit disable via batch_size==0; update the
logic around the raise to allow the explicit disable case by: compute whether
only batch_size was provided (e.g., only_batch_size_provided = audio_options_set
and all other audio option fields like audio_length_mean are unset/None), then
change the condition to raise only when not self.audio_enabled() and
audio_options_set and not (self.batch_size == 0 and only_batch_size_provided);
reference the existing audio_enabled(), audio_options_set, and self.batch_size
symbols when making this change.
Comment on lines
+167
to
+170
| if not self.images_enabled() and image_options_set: | ||
| raise ValueError( | ||
| "Image generation is disabled but image options were provided. Please set `--image-width-mean` and `--image-height-mean` to enable image generation." | ||
| ) |
There was a problem hiding this comment.
Allow explicit disable via --image-batch-size 0.
Line 167 currently rejects explicitly setting batch_size=0, even though Line 116 documents it as a disable mechanism. This makes a valid disable flow fail at validation time.
Proposed fix
`@model_validator`(mode="after")
def _validate_image_options(self) -> Self:
"""Validate the image options."""
+ if self.batch_size == 0:
+ return self
+
image_options_set = {
*self.model_fields_set,
*self.width.model_fields_set,
*self.height.model_fields_set,
}
if not self.images_enabled() and image_options_set:
raise ValueError(
- "Image generation is disabled but image options were provided. Please set `--image-width-mean` and `--image-height-mean` to enable image generation."
+ "Image generation is disabled but image options were provided. Please set `--image-batch-size` (>0), `--image-width-mean`, and `--image-height-mean` to enable image generation."
)
return self🧰 Tools
🪛 Ruff (0.15.2)
[warning] 168-170: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/aiperf/common/config/image_config.py` around lines 167 - 170, The
validator is rejecting an explicit disable via "--image-batch-size 0"; update
the check around images_enabled() so providing image options is allowed when the
user explicitly set batch size to 0. Concretely, in the block that raises when
not self.images_enabled() and image_options_set, add a condition to skip the
error if the batch-size flag is explicitly zero (e.g. check
self.image_batch_size == 0 or the equivalent attribute/flag used to store
--image-batch-size); alternatively implement images_enabled() to return False
when batch size == 0 and ensure the raise only happens if images are truly
enabled. Reference images_enabled(), image_options_set and the image-batch-size
flag/attribute when making the change.
Comment on lines
+35
to
+38
| if not self.videos_enabled() and video_options_set: | ||
| raise ValueError( | ||
| "Video generation is disabled but video options were provided. Please set `--video-width` and `--video-height` to enable video generation." | ||
| ) |
There was a problem hiding this comment.
Permit explicit disable with --video-batch-size 0.
Line 35 rejects disabled configs whenever any field was explicitly set. That makes --video-batch-size 0 invalid, despite Line 55 documenting it as the disable path.
Proposed fix
`@model_validator`(mode="after")
def _validate_video_options(self) -> Self:
"""Validate the video options."""
+ if self.batch_size == 0:
+ return self
+
video_options_set = self.model_fields_set
if not self.videos_enabled() and video_options_set:
raise ValueError(
- "Video generation is disabled but video options were provided. Please set `--video-width` and `--video-height` to enable video generation."
+ "Video generation is disabled but video options were provided. Please set `--video-batch-size` (>0), `--video-width`, and `--video-height` to enable video generation."
)
return self🧰 Tools
🪛 Ruff (0.15.2)
[warning] 36-38: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/aiperf/common/config/video_config.py` around lines 35 - 38, The current
check raises when video generation is disabled but any video option was
provided, which incorrectly rejects the explicit disable flag
`--video-batch-size 0`; update the conditional in the validator so it does not
raise if the user explicitly set video_batch_size to 0. Concretely, change the
guard using videos_enabled() and video_options_set to also allow the case
self.video_batch_size == 0 (e.g., only raise if not self.videos_enabled() and
video_options_set and self.video_batch_size != 0), referencing videos_enabled(),
video_options_set, and the video_batch_size attribute.
Comment on lines
+122
to
+134
| def integers(self, low: int, high: int | None = None, size=None, dtype=np.int64): | ||
| """Generate random integers from [low, high) using NumPy. | ||
|
|
||
| Args: | ||
| low: Lowest integer (inclusive), or if high is None, then [0, low) | ||
| high: Highest integer (exclusive), optional | ||
| size: Output shape, optional | ||
| dtype: Desired NumPy dtype for the result (default: np.int64) | ||
|
|
||
| Returns: | ||
| Random integer or array of integers | ||
| """ | ||
| return self._numpy_rng.integers(low, high, size=size) | ||
| return self._numpy_rng.integers(low, high, size=size, dtype=dtype) |
There was a problem hiding this comment.
Add missing type hints to integers signature and return type.
This updated method still leaves size, dtype, and the return type unannotated, which breaks the repo’s Python typing rule.
Proposed fix
+from numpy.typing import DTypeLike
+
- def integers(self, low: int, high: int | None = None, size=None, dtype=np.int64):
+ def integers(
+ self,
+ low: int,
+ high: int | None = None,
+ size: int | tuple[int, ...] | None = None,
+ dtype: DTypeLike = np.int64,
+ ) -> int | np.ndarray:
"""Generate random integers from [low, high) using NumPy.As per coding guidelines: "Add type hints on ALL functions (params and return)".
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def integers(self, low: int, high: int | None = None, size=None, dtype=np.int64): | |
| """Generate random integers from [low, high) using NumPy. | |
| Args: | |
| low: Lowest integer (inclusive), or if high is None, then [0, low) | |
| high: Highest integer (exclusive), optional | |
| size: Output shape, optional | |
| dtype: Desired NumPy dtype for the result (default: np.int64) | |
| Returns: | |
| Random integer or array of integers | |
| """ | |
| return self._numpy_rng.integers(low, high, size=size) | |
| return self._numpy_rng.integers(low, high, size=size, dtype=dtype) | |
| from numpy.typing import DTypeLike | |
| def integers( | |
| self, | |
| low: int, | |
| high: int | None = None, | |
| size: int | tuple[int, ...] | None = None, | |
| dtype: DTypeLike = np.int64, | |
| ) -> int | np.ndarray: | |
| """Generate random integers from [low, high) using NumPy. | |
| Args: | |
| low: Lowest integer (inclusive), or if high is None, then [0, low) | |
| high: Highest integer (exclusive), optional | |
| size: Output shape, optional | |
| dtype: Desired NumPy dtype for the result (default: np.int64) | |
| Returns: | |
| Random integer or array of integers | |
| """ | |
| return self._numpy_rng.integers(low, high, size=size, dtype=dtype) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/aiperf/common/random_generator.py` around lines 122 - 134, The integers
method is missing type annotations for size, dtype and the return value; update
the integers signature in class RandomGenerator (method integers) to annotate
size as Optional[Union[int, Tuple[int, ...]]], dtype as
Optional[np.typing.DTypeLike] (or np.dtype | None), and the return type as
Union[int, np.ndarray]; also add the necessary imports (Optional, Union, Tuple
from typing and np.typing as needed) so the function has full parameter and
return type hints.
Comment on lines
+34
to
+37
| if not config.images_enabled(): | ||
| self.debug(lambda: "Images are disabled, skipping image generation") | ||
| return | ||
|
|
There was a problem hiding this comment.
Avoid returning from __init__ before required state is initialized.
Line 34 exits constructor before key members are guaranteed (self.config, source factory state). If generate() is called later, this can fail with AttributeError instead of a controlled disabled-path error.
Proposed fix
def __init__(self, config: ImageConfig, **kwargs):
super().__init__(**kwargs)
+ self.config = config
+ self._images_enabled = config.images_enabled()
- if not config.images_enabled():
+ if not self._images_enabled:
self.debug(lambda: "Images are disabled, skipping image generation")
return
# Separate RNGs for independent concerns
self._dimensions_rng = rng.derive("dataset.image.dimensions")
self._format_rng = rng.derive("dataset.image.format")
- self.config = config def generate(self, *args, **kwargs) -> str:
+ if not self._images_enabled:
+ raise RuntimeError("Image generation is disabled for the current configuration")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/aiperf/dataset/generator/image.py` around lines 34 - 37, In the
ImageGenerator constructor (__init__) do not return before required instance
state is set; instead always initialize self.config and any source/factory
attributes (e.g., image source factory) and set an enabled flag (e.g.,
self.images_enabled = config.images_enabled()) so generate() can safely inspect
state; if images are disabled, log via self.debug and early-exit behavior should
rely on the flag rather than leaving members uninitialized so later calls to
generate() won't raise AttributeError.
Comment on lines
+60
to
+79
| def _load_source_images_from_disk(self, source_path: Path) -> list[Image.Image]: | ||
| """Load source images from the given directory.""" | ||
| if not source_path.exists(): | ||
| raise FileNotFoundError(f"The directory '{source_path}' does not exist.") | ||
| if not source_path.is_dir(): | ||
| raise NotADirectoryError(f"The path '{source_path}' is not a directory.") | ||
|
|
||
| image_paths = sorted(glob.glob(str(source_path / "*"))) | ||
| if not image_paths: | ||
| raise ValueError( | ||
| f"No source images found in '{source_images_dir}'. " | ||
| "Please ensure the source_images directory contains at least one image file." | ||
| f"No source images found in '{source_path}'. " | ||
| "Please ensure the directory contains at least one image file." | ||
| ) | ||
|
|
||
| self._source_images = [] | ||
| images: list[Image.Image] = [] | ||
| for path in image_paths: | ||
| with Image.open(path) as img: | ||
| self._source_images.append(img.copy()) | ||
| self.debug( | ||
| lambda: f"Pre-loaded {len(self._source_images)} source images into memory" | ||
| ) | ||
| images.append(img.copy()) | ||
| self.debug(lambda: f"Pre-loaded {len(images)} source images from disk") | ||
| return images |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# First, let's check if the file exists and read the relevant lines
if [ -f "src/aiperf/dataset/generator/image.py" ]; then
wc -l src/aiperf/dataset/generator/image.py
echo "=== Lines 50-90 ==="
sed -n '50,90p' src/aiperf/dataset/generator/image.py
else
echo "File not found, searching for it..."
find . -name "image.py" -path "*/dataset/generator/*"
fiRepository: ai-dynamo/aiperf
Length of output: 1998
🏁 Script executed:
#!/bin/bash
# Check the full class structure to understand the context
if [ -f "src/aiperf/dataset/generator/image.py" ]; then
echo "=== Full file (checking size first) ==="
wc -l src/aiperf/dataset/generator/image.py
echo "=== Class definition and method signatures ==="
sed -n '1,100p' src/aiperf/dataset/generator/image.py
fiRepository: ai-dynamo/aiperf
Length of output: 4342
🏁 Script executed:
#!/bin/bash
# Search for where _load_source_images_from_disk is called
rg "_load_source_images_from_disk" --type py -B 3 -A 3Repository: ai-dynamo/aiperf
Length of output: 1858
🏁 Script executed:
#!/bin/bash
# Find BaseGenerator to understand the class structure
fd "base\.py" dataset/generator/ --type fRepository: ai-dynamo/aiperf
Length of output: 165
🏁 Script executed:
#!/bin/bash
# Read the base generator file
if [ -f "src/aiperf/dataset/generator/base.py" ]; then
wc -l src/aiperf/dataset/generator/base.py
cat -n src/aiperf/dataset/generator/base.py
fiRepository: ai-dynamo/aiperf
Length of output: 1261
🏁 Script executed:
#!/bin/bash
# Check if there are any async patterns or lifecycle hooks in BaseGenerator or related classes
rg "class BaseGenerator|@on_init|@on_start|async def" src/aiperf/dataset/generator/ --type py -B 2 -A 5Repository: ai-dynamo/aiperf
Length of output: 623
🏁 Script executed:
#!/bin/bash
# Also check if ImageGenerator or similar components use AIPerfLifecycleMixin
rg "AIPerfLifecycleMixin|BaseComponentService" src/aiperf/dataset/ --type pyRepository: ai-dynamo/aiperf
Length of output: 553
🏁 Script executed:
#!/bin/bash
# Search for ImageGenerator usage and instantiation
rg "ImageGenerator" --type py -B 3 -A 3Repository: ai-dynamo/aiperf
Length of output: 18572
🏁 Script executed:
#!/bin/bash
# Check DatasetManager to see how it uses generators
rg "class DatasetManager|ImageGenerator" src/aiperf/dataset/dataset_manager.py --type py -A 10 -B 2Repository: ai-dynamo/aiperf
Length of output: 628
🏁 Script executed:
#!/bin/bash
# Find where DatasetManager or ImageGenerator is instantiated in async context
rg "ImageGenerator\(|DatasetManager" --type py | head -20Repository: ai-dynamo/aiperf
Length of output: 2193
🏁 Script executed:
#!/bin/bash
# Check how BaseComposer is used and if it's instantiated in async context
rg "BaseComposer|class.*Composer" src/aiperf/dataset/composer/ --type py -B 2 -A 10 | head -80Repository: ai-dynamo/aiperf
Length of output: 4205
🏁 Script executed:
#!/bin/bash
# Check if DatasetManager instantiation happens in async context
rg "class DatasetManager" src/aiperf/dataset/dataset_manager.py -A 30 | head -50Repository: ai-dynamo/aiperf
Length of output: 1313
🏁 Script executed:
#!/bin/bash
# Check if DatasetManager.__init__ is async or if instantiation happens in async lifecycle
rg "@on_init|async def __init__|def __init__" src/aiperf/dataset/dataset_manager.py -A 5Repository: ai-dynamo/aiperf
Length of output: 217
Move image loading out of __init__ to avoid blocking I/O during initialization.
The _load_source_images_from_disk method performs synchronous I/O operations (glob and image decoding) in the constructor. This blocks initialization and violates the requirement to use async/await for all I/O operations. Move image loading to a deferred initialization pattern or lifecycle hook that can properly handle asynchronous operations.
🧰 Tools
🪛 Ruff (0.15.2)
[warning] 63-63: Avoid specifying long messages outside the exception class
(TRY003)
[warning] 65-65: Avoid specifying long messages outside the exception class
(TRY003)
[warning] 69-72: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/aiperf/dataset/generator/image.py` around lines 60 - 79, The synchronous
disk I/O in _load_source_images_from_disk (glob + Image.open + img.copy) must be
moved out of __init__ into an async-friendly lifecycle: create an async method
(e.g., async_load_source_images or load_source_images_async) that performs the
work on a threadpool (use asyncio.to_thread to call the existing
_load_source_images_from_disk logic or move the per-file open/copy into the
thread) and update any startup code to await that method from the appropriate
async initialization hook; ensure the debug(...) call remains after loading and
that exceptions (FileNotFoundError/NotADirectoryError/ValueError) are preserved
and raised from the async loader so callers can handle them.
Comment on lines
+67
to
+77
| image_paths = sorted(glob.glob(str(source_path / "*"))) | ||
| if not image_paths: | ||
| raise ValueError( | ||
| f"No source images found in '{source_images_dir}'. " | ||
| "Please ensure the source_images directory contains at least one image file." | ||
| f"No source images found in '{source_path}'. " | ||
| "Please ensure the directory contains at least one image file." | ||
| ) | ||
|
|
||
| self._source_images = [] | ||
| images: list[Image.Image] = [] | ||
| for path in image_paths: | ||
| with Image.open(path) as img: | ||
| self._source_images.append(img.copy()) | ||
| self.debug( | ||
| lambda: f"Pre-loaded {len(self._source_images)} source images into memory" | ||
| ) | ||
| images.append(img.copy()) |
There was a problem hiding this comment.
Filter/skip non-image entries for custom source directories.
Line 67 includes every directory entry (*), and Line 76 tries to open each as an image. A custom folder containing subdirectories or non-image files will currently hard-fail generation.
Proposed fix
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
...
- image_paths = sorted(glob.glob(str(source_path / "*")))
+ image_paths = sorted(glob.glob(str(source_path / "*")))
if not image_paths:
raise ValueError(
f"No source images found in '{source_path}'. "
"Please ensure the directory contains at least one image file."
)
images: list[Image.Image] = []
for path in image_paths:
- with Image.open(path) as img:
- images.append(img.copy())
+ path_obj = Path(path)
+ if not path_obj.is_file():
+ continue
+ try:
+ with Image.open(path_obj) as img:
+ images.append(img.copy())
+ except UnidentifiedImageError:
+ self.debug(lambda: f"Skipping non-image file '{path_obj}'")
+
+ if not images:
+ raise ValueError(
+ f"No readable image files found in '{source_path}'. "
+ "Please ensure the directory contains valid image files."
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| image_paths = sorted(glob.glob(str(source_path / "*"))) | |
| if not image_paths: | |
| raise ValueError( | |
| f"No source images found in '{source_images_dir}'. " | |
| "Please ensure the source_images directory contains at least one image file." | |
| f"No source images found in '{source_path}'. " | |
| "Please ensure the directory contains at least one image file." | |
| ) | |
| self._source_images = [] | |
| images: list[Image.Image] = [] | |
| for path in image_paths: | |
| with Image.open(path) as img: | |
| self._source_images.append(img.copy()) | |
| self.debug( | |
| lambda: f"Pre-loaded {len(self._source_images)} source images into memory" | |
| ) | |
| images.append(img.copy()) | |
| image_paths = sorted(glob.glob(str(source_path / "*"))) | |
| if not image_paths: | |
| raise ValueError( | |
| f"No source images found in '{source_path}'. " | |
| "Please ensure the directory contains at least one image file." | |
| ) | |
| images: list[Image.Image] = [] | |
| for path in image_paths: | |
| path_obj = Path(path) | |
| if not path_obj.is_file(): | |
| continue | |
| try: | |
| with Image.open(path_obj) as img: | |
| images.append(img.copy()) | |
| except UnidentifiedImageError: | |
| self.debug(lambda: f"Skipping non-image file '{path_obj}'") | |
| if not images: | |
| raise ValueError( | |
| f"No readable image files found in '{source_path}'. " | |
| "Please ensure the directory contains valid image files." | |
| ) |
🧰 Tools
🪛 Ruff (0.15.2)
[warning] 69-72: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/aiperf/dataset/generator/image.py` around lines 67 - 77, The code
currently collects all directory entries into image_paths and attempts to open
each with Image.open causing failures on subdirectories or non-image files;
update the logic in the generator so image_paths only includes regular files
with image-like extensions or add a safe open loop that checks
Path(path).is_file() and wraps Image.open(path) in a try/except (catching
OSError/UnidentifiedImageError) to skip non-image files, appending only
successfully opened images into the images list; reference image_paths,
source_path, Image.open, and images to locate where to apply the filtering and
error-handling change.
ajcasagrande
commented
Feb 25, 2026
Comment on lines
+151
to
+163
| @model_validator(mode="after") | ||
| def _validate_audio_options(self) -> Self: | ||
| """Validate the audio options.""" | ||
| audio_options_set = {*self.model_fields_set, *self.length.model_fields_set} | ||
| if not self.audio_enabled() and audio_options_set: | ||
| raise ValueError( | ||
| "Audio generation is disabled but audio options were provided. Please set `--audio-batch-size` and `--audio-length-mean` to enable audio generation." | ||
| ) | ||
| return self | ||
|
|
||
| def audio_enabled(self) -> bool: | ||
| """Check if audio is enabled.""" | ||
| return self.length.mean > 0 and self.batch_size > 0 |
Contributor
Author
There was a problem hiding this comment.
Added basic sanity checks to prevent users from using multi modal options without enabling that modality
| ] | ||
|
|
||
| @model_validator(mode="after") | ||
| def _validate_image_options(self) -> Self: |
Contributor
There was a problem hiding this comment.
do we have warnings for if someone supplies a --input-file that, say, has images, and they also try to use some synthetic image options?
Contributor
matthewkotila
left a comment
matthewkotila
left a comment
There was a problem hiding this comment.
Looks good overall, just some small questions
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Move AudioDefaults, ImageDefaults, and VideoDefaults out of the shared config_defaults module and into their respective config classes to improve cohesion. Add ImageSource enum and --image-source CLI option to support noise-based and custom directory image generation modes alongside the existing assets-based approach.
Summary by CodeRabbit
Release Notes
New Features
--image-sourceCLI option supporting multiple image input modes: bundled assets, dynamically generated noise, or custom directory paths.Documentation
Bug Fixes