Add a new method¶
Before starting to add a new method, we highly recommend reading the Inference pipeline overview and Config system pages first to obtain an overview of the system’s architecture.
FlashDreams aims to offer researchers a codebase that they can utilize to extend and develop novel video and world models. Our vision is for users to establish a standalone repository that imports FlashDreams as a dependency and overrides pipeline components (such as encoders, transformers, or decoders) to cater to specific functionality requirements of the new approach. We encourage to maintain methods externally rather than pushing changes directly into the integrations/ directory of this repository.
However, if any new features requires modifications to the core FlashDreams infra or introduce generally useful components (such as the TAEHV decoder), we encourage you to submit a PR to enable others to benefit from them.
File structure¶
We recommend the following file structure for new methods:
customized_method/
├── customized_method/
│ ├── __init__.py
│ ├── runner.py # Custom runner and runner config
│ ├── config.py # Pipeline config literals
│ ├── pipeline.py # [optional] custom pipeline behavior
│ ├── transformer.py # [optional] DiT network for diffusion.
│ ├── encoder.py # [optional] custom control encoder
│ ├── decoder.py # [optional] custom outputdecoder
│ └── ...
└── pyproject.toml
Add optional files are only if customization beyond what is available in FlashDreams is required.
Most integrations can use the base StreamInferencePipeline directly and only provide model components plus config literals.
As explained in Config system, a method typically defines a pipeline config and a runner config.
The runner handles CLI-facing I/O and runtime loops.
from dataclasses import dataclass, field
from flashdreams.infra.pipeline import StreamInferencePipeline
from flashdreams.infra.runner import Runner, RunnerConfig
@dataclass(kw_only=True)
class CustomizedMethodRunnerConfig(RunnerConfig):
"""Runner config for my method."""
_target: type["CustomizedMethodRunner"] = field(default_factory=lambda: CustomizedMethodRunner)
prompt: str = "A cat surfing."
total_blocks: int = 60
class CustomizedMethodRunner(Runner[CustomizedMethodRunnerConfig, StreamInferencePipeline]):
def run(self) -> None:
cfg = self.config
# 1. Initialize the autoregressive cache
cache = self.pipeline.initialize_cache(text=[cfg.prompt])
# 2. Drive the autoregressive rollout
for i in range(cfg.total_blocks):
video_chunk = self.pipeline.generate(autoregressive_index=i, cache=cache)
self.pipeline.finalize(autoregressive_index=i, cache=cache)
if self.is_rank_zero:
# 3. Write outputs only on the main process
...
from flashdreams.infra.diffusion.model import DiffusionModelConfig
from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig
from flashdreams.infra.pipeline import StreamInferencePipelineConfig
from flashdreams.infra.runner import RunnerConfig
from customized_method.runner import CustomizedMethodRunnerConfig
from customized_method.transformer import MyTransformerConfig
# Define your pipeline config literal
CUSTOMIZED_PIPELINE_CONFIG = StreamInferencePipelineConfig(
name="customized-method",
diffusion_model=DiffusionModelConfig(
transformer=MyTransformerConfig(),
scheduler=FlowMatchSchedulerConfig(),
),
)
# Define your runner config
CUSTOMIZED_METHOD_RUNNER = CustomizedMethodRunnerConfig(
runner_name=CUSTOMIZED_PIPELINE_CONFIG.name,
description="Custom description for my method.",
pipeline=CUSTOMIZED_PIPELINE_CONFIG,
)
Existing integrations under the integrations/ can be used as a minimal guide. These folders are simple examples of what mini standalone repositories that depend on FlashDreams look like. Examples are often the best way to learn; take a look at the LingBot-World and Self-Forcing integrations for good references on how to extend and use FlashDreams in a custom project.
Registering methods¶
After registering a method one should be able to see it in the CLI helptext and run it, leveraging the existing CLI for complete command line control:
# List all available models, including new ones
flashdreams-run --help
# See configurable parameters for new models
flashdreams-run customized-method --help
# Run the new model
flashdreams-run customized-method --prompt "A beautiful custom generation."
In order to extend FlashDreams and register own models, one can package code as a Python package and register it with FlashDreams via an entrypoint in the pyproject.toml file. FlashDreams will automatically search for all registered runners and will register them to be used by the flashdreams-run CLI.
A pyproject.toml file is where the entrypoint to custom methods is specified and additional dependencies are specified. For example:
[project]
name = "customized_method"
version = "0.1.0"
dependencies = [
"flashdreams", # consider pinning the version, ie "flashdreams==0.1.0"
"mediapy>=1.1",
]
[tool.setuptools.packages.find]
include = ["customized_method*"]
[project.entry-points."flashdreams.runner_configs"]
customized-method = "customized_method.config:CUSTOMIZED_METHOD_RUNNER"
Finally, run the following to register the method:
pip install -e .
When developing a new methods they don’t always need to be installed as a package. Instead, the FLASHDREAMS_RUNNER_CONFIGS environment variable can be usedto temporarily register custom custom method.
export FLASHDREAMS_RUNNER_CONFIGS="customized-method=customized_method.config:RUNNER_CONFIGS"
The FLASHDREAMS_RUNNER_CONFIGS environment variable additionally accepts a function (a zero-arg callable) to temporarily register custom runners if construction has side effects to be deferred until CLI time.
Adding to the FlashDreams documentation¶
We invite researchers to contribute their own integrations to our official codebase and documentation. More information on how to do this can be found in the repository’s CONTRIBUTING.md. See the existing model pages under the docs/source/models/ (e.g., Self-Forcing) as templates for documenting new methods.