change the ae related code due to the latest update of DCAE branch;

Author: lawrence-cj

scripts/convert_sana_pag_to_diffusers.py

@@ -11,7 +11,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from diffusers import (
-    DCAE,
+    AutoencoderDC,
     DPMSolverMultistepScheduler,
     FlowMatchEulerDiscreteScheduler,
     SanaPAGPipeline,
@@ -187,7 +187,7 @@ def main(args):
     else:
         print(colored(f"Saving the whole SanaPAGPipeline containing {args.model_type}", "green", attrs=["bold"]))
         # VAE
-        dc_ae = DCAE.from_pretrained(
+        ae = AutoencoderDC.from_pretrained(
             "Efficient-Large-Model/dc_ae_f32c32_sana_1.0_diffusers",
             torch_dtype=torch.float32,
         ).to(device)
@@ -221,7 +221,7 @@ def main(args):
             tokenizer=tokenizer,
             text_encoder=text_encoder,
             transformer=transformer,
-            vae=dc_ae,
+            vae=ae,
             scheduler=scheduler,
             pag_applied_layers="blocks.8",
         )

scripts/convert_sana_to_diffusers.py

@@ -11,7 +11,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from diffusers import (
-    DCAE,
+    AutoencoderDC,
     DPMSolverMultistepScheduler,
     FlowMatchEulerDiscreteScheduler,
     SanaPipeline,
@@ -187,7 +187,7 @@ def main(args):
     else:
         print(colored(f"Saving the whole SanaPipeline containing {args.model_type}", "green", attrs=["bold"]))
         # VAE
-        dc_ae = DCAE.from_pretrained(
+        ae = AutoencoderDC.from_pretrained(
             "Efficient-Large-Model/dc_ae_f32c32_sana_1.0_diffusers",
             torch_dtype=torch.float32,
         ).to(device)
@@ -221,7 +221,7 @@ def main(args):
             tokenizer=tokenizer,
             text_encoder=text_encoder,
             transformer=transformer,
-            vae=dc_ae,
+            vae=ae,
             scheduler=scheduler,
         )
 

src/diffusers/__init__.py

@@ -130,7 +130,6 @@
             "UNetSpatioTemporalConditionModel",
             "UVit2DModel",
             "VQModel",
-            "DCAE",
         ]
     )
     _import_structure["optimization"] = [
@@ -576,7 +575,6 @@
         from .utils.dummy_pt_objects import *  # noqa F403
     else:
         from .models import (
-            DCAE,
             AllegroTransformer3DModel,
             AsymmetricAutoencoderKL,
             AuraFlowTransformer2DModel,

src/diffusers/models/transformers/sana_transformer_2d.py

@@ -26,7 +26,7 @@
     FusedAttnProcessor2_0,
     SanaLinearAttnProcessor2_0,
 )
-from ..autoencoders.dc_ae import GLUMBConv
+from ..autoencoders.autoencoder_dc import GLUMBConv
 from ..embeddings import PatchEmbed, PixArtAlphaTextProjection, SinusoidalPositionalEmbedding
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin

src/diffusers/pipelines/pag/pipeline_pag_sana.py

@@ -22,7 +22,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from ...image_processor import PixArtImageProcessor
-from ...models import DCAE, SanaTransformer2DModel
+from ...models import AutoencoderDC, SanaTransformer2DModel
 from ...models.attention_processor import PAGCFGSanaLinearAttnProcessor2_0, PAGIdentitySanaLinearAttnProcessor2_0
 from ...schedulers import FlowDPMSolverMultistepScheduler
 from ...utils import (
@@ -161,7 +161,7 @@ def __init__(
         self,
         tokenizer: AutoTokenizer,
         text_encoder: AutoModelForCausalLM,
-        vae: DCAE,
+        vae: AutoencoderDC,
         transformer: SanaTransformer2DModel,
         scheduler: FlowDPMSolverMultistepScheduler,
         pag_applied_layers: Union[str, List[str]] = "blocks.1",  # 1st transformer block

src/diffusers/pipelines/sana/pipeline_sana.py

@@ -22,7 +22,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from ...image_processor import PixArtImageProcessor
-from ...models import DCAE, SanaTransformer2DModel
+from ...models import AutoencoderDC, SanaTransformer2DModel
 from ...schedulers import FlowDPMSolverMultistepScheduler
 from ...utils import (
     BACKENDS_MAPPING,
@@ -157,7 +157,7 @@ def __init__(
         self,
         tokenizer: AutoTokenizer,
         text_encoder: AutoModelForCausalLM,
-        vae: DCAE,
+        vae: AutoencoderDC,
         transformer: SanaTransformer2DModel,
         scheduler: FlowDPMSolverMultistepScheduler,
     ):