Graduate flux from experiment folder to core torchtitan

[wwwjn]

FLUX.1 model is a diffusion model, which is different from language models and needs to extend train.py as needed.

• dataset.py
  • moved to datasets/flux_dataset.py
• tokenizer.py
  • Keep it under models/flux folder
• integration_test.py
  • moved to tests/integration_tests/flux.py
  • Because FLUX used a separate train.py and run_train.sh, I kept a copy of run_tests() instead of generalizing integration_tests/run_test.py
• train.py
  • Keep it under models/flux folder
• validate.py
• sample.py
  • moved to inference/

[tianyu-l]

Thanks! Left some comments.

[tianyu-l]

could you also rebase onto #1871

[wwwjn]

could you also rebase onto #1871

Sure, will rebase later

[tianyu-l]

Looks really good. Thanks for the refactor! Had some minor comments.

[tianyu-l]

For more logical organization of files, you think we can put clip and t5 config.json into tests/assets/flux_test_encoders/?

[tianyu-l]

need to rebase onto #1851