[Feature] Pure python config

diffengine/configs/__init__.py

@@ -3,7 +3,7 @@
 import os
 
 
-def get_cfgs_name_path():
+def get_cfgs_name_path() -> dict:
     path = os.path.dirname(__file__)
     mapping = {}
     for root, _, files in os.walk(path):

diffengine/configs/_base_/datasets/cartoonization_xl.py

@@ -1,33 +1,51 @@
+import torchvision
+from mmengine.dataset import DefaultSampler
+
+from diffengine.datasets import HFControlNetDataset
+from diffengine.datasets.transforms import (
+    ComputeTimeIds,
+    DumpImage,
+    PackInputs,
+    RandomCrop,
+    RandomHorizontalFlip,
+    SaveImageShape,
+    TorchVisonTransformWrapper,
+)
+from diffengine.engine.hooks import SDCheckpointHook, VisualizationHook
+
 train_pipeline = [
-    dict(type="SaveImageShape"),
+    dict(type=SaveImageShape),
     dict(
-        type="torchvision/Resize",
+        type=TorchVisonTransformWrapper,
+        transform=torchvision.transforms.Resize,
         size=768,
         interpolation="bilinear",
         keys=["img", "condition_img"]),
-    dict(type="RandomCrop", size=768, keys=["img", "condition_img"],
+    dict(type=RandomCrop, size=768, keys=["img", "condition_img"],
          force_same_size=False),
-    dict(type="RandomHorizontalFlip", p=0.5, keys=["img", "condition_img"]),
-    dict(type="ComputeTimeIds"),
-    dict(type="torchvision/ToTensor", keys=["img", "condition_img"]),
-    dict(type="DumpImage", max_imgs=10, dump_dir="work_dirs/dump"),
-    dict(type="torchvision/Normalize", mean=[0.5], std=[0.5],
+    dict(type=RandomHorizontalFlip, p=0.5, keys=["img", "condition_img"]),
+    dict(type=ComputeTimeIds),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.ToTensor, keys=["img", "condition_img"]),
+    dict(type=DumpImage, max_imgs=10, dump_dir="work_dirs/dump"),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.Normalize, mean=[0.5], std=[0.5],
         keys=["img", "condition_img"]),
     dict(
-        type="PackInputs",
+        type=PackInputs,
         input_keys=["img", "condition_img", "text", "time_ids"]),
 ]
 train_dataloader = dict(
     batch_size=1,
     num_workers=4,
     dataset=dict(
-        type="HFControlNetDataset",
+        type=HFControlNetDataset,
         dataset="instruction-tuning-sd/cartoonization",
         image_column="cartoonized_image",
         condition_column="original_image",
         caption_column="edit_prompt",
         pipeline=train_pipeline),
-    sampler=dict(type="DefaultSampler", shuffle=True),
+    sampler=dict(type=DefaultSampler, shuffle=True),
 )
 
 val_dataloader = None
@@ -37,12 +55,12 @@
 
 custom_hooks = [
     dict(
-        type="VisualizationHook",
+        type=VisualizationHook,
         prompt=["Generate a cartoonized version of the natural image"] * 4,
         condition_image=[
             'https://hf.co/datasets/diffusers/diffusers-images-docs/resolve/main/mountain.png'  # noqa
         ] * 4,
         height=768,
         width=768),
-    dict(type="SDCheckpointHook"),
+    dict(type=SDCheckpointHook),
 ]

diffengine/configs/_base_/datasets/cat_waterpainting_dreambooth_pixart_1024.py

@@ -1,25 +1,45 @@
+import torchvision
+from mmengine.dataset import InfiniteSampler
+
+from diffengine.datasets import HFDreamBoothDataset
+from diffengine.datasets.transforms import (
+    ComputePixArtImgInfo,
+    DumpImage,
+    PackInputs,
+    RandomCrop,
+    RandomHorizontalFlip,
+    SaveImageShape,
+    T5TextPreprocess,
+    TorchVisonTransformWrapper,
+)
+from diffengine.engine.hooks import PeftSaveHook, VisualizationHook
+
 train_pipeline = [
-    dict(type="SaveImageShape"),
-    dict(type="torchvision/Resize", size=1024, interpolation="bilinear"),
-    dict(type="RandomCrop", size=1024),
-    dict(type="RandomHorizontalFlip", p=0.5),
-    dict(type="ComputePixArtImgInfo"),
-    dict(type="torchvision/ToTensor"),
-    dict(type="DumpImage", max_imgs=5, dump_dir="work_dirs/dump"),
-    dict(type="torchvision/Normalize", mean=[0.5], std=[0.5]),
-    dict(type="T5TextPreprocess"),
-    dict(type="PackInputs", input_keys=["img", "text", "resolution", "aspect_ratio"]),
+    dict(type=SaveImageShape),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.Resize,
+         size=1024, interpolation="bilinear"),
+    dict(type=RandomCrop, size=1024),
+    dict(type=RandomHorizontalFlip, p=0.5),
+    dict(type=ComputePixArtImgInfo),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.ToTensor),
+    dict(type=DumpImage, max_imgs=5, dump_dir="work_dirs/dump"),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.Normalize, mean=[0.5], std=[0.5]),
+    dict(type=T5TextPreprocess),
+    dict(type=PackInputs, input_keys=["img", "text", "resolution", "aspect_ratio"]),
 ]
 train_dataloader = dict(
     batch_size=1,
     num_workers=4,
     dataset=dict(
-        type="HFDreamBoothDataset",
+        type=HFDreamBoothDataset,
         dataset="data/cat_waterpainting",
         instance_prompt="A cat in szn style",
         pipeline=train_pipeline,
         class_prompt=None),
-    sampler=dict(type="InfiniteSampler", shuffle=True),
+    sampler=dict(type=InfiniteSampler, shuffle=True),
 )
 
 val_dataloader = None
@@ -29,11 +49,11 @@
 
 custom_hooks = [
     dict(
-        type="VisualizationHook",
+        type=VisualizationHook,
         prompt=["A man in szn style"] * 4,
         by_epoch=False,
         interval=100,
         height=1024,
         width=1024),
-    dict(type="PeftSaveHook"),
+    dict(type=PeftSaveHook),
 ]

diffengine/configs/_base_/datasets/cat_waterpainting_dreambooth_xl.py

@@ -1,24 +1,43 @@
+import torchvision
+from mmengine.dataset import InfiniteSampler
+
+from diffengine.datasets import HFDreamBoothDataset
+from diffengine.datasets.transforms import (
+    ComputeTimeIds,
+    DumpImage,
+    PackInputs,
+    RandomCrop,
+    RandomHorizontalFlip,
+    SaveImageShape,
+    TorchVisonTransformWrapper,
+)
+from diffengine.engine.hooks import PeftSaveHook, VisualizationHook
+
 train_pipeline = [
-    dict(type="SaveImageShape"),
-    dict(type="torchvision/Resize", size=1024, interpolation="bilinear"),
-    dict(type="RandomCrop", size=1024),
-    dict(type="RandomHorizontalFlip", p=0.5),
-    dict(type="ComputeTimeIds"),
-    dict(type="torchvision/ToTensor"),
-    dict(type="DumpImage", max_imgs=5, dump_dir="work_dirs/dump"),
-    dict(type="torchvision/Normalize", mean=[0.5], std=[0.5]),
-    dict(type="PackInputs", input_keys=["img", "text", "time_ids"]),
+    dict(type=SaveImageShape),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.Resize,
+         size=1024, interpolation="bilinear"),
+    dict(type=RandomCrop, size=1024),
+    dict(type=RandomHorizontalFlip, p=0.5),
+    dict(type=ComputeTimeIds),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.ToTensor),
+    dict(type=DumpImage, max_imgs=5, dump_dir="work_dirs/dump"),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.Normalize, mean=[0.5], std=[0.5]),
+    dict(type=PackInputs, input_keys=["img", "text", "time_ids"]),
 ]
 train_dataloader = dict(
     batch_size=1,
     num_workers=4,
     dataset=dict(
-        type="HFDreamBoothDataset",
+        type=HFDreamBoothDataset,
         dataset="data/cat_waterpainting",
         instance_prompt="A cat in szn style",
         pipeline=train_pipeline,
         class_prompt=None),
-    sampler=dict(type="InfiniteSampler", shuffle=True),
+    sampler=dict(type=InfiniteSampler, shuffle=True),
 )
 
 val_dataloader = None
@@ -28,11 +47,11 @@
 
 custom_hooks = [
     dict(
-        type="VisualizationHook",
+        type=VisualizationHook,
         prompt=["A man in szn style"] * 4,
         by_epoch=False,
         interval=100,
         height=1024,
         width=1024),
-    dict(type="PeftSaveHook"),
+    dict(type=PeftSaveHook),
 ]

diffengine/configs/_base_/datasets/dog_dreambooth.py

@@ -1,21 +1,37 @@
+import torchvision
+from mmengine.dataset import InfiniteSampler
+
+from diffengine.datasets import HFDreamBoothDataset
+from diffengine.datasets.transforms import (
+    PackInputs,
+    RandomCrop,
+    RandomHorizontalFlip,
+    TorchVisonTransformWrapper,
+)
+from diffengine.engine.hooks import PeftSaveHook, VisualizationHook
+
 train_pipeline = [
-    dict(type="torchvision/Resize", size=512, interpolation="bilinear"),
-    dict(type="RandomCrop", size=512),
-    dict(type="RandomHorizontalFlip", p=0.5),
-    dict(type="torchvision/ToTensor"),
-    dict(type="torchvision/Normalize", mean=[0.5], std=[0.5]),
-    dict(type="PackInputs"),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.Resize,
+         size=512, interpolation="bilinear"),
+    dict(type=RandomCrop, size=512),
+    dict(type=RandomHorizontalFlip, p=0.5),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.ToTensor),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.Normalize, mean=[0.5], std=[0.5]),
+    dict(type=PackInputs),
 ]
 train_dataloader = dict(
     batch_size=4,
     num_workers=4,
     dataset=dict(
-        type="HFDreamBoothDataset",
+        type=HFDreamBoothDataset,
         dataset="diffusers/dog-example",
         instance_prompt="a photo of sks dog",
         pipeline=train_pipeline,
         class_prompt="a photo of dog"),
-    sampler=dict(type="InfiniteSampler", shuffle=True),
+    sampler=dict(type=InfiniteSampler, shuffle=True),
 )
 
 val_dataloader = None
@@ -25,9 +41,9 @@
 
 custom_hooks = [
     dict(
-        type="VisualizationHook",
+        type=VisualizationHook,
         prompt=["A photo of sks dog in a bucket"] * 4,
         by_epoch=False,
         interval=100),
-    dict(type="PeftSaveHook"),
+    dict(type=PeftSaveHook),
 ]

diffengine/configs/_base_/datasets/dog_dreambooth_if.py

@@ -1,20 +1,36 @@
+import torchvision
+from mmengine.dataset import InfiniteSampler
+
+from diffengine.datasets import HFDreamBoothDataset
+from diffengine.datasets.transforms import (
+    PackInputs,
+    RandomCrop,
+    RandomHorizontalFlip,
+    TorchVisonTransformWrapper,
+)
+from diffengine.engine.hooks import PeftSaveHook, VisualizationHook
+
 train_pipeline = [
-    dict(type="torchvision/Resize", size=64, interpolation="bilinear"),
-    dict(type="RandomCrop", size=64),
-    dict(type="RandomHorizontalFlip", p=0.5),
-    dict(type="torchvision/ToTensor"),
-    dict(type="torchvision/Normalize", mean=[0.5], std=[0.5]),
-    dict(type="PackInputs"),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.Resize,
+         size=64, interpolation="bilinear"),
+    dict(type=RandomCrop, size=64),
+    dict(type=RandomHorizontalFlip, p=0.5),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.ToTensor),
+    dict(type=TorchVisonTransformWrapper,
+         transform=torchvision.transforms.Normalize, mean=[0.5], std=[0.5]),
+    dict(type=PackInputs),
 ]
 train_dataloader = dict(
     batch_size=4,
     num_workers=4,
     dataset=dict(
-        type="HFDreamBoothDataset",
+        type=HFDreamBoothDataset,
         dataset="diffusers/dog-example",
         instance_prompt="a photo of sks dog",
         pipeline=train_pipeline),
-    sampler=dict(type="InfiniteSampler", shuffle=True),
+    sampler=dict(type=InfiniteSampler, shuffle=True),
 )
 
 val_dataloader = None
@@ -24,9 +40,9 @@
 
 custom_hooks = [
     dict(
-        type="VisualizationHook",
+        type=VisualizationHook,
         prompt=["A photo of sks dog in a bucket"] * 4,
         by_epoch=False,
         interval=100),
-    dict(type="PeftSaveHook"),
+    dict(type=PeftSaveHook),
 ]