flashdreams-postprocessing

FlashDreams Post-Processing

Use this skill when adding a video post-processor or changing the runner post-processing stream. The reference implementation is integrations_v2/flashvsr/impl/postprocess.py.

Mental Model

A post-processor is usually three classes, not one class inheriting everything:

• VideoPostProcessorConfig: serializable config and CLI surface. It sets _target to the processor factory and declares fields, output_spec(), requires_all_ranks(), and validate_execution().
• VideoPostProcessor: lightweight factory created from config. Its job is start(spec) -> VideoPostProcessorSession.
• VideoPostProcessorSession: mutable per-stream runtime. It owns buffers, caches, lazy model instances, counters, and process() / flush().

Keep stream state in the session. Do not store per-rollout mutable state on the config or processor factory.

Implementation Steps

1. Pick a home:

   • Generic reusable post-processing belongs under flashdreams/flashdreams/infra/postprocess/.
   • Model-specific processors belong in their integration, for example integrations/<name>/<pkg>/postprocess.py.

2. Define a config subclass:

   @dataclass(kw_only=True)
   class MyPostProcessorConfig(VideoPostProcessorConfig):
       _target: type["MyPostProcessor"] = field(
           default_factory=lambda: MyPostProcessor
       )
    
       scale: int = 2
    
       def output_spec(self, input_spec: VideoSpec) -> VideoSpec:
           return VideoSpec(
               height=input_spec.height * self.scale,
               width=input_spec.width * self.scale,
               fps=input_spec.fps,
               channels=input_spec.channels,
           )

   Override:

   • output_spec() when spatial size, channels, or timing changes.
   • requires_all_ranks() when the processor must run on nonzero ranks under torchrun.
   • validate_execution() to reject unsupported distributed or shape modes early.

   Every config inherits device from VideoPostProcessorConfig. Consumers copy cached preset configs, assign device before setup(), and keep it fixed for the processor's lifetime. Backends that need a native device representation convert the shared string when they construct their runtime resources. Never mutate the preset object returned by the registry.

3. Define the processor factory:

   class MyPostProcessor(VideoPostProcessor[MyPostProcessorConfig]):
       def start(self, spec: VideoSpec) -> VideoPostProcessorSession:
           return _MyPostProcessorSession(self.config, spec)

4. Define the session:

   class _MyPostProcessorSession(VideoPostProcessorSession):
       def __init__(self, config: MyPostProcessorConfig, spec: VideoSpec) -> None:
           self._config = config
           self._spec = spec
           self._buffer: Tensor | None = None
    
       def process(self, chunk: VideoChunk) -> list[VideoChunk]:
           ...
    
       def flush(self) -> list[VideoChunk]:
           ...

   process() is synchronous but may return []: that means it consumed the input chunk and is buffering frames until a later chunk or flush() can complete an output window.

5. Handle layouts at the boundary:

   • Accept VideoChunk.tensor in chunk.layout.
   • Use to_bvtchw() only as a generic boundary helper.
   • Convert once into the processor's native layout, make it contiguous if the model kernels require that, and keep internal buffers in that native layout.
   • Document any forced .contiguous() because it can copy.

6. Return VideoChunks:

   • Preserve [-1, 1] value range unless the API is intentionally changed.
   • Set the correct layout.
   • Carry metadata only if it helps downstream processors or provenance.

7. Register presets when users should select it from CLI:

   [project.entry-points."flashdreams.postprocess_presets"]
   "my-postprocessor-v1" = "my_pkg.postprocess:POSTPROCESS_PRESET_MY_V1"

   The exported object must be a VideoPostProcessorConfig, for example:

   POSTPROCESS_PRESET_MY_V1 = MyPostProcessorConfig(...)

   Users select it with --postprocess.preset my-postprocessor-v1.

   V2 applications that also expose device placement should resolve both at launch time through the shared factory:

   postprocess = VideoPostprocessChainConfig.from_preset(
       args.postprocess_preset,
       device=args.postprocess_device,
   )

   This copies the registered preset, assigns its shared device field, and stores it before setup() or prepare() can load model resources. Keep output-channel selection and application metadata pairing in the application; the session runner cannot infer those semantics from an ordered list of model results.

Runner Interaction

Runners create a VideoPostprocessStream through create_runner_postprocess_stream(). The stream:

• creates one chain session for whole-stream processing, or one session per view when postprocess_per_view=True;
• calls session.process(VideoChunk(...)) for each AR output;
• turns [] into a zero-frame tensor so process() remains tensor-only;
• skips collecting zero-time tensors in _append_if_nonempty();
• calls flush() once at end-of-stream and appends any tail output.

Use postprocess_output_layout to describe the runner's decoded output layout. Use postprocess_per_view=True for bvtchw outputs when each camera/view needs an independent processor session.

Tests

Add CPU-safe tests unless the behavior genuinely requires a GPU:

• Config/preset discovery: flashdreams/tests/test_postprocess_presets.py.
• Stream contract and buffering: flashdreams/tests/test_postprocess_stream.py.
• Processor-specific CPU fakes: integrations/<name>/tests/test_postprocess.py.
• Runner distributed skip/all-rank behavior: flashdreams/tests/test_runner_postprocess.py.

Every pytest test must use exactly one marker: ci_cpu, ci_gpu, or manual. Prefer fake processor builders for CPU tests instead of loading checkpoints.

Useful focused validation:

uv run pytest flashdreams/tests/test_runner_postprocess.py \
  flashdreams/tests/test_postprocess_stream.py \
  flashdreams/tests/test_postprocess_presets.py \
  integrations/<name>/tests/test_postprocess.py