Nerogar · dxqb · Jun 15, 2026 · Jun 4, 2026
diff --git a/modules/dataLoader/ChromaBaseDataLoader.py b/modules/dataLoader/ChromaBaseDataLoader.py
@@ -26,6 +26,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.CHROMA_1)
 class ChromaBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -154,5 +155,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-factory.register(BaseDataLoader, ChromaBaseDataLoader, ModelType.CHROMA_1)
diff --git a/modules/dataLoader/ErnieBaseDataLoader.py b/modules/dataLoader/ErnieBaseDataLoader.py
@@ -22,6 +22,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.ERNIE)
 class ErnieBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -133,6 +134,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-
-factory.register(BaseDataLoader, ErnieBaseDataLoader, ModelType.ERNIE)
diff --git a/modules/dataLoader/Flux2BaseDataLoader.py b/modules/dataLoader/Flux2BaseDataLoader.py
@@ -30,6 +30,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.FLUX_2)
 class Flux2BaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -161,6 +162,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-
-factory.register(BaseDataLoader, Flux2BaseDataLoader, ModelType.FLUX_2)
diff --git a/modules/dataLoader/FluxBaseDataLoader.py b/modules/dataLoader/FluxBaseDataLoader.py
@@ -26,6 +26,8 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.FLUX_DEV_1)
+@factory.register(BaseDataLoader, ModelType.FLUX_FILL_DEV_1)
 class FluxBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -181,6 +183,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-factory.register(BaseDataLoader, FluxBaseDataLoader, ModelType.FLUX_DEV_1)
-factory.register(BaseDataLoader, FluxBaseDataLoader, ModelType.FLUX_FILL_DEV_1)
diff --git a/modules/dataLoader/HiDreamBaseDataLoader.py b/modules/dataLoader/HiDreamBaseDataLoader.py
@@ -27,6 +27,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.HI_DREAM_FULL)
 class HiDreamBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -219,5 +220,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-factory.register(BaseDataLoader, HiDreamBaseDataLoader, ModelType.HI_DREAM_FULL)
diff --git a/modules/dataLoader/HunyuanVideoBaseDataLoader.py b/modules/dataLoader/HunyuanVideoBaseDataLoader.py
@@ -29,6 +29,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.HUNYUAN_VIDEO)
 class HunyuanVideoBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -183,5 +184,3 @@ def _create_dataset(
             allow_video_files=True,
             vae_frame_dim=True,
         )
-
-factory.register(BaseDataLoader, HunyuanVideoBaseDataLoader, ModelType.HUNYUAN_VIDEO)
diff --git a/modules/dataLoader/PixArtAlphaBaseDataLoader.py b/modules/dataLoader/PixArtAlphaBaseDataLoader.py
@@ -24,6 +24,8 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.PIXART_ALPHA)
+@factory.register(BaseDataLoader, ModelType.PIXART_SIGMA)
 class PixArtAlphaBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -159,6 +161,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=16,
         )
-
-factory.register(BaseDataLoader, PixArtAlphaBaseDataLoader, ModelType.PIXART_ALPHA)
-factory.register(BaseDataLoader, PixArtAlphaBaseDataLoader, ModelType.PIXART_SIGMA)
diff --git a/modules/dataLoader/QwenBaseDataLoader.py b/modules/dataLoader/QwenBaseDataLoader.py
@@ -30,6 +30,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.QWEN)
 class QwenBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -166,5 +167,3 @@ def _create_dataset(
             allow_video_files=False, #don't allow video files, but...
             vae_frame_dim=True,  #...Qwen has a video-capable VAE. convert images to video dimensions
         )
-
-factory.register(BaseDataLoader, QwenBaseDataLoader, ModelType.QWEN)
diff --git a/modules/dataLoader/SanaBaseDataLoader.py b/modules/dataLoader/SanaBaseDataLoader.py
@@ -23,6 +23,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.SANA)
 class SanaBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -152,5 +153,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=32,
         )
-
-factory.register(BaseDataLoader, SanaBaseDataLoader, ModelType.SANA)
diff --git a/modules/dataLoader/StableDiffusion3BaseDataLoader.py b/modules/dataLoader/StableDiffusion3BaseDataLoader.py
@@ -25,6 +25,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_35)
 class StableDiffusion3BaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -198,5 +199,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-factory.register(BaseDataLoader, StableDiffusion3BaseDataLoader, ModelType.STABLE_DIFFUSION_35)
diff --git a/modules/dataLoader/StableDiffusionBaseDataLoader.py b/modules/dataLoader/StableDiffusionBaseDataLoader.py
@@ -24,6 +24,14 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_15)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_15_INPAINTING)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_BASE)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_INPAINTING)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_DEPTH)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_21)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_21_BASE)
 class StableDiffusionBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -164,12 +172,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=8,
         )
-
-factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_15)
-factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_15_INPAINTING)
-factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_20)
-factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_20_BASE)
-factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_20_INPAINTING)
-factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_20_DEPTH)
-factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_21)
-factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_21_BASE)
diff --git a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py
@@ -39,6 +39,14 @@
 import torch
 
 
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_15, TrainingMethod.FINE_TUNE_VAE)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.FINE_TUNE_VAE)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20, TrainingMethod.FINE_TUNE_VAE)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.FINE_TUNE_VAE)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.FINE_TUNE_VAE)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.FINE_TUNE_VAE)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_21, TrainingMethod.FINE_TUNE_VAE)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.FINE_TUNE_VAE)
 class StableDiffusionFineTuneVaeDataLoader(BaseDataLoader):
     def _setup_cache_device(
             self,
@@ -293,12 +301,3 @@ def _create_dataset(
             train_progress,
             is_validation,
         )
-
-factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_15, TrainingMethod.FINE_TUNE_VAE)
-factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.FINE_TUNE_VAE)
-factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_20, TrainingMethod.FINE_TUNE_VAE)
-factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.FINE_TUNE_VAE)
-factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.FINE_TUNE_VAE)
-factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.FINE_TUNE_VAE)
-factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_21, TrainingMethod.FINE_TUNE_VAE)
-factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.FINE_TUNE_VAE)
diff --git a/modules/dataLoader/StableDiffusionXLBaseDataLoader.py b/modules/dataLoader/StableDiffusionXLBaseDataLoader.py
@@ -24,6 +24,8 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_XL_10_BASE)
+@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING)
 class StableDiffusionXLBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -176,5 +178,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-factory.register(BaseDataLoader, StableDiffusionXLBaseDataLoader, ModelType.STABLE_DIFFUSION_XL_10_BASE)
-factory.register(BaseDataLoader, StableDiffusionXLBaseDataLoader, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING)
diff --git a/modules/dataLoader/WuerstchenBaseDataLoader.py b/modules/dataLoader/WuerstchenBaseDataLoader.py
@@ -23,6 +23,8 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.WUERSTCHEN_2)
+@factory.register(BaseDataLoader, ModelType.STABLE_CASCADE_1)
 class WuerstchenBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -156,6 +158,3 @@ def _create_dataset(
             aspect_bucketing_quantization=128,
             supports_inpainting=False,
         )
-
-factory.register(BaseDataLoader, WuerstchenBaseDataLoader, ModelType.WUERSTCHEN_2)
-factory.register(BaseDataLoader, WuerstchenBaseDataLoader, ModelType.STABLE_CASCADE_1)
diff --git a/modules/dataLoader/ZImageBaseDataLoader.py b/modules/dataLoader/ZImageBaseDataLoader.py
@@ -26,6 +26,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.Z_IMAGE)
 class ZImageBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -153,5 +154,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-factory.register(BaseDataLoader, ZImageBaseDataLoader, ModelType.Z_IMAGE)
diff --git a/modules/modelSampler/ChromaSampler.py b/modules/modelSampler/ChromaSampler.py
@@ -19,6 +19,7 @@
 from tqdm import tqdm
 
 
+@factory.register(BaseModelSampler, ModelType.CHROMA_1)
 class ChromaSampler(BaseModelSampler):
     def __init__(
             self,
@@ -189,5 +190,3 @@ def sample(
         )
 
         on_sample(sampler_output)
-
-factory.register(BaseModelSampler, ChromaSampler, ModelType.CHROMA_1)
diff --git a/modules/modelSampler/ErnieSampler.py b/modules/modelSampler/ErnieSampler.py
@@ -20,6 +20,7 @@
 from tqdm import tqdm
 
 
+@factory.register(BaseModelSampler, ModelType.ERNIE)
 class ErnieSampler(BaseModelSampler):
     def __init__(
             self,
@@ -162,6 +163,3 @@ def sample(
         )
 
         on_sample(sampler_output)
-
-
-factory.register(BaseModelSampler, ErnieSampler, ModelType.ERNIE)
diff --git a/modules/modelSampler/Flux2Sampler.py b/modules/modelSampler/Flux2Sampler.py
@@ -22,6 +22,7 @@
 from tqdm import tqdm
 
 
+@factory.register(BaseModelSampler, ModelType.FLUX_2)
 class Flux2Sampler(BaseModelSampler):
     def __init__(
             self,
@@ -190,5 +191,3 @@ def sample(
         )
 
         on_sample(sampler_output)
-
-factory.register(BaseModelSampler, Flux2Sampler, ModelType.FLUX_2)
diff --git a/modules/modelSampler/FluxSampler.py b/modules/modelSampler/FluxSampler.py
@@ -23,6 +23,8 @@
 from tqdm import tqdm
 
 
+@factory.register(BaseModelSampler, ModelType.FLUX_DEV_1)
+@factory.register(BaseModelSampler, ModelType.FLUX_FILL_DEV_1)
 class FluxSampler(BaseModelSampler):
     def __init__(
             self,
@@ -450,6 +452,3 @@ def sample(
         )
 
         on_sample(sampler_output)
-
-factory.register(BaseModelSampler, FluxSampler, ModelType.FLUX_DEV_1)
-factory.register(BaseModelSampler, FluxSampler, ModelType.FLUX_FILL_DEV_1)
diff --git a/modules/modelSampler/HiDreamSampler.py b/modules/modelSampler/HiDreamSampler.py
@@ -19,6 +19,7 @@
 from tqdm import tqdm
 
 
+@factory.register(BaseModelSampler, ModelType.HI_DREAM_FULL)
 class HiDreamSampler(BaseModelSampler):
     def __init__(
             self,
@@ -192,5 +193,3 @@ def sample(
         )
 
         on_sample(sampler_output)
-
-factory.register(BaseModelSampler, HiDreamSampler, ModelType.HI_DREAM_FULL)
diff --git a/modules/modelSampler/HunyuanVideoSampler.py b/modules/modelSampler/HunyuanVideoSampler.py
@@ -20,6 +20,7 @@
 from tqdm import tqdm
 
 
+@factory.register(BaseModelSampler, ModelType.HUNYUAN_VIDEO)
 class HunyuanVideoSampler(BaseModelSampler):
     def __init__(
             self,
@@ -209,5 +210,3 @@ def sample(
         )
 
         on_sample(sampler_output)
-
-factory.register(BaseModelSampler, HunyuanVideoSampler, ModelType.HUNYUAN_VIDEO)
diff --git a/modules/modelSampler/PixArtAlphaSampler.py b/modules/modelSampler/PixArtAlphaSampler.py
@@ -18,6 +18,8 @@
 from tqdm import tqdm
 
 
+@factory.register(BaseModelSampler, ModelType.PIXART_ALPHA)
+@factory.register(BaseModelSampler, ModelType.PIXART_SIGMA)
 class PixArtAlphaSampler(BaseModelSampler):
     def __init__(
             self,
@@ -191,6 +193,3 @@ def sample(
         )
 
         on_sample(sampler_output)
-
-factory.register(BaseModelSampler, PixArtAlphaSampler, ModelType.PIXART_ALPHA)
-factory.register(BaseModelSampler, PixArtAlphaSampler, ModelType.PIXART_SIGMA)
diff --git a/modules/modelSampler/QwenSampler.py b/modules/modelSampler/QwenSampler.py
@@ -20,6 +20,7 @@
 from tqdm import tqdm
 
 
+@factory.register(BaseModelSampler, ModelType.QWEN)
 class QwenSampler(BaseModelSampler):
     def __init__(
             self,
@@ -188,5 +189,3 @@ def sample(
         )
 
         on_sample(sampler_output)
-
-factory.register(BaseModelSampler, QwenSampler, ModelType.QWEN)
diff --git a/modules/modelSampler/SanaSampler.py b/modules/modelSampler/SanaSampler.py
@@ -19,6 +19,7 @@
 from tqdm import tqdm
 
 
+@factory.register(BaseModelSampler, ModelType.SANA)
 class SanaSampler(BaseModelSampler):
     def __init__(
             self,
@@ -177,5 +178,3 @@ def sample(
         )
 
         on_sample(sampler_output)
-
-factory.register(BaseModelSampler, SanaSampler, ModelType.SANA)
diff --git a/modules/modelSampler/StableDiffusion3Sampler.py b/modules/modelSampler/StableDiffusion3Sampler.py
@@ -19,6 +19,8 @@
 from tqdm import tqdm
 
 
+@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_3)
+@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_35)
 class StableDiffusion3Sampler(BaseModelSampler):
     def __init__(
             self,
@@ -191,6 +193,3 @@ def sample(
         )
 
         on_sample(sampler_output)
-
-factory.register(BaseModelSampler, StableDiffusion3Sampler, ModelType.STABLE_DIFFUSION_3)
-factory.register(BaseModelSampler, StableDiffusion3Sampler, ModelType.STABLE_DIFFUSION_35)