Update src/pipeline.py

src/pipeline.py

@@ -27,7 +27,7 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
     
     pipeline.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipeline.scheduler.config))
 
-    # pipeline = compile_pipe(pipeline)
+    pipeline = compile_pipe(pipeline)
     load_pipe(pipeline, dir="/home/sandbox/.cache/huggingface/hub/models--RobertML--cached-pipe-02/snapshots/58d70deae87034cce351b780b48841f9746d4ad7")
 
     for _ in range(1):
@@ -64,31 +64,31 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
             return sched._step(*args, **kwargs)
         sched.step = step_and_preview
 
-    image_preview = None
-    @torch.no_grad()
-    def add_final_speed_comparison(pipe, compress):
-        if not hasattr(pipe.vae, "_decode_without_taesd_preview"):
-            pipe.vae._decode_without_taesd_preview = pipe.vae.decode
-        def decode_latents_and_compare_speeds(latents, *args, **kwargs):
-            global image_preview
-            tick_sd = time.time()
-            res_sd = pipe.vae._decode_without_taesd_preview(latents, *args, **kwargs)[0].cpu()
-            tock_sd = tick_taesd = time.time()
-            latents = latents.float()
-            res_taesd = compress(latents.mul(pipe.vae.config.scaling_factor)).mul_(2).sub_(1).cpu()
-            tock_taesd = time.time()
-            # Image.fromarray(latents[0, :3].mul(0.25).add(0.5).clamp(0, 1).mul(255).round().byte().permute(1, 2, 0).cpu().numpy())
-            print(f"SD-VAE (left) takes \033[34m{tock_sd - tick_sd:.4f}s\033[0m; TAESD (right) takes \033[36m{tock_taesd - tick_taesd:.4f}s\033[0m")
-            image_preview = torch.cat([res_sd, res_taesd], -1)
-            return (res_sd,)
-        pipe.vae.decode = decode_latents_and_compare_speeds
+    # image_preview = None
+    # @torch.no_grad()
+    # def add_final_speed_comparison(pipe, compress):
+    #     if not hasattr(pipe.vae, "_decode_without_taesd_preview"):
+    #         pipe.vae._decode_without_taesd_preview = pipe.vae.decode
+    #     def decode_latents_and_compare_speeds(latents, *args, **kwargs):
+    #         global image_preview
+    #         tick_sd = time.time()
+    #         res_sd = pipe.vae._decode_without_taesd_preview(latents, *args, **kwargs)[0].cpu()
+    #         tock_sd = tick_taesd = time.time()
+    #         latents = latents.float()
+    #         res_taesd = compress(latents.mul(pipe.vae.config.scaling_factor)).mul_(2).sub_(1).cpu()
+    #         tock_taesd = time.time()
+    #         # Image.fromarray(latents[0, :3].mul(0.25).add(0.5).clamp(0, 1).mul(255).round().byte().permute(1, 2, 0).cpu().numpy())
+    #         print(f"SD-VAE (left) takes \033[34m{tock_sd - tick_sd:.4f}s\033[0m; TAESD (right) takes \033[36m{tock_taesd - tick_taesd:.4f}s\033[0m")
+    #         image_preview = torch.cat([res_sd, res_taesd], -1)
+    #         return (res_sd,)
+    #     pipe.vae.decode = decode_latents_and_compare_speeds
     
-    add_final_speed_comparison(pipeline, instance)
+    # add_final_speed_comparison(pipeline, instance)
     add_taesd_previewing(pipeline, instance)
     decode = pipeline(prompt="telestereography, unstrengthen, preadministrator, copatroness, hyperpersonal, paramountness, paranoid, guaniferous", output_type="pil", num_inference_steps=20,guidance_scale = 5.0).images[0]
     preview_images_new[0].save("../preview_images_new_1.gif", save_all=True, append_images=preview_images_new[1:], duration=100, loop=0)
     # preview_images_original[0].save("../preview_images_original_1.gif", save_all=True, append_images=preview_images_original[1:], duration=100, loop=0)
-    image_preview.save("../preview_image_decoded_1.jpg")
+    # image_preview.save("../preview_image_decoded_1.jpg")
     
     import sys
     sys.exit(1)