Duplicate from Pika-Labs/Pika-H3-Flex

0c019e1 3 days ago

7.94 kB

# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 Mellis, Inc.
"""Shared loading, memory management and output checks for the two examples."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import subprocess
import time

os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")

CODE_REVISION = "d0b9b1b5a540e5f9df2f8ec236347bf9c5aa7d1f"
BUCKETS = {"480p": (832, 480), "768p": (1344, 768)}


def arguments(task):
    parser = argparse.ArgumentParser(description=f"Pika H3 Flex {task} example (16:9)")
    parser.add_argument("--repo", default="Pika-Labs/Pika-H3-Flex")
    parser.add_argument("--revision", required=True, help="Immutable 40-character HF commit SHA")
    parser.add_argument("--snapshot", type=Path, help="Optional complete local snapshot of that revision")
    parser.add_argument("--assets-dir", type=Path, default=Path(__file__).parent / "assets",
                        help="Directory containing fl2va/ and ref2va/ assets; missing assets download from the same HF revision")
    parser.add_argument("--resolution", choices=BUCKETS, default="480p")
    parser.add_argument("--duration", type=int, choices=range(5, 16), default=15 if task == "fl2va" else 10)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output", type=Path, default=Path("outputs") / f"{task}.mp4")
    parser.add_argument("--resident-transformer", action="store_true",
                        help="Disable transformer block offload; uses more GPU memory")
    parser.add_argument("--text-encoder-offload", action="store_true",
                        help="Stream text/vision encoder layers from CPU for larger reference inputs")
    args = parser.parse_args()
    if not re.fullmatch(r"[a-f0-9]{40}", args.revision):
        parser.error("--revision must be an immutable 40-character HF commit SHA")
    return args


def asset(args, task, filename):
    path = args.assets_dir / task / filename
    if path.is_file():
        return path
    from huggingface_hub import hf_hub_download
    return Path(hf_hub_download(args.repo, f"examples/assets/{task}/{filename}", revision=args.revision))


def offload_other_components(hooks, model_id, model, execution_device):
    # A full text encoder and transformer cannot coexist on an H100 80GB.
    # Move all other resident components aside when the next component starts.
    # The manager itself excludes group-offloaded components from this list.
    return hooks


def build_pipeline(args, task):
    import torch
    from diffusers import ComponentsManager, ModularPipeline
    if not torch.cuda.is_available():
        raise RuntimeError("An NVIDIA CUDA GPU is required for these examples")
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = True
    manager = ComponentsManager()
    manager.enable_auto_cpu_offload(device="cuda:0", offload_strategy=offload_other_components)
    source = str(args.snapshot) if args.snapshot else args.repo
    pipe = ModularPipeline.from_pretrained(source, revision=args.revision,
                                          workflow=task, components_manager=manager)
    pipe.load_components(
        pretrained_model_name_or_path=source,
        revision=args.revision,
        dtype={"vae": torch.float32, "audio_vae": torch.float32, "default": torch.bfloat16},
    )
    pipe.load_lora_weights(source, revision=args.revision, weight_name=f"{task}/adapter.safetensors")
    if args.text_encoder_offload:
        from accelerate import cpu_offload
        from accelerate.hooks import remove_hook_from_module
        # The pipeline calls text_encoder.model directly and does not use the
        # language-model head. Let Accelerate place those actual computation
        # layers, rather than moving the entire encoder onto the GPU at once.
        remove_hook_from_module(pipe.text_encoder)
        cpu_offload(pipe.text_encoder.model, execution_device=torch.device("cuda:0"))
    transformer = pipe.transformer_ref if task == "ref2va" else pipe.transformer
    def release_other_components(module, inputs):
        # Run before group-offload's wrapped forward: once its first group is on
        # CUDA, the manager can otherwise mistake the whole model for resident.
        for hook in manager.model_hooks:
            if hook.model is not module and hook.model.device.type == "cuda":
                hook.offload()
    transformer.register_forward_pre_hook(release_other_components, prepend=True)
    if not args.resident_transformer:
        transformer.enable_group_offload(
            onload_device=torch.device("cuda:0"), offload_device=torch.device("cpu"),
            offload_type="block_level", num_blocks_per_group=1, use_stream=True, record_stream=False,
        )
    return pipe


def generate(args, task, media_kwargs, inputs):
    import torch
    import diffusers
    from diffusers.utils.export_utils import encode_video
    prompt_path = asset(args, task, "prompt.txt")
    prompt = prompt_path.read_text(encoding="utf-8").strip()
    width, height = BUCKETS[args.resolution]
    expected_frames = args.duration * 24
    while expected_frames % 17 != 5:
        expected_frames += 1
    start = time.monotonic()
    pipe = build_pipeline(args, task)
    loaded = time.monotonic()
    torch.cuda.reset_peak_memory_stats()
    result = pipe(prompt=prompt, height=height, width=width, num_frames=args.duration * 24,
                  generator=torch.Generator("cuda:0").manual_seed(args.seed),
                  output=["videos", "audio", "sampling_rate"], **media_kwargs)
    torch.cuda.synchronize()
    generated = time.monotonic()
    frames = result["videos"][0]
    audio = result["audio"][0].detach().float().cpu()
    assert len(frames) == expected_frames, (len(frames), expected_frames)
    assert frames[0].size == (width, height), frames[0].size
    assert torch.isfinite(audio).all(), "Output audio contains NaN or Inf"
    assert result["sampling_rate"] == 32000 and audio.shape[0] == 2
    args.output.parent.mkdir(parents=True, exist_ok=True)
    encode_video(frames, fps=24, output_path=str(args.output), audio=audio,
                 audio_sample_rate=result["sampling_rate"])
    probe = json.loads(subprocess.check_output(
        ["ffprobe", "-v", "error", "-show_streams", "-show_format", "-of", "json", str(args.output)], text=True))
    video = next(s for s in probe["streams"] if s["codec_type"] == "video")
    sound = next(s for s in probe["streams"] if s["codec_type"] == "audio")
    assert (video["width"], video["height"]) == (width, height)
    assert int(video["nb_frames"]) == expected_frames and video["r_frame_rate"] == "24/1"
    assert int(sound["sample_rate"]) == 32000 and sound["channels"] == 2
    record = {
        "repo": args.repo, "revision": args.revision, "diffusers_code_revision": CODE_REVISION,
        "diffusers_version": diffusers.__version__, "task": task, "seed": args.seed,
        "resolution": args.resolution, "width": width, "height": height,
        "requested_duration": args.duration, "frames": expected_frames, "fps": 24,
        "transformer_offload": not args.resident_transformer,
        "text_encoder_offload": args.text_encoder_offload,
        "load_seconds": loaded - start, "generation_seconds": generated - loaded,
        "peak_allocated_gib": torch.cuda.max_memory_allocated() / 2**30,
        "input_sha256": {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in [prompt_path, *inputs]},
        "output_sha256": hashlib.sha256(args.output.read_bytes()).hexdigest(),
        "media": probe,
    }
    args.output.with_suffix(".json").write_text(json.dumps(record, indent=2) + "\n")
    print(json.dumps({"output": str(args.output), "frames": expected_frames,
                      "generation_seconds": record["generation_seconds"],
                      "peak_allocated_gib": record["peak_allocated_gib"]}), flush=True)

	# SPDX-License-Identifier: Apache-2.0
	# Copyright 2026 Mellis, Inc.
	"""Shared loading, memory management and output checks for the two examples."""
	import argparse
	import hashlib
	import json
	import os
	from pathlib import Path
	import re
	import subprocess
	import time

	os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")

	CODE_REVISION = "d0b9b1b5a540e5f9df2f8ec236347bf9c5aa7d1f"
	BUCKETS = {"480p": (832, 480), "768p": (1344, 768)}


	def arguments(task):
	parser = argparse.ArgumentParser(description=f"Pika H3 Flex {task} example (16:9)")
	parser.add_argument("--repo", default="Pika-Labs/Pika-H3-Flex")
	parser.add_argument("--revision", required=True, help="Immutable 40-character HF commit SHA")
	parser.add_argument("--snapshot", type=Path, help="Optional complete local snapshot of that revision")
	parser.add_argument("--assets-dir", type=Path, default=Path(__file__).parent / "assets",
	help="Directory containing fl2va/ and ref2va/ assets; missing assets download from the same HF revision")
	parser.add_argument("--resolution", choices=BUCKETS, default="480p")
	parser.add_argument("--duration", type=int, choices=range(5, 16), default=15 if task == "fl2va" else 10)
	parser.add_argument("--seed", type=int, default=42)
	parser.add_argument("--output", type=Path, default=Path("outputs") / f"{task}.mp4")
	parser.add_argument("--resident-transformer", action="store_true",
	help="Disable transformer block offload; uses more GPU memory")
	parser.add_argument("--text-encoder-offload", action="store_true",
	help="Stream text/vision encoder layers from CPU for larger reference inputs")
	args = parser.parse_args()
	if not re.fullmatch(r"[a-f0-9]{40}", args.revision):
	parser.error("--revision must be an immutable 40-character HF commit SHA")
	return args


	def asset(args, task, filename):
	path = args.assets_dir / task / filename
	if path.is_file():
	return path
	from huggingface_hub import hf_hub_download
	return Path(hf_hub_download(args.repo, f"examples/assets/{task}/{filename}", revision=args.revision))


	def offload_other_components(hooks, model_id, model, execution_device):
	# A full text encoder and transformer cannot coexist on an H100 80GB.
	# Move all other resident components aside when the next component starts.
	# The manager itself excludes group-offloaded components from this list.
	return hooks


	def build_pipeline(args, task):
	import torch
	from diffusers import ComponentsManager, ModularPipeline
	if not torch.cuda.is_available():
	raise RuntimeError("An NVIDIA CUDA GPU is required for these examples")
	torch.backends.cuda.matmul.allow_tf32 = False
	torch.backends.cudnn.allow_tf32 = True
	manager = ComponentsManager()
	manager.enable_auto_cpu_offload(device="cuda:0", offload_strategy=offload_other_components)
	source = str(args.snapshot) if args.snapshot else args.repo
	pipe = ModularPipeline.from_pretrained(source, revision=args.revision,
	workflow=task, components_manager=manager)
	pipe.load_components(
	pretrained_model_name_or_path=source,
	revision=args.revision,
	dtype={"vae": torch.float32, "audio_vae": torch.float32, "default": torch.bfloat16},
	)
	pipe.load_lora_weights(source, revision=args.revision, weight_name=f"{task}/adapter.safetensors")
	if args.text_encoder_offload:
	from accelerate import cpu_offload
	from accelerate.hooks import remove_hook_from_module
	# The pipeline calls text_encoder.model directly and does not use the
	# language-model head. Let Accelerate place those actual computation
	# layers, rather than moving the entire encoder onto the GPU at once.
	remove_hook_from_module(pipe.text_encoder)
	cpu_offload(pipe.text_encoder.model, execution_device=torch.device("cuda:0"))
	transformer = pipe.transformer_ref if task == "ref2va" else pipe.transformer
	def release_other_components(module, inputs):
	# Run before group-offload's wrapped forward: once its first group is on
	# CUDA, the manager can otherwise mistake the whole model for resident.
	for hook in manager.model_hooks:
	if hook.model is not module and hook.model.device.type == "cuda":
	hook.offload()
	transformer.register_forward_pre_hook(release_other_components, prepend=True)
	if not args.resident_transformer:
	transformer.enable_group_offload(
	onload_device=torch.device("cuda:0"), offload_device=torch.device("cpu"),
	offload_type="block_level", num_blocks_per_group=1, use_stream=True, record_stream=False,
	)
	return pipe


	def generate(args, task, media_kwargs, inputs):
	import torch
	import diffusers
	from diffusers.utils.export_utils import encode_video
	prompt_path = asset(args, task, "prompt.txt")
	prompt = prompt_path.read_text(encoding="utf-8").strip()
	width, height = BUCKETS[args.resolution]
	expected_frames = args.duration * 24
	while expected_frames % 17 != 5:
	expected_frames += 1
	start = time.monotonic()
	pipe = build_pipeline(args, task)
	loaded = time.monotonic()
	torch.cuda.reset_peak_memory_stats()
	result = pipe(prompt=prompt, height=height, width=width, num_frames=args.duration * 24,
	generator=torch.Generator("cuda:0").manual_seed(args.seed),
	output=["videos", "audio", "sampling_rate"], **media_kwargs)
	torch.cuda.synchronize()
	generated = time.monotonic()
	frames = result["videos"][0]
	audio = result["audio"][0].detach().float().cpu()
	assert len(frames) == expected_frames, (len(frames), expected_frames)
	assert frames[0].size == (width, height), frames[0].size
	assert torch.isfinite(audio).all(), "Output audio contains NaN or Inf"
	assert result["sampling_rate"] == 32000 and audio.shape[0] == 2
	args.output.parent.mkdir(parents=True, exist_ok=True)
	encode_video(frames, fps=24, output_path=str(args.output), audio=audio,
	audio_sample_rate=result["sampling_rate"])
	probe = json.loads(subprocess.check_output(
	["ffprobe", "-v", "error", "-show_streams", "-show_format", "-of", "json", str(args.output)], text=True))
	video = next(s for s in probe["streams"] if s["codec_type"] == "video")
	sound = next(s for s in probe["streams"] if s["codec_type"] == "audio")
	assert (video["width"], video["height"]) == (width, height)
	assert int(video["nb_frames"]) == expected_frames and video["r_frame_rate"] == "24/1"
	assert int(sound["sample_rate"]) == 32000 and sound["channels"] == 2
	record = {
	"repo": args.repo, "revision": args.revision, "diffusers_code_revision": CODE_REVISION,
	"diffusers_version": diffusers.__version__, "task": task, "seed": args.seed,
	"resolution": args.resolution, "width": width, "height": height,
	"requested_duration": args.duration, "frames": expected_frames, "fps": 24,
	"transformer_offload": not args.resident_transformer,
	"text_encoder_offload": args.text_encoder_offload,
	"load_seconds": loaded - start, "generation_seconds": generated - loaded,
	"peak_allocated_gib": torch.cuda.max_memory_allocated() / 2**30,
	"input_sha256": {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in [prompt_path, *inputs]},
	"output_sha256": hashlib.sha256(args.output.read_bytes()).hexdigest(),
	"media": probe,
	}
	args.output.with_suffix(".json").write_text(json.dumps(record, indent=2) + "\n")
	print(json.dumps({"output": str(args.output), "frames": expected_frames,
	"generation_seconds": record["generation_seconds"],
	"peak_allocated_gib": record["peak_allocated_gib"]}), flush=True)