Migrated files batch 32

diffsynth/tokenizer_configs/stable_diffusion_3/tokenizer_1/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

diffsynth/tokenizer_configs/stable_diffusion_3/tokenizer_1/tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "49406": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

diffsynth/tokenizer_configs/stable_diffusion_3/tokenizer_2/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "!",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

diffsynth/tokenizer_configs/stable_diffusion_3/tokenizer_2/tokenizer_config.json

@@ -0,0 +1,38 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "!",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49406": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "!",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

diffsynth/tokenizer_configs/stable_diffusion_3/tokenizer_3/special_tokens_map.json

@@ -0,0 +1,125 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

diffsynth/tokenizer_configs/stable_diffusion_3/tokenizer_3/spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d60acb128cf7b7f2536e8f38a5b18a05535c9e14c7a355904270e15b0945ea86
+size 791656

diffsynth/tokenizer_configs/stable_diffusion_3/tokenizer_3/tokenizer_config.json

@@ -0,0 +1,940 @@
+{
+  "add_prefix_space": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32000": {
+      "content": "<extra_id_99>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32001": {
+      "content": "<extra_id_98>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32002": {
+      "content": "<extra_id_97>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32003": {
+      "content": "<extra_id_96>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32004": {
+      "content": "<extra_id_95>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32005": {
+      "content": "<extra_id_94>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32006": {
+      "content": "<extra_id_93>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32007": {
+      "content": "<extra_id_92>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32008": {
+      "content": "<extra_id_91>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32009": {
+      "content": "<extra_id_90>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32010": {
+      "content": "<extra_id_89>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32011": {
+      "content": "<extra_id_88>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32012": {
+      "content": "<extra_id_87>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32013": {
+      "content": "<extra_id_86>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32014": {
+      "content": "<extra_id_85>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32015": {
+      "content": "<extra_id_84>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32016": {
+      "content": "<extra_id_83>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32017": {
+      "content": "<extra_id_82>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32018": {
+      "content": "<extra_id_81>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32019": {
+      "content": "<extra_id_80>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32020": {
+      "content": "<extra_id_79>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32021": {
+      "content": "<extra_id_78>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32022": {
+      "content": "<extra_id_77>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32023": {
+      "content": "<extra_id_76>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32024": {
+      "content": "<extra_id_75>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32025": {
+      "content": "<extra_id_74>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32026": {
+      "content": "<extra_id_73>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32027": {
+      "content": "<extra_id_72>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32028": {
+      "content": "<extra_id_71>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32029": {
+      "content": "<extra_id_70>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32030": {
+      "content": "<extra_id_69>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32031": {
+      "content": "<extra_id_68>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32032": {
+      "content": "<extra_id_67>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32033": {
+      "content": "<extra_id_66>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32034": {
+      "content": "<extra_id_65>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32035": {
+      "content": "<extra_id_64>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32036": {
+      "content": "<extra_id_63>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32037": {
+      "content": "<extra_id_62>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32038": {
+      "content": "<extra_id_61>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32039": {
+      "content": "<extra_id_60>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32040": {
+      "content": "<extra_id_59>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32041": {
+      "content": "<extra_id_58>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32042": {
+      "content": "<extra_id_57>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32043": {
+      "content": "<extra_id_56>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32044": {
+      "content": "<extra_id_55>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32045": {
+      "content": "<extra_id_54>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32046": {
+      "content": "<extra_id_53>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32047": {
+      "content": "<extra_id_52>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32048": {
+      "content": "<extra_id_51>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32049": {
+      "content": "<extra_id_50>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32050": {
+      "content": "<extra_id_49>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32051": {
+      "content": "<extra_id_48>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32052": {
+      "content": "<extra_id_47>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32053": {
+      "content": "<extra_id_46>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32054": {
+      "content": "<extra_id_45>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32055": {
+      "content": "<extra_id_44>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32056": {
+      "content": "<extra_id_43>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32057": {
+      "content": "<extra_id_42>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32058": {
+      "content": "<extra_id_41>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32059": {
+      "content": "<extra_id_40>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32060": {
+      "content": "<extra_id_39>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32061": {
+      "content": "<extra_id_38>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32062": {
+      "content": "<extra_id_37>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32063": {
+      "content": "<extra_id_36>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32064": {
+      "content": "<extra_id_35>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32065": {
+      "content": "<extra_id_34>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32066": {
+      "content": "<extra_id_33>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32067": {
+      "content": "<extra_id_32>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32068": {
+      "content": "<extra_id_31>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32069": {
+      "content": "<extra_id_30>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32070": {
+      "content": "<extra_id_29>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32071": {
+      "content": "<extra_id_28>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32072": {
+      "content": "<extra_id_27>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32073": {
+      "content": "<extra_id_26>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32074": {
+      "content": "<extra_id_25>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32075": {
+      "content": "<extra_id_24>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32076": {
+      "content": "<extra_id_23>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32077": {
+      "content": "<extra_id_22>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32078": {
+      "content": "<extra_id_21>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32079": {
+      "content": "<extra_id_20>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32080": {
+      "content": "<extra_id_19>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32081": {
+      "content": "<extra_id_18>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32082": {
+      "content": "<extra_id_17>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32083": {
+      "content": "<extra_id_16>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32084": {
+      "content": "<extra_id_15>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32085": {
+      "content": "<extra_id_14>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32086": {
+      "content": "<extra_id_13>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32087": {
+      "content": "<extra_id_12>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32088": {
+      "content": "<extra_id_11>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32089": {
+      "content": "<extra_id_10>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32090": {
+      "content": "<extra_id_9>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32091": {
+      "content": "<extra_id_8>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32092": {
+      "content": "<extra_id_7>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32093": {
+      "content": "<extra_id_6>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32094": {
+      "content": "<extra_id_5>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32095": {
+      "content": "<extra_id_4>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32096": {
+      "content": "<extra_id_3>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32097": {
+      "content": "<extra_id_2>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32098": {
+      "content": "<extra_id_1>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32099": {
+      "content": "<extra_id_0>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "legacy": true,
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

diffsynth/tokenizer_configs/stable_diffusion_xl/tokenizer_2/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+    "bos_token": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    "eos_token": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    "pad_token": "!",
+    "unk_token": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    }
+}

diffsynth/tokenizer_configs/stable_diffusion_xl/tokenizer_2/tokenizer_config.json

@@ -0,0 +1,38 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "!",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49406": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "!",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

diffsynth/trainers/text_to_image.py

@@ -0,0 +1,318 @@
+import lightning as pl
+from peft import LoraConfig, inject_adapter_in_model
+import torch, os
+from ..data.simple_text_image import TextImageDataset
+from modelscope.hub.api import HubApi
+from ..models.utils import load_state_dict
+
+
+
+class LightningModelForT2ILoRA(pl.LightningModule):
+    def __init__(
+        self,
+        learning_rate=1e-4,
+        use_gradient_checkpointing=True,
+        state_dict_converter=None,
+    ):
+        super().__init__()
+        # Set parameters
+        self.learning_rate = learning_rate
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+        self.state_dict_converter = state_dict_converter
+        self.lora_alpha = None
+
+
+    def load_models(self):
+        # This function is implemented in other modules
+        self.pipe = None
+
+
+    def freeze_parameters(self):
+        # Freeze parameters
+        self.pipe.requires_grad_(False)
+        self.pipe.eval()
+        self.pipe.denoising_model().train()
+
+    
+    def add_lora_to_model(self, model, lora_rank=4, lora_alpha=4, lora_target_modules="to_q,to_k,to_v,to_out", init_lora_weights="gaussian", pretrained_lora_path=None, state_dict_converter=None):
+        # Add LoRA to UNet
+        self.lora_alpha = lora_alpha
+        if init_lora_weights == "kaiming":
+            init_lora_weights = True
+            
+        lora_config = LoraConfig(
+            r=lora_rank,
+            lora_alpha=lora_alpha,
+            init_lora_weights=init_lora_weights,
+            target_modules=lora_target_modules.split(","),
+        )
+        model = inject_adapter_in_model(lora_config, model)
+        for param in model.parameters():
+            # Upcast LoRA parameters into fp32
+            if param.requires_grad:
+                param.data = param.to(torch.float32)
+
+        # Lora pretrained lora weights
+        if pretrained_lora_path is not None:
+            state_dict = load_state_dict(pretrained_lora_path)
+            if state_dict_converter is not None:
+                state_dict = state_dict_converter(state_dict)
+            missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+            all_keys = [i for i, _ in model.named_parameters()]
+            num_updated_keys = len(all_keys) - len(missing_keys)
+            num_unexpected_keys = len(unexpected_keys)
+            print(f"{num_updated_keys} parameters are loaded from {pretrained_lora_path}. {num_unexpected_keys} parameters are unexpected.")
+
+
+    def training_step(self, batch, batch_idx):
+        # Data
+        text, image = batch["text"], batch["image"]
+
+        # Prepare input parameters
+        self.pipe.device = self.device
+        prompt_emb = self.pipe.encode_prompt(text, positive=True)
+        if "latents" in batch:
+            latents = batch["latents"].to(dtype=self.pipe.torch_dtype, device=self.device)
+        else:
+            latents = self.pipe.vae_encoder(image.to(dtype=self.pipe.torch_dtype, device=self.device))
+        noise = torch.randn_like(latents)
+        timestep_id = torch.randint(0, self.pipe.scheduler.num_train_timesteps, (1,))
+        timestep = self.pipe.scheduler.timesteps[timestep_id].to(self.device)
+        extra_input = self.pipe.prepare_extra_input(latents)
+        noisy_latents = self.pipe.scheduler.add_noise(latents, noise, timestep)
+        training_target = self.pipe.scheduler.training_target(latents, noise, timestep)
+
+        # Compute loss
+        noise_pred = self.pipe.denoising_model()(
+            noisy_latents, timestep=timestep, **prompt_emb, **extra_input,
+            use_gradient_checkpointing=self.use_gradient_checkpointing
+        )
+        loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float())
+        loss = loss * self.pipe.scheduler.training_weight(timestep)
+
+        # Record log
+        self.log("train_loss", loss, prog_bar=True)
+        return loss
+
+
+    def configure_optimizers(self):
+        trainable_modules = filter(lambda p: p.requires_grad, self.pipe.denoising_model().parameters())
+        optimizer = torch.optim.AdamW(trainable_modules, lr=self.learning_rate)
+        return optimizer
+    
+
+    def on_save_checkpoint(self, checkpoint):
+        checkpoint.clear()
+        trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.pipe.denoising_model().named_parameters()))
+        trainable_param_names = set([named_param[0] for named_param in trainable_param_names])
+        state_dict = self.pipe.denoising_model().state_dict()
+        lora_state_dict = {}
+        for name, param in state_dict.items():
+            if name in trainable_param_names:
+                lora_state_dict[name] = param
+        if self.state_dict_converter is not None:
+            lora_state_dict = self.state_dict_converter(lora_state_dict, alpha=self.lora_alpha)
+        checkpoint.update(lora_state_dict)
+
+
+
+def add_general_parsers(parser):
+    parser.add_argument(
+        "--dataset_path",
+        type=str,
+        default=None,
+        required=True,
+        help="The path of the Dataset.",
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        default="./",
+        help="Path to save the model.",
+    )
+    parser.add_argument(
+        "--steps_per_epoch",
+        type=int,
+        default=500,
+        help="Number of steps per epoch.",
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=1024,
+        help="Image height.",
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=1024,
+        help="Image width.",
+    )
+    parser.add_argument(
+        "--center_crop",
+        default=False,
+        action="store_true",
+        help=(
+            "Whether to center crop the input images to the resolution. If not set, the images will be randomly"
+            " cropped. The images will be resized to the resolution first before cropping."
+        ),
+    )
+    parser.add_argument(
+        "--random_flip",
+        default=False,
+        action="store_true",
+        help="Whether to randomly flip images horizontally",
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=1,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+    )
+    parser.add_argument(
+        "--precision",
+        type=str,
+        default="16-mixed",
+        choices=["32", "16", "16-mixed", "bf16"],
+        help="Training precision",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-4,
+        help="Learning rate.",
+    )
+    parser.add_argument(
+        "--lora_rank",
+        type=int,
+        default=4,
+        help="The dimension of the LoRA update matrices.",
+    )
+    parser.add_argument(
+        "--lora_alpha",
+        type=float,
+        default=4.0,
+        help="The weight of the LoRA update matrices.",
+    )
+    parser.add_argument(
+        "--init_lora_weights",
+        type=str,
+        default="kaiming",
+        choices=["gaussian", "kaiming"],
+        help="The initializing method of LoRA weight.",
+    )
+    parser.add_argument(
+        "--use_gradient_checkpointing",
+        default=False,
+        action="store_true",
+        help="Whether to use gradient checkpointing.",
+    )
+    parser.add_argument(
+        "--accumulate_grad_batches",
+        type=int,
+        default=1,
+        help="The number of batches in gradient accumulation.",
+    )
+    parser.add_argument(
+        "--training_strategy",
+        type=str,
+        default="auto",
+        choices=["auto", "deepspeed_stage_1", "deepspeed_stage_2", "deepspeed_stage_3"],
+        help="Training strategy",
+    )
+    parser.add_argument(
+        "--max_epochs",
+        type=int,
+        default=1,
+        help="Number of epochs.",
+    )
+    parser.add_argument(
+        "--modelscope_model_id",
+        type=str,
+        default=None,
+        help="Model ID on ModelScope (https://www.modelscope.cn/). The model will be uploaded to ModelScope automatically if you provide a Model ID.",
+    )
+    parser.add_argument(
+        "--modelscope_access_token",
+        type=str,
+        default=None,
+        help="Access key on ModelScope (https://www.modelscope.cn/). Required if you want to upload the model to ModelScope.",
+    )
+    parser.add_argument(
+        "--pretrained_lora_path",
+        type=str,
+        default=None,
+        help="Pretrained LoRA path. Required if the training is resumed.",
+    )
+    parser.add_argument(
+        "--use_swanlab",
+        default=False,
+        action="store_true",
+        help="Whether to use SwanLab logger.",
+    )
+    parser.add_argument(
+        "--swanlab_mode",
+        default=None,
+        help="SwanLab mode (cloud or local).",
+    )
+    return parser
+
+
+def launch_training_task(model, args):
+    # dataset and data loader
+    dataset = TextImageDataset(
+        args.dataset_path,
+        steps_per_epoch=args.steps_per_epoch * args.batch_size,
+        height=args.height,
+        width=args.width,
+        center_crop=args.center_crop,
+        random_flip=args.random_flip
+    )
+    train_loader = torch.utils.data.DataLoader(
+        dataset,
+        shuffle=True,
+        batch_size=args.batch_size,
+        num_workers=args.dataloader_num_workers
+    )
+    # train
+    if args.use_swanlab:
+        from swanlab.integration.pytorch_lightning import SwanLabLogger
+        swanlab_config = {"UPPERFRAMEWORK": "DiffSynth-Studio"}
+        swanlab_config.update(vars(args))
+        swanlab_logger = SwanLabLogger(
+            project="diffsynth_studio", 
+            name="diffsynth_studio",
+            config=swanlab_config,
+            mode=args.swanlab_mode,
+            logdir=os.path.join(args.output_path, "swanlog"),
+        )
+        logger = [swanlab_logger]
+    else:
+        logger = None
+    trainer = pl.Trainer(
+        max_epochs=args.max_epochs,
+        accelerator="gpu",
+        devices="auto",
+        precision=args.precision,
+        strategy=args.training_strategy,
+        default_root_dir=args.output_path,
+        accumulate_grad_batches=args.accumulate_grad_batches,
+        callbacks=[pl.pytorch.callbacks.ModelCheckpoint(save_top_k=-1)],
+        logger=logger,
+    )
+    trainer.fit(model=model, train_dataloaders=train_loader)
+
+    # Upload models
+    if args.modelscope_model_id is not None and args.modelscope_access_token is not None:
+        print(f"Uploading models to modelscope. model_id: {args.modelscope_model_id} local_path: {trainer.log_dir}")
+        with open(os.path.join(trainer.log_dir, "configuration.json"), "w", encoding="utf-8") as f:
+            f.write('{"framework":"Pytorch","task":"text-to-image-synthesis"}\n')
+        api = HubApi()
+        api.login(args.modelscope_access_token)
+        api.push_model(model_id=args.modelscope_model_id, model_dir=trainer.log_dir)

diffsynth/trainers/utils.py

@@ -0,0 +1,506 @@
+import imageio, os, torch, warnings, torchvision, argparse, json
+from peft import LoraConfig, inject_adapter_in_model
+from PIL import Image
+import pandas as pd
+from tqdm import tqdm
+from accelerate import Accelerator
+
+
+
+class ImageDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        base_path=None, metadata_path=None,
+        max_pixels=1920*1080, height=None, width=None,
+        height_division_factor=16, width_division_factor=16,
+        data_file_keys=("image",),
+        image_file_extension=("jpg", "jpeg", "png", "webp"),
+        repeat=1,
+        args=None,
+    ):
+        if args is not None:
+            base_path = args.dataset_base_path
+            metadata_path = args.dataset_metadata_path
+            height = args.height
+            width = args.width
+            max_pixels = args.max_pixels
+            data_file_keys = args.data_file_keys.split(",")
+            repeat = args.dataset_repeat
+            
+        self.base_path = base_path
+        self.max_pixels = max_pixels
+        self.height = height
+        self.width = width
+        self.height_division_factor = height_division_factor
+        self.width_division_factor = width_division_factor
+        self.data_file_keys = data_file_keys
+        self.image_file_extension = image_file_extension
+        self.repeat = repeat
+
+        if height is not None and width is not None:
+            print("Height and width are fixed. Setting `dynamic_resolution` to False.")
+            self.dynamic_resolution = False
+        elif height is None and width is None:
+            print("Height and width are none. Setting `dynamic_resolution` to True.")
+            self.dynamic_resolution = True
+            
+        if metadata_path is None:
+            print("No metadata. Trying to generate it.")
+            metadata = self.generate_metadata(base_path)
+            print(f"{len(metadata)} lines in metadata.")
+            self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))]
+        elif metadata_path.endswith(".json"):
+            with open(metadata_path, "r") as f:
+                metadata = json.load(f)
+            self.data = metadata
+        elif metadata_path.endswith(".jsonl"):
+            metadata = []
+            with open(metadata_path, 'r') as f:
+                for line in tqdm(f):
+                    metadata.append(json.loads(line.strip()))
+            self.data = metadata
+        else:
+            metadata = pd.read_csv(metadata_path)
+            self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))]
+
+
+    def generate_metadata(self, folder):
+        image_list, prompt_list = [], []
+        file_set = set(os.listdir(folder))
+        for file_name in file_set:
+            if "." not in file_name:
+                continue
+            file_ext_name = file_name.split(".")[-1].lower()
+            file_base_name = file_name[:-len(file_ext_name)-1]
+            if file_ext_name not in self.image_file_extension:
+                continue
+            prompt_file_name = file_base_name + ".txt"
+            if prompt_file_name not in file_set:
+                continue
+            with open(os.path.join(folder, prompt_file_name), "r", encoding="utf-8") as f:
+                prompt = f.read().strip()
+            image_list.append(file_name)
+            prompt_list.append(prompt)
+        metadata = pd.DataFrame()
+        metadata["image"] = image_list
+        metadata["prompt"] = prompt_list
+        return metadata
+    
+    
+    def crop_and_resize(self, image, target_height, target_width):
+        width, height = image.size
+        scale = max(target_width / width, target_height / height)
+        image = torchvision.transforms.functional.resize(
+            image,
+            (round(height*scale), round(width*scale)),
+            interpolation=torchvision.transforms.InterpolationMode.BILINEAR
+        )
+        image = torchvision.transforms.functional.center_crop(image, (target_height, target_width))
+        return image
+    
+    
+    def get_height_width(self, image):
+        if self.dynamic_resolution:
+            width, height = image.size
+            if width * height > self.max_pixels:
+                scale = (width * height / self.max_pixels) ** 0.5
+                height, width = int(height / scale), int(width / scale)
+            height = height // self.height_division_factor * self.height_division_factor
+            width = width // self.width_division_factor * self.width_division_factor
+        else:
+            height, width = self.height, self.width
+        return height, width
+    
+    
+    def load_image(self, file_path):
+        image = Image.open(file_path).convert("RGB")
+        image = self.crop_and_resize(image, *self.get_height_width(image))
+        return image
+    
+    
+    def load_data(self, file_path):
+        return self.load_image(file_path)
+
+
+    def __getitem__(self, data_id):
+        data = self.data[data_id % len(self.data)].copy()
+        for key in self.data_file_keys:
+            if key in data:
+                if isinstance(data[key], list):
+                    path = [os.path.join(self.base_path, p) for p in data[key]]
+                    data[key] = [self.load_data(p) for p in path]
+                else:
+                    path = os.path.join(self.base_path, data[key])
+                    data[key] = self.load_data(path)
+                if data[key] is None:
+                    warnings.warn(f"cannot load file {data[key]}.")
+                    return None
+        return data
+    
+
+    def __len__(self):
+        return len(self.data) * self.repeat
+
+
+
+class VideoDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        base_path=None, metadata_path=None,
+        num_frames=81,
+        time_division_factor=4, time_division_remainder=1,
+        max_pixels=1920*1080, height=None, width=None,
+        height_division_factor=16, width_division_factor=16,
+        data_file_keys=("video",),
+        image_file_extension=("jpg", "jpeg", "png", "webp"),
+        video_file_extension=("mp4", "avi", "mov", "wmv", "mkv", "flv", "webm"),
+        repeat=1,
+        args=None,
+    ):
+        if args is not None:
+            base_path = args.dataset_base_path
+            metadata_path = args.dataset_metadata_path
+            height = args.height
+            width = args.width
+            max_pixels = args.max_pixels
+            num_frames = args.num_frames
+            data_file_keys = args.data_file_keys.split(",")
+            repeat = args.dataset_repeat
+        
+        self.base_path = base_path
+        self.num_frames = num_frames
+        self.time_division_factor = time_division_factor
+        self.time_division_remainder = time_division_remainder
+        self.max_pixels = max_pixels
+        self.height = height
+        self.width = width
+        self.height_division_factor = height_division_factor
+        self.width_division_factor = width_division_factor
+        self.data_file_keys = data_file_keys
+        self.image_file_extension = image_file_extension
+        self.video_file_extension = video_file_extension
+        self.repeat = repeat
+        
+        if height is not None and width is not None:
+            print("Height and width are fixed. Setting `dynamic_resolution` to False.")
+            self.dynamic_resolution = False
+        elif height is None and width is None:
+            print("Height and width are none. Setting `dynamic_resolution` to True.")
+            self.dynamic_resolution = True
+            
+        if metadata_path is None:
+            print("No metadata. Trying to generate it.")
+            metadata = self.generate_metadata(base_path)
+            print(f"{len(metadata)} lines in metadata.")
+            self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))]
+        elif metadata_path.endswith(".json"):
+            with open(metadata_path, "r") as f:
+                metadata = json.load(f)
+            self.data = metadata
+        else:
+            metadata = pd.read_csv(metadata_path)
+            self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))]
+            
+    
+    def generate_metadata(self, folder):
+        video_list, prompt_list = [], []
+        file_set = set(os.listdir(folder))
+        for file_name in file_set:
+            if "." not in file_name:
+                continue
+            file_ext_name = file_name.split(".")[-1].lower()
+            file_base_name = file_name[:-len(file_ext_name)-1]
+            if file_ext_name not in self.image_file_extension and file_ext_name not in self.video_file_extension:
+                continue
+            prompt_file_name = file_base_name + ".txt"
+            if prompt_file_name not in file_set:
+                continue
+            with open(os.path.join(folder, prompt_file_name), "r", encoding="utf-8") as f:
+                prompt = f.read().strip()
+            video_list.append(file_name)
+            prompt_list.append(prompt)
+        metadata = pd.DataFrame()
+        metadata["video"] = video_list
+        metadata["prompt"] = prompt_list
+        return metadata
+        
+        
+    def crop_and_resize(self, image, target_height, target_width):
+        width, height = image.size
+        scale = max(target_width / width, target_height / height)
+        image = torchvision.transforms.functional.resize(
+            image,
+            (round(height*scale), round(width*scale)),
+            interpolation=torchvision.transforms.InterpolationMode.BILINEAR
+        )
+        image = torchvision.transforms.functional.center_crop(image, (target_height, target_width))
+        return image
+    
+    
+    def get_height_width(self, image):
+        if self.dynamic_resolution:
+            width, height = image.size
+            if width * height > self.max_pixels:
+                scale = (width * height / self.max_pixels) ** 0.5
+                height, width = int(height / scale), int(width / scale)
+            height = height // self.height_division_factor * self.height_division_factor
+            width = width // self.width_division_factor * self.width_division_factor
+        else:
+            height, width = self.height, self.width
+        return height, width
+    
+    
+    def get_num_frames(self, reader):
+        num_frames = self.num_frames
+        if int(reader.count_frames()) < num_frames:
+            num_frames = int(reader.count_frames())
+            while num_frames > 1 and num_frames % self.time_division_factor != self.time_division_remainder:
+                num_frames -= 1
+        return num_frames
+    
+
+    def load_video(self, file_path):
+        reader = imageio.get_reader(file_path)
+        num_frames = self.get_num_frames(reader)
+        frames = []
+        for frame_id in range(num_frames):
+            frame = reader.get_data(frame_id)
+            frame = Image.fromarray(frame)
+            frame = self.crop_and_resize(frame, *self.get_height_width(frame))
+            frames.append(frame)
+        reader.close()
+        return frames
+    
+    
+    def load_image(self, file_path):
+        image = Image.open(file_path).convert("RGB")
+        image = self.crop_and_resize(image, *self.get_height_width(image))
+        frames = [image]
+        return frames
+    
+    
+    def is_image(self, file_path):
+        file_ext_name = file_path.split(".")[-1]
+        return file_ext_name.lower() in self.image_file_extension
+    
+    
+    def is_video(self, file_path):
+        file_ext_name = file_path.split(".")[-1]
+        return file_ext_name.lower() in self.video_file_extension
+    
+    
+    def load_data(self, file_path):
+        if self.is_image(file_path):
+            return self.load_image(file_path)
+        elif self.is_video(file_path):
+            return self.load_video(file_path)
+        else:
+            return None
+
+
+    def __getitem__(self, data_id):
+        data = self.data[data_id % len(self.data)].copy()
+        for key in self.data_file_keys:
+            if key in data:
+                path = os.path.join(self.base_path, data[key])
+                data[key] = self.load_data(path)
+                if data[key] is None:
+                    warnings.warn(f"cannot load file {data[key]}.")
+                    return None
+        return data
+    
+
+    def __len__(self):
+        return len(self.data) * self.repeat
+
+
+
+class DiffusionTrainingModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        
+        
+    def to(self, *args, **kwargs):
+        for name, model in self.named_children():
+            model.to(*args, **kwargs)
+        return self
+        
+        
+    def trainable_modules(self):
+        trainable_modules = filter(lambda p: p.requires_grad, self.parameters())
+        return trainable_modules
+    
+    
+    def trainable_param_names(self):
+        trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.named_parameters()))
+        trainable_param_names = set([named_param[0] for named_param in trainable_param_names])
+        return trainable_param_names
+    
+    
+    def add_lora_to_model(self, model, target_modules, lora_rank, lora_alpha=None):
+        if lora_alpha is None:
+            lora_alpha = lora_rank
+        lora_config = LoraConfig(r=lora_rank, lora_alpha=lora_alpha, target_modules=target_modules)
+        model = inject_adapter_in_model(lora_config, model)
+        return model
+    
+    
+    def export_trainable_state_dict(self, state_dict, remove_prefix=None):
+        trainable_param_names = self.trainable_param_names()
+        state_dict = {name: param for name, param in state_dict.items() if name in trainable_param_names}
+        if remove_prefix is not None:
+            state_dict_ = {}
+            for name, param in state_dict.items():
+                if name.startswith(remove_prefix):
+                    name = name[len(remove_prefix):]
+                state_dict_[name] = param
+            state_dict = state_dict_
+        return state_dict
+
+
+
+class ModelLogger:
+    def __init__(self, output_path, remove_prefix_in_ckpt=None, state_dict_converter=lambda x:x):
+        self.output_path = output_path
+        self.remove_prefix_in_ckpt = remove_prefix_in_ckpt
+        self.state_dict_converter = state_dict_converter
+        
+    
+    def on_step_end(self, loss):
+        pass
+    
+    
+    def on_epoch_end(self, accelerator, model, epoch_id):
+        accelerator.wait_for_everyone()
+        if accelerator.is_main_process:
+            state_dict = accelerator.get_state_dict(model)
+            state_dict = accelerator.unwrap_model(model).export_trainable_state_dict(state_dict, remove_prefix=self.remove_prefix_in_ckpt)
+            state_dict = self.state_dict_converter(state_dict)
+            os.makedirs(self.output_path, exist_ok=True)
+            path = os.path.join(self.output_path, f"epoch-{epoch_id}.safetensors")
+            accelerator.save(state_dict, path, safe_serialization=True)
+
+
+
+def launch_training_task(
+    dataset: torch.utils.data.Dataset,
+    model: DiffusionTrainingModule,
+    model_logger: ModelLogger,
+    optimizer: torch.optim.Optimizer,
+    scheduler: torch.optim.lr_scheduler.LRScheduler,
+    num_epochs: int = 1,
+    gradient_accumulation_steps: int = 1,
+):
+    dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, collate_fn=lambda x: x[0])
+    accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
+    model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler)
+    
+    for epoch_id in range(num_epochs):
+        for data in tqdm(dataloader):
+            with accelerator.accumulate(model):
+                optimizer.zero_grad()
+                loss = model(data)
+                accelerator.backward(loss)
+                optimizer.step()
+                model_logger.on_step_end(loss)
+                scheduler.step()
+        model_logger.on_epoch_end(accelerator, model, epoch_id)
+
+
+
+def launch_data_process_task(model: DiffusionTrainingModule, dataset, output_path="./models"):
+    dataloader = torch.utils.data.DataLoader(dataset, shuffle=False, collate_fn=lambda x: x[0])
+    accelerator = Accelerator()
+    model, dataloader = accelerator.prepare(model, dataloader)
+    os.makedirs(os.path.join(output_path, "data_cache"), exist_ok=True)
+    for data_id, data in enumerate(tqdm(dataloader)):
+        with torch.no_grad():
+            inputs = model.forward_preprocess(data)
+            inputs = {key: inputs[key] for key in model.model_input_keys if key in inputs}
+            torch.save(inputs, os.path.join(output_path, "data_cache", f"{data_id}.pth"))
+
+
+
+def wan_parser():
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument("--dataset_base_path", type=str, default="", required=True, help="Base path of the dataset.")
+    parser.add_argument("--dataset_metadata_path", type=str, default=None, help="Path to the metadata file of the dataset.")
+    parser.add_argument("--max_pixels", type=int, default=1280*720, help="Maximum number of pixels per frame, used for dynamic resolution..")
+    parser.add_argument("--height", type=int, default=None, help="Height of images or videos. Leave `height` and `width` empty to enable dynamic resolution.")
+    parser.add_argument("--width", type=int, default=None, help="Width of images or videos. Leave `height` and `width` empty to enable dynamic resolution.")
+    parser.add_argument("--num_frames", type=int, default=81, help="Number of frames per video. Frames are sampled from the video prefix.")
+    parser.add_argument("--data_file_keys", type=str, default="image,video", help="Data file keys in the metadata. Comma-separated.")
+    parser.add_argument("--dataset_repeat", type=int, default=1, help="Number of times to repeat the dataset per epoch.")
+    parser.add_argument("--model_paths", type=str, default=None, help="Paths to load models. In JSON format.")
+    parser.add_argument("--model_id_with_origin_paths", type=str, default=None, help="Model ID with origin paths, e.g., Wan-AI/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors. Comma-separated.")
+    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate.")
+    parser.add_argument("--num_epochs", type=int, default=1, help="Number of epochs.")
+    parser.add_argument("--output_path", type=str, default="./models", help="Output save path.")
+    parser.add_argument("--remove_prefix_in_ckpt", type=str, default="pipe.dit.", help="Remove prefix in ckpt.")
+    parser.add_argument("--trainable_models", type=str, default=None, help="Models to train, e.g., dit, vae, text_encoder.")
+    parser.add_argument("--lora_base_model", type=str, default=None, help="Which model LoRA is added to.")
+    parser.add_argument("--lora_target_modules", type=str, default="q,k,v,o,ffn.0,ffn.2", help="Which layers LoRA is added to.")
+    parser.add_argument("--lora_rank", type=int, default=32, help="Rank of LoRA.")
+    parser.add_argument("--extra_inputs", default=None, help="Additional model inputs, comma-separated.")
+    parser.add_argument("--use_gradient_checkpointing_offload", default=False, action="store_true", help="Whether to offload gradient checkpointing to CPU memory.")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps.")
+    parser.add_argument("--max_timestep_boundary", type=float, default=1.0, help="Max timestep boundary (for mixed models, e.g., Wan-AI/Wan2.2-I2V-A14B).")
+    parser.add_argument("--min_timestep_boundary", type=float, default=0.0, help="Min timestep boundary (for mixed models, e.g., Wan-AI/Wan2.2-I2V-A14B).")
+    return parser
+
+
+
+def flux_parser():
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument("--dataset_base_path", type=str, default="", required=True, help="Base path of the dataset.")
+    parser.add_argument("--dataset_metadata_path", type=str, default=None, help="Path to the metadata file of the dataset.")
+    parser.add_argument("--max_pixels", type=int, default=1024*1024, help="Maximum number of pixels per frame, used for dynamic resolution..")
+    parser.add_argument("--height", type=int, default=None, help="Height of images. Leave `height` and `width` empty to enable dynamic resolution.")
+    parser.add_argument("--width", type=int, default=None, help="Width of images. Leave `height` and `width` empty to enable dynamic resolution.")
+    parser.add_argument("--data_file_keys", type=str, default="image", help="Data file keys in the metadata. Comma-separated.")
+    parser.add_argument("--dataset_repeat", type=int, default=1, help="Number of times to repeat the dataset per epoch.")
+    parser.add_argument("--model_paths", type=str, default=None, help="Paths to load models. In JSON format.")
+    parser.add_argument("--model_id_with_origin_paths", type=str, default=None, help="Model ID with origin paths, e.g., Wan-AI/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors. Comma-separated.")
+    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate.")
+    parser.add_argument("--num_epochs", type=int, default=1, help="Number of epochs.")
+    parser.add_argument("--output_path", type=str, default="./models", help="Output save path.")
+    parser.add_argument("--remove_prefix_in_ckpt", type=str, default="pipe.dit.", help="Remove prefix in ckpt.")
+    parser.add_argument("--trainable_models", type=str, default=None, help="Models to train, e.g., dit, vae, text_encoder.")
+    parser.add_argument("--lora_base_model", type=str, default=None, help="Which model LoRA is added to.")
+    parser.add_argument("--lora_target_modules", type=str, default="q,k,v,o,ffn.0,ffn.2", help="Which layers LoRA is added to.")
+    parser.add_argument("--lora_rank", type=int, default=32, help="Rank of LoRA.")
+    parser.add_argument("--extra_inputs", default=None, help="Additional model inputs, comma-separated.")
+    parser.add_argument("--align_to_opensource_format", default=False, action="store_true", help="Whether to align the lora format to opensource format. Only for DiT's LoRA.")
+    parser.add_argument("--use_gradient_checkpointing", default=False, action="store_true", help="Whether to use gradient checkpointing.")
+    parser.add_argument("--use_gradient_checkpointing_offload", default=False, action="store_true", help="Whether to offload gradient checkpointing to CPU memory.")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps.")
+    return parser
+
+
+
+def qwen_image_parser():
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument("--dataset_base_path", type=str, default="", required=True, help="Base path of the dataset.")
+    parser.add_argument("--dataset_metadata_path", type=str, default=None, help="Path to the metadata file of the dataset.")
+    parser.add_argument("--max_pixels", type=int, default=1024*1024, help="Maximum number of pixels per frame, used for dynamic resolution..")
+    parser.add_argument("--height", type=int, default=None, help="Height of images. Leave `height` and `width` empty to enable dynamic resolution.")
+    parser.add_argument("--width", type=int, default=None, help="Width of images. Leave `height` and `width` empty to enable dynamic resolution.")
+    parser.add_argument("--data_file_keys", type=str, default="image", help="Data file keys in the metadata. Comma-separated.")
+    parser.add_argument("--dataset_repeat", type=int, default=1, help="Number of times to repeat the dataset per epoch.")
+    parser.add_argument("--model_paths", type=str, default=None, help="Paths to load models. In JSON format.")
+    parser.add_argument("--model_id_with_origin_paths", type=str, default=None, help="Model ID with origin paths, e.g., Wan-AI/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors. Comma-separated.")
+    parser.add_argument("--tokenizer_path", type=str, default=None, help="Paths to tokenizer.")
+    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate.")
+    parser.add_argument("--num_epochs", type=int, default=1, help="Number of epochs.")
+    parser.add_argument("--output_path", type=str, default="./models", help="Output save path.")
+    parser.add_argument("--remove_prefix_in_ckpt", type=str, default="pipe.dit.", help="Remove prefix in ckpt.")
+    parser.add_argument("--trainable_models", type=str, default=None, help="Models to train, e.g., dit, vae, text_encoder.")
+    parser.add_argument("--lora_base_model", type=str, default=None, help="Which model LoRA is added to.")
+    parser.add_argument("--lora_target_modules", type=str, default="q,k,v,o,ffn.0,ffn.2", help="Which layers LoRA is added to.")
+    parser.add_argument("--lora_rank", type=int, default=32, help="Rank of LoRA.")
+    parser.add_argument("--extra_inputs", default=None, help="Additional model inputs, comma-separated.")
+    parser.add_argument("--align_to_opensource_format", default=False, action="store_true", help="Whether to align the lora format to opensource format. Only for DiT's LoRA.")
+    parser.add_argument("--use_gradient_checkpointing", default=False, action="store_true", help="Whether to use gradient checkpointing.")
+    parser.add_argument("--use_gradient_checkpointing_offload", default=False, action="store_true", help="Whether to offload gradient checkpointing to CPU memory.")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps.")
+    return parser

diffsynth/utils/__init__.py

@@ -0,0 +1,261 @@
+import torch, warnings, glob, os
+import numpy as np
+from PIL import Image
+from einops import repeat, reduce
+from typing import Optional, Union
+from dataclasses import dataclass
+from modelscope import snapshot_download
+import numpy as np
+from PIL import Image
+from typing import Optional
+
+
+class BasePipeline(torch.nn.Module):
+
+    def __init__(
+        self,
+        device="cuda", torch_dtype=torch.float16,
+        height_division_factor=64, width_division_factor=64,
+        time_division_factor=None, time_division_remainder=None,
+    ):
+        super().__init__()
+        # The device and torch_dtype is used for the storage of intermediate variables, not models.
+        self.device = device
+        self.torch_dtype = torch_dtype
+        # The following parameters are used for shape check.
+        self.height_division_factor = height_division_factor
+        self.width_division_factor = width_division_factor
+        self.time_division_factor = time_division_factor
+        self.time_division_remainder = time_division_remainder
+        self.vram_management_enabled = False
+        
+        
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
+        if device is not None:
+            self.device = device
+        if dtype is not None:
+            self.torch_dtype = dtype
+        super().to(*args, **kwargs)
+        return self
+
+
+    def check_resize_height_width(self, height, width, num_frames=None):
+        # Shape check
+        if height % self.height_division_factor != 0:
+            height = (height + self.height_division_factor - 1) // self.height_division_factor * self.height_division_factor
+            print(f"height % {self.height_division_factor} != 0. We round it up to {height}.")
+        if width % self.width_division_factor != 0:
+            width = (width + self.width_division_factor - 1) // self.width_division_factor * self.width_division_factor
+            print(f"width % {self.width_division_factor} != 0. We round it up to {width}.")
+        if num_frames is None:
+            return height, width
+        else:
+            if num_frames % self.time_division_factor != self.time_division_remainder:
+                num_frames = (num_frames + self.time_division_factor - 1) // self.time_division_factor * self.time_division_factor + self.time_division_remainder
+                print(f"num_frames % {self.time_division_factor} != {self.time_division_remainder}. We round it up to {num_frames}.")
+            return height, width, num_frames
+
+
+    def preprocess_image(self, image, torch_dtype=None, device=None, pattern="B C H W", min_value=-1, max_value=1):
+        # Transform a PIL.Image to torch.Tensor
+        image = torch.Tensor(np.array(image, dtype=np.float32))
+        image = image.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device)
+        image = image * ((max_value - min_value) / 255) + min_value
+        image = repeat(image, f"H W C -> {pattern}", **({"B": 1} if "B" in pattern else {}))
+        return image
+
+
+    def preprocess_video(self, video, torch_dtype=None, device=None, pattern="B C T H W", min_value=-1, max_value=1):
+        # Transform a list of PIL.Image to torch.Tensor
+        video = [self.preprocess_image(image, torch_dtype=torch_dtype, device=device, min_value=min_value, max_value=max_value) for image in video]
+        video = torch.stack(video, dim=pattern.index("T") // 2)
+        return video
+
+
+    def vae_output_to_image(self, vae_output, pattern="B C H W", min_value=-1, max_value=1):
+        # Transform a torch.Tensor to PIL.Image
+        if pattern != "H W C":
+            vae_output = reduce(vae_output, f"{pattern} -> H W C", reduction="mean")
+        image = ((vae_output - min_value) * (255 / (max_value - min_value))).clip(0, 255)
+        image = image.to(device="cpu", dtype=torch.uint8)
+        image = Image.fromarray(image.numpy())
+        return image
+
+
+    def vae_output_to_video(self, vae_output, pattern="B C T H W", min_value=-1, max_value=1):
+        # Transform a torch.Tensor to list of PIL.Image
+        if pattern != "T H W C":
+            vae_output = reduce(vae_output, f"{pattern} -> T H W C", reduction="mean")
+        video = [self.vae_output_to_image(image, pattern="H W C", min_value=min_value, max_value=max_value) for image in vae_output]
+        return video
+
+
+    def load_models_to_device(self, model_names=[]):
+        if self.vram_management_enabled:
+            # offload models
+            for name, model in self.named_children():
+                if name not in model_names:
+                    if hasattr(model, "vram_management_enabled") and model.vram_management_enabled:
+                        for module in model.modules():
+                            if hasattr(module, "offload"):
+                                module.offload()
+                    else:
+                        model.cpu()
+            torch.cuda.empty_cache()
+            # onload models
+            for name, model in self.named_children():
+                if name in model_names:
+                    if hasattr(model, "vram_management_enabled") and model.vram_management_enabled:
+                        for module in model.modules():
+                            if hasattr(module, "onload"):
+                                module.onload()
+                    else:
+                        model.to(self.device)
+
+
+    def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None):
+        # Initialize Gaussian noise
+        generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed)
+        noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype)
+        noise = noise.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device)
+        return noise
+
+
+    def enable_cpu_offload(self):
+        warnings.warn("`enable_cpu_offload` will be deprecated. Please use `enable_vram_management`.")
+        self.vram_management_enabled = True
+        
+        
+    def get_vram(self):
+        return torch.cuda.mem_get_info(self.device)[1] / (1024 ** 3)
+    
+    
+    def freeze_except(self, model_names):
+        for name, model in self.named_children():
+            if name in model_names:
+                model.train()
+                model.requires_grad_(True)
+            else:
+                model.eval()
+                model.requires_grad_(False)
+
+
+@dataclass
+class ModelConfig:
+    path: Union[str, list[str]] = None
+    model_id: str = None
+    origin_file_pattern: Union[str, list[str]] = None
+    download_resource: str = "ModelScope"
+    offload_device: Optional[Union[str, torch.device]] = None
+    offload_dtype: Optional[torch.dtype] = None
+    local_model_path: str = None
+    skip_download: bool = False
+
+    def download_if_necessary(self, use_usp=False):
+        if self.path is None:
+            # Check model_id and origin_file_pattern
+            if self.model_id is None:
+                raise ValueError(f"""No valid model files. Please use `ModelConfig(path="xxx")` or `ModelConfig(model_id="xxx/yyy", origin_file_pattern="zzz")`.""")
+            
+            # Skip if not in rank 0
+            if use_usp:
+                import torch.distributed as dist
+                skip_download = self.skip_download or dist.get_rank() != 0
+            else:
+                skip_download = self.skip_download
+                
+            # Check whether the origin path is a folder
+            if self.origin_file_pattern is None or self.origin_file_pattern == "":
+                self.origin_file_pattern = ""
+                allow_file_pattern = None
+                is_folder = True
+            elif isinstance(self.origin_file_pattern, str) and self.origin_file_pattern.endswith("/"):
+                allow_file_pattern = self.origin_file_pattern + "*"
+                is_folder = True
+            else:
+                allow_file_pattern = self.origin_file_pattern
+                is_folder = False
+            
+            # Download
+            if self.local_model_path is None:
+                self.local_model_path = "./models"
+            if not skip_download:
+                downloaded_files = glob.glob(self.origin_file_pattern, root_dir=os.path.join(self.local_model_path, self.model_id))
+                snapshot_download(
+                    self.model_id,
+                    local_dir=os.path.join(self.local_model_path, self.model_id),
+                    allow_file_pattern=allow_file_pattern,
+                    ignore_file_pattern=downloaded_files,
+                    local_files_only=False
+                )
+            
+            # Let rank 1, 2, ... wait for rank 0
+            if use_usp:
+                import torch.distributed as dist
+                dist.barrier(device_ids=[dist.get_rank()])
+                
+            # Return downloaded files
+            if is_folder:
+                self.path = os.path.join(self.local_model_path, self.model_id, self.origin_file_pattern)
+            else:
+                self.path = glob.glob(os.path.join(self.local_model_path, self.model_id, self.origin_file_pattern))
+            if isinstance(self.path, list) and len(self.path) == 1:
+                self.path = self.path[0]
+
+
+
+class PipelineUnit:
+    def __init__(
+        self,
+        seperate_cfg: bool = False,
+        take_over: bool = False,
+        input_params: tuple[str] = None,
+        input_params_posi: dict[str, str] = None,
+        input_params_nega: dict[str, str] = None,
+        onload_model_names: tuple[str] = None
+    ):
+        self.seperate_cfg = seperate_cfg
+        self.take_over = take_over
+        self.input_params = input_params
+        self.input_params_posi = input_params_posi
+        self.input_params_nega = input_params_nega
+        self.onload_model_names = onload_model_names
+
+
+    def process(self, pipe: BasePipeline, inputs: dict, positive=True, **kwargs) -> dict:
+        raise NotImplementedError("`process` is not implemented.")
+
+
+
+class PipelineUnitRunner:
+    def __init__(self):
+        pass
+
+    def __call__(self, unit: PipelineUnit, pipe: BasePipeline, inputs_shared: dict, inputs_posi: dict, inputs_nega: dict) -> tuple[dict, dict]:
+        if unit.take_over:
+            # Let the pipeline unit take over this function.
+            inputs_shared, inputs_posi, inputs_nega = unit.process(pipe, inputs_shared=inputs_shared, inputs_posi=inputs_posi, inputs_nega=inputs_nega)
+        elif unit.seperate_cfg:
+            # Positive side
+            processor_inputs = {name: inputs_posi.get(name_) for name, name_ in unit.input_params_posi.items()}
+            if unit.input_params is not None:
+                for name in unit.input_params:
+                    processor_inputs[name] = inputs_shared.get(name)
+            processor_outputs = unit.process(pipe, **processor_inputs)
+            inputs_posi.update(processor_outputs)
+            # Negative side
+            if inputs_shared["cfg_scale"] != 1:
+                processor_inputs = {name: inputs_nega.get(name_) for name, name_ in unit.input_params_nega.items()}
+                if unit.input_params is not None:
+                    for name in unit.input_params:
+                        processor_inputs[name] = inputs_shared.get(name)
+                processor_outputs = unit.process(pipe, **processor_inputs)
+                inputs_nega.update(processor_outputs)
+            else:
+                inputs_nega.update(processor_outputs)
+        else:
+            processor_inputs = {name: inputs_shared.get(name) for name in unit.input_params}
+            processor_outputs = unit.process(pipe, **processor_inputs)
+            inputs_shared.update(processor_outputs)
+        return inputs_shared, inputs_posi, inputs_nega

diffsynth/vram_management/__init__.py

@@ -0,0 +1,2 @@
+from .layers import *
+from .gradient_checkpointing import *

diffsynth/vram_management/gradient_checkpointing.py

@@ -0,0 +1,34 @@
+import torch
+
+
+def create_custom_forward(module):
+    def custom_forward(*inputs, **kwargs):
+        return module(*inputs, **kwargs)
+    return custom_forward
+
+
+def gradient_checkpoint_forward(
+    model,
+    use_gradient_checkpointing,
+    use_gradient_checkpointing_offload,
+    *args,
+    **kwargs,
+):
+    if use_gradient_checkpointing_offload:
+        with torch.autograd.graph.save_on_cpu():
+            model_output = torch.utils.checkpoint.checkpoint(
+                create_custom_forward(model),
+                *args,
+                **kwargs,
+                use_reentrant=False,
+            )
+    elif use_gradient_checkpointing:
+        model_output = torch.utils.checkpoint.checkpoint(
+            create_custom_forward(model),
+            *args,
+            **kwargs,
+            use_reentrant=False,
+        )
+    else:
+        model_output = model(*args, **kwargs)
+    return model_output

diffsynth/vram_management/layers.py

@@ -0,0 +1,167 @@
+import torch, copy
+from ..models.utils import init_weights_on_device
+
+
+def cast_to(weight, dtype, device):
+    r = torch.empty_like(weight, dtype=dtype, device=device)
+    r.copy_(weight)
+    return r
+
+
+class AutoTorchModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        
+    def check_free_vram(self):
+        gpu_mem_state = torch.cuda.mem_get_info(self.computation_device)
+        used_memory = (gpu_mem_state[1] - gpu_mem_state[0]) / (1024 ** 3)
+        return used_memory < self.vram_limit
+
+    def offload(self):
+        if self.state != 0:
+            self.to(dtype=self.offload_dtype, device=self.offload_device)
+            self.state = 0
+
+    def onload(self):
+        if self.state != 1:
+            self.to(dtype=self.onload_dtype, device=self.onload_device)
+            self.state = 1
+            
+    def keep(self):
+        if self.state != 2:
+            self.to(dtype=self.computation_dtype, device=self.computation_device)
+            self.state = 2
+
+
+class AutoWrappedModule(AutoTorchModule):
+    def __init__(self, module: torch.nn.Module, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, **kwargs):
+        super().__init__()
+        self.module = module.to(dtype=offload_dtype, device=offload_device)
+        self.offload_dtype = offload_dtype
+        self.offload_device = offload_device
+        self.onload_dtype = onload_dtype
+        self.onload_device = onload_device
+        self.computation_dtype = computation_dtype
+        self.computation_device = computation_device
+        self.vram_limit = vram_limit
+        self.state = 0
+
+    def forward(self, *args, **kwargs):
+        if self.state == 2:
+            module = self.module
+        else:
+            if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device:
+                module = self.module
+            elif self.vram_limit is not None and self.check_free_vram():
+                self.keep()
+                module = self.module
+            else:
+                module = copy.deepcopy(self.module).to(dtype=self.computation_dtype, device=self.computation_device)
+        return module(*args, **kwargs)
+    
+
+class WanAutoCastLayerNorm(torch.nn.LayerNorm, AutoTorchModule):
+    def __init__(self, module: torch.nn.LayerNorm, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, **kwargs):
+        with init_weights_on_device(device=torch.device("meta")):
+            super().__init__(module.normalized_shape, eps=module.eps, elementwise_affine=module.elementwise_affine, bias=module.bias is not None, dtype=offload_dtype, device=offload_device)
+        self.weight = module.weight
+        self.bias = module.bias
+        self.offload_dtype = offload_dtype
+        self.offload_device = offload_device
+        self.onload_dtype = onload_dtype
+        self.onload_device = onload_device
+        self.computation_dtype = computation_dtype
+        self.computation_device = computation_device
+        self.vram_limit = vram_limit
+        self.state = 0
+
+    def forward(self, x, *args, **kwargs):
+        if self.state == 2:
+            weight, bias = self.weight, self.bias
+        else:
+            if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device:
+                weight, bias = self.weight, self.bias
+            elif self.vram_limit is not None and self.check_free_vram():
+                self.keep()
+                weight, bias = self.weight, self.bias
+            else:
+                weight = None if self.weight is None else cast_to(self.weight, self.computation_dtype, self.computation_device)
+                bias = None if self.bias is None else cast_to(self.bias, self.computation_dtype, self.computation_device)
+        with torch.amp.autocast(device_type=x.device.type):
+            x = torch.nn.functional.layer_norm(x.float(), self.normalized_shape, weight, bias, self.eps).type_as(x)
+        return x
+    
+
+class AutoWrappedLinear(torch.nn.Linear, AutoTorchModule):
+    def __init__(self, module: torch.nn.Linear, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, name="", **kwargs):
+        with init_weights_on_device(device=torch.device("meta")):
+            super().__init__(in_features=module.in_features, out_features=module.out_features, bias=module.bias is not None, dtype=offload_dtype, device=offload_device)
+        self.weight = module.weight
+        self.bias = module.bias
+        self.offload_dtype = offload_dtype
+        self.offload_device = offload_device
+        self.onload_dtype = onload_dtype
+        self.onload_device = onload_device
+        self.computation_dtype = computation_dtype
+        self.computation_device = computation_device
+        self.vram_limit = vram_limit
+        self.state = 0
+        self.name = name
+        self.lora_A_weights = []
+        self.lora_B_weights = []
+        self.lora_merger = None
+
+    def forward(self, x, *args, **kwargs):
+        if self.state == 2:
+            weight, bias = self.weight, self.bias
+        else:
+            if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device:
+                weight, bias = self.weight, self.bias
+            elif self.vram_limit is not None and self.check_free_vram():
+                self.keep()
+                weight, bias = self.weight, self.bias
+            else:
+                weight = cast_to(self.weight, self.computation_dtype, self.computation_device)
+                bias = None if self.bias is None else cast_to(self.bias, self.computation_dtype, self.computation_device)
+        out = torch.nn.functional.linear(x, weight, bias)
+        
+        if len(self.lora_A_weights) == 0:
+            # No LoRA
+            return out
+        elif self.lora_merger is None:
+            # Native LoRA inference
+            for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights):
+                out = out + x @ lora_A.T @ lora_B.T
+        else:
+            # LoRA fusion
+            lora_output = []
+            for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights):
+                lora_output.append(x @ lora_A.T @ lora_B.T)
+            lora_output = torch.stack(lora_output)
+            out = self.lora_merger(out, lora_output)
+        return out
+
+
+def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, total_num_param=0, vram_limit=None, name_prefix=""):
+    for name, module in model.named_children():
+        layer_name = name if name_prefix == "" else name_prefix + "." + name
+        for source_module, target_module in module_map.items():
+            if isinstance(module, source_module):
+                num_param = sum(p.numel() for p in module.parameters())
+                if max_num_param is not None and total_num_param + num_param > max_num_param:
+                    module_config_ = overflow_module_config
+                else:
+                    module_config_ = module_config
+                module_ = target_module(module, **module_config_, vram_limit=vram_limit, name=layer_name)
+                setattr(model, name, module_)
+                total_num_param += num_param
+                break
+        else:
+            total_num_param = enable_vram_management_recursively(module, module_map, module_config, max_num_param, overflow_module_config, total_num_param, vram_limit=vram_limit, name_prefix=layer_name)
+    return total_num_param
+
+
+def enable_vram_management(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, vram_limit=None):
+    enable_vram_management_recursively(model, module_map, module_config, max_num_param, overflow_module_config, total_num_param=0, vram_limit=vram_limit)
+    model.vram_management_enabled = True
+

tools/full_inference_modules_gradio.sh

@@ -0,0 +1,162 @@
+#!/bin/bash
+which python
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}
+
+export CUDA_HOME=/usr/local/cuda-12.4
+export PATH=$CUDA_HOME/bin:$PATH
+export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
+
+# export HF_HUB_DOWNLOAD_TIMEOUT=30
+# export HF_HOME=${HOME}/.cache/huggingface/
+
+export CONDA_PATH=${HOME}/miniconda3 # ! Replace with your CONDA path
+export CONDA_PREFIX_1=${CONDA_PATH}/envs/diffsynth # ! Replace with your diffsynth environment path
+export CONDA_PREFIX_2=${CONDA_PATH}/envs/cosmos-predict1 # ! Replace with your cosmos-predict1 environment path
+export DR_REPO_PATH=${REPO_PATH}/third_party/cosmos-transfer1-diffusion-renderer
+
+# if input image, set FRAME=1
+if [ "$TEST_TYPE" -eq 1 ]; then
+    export FRAME=1
+fi
+
+# -------------------------------------- module 1 -----------------------------------------
+if [ "$INFER_1" -eq 1 ]; then
+    source ${CONDA_PATH}/bin/activate ${CONDA_PREFIX_2}
+    cd ${DR_REPO_PATH}
+    mkdir -p ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID
+    rm -rf ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/frames
+    if [ "$TEST_TYPE" -eq 0 ]; then
+        # If the input is a video, split it into multiple frames and save them in the output_folder/frames/$TEST_ID folder
+        CUDA_HOME=$CONDA_PREFIX_2 PYTHONPATH=$(pwd) python scripts/dataproc_extract_frames_from_video.py \
+        --input_folder ${REPO_PATH}/datasets/gradio_data/upload_data/$TEST_ID \
+        --output_folder ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/frames/ --frame_rate=$FRAME_RATE --max_frames=$FRAME
+    elif [ "$TEST_TYPE" -eq 1 ]; then
+        # If the input is an image, directly copy it to the output_folder/frames/$TEST_ID folder
+        mkdir -p ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/frames/$TEST_ID
+        cp ${REPO_PATH}/datasets/gradio_data/upload_data/$TEST_ID/$TEST_ID.jpg ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/frames/$TEST_ID/frame_00000.jpg
+    fi
+
+    if [ "$FRAME" -eq 1 ]; then
+        export COSMOS_FRAME=$FRAME
+    else
+        export COSMOS_FRAME=57
+    fi
+    rm -rf ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/frames_delighting
+    CUDA_HOME=$CONDA_PREFIX_2 PYTHONPATH=$(pwd) python cosmos_predict1/diffusion/inference/inference_inverse_renderer.py \
+        --checkpoint_dir checkpoints --diffusion_transformer_dir Diffusion_Renderer_Inverse_Cosmos_7B \
+        --dataset_path="${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/frames/" --num_video_frames $COSMOS_FRAME --group_mode folder \
+        --video_save_folder="${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/frames_delighting/" --chunk_mode 'first' --overlap_n_frames 0 --save_video 'False'
+
+    cd ${REPO_PATH}
+    CUDA_HOME=$CONDA_PREFIX_2 PYTHONPATH=$(pwd) python tools/reorg_gbuffer_from_dr_delighting.py \
+        --frame_dir "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/frames/" \
+        --gbuffer_dir "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/frames_delighting/gbuffer_frames/" \
+        --tgt_dir "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/rego/"
+fi
+
+# -------------------------------------- module 2 -----------------------------------------
+if [ "$INFER_2" -eq 1 ]; then
+    source ${CONDA_PATH}/bin/activate ${CONDA_PREFIX_2}
+    cd ${REPO_PATH}
+    mkdir -p ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/envs
+    if [ "$USE_OFFICE_ENV" -eq 0 ]; then
+        CUDA_HOME=$CONDA_PREFIX_2 PYTHONPATH=$(pwd) python tools/process_env_maps.py \
+            --env_dir ${REPO_PATH}/datasets/gradio_data/upload_data/$TEST_ID/$TEST_ENV \
+            --save_path ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/envs/$TEST_ENV \
+            --env_strength $ENV_STRENGTH
+    else
+        CUDA_HOME=$CONDA_PREFIX_2 PYTHONPATH=$(pwd) python tools/process_env_maps.py \
+            --env_dir ${REPO_PATH}/datasets/gradio_data/assets/envs_demo/$TEST_ENV \
+            --save_path ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/envs/$TEST_ENV \
+            --env_strength $ENV_STRENGTH
+    fi
+fi
+
+# -------------------------------------- module 3 -----------------------------------------
+if [ "$INFER_3" -eq 1 ]; then
+    source ${CONDA_PATH}/bin/activate ${CONDA_PREFIX_1}
+    cd ${REPO_PATH}
+
+    # LIGHT_TYPE non-zero forces FRAME ≠ 1
+    if [ "$LIGHT_TYPE" -ne 0 ] && [ "$FRAME" -eq 1 ]; then
+        export INFER_FRAME=25
+    else
+        export INFER_FRAME=$FRAME
+    fi
+
+    if [ "$INFER_FRAME" -eq 1 ]; then
+        SAVE_EXT='png'
+    else
+        SAVE_EXT='mp4'
+    fi
+
+    if [ "$INFER_FRAME" -eq 1 ]; then
+        export HEIGHT=1024
+        export WIDTH=1472
+        export CKT="${REPO_PATH}/checkpoints/model_frame1_1024_1472.ckpt"
+    elif [ "$INFER_FRAME" -eq 25 ]; then
+        export HEIGHT=480
+        export WIDTH=832
+        export CKT="${REPO_PATH}/checkpoints/model_frame25_480_832.ckpt"
+    else
+        export HEIGHT=480
+        export WIDTH=832
+        export CKT="${REPO_PATH}/checkpoints/model_frame57_480_832.ckpt"
+    fi
+
+
+    if [ "$LIGHT_TYPE" -eq 0 ]; then
+        CUDA_HOME=$CONDA_PREFIX_1 PYTHONPATH=$(pwd) python relit_inference.py \
+            --dataset_path "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/rego" \
+            --ckpt_path=$CKT \
+            --output_dir "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/relighting.$TEST_ENV" \
+            --output_path "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/$TEST_ID.$TEST_ENV.$SAVE_EXT" \
+            --cfg_scale 1.0 \
+            --height $HEIGHT \
+            --width $WIDTH \
+            --num_frames $INFER_FRAME \
+            --padding_resolution \
+            --use_ref_image \
+            --env_map_path ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/envs/$TEST_ENV \
+            --frame_interval 1 \
+            --num_inference_steps $NUM_INFER_STEPS \
+            --wo_ref_weight $WORW \
+            --quality 10
+    elif [ "$LIGHT_TYPE" -eq 1 ]; then
+        CUDA_HOME=$CONDA_PREFIX_1 PYTHONPATH=$(pwd) python relit_inference.py \
+            --dataset_path "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/rego" \
+            --ckpt_path=$CKT \
+            --output_dir "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/relighting.$TEST_ENV" \
+            --output_path "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/$TEST_ID.$TEST_ENV.$SAVE_EXT" \
+            --cfg_scale 1.0 \
+            --height $HEIGHT \
+            --width $WIDTH \
+            --num_frames $INFER_FRAME \
+            --padding_resolution \
+            --use_ref_image \
+            --env_map_path ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/envs/$TEST_ENV \
+            --frame_interval 1 \
+            --num_inference_steps $NUM_INFER_STEPS \
+            --wo_ref_weight $WORW \
+            --use_rotate_light \
+            --quality 10
+    elif [ "$LIGHT_TYPE" -eq 2 ]; then
+        CUDA_HOME=$CONDA_PREFIX_1 PYTHONPATH=$(pwd) python relit_inference.py \
+            --dataset_path "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/rego" \
+            --ckpt_path=$CKT \
+            --output_dir "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/relighting.$TEST_ENV" \
+            --output_path "${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/$TEST_ID.$TEST_ENV.$SAVE_EXT" \
+            --cfg_scale 1.0 \
+            --height $HEIGHT \
+            --width $WIDTH \
+            --num_frames $INFER_FRAME \
+            --padding_resolution \
+            --use_ref_image \
+            --env_map_path ${REPO_PATH}/datasets/gradio_data/results/$TEST_ID/envs/$TEST_ENV \
+            --frame_interval 1 \
+            --num_inference_steps $NUM_INFER_STEPS \
+            --wo_ref_weight $WORW \
+            --use_fixed_frame_and_w_rotate_light \
+            --quality 10
+    fi
+fi

tools/process_env_maps.py

@@ -0,0 +1,599 @@
+import os
+import glob
+import json
+import cv2
+import numpy as np
+import argparse
+import torch
+import imageio.v2 as imageio
+import imageio.v3 as imageio_v3
+
+import nvdiffrast.torch as dr
+
+# Enable OpenEXR support in OpenCV
+os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
+
+def swap_yz_in_extrinsic_matrix(matrix):
+    assert matrix.shape == (4, 4), "Input must be a 4x4 matrix"
+
+    new_matrix = matrix.copy()
+    new_matrix[1, :], new_matrix[2, :] = new_matrix[2, :].copy(), new_matrix[1, :].copy()
+    new_matrix[:, 1], new_matrix[:, 2] = new_matrix[:, 2].copy(), new_matrix[:, 1].copy()
+
+    return new_matrix
+
+def swap_yz_output_in_extrinsic_matrix(matrix):
+    assert matrix.shape == (4, 4), "Input must be a 4x4 matrix"
+
+    new_matrix = matrix.copy()
+    new_matrix[1, :], new_matrix[2, :] = new_matrix[2, :].copy(), new_matrix[1, :].copy()
+
+    return new_matrix
+
+def euler_to_rotation_matrix(euler_angles, inverse_y=True, y_bias=0):
+    alpha, gamma, beta = euler_angles
+    beta += y_bias
+    if inverse_y:
+        beta = -beta
+
+    R_x = np.array([[1, 0, 0],
+                    [0, np.cos(alpha), -np.sin(alpha)],
+                    [0, np.sin(alpha), np.cos(alpha)]])
+    
+    R_y = np.array([[np.cos(beta), 0, np.sin(beta)],
+                    [0, 1, 0],
+                    [-np.sin(beta), 0, np.cos(beta)]])
+    
+    R_z = np.array([[np.cos(gamma), -np.sin(gamma), 0],
+                    [np.sin(gamma), np.cos(gamma), 0],
+                    [0, 0, 1]])
+    
+    R = np.dot(R_z, np.dot(R_y, R_x))
+    return R
+
+def remove_yaw_rotation(c2w_list):
+    c2w_0 = c2w_list[0]
+    rotation_0 = c2w_0[:3, :3]
+
+    yaw_0 = np.arctan2(rotation_0[2, 0], rotation_0[0, 0])
+
+    yaw_rotation_matrix = np.array([
+        [np.cos(-yaw_0), 0, np.sin(-yaw_0), 0],
+        [0, 1, 0, 0],
+        [-np.sin(-yaw_0), 0, np.cos(-yaw_0), 0],
+        [0, 0, 0, 1]
+    ])
+
+    c2w_list_adjusted = [yaw_rotation_matrix @ c2w_i for c2w_i in c2w_list]
+
+    return c2w_list_adjusted
+
+def adjust_yaw_rotation(c2w_list, yaw_0):
+
+    yaw_rotation_matrix = np.array([
+        [np.cos(-yaw_0), 0, np.sin(-yaw_0), 0],
+        [0, 1, 0, 0],
+        [-np.sin(-yaw_0), 0, np.cos(-yaw_0), 0],
+        [0, 0, 0, 1]
+    ])
+
+    c2w_list_adjusted = [yaw_rotation_matrix @ c2w_i for c2w_i in c2w_list]
+
+    return c2w_list_adjusted
+
+def reverse_yaw_rotation(c2w_list):
+    c2w_list_reversed = []
+
+    for c2w in c2w_list:
+        rotation = c2w[:3, :3]
+
+        yaw = np.arctan2(rotation[2, 0], rotation[0, 0])
+
+        reversed_yaw = 2 * yaw
+
+        reverse_yaw_rotation_matrix = np.array([
+            [np.cos(reversed_yaw), 0, np.sin(reversed_yaw), 0],
+            [0, 1, 0, 0],
+            [-np.sin(reversed_yaw), 0, np.cos(reversed_yaw), 0],
+            [0, 0, 0, 1]
+        ])
+
+        c2w_reversed = reverse_yaw_rotation_matrix @ c2w
+        c2w_list_reversed.append(c2w_reversed)
+
+    return c2w_list_reversed
+
+def prepare_camera_poses(num_frames, fixed_pose, pose_file, pose_offset, pose_reset, device, ign_camera_pose=True, swap_type=0, load_w2c=False, remove_y_rotation=False, reverse_y_rotation=False, yaw_0=0, pose_list=None, rotation_euler=None):
+    """Prepare camera poses based on the provided arguments."""
+    if pose_list is not None:
+        c2w_list = pose_list
+        for frame_idx, transform_matrix in enumerate(c2w_list):
+            if swap_type == 0:
+                new_transform_matrix = transform_matrix
+            elif swap_type == 1:
+                new_transform_matrix = swap_yz_in_extrinsic_matrix(transform_matrix)
+            elif swap_type == 2:
+                new_transform_matrix = swap_yz_output_in_extrinsic_matrix(transform_matrix)
+            if load_w2c:
+                new_transform_matrix = np.linalg.inv(new_transform_matrix)
+            c2w_list[frame_idx] = new_transform_matrix
+
+        if pose_reset:
+            w2c_0 = np.linalg.inv(c2w_list[0])
+            c2w_list = [w2c_0 @ c2w_i for c2w_i in c2w_list]
+        
+        if remove_y_rotation:
+            c2w_list = remove_yaw_rotation(c2w_list)
+
+        if reverse_y_rotation:
+            c2w_list = reverse_yaw_rotation(c2w_list)
+
+        if rotation_euler is not None:
+            for frame_idx, transform_matrix in enumerate(c2w_list):
+                rotation_matrix = euler_to_rotation_matrix(rotation_euler, inverse_y=True, y_bias=np.pi/2)
+
+                if ign_camera_pose:
+                    transform_matrix = c2w_list[0]
+                else:
+                    transform_matrix = c2w
+
+                rotation_matrix_4x4 = np.eye(4)
+                rotation_matrix_4x4[:3, :3] = rotation_matrix
+
+                new_transform_matrix = np.dot(rotation_matrix_4x4, transform_matrix)
+                c2w_list[frame_idx] = new_transform_matrix
+
+        return c2w_list
+
+    elif fixed_pose or pose_file is None:
+        return [np.eye(4) for _ in range(num_frames)]
+    
+    with open(pose_file, 'r') as f:
+        meta = json.load(f)
+    frames = meta['frames'][pose_offset:pose_offset + num_frames]
+
+    for frame_idx, data in enumerate(frames):
+        transform_matrix = np.array(data["transform_matrix"])
+        if swap_type == 0:
+            new_transform_matrix = transform_matrix
+        elif swap_type == 1:
+            new_transform_matrix = swap_yz_in_extrinsic_matrix(transform_matrix)
+        if load_w2c:
+            new_transform_matrix = np.linalg.inv(new_transform_matrix)
+        data["transform_matrix"] = new_transform_matrix.tolist()
+        frames[frame_idx] = data
+
+    if ign_camera_pose:
+        c2w_list = [np.array(frames[0]["transform_matrix"]) for frame in frames]
+    else:
+        c2w_list = [np.array(frame['transform_matrix']) for frame in frames]
+    
+    if pose_reset:
+        w2c_0 = np.linalg.inv(c2w_list[0])
+        c2w_list = [w2c_0 @ c2w_i for c2w_i in c2w_list] # compute c2c0
+    
+    if remove_y_rotation:
+        c2w_list = remove_yaw_rotation(c2w_list)
+
+    if reverse_y_rotation:
+        c2w_list = reverse_yaw_rotation(c2w_list)
+
+    if yaw_0 != 0:
+        c2w_list = adjust_yaw_rotation(c2w_list, yaw_0)
+
+    for frame_idx, (data, c2w) in enumerate(zip(frames, c2w_list)):
+        if "hdri_euler" in data.keys():
+            rotation_matrix = euler_to_rotation_matrix(data["hdri_euler"], inverse_y=False)
+
+            if ign_camera_pose:
+                transform_matrix = c2w_list[0]
+            else:
+                transform_matrix = c2w
+
+            rotation_matrix_4x4 = np.eye(4)
+            rotation_matrix_4x4[:3, :3] = rotation_matrix
+            new_transform_matrix = np.dot(rotation_matrix_4x4, transform_matrix)
+            c2w_list[frame_idx] = new_transform_matrix
+        else:
+            print(f"Warning: 'hdri_euler' not found in frame {frame_idx} of {pose_file}. Using original transform matrix.")
+            break
+
+    return c2w_list
+
+def latlong_vec(res, device=None):
+    gy, gx = torch.meshgrid(torch.linspace( 0.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device=device), 
+                            torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device=device),
+                            indexing='ij')
+    
+    sintheta, costheta = torch.sin(gy*np.pi), torch.cos(gy*np.pi)
+    sinphi, cosphi     = torch.sin(gx*np.pi), torch.cos(gx*np.pi)
+    
+    dir_vec = torch.stack((
+        sintheta*sinphi, 
+        costheta, 
+        -sintheta*cosphi
+        ), dim=-1)
+    # return dr.texture(cubemap[None, ...], dir_vec[None, ...].contiguous(), filter_mode='linear', boundary_mode='cube')[0]
+    return dir_vec #[H, W, 3]
+
+def envmap_vec(res, device=None):
+    return -latlong_vec(res, device).flip(0).flip(1) #[H, W, 3]
+
+def dot(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    return torch.sum(x*y, -1, keepdim=True)
+
+def length(x: torch.Tensor, eps: float =1e-20) -> torch.Tensor:
+    return torch.sqrt(torch.clamp(dot(x,x), min=eps)) # Clamp to avoid nan gradients because grad(sqrt(0)) = NaN
+
+def safe_normalize(x: torch.Tensor, eps: float =1e-20) -> torch.Tensor:
+    return x / length(x, eps)
+
+def cube_to_dir(s, x, y):
+    if s == 0:   rx, ry, rz = torch.ones_like(x), -y, -x
+    elif s == 1: rx, ry, rz = -torch.ones_like(x), -y, x
+    elif s == 2: rx, ry, rz = x, torch.ones_like(x), y
+    elif s == 3: rx, ry, rz = x, -torch.ones_like(x), -y
+    elif s == 4: rx, ry, rz = x, -y, torch.ones_like(x)
+    elif s == 5: rx, ry, rz = -x, -y, -torch.ones_like(x)
+    return torch.stack((rx, ry, rz), dim=-1)
+
+def latlong_to_cubemap(latlong_map, res):
+    cubemap = torch.zeros(6, res[0], res[1], latlong_map.shape[-1], dtype=torch.float32, device='cuda')
+    for s in range(6):
+        gy, gx = torch.meshgrid(torch.linspace(-1.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device='cuda'), 
+                                torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device='cuda'),
+                                indexing='ij')
+        v = safe_normalize(cube_to_dir(s, gx, gy))
+
+        tu = torch.atan2(v[..., 0:1], -v[..., 2:3]) / (2 * np.pi) + 0.5
+        tv = torch.acos(torch.clamp(v[..., 1:2], min=-1, max=1)) / np.pi
+        texcoord = torch.cat((tu, tv), dim=-1)
+
+        cubemap[s, ...] = dr.texture(latlong_map[None, ...], texcoord[None, ...], filter_mode='linear')[0]
+    return cubemap
+
+def load_and_preprocess_hdr(hdr_dir, env_strength, env_flip, env_rot, device, rotation180=False, inverse_env=False, flip_env=False):
+    """Load and preprocess the HDR environment map."""
+    if hdr_dir.endswith('.hdr') or hdr_dir.endswith('.exr'):
+        latlong_img = imageio_v3.imread(hdr_dir, flags=cv2.IMREAD_UNCHANGED, plugin='opencv')
+    elif hdr_dir.endswith('.jpg') or hdr_dir.endswith('.png'):
+        import skimage
+        latlong_img = skimage.io.imread(hdr_dir)[..., :3]
+        latlong_img = skimage.img_as_float(latlong_img)
+        latlong_img = np.power(latlong_img, 2.4).astype(np.float32)
+        latlong_img *= 2 # for mit dataset
+
+    if rotation180:
+        height, width, channels = latlong_img.shape
+        shift_amount = width // 2
+        shifted_hdr = np.zeros_like(latlong_img)
+        shifted_hdr[:, -shift_amount:, :] = latlong_img[:, :shift_amount, :]
+        shifted_hdr[:, :-shift_amount, :] = latlong_img[:, shift_amount:, :]
+        latlong_img = shifted_hdr
+
+    if inverse_env:
+        latlong_img = latlong_img[:, ::-1, :]
+        height, width, channels = latlong_img.shape
+        shift_amount = width // 2
+        shifted_hdr = np.zeros_like(latlong_img)
+        shifted_hdr[:, -shift_amount:, :] = latlong_img[:, :shift_amount, :]
+        shifted_hdr[:, :-shift_amount, :] = latlong_img[:, shift_amount:, :]
+        latlong_img = shifted_hdr
+
+    if flip_env:
+        latlong_img = latlong_img[:, ::-1, :].copy()
+
+    latlong_img = torch.tensor(latlong_img, dtype=torch.float32, device=device)
+    latlong_img *= env_strength
+
+    # Cleanup NaNs and Infs
+    latlong_img = torch.nan_to_num(latlong_img, nan=0.0, posinf=65504.0, neginf=0.0)
+    latlong_img = latlong_img.clamp(0.0, 65504.0)
+
+    if env_flip:
+        latlong_img = torch.flip(latlong_img, dims=[1])
+
+    if env_rot != 0:
+        lat_h, lat_w = latlong_img.shape[:2]
+        pixel_rot = int(lat_w * env_rot / 360)
+        latlong_img = torch.roll(latlong_img, shifts=pixel_rot, dims=1)
+
+    # Convert to cubemap
+    cubemap = latlong_to_cubemap(latlong_img, [512, 512])
+    return cubemap
+
+def prepare_metadata(hdr_dir, env_rot, env_flip, env_strength, fixed_pose, rotate_envlight, save_dir, prefix):
+    """Prepare metadata about the environment map processing."""
+    env_meta = {
+        'envmap': os.path.basename(hdr_dir),
+        'envmap_rot': env_rot,
+        'envmap_flip': env_flip,
+        'envmap_strength': env_strength,
+        'fixed_pose': fixed_pose,
+        'rotate_envlight': rotate_envlight,
+    }
+
+    if save_dir:
+        os.makedirs(save_dir, exist_ok=True)
+        meta_path = os.path.join(save_dir, f'{prefix}.meta.json')
+        with open(meta_path, 'w') as f:
+            json.dump(env_meta, f, indent=4)
+
+    return env_meta
+
+def rotate_y(a, device=None):
+    s, c = np.sin(a), np.cos(a)
+    return torch.tensor([[ c, 0, s, 0], 
+                         [ 0, 1, 0, 0], 
+                         [-s, 0, c, 0], 
+                         [ 0, 0, 0, 1]], dtype=torch.float32, device=device)
+
+def process_projected_envmap(cubemap, vec, c2w, y_rot, H, W):
+    """Process the camera-oriented projected environment map."""
+    vec_cam = vec.view(-1, 3) @ c2w[:3, :3].T
+    vec_query = (vec_cam @ y_rot[:3, :3].T).view(1, H, W, 3)
+    env_proj = dr.texture(cubemap.unsqueeze(0), -vec_query.contiguous(),
+                          filter_mode='linear', boundary_mode='cube')[0]
+    env_proj = torch.flip(env_proj, dims=[0, 1])
+    return env_proj
+
+def rgb2srgb(rgb):
+    return torch.where(rgb <= 0.0031308, 12.92 * rgb, 1.055 * rgb**(1/2.4) - 0.055)
+
+def reinhard(x, max_point=16):
+    y_rein = x * (1 + x / (max_point ** 2)) / (1 + x)
+    return y_rein
+
+def hdr_mapping(env_hdr, log_scale):
+    """Map HDR environment maps to LDR and logarithmic representations."""
+    env_ev0 = rgb2srgb(reinhard(env_hdr, max_point=16).clamp(0, 1))
+    env_log = rgb2srgb(torch.log1p(env_hdr) / np.log1p(log_scale)).clamp(0, 1)
+    return {
+        'env_hdr': env_hdr,    # Original HDR image
+        'env_ev0': env_ev0,    # LDR image after tone mapping
+        'env_log': env_log,    # Logarithmic scaling
+    }
+
+def process_environment_map(
+    hdr_dir,
+    resolution=(512, 512),
+    num_frames=1,
+    fixed_pose=True,
+    pose_file=None,
+    pose_list=None,
+    rotation_euler=None,
+    pose_offset=0,
+    pose_reset=False,
+    rotate_envlight=False,
+    env_format=['proj'],
+    log_scale=10000,
+    env_strength=1.0,
+    env_flip=True,
+    env_rot=180.0,
+    save_dir=None,
+    prefix='0000',
+    device=None,
+    rotation180=False,
+    ign_camera_pose=True,
+    inverse_env=False,
+    swap_type=0,
+    load_w2c=False,
+    flip_env=False,
+    remove_y_rotation=False,
+    reverse_y_rotation=False,
+    yaw_0=0,
+):
+    """
+    Preprocess HDR environment maps for rendering.
+    FIXME: Note that this function bakes in a flip and rotate operation for the environment light. Set to env_flip=True and env_rot=180 is considered as loading the original environment map.
+
+    Args:
+        hdr_dir (str): Path to the HDR environment map file.
+        resolution (tuple of int): Resolution of the output images (H, W).
+        num_frames (int): Number of frames to process.
+        fixed_pose (bool): Use a fixed camera pose (identity matrix) if True.
+        pose_file (str): Path to the camera pose file (JSON).
+        pose_offset (int): Offset for the pose frames in the pose file.
+        pose_reset (bool): Reset camera poses to be relative to the first frame.
+        rotate_envlight (bool): Rotate the environment light over frames if True.
+        env_format (list of str): Formats of the environment maps to generate ('proj', 'fixed', 'ball').
+        log_scale (int): Log scale factor for HDR mapping.
+        env_strength (float): Strength multiplier for the environment map.
+        env_flip (bool): Flip the environment map horizontally if True.
+        env_rot (float): Rotation angle for the environment map in degrees.
+        save_dir (str): Directory to save the processed images (optional).
+        prefix (str): Prefix for the output files (used if saving images).
+
+    Returns:
+        dict: A dictionary containing the processed environment maps and metadata.
+        {
+            'metadata': env_meta,
+            'fixed': mapping_results_for_fixed_envmap,  # Only if 'fixed' in env_format
+            'env_ldr': stacked_tensor_of_proj_env_ldr,  # Only if 'proj' in env_format
+            'env_log': stacked_tensor_of_proj_env_log,  # Only if 'proj' in env_format
+            'ball_env_ldr': stacked_tensor_of_ball_env_ldr,  # Only if 'ball' in env_format
+            'ball_env_log': stacked_tensor_of_ball_env_log,  # Only if 'ball' in env_format
+        }
+        Tensors are with shape (T, H, W, 3) in [0, 1]
+    """
+    H, W = resolution # (704, 1280)
+    if device is None:
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    vec = latlong_vec((H, W), device=device)
+
+    # Prepare camera poses
+    poses = prepare_camera_poses(
+        num_frames=num_frames,
+        fixed_pose=fixed_pose,
+        pose_file=pose_file,
+        pose_offset=pose_offset,
+        pose_reset=pose_reset,
+        device=device,
+        ign_camera_pose=ign_camera_pose,
+        swap_type=swap_type,
+        load_w2c=load_w2c,
+        remove_y_rotation=remove_y_rotation,
+        reverse_y_rotation=reverse_y_rotation,
+        yaw_0=yaw_0,
+        pose_list=pose_list,
+        rotation_euler=rotation_euler,
+    )
+
+    # Prepare rotations for the environment light # 57 * rot
+    rots = np.linspace(0, 2 * np.pi, num_frames) if rotate_envlight else [0] * num_frames
+
+    # Load and preprocess the HDR environment map
+    cubemap = load_and_preprocess_hdr(
+        hdr_dir=hdr_dir,
+        env_strength=env_strength,
+        env_flip=env_flip,
+        env_rot=env_rot,
+        device=device,
+        rotation180=rotation180,
+        inverse_env=inverse_env,
+        flip_env=flip_env,
+    )
+
+    # Prepare metadata
+    env_meta = prepare_metadata(
+        hdr_dir=hdr_dir,
+        env_rot=env_rot,
+        env_flip=env_flip,
+        env_strength=env_strength,
+        fixed_pose=fixed_pose,
+        rotate_envlight=rotate_envlight,
+        save_dir=save_dir,
+        prefix=prefix
+    )
+
+    # Initialize result dictionary
+    results = {
+        'metadata': env_meta,
+    }
+
+    # Prepare lists to collect per-frame tensors
+    if 'proj' in env_format:
+        proj_env_ldr_list = []
+        proj_env_log_list = []
+
+    # Process per-frame environment maps
+    for i in range(num_frames):
+        c2w = torch.from_numpy(poses[i]).float().to(device)
+        y_rot = rotate_y(rots[i], device=device)
+
+        if 'proj' in env_format:
+            env_proj = process_projected_envmap(cubemap, vec, c2w, y_rot, H, W)
+            mapping_results = hdr_mapping(env_proj, log_scale=log_scale)
+            proj_env_ldr_list.append(mapping_results['env_ev0'])
+            proj_env_log_list.append(mapping_results['env_log'])
+
+    if 'proj' in env_format:
+        results['env_ldr'] = torch.stack(proj_env_ldr_list, dim=0)
+        results['env_log'] = torch.stack(proj_env_log_list, dim=0)
+
+    return results
+
+def save_array_as_video(video_array, output_path: str, fps: int = 24):
+    """
+    video_array: t h w c, np.array or tensors
+    """
+    if isinstance(video_array, torch.Tensor):
+        video_array = video_array.cpu().numpy()
+    if video_array.dtype != np.uint8:
+        print("float 2 uint8")
+        # If the data range is [-1, 1]
+        if video_array.min() < 0:
+            video_array = ((video_array + 1) * 127.5).clip(0, 255).astype(np.uint8)
+        # If the data range is [0, 1]
+        else:
+            video_array = (video_array * 255).clip(0, 255).astype(np.uint8)
+    
+    try:
+        if not os.path.isfile(output_path):
+            imageio.mimsave(
+                output_path, 
+                [frame for frame in video_array], 
+                fps=fps, 
+                codec='libx264'
+            )
+            print("succeed to save vide")
+        print(f"video already exists in {output_path}")
+    except Exception as e:
+        print(f"fail to save video: {e}")
+
+def process_hdr(hdr_path, save_path, env_strength=1.0, inverse_env=False):
+    hdr_path = glob.glob(f'{hdr_path}*')[0]
+
+    if '.hdr' in hdr_path:
+        env_strength = env_strength / 3.0
+
+    num_of_frames = 57
+    ldr_list = []
+    hdr_log_list = []
+    env_dir_list = []
+
+    envlight_dict = process_environment_map(
+                        hdr_dir=hdr_path,
+                        resolution=(320, 576),
+                        num_frames=num_of_frames, # 1 for mit dataset
+                        fixed_pose=True,
+                        rotate_envlight=False,
+                        env_format=['proj', ],
+                        device='cuda',
+                        rotation180=False,
+                        inverse_env=inverse_env, # True for mit dataset, False for others
+                        log_scale=60000,
+                        env_strength=env_strength, # 1.0 for mit dataset
+                    ) # Tensors are with shape (T, H, W, 3) in [0, 1]
+    ldr_list = (envlight_dict['env_ldr'].cpu().numpy() * 255).astype(np.uint8)
+    hdr_log_list = (envlight_dict['env_log'].cpu().numpy() * 255).astype(np.uint8)
+    env_nrm = ((envmap_vec([320, 576], device='cpu').cpu().numpy()*0.5 + 0.5) * 255).astype(np.uint8)
+    for _ in range(num_of_frames):
+        env_dir_list.append(env_nrm)
+
+    os.makedirs(save_path, exist_ok=True)
+    ldr_video_path = os.path.join(save_path, "ldr_video_fix_first_frame.mp4")
+    hdr_log_video_path = os.path.join(save_path, "hdr_log_video_fix_first_frame.mp4")
+    env_dir_video_path = os.path.join(save_path, "env_dir_video_fix_first_frame.mp4")
+    
+    if os.path.exists(ldr_video_path):
+        os.remove(ldr_video_path)
+    if os.path.exists(hdr_log_video_path):
+        os.remove(hdr_log_video_path)
+    if os.path.exists(env_dir_video_path):
+        os.remove(env_dir_video_path)
+    
+    save_array_as_video(np.array(ldr_list),ldr_video_path)
+    save_array_as_video(np.array(hdr_log_list),hdr_log_video_path)
+    save_array_as_video(np.array(env_dir_list),env_dir_video_path)
+
+def parse_arguments() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Env maps processing script")
+    # Specific arguments
+    parser.add_argument(
+        "--env_dir",
+        type=str,
+        default=None,
+        help="Path to the directory containing the environment map."
+    )
+    parser.add_argument(
+        "--save_path",
+        type=str,
+        default=None,
+        help="Path to the directory where the processed environment maps will be saved."
+    )
+    parser.add_argument(
+        "--env_strength",
+        type=float,
+        default=1.0,
+        help="Strength of the environment map."
+    )
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_arguments()
+    process_hdr(args.env_dir, save_path=args.save_path, env_strength=args.env_strength)
+
+
+
+

tools/reorg_gbuffer_from_dr_delighting.py

@@ -0,0 +1,129 @@
+import os
+import pdb
+import json
+from tqdm import tqdm
+import shutil, glob
+import argparse
+
+def parse_arguments() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Text to world generation demo script")
+
+    parser.add_argument(
+        "--frame_dir",
+        type=str,
+        default=None,
+        help="."
+    )
+    parser.add_argument(
+        "--gbuffer_dir",
+        type=str,
+        default=None,
+        help="."
+    )
+    parser.add_argument(
+        "--tgt_dir",
+        type=str,
+        default=None,
+        help="."
+    )
+
+    return parser.parse_args()
+
+
+def collect_available_chunk_ids(gbuffer_seq_dir):
+    chunk_ids = set()
+    for path in glob.glob(os.path.join(gbuffer_seq_dir, "*.jpg")):
+        name = os.path.basename(path)
+        parts = name.split(".")
+        if len(parts) < 4:
+            continue
+        if parts[0].isdigit():
+            chunk_ids.add(int(parts[0]))
+    return sorted(chunk_ids)
+
+
+args = parse_arguments()
+frame_dir = args.frame_dir
+gbuffer_dir = args.gbuffer_dir
+tgt_dir = args.tgt_dir
+os.makedirs(tgt_dir, exist_ok=True)
+
+seq_list = sorted(os.listdir(frame_dir))
+for seq_name in tqdm(seq_list):
+    chunk_size = 57
+    overlap_n_frames = 0
+    step = chunk_size - overlap_n_frames
+
+    rgb_path_list = sorted(
+        glob.glob(os.path.join(frame_dir, f"{seq_name}/*.jpg")) +
+        glob.glob(os.path.join(frame_dir, f"{seq_name}/*.png"))
+    )
+    images_num = len(rgb_path_list)
+
+    gbuffer_file_path = f"{gbuffer_dir}/{seq_name}"
+    available_chunk_ids = collect_available_chunk_ids(gbuffer_file_path)
+
+    if len(available_chunk_ids) == 0:
+        print(f"Skip {seq_name}: no gbuffer files found in {gbuffer_file_path}")
+        continue
+
+    max_valid_chunk_idx = max(available_chunk_ids)
+
+    print("Processing", seq_name)
+    for t_gbuffer in ["Base Color", "normal", "depth", "Roughness", "Metallic", "env", "images_4"]:
+        os.system(f"rm -rf {tgt_dir}/{seq_name}*/{t_gbuffer}")
+
+    for image_idx in range(images_num):
+        raw_chunk_idx = image_idx // step
+        frame_in_chunk = image_idx % step
+
+        if raw_chunk_idx > max_valid_chunk_idx:
+            print(
+                f"Stop {seq_name} at image_idx={image_idx}: "
+                f"raw_chunk_idx={raw_chunk_idx} exceeds max_valid_chunk_idx={max_valid_chunk_idx}"
+            )
+            break
+
+        chunk_idx = raw_chunk_idx
+        frame_idx = frame_in_chunk
+        seq_name_now = seq_name + f".{chunk_idx}"
+
+        gbuffer_src_map = {}
+        missing_file = False
+        for gbuffer, t_gbuffer in [
+            ("basecolor", "Base Color"),
+            ("normal", "normal"),
+            ("depth", "depth"),
+            ("roughness", "Roughness"),
+            ("metallic", "Metallic"),
+        ]:
+            img_name = f"{chunk_idx:04d}.{frame_in_chunk:04d}.{gbuffer}.jpg"
+            src_path = os.path.join(gbuffer_file_path, img_name)
+            if not os.path.exists(src_path):
+                print(f"Stop {seq_name_now}: missing gbuffer file {src_path}")
+                missing_file = True
+                break
+            gbuffer_src_map[t_gbuffer] = src_path
+
+        if missing_file:
+            break
+
+        for t_gbuffer in ["Base Color", "normal", "depth", "Roughness", "Metallic", "env", "images_4"]:
+            if os.path.exists(f"{tgt_dir}/{seq_name_now}/{t_gbuffer}") and frame_idx == 0:
+                shutil.rmtree(f"{tgt_dir}/{seq_name_now}/{t_gbuffer}")
+            os.makedirs(f"{tgt_dir}/{seq_name_now}/{t_gbuffer}", exist_ok=True)
+
+        for t_gbuffer, src_path in gbuffer_src_map.items():
+            shutil.copy(
+                src_path,
+                f"{tgt_dir}/{seq_name_now}/{t_gbuffer}/frame_{frame_idx:04d}.jpg"
+            )
+
+        rgb_ext = os.path.basename(rgb_path_list[image_idx]).split(".")[-1]
+        shutil.copy(
+            rgb_path_list[image_idx],
+            f"{tgt_dir}/{seq_name_now}/images_4/frame_{frame_idx:04d}.{rgb_ext}"
+        )
+
+
+