135 lines
5.4 KiB
Python
135 lines
5.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
import os
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from transformers import AutoConfig, AutoTokenizer, CLIPTextModel
|
|
import gc
|
|
from fastvideo.configs.pipelines import HunyuanConfig, PipelineConfig
|
|
from fastvideo.forward_context import set_forward_context
|
|
from fastvideo.fastvideo_args import FastVideoArgs
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.utils import maybe_download_model
|
|
from fastvideo.configs.models.encoders import CLIPTextConfig
|
|
from torch.distributed.tensor import DTensor
|
|
from torch.testing import assert_close
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
os.environ["MASTER_ADDR"] = "localhost"
|
|
os.environ["MASTER_PORT"] = "29503"
|
|
|
|
BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
|
|
MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
|
|
TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder_2")
|
|
TOKENIZER_PATH = os.path.join(MODEL_PATH, "tokenizer_2")
|
|
|
|
|
|
@pytest.mark.usefixtures("distributed_setup")
|
|
def test_clip_encoder():
|
|
"""
|
|
Tests compatibility between two different implementations for loading text encoders:
|
|
1. load_text_encoder from fastvideo.models.hunyuan.text_encoder
|
|
2. TextEncoderLoader from fastvideo.models.loader
|
|
|
|
The test verifies that both implementations:
|
|
- Load models with the same weights and parameters
|
|
- Produce nearly identical outputs for the same input prompts
|
|
"""
|
|
args = FastVideoArgs(model_path="openai/clip-vit-large-patch14", pipeline_config=HunyuanConfig())
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
|
|
logger.info("Loading models from %s", args.model_path)
|
|
|
|
# config = json.load(open(os.path.join(model_path, "config.json")))
|
|
|
|
hf_config = AutoConfig.from_pretrained(TEXT_ENCODER_PATH)
|
|
print(hf_config)
|
|
print(hf_config.use_return_dict)
|
|
|
|
# Load HuggingFace implementation
|
|
model1 = CLIPTextModel.from_pretrained(TEXT_ENCODER_PATH).to(torch.float16).to(device).eval()
|
|
|
|
from fastvideo.models.loader.component_loader import TextEncoderLoader
|
|
|
|
loader = TextEncoderLoader()
|
|
model2 = loader.load(TEXT_ENCODER_PATH, args)
|
|
|
|
# Load the HuggingFace implementation directly
|
|
# model2 = CLIPTextModel(hf_config)
|
|
# model2 = model2.to(torch.float16)
|
|
model2.eval()
|
|
|
|
# Sanity check weights between the two models
|
|
logger.info("Comparing model weights for sanity check...")
|
|
params1 = dict(model1.named_parameters())
|
|
params2 = dict(model2.named_parameters())
|
|
|
|
# Check number of parameters
|
|
logger.info("Model1 has %d parameters", len(params1))
|
|
logger.info("Model2 has %d parameters", len(params2))
|
|
|
|
for name1, param1 in sorted(params1.items()):
|
|
name2 = name1
|
|
skip = False
|
|
for param_name, weight_name, shard_id in model2.config.arch_config.stacked_params_mapping:
|
|
if weight_name not in name1:
|
|
skip = True
|
|
# stacked params are more troublesome
|
|
if skip:
|
|
continue
|
|
param2 = params2[name2]
|
|
param2 = param2.to_local().to(device) if isinstance(param2, DTensor) else param2.to(device)
|
|
assert_close(param1, param2, atol=1e-4, rtol=1e-4)
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
# Load tokenizer
|
|
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
|
|
|
|
# Test with some sample prompts
|
|
prompts = ["a photo of a cat", "a beautiful landscape with mountains", "an astronaut riding a horse on the moon"]
|
|
|
|
logger.info("Testing CLIP text encoder with sample prompts")
|
|
|
|
with torch.no_grad():
|
|
for prompt in prompts:
|
|
logger.info("Testing prompt: '%s'", prompt)
|
|
|
|
# Tokenize the prompt
|
|
tokens = tokenizer(prompt, padding="max_length", max_length=77, truncation=True, return_tensors="pt").to(
|
|
device
|
|
)
|
|
# Get embeddings from our implementation
|
|
outputs1 = model1(input_ids=tokens.input_ids, output_hidden_states=True)
|
|
|
|
logger.info("Testing model2")
|
|
print("--------------------------------")
|
|
# Get embeddings from HuggingFace implementation
|
|
with set_forward_context(current_timestep=0, attn_metadata=None):
|
|
outputs2 = model2(
|
|
input_ids=tokens.input_ids,
|
|
# attention_mask=tokens.attention_mask,
|
|
output_hidden_states=True,
|
|
)
|
|
|
|
# Compare last hidden states
|
|
last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask == 1]
|
|
last_hidden_state2 = outputs2.last_hidden_state[tokens.attention_mask == 1]
|
|
# print("last_hidden_state1", last_hidden_state1)
|
|
# print("last_hidden_state2", last_hidden_state2)
|
|
|
|
assert last_hidden_state1.shape == last_hidden_state2.shape, (
|
|
f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}"
|
|
)
|
|
# Compare pooler outputs
|
|
pooler_output1 = outputs1.pooler_output
|
|
pooler_output2 = outputs2.pooler_output
|
|
|
|
assert pooler_output1.shape == pooler_output2.shape, (
|
|
f"Pooler output shapes don't match: {pooler_output1.shape} vs {pooler_output2.shape}"
|
|
)
|
|
|
|
assert_close(pooler_output1, pooler_output2, atol=1e-2, rtol=1e-3)
|
|
assert_close(last_hidden_state1, last_hidden_state2, atol=1e-2, rtol=1e-3)
|