[FLUX] Add single prompt mode for FLUX inference and inference CI test

tests/integration_tests/flux.py

@@ -41,6 +41,17 @@ def build_flux_test_list() -> list[OverrideDefinitions]:
             "Flux Validation Test",
             "validation",
         ),
+        OverrideDefinitions(
+            [
+                [
+                    "--checkpoint.enable",
+                ],
+                [],
+            ],
+            "Flux Generation script test",
+            "test_generate",
+            ngpu=2,
+        ),
     ]
     return integration_tests_flavors
 

torchtitan/models/flux/inference/infer.py

@@ -28,6 +28,13 @@ def inference(config: JobConfig):
     original_prompts = open(config.inference.prompts_path).readlines()
     total_prompts = len(original_prompts)
 
+    if total_prompts < world_size:
+        raise ValueError(
+            f"Number of prompts ({total_prompts}) must be >= number of ranks ({world_size}). "
+            f"FSDP all-gather will hang if some ranks have no prompts to process."
+        )
+
+    bs = config.inference.local_batch_size
     # Distribute prompts across processes using round-robin assignment
     prompts = original_prompts[global_rank::world_size]
 
@@ -39,13 +46,11 @@ def inference(config: JobConfig):
 
     if prompts:
         # Generate images for this process's assigned prompts
-        bs = config.inference.local_batch_size
 
         output_dir = os.path.join(
             config.job.dump_folder,
             config.inference.save_img_folder,
         )
-
         # Create mapping from local indices to global prompt indices
         global_ids = list(range(global_rank, total_prompts, world_size))
 

torchtitan/models/flux/inference/prompts.txt

@@ -2,27 +2,3 @@ A serene mountain landscape at sunset with a crystal clear lake reflecting the g
 A futuristic cityscape with flying cars and neon lights illuminating the night sky
 A cozy cafe interior with steam rising from coffee cups and warm lighting
 A magical forest with glowing mushrooms and fireflies dancing between ancient trees
-A peaceful beach scene with turquoise waves and palm trees swaying in the breeze
-A steampunk-inspired mechanical dragon soaring through clouds
-A mystical library with floating books and magical artifacts
-A Japanese garden in spring with cherry blossoms falling gently
-A space station orbiting a colorful nebula
-A medieval castle on a hilltop during a dramatic thunderstorm
-A underwater scene with bioluminescent creatures and coral reefs
-A desert oasis with a majestic palace and palm trees
-A cyberpunk street market with holographic signs and diverse crowds
-A cozy winter cabin surrounded by snow-covered pine trees
-A fantasy tavern filled with unique characters and magical atmosphere
-A tropical rainforest with exotic birds and waterfalls
-A steampunk airship navigating through storm clouds
-A peaceful zen garden with a traditional Japanese tea house
-A magical potion shop with bubbling cauldrons and mysterious ingredients
-A futuristic space colony on Mars with domed habitats
-A mystical temple hidden in the clouds
-A vintage train station with steam locomotives and period architecture
-A magical bakery with floating pastries and enchanted ingredients
-A peaceful countryside scene with rolling hills and a rustic farmhouse
-A underwater city with advanced technology and marine life
-A fantasy marketplace with magical creatures and exotic goods
-A peaceful meditation garden with lotus flowers and koi ponds
-A steampunk laboratory with intricate machinery and glowing elements