Fix deprecation warning for torch.utils._pytree._register_pytree_node in PyTorch 2.2 (#7008)

Fixed deprecation warning for torch.utils._pytree._register_pytree_node in PyTorch 2.2

Co-authored-by: Yinghua <[email protected]>

Author: zyinghua

src/diffusers/utils/outputs.py

@@ -21,7 +21,7 @@
 
 import numpy as np
 
-from .import_utils import is_torch_available
+from .import_utils import is_torch_available, is_torch_version
 
 
 def is_tensor(x) -> bool:
@@ -60,11 +60,18 @@ def __init_subclass__(cls) -> None:
         if is_torch_available():
             import torch.utils._pytree
 
-            torch.utils._pytree._register_pytree_node(
-                cls,
-                torch.utils._pytree._dict_flatten,
-                lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)),
-            )
+            if is_torch_version("<", "2.2"):
+                torch.utils._pytree._register_pytree_node(
+                    cls,
+                    torch.utils._pytree._dict_flatten,
+                    lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)),
+                )
+            else:
+                torch.utils._pytree.register_pytree_node(
+                    cls,
+                    torch.utils._pytree._dict_flatten,
+                    lambda values, context: cls(**torch.utils._pytree._dict_unflatten(values, context)),
+                )
 
     def __post_init__(self) -> None:
         class_fields = fields(self)