Update README.md

README.md

@@ -8,7 +8,8 @@ A Generative diffusion model for 2D B-mode cardiac ultrasound images.
 This model can be used for unconditional sampling as well as posterior sampling and inpainting tasks.
 
 - The model is trained on EchonetLVH at a resolution of `256x256`
-- The three folders contain a 1-frame, 3-frame and 12-frame model respectively, all with the same unet architecture, but producing different length sequences.
+- The three folders `1ch, 3ch, 12ch` contain a 1-frame, 3-frame and 12-frame model respectively, all with the same unet architecture, but producing different length sequences.
+- The three folders `1ch-dit, 3ch-dit, 7ch-dit` are similarly named, with a Diffusion Transformer (DiT) model, instead of a UNet.
 
 ---
 
@@ -18,6 +19,11 @@ from zea.models.flow_matching import FlowMatchingModel
 model_1ch = FlowMatchingModel.from_preset("flowmatching-echonetlvh")
 model_3ch = FlowMatchingModel.from_preset("flowmatching-echonetlvh-3ch")
 model_12ch = FlowMatchingModel.from_preset("flowmatching-echonetlvh-12ch")
+
+model_1ch = FlowMatchingModel.from_preset("flowmatching-echonetlvh-dit")
+model_3ch = FlowMatchingModel.from_preset("flowmatching-echonetlvh-dit-3ch")
+model_12ch = FlowMatchingModel.from_preset("flowmatching-echonetlvh-dit-12ch")
 ```
 
 
+