Muyang Li
db978b4757
* better linter * update deprecation information * support both model loading * finished merge models * update a fill removal workflow * add legacy test scripts * add a conversion workflow * fix the suffix * update workflows * download the new models * start implementing the 4-bit t5 * figuring out how to overwrite * figuring out how to overwrite * default workflow is runnable * my customized loader for work naive clip * ready to add tests * ready for tests
73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
import gc
|
|
import os
|
|
import subprocess
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from diffusers.utils import load_image
|
|
from PIL import Image
|
|
from torchmetrics.image import LearnedPerceptualImagePatchSimilarity, PeakSignalNoiseRatio
|
|
from torchmetrics.multimodal import CLIPImageQualityAssessment
|
|
|
|
from nunchaku.utils import get_precision
|
|
|
|
script_dir = os.path.join(os.path.dirname(__file__), "scripts")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"script_name, expected_clip_iqa, expected_lpips, expected_psnr",
|
|
[
|
|
("nunchaku_flux1_redux_dev.py", 0.9, 0.137, 21.7),
|
|
("nunchaku_flux1_dev_controlnet_upscaler.py", 0.9, 0.1, 26),
|
|
("nunchaku_flux1_dev_controlnet_union_pro2.py", 0.9, 0.1, 26),
|
|
("nunchaku_flux1_depth_lora.py", 0.7, 0.1, 26),
|
|
("nunchaku_flux1_canny.py", 0.9, 0.1, 26),
|
|
("nunchaku_flux1_schnell.py", 0.9, 0.29, 19.3),
|
|
("nunchaku_flux1_depth.py", 0.9, 0.1, 26),
|
|
("nunchaku_flux1_shuttle_jaguar.py", 0.9, 0.14, 23.9),
|
|
("nunchaku_flux1_fill.py", 0.9, 0.1, 26),
|
|
("nunchaku_flux1_fill_removalV2.py", 0.63, 0.13, 26),
|
|
("nunchaku_flux1_dev.py", 0.9, 0.21, 19.7),
|
|
("nunchaku_flux1_canny_lora.py", 0.9, 0.1, 25),
|
|
("nunchaku_flux1_dev_qencoder.py", 0.9, 0.1, 22.8),
|
|
],
|
|
)
|
|
def test_workflows(script_name: str, expected_clip_iqa: float, expected_lpips: float, expected_psnr: float):
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
script_path = os.path.join(script_dir, script_name)
|
|
|
|
result = subprocess.run(["python", script_path], capture_output=True, text=True)
|
|
print(f"Running {script_path} -> Return code: {result.returncode}")
|
|
assert result.returncode == 0, f"{script_path} failed with code {result.returncode}"
|
|
|
|
path = open("image_path.txt", "r").read().strip()
|
|
|
|
# clip_iqa metric
|
|
metric = CLIPImageQualityAssessment(model_name_or_path="openai/clip-vit-large-patch14").to("cuda")
|
|
image = Image.open(path).convert("RGB")
|
|
gen_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).to(torch.float32).unsqueeze(0).to("cuda")
|
|
clip_iqa = metric(gen_tensor).item()
|
|
print(f"CLIP-IQA: {clip_iqa}")
|
|
|
|
# lpips metric
|
|
ref_image_url = (
|
|
f"https://huggingface.co/mit-han-lab/nunchaku-artifacts/resolve/main/ComfyUI-nunchaku/ref_images/"
|
|
f"{get_precision()}/{script_name.replace('.py', '.png')}"
|
|
)
|
|
ref_image = load_image(ref_image_url).convert("RGB")
|
|
metric = LearnedPerceptualImagePatchSimilarity().to("cuda")
|
|
ref_tensor = torch.from_numpy(np.array(ref_image)).permute(2, 0, 1).to(torch.float32)
|
|
ref_tensor = ref_tensor.unsqueeze(0).to("cuda")
|
|
lpips = metric(gen_tensor / 255, ref_tensor / 255).item()
|
|
print(f"LPIPS: {lpips}")
|
|
|
|
metric = PeakSignalNoiseRatio(data_range=(0, 255)).cuda()
|
|
psnr = metric(gen_tensor, ref_tensor).item()
|
|
print(f"PSNR: {psnr}")
|
|
|
|
assert clip_iqa >= expected_clip_iqa * 0.9
|
|
assert lpips <= expected_lpips * 1.1
|
|
assert psnr >= expected_psnr * 0.9
|