Refactoring the diffusion workspace to make it more modular