revert

Author: wwwjn

torchtitan/models/flux/inference/infer.py

@@ -26,21 +26,17 @@ def inference(config: JobConfig):
     world_size = int(os.environ["WORLD_SIZE"])
     global_rank = int(os.environ["RANK"])
 
-    single_prompt_mode = config.inference.prompt is not None
-
-    # Use single prompt if specified, otherwise read from file
-    if single_prompt_mode:
-        original_prompts = [config.inference.prompt]
-        logger.info(f"Using single prompt: {config.inference.prompt}")
-        bs = 1
-        # If only single prompt, each rank will generate an image with the same prompt
-        prompts = original_prompts
-    else:
-        original_prompts = open(config.inference.prompts_path).readlines()
-        logger.info(f"Reading prompts from: {config.inference.prompts_path}")
-        bs = config.inference.local_batch_size
-        # Distribute prompts across processes using round-robin assignment
-        prompts = original_prompts[global_rank::world_size]
+    original_prompts = open(config.inference.prompts_path).readlines()
+    logger.info(f"Reading prompts from: {config.inference.prompts_path}")
+    if len(original_prompts) < world_size:
+        raise ValueError(
+            f"Number of prompts ({len(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]
 
     total_prompts = len(original_prompts)
 
@@ -58,13 +54,7 @@ def inference(config: JobConfig):
             config.inference.save_img_folder,
         )
         # Create mapping from local indices to global prompt indices
-        if single_prompt_mode:
-            # In single prompt mode, all ranks process the same prompt (index 0)
-            # But each rank generates a different image (different seed/rank)
-            global_ids = [0] * len(prompts)
-        else:
-            # In multi-prompt mode, use round-robin distribution
-            global_ids = list(range(global_rank, total_prompts, world_size))
+        global_ids = list(range(global_rank, total_prompts, world_size))
 
         for i in range(0, len(prompts), bs):
             images = generate_image(

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

torchtitan/models/flux/job_config.py

@@ -62,8 +62,6 @@ class Inference:
     """Path to save the inference results"""
     prompts_path: str = "./torchtitan/experiments/flux/inference/prompts.txt"
     """Path to file with newline separated prompts to generate images for"""
-    prompt: str = ""
-    """Single prompt to generate image for. If specified, takes precedence over prompts_path"""
     local_batch_size: int = 2
     """Batch size for inference"""
     img_size: int = 256

torchtitan/models/flux/run_infer.sh

@@ -10,7 +10,7 @@ set -ex
 # use envs as local overrides for convenience
 # e.g.
 # LOG_RANK=0,1 NGPU=4 ./torchtitan/models/flux/run_train.sh
-NGPU=${NGPU:-"4"}
+NGPU=${NGPU:-"8"}
 export LOG_RANK=${LOG_RANK:-0}
 CONFIG_FILE=${CONFIG_FILE:-"./torchtitan/models/flux/train_configs/debug_model.toml"}
 

torchtitan/models/flux/run_train.sh

@@ -10,7 +10,7 @@ set -ex
 # use envs as local overrides for convenience
 # e.g.
 # LOG_RANK=0,1 NGPU=4 ./torchtitan/experiments/flux/run_train.sh
-NGPU=${NGPU:-"4"}
+NGPU=${NGPU:-"8"}
 export LOG_RANK=${LOG_RANK:-0}
 CONFIG_FILE=${CONFIG_FILE:-"./torchtitan/models/flux/train_configs/debug_model.toml"}