causal_fast_pipeline ¶

@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    """Execute LingBot World 2 causal-fast generation and store final latents on the batch."""
    device = get_local_torch_device()
    cfg = fastvideo_args.pipeline_config.dit_config.arch_config
    chunk_size = int(cfg.chunk_size)
    max_sequence_length = int(batch.max_sequence_length or cfg.text_len)
    action_path = batch.action_path
    if action_path is None:
        raise ValueError("LingBot World 2 causal-fast requires `action_path`.")

    c2ws = np.load(os.path.join(action_path, "poses.npy"))
    len_c2ws = ((len(c2ws) - 1) // 4) * 4 + 1
    frame_num = ((int(batch.num_frames) - 1) // 4) * 4 + 1
    frame_num = min(frame_num, len_c2ws)
    c2ws = c2ws[:frame_num]

    img = self._prepare_image_tensor(batch, device)
    h0, w0 = img.shape[1:]
    aspect_ratio = h0 / w0
    lat_h = round(np.sqrt(cfg.max_area * aspect_ratio) // 8 // cfg.patch_size[1] * cfg.patch_size[1])
    lat_w = round(np.sqrt(cfg.max_area / aspect_ratio) // 8 // cfg.patch_size[2] * cfg.patch_size[2])
    h = lat_h * 8
    w = lat_w * 8
    lat_f = (frame_num - 1) // 4 + 1
    lat_f = int(lat_f - (lat_f % chunk_size))
    frames = (lat_f - 1) * 4 + 1
    batch.height = h
    batch.width = w
    batch.num_frames = frames

    seed = int(batch.seed if batch.seed is not None else 42)
    seed_g = torch.Generator(device=device)
    seed_g.manual_seed(seed)
    noise = torch.randn(16, lat_f, lat_h, lat_w, dtype=torch.float32, generator=seed_g, device=device)

    mask = torch.ones(1, frames, lat_h, lat_w, device=device)
    mask[:, 1:] = 0
    mask = torch.concat([torch.repeat_interleave(mask[:, 0:1], repeats=4, dim=1), mask[:, 1:]], dim=1)
    mask = mask.view(1, mask.shape[1] // 4, 4, lat_h, lat_w).transpose(1, 2)[0]

    self.scheduler.set_timesteps(cfg.num_train_timesteps, shift=cfg.sample_shift)
    timesteps = self.scheduler.timesteps[list(cfg.timesteps_index)].to(device)
    context = self._prompt_context(batch, device)
    c2ws_plucker_emb = self._prepare_camera(
        action_path,
        c2ws,
        h,
        w,
        lat_f,
        lat_h,
        lat_w,
        chunk_size,
        torch.bfloat16,
        device,
    )
    y = self._encode_condition_video(img, h, w, frames, mask, fastvideo_args).to(device=device,
                                                                                 dtype=torch.bfloat16)

    transformer_dtype = torch.bfloat16
    frame_seqlen = int(noise.shape[-2] * noise.shape[-1] // 4)
    kv_size = frame_seqlen * cfg.local_attn_size if cfg.local_attn_size > -1 else frame_seqlen * lat_f
    self_kv_cache = self._initialize_self_kv_cache(1, kv_size, transformer_dtype, device)
    cross_kv_cache = self._initialize_crossattn_cache(1, max_sequence_length, transformer_dtype, device)

    self.transformer = self.transformer.to(device)
    self._cross_attn_initialized = False
    pred_latent_chunks = []
    latents_chunk = noise.split(chunk_size, dim=1)
    condition_chunk = y.split(chunk_size, dim=1)
    c2ws_plucker_emb_chunk = c2ws_plucker_emb.split(chunk_size, dim=2)
    max_seq_len = int(math.ceil(chunk_size * lat_h * lat_w // 4))

    with torch.amp.autocast("cuda", dtype=transformer_dtype):
        for chunk_id, current_latent in enumerate(latents_chunk):
            current_condition = condition_chunk[chunk_id]
            current_c2ws_plucker_emb = c2ws_plucker_emb_chunk[chunk_id]
            dit_cond_dict = {"c2ws_plucker_emb": current_c2ws_plucker_emb.chunk(1, dim=0)}
            kwargs = {
                "context": [context[0]],
                "seq_len": max_seq_len,
                "y": [current_condition],
                "dit_cond_dict": dit_cond_dict,
                "kv_cache": self_kv_cache,
                "crossattn_cache": cross_kv_cache,
                "current_start": chunk_id * chunk_size * frame_seqlen,
                "max_attention_size": kv_size,
                "frame_seqlen": frame_seqlen,
            }
            x0 = current_latent
            for timestep_idx, timestep_value in enumerate(timesteps):
                timestep = torch.stack([timestep_value]).to(device)
                noise_pred = self.transformer(
                    x=[current_latent.to(device)],
                    t=timestep,
                    cross_attn_first_call=not self._cross_attn_initialized,
                    **kwargs,
                )[0]
                self._cross_attn_initialized = True
                x0 = self._convert_flow_pred_to_x0(noise_pred, current_latent, timestep_value)
                if timestep_idx < len(timesteps) - 1:
                    next_timestep = timesteps[timestep_idx + 1].reshape(1)
                    current_latent = self.scheduler.add_noise(
                        x0,
                        torch.randn(x0.shape, generator=seed_g, device=x0.device, dtype=x0.dtype),
                        next_timestep,
                    )
            pred_latent_chunks.append(x0)
            context_timestep = torch.stack([timesteps[-1] * 0.0]).to(device)
            self.transformer(x=[x0], t=context_timestep, cross_attn_first_call=False, **kwargs)

    batch.latents = torch.cat(pred_latent_chunks, dim=1).unsqueeze(0)
    return batch

def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
    """Set up the LingBot World 2 causal-fast pipeline stages."""
    self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
    self.add_stage(
        stage_name="prompt_encoding_stage",
        stage=LingBotWorld2TextEncodingStage(
            text_encoders=[self.get_module("text_encoder")],
            tokenizers=[self.get_module("tokenizer")],
        ),
    )
    self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
    self.add_stage(
        stage_name="lingbotworld2_causal_fast_generation_stage",
        stage=LingBotWorld2CausalFastGenerationStage(
            transformer=self.get_module("transformer"),
            scheduler=self.get_module("scheduler"),
            vae=self.get_module("vae"),
        ),
    )
    self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))

def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    """Initialize mask storage before running the shared text-encoding stage."""
    if batch.prompt_attention_mask is None:
        batch.prompt_attention_mask = []
    return super().forward(batch, fastvideo_args)