DreamX-World autoregressive causal denoising stage.
Classes
DreamXWorldARCausalDenoisingStage(transformer, scheduler, pipeline=None, vae=None)
Bases: DenoisingStage
Official DreamX AR-forcing denoising loop with KV cache.
Source code in fastvideo/pipelines/basic/dreamx_world/ar_denoising.py
| def __init__(self, transformer, scheduler, pipeline=None, vae=None) -> None:
super().__init__(transformer=transformer, scheduler=scheduler, pipeline=pipeline, vae=vae)
self.num_transformer_blocks = len(self.transformer.blocks)
self.num_frame_per_block = int(getattr(self.transformer, "num_frame_per_block", 3))
self.local_attn_size = int(getattr(self.transformer, "local_attn_size", 12))
|