Skip to content

fix: Register root utility submodules in _import_structure - #454

Open
richaguptaa17 wants to merge 1 commit into
AI-Hypercomputer:mainfrom
richaguptaa17:fix-lazy-module-import
Open

fix: Register root utility submodules in _import_structure #454
richaguptaa17 wants to merge 1 commit into
AI-Hypercomputer:mainfrom
richaguptaa17:fix-lazy-module-import

Conversation

@richaguptaa17

Copy link
Copy Markdown

Summary of Changes

  • Registered root utility submodules (max_logging, max_utils, pyconfig, maxdiffusion_utils, aot_cache, multihost_dataloading, train_utils, tpu_utils, checkpointing, common_types) in _import_structure in src/maxdiffusion/__init__.py.
  • Added unit test coverage in src/maxdiffusion/tests/models_import_test.py to verify lazy resolution for FlaxAutoencoderKL, FlaxUNet2DConditionModel, BaseStableDiffusionCheckpointer, and top-level utility imports (max_logging, max_utils, pyconfig, maxdiffusion_utils).

Issues Fixed

Fixes #453

… lazy loading AttributeError

Register root utility submodules (`max_logging`, `max_utils`, `pyconfig`, `maxdiffusion_utils`, `aot_cache`, `multihost_dataloading`, `train_utils`, `tpu_utils`, `checkpointing`, `common_types`) in `_import_structure` in `src/maxdiffusion/__init__.py`.

During `_LazyModule` lazy module resolution, submodules importing top-level utility modules like `max_logging` via `from maxdiffusion import max_logging` triggered an `AttributeError` because these utility submodules were not registered in `_import_structure`. Python caught this `AttributeError` during submodule loading and raised `ImportError: cannot import name 'FlaxAutoencoderKL' from 'maxdiffusion.models'`.

Adding these utility submodules to `_import_structure` resolves root attribute lookups lazily, allowing imports of `FlaxAutoencoderKL`, `FlaxUNet2DConditionModel`, and `BaseStableDiffusionCheckpointer` to succeed.

Fixes AI-Hypercomputer#453
@richaguptaa17
richaguptaa17 requested a review from entrpn as a code owner July 31, 2026 09:07
@gemini-code-assist

Copy link
Copy Markdown

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] ImportError: cannot import name 'FlaxAutoencoderKL' from 'maxdiffusion.models'

1 participant