path fix.

Author: sayakpaul

base.png

@@ -0,0 +1,8 @@
+Loading pipeline components...:   0%|          | 0/7 [00:00<?, ?it/s]You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
+Loading pipeline components...:  29%|██▊       | 2/7 [00:00<00:00,  9.52it/s]
+Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]Loading checkpoint shards: 100%|██████████| 3/3 [00:00<00:00, 51.82it/s]
+Loading pipeline components...:  71%|███████▏  | 5/7 [00:00<00:00, 13.99it/s]
+Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]Loading checkpoint shards: 100%|██████████| 2/2 [00:00<00:00, 65.67it/s]
+Loading pipeline components...: 100%|██████████| 7/7 [00:00<00:00, 15.05it/s]
+time mean/var: tensor([4.3879, 4.3965, 4.3994, 4.3941, 4.3931, 4.3900, 4.3813, 4.3718, 4.3671,
+        4.3671]) 4.384819507598877 0.00015030684880912304

compile_last.png

@@ -0,0 +1,8 @@
+Loading pipeline components...:   0%|          | 0/7 [00:00<?, ?it/s]You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
+Loading pipeline components...:  43%|████▎     | 3/7 [00:00<00:00, 11.90it/s]
+Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]Loading checkpoint shards: 100%|██████████| 3/3 [00:00<00:00, 50.62it/s]
+Loading pipeline components...:  71%|███████▏  | 5/7 [00:00<00:00, 12.37it/s]
+Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]Loading checkpoint shards: 100%|██████████| 2/2 [00:00<00:00, 63.95it/s]
+Loading pipeline components...: 100%|██████████| 7/7 [00:00<00:00, 14.82it/s]
+time mean/var: tensor([6.6733, 6.6850, 6.6962, 6.7148, 6.7146, 6.7150, 6.7231, 6.7184, 6.7057,
+        6.7117]) 6.705780982971191 0.00025744145386852324

compile_last.txt

@@ -12,10 +12,10 @@ def load_pipeline(args):
     is_timestep_distilled = not pipeline.transformer.config.guidance_embeds
 
     transformer_package_path = os.path.join(
-        args.cache_dir, "exported_transformer.pt2" if is_timestep_distilled else "exported_dev_transformer.pt"
+        args.cache_dir, "exported_transformer.pt2" if is_timestep_distilled else "exported_dev_transformer.pt2"
     )
     decoder_package_path = os.path.join(
-        args.cache_dir, "exported_decoder.pt2" if is_timestep_distilled else "exported_dev_decoder.pt"
+        args.cache_dir, "exported_decoder.pt2" if is_timestep_distilled else "exported_dev_decoder.pt2"
     )
     loaded_transformer = load_package(transformer_package_path)
     loaded_decoder = load_package(decoder_package_path)

just_compile.png

@@ -270,7 +270,9 @@ def _example_tensor(*shape):
     }
 
     # Possibly serialize model out
-    transformer_package_path = os.path.join(cache_dir, "exported_transformer.pt2")
+    transformer_package_path = os.path.join(
+        cache_dir, "exported_transformer.pt2" if is_timestep_distilled else "exported_dev_transformer.pt2"
+    )
     if serialize:
         # Apply export
         exported_transformer: torch.export.ExportedProgram = torch.export.export(
@@ -305,7 +307,9 @@ def _example_tensor(*shape):
     vae_decode_kwargs = {"return_dict": False}
 
     # Possibly serialize model out
-    decoder_package_path = os.path.join(cache_dir, "exported_decoder.pt2")
+    decoder_package_path = os.path.join(
+        cache_dir, "exported_decoder.pt2" if is_timestep_distilled else "exported_dev_decoder.pt2"
+    )
     if serialize:
         # Apply export
         exported_decoder: torch.export.ExportedProgram = torch.export.export(