mirror of
https://github.com/finegrain-ai/refiners.git
synced 2024-11-09 23:12:02 +00:00
54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
from transformers import CLIPVisionModelWithProjection # type: ignore
|
|
|
|
from refiners.fluxion.utils import load_from_safetensors, no_grad
|
|
from refiners.foundationals.clip.image_encoder import CLIPImageEncoderH
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def our_encoder(
|
|
clip_image_encoder_huge_weights_path: Path,
|
|
test_device: torch.device,
|
|
test_dtype_fp32_bf16_fp16: torch.dtype,
|
|
) -> CLIPImageEncoderH:
|
|
encoder = CLIPImageEncoderH(device=test_device, dtype=test_dtype_fp32_bf16_fp16)
|
|
tensors = load_from_safetensors(clip_image_encoder_huge_weights_path)
|
|
encoder.load_state_dict(tensors)
|
|
return encoder
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def ref_encoder(
|
|
unclip21_transformers_stabilityai_path: str,
|
|
test_device: torch.device,
|
|
test_dtype_fp32_bf16_fp16: torch.dtype,
|
|
use_local_weights: bool,
|
|
) -> CLIPVisionModelWithProjection:
|
|
return CLIPVisionModelWithProjection.from_pretrained( # type: ignore
|
|
unclip21_transformers_stabilityai_path,
|
|
local_files_only=use_local_weights,
|
|
subfolder="image_encoder",
|
|
).to(device=test_device, dtype=test_dtype_fp32_bf16_fp16) # type: ignore
|
|
|
|
|
|
@no_grad()
|
|
@pytest.mark.flaky(reruns=3)
|
|
def test_encoder(
|
|
ref_encoder: CLIPVisionModelWithProjection,
|
|
our_encoder: CLIPImageEncoderH,
|
|
):
|
|
assert ref_encoder.dtype == our_encoder.dtype
|
|
assert ref_encoder.device == our_encoder.device
|
|
x = torch.randn((1, 3, 224, 224), dtype=ref_encoder.dtype, device=ref_encoder.device)
|
|
|
|
ref_embeddings = ref_encoder(x).image_embeds
|
|
our_embeddings = our_encoder(x)
|
|
|
|
assert ref_embeddings.shape == (1, 1024)
|
|
assert our_embeddings.shape == (1, 1024)
|
|
|
|
assert torch.allclose(our_embeddings, ref_embeddings, atol=0.05)
|