fix(api): set VAE attn processor during conversion

ssube · Sep 24, 2023 · 0ecae65 · 0ecae65
1 parent 6e2896f
commit 0ecae65
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/api/onnx_web/convert/diffusion/diffusion.py b/api/onnx_web/convert/diffusion/diffusion.py
@@ -380,6 +380,10 @@ def convert_diffusion_diffusers(
         else:
             pipeline.vae = AutoencoderKL.from_pretrained(vae_path)
 
+    if is_torch_2_0:
+        pipeline.unet.set_attn_processor(AttnProcessor())
+        pipeline.vae.set_attn_processor(AttnProcessor())
+
     optimize_pipeline(conversion, pipeline)
 
     output_path = Path(dest_path)
@@ -430,9 +434,6 @@ def convert_diffusion_diffusers(
         unet_inputs = ["sample", "timestep", "encoder_hidden_states", "return_dict"]
         unet_scale = torch.tensor(False).to(device=device, dtype=torch.bool)
 
-    if is_torch_2_0:
-        pipeline.unet.set_attn_processor(AttnProcessor())
-
     unet_in_channels = pipeline.unet.config.in_channels
     unet_sample_size = pipeline.unet.config.sample_size
     unet_path = output_path / "unet" / ONNX_MODEL