Update src/pipeline.py

src/pipeline.py

@@ -70,18 +70,19 @@ def decode_latents_to_image(latents, height: int, width: int, vae):
         vae_scale_factor = 1
     image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
     
-     # Try to load the traced model; trace and save if not found
-    if os.path.exists(traced_vae_decode_path):
-        try:
-            traced_vae_decode = torch.jit.load(traced_vae_decode_path)
-            # print("Loaded traced VAE decoder from file.")
-        except Exception as e:
-            # print(f"Error loading traced VAE decoder: {e}. Retracing...")
-            traced_vae_decode = trace_and_save_vae_decoder(vae, latents)
+    #  # Try to load the traced model; trace and save if not found
+    # if os.path.exists(traced_vae_decode_path):
+    #     try:
+    #         traced_vae_decode = torch.jit.load(traced_vae_decode_path)
+    #         # print("Loaded traced VAE decoder from file.")
+    #     except Exception as e:
+    #         # print(f"Error loading traced VAE decoder: {e}. Retracing...")
+    #         traced_vae_decode = trace_and_save_vae_decoder(vae, latents)
 
-    else:
-        traced_vae_decode = trace_and_save_vae_decoder(vae, latents)
+    # else:
+    #     traced_vae_decode = trace_and_save_vae_decoder(vae, latents)
 
+    traced_vae_decode = vae.decode
     with torch.no_grad():
         latents = FluxPipeline._unpack_latents(latents.unsqueeze(0), height, width, vae_scale_factor)
         latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor