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

[richaguptaa17]

Problem Description

When importing FlaxAutoencoderKL, FlaxUNet2DConditionModel, or BaseStableDiffusionCheckpointer, initialization fails with the following traceback:

Traceback (most recent call last):
  File ".../maxdiffusion/checkpointing/base_stable_diffusion_checkpointer.py", line 26, in <module>
    from maxdiffusion import (
  File ".../maxdiffusion/utils/import_utils.py", line 739, in __getattr__
    value = getattr(module, name)
  File ".../maxdiffusion/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py", line 26, in <module>
    from ...models import FlaxAutoencoderKL, FlaxUNet2DConditionModel
ImportError: cannot import name 'FlaxAutoencoderKL' from 'maxdiffusion.models'

Root Cause

During lazy loading via _LazyModule, src/maxdiffusion/models/quantizations.py executes:

from maxdiffusion import max_logging

Because top-level utility modules (max_logging, max_utils, pyconfig, maxdiffusion_utils, aot_cache, multihost_dataloading, train_utils, tpu_utils, checkpointing, common_types) were not registered in _import_structure in src/maxdiffusion/__init__.py, _LazyModule.__getattr__("max_logging") raised an AttributeError. Python's import machinery intercepted this AttributeError inside the submodule import and re-raised it as ImportError: cannot import name 'FlaxAutoencoderKL' from 'maxdiffusion.models'

[Perseus14]

Thanks, could you create a PR to fix this?