Added the num-iteration flag and cleaned up benchmark timing

[lvaidya2910]

Summary

Adds a new --num-iterations option to Cosmos3 inference for collecting multiple benchmark measurements in one invocation. Two additions that make benchmark timing clean and attributable

Finer-grained generate_batch sub-timers — prepare_inputs, resolve_upsample_tasks, save_outputs. These pinpoint where per-iteration time goes (and originally exposed the ~0.7s first-iteration overhead as output persistence, not compute).
Save the artifact during the last warmup pass — in benchmark mode with warmup >= 1, the final warmup pass runs the real save path (dir creation, guardrail lazy-load, first ffmpeg spawn, video write). Every measured iteration is then pure generate+decode, so iterations are uniform with no first-iteration skew. The artifact is still produced (from the deterministic, identical-seed warmup pass). --profile wraps only the (non-measured) save pass, so it never perturbs a measured iteration.
Behavior Defaults to one measured iteration. Values greater than one require --benchmark. Warmup runs remain separate from measured iterations. Each iteration executes the complete generation and decode workload. Only the first measured iteration writes output artifacts. Benchmark timing contains one entry per measured iteration. Copies mutable batch containers between runs to prevent one iteration from affecting another. Supports standard, transfer, and reasoner inference paths. Example

torchrun --nproc_per_node=1 \
  -m cosmos3.scripts.inference \
  -i inputs/omni/t2v.json \
  -o outputs/omni \
  --checkpoint-path Cosmos3-Edge \
  --warmup 1 \
  --benchmark \
  --num-iterations 4

Averaging policy (get_timer_results)

warmup >= 1: measured iterations are uniform → average is the plain mean.
warmup == 0: there is no warmup pass, so the first measured iteration is the cold pass (it saves the artifact + pays compile/setup) → it is dropped from average. Raw per-sample values are always kept in all.

Validation (RTX Pro 6000, Cosmos3-Edge, --warmup 1 --num-iterations 4 --benchmark)

Measured generate_batch: [51.5, 52.7, 53.1, 53.5]s — no first-iteration spike (first is fastest; remaining creep is thermal).
save_outputs for measured iterations: empty; save cost (1.7s) is recorded once under [warmup] save_outputs.
Artifact (vision.mp4) produced; exit 0.

[lvaidya2910]

Error summary caused by running action modatilies

The failure is caused by reusing the same input batch across multiple warmup and benchmark iterations.
During the first inference pass, the video is correctly loaded as torch.uint8 and normalized to floating point in the [-1, 1] range. However, _finalize_data_batch() previously made only a shallow copy of the input dictionary. Its video list remained shared with the original reusable batch.

As a result:

1. The first pass replaced the shared uint8 video with a normalized floating-point tensor.
2. The is_preprocessed=True flag was set only on the copied dictionary.
3. The original batch therefore contained floating-point video but did not contain the preprocessing flag.
4. On the next warmup or measured iteration, the model treated the video as raw input and asserted that it must be uint8:

AssertionError: Video data is not in uint8 format.

The successful Saved sample outputs message confirms that downloading and the first inference pass completed correctly.

Relationship to num-iterations

num-iterations did not create the underlying bug, but additional iterations made it easier to reproduce. The same issue could occur whenever the total number of inference passes exceeded one, including warmup passes.
The save/profile changes requested during review only changed which pass saved output. They did not cause the dtype conversion issue.

Fix

_finalize_data_batch() now copies mutable list containers before preprocessing. Each inference pass can replace its local video-list entries without modifying the raw batch reused by later passes.
A regression test was added to verify that normalizing the first finalized batch leaves the source batch as torch.uint8 for the next iteration.