Update src/pipeline.py

src/pipeline.py

@@ -726,16 +726,16 @@ class FluxPipeline(
                     self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds
                 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                 timestep = t.expand(latents.shape[0]).to(latents.dtype)
-                print("=============== printing all the shapes right now ======================")
-                print(latents.shape)
-                print(timestep)
-                print(guidance)
-                print(pooled_prompt_embeds.shape)
-                print(prompt_embeds.shape)
-                print(text_ids.shape)
-                print(latent_image_ids.shape)
-                print("=================== thats all folks for now ============================")
-                exit()
+                # print("=============== printing all the shapes right now ======================")
+                # print(latents.shape)
+                # print(timestep)
+                # print(guidance)
+                # print(pooled_prompt_embeds.shape)
+                # print(prompt_embeds.shape)
+                # print(text_ids.shape)
+                # print(latent_image_ids.shape)
+                # print("=================== thats all folks for now ============================")
+                # exit()
                 noise_pred = self.transformer(
                     hidden_states=latents,
                     timestep=timestep / 1000,
@@ -858,45 +858,45 @@ def load_pipeline() -> Pipeline:
     model_name = "manbeast3b/Flux.1.Schnell-full-quant1"
     revision = "e7ddf488a4ea8a3cba05db5b8d06e7e0feb826a2"
     
-    hub_model_dir = os.path.join(
-        HF_HUB_CACHE,
-        f"models--{model_name.replace('/', '--')}",
-        "snapshots",
-        revision,
-        "transformer"
-    )
-    transformer = FluxTransformer2DModel.from_pretrained(
-        hub_model_dir,
-        torch_dtype=torch.bfloat16,
-        use_safetensors=False
-    ).to(memory_format=torch.channels_last)
+    # hub_model_dir = os.path.join(
+    #     HF_HUB_CACHE,
+    #     f"models--{model_name.replace('/', '--')}",
+    #     "snapshots",
+    #     revision,
+    #     "transformer"
+    # )
+    # transformer = FluxTransformer2DModel.from_pretrained(
+    #     hub_model_dir,
+    #     torch_dtype=torch.bfloat16,
+    #     use_safetensors=False
+    # ).to(memory_format=torch.channels_last)
     
     pipeline = FluxPipeline.from_pretrained(
         ckpt_id,
         revision=ckpt_revision,
         # text_encoder_2=text_encoder_2,
-        transformer=transformer,
+        transformer=None, #transformer,
         # vae=vae,
         torch_dtype=torch.bfloat16
     )
     # pipeline.vae = torch.compile(vae)
     pipeline.to("cuda")
     
-    # path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.la_schnella_transformer/snapshots/8bfd89bd9e2099e70a1155403bf8aabb0a3177df/flux_la_schnell_aten.so.pt2")
-    # inputs1 = get_example_inputs()
-    # print(f"AoT pre compiled path is {path}")
+    path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.la_schnella_transformer/snapshots/8e704256516ef1cbd6730fedd019f5b1a71d38d3/flux_la_schnell_ez.so.pt2")
+    inputs1 = get_example_inputs()
+    print(f"AoT pre compiled path is {path}")
     
-    # # transformer = torch._inductor.aoti_load_package(path)
     # transformer = torch._inductor.aoti_load_package(path)
-    # print(f"{transformer(**inputs1)[0].shape=}")
+    transformer = torch._inductor.aoti_load_package(path)
+    print(f"{transformer(**inputs1)[0].shape=}")
 
-    # for _ in range(3):
-    #     _ = transformer(**inputs1)[0]
+    for _ in range(3):
+        _ = transformer(**inputs1)[0]
 
-    # time = benchmark_fn(f, transformer, **inputs1)
-    # print(f"{time=} seconds.")
+    time = benchmark_fn(f, transformer, **inputs1)
+    print(f"{time=} seconds.")
 
-    # pipeline.transformer = transformer
+    pipeline.transformer = transformer
     
     warmup_ = "controllable varied focus thai warriors entertainment blue golden pink soft tough padthai"
     for _ in range(1):