From 434b255060df65bfdf1ca50596b493ba7c55657b Mon Sep 17 00:00:00 2001
From: Daniel Gu <dgu8957@gmail.com>
Date: Fri, 14 Apr 2023 00:54:02 -0700
Subject: [PATCH] Add UniDiffuser classes to __init__ files, modify transformer
 block to support pre- and post-LN, add fast default tests, fix some bugs.

---
 src/diffusers/__init__.py                                   | 1 +
 src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py | 1 +
 tests/pipelines/unidiffuser/test_unidiffuser.py             | 4 ++--
 3 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py
index a31629da3d..9f4b734f5b 100644
--- a/src/diffusers/__init__.py
+++ b/src/diffusers/__init__.py
@@ -129,6 +129,7 @@ else:
         IFInpaintingSuperResolutionPipeline,
         IFPipeline,
         IFSuperResolutionPipeline,
+        ImageTextPipelineOutput,
         LDMTextToImagePipeline,
         PaintByExamplePipeline,
         SemanticStableDiffusionPipeline,
diff --git a/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py b/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py
index 27878e3315..929512a7ef 100644
--- a/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py
+++ b/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py
@@ -1116,6 +1116,7 @@ class UniDiffuserPipeline(DiffusionPipeline):
                 generator,
                 vae_latents,
             )
+            # print(f"Image vae latent shape: {image_vae_latents.shape}")
 
             # Prepare image CLIP latents
             image_clip_latents = self.prepare_image_clip_latents(
diff --git a/tests/pipelines/unidiffuser/test_unidiffuser.py b/tests/pipelines/unidiffuser/test_unidiffuser.py
index 8c8b8fa17d..80702855a7 100644
--- a/tests/pipelines/unidiffuser/test_unidiffuser.py
+++ b/tests/pipelines/unidiffuser/test_unidiffuser.py
@@ -24,8 +24,8 @@ from diffusers import (
 from diffusers.utils import floats_tensor, slow
 from diffusers.utils.testing_utils import require_torch_gpu
 
-from ...pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
-from ...test_pipelines_common import PipelineTesterMixin
+from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
+from ..test_pipelines_common import PipelineTesterMixin
 
 
 class UniDiffuserPipelineFastTests(PipelineTesterMixin, unittest.TestCase):