[Flux] Remove generate_image script and add CI test for inference (#1746)

As follow up of #1726

Author: wwwjn

torchtitan/experiments/flux/tests/integration_tests.py

@@ -87,6 +87,19 @@ def build_flux_test_list() -> list[OverrideDefinitions]:
             "Flux Validation Test",
             "validation",
         ),
+        OverrideDefinitions(
+            [
+                [
+                    "--checkpoint.enable",
+                ],
+                [
+                    # placeholder for the generation script's generate step
+                ],
+            ],
+            "Flux Generation script test",
+            "test_generate",
+            ngpu=2,
+        ),
     ]
     return integration_tests_flavors
 
@@ -116,6 +129,15 @@ def run_single_test(test_flavor: OverrideDefinitions, full_path: str, output_dir
         cmd = f"CONFIG_FILE={full_path} NGPU={test_flavor.ngpu} LOG_RANK={all_ranks} ./torchtitan/experiments/flux/run_train.sh"
         # dump compile trace for debugging purpose
         cmd = f'TORCH_TRACE="{output_dir}/{test_name}/compile_trace" ' + cmd
+
+        # save checkpoint (idx == 0) and load it for generation (idx == 1)
+        if test_name == "test_generate" and idx == 1:
+            # For flux generation, test using inference script
+            cmd = (
+                f"CONFIG_FILE={full_path} NGPU={test_flavor.ngpu} LOG_RANK={all_ranks} "
+                f"./torchtitan/experiments/flux/inference/run_infer.sh"
+            )
+
         cmd += " " + model_arg
         cmd += " " + dump_folder_arg
         cmd += " " + random_init_encoder_arg
@@ -124,10 +146,10 @@ def run_single_test(test_flavor: OverrideDefinitions, full_path: str, output_dir
         cmd += " " + tokenzier_path_arg
         if override_arg:
             cmd += " " + " ".join(override_arg)
+
         logger.info(
             f"=====Flux Integration test, flavor : {test_flavor.test_descr}, command : {cmd}====="
         )
-
         result = _run_cmd(cmd)
         logger.info(result.stdout)
         if result.returncode != 0: