Add files using upload-large-folder tool

f87692b verified about 1 month ago

9.96 kB

"""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

	"""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