Apply suggestions from code review

Co-authored-by: Linoy Tsaban <[email protected]>

Author: sayakpaul

examples/dreambooth/README_flux.md

@@ -319,8 +319,9 @@ accelerate launch train_dreambooth_lora_flux_kontext.py \
   --cache_latents \
   --learning_rate=1e-4 \
   --lr_scheduler="constant" \
-  --lr_warmup_steps=0 \
-  --max_train_steps=500 \
+  --lr_warmup_steps=200 \
+  --max_train_steps=1000 \
+  --rank=16\
   --seed="0" 
 ```
 

examples/dreambooth/train_dreambooth_lora_flux_kontext.py

@@ -795,7 +795,8 @@ def __init__(
             column_names = dataset["train"].column_names
 
             # 6. Get the column names for input/target.
-            # TODO: add validation for `cond_image_column`
+if args.cond_image_column is not None and not args.cond_image_column in column_names:
+      raise ValueError(f"`--cond_image_column` value '{args.cond_image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}")
             if args.image_column is None:
                 image_column = column_names[0]
                 logger.info(f"image column defaulting to {image_column}")