mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-24 20:59:27 +03:00
448 lines
20 KiB
Python
448 lines
20 KiB
Python
# coding=utf-8
|
|
# Copyright 2026 HuggingFace Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import re
|
|
|
|
import numpy as np
|
|
import PIL.Image
|
|
import pytest
|
|
import torch
|
|
|
|
from diffusers.image_processor import VaeImageProcessor
|
|
from diffusers.utils.constants import (
|
|
DECODE_ENDPOINT_FLUX,
|
|
DECODE_ENDPOINT_HUNYUAN_VIDEO,
|
|
DECODE_ENDPOINT_SD_V1,
|
|
DECODE_ENDPOINT_SD_XL,
|
|
)
|
|
from diffusers.utils.remote_utils import (
|
|
remote_decode,
|
|
)
|
|
from diffusers.video_processor import VideoProcessor
|
|
|
|
from ..testing_utils import (
|
|
enable_full_determinism,
|
|
slow,
|
|
torch_all_close,
|
|
torch_device,
|
|
)
|
|
|
|
|
|
enable_full_determinism()
|
|
|
|
|
|
class RemoteAutoencoderKLMixin:
|
|
shape: tuple[int, ...] = None
|
|
out_hw: tuple[int, int] = None
|
|
endpoint: str = None
|
|
dtype: torch.dtype = None
|
|
scaling_factor: float = None
|
|
shift_factor: float = None
|
|
processor_cls: VaeImageProcessor | VideoProcessor = None
|
|
output_pil_slice: torch.Tensor = None
|
|
output_pt_slice: torch.Tensor = None
|
|
partial_postprocess_return_pt_slice: torch.Tensor = None
|
|
return_pt_slice: torch.Tensor = None
|
|
width: int = None
|
|
height: int = None
|
|
|
|
def get_dummy_inputs(self):
|
|
# The tensor is serialized and decoded remotely, so the local device has no bearing on the result.
|
|
# Seed it on CPU and move it afterwards, otherwise `torch.randn` yields different latents on a CUDA
|
|
# runner than on a CPU one and the reference slices below only match on whichever recorded them.
|
|
inputs = {
|
|
"endpoint": self.endpoint,
|
|
"tensor": torch.randn(
|
|
self.shape,
|
|
device="cpu",
|
|
dtype=self.dtype,
|
|
generator=torch.Generator("cpu").manual_seed(13),
|
|
).to(torch_device),
|
|
"scaling_factor": self.scaling_factor,
|
|
"shift_factor": self.shift_factor,
|
|
"height": self.height,
|
|
"width": self.width,
|
|
}
|
|
return inputs
|
|
|
|
def test_no_scaling(self):
|
|
inputs = self.get_dummy_inputs()
|
|
if inputs["scaling_factor"] is not None:
|
|
inputs["tensor"] = inputs["tensor"] / inputs["scaling_factor"]
|
|
inputs["scaling_factor"] = None
|
|
if inputs["shift_factor"] is not None:
|
|
inputs["tensor"] = inputs["tensor"] + inputs["shift_factor"]
|
|
inputs["shift_factor"] = None
|
|
processor = self.processor_cls()
|
|
output = remote_decode(
|
|
output_type="pt",
|
|
# required for now, will be removed in next update
|
|
do_scaling=False,
|
|
processor=processor,
|
|
**inputs,
|
|
)
|
|
assert isinstance(output, PIL.Image.Image), f"Expected `PIL.Image.Image` output, got {type(output)}"
|
|
assert output.height == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.height}"
|
|
assert output.width == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.width}"
|
|
output_slice = torch.from_numpy(np.array(output)[0, -3:, -3:].flatten())
|
|
# Increased tolerance for Flux Packed diff [1, 0, 1, 0, 0, 0, 0, 0, 0]
|
|
assert torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1, atol=1)
|
|
|
|
def test_output_type_pt(self):
|
|
inputs = self.get_dummy_inputs()
|
|
processor = self.processor_cls()
|
|
output = remote_decode(output_type="pt", processor=processor, **inputs)
|
|
assert isinstance(output, PIL.Image.Image), f"Expected `PIL.Image.Image` output, got {type(output)}"
|
|
assert output.height == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.height}"
|
|
assert output.width == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.width}"
|
|
output_slice = torch.from_numpy(np.array(output)[0, -3:, -3:].flatten())
|
|
assert torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1e-2)
|
|
|
|
# output is visually the same, slice is flaky?
|
|
def test_output_type_pil(self):
|
|
inputs = self.get_dummy_inputs()
|
|
output = remote_decode(output_type="pil", **inputs)
|
|
assert isinstance(output, PIL.Image.Image), f"Expected `PIL.Image.Image` output, got {type(output)}"
|
|
assert output.height == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.height}"
|
|
assert output.width == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.width}"
|
|
|
|
def test_output_type_pil_image_format(self):
|
|
inputs = self.get_dummy_inputs()
|
|
output = remote_decode(output_type="pil", image_format="png", **inputs)
|
|
assert isinstance(output, PIL.Image.Image), f"Expected `PIL.Image.Image` output, got {type(output)}"
|
|
assert output.height == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.height}"
|
|
assert output.width == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.width}"
|
|
assert output.format == "png", f"Expected image format `png`, got {output.format}"
|
|
output_slice = torch.from_numpy(np.array(output)[0, -3:, -3:].flatten())
|
|
assert torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1e-2)
|
|
|
|
def test_output_type_pt_partial_postprocess(self):
|
|
inputs = self.get_dummy_inputs()
|
|
output = remote_decode(output_type="pt", partial_postprocess=True, **inputs)
|
|
assert isinstance(output, PIL.Image.Image), f"Expected `PIL.Image.Image` output, got {type(output)}"
|
|
assert output.height == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.height}"
|
|
assert output.width == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.width}"
|
|
output_slice = torch.from_numpy(np.array(output)[0, -3:, -3:].flatten())
|
|
assert torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1e-2)
|
|
|
|
def test_output_type_pt_return_type_pt(self):
|
|
inputs = self.get_dummy_inputs()
|
|
output = remote_decode(output_type="pt", return_type="pt", **inputs)
|
|
assert isinstance(output, torch.Tensor), f"Expected `torch.Tensor` output, got {type(output)}"
|
|
assert output.shape[2] == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.shape[2]}"
|
|
assert output.shape[3] == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.shape[3]}"
|
|
output_slice = output[0, 0, -3:, -3:].flatten()
|
|
assert torch_all_close(output_slice, self.return_pt_slice.to(output_slice.dtype), rtol=1e-3, atol=1e-3)
|
|
|
|
def test_output_type_pt_partial_postprocess_return_type_pt(self):
|
|
inputs = self.get_dummy_inputs()
|
|
output = remote_decode(output_type="pt", partial_postprocess=True, return_type="pt", **inputs)
|
|
assert isinstance(output, torch.Tensor), f"Expected `torch.Tensor` output, got {type(output)}"
|
|
assert output.shape[1] == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.shape[1]}"
|
|
assert output.shape[2] == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.shape[2]}"
|
|
output_slice = output[0, -3:, -3:, 0].flatten().cpu()
|
|
assert torch_all_close(
|
|
output_slice, self.partial_postprocess_return_pt_slice.to(output_slice.dtype), rtol=1e-2
|
|
)
|
|
|
|
def test_do_scaling_deprecation(self):
|
|
inputs = self.get_dummy_inputs()
|
|
inputs.pop("scaling_factor", None)
|
|
inputs.pop("shift_factor", None)
|
|
expected = "`do_scaling` is deprecated, pass `scaling_factor` and `shift_factor` if required."
|
|
with pytest.warns(FutureWarning, match=re.escape(expected)):
|
|
_ = remote_decode(output_type="pt", partial_postprocess=True, **inputs)
|
|
|
|
def test_input_tensor_type_base64_deprecation(self):
|
|
inputs = self.get_dummy_inputs()
|
|
expected = "input_tensor_type='base64' is deprecated. Using `binary`."
|
|
with pytest.warns(FutureWarning, match=re.escape(expected)):
|
|
_ = remote_decode(output_type="pt", input_tensor_type="base64", partial_postprocess=True, **inputs)
|
|
|
|
def test_output_tensor_type_base64_deprecation(self):
|
|
inputs = self.get_dummy_inputs()
|
|
expected = "output_tensor_type='base64' is deprecated. Using `binary`."
|
|
with pytest.warns(FutureWarning, match=re.escape(expected)):
|
|
_ = remote_decode(output_type="pt", output_tensor_type="base64", partial_postprocess=True, **inputs)
|
|
|
|
|
|
class RemoteAutoencoderKLHunyuanVideoMixin(RemoteAutoencoderKLMixin):
|
|
def test_no_scaling(self):
|
|
inputs = self.get_dummy_inputs()
|
|
if inputs["scaling_factor"] is not None:
|
|
inputs["tensor"] = inputs["tensor"] / inputs["scaling_factor"]
|
|
inputs["scaling_factor"] = None
|
|
if inputs["shift_factor"] is not None:
|
|
inputs["tensor"] = inputs["tensor"] + inputs["shift_factor"]
|
|
inputs["shift_factor"] = None
|
|
processor = self.processor_cls()
|
|
output = remote_decode(
|
|
output_type="pt",
|
|
# required for now, will be removed in next update
|
|
do_scaling=False,
|
|
processor=processor,
|
|
**inputs,
|
|
)
|
|
assert isinstance(output, list) and isinstance(output[0], PIL.Image.Image), (
|
|
f"Expected `List[PIL.Image.Image]` output, got {type(output)}"
|
|
)
|
|
assert output[0].height == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output[0].height}"
|
|
assert output[0].width == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output[0].width}"
|
|
output_slice = torch.from_numpy(np.array(output[0])[0, -3:, -3:].flatten())
|
|
assert torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1, atol=1)
|
|
|
|
def test_output_type_pt(self):
|
|
inputs = self.get_dummy_inputs()
|
|
processor = self.processor_cls()
|
|
output = remote_decode(output_type="pt", processor=processor, **inputs)
|
|
assert isinstance(output, list) and isinstance(output[0], PIL.Image.Image), (
|
|
f"Expected `List[PIL.Image.Image]` output, got {type(output)}"
|
|
)
|
|
assert output[0].height == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output[0].height}"
|
|
assert output[0].width == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output[0].width}"
|
|
output_slice = torch.from_numpy(np.array(output[0])[0, -3:, -3:].flatten())
|
|
assert torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1, atol=1)
|
|
|
|
# output is visually the same, slice is flaky?
|
|
def test_output_type_pil(self):
|
|
inputs = self.get_dummy_inputs()
|
|
processor = self.processor_cls()
|
|
output = remote_decode(output_type="pil", processor=processor, **inputs)
|
|
assert isinstance(output, list) and isinstance(output[0], PIL.Image.Image), (
|
|
f"Expected `List[PIL.Image.Image]` output, got {type(output)}"
|
|
)
|
|
assert output[0].height == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output[0].height}"
|
|
assert output[0].width == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output[0].width}"
|
|
|
|
def test_output_type_pil_image_format(self):
|
|
inputs = self.get_dummy_inputs()
|
|
processor = self.processor_cls()
|
|
output = remote_decode(output_type="pil", processor=processor, image_format="png", **inputs)
|
|
assert isinstance(output, list) and isinstance(output[0], PIL.Image.Image), (
|
|
f"Expected `List[PIL.Image.Image]` output, got {type(output)}"
|
|
)
|
|
assert output[0].height == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output[0].height}"
|
|
assert output[0].width == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output[0].width}"
|
|
output_slice = torch.from_numpy(np.array(output[0])[0, -3:, -3:].flatten())
|
|
assert torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1, atol=1)
|
|
|
|
def test_output_type_pt_partial_postprocess(self):
|
|
inputs = self.get_dummy_inputs()
|
|
output = remote_decode(output_type="pt", partial_postprocess=True, **inputs)
|
|
assert isinstance(output, list) and isinstance(output[0], PIL.Image.Image), (
|
|
f"Expected `List[PIL.Image.Image]` output, got {type(output)}"
|
|
)
|
|
assert output[0].height == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output[0].height}"
|
|
assert output[0].width == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output[0].width}"
|
|
output_slice = torch.from_numpy(np.array(output[0])[0, -3:, -3:].flatten())
|
|
assert torch_all_close(output_slice, self.output_pt_slice.to(output_slice.dtype), rtol=1, atol=1)
|
|
|
|
def test_output_type_pt_return_type_pt(self):
|
|
inputs = self.get_dummy_inputs()
|
|
output = remote_decode(output_type="pt", return_type="pt", **inputs)
|
|
assert isinstance(output, torch.Tensor), f"Expected `torch.Tensor` output, got {type(output)}"
|
|
assert output.shape[3] == self.out_hw[0], f"Expected image height {self.out_hw[0]}, got {output.shape[3]}"
|
|
assert output.shape[4] == self.out_hw[1], f"Expected image width {self.out_hw[0]}, got {output.shape[4]}"
|
|
output_slice = output[0, 0, 0, -3:, -3:].flatten()
|
|
assert torch_all_close(output_slice, self.return_pt_slice.to(output_slice.dtype), rtol=1e-3, atol=1e-3)
|
|
|
|
def test_output_type_mp4(self):
|
|
inputs = self.get_dummy_inputs()
|
|
output = remote_decode(output_type="mp4", return_type="mp4", **inputs)
|
|
assert isinstance(output, bytes), f"Expected `bytes` output, got {type(output)}"
|
|
|
|
|
|
# The tests below hit live HF Inference Endpoints, which are not part of the fast CI contract, so they are
|
|
# gated behind `RUN_SLOW` instead of running on every push. The reference slices are recorded against the
|
|
# deployed `DECODE_ENDPOINT_*` models and drift if those are redeployed.
|
|
@slow
|
|
class TestRemoteAutoencoderKLSDv1(RemoteAutoencoderKLMixin):
|
|
shape = (
|
|
1,
|
|
4,
|
|
64,
|
|
64,
|
|
)
|
|
out_hw = (
|
|
512,
|
|
512,
|
|
)
|
|
endpoint = DECODE_ENDPOINT_SD_V1
|
|
dtype = torch.float16
|
|
scaling_factor = 0.18215
|
|
shift_factor = None
|
|
processor_cls = VaeImageProcessor
|
|
output_pt_slice = torch.tensor([162, 131, 123, 158, 131, 124, 148, 128, 115], dtype=torch.uint8)
|
|
partial_postprocess_return_pt_slice = torch.tensor(
|
|
[133, 146, 130, 161, 141, 132, 148, 134, 135], dtype=torch.uint8
|
|
)
|
|
return_pt_slice = torch.tensor([0.0400, 0.1448, 0.0217, 0.2649, 0.1049, 0.0376, 0.1622, 0.0552, 0.0574])
|
|
|
|
|
|
@slow
|
|
class TestRemoteAutoencoderKLSDXL(RemoteAutoencoderKLMixin):
|
|
shape = (
|
|
1,
|
|
4,
|
|
128,
|
|
128,
|
|
)
|
|
out_hw = (
|
|
1024,
|
|
1024,
|
|
)
|
|
endpoint = DECODE_ENDPOINT_SD_XL
|
|
dtype = torch.float16
|
|
scaling_factor = 0.13025
|
|
shift_factor = None
|
|
processor_cls = VaeImageProcessor
|
|
output_pt_slice = torch.tensor([133, 182, 166, 131, 179, 163, 122, 173, 151], dtype=torch.uint8)
|
|
partial_postprocess_return_pt_slice = torch.tensor(
|
|
[172, 160, 147, 149, 136, 126, 118, 112, 107], dtype=torch.uint8
|
|
)
|
|
return_pt_slice = torch.tensor([0.3511, 0.2539, 0.1550, 0.1720, 0.0688, -0.0117, -0.0730, -0.1215, -0.1581])
|
|
|
|
|
|
@slow
|
|
class TestRemoteAutoencoderKLFlux(RemoteAutoencoderKLMixin):
|
|
shape = (
|
|
1,
|
|
16,
|
|
128,
|
|
128,
|
|
)
|
|
out_hw = (
|
|
1024,
|
|
1024,
|
|
)
|
|
endpoint = DECODE_ENDPOINT_FLUX
|
|
dtype = torch.bfloat16
|
|
scaling_factor = 0.3611
|
|
shift_factor = 0.1159
|
|
processor_cls = VaeImageProcessor
|
|
output_pt_slice = torch.tensor([236, 222, 171, 226, 208, 160, 210, 193, 154], dtype=torch.uint8)
|
|
partial_postprocess_return_pt_slice = torch.tensor(
|
|
[173, 174, 173, 188, 182, 178, 183, 175, 167], dtype=torch.uint8
|
|
)
|
|
return_pt_slice = torch.tensor([0.3538, 0.3613, 0.3582, 0.4749, 0.4294, 0.3977, 0.4336, 0.3738, 0.3098])
|
|
|
|
|
|
@slow
|
|
class TestRemoteAutoencoderKLFluxPacked(RemoteAutoencoderKLMixin):
|
|
shape = (
|
|
1,
|
|
4096,
|
|
64,
|
|
)
|
|
out_hw = (
|
|
1024,
|
|
1024,
|
|
)
|
|
height = 1024
|
|
width = 1024
|
|
endpoint = DECODE_ENDPOINT_FLUX
|
|
dtype = torch.bfloat16
|
|
scaling_factor = 0.3611
|
|
shift_factor = 0.1159
|
|
processor_cls = VaeImageProcessor
|
|
# slices are different due to randn on different shape. we can pack the latent instead if we want the same
|
|
output_pt_slice = torch.tensor([203, 188, 127, 163, 140, 89, 113, 86, 50], dtype=torch.uint8)
|
|
partial_postprocess_return_pt_slice = torch.tensor(
|
|
[194, 173, 131, 192, 168, 126, 187, 158, 129], dtype=torch.uint8
|
|
)
|
|
return_pt_slice = torch.tensor([0.5259, 0.3564, 0.0272, 0.5093, 0.3171, -0.0111, 0.4639, 0.2371, 0.0083])
|
|
|
|
|
|
@slow
|
|
class TestRemoteAutoencoderKLHunyuanVideo(RemoteAutoencoderKLHunyuanVideoMixin):
|
|
shape = (
|
|
1,
|
|
16,
|
|
3,
|
|
40,
|
|
64,
|
|
)
|
|
out_hw = (
|
|
320,
|
|
512,
|
|
)
|
|
endpoint = DECODE_ENDPOINT_HUNYUAN_VIDEO
|
|
dtype = torch.float16
|
|
scaling_factor = 0.476986
|
|
processor_cls = VideoProcessor
|
|
output_pt_slice = torch.tensor([126, 129, 195, 128, 130, 196, 129, 130, 196], dtype=torch.uint8)
|
|
partial_postprocess_return_pt_slice = torch.tensor([19, 16, 13, 30, 25, 21, 39, 35, 31], dtype=torch.uint8)
|
|
return_pt_slice = torch.tensor([-0.8501, -0.8750, -0.8979, -0.7681, -0.8027, -0.8389, -0.6958, -0.7222, -0.7593])
|
|
|
|
|
|
class RemoteAutoencoderKLSlowTestMixin:
|
|
channels: int = 4
|
|
endpoint: str = None
|
|
dtype: torch.dtype = None
|
|
scaling_factor: float = None
|
|
shift_factor: float = None
|
|
width: int = None
|
|
height: int = None
|
|
|
|
def get_dummy_inputs(self):
|
|
inputs = {
|
|
"endpoint": self.endpoint,
|
|
"scaling_factor": self.scaling_factor,
|
|
"shift_factor": self.shift_factor,
|
|
"height": self.height,
|
|
"width": self.width,
|
|
}
|
|
return inputs
|
|
|
|
def test_multi_res(self, tmp_path):
|
|
inputs = self.get_dummy_inputs()
|
|
for height in {320, 512, 640, 704, 896, 1024, 1208, 1384, 1536, 1608, 1864, 2048}:
|
|
for width in {320, 512, 640, 704, 896, 1024, 1208, 1384, 1536, 1608, 1864, 2048}:
|
|
inputs["tensor"] = torch.randn(
|
|
(1, self.channels, height // 8, width // 8),
|
|
device="cpu",
|
|
dtype=self.dtype,
|
|
generator=torch.Generator("cpu").manual_seed(13),
|
|
).to(torch_device)
|
|
inputs["height"] = height
|
|
inputs["width"] = width
|
|
output = remote_decode(output_type="pt", partial_postprocess=True, **inputs)
|
|
output.save(tmp_path / f"test_multi_res_{height}_{width}.png")
|
|
|
|
|
|
@slow
|
|
class TestRemoteAutoencoderKLSDv1Slow(RemoteAutoencoderKLSlowTestMixin):
|
|
endpoint = DECODE_ENDPOINT_SD_V1
|
|
dtype = torch.float16
|
|
scaling_factor = 0.18215
|
|
shift_factor = None
|
|
|
|
|
|
@slow
|
|
class TestRemoteAutoencoderKLSDXLSlow(RemoteAutoencoderKLSlowTestMixin):
|
|
endpoint = DECODE_ENDPOINT_SD_XL
|
|
dtype = torch.float16
|
|
scaling_factor = 0.13025
|
|
shift_factor = None
|
|
|
|
|
|
@slow
|
|
class TestRemoteAutoencoderKLFluxSlow(RemoteAutoencoderKLSlowTestMixin):
|
|
channels = 16
|
|
endpoint = DECODE_ENDPOINT_FLUX
|
|
dtype = torch.bfloat16
|
|
scaling_factor = 0.3611
|
|
shift_factor = 0.1159
|