"""Chunked denoising loop with cache hooks and denoise-only DiT timing. Mirrors ``CausalInferencePipeline.inference`` from both base repos (they are identical apart from comments), with three additions: * the controller is told which (chunk, step) each forward belongs to; * only the four denoising forwards are timed -- the KV-cache refresh pass (``context_noise``) and any initial-latent context passes are excluded, which is what "只考虑 denoise DiT 耗时" asks for; * VAE decode and text encoding are outside the measured window. """ import time import torch def _reset_caches(pipeline, noise, batch_size): if pipeline.kv_cache1 is None: pipeline._initialize_kv_cache(batch_size=batch_size, dtype=noise.dtype, device=noise.device) pipeline._initialize_crossattn_cache(batch_size=batch_size, dtype=noise.dtype, device=noise.device) else: for i in range(pipeline.num_transformer_blocks): pipeline.crossattn_cache[i]["is_init"] = False for i in range(len(pipeline.kv_cache1)): pipeline.kv_cache1[i]["global_end_index"] = torch.tensor( [0], dtype=torch.long, device=noise.device) pipeline.kv_cache1[i]["local_end_index"] = torch.tensor( [0], dtype=torch.long, device=noise.device) @torch.no_grad() def _sigma_for_timestep(scheduler, timestep, device): """sigma looked up exactly the way the wrapper's x0 conversion does (nearest entry of the scheduler's timestep table).""" timesteps = scheduler.timesteps.to(device=device, dtype=torch.float64) sigmas = scheduler.sigmas.to(device=device, dtype=torch.float64) t = torch.as_tensor(timestep, device=device, dtype=torch.float64).reshape(-1) idx = torch.argmin((timesteps.unsqueeze(0) - t.unsqueeze(1)).abs(), dim=1) return sigmas[idx][0] def _flow_euler_update(scheduler, x_cur, flow_pred, current_timestep, next_timestep): sigma_cur = _sigma_for_timestep(scheduler, current_timestep, x_cur.device) sigma_next = _sigma_for_timestep(scheduler, next_timestep, x_cur.device) x_next = x_cur.double() + (sigma_next - sigma_cur) * flow_pred.double() return x_next.to(x_cur.dtype) SAMPLERS = ("renoise", "flow_euler") def cached_inference(pipeline, controller, noise, text_prompts, initial_latent=None, decode=True, first_chunk_steps=None, sampler="renoise"): if sampler not in SAMPLERS: raise ValueError(f"unknown sampler {sampler!r}; have {SAMPLERS}") batch_size, num_frames, num_channels, height, width = noise.shape nfpb = pipeline.num_frame_per_block if not pipeline.independent_first_frame or initial_latent is not None: assert num_frames % nfpb == 0 num_blocks = num_frames // nfpb else: assert (num_frames - 1) % nfpb == 0 num_blocks = (num_frames - 1) // nfpb num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0 num_output_frames = num_frames + num_input_frames conditional_dict = pipeline.text_encoder(text_prompts=text_prompts) output = torch.zeros([batch_size, num_output_frames, num_channels, height, width], device=noise.device, dtype=noise.dtype) _reset_caches(pipeline, noise, batch_size) controller.reset_video() current_start_frame = 0 if initial_latent is not None: timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0 if pipeline.independent_first_frame: assert (num_input_frames - 1) % nfpb == 0 num_input_blocks = (num_input_frames - 1) // nfpb output[:, :1] = initial_latent[:, :1] pipeline.generator( noisy_image_or_video=initial_latent[:, :1], conditional_dict=conditional_dict, timestep=timestep * 0, kv_cache=pipeline.kv_cache1, crossattn_cache=pipeline.crossattn_cache, current_start=current_start_frame * pipeline.frame_seq_length) current_start_frame += 1 else: assert num_input_frames % nfpb == 0 num_input_blocks = num_input_frames // nfpb for _ in range(num_input_blocks): refs = initial_latent[:, current_start_frame:current_start_frame + nfpb] output[:, current_start_frame:current_start_frame + nfpb] = refs pipeline.generator( noisy_image_or_video=refs, conditional_dict=conditional_dict, timestep=timestep * 0, kv_cache=pipeline.kv_cache1, crossattn_cache=pipeline.crossattn_cache, current_start=current_start_frame * pipeline.frame_seq_length) current_start_frame += nfpb all_num_frames = [nfpb] * num_blocks if pipeline.independent_first_frame and initial_latent is None: all_num_frames = [1] + all_num_frames denoise_events = [] context_events = [] per_chunk_ms = [] for block_index, current_num_frames in enumerate(all_num_frames): controller.begin_chunk(block_index) noisy_input = noise[:, current_start_frame - num_input_frames: current_start_frame + current_num_frames - num_input_frames] # Chunk 0 may run its own (longer) schedule, e.g. the full four steps while # later chunks run a reduced naive-N-step list. steps = (first_chunk_steps if first_chunk_steps is not None and block_index == 0 else pipeline.denoising_step_list) chunk_events = [] for index, current_timestep in enumerate(steps): timestep = torch.ones([batch_size, current_num_frames], device=noise.device, dtype=torch.int64) * current_timestep start_ev = torch.cuda.Event(enable_timing=True) end_ev = torch.cuda.Event(enable_timing=True) controller.denoise_step(index) start_ev.record() flow_pred, denoised_pred = pipeline.generator( noisy_image_or_video=noisy_input, conditional_dict=conditional_dict, timestep=timestep, kv_cache=pipeline.kv_cache1, crossattn_cache=pipeline.crossattn_cache, current_start=current_start_frame * pipeline.frame_seq_length) end_ev.record() controller.end_step() denoise_events.append((start_ev, end_ev)) chunk_events.append((start_ev, end_ev)) if index < len(steps) - 1: next_timestep = steps[index + 1] if sampler == "flow_euler": # Deterministic flow-matching Euler step on the flow field # the DiT just predicted: x_next = x_cur + (s_next - s_cur) * v. # Replaces the stochastic "x0 + fresh noise" re-noising; the # DiT call, timesteps, KV cache and final x0 output are the same. noisy_input = _flow_euler_update( pipeline.scheduler, noisy_input, flow_pred, current_timestep, next_timestep) else: noisy_input = pipeline.scheduler.add_noise( denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), next_timestep * torch.ones([batch_size * current_num_frames], device=noise.device, dtype=torch.long) ).unflatten(0, denoised_pred.shape[:2]) output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred # KV-cache refresh with the clean context: always the full DiT. Timed # separately and never added to the denoise total -- the evaluation # protocol excludes the context DiT from policy latency but requires it to # be reported, as evidence it was excluded. context_timestep = torch.ones_like(timestep) * pipeline.args.context_noise ctx_start = torch.cuda.Event(enable_timing=True) ctx_end = torch.cuda.Event(enable_timing=True) ctx_start.record() pipeline.generator( noisy_image_or_video=denoised_pred, conditional_dict=conditional_dict, timestep=context_timestep, kv_cache=pipeline.kv_cache1, crossattn_cache=pipeline.crossattn_cache, current_start=current_start_frame * pipeline.frame_seq_length) ctx_end.record() context_events.append((ctx_start, ctx_end)) per_chunk_ms.append(chunk_events) current_start_frame += current_num_frames torch.cuda.synchronize() per_step_ms = [s.elapsed_time(e) for s, e in denoise_events] denoise_ms = sum(per_step_ms) chunk_ms = [sum(s.elapsed_time(e) for s, e in evs) for evs in per_chunk_ms] context_ms = sum(s.elapsed_time(e) for s, e in context_events) timing = { "denoise_dit_ms": denoise_ms, "first_chunk_denoise_ms": chunk_ms[0] if chunk_ms else 0.0, "per_chunk_denoise_ms": chunk_ms, # One entry per denoising forward, in order, aligned with # ``controller.records``. Per-step times give hundreds of samples per run, # so their minima survive contention that ruins whole-video timings. "per_step_denoise_ms": per_step_ms, "num_denoise_forwards": len(denoise_events), # Excluded from the denoise total on purpose; reported so the exclusion # is auditable (evaluation protocol section 9.2). "context_kv_dit_ms": context_ms, "num_context_forwards": len(context_events), } video = None if decode: video = pipeline.vae.decode_to_pixel(output, use_cache=False) video = (video * 0.5 + 0.5).clamp(0, 1) return video, output, timing