stages ¶

def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None) -> None:
    super().__init__()
    self.transformer = transformer
    self.transformer_2 = transformer_2
    self.scheduler = scheduler
    self.vae = vae
    self.pipeline = weakref.ref(pipeline) if pipeline else None
    attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
    self.attn_backend = get_attn_backend(
        head_size=attn_head_size,
        dtype=torch.float16,  # TODO(will): hack
        supported_attention_backends=(AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.BSA_ATTN,
                                      AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.FLASH_ATTN,
                                      AttentionBackendEnum.TORCH_SDPA,
                                      AttentionBackendEnum.SAGE_ATTN_THREE),  # hack
        # Build metadata for the backend this transformer actually resolved
        # instead of re-deriving it from the environment. The two agreed
        # only when the request arrived via the env var: a request passed as
        # `attention_backend` reached the layers but never this stage.
        requested=component_attention_backend(self.transformer),
    )