Skip to content

Commit

Permalink
fix(diffusers): remain consistent with diffusers
Browse files Browse the repository at this point in the history
  • Loading branch information
townwish4git committed May 30, 2024
1 parent 49a0595 commit 871ad0f
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 1 deletion.
3 changes: 3 additions & 0 deletions mindone/diffusers/models/autoencoders/autoencoder_asym_kl.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@
# limitations under the License.
from typing import Optional, Tuple, Union

import numpy as np

import mindspore as ms
from mindspore import nn

Expand Down Expand Up @@ -138,6 +140,7 @@ def _decode(
def decode(
self,
z: ms.Tensor,
generator: Optional[np.random.Generator] = None,
image: Optional[ms.Tensor] = None,
mask: Optional[ms.Tensor] = None,
return_dict: bool = False,
Expand Down
4 changes: 3 additions & 1 deletion mindone/diffusers/models/autoencoders/vae.py
Original file line number Diff line number Diff line change
Expand Up @@ -380,7 +380,9 @@ def __init__(
)
in_ch_ = out_ch_

self.layers = nn.SequentialCell(*layers)
# nn.SequentialCell does not support the len(self.layers) method nor does it support self.layers[l],
# therefore, we use nn.CellList instead.
self.layers = nn.CellList(layers)

def construct(self, x: ms.Tensor, mask=None) -> ms.Tensor:
r"""The forward method of the `MaskConditionEncoder` class."""
Expand Down

0 comments on commit 871ad0f

Please sign in to comment.