⛰️[EPIC]: Enhance diffusion models

[mnabian]

Placeholder for tracking the issues related to enhancing diffusion models

### Tasks
- [ ] Improve unit tests for loss functions, and add tests for dataset, FID, EDMPrecondSR
- [ ] Move `print_module_summary` to core and generalize
- [ ] Fix grad stride UserWarning for DDP
- [ ] Move the Dataset classes out of the `examples` folder and generalize
- [ ] https://github.com/NVIDIA/modulus/issues/205
- [ ] https://github.com/NVIDIA/modulus/issues/206
- [ ] Consider whether `models/diffusion`, and `diffusion/layers` should live in a more general-purpose directory
- [ ] Consider combining the `Linear` class in `diffusion/layers` with the existing class in `models/layers`
- [ ] https://github.com/NVIDIA/modulus/issues/207
- [ ] https://github.com/NVIDIA/modulus/issues/208
- [ ] Evaluate whether the `weight_init` in `models/diffusuon/utils.py` is necessary.
- [ ] Use Modulus `save/load` for checkpoints