Update stable_diffusion_pipeline_compiler.py

Do not trace vae for pytorch 2
chengzeyi · Nov 6, 2023 · f255ed8 · f255ed8
1 parent 3058d09
commit f255ed8
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/sfast/compilers/stable_diffusion_pipeline_compiler.py b/sfast/compilers/stable_diffusion_pipeline_compiler.py
@@ -131,14 +131,14 @@ def unet_forward_wrapper(sample, t, *args, **kwargs):
 
             m.unet.forward = unet_forward_wrapper
 
-            # if packaging.version.parse(
-            #         torch.__version__) < packaging.version.parse('2.0.0'):
-            m.vae.decode = lazy_trace_(to_module(m.vae.decode))
-            # For img2img
-            m.vae.encoder.forward = lazy_trace_(
-                to_module(m.vae.encoder.forward))
-            m.vae.quant_conv.forward = lazy_trace_(
-                to_module(m.vae.quant_conv.forward))
+            if packaging.version.parse(
+                    torch.__version__) < packaging.version.parse('2.0.0'):
+                m.vae.decode = lazy_trace_(to_module(m.vae.decode))
+                # For img2img
+                m.vae.encoder.forward = lazy_trace_(
+                    to_module(m.vae.encoder.forward))
+                m.vae.quant_conv.forward = lazy_trace_(
+                    to_module(m.vae.quant_conv.forward))
 
             if config.trace_scheduler:
                 m.scheduler.scale_model_input = lazy_trace_(