Merge pull request #18 from LightricksResearch/delete-pixart

xora/examples/image_to_video.py

@@ -3,7 +3,7 @@ from xora.models.autoencoders.causal_video_autoencoder import CausalVideoAutoenc
 from xora.models.transformers.transformer3d import Transformer3DModel
 from xora.models.transformers.symmetric_patchifier import SymmetricPatchifier
 from xora.schedulers.rf import RectifiedFlowScheduler
-from xora.pipelines.pipeline_video_pixart_alpha import VideoPixArtAlphaPipeline
+from xora.pipelines.pipeline_xora_video import XoraVideoPipeline
 from pathlib import Path
 from transformers import T5EncoderModel, T5Tokenizer
 import safetensors.torch
@@ -180,7 +180,7 @@ def main():
         "vae": vae,
     }
 
-    pipeline = VideoPixArtAlphaPipeline(**submodel_dict).to("cuda")
+    pipeline = XoraVideoPipeline(**submodel_dict).to("cuda")
 
     # Load media (video or image)
     if args.video_path:

xora/examples/text_to_video.py

@@ -3,7 +3,7 @@ from xora.models.autoencoders.causal_video_autoencoder import CausalVideoAutoenc
 from xora.models.transformers.transformer3d import Transformer3DModel
 from xora.models.transformers.symmetric_patchifier import SymmetricPatchifier
 from xora.schedulers.rf import RectifiedFlowScheduler
-from xora.pipelines.pipeline_video_pixart_alpha import VideoPixArtAlphaPipeline
+from xora.pipelines.pipeline_xora_video import XoraVideoPipeline
 from pathlib import Path
 from transformers import T5EncoderModel, T5Tokenizer
 import safetensors.torch
@@ -82,7 +82,7 @@ def main():
         "vae": vae,
     }
 
-    pipeline = VideoPixArtAlphaPipeline(**submodel_dict).to("cuda")
+    pipeline = XoraVideoPipeline(**submodel_dict).to("cuda")
 
     # Sample input
     num_inference_steps = 20

xora/models/transformers/symmetric_patchifier.py

@@ -60,26 +60,19 @@ class Patchifier(ConfigMixin, ABC):
         return grid
 
 
-def pixart_alpha_patchify(
-    latents: Tensor,
-    patch_size: int,
-) -> Tuple[Tensor, Tensor]:
-    latents = rearrange(
-        latents,
-        "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)",
-        p1=patch_size[0],
-        p2=patch_size[1],
-        p3=patch_size[2],
-    )
-    return latents
-
-
 class SymmetricPatchifier(Patchifier):
     def patchify(
         self,
         latents: Tensor,
     ) -> Tuple[Tensor, Tensor]:
-        return pixart_alpha_patchify(latents, self._patch_size)
+        latents = rearrange(
+            latents,
+            "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)",
+            p1=self._patch_size[0],
+            p2=self._patch_size[1],
+            p3=self._patch_size[2],
+        )
+        return latents
 
     def unpatchify(
         self,

xora/models/transformers/transformer3d.py

@@ -141,12 +141,10 @@ class Transformer3DModel(ModelMixin, ConfigMixin):
         )
         self.proj_out = nn.Linear(inner_dim, self.out_channels)
 
-        # 5. PixArt-Alpha blocks.
         self.adaln_single = AdaLayerNormSingle(
             inner_dim, use_additional_conditions=False
         )
         if adaptive_norm == "single_scale":
-            # Use 4 channels instead of the 6 for the PixArt-Alpha scale + shift ada norm.
             self.adaln_single.linear = nn.Linear(inner_dim, 4 * inner_dim, bias=True)
 
         self.caption_projection = None
@@ -170,7 +168,7 @@ class Transformer3DModel(ModelMixin, ConfigMixin):
             for block in self.transformer_blocks:
                 block.set_use_tpu_flash_attention(self.device.type)
 
-    def initialize(self, embedding_std: float, mode: Literal["xora", "pixart"]):
+    def initialize(self, embedding_std: float, mode: Literal["xora", "legacy"]):
         def _basic_init(module):
             if isinstance(module, nn.Linear):
                 torch.nn.init.xavier_uniform_(module.weight)
@@ -211,7 +209,6 @@ class Transformer3DModel(ModelMixin, ConfigMixin):
         nn.init.normal_(self.caption_projection.linear_1.weight, std=embedding_std)
         nn.init.normal_(self.caption_projection.linear_1.weight, std=embedding_std)
 
-        # Zero-out adaLN modulation layers in PixArt blocks:
         for block in self.transformer_blocks:
             if mode.lower() == "xora":
                 nn.init.constant_(block.attn1.to_out[0].weight, 0)

xora/pipelines/{pipeline_video_pixart_alpha.py → pipeline_xora_video.py}

@@ -1,4 +1,4 @@
-# # Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py
+# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py
 import html
 import inspect
 import math
@@ -19,7 +19,6 @@ from diffusers.utils import (
     is_bs4_available,
     is_ftfy_available,
     logging,
-    replace_example_docstring,
 )
 from diffusers.utils.torch_utils import randn_tensor
 from einops import rearrange
@@ -44,22 +43,6 @@ if is_bs4_available():
 if is_ftfy_available():
     import ftfy
 
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> import torch
-        >>> from diffusers import PixArtAlphaPipeline
-
-        >>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-XL-2-512x512" too.
-        >>> pipe = PixArtAlphaPipeline.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", torch_dtype=torch.float16)
-        >>> # Enable memory optimizations.
-        >>> pipe.enable_model_cpu_offload()
-
-        >>> prompt = "A small cactus with a happy face in the Sahara desert."
-        >>> image = pipe(prompt).images[0]
-        ```
-"""
-
 ASPECT_RATIO_1024_BIN = {
     "0.25": [512.0, 2048.0],
     "0.28": [512.0, 1856.0],
@@ -180,9 +163,9 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-class VideoPixArtAlphaPipeline(DiffusionPipeline):
+class XoraVideoPipeline(DiffusionPipeline):
     r"""
-    Pipeline for text-to-image generation using PixArt-Alpha.
+    Pipeline for text-to-image generation using Xora.
 
     This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
     library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
@@ -191,7 +174,7 @@ class VideoPixArtAlphaPipeline(DiffusionPipeline):
         vae ([`AutoencoderKL`]):
             Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
         text_encoder ([`T5EncoderModel`]):
-            Frozen text-encoder. PixArt-Alpha uses
+            Frozen text-encoder. This uses
             [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the
             [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant.
         tokenizer (`T5Tokenizer`):
@@ -247,7 +230,6 @@ class VideoPixArtAlphaPipeline(DiffusionPipeline):
         )
         self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
 
-    # Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/utils.py
     def mask_text_embeddings(self, emb, mask):
         if emb.shape[0] == 1:
             keep_index = mask.sum().item()
@@ -280,7 +262,7 @@ class VideoPixArtAlphaPipeline(DiffusionPipeline):
             negative_prompt (`str` or `List[str]`, *optional*):
                 The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
                 instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
-                PixArt-Alpha, this should be "".
+                This should be "".
             do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
                 whether to use classifier free guidance or not
             num_images_per_prompt (`int`, *optional*, defaults to 1):
@@ -291,8 +273,7 @@ class VideoPixArtAlphaPipeline(DiffusionPipeline):
                 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                 provided, text embeddings will be generated from `prompt` input argument.
             negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. For PixArt-Alpha, it's should be the embeddings of the ""
-                string.
+                Pre-generated negative text embeddings.
             clean_caption (bool, defaults to `False`):
                 If `True`, the function will preprocess and clean the provided caption before encoding.
         """
@@ -753,7 +734,6 @@ class VideoPixArtAlphaPipeline(DiffusionPipeline):
         return samples
 
     @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
     def __call__(
         self,
         height: int,
@@ -824,7 +804,7 @@ class VideoPixArtAlphaPipeline(DiffusionPipeline):
                 provided, text embeddings will be generated from `prompt` input argument.
             prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.
             negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. For PixArt-Alpha this negative prompt should be "". If not
+                Pre-generated negative text embeddings. This negative prompt should be "". If not
                 provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
             negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
                 Pre-generated attention mask for negative text embeddings.