Support multiple prompts

[jitendra-jalwaniya]

This PR adds supports for enabling inference scripts to provide multiple prompt-inputs to generate videos via txt files or GCS locations.

It will be useful for integrating VBench evals into this repo.

[gemini-code-assist]

Code Review

This pull request introduces support for loading generation prompts from a file (prompt_file) across multiple video generation scripts (generate_ltx2.py, generate_ltx_video.py, and generate_wan.py), alongside updating configuration files and improving GCS upload paths. It also fixes latent decoding in wan_pipeline.py by concatenating addressable shards. The review feedback highlights a critical performance bottleneck and a NameError in generate_ltx_video.py caused by initializing the pipeline inside the prompt loop. Additionally, the reviewer recommends raising a ValueError in max_utils.py to prevent downstream IndexError crashes when no prompts are found, and suggests removing an unused last_out variable in generate_ltx2.py.

[mbohlool]

Summary & Thanks

Thanks for putting this together! Adding prompt-file support across Wan and LTX will be very valuable for automated evaluation workflows like VBench.

The recent update guarding the shard concatenation in WanPipeline._decode_latents_to_video (video.addressable_shards[0].data.shape[0] < video.shape[0]) nicely resolves the replicated tensor duplication issue that caused the WanKvCacheTest failure.

There are still two blocking issues (one runtime IAM permission issue with GCS and one pipeline API signature inconsistency) along with a few important improvements before this can be merged.

---

🔴 Required Changes (Blockers)

1. GCS IAM Permission failure: storage_client.get_bucket

Location: src/maxdiffusion/max_utils.py (around line 383)

bucket_name, prefix_name = parse_gcs_bucket_and_prefix(prompt_file_path)
storage_client = storage.Client()
bucket = storage_client.get_bucket(bucket_name)
blob = bucket.blob(prefix_name)

Issue:
storage_client.get_bucket(bucket_name) performs an HTTP GET request to retrieve bucket metadata, which requires storage.buckets.get permission. Service accounts and TPU VM workloads are typically granted object-level reader access (storage.objects.get / Storage Object Viewer) but lack bucket metadata permissions. This will trigger a 403 Forbidden error when attempting to download prompts from gs://....

Fix:
Use client-side bucket instantiation storage_client.bucket(bucket_name) instead (matching how upload_file_to_gcs works):

storage_client = storage.Client()
bucket = storage_client.bucket(bucket_name)
blob = bucket.blob(prefix_name)
content = blob.download_as_text()

---

2. Pipeline API signature inconsistency in LTXVideoPipeline

Location: src/maxdiffusion/pipelines/ltx_video/ltx_video_pipeline.py (around line 630 & line 1157)

Issue:
In LTXMultiScalePipeline.__call__, prompt was explicitly added as a typed parameter:

def __call__(
    self,
    ...,
    prompt: Optional[Union[str, List[str]]] = None,
) -> Any:

However, in LTXVideoPipeline.__call__, prompt was not added to the method signature and is extracted via kwargs.get("prompt", self.config.prompt). This creates an inconsistent public interface between the base and multi-scale pipelines and prevents type checkers/IDEs from discovering the parameter.

Fix:
Explicitly add prompt: Optional[Union[str, List[str]]] = None to the parameter list of LTXVideoPipeline.__call__.

---

🟡 Important Improvements

1. Add explicit timeout to requests.get (src/maxdiffusion/max_utils.py:393):

   response = requests.get(prompt_file_path, timeout=30)

   Without a timeout, dropped packets or unresponsive HTTP servers can cause TPU worker threads to hang indefinitely.

2. Defensive handling against None config values in get_gcs_output_path (src/maxdiffusion/max_utils.py:347):
   If output_dir or base_output_directory is None in the config, getattr(config, "output_dir", "") returns None, leading to AttributeError: 'NoneType' object has no attribute 'startswith'.

   output_dir = getattr(config, "output_dir", "") or ""
   base_output_dir = getattr(config, "base_output_directory", "") or ""

3. Deduplicate batch chunking and padding:
   The exact same prompt slicing and tail-batch replication logic is duplicated across generate_wan.py (inference_generate_video and run) and generate_ltx2.py. Consider extracting a reusable generator in max_utils.py, e.g.:

   def chunk_and_pad(items: list, batch_size: int):
     for i in range(0, len(items), batch_size):
       chunk = items[i : i + batch_size]
       actual_len = len(chunk)
       padded_chunk = chunk + [chunk[-1]] * (batch_size - actual_len) if actual_len < batch_size else chunk
       yield i, padded_chunk, actual_len

4. Expand unit test coverage in tests/maxdiffusion_utils_test.py:
   Currently only the empty string ValueError is tested. Please consider adding unit tests for:

   • Reading valid multi-line prompt files (handling newlines, empty lines, and trimming whitespace).
   • Fallback to default_prompt when prompt file contains only whitespace.
   • get_gcs_output_path path resolution.

---

🟢 Nits & Suggestions

• 💡 Support # comments in prompt files: In load_prompts, consider ignoring comment lines: if line.strip() and not line.strip().startswith('#'). This makes prompt files much friendlier for benchmark suites.
• 🟢 Dead config in base_wan_animate.yml: prompt_file: "" was added to base_wan_animate.yml, but generate_wan_animate.py does not load or use prompt_file. Either implement prompt file support in generate_wan_animate.py or remove it from the YAML.

[jitendra-jalwaniya]

Thanks for the extensive review @mbohlool. I've addressed all your comments.

For wan_animate, earlier prompt_file was added to the config but it wasn't used. I've added its usage in generate_wan_animate now.