From 1d9a6a81b9a00b2dfab992ba75eb55bf8afb4eae Mon Sep 17 00:00:00 2001 From: Jianqi Pan Date: Mon, 10 Jun 2024 20:11:41 +0900 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20chore:=20use=20modeling=5Foutput?= =?UTF-8?q?s.Transformer2DModelOutput=20(#8436)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * :wrench: chore: use modeling_outputs.Transformer2DModelOutput * :wrench: chore: isort * :wrench: chore: isort * style --------- Co-authored-by: sayakpaul --- src/diffusers/models/transformers/dual_transformer_2d.py | 3 ++- src/diffusers/pipelines/unidiffuser/modeling_uvit.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/transformers/dual_transformer_2d.py b/src/diffusers/models/transformers/dual_transformer_2d.py index bbf2d387e4..1c48c4e3db 100644 --- a/src/diffusers/models/transformers/dual_transformer_2d.py +++ b/src/diffusers/models/transformers/dual_transformer_2d.py @@ -15,7 +15,8 @@ from typing import Optional from torch import nn -from .transformer_2d import Transformer2DModel, Transformer2DModelOutput +from ..modeling_outputs import Transformer2DModelOutput +from .transformer_2d import Transformer2DModel class DualTransformer2DModel(nn.Module): diff --git a/src/diffusers/pipelines/unidiffuser/modeling_uvit.py b/src/diffusers/pipelines/unidiffuser/modeling_uvit.py index abc51edf6d..cb1514b153 100644 --- a/src/diffusers/pipelines/unidiffuser/modeling_uvit.py +++ b/src/diffusers/pipelines/unidiffuser/modeling_uvit.py @@ -9,8 +9,8 @@ from ...models import ModelMixin from ...models.attention import FeedForward from ...models.attention_processor import Attention from ...models.embeddings import TimestepEmbedding, Timesteps, get_2d_sincos_pos_embed +from ...models.modeling_outputs import Transformer2DModelOutput from ...models.normalization import AdaLayerNorm -from ...models.transformers.transformer_2d import Transformer2DModelOutput from ...utils import logging