fix(diffusers): correct usage of Cell.name_cells()

mindspore-lab · Jun 6, 2024 · 29bdbb9 · 29bdbb9
1 parent d5176fc
commit 29bdbb9
Show file tree

Hide file tree

Showing 12 changed files with 62 additions and 62 deletions.
diff --git a/mindone/diffusers/models/autoencoders/autoencoder_kl.py b/mindone/diffusers/models/autoencoders/autoencoder_kl.py
@@ -136,12 +136,12 @@ def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str
             if hasattr(module, "get_processor"):
                 processors[f"{name}.processor"] = module.get_processor()
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
 
             return processors
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_add_processors(name, module, processors)
 
         return processors
@@ -175,10 +175,10 @@ def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
                 else:
                     module.set_processor(processor.pop(f"{name}.processor"))
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_attn_processor(name, module, processor)
 
     # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor

diff --git a/mindone/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py b/mindone/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py
@@ -219,12 +219,12 @@ def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str
             if hasattr(module, "get_processor"):
                 processors[f"{name}.processor"] = module.get_processor()
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
 
             return processors
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_add_processors(name, module, processors)
 
         return processors
@@ -255,10 +255,10 @@ def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
                 else:
                     module.set_processor(processor.pop(f"{name}.processor"))
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_attn_processor(name, module, processor)
 
     def set_default_attn_processor(self):

diff --git a/mindone/diffusers/models/autoencoders/consistency_decoder_vae.py b/mindone/diffusers/models/autoencoders/consistency_decoder_vae.py
@@ -157,12 +157,12 @@ def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str
             if hasattr(module, "get_processor"):
                 processors[f"{name}.processor"] = module.get_processor()
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
 
             return processors
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_add_processors(name, module, processors)
 
         return processors
@@ -193,10 +193,10 @@ def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
                 else:
                     module.set_processor(processor.pop(f"{name}.processor"))
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_attn_processor(name, module, processor)
 
     # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor

diff --git a/mindone/diffusers/models/controlnet.py b/mindone/diffusers/models/controlnet.py
@@ -549,12 +549,12 @@ def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str
             if hasattr(module, "get_processor"):
                 processors[f"{name}.processor"] = module.get_processor()
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
 
             return processors
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_add_processors(name, module, processors)
 
         return processors
@@ -585,10 +585,10 @@ def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
                 else:
                     module.set_processor(processor.pop(f"{name}.processor"))
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_attn_processor(name, module, processor)
 
     # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor

diff --git a/mindone/diffusers/models/transformers/prior_transformer.py b/mindone/diffusers/models/transformers/prior_transformer.py
@@ -176,12 +176,12 @@ def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str
             if hasattr(module, "get_processor"):
                 processors[f"{name}.processor"] = module.get_processor()
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
 
             return processors
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_add_processors(name, module, processors)
 
         return processors
@@ -212,10 +212,10 @@ def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
                 else:
                     module.set_processor(processor.pop(f"{name}.processor"))
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_attn_processor(name, module, processor)
 
     # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor

diff --git a/mindone/diffusers/models/unets/unet_2d_condition.py b/mindone/diffusers/models/unets/unet_2d_condition.py
@@ -708,12 +708,12 @@ def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str
             if hasattr(module, "get_processor"):
                 processors[f"{name}.processor"] = module.get_processor()
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
 
             return processors
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_add_processors(name, module, processors)
 
         return processors
@@ -746,10 +746,10 @@ def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
                 else:
                     module.set_processor(processor.pop(f"{name}.processor"))
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_attn_processor(name, module, processor)
 
     def set_default_attn_processor(self):

diff --git a/mindone/diffusers/models/unets/unet_3d_condition.py b/mindone/diffusers/models/unets/unet_3d_condition.py
@@ -308,12 +308,12 @@ def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str
             if hasattr(module, "get_processor"):
                 processors[f"{name}.processor"] = module.get_processor()
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
 
             return processors
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_add_processors(name, module, processors)
 
         return processors
@@ -344,10 +344,10 @@ def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
                 else:
                     module.set_processor(processor.pop(f"{name}.processor"))
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_attn_processor(name, module, processor)
 
     def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
@@ -373,21 +373,21 @@ def fn_recursive_feed_forward(module: nn.Cell, chunk_size: int, dim: int):
             if hasattr(module, "set_chunk_feed_forward"):
                 module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
 
-            for child in module.name_cells():
+            for child in module.name_cells().items():
                 fn_recursive_feed_forward(child, chunk_size, dim)
 
-        for module in self.name_cells():
+        for module in self.name_cells().items():
             fn_recursive_feed_forward(module, chunk_size, dim)
 
     def disable_forward_chunking(self):
         def fn_recursive_feed_forward(module: nn.Cell, chunk_size: int, dim: int):
             if hasattr(module, "set_chunk_feed_forward"):
                 module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
 
-            for child in module.name_cells():
+            for child in module.name_cells().items():
                 fn_recursive_feed_forward(child, chunk_size, dim)
 
-        for module in self.name_cells():
+        for module in self.name_cells().items():
             fn_recursive_feed_forward(module, None, 0)
 
     # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor

diff --git a/mindone/diffusers/models/unets/unet_i2vgen_xl.py b/mindone/diffusers/models/unets/unet_i2vgen_xl.py
@@ -336,12 +336,12 @@ def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str
             if hasattr(module, "get_processor"):
                 processors[f"{name}.processor"] = module.get_processor()
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
 
             return processors
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_add_processors(name, module, processors)
 
         return processors
@@ -372,10 +372,10 @@ def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
                 else:
                     module.set_processor(processor.pop(f"{name}.processor"))
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_attn_processor(name, module, processor)
 
     # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
@@ -402,10 +402,10 @@ def fn_recursive_feed_forward(module: nn.Cell, chunk_size: int, dim: int):
             if hasattr(module, "set_chunk_feed_forward"):
                 module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
 
-            for child in module.name_cells():
+            for child in module.name_cells().items():
                 fn_recursive_feed_forward(child, chunk_size, dim)
 
-        for module in self.name_cells():
+        for module in self.name_cells().items():
             fn_recursive_feed_forward(module, chunk_size, dim)
 
     # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.disable_forward_chunking
@@ -414,10 +414,10 @@ def fn_recursive_feed_forward(module: nn.Cell, chunk_size: int, dim: int):
             if hasattr(module, "set_chunk_feed_forward"):
                 module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
 
-            for child in module.name_cells():
+            for child in module.name_cells().items():
                 fn_recursive_feed_forward(child, chunk_size, dim)
 
-        for module in self.name_cells():
+        for module in self.name_cells().items():
             fn_recursive_feed_forward(module, None, 0)
 
     # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor

diff --git a/mindone/diffusers/models/unets/unet_kandinsky3.py b/mindone/diffusers/models/unets/unet_kandinsky3.py
@@ -157,12 +157,12 @@ def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str
             if hasattr(module, "get_processor"):
                 processors[f"{name}.processor"] = module.get_processor()
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
 
             return processors
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_add_processors(name, module, processors)
 
         return processors
@@ -192,10 +192,10 @@ def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
                 else:
                     module.set_processor(processor.pop(f"{name}.processor"))
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_attn_processor(name, module, processor)
 
     def set_default_attn_processor(self):

diff --git a/mindone/diffusers/models/unets/unet_motion_model.py b/mindone/diffusers/models/unets/unet_motion_model.py
@@ -573,12 +573,12 @@ def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str
             if hasattr(module, "get_processor"):
                 processors[f"{name}.processor"] = module.get_processor()
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
 
             return processors
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_add_processors(name, module, processors)
 
         return processors
@@ -609,10 +609,10 @@ def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
                 else:
                     module.set_processor(processor.pop(f"{name}.processor"))
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_attn_processor(name, module, processor)
 
     # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
@@ -639,10 +639,10 @@ def fn_recursive_feed_forward(module: nn.Cell, chunk_size: int, dim: int):
             if hasattr(module, "set_chunk_feed_forward"):
                 module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
 
-            for child in module.name_cells():
+            for child in module.name_cells().items():
                 fn_recursive_feed_forward(child, chunk_size, dim)
 
-        for module in self.name_cells():
+        for module in self.name_cells().items():
             fn_recursive_feed_forward(module, chunk_size, dim)
 
     # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.disable_forward_chunking
@@ -651,10 +651,10 @@ def fn_recursive_feed_forward(module: nn.Cell, chunk_size: int, dim: int):
             if hasattr(module, "set_chunk_feed_forward"):
                 module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
 
-            for child in module.name_cells():
+            for child in module.name_cells().items():
                 fn_recursive_feed_forward(child, chunk_size, dim)
 
-        for module in self.name_cells():
+        for module in self.name_cells().items():
             fn_recursive_feed_forward(module, None, 0)
 
     # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor

diff --git a/mindone/diffusers/models/unets/unet_spatio_temporal_condition.py b/mindone/diffusers/models/unets/unet_spatio_temporal_condition.py
@@ -273,12 +273,12 @@ def fn_recursive_add_processors(name: str, module: nn.Cell, processors: Dict[str
             if hasattr(module, "get_processor"):
                 processors[f"{name}.processor"] = module.get_processor()
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
 
             return processors
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_add_processors(name, module, processors)
 
         return processors
@@ -308,10 +308,10 @@ def fn_recursive_attn_processor(name: str, module: nn.Cell, processor):
                 else:
                     module.set_processor(processor.pop(f"{name}.processor"))
 
-            for sub_name, child in module.name_cells():
+            for sub_name, child in module.name_cells().items():
                 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
 
-        for name, module in self.name_cells():
+        for name, module in self.name_cells().items():
             fn_recursive_attn_processor(name, module, processor)
 
     def set_default_attn_processor(self):
@@ -355,10 +355,10 @@ def fn_recursive_feed_forward(module: nn.Cell, chunk_size: int, dim: int):
             if hasattr(module, "set_chunk_feed_forward"):
                 module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
 
-            for child in module.name_cells():
+            for child in module.name_cells().items():
                 fn_recursive_feed_forward(child, chunk_size, dim)
 
-        for module in self.name_cells():
+        for module in self.name_cells().items():
             fn_recursive_feed_forward(module, chunk_size, dim)
 
     def construct(