test_modeling_common_flax.py

importinspect

fromdiffusers.utilsimportis_flax_available
fromdiffusers.utils.testing_utilsimportrequire_flax


ifis_flax_available():
importjax


@require_flax
classFlaxModelTesterMixin:
deftest_output(self):
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()

model=self.model_class(**init_dict)
variables=model.init(inputs_dict["prng_key"], inputs_dict["sample"])
jax.lax.stop_gradient(variables)

output=model.apply(variables, inputs_dict["sample"])

ifisinstance(output, dict):
output=output.sample

self.assertIsNotNone(output)
expected_shape=inputs_dict["sample"].shape
self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")

deftest_forward_with_norm_groups(self):
init_dict, inputs_dict=self.prepare_init_args_and_inputs_for_common()

init_dict["norm_num_groups"] =16
init_dict["block_out_channels"] = (16, 32)

model=self.model_class(**init_dict)
variables=model.init(inputs_dict["prng_key"], inputs_dict["sample"])
jax.lax.stop_gradient(variables)

output=model.apply(variables, inputs_dict["sample"])

ifisinstance(output, dict):
output=output.sample

self.assertIsNotNone(output)
expected_shape=inputs_dict["sample"].shape
self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")

deftest_deprecated_kwargs(self):
has_kwarg_in_model_class="kwargs"ininspect.signature(self.model_class.__init__).parameters
has_deprecated_kwarg=len(self.model_class._deprecated_kwargs) >0

ifhas_kwarg_in_model_classandnothas_deprecated_kwarg:
raiseValueError(
f"{self.model_class} has `**kwargs` in its __init__ method but has not defined any deprecated kwargs"
" under the `_deprecated_kwargs` class attribute. Make sure to either remove `**kwargs` if there are"
" no deprecated arguments or add the deprecated argument with `_deprecated_kwargs ="
" [<deprecated_argument>]`"
 )

ifnothas_kwarg_in_model_classandhas_deprecated_kwarg:
raiseValueError(
f"{self.model_class} doesn't have `**kwargs` in its __init__ method but has defined deprecated kwargs"
" under the `_deprecated_kwargs` class attribute. Make sure to either add the `**kwargs` argument to"
f" {self.model_class}.__init__ if there are deprecated arguments or remove the deprecated argument"
" from `_deprecated_kwargs = [<deprecated_argument>]`"
 )