From f8b292d67da58fd4f1dd205a85064b4f22f05acd Mon Sep 17 00:00:00 2001 From: Aamir Nazir Date: Wed, 8 Jan 2025 22:13:21 +0400 Subject: [PATCH] Update outputs.py --- src/diffusers/utils/outputs.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/diffusers/utils/outputs.py b/src/diffusers/utils/outputs.py index 6080a86b871a..2e0cdd56e3d7 100644 --- a/src/diffusers/utils/outputs.py +++ b/src/diffusers/utils/outputs.py @@ -65,12 +65,14 @@ def __init_subclass__(cls) -> None: cls, torch.utils._pytree._dict_flatten, lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)), + serialized_type_name=f"{cls.__module__}.{cls.__name__}" ) else: torch.utils._pytree.register_pytree_node( cls, torch.utils._pytree._dict_flatten, lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)), + serialized_type_name=f"{cls.__module__}.{cls.__name__}" ) def __post_init__(self) -> None: