diffusers/tests/others/test_outputs.py at main · huggingface/diffusers · GitHub

Name: diffusers/tests/others/test_outputs.py at main · huggingface/diffusers · GitHub
Rating: 4.8 (2421 reviews)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
importpickleaspkl
importunittest
fromdataclassesimportdataclass
fromtypingimportList, Union

importnumpyasnp
importPIL.Image

fromdiffusers.utils.outputsimportBaseOutput
fromdiffusers.utils.testing_utilsimportrequire_torch


@dataclass
classCustomOutput(BaseOutput):
images: Union[List[PIL.Image.Image], np.ndarray]


classConfigTester(unittest.TestCase):
deftest_outputs_single_attribute(self):
outputs=CustomOutput(images=np.random.rand(1, 3, 4, 4))

# check every way of getting the attribute
assertisinstance(outputs.images, np.ndarray)
assertoutputs.images.shape== (1, 3, 4, 4)
assertisinstance(outputs["images"], np.ndarray)
assertoutputs["images"].shape== (1, 3, 4, 4)
assertisinstance(outputs[0], np.ndarray)
assertoutputs[0].shape== (1, 3, 4, 4)

# test with a non-tensor attribute
outputs=CustomOutput(images=[PIL.Image.new("RGB", (4, 4))])

# check every way of getting the attribute
assertisinstance(outputs.images, list)
assertisinstance(outputs.images[0], PIL.Image.Image)
assertisinstance(outputs["images"], list)
assertisinstance(outputs["images"][0], PIL.Image.Image)
assertisinstance(outputs[0], list)
assertisinstance(outputs[0][0], PIL.Image.Image)

deftest_outputs_dict_init(self):
# test output reinitialization with a `dict` for compatibility with `accelerate`
outputs=CustomOutput({"images": np.random.rand(1, 3, 4, 4)})

# check every way of getting the attribute
assertisinstance(outputs.images, np.ndarray)
assertoutputs.images.shape== (1, 3, 4, 4)
assertisinstance(outputs["images"], np.ndarray)
assertoutputs["images"].shape== (1, 3, 4, 4)
assertisinstance(outputs[0], np.ndarray)
assertoutputs[0].shape== (1, 3, 4, 4)

# test with a non-tensor attribute
outputs=CustomOutput({"images": [PIL.Image.new("RGB", (4, 4))]})

# check every way of getting the attribute
assertisinstance(outputs.images, list)
assertisinstance(outputs.images[0], PIL.Image.Image)
assertisinstance(outputs["images"], list)
assertisinstance(outputs["images"][0], PIL.Image.Image)
assertisinstance(outputs[0], list)
assertisinstance(outputs[0][0], PIL.Image.Image)

deftest_outputs_serialization(self):
outputs_orig=CustomOutput(images=[PIL.Image.new("RGB", (4, 4))])
serialized=pkl.dumps(outputs_orig)
outputs_copy=pkl.loads(serialized)

# Check original and copy are equal
assertdir(outputs_orig) ==dir(outputs_copy)
assertdict(outputs_orig) ==dict(outputs_copy)
assertvars(outputs_orig) ==vars(outputs_copy)

@require_torch
deftest_torch_pytree(self):
# ensure torch.utils._pytree treats ModelOutput subclasses as nodes (and not leaves)
# this is important for DistributedDataParallel gradient synchronization with static_graph=True
importtorch
importtorch.utils._pytree

data=np.random.rand(1, 3, 4, 4)
x=CustomOutput(images=data)
self.assertFalse(torch.utils._pytree._is_leaf(x))

expected_flat_outs= [data]
expected_tree_spec=torch.utils._pytree.TreeSpec(CustomOutput, ["images"], [torch.utils._pytree.LeafSpec()])

actual_flat_outs, actual_tree_spec=torch.utils._pytree.tree_flatten(x)
self.assertEqual(expected_flat_outs, actual_flat_outs)
self.assertEqual(expected_tree_spec, actual_tree_spec)

unflattened_x=torch.utils._pytree.tree_unflatten(actual_flat_outs, actual_tree_spec)
self.assertEqual(x, unflattened_x)