Update src/pipeline.py

src/pipeline.py

@@ -595,7 +595,7 @@ def load_pipeline() -> Pipeline:
     ).to(memory_format=torch.channels_last)
     
     vae = AutoencoderTiny.from_pretrained("RobertML/FLUX.1-schnell-vae_e3m2", revision="da0d2cd7815792fb40d084dbd8ed32b63f153d8d", torch_dtype=dtype)
-    pipeline.vae = torch.compile(pipeline.vae)
+    
     path = os.path.join(HF_HUB_CACHE, "models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a")
     generator = torch.Generator(device=device)
     model = FluxTransformer2DModel.from_pretrained(path, torch_dtype=dtype, use_safetensors=False, generator= generator).to(memory_format=torch.channels_last)
@@ -608,7 +608,8 @@ def load_pipeline() -> Pipeline:
         transformer=model,
         text_encoder_2=text_encoder_2,
         torch_dtype=dtype,
-        ).to(device)
+        ).to(device
+    pipeline.vae = torch.compile(pipeline.vae)
     for _ in range(3):
         pipeline(prompt="blah blah waah waah oneshot oneshot gang gang", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)