Custom UI for bodied macro in edit mode

When using layout: bodied to get a macro body, the Custom UI html is not used when in page edit mode. How can do I enable my own rendering of the body data when in edit mode like it was on Connect so that the user can see how the, in this case diagram, will look on the page before publishing?