feat: Add support for flux (#94)

* feat: Add DDP support in the Automodel path

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Fix linting issues

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* add cicd nightly tests

Signed-off-by: linnan wang <linnanw@nvidia.com>

* feat: Flux training support

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Formatting changes

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Changing import order

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Linting fix

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Fixing imports

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Init fix

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Code fixes

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Import order fix

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Fix linting

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Fix broken unit tests

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Fix lint issue

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Fix failing GPU tests

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Ensure flow matching pipeline is backward compatible for megatron path

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Update secrets baseline

Signed-off-by: Charlie Truong <chtruong@nvidia.com>

---------

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
Signed-off-by: linnan wang <linnanw@nvidia.com>
Signed-off-by: Charlie Truong <chtruong@nvidia.com>
Co-authored-by: linnan wang <linnanw@nvidia.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>

Author: 3 people

.github/workflows/config/.secrets.baseline

@@ -139,10 +139,10 @@
         "filename": "examples/automodel/pretrain/cicd/wan21_cicd_nightly_video.yaml",
         "hashed_secret": "c70f071570ba65f9c4079d6051e955ff4f802eea",
         "is_verified": false,
-        "line_number": 61,
+        "line_number": 67,
         "is_secret": false
       }
     ]
   },
-  "generated_at": "2026-01-29T18:49:17Z"
+  "generated_at": "2026-01-30T18:50:34Z"
 }

dfm/src/automodel/_diffusers/__init__.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .auto_diffusion_pipeline import NeMoAutoDiffusionPipeline
+from .auto_diffusion_pipeline import NeMoAutoDiffusionPipeline, PipelineSpec
 
 
 __all__ = [
     "NeMoAutoDiffusionPipeline",
+    "PipelineSpec",
 ]