Try to include torchscript tests for the latent diffusion unet networks. For this, it might be necessary to remove the use of TimestepEmbedSequential and TimestepBlock by creating blocks exclusive to when we have the time embedding included. A similar solution to Huggingface diffuser might be useful: https://github.com/huggingface/diffusers/blob/2d35f6733a2d698e8917896071444a5923993ae7/src/diffusers/models/unet_blocks.py#L461 https://github.com/huggingface/diffusers/blob/2d35f6733a2d698e8917896071444a5923993ae7/src/diffusers/models/unet_blocks.py#L379 https://github.com/huggingface/diffusers/blob/2d35f6733a2d698e8917896071444a5923993ae7/src/diffusers/models/unet_blocks.py#L576
Try to include torchscript tests for the latent diffusion unet networks. For this, it might be necessary to remove the use of TimestepEmbedSequential and TimestepBlock by creating blocks exclusive to when we have the time embedding included.
A similar solution to Huggingface diffuser might be useful:
https://github.com/huggingface/diffusers/blob/2d35f6733a2d698e8917896071444a5923993ae7/src/diffusers/models/unet_blocks.py#L461
https://github.com/huggingface/diffusers/blob/2d35f6733a2d698e8917896071444a5923993ae7/src/diffusers/models/unet_blocks.py#L379
https://github.com/huggingface/diffusers/blob/2d35f6733a2d698e8917896071444a5923993ae7/src/diffusers/models/unet_blocks.py#L576