Port VQ-diffusion text embeddings (CLIP)

Port the text embeddings (CLIP) for the ITHQ dataset

`convert_vq_diffusion_to_diffusers.py` script now uses transformers
to pull CLIP and save it along with the rest of the model. Note that
in VQ-diffusion, the output text embeddings are additionally normalized.
The additional normalization will be added to the pipeline as a part of
inference in a later commit.

Author: williamberman

scripts/convert_vq_diffusion_to_diffusers.py

@@ -35,6 +35,7 @@
 from diffusers import VQModel
 from diffusers.models.vq_diffusion_attention import VQDiffusionTransformer
 from diffusers.pipelines import VQDiffusionPipeline
+from transformers import CLIPTextModel, CLIPTokenizer
 from yaml.loader import FullLoader
 
 
@@ -819,9 +820,37 @@ def read_config_file(filename):
 
     # done transformer_model
 
+    # text encoder
+
+    print("loading CLIP text encoder")
+
+    clip_name = "openai/clip-vit-base-patch32"
+
+    # The original VQ-Diffusion specifies the pad value by the int used in the
+    # returned tokens. Each model uses `0` as the pad value. The transformers clip api
+    # specifies the pad value via the token before it has been tokenized. The `!` pad
+    # token is the same as padding with the `0` pad value.
+    pad_token = "!"
+
+    tokenizer_model = CLIPTokenizer.from_pretrained(clip_name, pad_token=pad_token, device_map="auto")
+
+    assert tokenizer_model.convert_tokens_to_ids(pad_token) == 0
+
+    text_encoder_model = CLIPTextModel.from_pretrained(
+        clip_name,
+        # `CLIPTextModel` does not support device_map="auto"
+        # device_map="auto"
+    )
+
+    print("done loading CLIP text encoder")
+
+    # done text encoder
+
     print(f"saving VQ diffusion model, path: {args.dump_path}")
 
-    pipe = VQDiffusionPipeline(vqvae=vqvae_model, transformer=transformer_model)
+    pipe = VQDiffusionPipeline(
+        vqvae=vqvae_model, transformer=transformer_model, tokenizer=tokenizer_model, text_encoder=text_encoder_model
+    )
     pipe.save_pretrained(args.dump_path)
 
     print("done writing VQ diffusion model")

src/diffusers/pipelines/vq_diffusion/pipeline_vq_diffusion.py

@@ -1,4 +1,5 @@
 from diffusers import VQDiffusionTransformer, VQModel
+from transformers import CLIPTextModel, CLIPTokenizer
 
 from ...pipeline_utils import DiffusionPipeline
 
@@ -14,6 +15,17 @@ class VQDiffusionPipeline(DiffusionPipeline):
     vqvae: VQModel
     transformer: VQDiffusionTransformer
 
-    def __init__(self, vqvae: VQModel, transformer: VQDiffusionTransformer):
+    def __init__(
+        self,
+        vqvae: VQModel,
+        transformer: VQDiffusionTransformer,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+    ):
         super().__init__()
-        self.register_modules(vqvae=vqvae, transformer=transformer)
+        self.register_modules(
+            vqvae=vqvae,
+            transformer=transformer,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+        )