nemo_automodel.components.datasets.diffusion.multi_tier_bucketing#

Module Contents#

Classes#

MultiTierBucketCalculator

Calculate resolution buckets constrained by a maximum pixel budget. Supports various aspect ratios, each scaled to fit within the pixel budget.

Data#

logger

API#

nemo_automodel.components.datasets.diffusion.multi_tier_bucketing.logger#

‘getLogger(…)’

class nemo_automodel.components.datasets.diffusion.multi_tier_bucketing.MultiTierBucketCalculator(
quantization: int = 64,
max_pixels: Optional[int] = None,
debug_mode: bool = False,
)[source]#

Calculate resolution buckets constrained by a maximum pixel budget. Supports various aspect ratios, each scaled to fit within the pixel budget.

Initialization

Parameters:

  • quantization – Resolution must be multiple of this (64 for Flux)

  • max_pixels – Maximum pixel count for buckets (default: 256*256 = 65536)

ASPECT_RATIOS#

[(1, 2), (9, 16), (2, 3), (3, 4), (4, 5), (1, 1), (5, 4), (4, 3), (3, 2), (16, 10), (5, 3), (16, 9),…

RESOLUTION_PRESETS#

None

_generate_all_buckets() → List[Dict][source]#

Generate all unique resolution buckets within the pixel budget.

_calculate_max_resolution(
aspect_ratio: float,
) → Optional[Tuple[int, int]][source]#

Calculate the maximum resolution for an aspect ratio within the pixel budget.

For a given aspect ratio r = w/h, and pixel budget P: w * h <= P w = r * h r * h * h <= P h <= sqrt(P / r)

Then w = r * h

_round_to_quantization(value: int) → int[source]#

Round value to nearest quantization multiple.

_build_lookup_structures()[source]#

Build efficient lookup structures.

_print_bucket_summary()[source]#

Print summary of generated buckets.

get_bucket_for_image(
image_width: int,
image_height: int,
) → Dict[source]#

Get the best bucket for an image.

Parameters:

  • image_width – Original image width

  • image_height – Original image height

  • max_pixels – Override max pixels for this query (deprecated, use constructor)

Returns:

Bucket dictionary with resolution and metadata

get_bucket_by_resolution(
width: int,
height: int,
) → Optional[Dict][source]#

Get bucket by exact resolution.

get_bucket_by_id(bucket_id: int) → Dict[source]#

Get bucket by ID.

get_all_buckets() → List[Dict][source]#

Get all buckets.

resize_and_crop(
image,
target_width: int,
target_height: int,
crop_mode: str = 'center',
) → Tuple[source]#

Resize and crop image to target resolution.

Parameters:

  • image – PIL Image or numpy array

  • target_width – Target width

  • target_height – Target height

  • crop_mode – ‘center’, ‘random’, or ‘smart’

Returns:

(resized_image, crop_offset_x, crop_offset_y)

get_dynamic_batch_size(
resolution: Tuple[int, int],
base_batch_size: int = 32,
base_resolution: Tuple[int, int] = (512, 512),
) → int[source]#

Calculate dynamic batch size based on resolution. Larger images get smaller batches to maintain GPU memory usage.

Parameters:

  • resolution – (width, height)

  • base_batch_size – Batch size for base resolution

  • base_resolution – Reference resolution

Returns:

Recommended batch size

classmethod from_preset(
preset: str,
quantization: int = 64,
) → nemo_automodel.components.datasets.diffusion.multi_tier_bucketing.MultiTierBucketCalculator[source]#

Create calculator from a named preset.

Parameters:

  • preset – One of ‘256p’, ‘512p’, ‘768p’, ‘1024p’, ‘1536p’

  • quantization – Resolution quantization

Returns:

MultiTierBucketCalculator instance