Allow images; Remove LLM generated prefixes; Allow JSON/JSONL; Fix bugs (#158)

* update

* update

* fix

* Update finetrainers/dataset.py

* update

* argument for enabling remove of common llm prefixes

* update

* make new generator for validation

* update

* update

* update

* update

* update

finetrainers/args.py

@@ -41,6 +41,7 @@ class Args:
     caption_dropout_p: float = 0.00
     caption_dropout_technique: str = "empty"
     precompute_conditions: bool = False
+    remove_common_llm_caption_prefixes: bool = False
 
     # Dataloader arguments
     dataloader_num_workers: int = 0
@@ -48,8 +49,8 @@ class Args:
 
     # Diffusion arguments
     flow_resolution_shifting: bool = False
-    flow_base_image_seq_len: int = 256
-    flow_max_image_seq_len: int = 4096
+    flow_base_seq_len: int = 256
+    flow_max_seq_len: int = 4096
     flow_base_shift: float = 0.5
     flow_max_shift: float = 1.15
     flow_shift: float = 1.0
@@ -143,11 +144,24 @@ class Args:
                 "caption_dropout_p": self.caption_dropout_p,
                 "caption_dropout_technique": self.caption_dropout_technique,
                 "precompute_conditions": self.precompute_conditions,
+                "remove_common_llm_caption_prefixes": self.remove_common_llm_caption_prefixes,
             },
             "dataloader_arguments": {
                 "dataloader_num_workers": self.dataloader_num_workers,
                 "pin_memory": self.pin_memory,
             },
+            "diffusion_arguments": {
+                "flow_resolution_shifting": self.flow_resolution_shifting,
+                "flow_base_seq_len": self.flow_base_seq_len,
+                "flow_max_seq_len": self.flow_max_seq_len,
+                "flow_base_shift": self.flow_base_shift,
+                "flow_max_shift": self.flow_max_shift,
+                "flow_shift": self.flow_shift,
+                "flow_weighting_scheme": self.flow_weighting_scheme,
+                "flow_logit_mean": self.flow_logit_mean,
+                "flow_logit_std": self.flow_logit_std,
+                "flow_mode_scale": self.flow_mode_scale,
+            },
             "training_arguments": {
                 "training_type": self.training_type,
                 "seed": self.seed,
@@ -351,6 +365,11 @@ def _add_dataset_arguments(parser: argparse.ArgumentParser) -> None:
         action="store_true",
         help="Whether or not to precompute the conditionings for the model.",
     )
+    parser.add_argument(
+        "--remove_common_llm_caption_prefixes",
+        action="store_true",
+        help="Whether or not to remove common LLM caption prefixes.",
+    )
 
 
 def _add_dataloader_arguments(parser: argparse.ArgumentParser) -> None:
@@ -373,6 +392,36 @@ def _add_diffusion_arguments(parser: argparse.ArgumentParser) -> None:
         action="store_true",
         help="Resolution-dependant shifting of timestep schedules.",
     )
+    parser.add_argument(
+        "--flow_base_seq_len",
+        type=int,
+        default=256,
+        help="Base image/video sequence length for the diffusion model.",
+    )
+    parser.add_argument(
+        "--flow_max_seq_len",
+        type=int,
+        default=4096,
+        help="Maximum image/video sequence length for the diffusion model.",
+    )
+    parser.add_argument(
+        "--flow_base_shift",
+        type=float,
+        default=0.5,
+        help="Base shift as described in [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://arxiv.org/abs/2403.03206)",
+    )
+    parser.add_argument(
+        "--flow_max_shift",
+        type=float,
+        default=1.15,
+        help="Maximum shift as described in [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://arxiv.org/abs/2403.03206)",
+    )
+    parser.add_argument(
+        "--flow_shift",
+        type=float,
+        default=1.0,
+        help="Shift value to use for the flow matching timestep schedule.",
+    )
     parser.add_argument(
         "--flow_weighting_scheme",
         type=str,
@@ -722,6 +771,7 @@ def _map_to_args_type(args: Dict[str, Any]) -> Args:
     result_args.caption_dropout_p = args.caption_dropout_p
     result_args.caption_dropout_technique = args.caption_dropout_technique
     result_args.precompute_conditions = args.precompute_conditions
+    result_args.remove_common_llm_caption_prefixes = args.remove_common_llm_caption_prefixes
 
     # Dataloader arguments
     result_args.dataloader_num_workers = args.dataloader_num_workers
@@ -729,6 +779,11 @@ def _map_to_args_type(args: Dict[str, Any]) -> Args:
 
     # Diffusion arguments
     result_args.flow_resolution_shifting = args.flow_resolution_shifting
+    result_args.flow_base_seq_len = args.flow_base_seq_len
+    result_args.flow_max_seq_len = args.flow_max_seq_len
+    result_args.flow_base_shift = args.flow_base_shift
+    result_args.flow_max_shift = args.flow_max_shift
+    result_args.flow_shift = args.flow_shift
     result_args.flow_weighting_scheme = args.flow_weighting_scheme
     result_args.flow_logit_mean = args.flow_logit_mean
     result_args.flow_logit_std = args.flow_logit_std
@@ -743,6 +798,7 @@ def _map_to_args_type(args: Dict[str, Any]) -> Args:
     result_args.train_steps = args.train_steps
     result_args.rank = args.rank
     result_args.lora_alpha = args.lora_alpha
+    result_args.target_modules = args.target_modules
     result_args.gradient_accumulation_steps = args.gradient_accumulation_steps
     result_args.gradient_checkpointing = args.gradient_checkpointing
     result_args.checkpointing_steps = args.checkpointing_steps