301 lines
16 KiB
Python
301 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""PAN2 reference implementation for local parity tests.
|
|
|
|
The diffusers ``PAN2Transformer3DModel`` and ``PAN2ModularPipeline`` are not public yet. The reference here is
|
|
assembled from diffusers' and transformers' own building blocks (``AdaLayerNormZero``, ``AdaLayerNormContinuous``,
|
|
``FeedForward`` with ``gelu-approximate``/``swiglu``, ``Timesteps``/``TimestepEmbedding``,
|
|
``get_1d_rotary_pos_embed``/``apply_rotary_emb``, ``Qwen3_5TextModel``, ``AutoencoderKLWan``,
|
|
``FlowMatchEulerDiscreteScheduler``) and follows the vLLM-Omni PAN2 port (vllm-project/vllm-omni#8602), which was
|
|
validated against the diffusers implementation (upstream revision ``e0c4c731``). Module names are the checkpoint's,
|
|
so checkpoints load strictly. Its unconditional SP padding mask is equivalent to unpadded reference attention.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import PIL.Image
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from diffusers.models.attention import FeedForward
|
|
from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
|
|
from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm
|
|
|
|
TINY_REPO_ID = "wuqing157/tiny-pan2-modular-pipe"
|
|
TINY_REVISION = "3ff3130b4527f61aaab7bde01465c9884dc84cdd"
|
|
SYSTEM_PROMPT = "You are a helpful assistant. Generate the video with the following instructions."
|
|
TEXT_ENCODER_LAYERS_TO_SKIP = 2
|
|
MAX_SEQUENCE_LENGTH = 1536
|
|
|
|
|
|
def tiny_model_dir() -> Path | None:
|
|
"""``PAN2_TINY_MODEL_DIR``, else the pinned tiny checkpoint in the local HF cache."""
|
|
local = os.getenv("PAN2_TINY_MODEL_DIR")
|
|
if local:
|
|
return Path(local)
|
|
try:
|
|
from huggingface_hub import snapshot_download
|
|
|
|
return Path(snapshot_download(TINY_REPO_ID, revision=TINY_REVISION, local_files_only=True))
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _sdpa(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
|
|
# [B, S, H, D] -> [B, S, H*D]
|
|
out = F.scaled_dot_product_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2))
|
|
return out.transpose(1, 2).flatten(2)
|
|
|
|
|
|
class _RefRefinerAttention(nn.Module):
|
|
|
|
def __init__(self, dim: int, heads: int, dim_head: int, eps: float):
|
|
super().__init__()
|
|
self.heads = heads
|
|
self.to_q = nn.Linear(dim, heads * dim_head, bias=False)
|
|
self.to_k = nn.Linear(dim, heads * dim_head, bias=False)
|
|
self.to_v = nn.Linear(dim, heads * dim_head, bias=False)
|
|
self.norm_q = RMSNorm(dim_head, eps=eps)
|
|
self.norm_k = RMSNorm(dim_head, eps=eps)
|
|
self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, dim, bias=False)])
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
query = self.norm_q(self.to_q(x).unflatten(-1, (self.heads, -1)))
|
|
key = self.norm_k(self.to_k(x).unflatten(-1, (self.heads, -1)))
|
|
value = self.to_v(x).unflatten(-1, (self.heads, -1))
|
|
return self.to_out[0](_sdpa(query, key, value))
|
|
|
|
|
|
class _RefRefinerBlock(nn.Module):
|
|
|
|
def __init__(self, dim: int, heads: int, dim_head: int, mlp_ratio: float, eps: float):
|
|
super().__init__()
|
|
self.norm1 = RMSNorm(dim, eps=eps)
|
|
self.attn = _RefRefinerAttention(dim, heads, dim_head, eps)
|
|
self.norm2 = RMSNorm(dim, eps=eps)
|
|
self.ff = FeedForward(dim, inner_dim=int(dim * mlp_ratio), activation_fn="swiglu", bias=False)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
x = x + self.attn(self.norm1(x))
|
|
return x + self.ff(self.norm2(x))
|
|
|
|
|
|
class _RefJointAttention(nn.Module):
|
|
|
|
def __init__(self, dim: int, heads: int, dim_head: int, eps: float):
|
|
super().__init__()
|
|
self.heads = heads
|
|
self.to_q = nn.Linear(dim, heads * dim_head)
|
|
self.to_k = nn.Linear(dim, heads * dim_head)
|
|
self.to_v = nn.Linear(dim, heads * dim_head)
|
|
self.norm_q = RMSNorm(dim_head, eps=eps)
|
|
self.norm_k = RMSNorm(dim_head, eps=eps)
|
|
self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, dim)])
|
|
self.add_q_proj = nn.Linear(dim, heads * dim_head)
|
|
self.add_k_proj = nn.Linear(dim, heads * dim_head)
|
|
self.add_v_proj = nn.Linear(dim, heads * dim_head)
|
|
self.norm_added_q = RMSNorm(dim_head, eps=eps)
|
|
self.norm_added_k = RMSNorm(dim_head, eps=eps)
|
|
self.to_add_out = nn.Linear(heads * dim_head, dim)
|
|
|
|
def forward(self, x: torch.Tensor, context: torch.Tensor,
|
|
rope: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
|
|
query = self.norm_q(self.to_q(x).unflatten(-1, (self.heads, -1)))
|
|
key = self.norm_k(self.to_k(x).unflatten(-1, (self.heads, -1)))
|
|
value = self.to_v(x).unflatten(-1, (self.heads, -1))
|
|
# Interleaved (GPT-J) rotary embedding on the video tokens only.
|
|
query = apply_rotary_emb(query, rope, use_real=True, use_real_unbind_dim=-1, sequence_dim=1)
|
|
key = apply_rotary_emb(key, rope, use_real=True, use_real_unbind_dim=-1, sequence_dim=1)
|
|
|
|
context_query = self.norm_added_q(self.add_q_proj(context).unflatten(-1, (self.heads, -1)))
|
|
context_key = self.norm_added_k(self.add_k_proj(context).unflatten(-1, (self.heads, -1)))
|
|
context_value = self.add_v_proj(context).unflatten(-1, (self.heads, -1))
|
|
|
|
out = _sdpa(torch.cat([query, context_query], dim=1), torch.cat([key, context_key], dim=1),
|
|
torch.cat([value, context_value], dim=1))
|
|
out, context_out = out.split([x.shape[1], context.shape[1]], dim=1)
|
|
return self.to_out[0](out), self.to_add_out(context_out)
|
|
|
|
|
|
class _RefBlock(nn.Module):
|
|
|
|
def __init__(self, dim: int, heads: int, dim_head: int, mlp_ratio: float, context_mlp_ratio: float):
|
|
super().__init__()
|
|
eps = 1e-6
|
|
self.norm1 = AdaLayerNormZero(dim)
|
|
self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
|
|
self.scale_shift_table_context = nn.Parameter(torch.zeros(6 * dim))
|
|
self.attn = _RefJointAttention(dim, heads, dim_head, eps)
|
|
self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
|
|
self.ff = FeedForward(dim, inner_dim=int(dim * mlp_ratio), activation_fn="gelu-approximate")
|
|
self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
|
|
self.ff_context = FeedForward(dim, inner_dim=int(dim * context_mlp_ratio), activation_fn="gelu-approximate")
|
|
|
|
def forward(self, x, context, temb, rope):
|
|
norm_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(x, emb=temb)
|
|
c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = (
|
|
self.scale_shift_table_context.chunk(6, dim=-1))
|
|
norm_context = self.norm1_context(context) * (1 + c_scale_msa) + c_shift_msa
|
|
attn_out, context_attn_out = self.attn(norm_x, norm_context, rope)
|
|
x = x + attn_out * gate_msa.unsqueeze(1)
|
|
context = context + context_attn_out * c_gate_msa
|
|
norm_x = self.norm2(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
|
|
norm_context = self.norm2_context(context) * (1 + c_scale_mlp) + c_shift_mlp
|
|
x = x + self.ff(norm_x) * gate_mlp.unsqueeze(1)
|
|
context = context + self.ff_context(norm_context) * c_gate_mlp
|
|
return x, context
|
|
|
|
|
|
class ReferencePAN2Transformer(nn.Module):
|
|
|
|
def __init__(self, patch_size=(1, 2, 2), in_channels=97, out_channels=48, num_attention_heads=32,
|
|
attention_head_dim=128, num_layers=44, num_refiner_layers=4, mlp_ratio=8.0, context_mlp_ratio=4.0,
|
|
refiner_mlp_ratio=4.0, text_embed_dim=4096, rope_axes_dim=(16, 56, 56), rope_theta=256.0,
|
|
rope_temporal_scale=16 / 24, **_):
|
|
super().__init__()
|
|
self.patch_size = tuple(patch_size)
|
|
self.rope_axes_dim = tuple(rope_axes_dim)
|
|
self.rope_theta = rope_theta
|
|
self.rope_temporal_scale = rope_temporal_scale
|
|
inner_dim = num_attention_heads * attention_head_dim
|
|
self.x_embedder = nn.Conv3d(in_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size)
|
|
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
|
|
self.time_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=inner_dim)
|
|
self.context_embedder = nn.Sequential(RMSNorm(text_embed_dim, eps=1e-5), nn.Linear(text_embed_dim, inner_dim))
|
|
self.context_refiner = nn.ModuleList([
|
|
_RefRefinerBlock(inner_dim, num_attention_heads, attention_head_dim, refiner_mlp_ratio, eps=1e-5)
|
|
for _ in range(num_refiner_layers)
|
|
])
|
|
self.context_type_embedding = nn.Parameter(torch.zeros(inner_dim))
|
|
self.transformer_blocks = nn.ModuleList([
|
|
_RefBlock(inner_dim, num_attention_heads, attention_head_dim, mlp_ratio, context_mlp_ratio)
|
|
for _ in range(num_layers)
|
|
])
|
|
self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
|
|
self.proj_out = nn.Linear(inner_dim, self.patch_size[0] * self.patch_size[1] * self.patch_size[2] * out_channels)
|
|
|
|
def _rope(self, num_frames: int, height: int, width: int, device) -> tuple[torch.Tensor, torch.Tensor]:
|
|
axes = [torch.arange(size, device=device, dtype=torch.float32) for size in (num_frames, height, width)]
|
|
grid = torch.stack(torch.meshgrid(*axes, indexing="ij"), dim=0)
|
|
grid[0] = grid[0] * self.rope_temporal_scale
|
|
cos, sin = zip(*[
|
|
get_1d_rotary_pos_embed(dim, grid[i].reshape(-1), self.rope_theta, use_real=True,
|
|
repeat_interleave_real=True) for i, dim in enumerate(self.rope_axes_dim)
|
|
])
|
|
return torch.cat(cos, dim=1).float(), torch.cat(sin, dim=1).float()
|
|
|
|
def forward(self, hidden_states, timestep, encoder_hidden_states):
|
|
batch_size, _, num_frames, height, width = hidden_states.shape
|
|
p_t, p_h, p_w = self.patch_size
|
|
t, h, w = num_frames // p_t, height // p_h, width // p_w
|
|
rope = self._rope(t, h, w, hidden_states.device)
|
|
x = self.x_embedder(hidden_states).flatten(2).transpose(1, 2)
|
|
temb = self.time_embedder(self.time_proj(timestep).to(x.dtype))
|
|
context = self.context_embedder(encoder_hidden_states)
|
|
for block in self.context_refiner:
|
|
context = block(context)
|
|
context = context + self.context_type_embedding
|
|
for block in self.transformer_blocks:
|
|
x, context = block(x, context, temb, rope)
|
|
x = self.proj_out(self.norm_out(x, temb))
|
|
x = x.reshape(batch_size, t, h, w, -1, p_t, p_h, p_w).permute(0, 4, 1, 5, 2, 6, 3, 7)
|
|
return x.flatten(6, 7).flatten(4, 5).flatten(2, 3)
|
|
|
|
|
|
# --------------------------------------------------------------------------------------------------------------------
|
|
# Reference pipeline (vLLM-Omni ``PAN2Pipeline.forward``)
|
|
# --------------------------------------------------------------------------------------------------------------------
|
|
|
|
|
|
def reference_encode_prompt(tokenizer, text_encoder, prompt: str, device) -> torch.Tensor:
|
|
messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt or " "}]
|
|
input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True,
|
|
return_dict=True).input_ids
|
|
header_ids = tokenizer("<|im_start|>user\n", add_special_tokens=False).input_ids
|
|
crop_start = next(i + len(header_ids) for i in range(len(input_ids))
|
|
if input_ids[i:i + len(header_ids)] == header_ids)
|
|
input_ids = torch.tensor([input_ids[:crop_start + MAX_SEQUENCE_LENGTH]], device=device)
|
|
hidden_states = text_encoder(input_ids=input_ids, output_hidden_states=True).hidden_states
|
|
return hidden_states[-(TEXT_ENCODER_LAYERS_TO_SKIP + 1)][:, crop_start:]
|
|
|
|
|
|
def reference_condition_latents(vae, latents: torch.Tensor, image: PIL.Image.Image | None, height: int,
|
|
width: int) -> torch.Tensor:
|
|
batch_size, num_channels, num_frames, latent_height, latent_width = latents.shape
|
|
condition = latents.new_zeros(batch_size, num_channels + 1, num_frames, latent_height, latent_width)
|
|
if image is None:
|
|
return condition
|
|
from diffusers.video_processor import VideoProcessor
|
|
|
|
image = image.convert("RGB")
|
|
scale = max(width / image.width, height / image.height)
|
|
image = image.resize((round(image.width * scale), round(image.height * scale)), resample=PIL.Image.LANCZOS)
|
|
left, top = round((image.width - width) / 2), round((image.height - height) / 2)
|
|
image = image.crop((left, top, left + width, top + height))
|
|
pixels = VideoProcessor(vae_scale_factor=16).preprocess(image, height=height, width=width)
|
|
pixels = pixels.unsqueeze(2).to(device=latents.device, dtype=vae.dtype)
|
|
image_latents = vae.encode(pixels).latent_dist.mode()
|
|
mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(image_latents)
|
|
std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(image_latents)
|
|
condition[:, :num_channels, :1] = (image_latents - mean) / std
|
|
condition[:, num_channels:, :1] = 1.0
|
|
return condition
|
|
|
|
|
|
@torch.no_grad()
|
|
def reference_generate(components: dict, *, prompt: str, negative_prompt: str, image: PIL.Image.Image | None,
|
|
height: int, width: int, num_frames: int, num_inference_steps: int, guidance_scale: float,
|
|
seed: int, dit_dtype: torch.dtype, device="cpu") -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Final latents and the decoded video in ``[0, 1]``, following the vLLM-Omni pipeline step by step."""
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
|
|
tokenizer, text_encoder = components["tokenizer"], components["text_encoder"]
|
|
transformer, vae, scheduler = components["transformer"], components["vae"], components["scheduler"]
|
|
do_cfg = guidance_scale > 1.0
|
|
prompt_embeds = reference_encode_prompt(tokenizer, text_encoder, prompt, device).to(dit_dtype)
|
|
negative_embeds = (reference_encode_prompt(tokenizer, text_encoder, negative_prompt, device).to(dit_dtype)
|
|
if do_cfg else None)
|
|
|
|
shape = (1, vae.config.z_dim, (num_frames - 1) // vae.config.scale_factor_temporal + 1,
|
|
height // vae.config.scale_factor_spatial, width // vae.config.scale_factor_spatial)
|
|
generator = [torch.Generator("cpu").manual_seed(seed)]
|
|
latents = randn_tensor(shape, generator=generator, device=torch.device(device), dtype=torch.float32)
|
|
condition = reference_condition_latents(vae, latents, image, height, width).to(dit_dtype)
|
|
|
|
scheduler.set_timesteps(sigmas=np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1], device=device)
|
|
for t in scheduler.timesteps:
|
|
model_input = torch.cat([latents.to(dit_dtype), condition], dim=1)
|
|
timestep = t.expand(model_input.shape[0])
|
|
noise_pred = transformer(model_input, timestep, prompt_embeds)
|
|
if do_cfg:
|
|
noise_uncond = transformer(model_input, timestep, negative_embeds)
|
|
noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond)
|
|
latents = scheduler.step(noise_pred.float(), t, latents, return_dict=False)[0]
|
|
|
|
mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latents)
|
|
std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(latents)
|
|
video = vae.decode((latents * std + mean).to(vae.dtype), return_dict=False)[0]
|
|
return latents, (video.float() / 2 + 0.5).clamp(0, 1)
|
|
|
|
|
|
def load_reference_components(model_dir: Path, dit_dtype: torch.dtype) -> dict:
|
|
import json
|
|
|
|
from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler
|
|
from safetensors.torch import load_file
|
|
from transformers import AutoTokenizer, Qwen3_5TextModel
|
|
|
|
transformer_dir = model_dir / "transformer"
|
|
transformer = ReferencePAN2Transformer(**json.loads((transformer_dir / "config.json").read_text()))
|
|
transformer.load_state_dict(load_file(str(transformer_dir / "diffusion_pytorch_model.safetensors")), strict=True)
|
|
return {
|
|
"tokenizer": AutoTokenizer.from_pretrained(model_dir / "tokenizer"),
|
|
"text_encoder": Qwen3_5TextModel.from_pretrained(model_dir / "text_encoder", torch_dtype=torch.float32).eval(),
|
|
"transformer": transformer.to(dit_dtype).eval(),
|
|
"vae": AutoencoderKLWan.from_pretrained(model_dir / "vae", torch_dtype=torch.float32).eval(),
|
|
"scheduler": FlowMatchEulerDiscreteScheduler.from_pretrained(model_dir / "scheduler"),
|
|
}
|