graph LR
micro_diffusion_models_dit_DiT["micro_diffusion.models.dit.DiT"]
micro_diffusion_models_dit___init__["micro_diffusion.models.dit.__init__"]
micro_diffusion_models_dit_forward["micro_diffusion.models.dit.forward"]
micro_diffusion_models_dit_initialize_weights["micro_diffusion.models.dit.initialize_weights"]
micro_diffusion_models_dit_forward_without_cfg["micro_diffusion.models.dit.forward_without_cfg"]
micro_diffusion_models_dit_forward_with_cfg["micro_diffusion.models.dit.forward_with_cfg"]
micro_diffusion_models_utils_TimestepEmbedder["micro_diffusion.models.utils.TimestepEmbedder"]
micro_diffusion_models_utils_CaptionProjection["micro_diffusion.models.utils.CaptionProjection"]
micro_diffusion_models_dit___init__ -- "initializes" --> micro_diffusion_models_dit_DiT
micro_diffusion_models_dit_DiT -- "exposes" --> micro_diffusion_models_dit_forward
micro_diffusion_models_dit___init__ -- "invokes" --> micro_diffusion_models_dit_initialize_weights
micro_diffusion_models_dit_forward -- "delegates" --> micro_diffusion_models_dit_forward_without_cfg
micro_diffusion_models_dit_forward -- "dispatches" --> micro_diffusion_models_dit_forward_with_cfg
micro_diffusion_models_dit_forward_with_cfg -- "reuses" --> micro_diffusion_models_dit_forward_without_cfg
micro_diffusion_models_dit___init__ -- "instantiates and configures" --> micro_diffusion_models_utils_TimestepEmbedder
micro_diffusion_models_dit___init__ -- "instantiates and integrates" --> micro_diffusion_models_utils_CaptionProjection
The Diffusion Transformer (DiT) Architecture subsystem is primarily encapsulated within the micro_diffusion.models.dit module, with key dependencies on shared utilities in micro_diffusion.models.utils. This subsystem defines the core neural network structure, its initialization, and its forward pass logic, aligning with the "Model Definition" and "Model-Centric View" patterns of an ML Toolkit/Research Project.
The top-level orchestrator and container for the entire Diffusion Transformer model. It serves as the primary interface for model inference and encapsulates the complete architecture.
Related Classes/Methods:
Manages the construction of the DiT network, including the instantiation of its various blocks, layers, and the integration of necessary utility components (e.g., TimestepEmbedder, CaptionProjection).
Related Classes/Methods:
Orchestrates the model's inference flow, including the decision to apply classifier-free guidance (CFG) and delegating to specific forward pass methods (forward_without_cfg or forward_with_cfg).
Related Classes/Methods:
Manages the initial setup and configuration of the model's parameters and weights, ensuring proper initialization for training stability and performance.
Related Classes/Methods:
Implements the fundamental forward pass logic of the DiT model when classifier-free guidance is not applied. This represents the core computational path for unconditional generation.
Related Classes/Methods:
Implements the forward pass logic that incorporates classifier-free guidance, typically by combining conditional and unconditional outputs to enhance generation quality.
Related Classes/Methods:
Transforms numerical timestep values into dense, learnable embeddings suitable for neural network consumption, crucial for informing the model about the diffusion process stage.
Related Classes/Methods:
Projects textual caption embeddings into a compatible dimension for integration into the DiT model, allowing text-based conditioning.
Related Classes/Methods: