fix imports

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Author: akoumpa

dfm/src/automodel/_diffusers/auto_diffusion_pipeline.py

@@ -19,7 +19,7 @@
 
 import torch
 import torch.nn as nn
-from Automodel.distributed.dfm_parallelizer import WanParallelizationStrategy
+from dfm.src.automodel.distributed.dfm_parallelizer import WanParallelizationStrategy
 from diffusers import DiffusionPipeline, WanPipeline
 from nemo_automodel.components.distributed import parallelizer
 from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager

dfm/src/automodel/datasets/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from Automodel.datasets.wan21 import (
+from dfm.src.automodel.datasets.wan21 import (
     MetaFilesDataset,
     build_node_parallel_sampler,
     build_wan21_dataloader,

dfm/src/automodel/flow_matching/training_step_t2v.py

@@ -19,7 +19,7 @@
 from typing import Dict, Tuple
 
 import torch
-from Automodel.flow_matching.time_shift_utils import (
+from dfm.src.automodel.flow_matching.time_shift_utils import (
     compute_density_for_timestep_sampling,
 )
 

dfm/src/automodel/recipes/train.py

@@ -22,8 +22,8 @@
 import torch
 import torch.distributed as dist
 import wandb
-from Automodel._diffusers.auto_diffusion_pipeline import NeMoWanPipeline
-from Automodel.flow_matching.training_step_t2v import (
+from dfm.src.automodel._diffusers.auto_diffusion_pipeline import NeMoWanPipeline
+from dfm.src.automodel.flow_matching.training_step_t2v import (
     step_fsdp_transformer_t2v,
 )
 from nemo_automodel.components.checkpoint.checkpointing import Checkpointer, CheckpointingConfig

examples/automodel/generate/wan_generate.py

@@ -18,7 +18,7 @@
 
 import torch
 import torch.distributed as dist
-from Automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline
+from dfm.src.automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline
 from diffusers import AutoencoderKLWan
 from diffusers.utils import export_to_video
 from nemo_automodel.components.distributed.init_utils import initialize_distributed